|
|
@@ -1,6 +1,15 @@
|
|
|
+use actix_web::{HttpResponse, web::Bytes};
|
|
|
+use actix_web_lab::sse;
|
|
|
use reqwest::Client;
|
|
|
use serde::Serialize;
|
|
|
-use serde_json::{json, Value};
|
|
|
+use serde_json::{Value, json, from_str};
|
|
|
+use futures_util::{StreamExt, TryStreamExt, Stream, stream};
|
|
|
+use tokio::sync::oneshot;
|
|
|
+use std::{
|
|
|
+ convert::Infallible,
|
|
|
+ sync::{Arc, Mutex},
|
|
|
+ pin::Pin
|
|
|
+};
|
|
|
use crate::{
|
|
|
http_error::HttpError,
|
|
|
models::{
|
|
|
@@ -27,9 +36,8 @@ enum PromptRole {
|
|
|
}
|
|
|
|
|
|
pub struct PromptResponse {
|
|
|
- pub message: String,
|
|
|
- pub input_tokens: i32,
|
|
|
- pub output_tokens: i32
|
|
|
+ pub message: oneshot::Receiver<(String, i32, i32)>,
|
|
|
+ pub sse: HttpResponse
|
|
|
}
|
|
|
|
|
|
pub async fn prompt(
|
|
|
@@ -49,7 +57,7 @@ pub async fn prompt(
|
|
|
content: p
|
|
|
});
|
|
|
}
|
|
|
- Ok(send_message(prompt).await?)
|
|
|
+ send_message(prompt).await
|
|
|
}
|
|
|
|
|
|
fn create_system_message(
|
|
|
@@ -112,7 +120,10 @@ async fn send_message(prompt: Vec<PromptMessage>) -> Result<PromptResponse, Http
|
|
|
let body = json!({
|
|
|
"model": "Qwen/Qwen3-32B",
|
|
|
"messages": prompt,
|
|
|
- "reasoning_effort": "none"
|
|
|
+ "reasoning_effort": "none",
|
|
|
+ "temperature": 1.2,
|
|
|
+ "stream": true,
|
|
|
+ "stream_options": {"include_usage": true}
|
|
|
});
|
|
|
|
|
|
let response = client
|
|
|
@@ -124,14 +135,50 @@ async fn send_message(prompt: Vec<PromptMessage>) -> Result<PromptResponse, Http
|
|
|
.await
|
|
|
.map_err(|e| HttpError::InternalError(e.to_string()))?;
|
|
|
|
|
|
- let data: Value = response
|
|
|
- .json()
|
|
|
- .await
|
|
|
- .map_err(|e| HttpError::InternalError(e.to_string()))?;
|
|
|
+ let (tx, rx) = oneshot::channel::<(String, i32, i32)>();
|
|
|
+ let full_text = Arc::new(Mutex::new(String::new()));
|
|
|
+ let acc = full_text.clone();
|
|
|
+
|
|
|
+ let stream = response.bytes_stream().map(move |chunk| -> Result<Bytes, HttpError> {
|
|
|
+ match chunk {
|
|
|
+ Ok(bytes) => {
|
|
|
+ acc.lock().unwrap().push_str(&String::from_utf8_lossy(&bytes));
|
|
|
+ Ok(bytes)
|
|
|
+ },
|
|
|
+ Err(e) => Err(HttpError::InternalError(e.to_string()))
|
|
|
+ }
|
|
|
+ }).chain(stream::once(async move {
|
|
|
+ let raw = full_text.lock().unwrap().clone();
|
|
|
+
|
|
|
+ let content = raw.lines()
|
|
|
+ .filter_map(|l| l.strip_prefix("data: "))
|
|
|
+ .filter(|d| *d != "[DONE]")
|
|
|
+ .filter_map(|d| from_str::<Value>(d).ok())
|
|
|
+ .filter_map(|v| v["choices"][0]["delta"]["content"].as_str().map(str::to_string))
|
|
|
+ .collect::<String>();
|
|
|
+
|
|
|
+ let (input_tokens, output_tokens) = raw.lines()
|
|
|
+ .filter_map(|l| l.strip_prefix("data: "))
|
|
|
+ .filter_map(|d| from_str::<Value>(d).ok())
|
|
|
+ .find_map(|v| {
|
|
|
+ let usage = v.get("usage")?;
|
|
|
+ if usage.is_null() { return None; }
|
|
|
+ Some(usage.clone())
|
|
|
+ })
|
|
|
+ .map(|u| (
|
|
|
+ u["prompt_tokens"].as_i64().unwrap_or(0) as i32,
|
|
|
+ u["completion_tokens"].as_i64().unwrap_or(0) as i32
|
|
|
+ ))
|
|
|
+ .unwrap_or((0, 0));
|
|
|
+
|
|
|
+ let _ = tx.send((content, input_tokens, output_tokens));
|
|
|
+ Ok::<Bytes, HttpError>(Bytes::new())
|
|
|
+ }));
|
|
|
|
|
|
Ok(PromptResponse {
|
|
|
- message: (&data["choices"][0]["message"]["content"]).to_string(),
|
|
|
- input_tokens: data["usage"]["prompt_tokens"].as_i64().unwrap_or(0) as i32,
|
|
|
- output_tokens: data["usage"]["completion_tokens"].as_i64().unwrap_or(0) as i32
|
|
|
+ message: rx,
|
|
|
+ sse: HttpResponse::Ok()
|
|
|
+ .content_type("text/event-stream")
|
|
|
+ .streaming(stream)
|
|
|
})
|
|
|
}
|
