Add the assembler that folds streamed chunks into a completion

Implemented-By: GLM-5.3 (z.ai, default settings) via OpenCode

crates/loopd/src/llama/assemble.rs

@@ -0,0 +1,221 @@
+//! 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(proto::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,
+        })
+    }
+}

crates/loopd/src/llama/mod.rs

@@ -1,5 +1,6 @@
 //! The llama inference server's chat-completions types and the request builder.
 
+pub mod assemble;
 pub mod request;
 
 #[derive(Debug, Clone, PartialEq, Eq)]