fix: streaming

src/providers/ollama/client.rs

@@ -285,41 +285,62 @@ impl OllamaProvider {
             .await?
             .bytes_stream();
 
-        let stream = byte_stream.flat_map(|chunk_result| {
-            let mut out: Vec<Result<super::types::OllamaChatStreamEvent, LlmError>> = Vec::new();
-
-            let chunk = match chunk_result {
-                Ok(b) => b,
-                Err(e) => {
-                    out.push(Err(LlmError::Http(e)));
-                    return futures::stream::iter(out);
-                }
-            };
-
-            for line in chunk.split(|&b| b == b'\n') {
-                if line.is_empty() {
-                    continue;
-                }
-
-                let parsed: ollama::types::OllamaChatResponse = match serde_json::from_slice(line) {
-                    Ok(v) => v,
-                    Err(_) => continue,
+        let stream = byte_stream
+            .flat_map(|chunk_result| {
+                let mut out: Vec<Result<super::types::OllamaChatStreamEvent, LlmError>> =
+                    Vec::new();
+                let chunk = match chunk_result {
+                    Ok(b) => b,
+                    Err(e) => {
+                        tracing::debug!("Error: {:?}", e);
+                        out.push(Err(LlmError::Http(e)));
+                        return futures::stream::iter(out);
+                    }
                 };
 
-                if !parsed.message.content.is_empty() && !parsed.done {
-                    out.push(Ok(super::types::OllamaChatStreamEvent::Token(
-                        parsed.message.content.clone(),
-                    )));
-                }
+                tracing::debug!("Raw line: {:?}", std::str::from_utf8(&chunk));
 
-                if parsed.done {
-                    out.push(Ok(super::types::OllamaChatStreamEvent::Final(parsed)));
-                    return futures::stream::iter(out);
-                }
-            }
+                for line in chunk.split(|&b| b == b'\n') {
+                    if line.is_empty() {
+                        continue;
+                    }
 
-            futures::stream::iter(out)
-        });
+                    let parsed: ollama::types::OllamaChatResponse =
+                        match serde_json::from_slice(line) {
+                            Ok(v) => v,
+                            Err(_) => continue,
+                        };
+
+                    tracing::debug!("Parsed: {:?}", parsed);
+
+                    if !parsed.message.content.is_empty() {
+                        out.push(Ok(super::types::OllamaChatStreamEvent::Token(
+                            parsed.message.content.clone(),
+                        )));
+                    }
+
+                    if parsed.done {
+                        out.push(Ok(super::types::OllamaChatStreamEvent::Final(parsed)));
+                        return futures::stream::iter(out);
+                    }
+                }
+                futures::stream::iter(out)
+            })
+            .scan(String::new(), |acc, event| {
+                let result = match event {
+                    Ok(ollama::types::OllamaChatStreamEvent::Token(ref tok)) => {
+                        acc.push_str(tok);
+                        Some(event)
+                    }
+                    Ok(ollama::types::OllamaChatStreamEvent::Final(mut resp)) => {
+                        tracing::debug!("Acc: {:?}", acc);
+                        resp.message.content = std::mem::take(acc);
+                        Some(Ok(ollama::types::OllamaChatStreamEvent::Final(resp)))
+                    }
+                    Err(_) => Some(event),
+                };
+                futures::future::ready(result)
+            });
 
         Ok(Box::pin(stream))
     }

src/providers/ollama/types.rs

@@ -117,13 +117,13 @@ pub struct OllamaChatResponse {
     pub message: OllamaMessage,
 
     pub done: bool,
-    pub done_reason: Option<String>,
+    pub done_reason: String,
 
-    pub total_duration: Option<u64>,
-    pub load_duration: Option<u64>,
+    pub total_duration: u64,
+    pub load_duration: u64,
 
-    pub prompt_eval_count: Option<u32>,
-    pub eval_count: Option<u32>,
+    pub prompt_eval_count: u32,
+    pub eval_count: u32,
 }
 
 #[derive(Debug)]