boxmaker/crates/loopd/src/llama/assemble.rs

//! Folding the server's streamed chunks into one completion and the `ChatEvent`s they stand for.

use serde::Deserialize;

use super::{ChatEvent, Completion, FinishReason, InferError, Timings};

// These structs parse the server's format, not ours, so unknown fields are ignored on purpose:
// newer servers send fields this code does not use.

#[derive(Deserialize)]
struct Chunk {
    id: Option<String>,
    choices: Vec<Choice>,
    timings: Option<Timings>,
    prompt_progress: Option<PromptProgress>,
}

#[derive(Deserialize)]
struct Choice {
    finish_reason: Option<String>,
    delta: Delta,
}

#[derive(Deserialize)]
struct Delta {
    reasoning_content: Option<String>,
    content: Option<String>,
    tool_calls: Option<Vec<ToolCallPiece>>,
}

#[derive(Deserialize)]
struct ToolCallPiece {
    index: u64,
    id: Option<String>,
    function: Option<FunctionPiece>,
}

#[derive(Deserialize)]
struct FunctionPiece {
    name: Option<String>,
    arguments: Option<String>,
}

#[derive(Deserialize)]
struct PromptProgress {
    total: u64,
    cache: u64,
    processed: u64,
}

/// A tool call being assembled from its pieces; `id` and `name` wait for the first piece
/// that carries them.
#[derive(Default)]
struct PendingCall {
    id: Option<String>,
    name: Option<String>,
    arguments: String,
}

#[derive(Default)]
pub struct Assembler {
    id: Option<String>,
    content: Option<String>,
    reasoning_content: Option<String>,
    reasoning_tokens: u64,
    in_reasoning: bool,
    finish_reason: Option<FinishReason>,
    timings: Timings,
    tool_calls: Vec<PendingCall>,
}

impl Assembler {
    pub fn new() -> Self {
        Self::default()
    }

    /// Takes one `data:` payload (JSON text) and returns the events it stands for, in order.
    pub fn push(&mut self, data: &str) -> Result<Vec<ChatEvent>, InferError> {
        let chunk: Chunk = serde_json::from_str(data)
            .map_err(|e| InferError::Protocol(format!("unreadable chunk: {e}")))?;
        // Timings first: the reasoning-token count below reads the latest one.
        if let Some(t) = chunk.timings {
            self.timings = t;
        }
        if let Some(id) = chunk.id {
            self.id = Some(id);
        }
        let mut events = Vec::new();
        if let Some(p) = chunk.prompt_progress {
            events.push(ChatEvent::Progress {
                total: p.total,
                cache: p.cache,
                processed: p.processed,
            });
        }
        let mut carries_reasoning = false;
        for choice in chunk.choices {
            if let Some(reason) = choice.finish_reason {
                self.finish_reason = Some(match reason.as_str() {
                    "stop" => FinishReason::Stop,
                    "tool_calls" => FinishReason::ToolCalls,
                    "length" => FinishReason::Length,
                    other => {
                        return Err(InferError::Protocol(format!(
                            "unknown finish_reason {other:?}"
                        )));
                    }
                });
            }
            if let Some(text) = choice.delta.reasoning_content
                && !text.is_empty()
            {
                carries_reasoning = true;
                self.in_reasoning = true;
                self.reasoning_content
                    .get_or_insert_with(String::new)
                    .push_str(&text);
                events.push(ChatEvent::Reasoning(text));
            }
            if let Some(text) = choice.delta.content
                && !text.is_empty()
            {
                self.in_reasoning = false;
                self.content.get_or_insert_with(String::new).push_str(&text);
                events.push(ChatEvent::Content(text));
            }
            for piece in choice.delta.tool_calls.unwrap_or_default() {
                self.in_reasoning = false;
                let index = u32::try_from(piece.index).map_err(|_| {
                    InferError::Protocol("tool call index out of range".to_string())
                })?;
                let (name, arguments) = match piece.function {
                    Some(f) => (f.name, f.arguments.unwrap_or_default()),
                    None => (None, String::new()),
                };
                events.push(ChatEvent::ToolCallDelta {
                    index,
                    id: piece.id.clone(),
                    name: name.clone(),
                    arguments: arguments.clone(),
                });
                let slot = usize::try_from(index).map_err(|_| {
                    InferError::Protocol("tool call index out of range".to_string())
                })?;
                if slot == self.tool_calls.len() {
                    self.tool_calls.push(PendingCall {
                        id: piece.id,
                        name,
                        arguments,
                    });
                } else {
                    let Some(call) = self.tool_calls.get_mut(slot) else {
                        return Err(InferError::Protocol(format!(
                            "tool call index {index} skips ahead"
                        )));
                    };
                    if call.id.is_none() {
                        call.id = piece.id;
                    }
                    if call.name.is_none() {
                        call.name = name;
                    }
                    call.arguments.push_str(&arguments);
                }
            }
        }
        if carries_reasoning {
            self.reasoning_tokens = self.timings.predicted_n;
        }
        Ok(events)
    }

    pub fn id(&self) -> Option<&str> {
        self.id.as_deref()
    }

    pub fn reasoning_tokens(&self) -> u64 {
        self.reasoning_tokens
    }

    pub fn in_reasoning(&self) -> bool {
        self.in_reasoning
    }

    pub fn is_finished(&self) -> bool {
        self.finish_reason.is_some()
    }

    pub fn finish(self, thinking_capped: bool) -> Result<Completion, InferError> {
        let finish_reason = self.finish_reason.ok_or(InferError::StreamClosedEarly)?;
        let id = self
            .id
            .ok_or_else(|| InferError::Protocol("no chunk carried a completion id".to_string()))?;
        let mut tool_calls = Vec::new();
        for call in self.tool_calls {
            let (id, name) = match (call.id, call.name) {
                (Some(id), Some(name)) => (id, name),
                _ => {
                    return Err(InferError::Protocol(
                        "a tool call never got an id or a name".to_string(),
                    ));
                }
            };
            tool_calls.push(bxcore::ToolCall {
                id,
                name,
                arguments: call.arguments,
            });
        }
        Ok(Completion {
            id,
            content: self.content,
            reasoning_content: self.reasoning_content,
            tool_calls,
            finish_reason,
            timings: self.timings,
            reasoning_tokens: self.reasoning_tokens,
            thinking_capped,
        })
    }
}