Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
25 kB · 663 lines
Rust
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664//! Shared, bounded upstream streaming. No synthetic text is emitted as a completion.use crate::{ domains::provider::{AgentMessage, ProviderSettingRow, decrypt_secret, provider_setting}, error::ApiError, state::AppState,};use futures_util::StreamExt;use serde_json::{Value, json};use std::time::Duration;use tokio::sync::mpsc;
pub(crate) struct CompletionRequest { pub provider: String, pub model: String, pub reasoning_effort: Option<String>, pub system: String, pub messages: Vec<Value>, pub max_output_tokens: u32,}
pub(crate) fn bounded_history(messages: &[AgentMessage], token_budget: usize) -> Vec<Value> { let mut size = 0; let mut history = Vec::new(); for m in messages.iter().rev() { if !matches!(m.role.as_str(), "user" | "assistant") || m.content.is_empty() { continue; } if m.role == "assistant" && m.metadata .get("status") .and_then(Value::as_str) .is_some_and(|s| s != "complete") { continue; } let mut content = m.content.clone(); if m.role == "user" { if let Some(selection) = m.metadata["context"]["selection"] .as_str() .filter(|s| !s.is_empty()) { content.push_str(&format!("\n\nSelected passage:\n{selection}")); } if let Some(page) = m.metadata["context"]["page"].as_i64() { content.push_str(&format!("\nBook page: {page}")); } } let cost = content.chars().count().div_ceil(3) + 16; if size + cost > token_budget { break; } size += cost; history.push(json!({"role":m.role,"content":content})); } history.reverse(); // A history truncated between a user and assistant starts at the next question. if history.first().is_some_and(|m| m["role"] == "assistant") { history.remove(0); } history}
pub(crate) fn endpoint_url(base: &str, endpoint: &str) -> String { let base = base.trim_end_matches('/'); let versioned = if base.ends_with("/v1") { base.to_owned() } else { format!("{base}/v1") }; format!("{versioned}/{endpoint}")}
pub(crate) fn protocol(setting: &ProviderSettingRow) -> &str { setting .options .get("protocol") .and_then(Value::as_str) .unwrap_or(match setting.provider.as_str() { "openai" => "responses", "anthropic" => "anthropic", _ => "chat", })}
pub(crate) async fn completion_stream( state: AppState, input: CompletionRequest,) -> Result<mpsc::Receiver<Result<String, String>>, ApiError> { let setting = provider_setting(&state, &input.provider).await?; if !setting.enabled { return Err(ApiError::BadRequest( "Enable this provider in Settings first.".into(), )); } if input.model.trim().is_empty() || input.model.len() > 200 { return Err(ApiError::BadRequest("Select a model.".into())); } if let Some(effort) = &input.reasoning_effort { if !matches!( effort.as_str(), "none" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max" ) { return Err(ApiError::BadRequest("Invalid reasoning effort.".into())); } } if input.reasoning_effort.is_some() && protocol(&setting) == "anthropic" && setting.options["reasoning"].as_bool() != Some(true) { return Err(ApiError::BadRequest("This Anthropic provider has no enabled effort control. Use Default reasoning or enable a supported effort parameter in its options.".into())); } let encrypted = setting .api_key_ciphertext .as_deref() .ok_or_else(|| ApiError::BadRequest("Save an API key in Settings first.".into()))?; let api_key = String::from_utf8(decrypt_secret(&state, encrypted)?).map_err(anyhow::Error::from)?; let (tx, rx) = mpsc::channel(32); tokio::spawn(async move { let result = run_upstream(&setting, &api_key, &input, &tx).await; if let Err(error) = result { let _ = tx.send(Err(error)).await; } }); Ok(rx)}
fn request_payload(setting: &ProviderSettingRow, input: &CompletionRequest) -> Value { let max = input.max_output_tokens.clamp(256, 32_000); let mut payload = match protocol(setting) { "responses" => { let messages:Vec<Value>=input.messages.iter().map(|m| { let mut m=m.clone(); if let Some(parts)=m.get_mut("content").and_then(Value::as_array_mut) { for part in parts { if part["type"]=="text" {*part=json!({"type":"input_text","text":part["text"]});} else if part["type"]=="image_url" {*part=json!({"type":"input_image","image_url":part["image_url"]["url"],"detail":"auto"});} } } m }).collect(); json!({"model":input.model,"instructions":input.system,"input":messages,"stream":true,"store":false,"max_output_tokens":max}) } "anthropic" => { let messages:Vec<Value>=input.messages.iter().map(|m| { let mut m=m.clone(); if let Some(parts)=m.get_mut("content").and_then(Value::as_array_mut) { for part in parts { if part["type"]=="image_url" { if let Some(url)=part["image_url"]["url"].as_str() { if let Some((media,data))=url.strip_prefix("data:").and_then(|v|v.split_once(";base64,")) { *part=json!({"type":"image","source":{"type":"base64","media_type":media,"data":data}}); } } } } } m }).collect(); json!({"model":input.model,"system":input.system,"messages":messages,"stream":true,"max_tokens":max}) } _ => { let mut messages = vec![json!({"role":"system","content":input.system})]; messages.extend(input.messages.clone()); json!({"model":input.model,"messages":messages,"stream":true,"max_tokens":max}) } }; if let Some(effort) = &input.reasoning_effort { match protocol(setting) { "responses" => payload["reasoning"] = json!({"effort":effort}), "anthropic" => { // Provider-specific optional controls are explicit, not inferred from model IDs. if setting.options["reasoning"].as_bool() == Some(true) { payload["output_config"] = json!({"effort":effort}); } } _ => payload["reasoning_effort"] = json!(effort), } } if protocol(setting) != "responses" { if let Some(t) = setting.options.get("temperature").and_then(Value::as_f64) { payload["temperature"] = json!(t.clamp(0.0, 2.0)); } } payload}
async fn run_upstream( setting: &ProviderSettingRow, key: &str, input: &CompletionRequest, tx: &mpsc::Sender<Result<String, String>>,) -> Result<(), String> { let p = protocol(setting); let base = setting.base_url.as_deref().unwrap_or(if p == "anthropic" { "https://api.anthropic.com" } else { "https://api.openai.com" }); let endpoint = match p { "anthropic" => "messages", "responses" => "responses", _ => "chat/completions", }; let client = reqwest::Client::builder() .connect_timeout(Duration::from_secs(15)) .timeout(Duration::from_secs(240)) .build() .map_err(|_| "Could not initialize the provider connection.".to_owned())?; let mut request = client .post(endpoint_url(base, endpoint)) .json(&request_payload(setting, input)); request = if p == "anthropic" { request .header("x-api-key", key) .header("anthropic-version", "2023-06-01") } else { request.bearer_auth(key) }; let response = tokio::select! { _=tx.closed()=>return Ok(()), r=request.send()=>r.map_err(|_|"Provider connection failed or timed out.".to_owned())?, }; if !response.status().is_success() { let status = response.status(); let mut chunks = response.bytes_stream(); let mut body = Vec::new(); while body.len() < 64 * 1024 { let chunk = tokio::select! {_=tx.closed()=>return Ok(()),r=chunks.next()=>r}; let Some(chunk) = chunk else { break }; let chunk = chunk.map_err(|_| "Provider rejected the request.".to_owned())?; body.extend_from_slice(&chunk[..chunk.len().min(64 * 1024 - body.len())]); } let payload: Value = serde_json::from_slice(&body).unwrap_or(Value::Null); let message = payload .pointer("/error/message") .or_else(|| payload.get("message")) .and_then(Value::as_str) .unwrap_or("Check the key, model ID, quota and base URL."); return Err(format!( "Provider HTTP {status}: {}", message .replace(key, "[redacted]") .chars() .take(350) .collect::<String>() )); } if !response .headers() .get(reqwest::header::CONTENT_TYPE) .and_then(|v| v.to_str().ok()) .is_some_and(|v| v.contains("text/event-stream")) { // Some compatible endpoints ignore stream=true. Emit their real complete text once. let mut chunks = response.bytes_stream(); let mut body = Vec::new(); while let Some(chunk) = tokio::select! {_=tx.closed()=>return Ok(()),c=chunks.next()=>c} { let chunk = chunk.map_err(|_| "Provider response was interrupted.".to_owned())?; if body.len() + chunk.len() > 1_000_000 { return Err("Response exceeded the output limit.".into()); } body.extend_from_slice(&chunk); } let payload: Value = serde_json::from_slice(&body) .map_err(|_| "Provider returned a malformed response.".to_owned())?; let text = full_response_text(p, &payload).ok_or_else(|| { "The provider returned no assistant text. Check the model and endpoint.".to_owned() })?; if text.trim().is_empty() { return Err("The provider returned no assistant text.".into()); } if tx.send(Ok(text)).await.is_err() { return Ok(()); } if payload["status"] == "incomplete" || payload["stop_reason"] == "max_tokens" || payload["choices"][0]["finish_reason"] == "length" { return Err( "The response reached its output limit; partial text was preserved.".into(), ); } return Ok(()); } let mut bytes = response.bytes_stream(); let mut decoder = SseDecoder::default(); let mut total = 0; let mut terminal = false; let mut output = false; loop { let chunk = tokio::select! {_=tx.closed()=>return Ok(()), c=bytes.next()=>c}; let Some(chunk) = chunk else { break }; let chunk = chunk.map_err(|_| "Provider stream was interrupted.".to_owned())?; let events = decoder.push(&chunk)?; for event in events { match parse_event(p, &event)? { StreamEvent::Text(text) => { total += text.len(); if total > 1_000_000 { return Err("Response exceeded the output limit.".into()); } output |= !text.trim().is_empty(); if tx.send(Ok(text)).await.is_err() { return Ok(()); } } StreamEvent::Complete => terminal = true, StreamEvent::Ignore => {} } } if terminal { break; } } if !terminal { return Err( "Provider closed the stream before completion; partial text was preserved.".into(), ); } if !output { return Err( "The model returned no text. Try another model or a larger output limit.".into(), ); } Ok(())}
fn full_response_text(protocol: &str, payload: &Value) -> Option<String> { if protocol == "chat" { return payload["choices"][0]["message"]["content"] .as_str() .map(ToOwned::to_owned); } let mut texts = Vec::new(); if protocol == "anthropic" { for part in payload["content"].as_array().into_iter().flatten() { if part["type"] == "text" { if let Some(text) = part["text"].as_str() { texts.push(text); } } } } else { for item in payload["output"].as_array().into_iter().flatten() { for part in item["content"].as_array().into_iter().flatten() { if part["type"] == "output_text" { if let Some(text) = part["text"].as_str() { texts.push(text); } } } } } if texts.is_empty() { None } else { Some(texts.join("")) }}
#[derive(Default)]struct SseDecoder { buffer: Vec<u8>, data: Vec<String>,}impl SseDecoder { fn push(&mut self, bytes: &[u8]) -> Result<Vec<String>, String> { self.buffer.extend_from_slice(bytes); if self.buffer.len() > 2_000_000 { return Err("Provider event exceeded the size limit.".into()); } let mut events = Vec::new(); while let Some(end) = self.buffer.iter().position(|b| *b == b'\n') { let raw: Vec<_> = self.buffer.drain(..=end).collect(); let line = std::str::from_utf8(&raw[..raw.len() - 1]) .map_err(|_| "Provider sent invalid UTF-8.".to_owned())? .trim_end_matches('\r'); if line.is_empty() { if !self.data.is_empty() { events.push(self.data.join("\n")); self.data.clear(); } } else if let Some(data) = line.strip_prefix("data:") { self.data .push(data.strip_prefix(' ').unwrap_or(data).to_owned()); } if self.data.iter().map(String::len).sum::<usize>() > 2_000_000 { return Err("Provider event exceeded the size limit.".into()); } } Ok(events) }}
#[derive(Debug, PartialEq)]enum StreamEvent { Text(String), Complete, Ignore,}fn parse_event(protocol: &str, data: &str) -> Result<StreamEvent, String> { if data.trim() == "[DONE]" { return Ok(StreamEvent::Complete); } let event: Value = serde_json::from_str(data) .map_err(|_| "Provider sent a malformed stream event.".to_owned())?; if event.get("error").is_some_and(|e| !e.is_null()) || event["type"] == "error" || event["type"] == "response.failed" { // Never surface upstream text verbatim here: it can contain credential or request echoes. return Err("The provider reported a generation error; partial text was preserved.".into()); } match protocol { "responses" => match event["type"].as_str().unwrap_or("") { "response.output_text.delta" => Ok(StreamEvent::Text( event["delta"].as_str().unwrap_or_default().to_owned(), )), "response.completed" => Ok(StreamEvent::Complete), "response.incomplete" => Err( "The response reached its limit before completion; partial text was preserved." .into(), ), _ => Ok(StreamEvent::Ignore), }, "anthropic" => match event["type"].as_str().unwrap_or("") { "content_block_delta" if event["delta"]["type"] == "text_delta" => { Ok(StreamEvent::Text( event["delta"]["text"] .as_str() .unwrap_or_default() .to_owned(), )) } "message_delta" if event["delta"]["stop_reason"] == "max_tokens" => { Err("The response reached its output limit; partial text was preserved.".into()) } "message_stop" => Ok(StreamEvent::Complete), _ => Ok(StreamEvent::Ignore), }, _ => { if event["choices"][0]["finish_reason"] == "length" { return Err( "The response reached its output limit; partial text was preserved.".into(), ); } if let Some(text) = event["choices"][0]["delta"]["content"].as_str() { return Ok(StreamEvent::Text(text.to_owned())); } Ok(StreamEvent::Ignore) } }}
#[cfg(test)]mod tests { use super::*; #[test] fn urls_never_duplicate_version() { assert_eq!( endpoint_url("https://api.aitunnel.ru/v1/", "chat/completions"), "https://api.aitunnel.ru/v1/chat/completions" ); assert_eq!( endpoint_url("https://api.openai.com", "responses"), "https://api.openai.com/v1/responses" ); } #[test] fn decoder_handles_utf8_boundaries_and_crlf() { let mut d = SseDecoder::default(); let bytes = "data: {\"choices\":[{\"delta\":{\"content\":\"я\"}}]}\r\n\r\n".as_bytes(); let split = bytes.iter().position(|b| *b == 0xd1).unwrap() + 1; assert!(d.push(&bytes[..split]).unwrap().is_empty()); assert_eq!( parse_event("chat", &d.push(&bytes[split..]).unwrap()[0]).unwrap(), StreamEvent::Text("я".into()) ); } #[test] fn all_protocols_extract_real_deltas_and_stop_errors() { assert_eq!( parse_event( "responses", r#"{"type":"response.output_text.delta","delta":"A"}"# ) .unwrap(), StreamEvent::Text("A".into()) ); assert_eq!( parse_event( "anthropic", r#"{"type":"content_block_delta","delta":{"type":"text_delta","text":"B"}}"# ) .unwrap(), StreamEvent::Text("B".into()) ); assert!(parse_event("chat", r#"{"choices":[{"finish_reason":"length"}]}"#).is_err()); assert!(parse_event("responses", r#"{"type":"response.incomplete"}"#).is_err()); } fn test_setting(protocol: &str, base: String) -> ProviderSettingRow { ProviderSettingRow { provider: "test".into(), model: "wrong-default-model".into(), api_key_ciphertext: None, base_url: Some(base), enabled: true, name: "Local test".into(), models: json!(["chosen-model"]), options: json!({"protocol":protocol}), updated_at: chrono::Utc::now(), } } fn test_input() -> CompletionRequest { CompletionRequest { provider: "test".into(), model: "chosen-model".into(), reasoning_effort: Some("high".into()), system: "rules".into(), messages: vec![json!({"role":"user","content":"question"})], max_output_tokens: 2048, } } #[test] fn selected_model_and_effort_reach_each_protocol() { for protocol in ["chat", "responses", "anthropic"] { let mut setting = test_setting(protocol, "http://127.0.0.1".into()); setting.options["reasoning"] = json!(true); let payload = request_payload(&setting, &test_input()); assert_eq!(payload["model"], "chosen-model"); let effort = match protocol { "responses" => &payload["reasoning"]["effort"], "anthropic" => &payload["output_config"]["effort"], _ => &payload["reasoning_effort"], }; assert_eq!(effort, "high"); } } #[tokio::test] async fn streams_real_http_events_and_cancels_on_drop() { use axum::{ Json, Router, response::sse::{Event, Sse}, routing::post, }; let captured = std::sync::Arc::new(tokio::sync::Mutex::new(Value::Null)); let copy = captured.clone(); let app = Router::new().route( "/v1/chat/completions", post(move |Json(input): Json<Value>| { let copy = copy.clone(); async move { *copy.lock().await = input; Sse::new(tokio_stream::iter(vec![ Ok::<_, std::convert::Infallible>( Event::default() .data(r#"{"choices":[{"delta":{"content":"actual "}}]}"#), ), Ok(Event::default() .data(r#"{"choices":[{"delta":{"content":"response"}}]}"#)), Ok(Event::default().data("[DONE]")), ])) } }), ); let app = app.route( "/v1/responses", post(|| async { let first = futures_util::stream::once(async { Ok::<_, std::convert::Infallible>( Event::default() .data(r#"{"type":"response.output_text.delta","delta":"first"}"#), ) }); Sse::new(first.chain(futures_util::stream::pending())) }), ); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); let setting = test_setting("chat", format!("http://{addr}/v1")); let input = test_input(); let (tx, mut rx) = mpsc::channel(32); run_upstream(&setting, "local-dummy-key", &input, &tx) .await .unwrap(); drop(tx); let mut output = String::new(); while let Some(text) = rx.recv().await { output.push_str(&text.unwrap()); } assert_eq!(output, "actual response"); assert_eq!(captured.lock().await["model"], "chosen-model"); let (tx, rx) = mpsc::channel(1); drop(rx); tokio::time::timeout( Duration::from_millis(500), run_upstream(&setting, "local-dummy-key", &input, &tx), ) .await .unwrap() .unwrap(); let setting = test_setting("responses", format!("http://{addr}")); let (tx, mut rx) = mpsc::channel(1); let input = test_input(); let run = tokio::spawn( async move { run_upstream(&setting, "local-dummy-key", &input, &tx).await }, ); assert_eq!(rx.recv().await.unwrap().unwrap(), "first"); drop(rx); tokio::time::timeout(Duration::from_millis(500), run) .await .unwrap() .unwrap() .unwrap(); server.abort(); } #[test] fn complete_json_fallback_and_history_keep_source_context() { assert_eq!( full_response_text( "chat", &json!({"choices":[{"message":{"content":"full text"}}]}) ), Some("full text".into()) ); assert_eq!( full_response_text( "responses", &json!({"output":[{"content":[{"type":"output_text","text":"a"},{"type":"output_text","text":"b"}]}]}) ), Some("ab".into()) ); let m = AgentMessage { id: uuid::Uuid::new_v4(), conversation_id: uuid::Uuid::new_v4(), role: "user".into(), content: "Explain this".into(), metadata: json!({"context":{"selection":"original selection","page":3}}), created_at: chrono::Utc::now(), }; let mut incomplete = m.clone(); incomplete.role = "assistant".into(); incomplete.content = "partial".into(); incomplete.metadata = json!({"status":"cancelled"}); let history = bounded_history(&[m, incomplete], 1000); assert_eq!(history.len(), 1); assert!( history[0]["content"] .as_str() .unwrap() .contains("original selection") ); assert!( history[0]["content"] .as_str() .unwrap() .contains("Book page: 3") ); }}