Parcourir la source

Add more data for SSE response.

Lee Morgan il y a 1 mois
Parent
commit
e3880aa53a
1 fichiers modifiés avec 12 ajouts et 2 suppressions
  1. 12 2
      src/logic/prompt.rs

+ 12 - 2
src/logic/prompt.rs

@@ -4,6 +4,7 @@ use serde::Serialize;
 use serde_json::{Value, json, from_str};
 use futures_util::{StreamExt, stream};
 use tokio::sync::oneshot;
+use chrono::Utc;
 use std::sync::{Arc, Mutex};
 use crate::{
     http_error::HttpError,
@@ -118,7 +119,8 @@ async fn send_message(prompt: Vec<PromptMessage>) -> Result<PromptResponse, Http
         "reasoning_effort": "none",
         "temperature": 0.7,
         "stream": true,
-        "stream_options": {"include_usage": true}
+        "stream_options": {"include_usage": true},
+        "max_tokens": 4096
     });
 
     let response = client
@@ -166,8 +168,16 @@ async fn send_message(prompt: Vec<PromptMessage>) -> Result<PromptResponse, Http
             ))
             .unwrap_or((0, 0));
 
+        let meta = json!({
+            "tokens": input_tokens + (output_tokens * 2),
+            "created_at": Utc::now(),
+            "something": "else"
+        });
+
+        let sse_chunk = format!("event: meta\ndata: {}\n\n", meta.to_string());
+
         let _ = tx.send((content, input_tokens, output_tokens));
-        Ok::<Bytes, HttpError>(Bytes::new())
+        Ok::<Bytes, HttpError>(Bytes::from(sse_chunk))
     }));
 
     Ok(PromptResponse {