diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 57a020c..b2b41c1 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -899,16 +899,16 @@ mod tests { #[async_trait::async_trait] impl LLMProvider for ToolMediaProvider { - async fn chat( + async fn stream( &self, request: ChatCompletionRequest, - ) -> Result> { + ) -> Result { let call_number = { let mut requests = self.requests.lock().unwrap(); requests.push(request); requests.len() }; - Ok(ChatCompletionResponse { + let response = ChatCompletionResponse { id: format!("response-{call_number}"), model: "vision-test".to_string(), content: if call_number == 1 { @@ -917,6 +917,7 @@ mod tests { "image seen".to_string() }, reasoning_content: None, + provider_state: None, tool_calls: if call_number == 1 { vec![ToolCall { id: "call-image".to_string(), @@ -934,7 +935,8 @@ mod tests { cache_read_input_tokens: None, cache_creation_input_tokens: None, }, - }) + }; + Ok(crate::providers::provider_stream_from_response(response)) } fn ptype(&self) -> &str { diff --git a/src/agent/context_compressor.rs b/src/agent/context_compressor.rs index eb0eddf..927ba5e 100644 --- a/src/agent/context_compressor.rs +++ b/src/agent/context_compressor.rs @@ -669,11 +669,11 @@ mod tests { #[async_trait] impl LLMProvider for MockProvider { - async fn chat( + async fn stream( &self, _request: ChatCompletionRequest, - ) -> Result> { - panic!("MockProvider.chat() called - not expected in test") + ) -> Result { + panic!("MockProvider.stream() called - not expected in test") } fn ptype(&self) -> &str { @@ -699,25 +699,28 @@ mod tests { #[async_trait] impl LLMProvider for MockSummarizer { - async fn chat( + async fn stream( &self, _request: ChatCompletionRequest, - ) -> Result> { - Ok(ChatCompletionResponse { - id: "mock".into(), - model: "mock".into(), - content: "[summarized]".into(), - reasoning_content: None, - tool_calls: vec![], - usage: Usage { - prompt_tokens: 0, - completion_tokens: 0, - total_tokens: 0, - cached_tokens: None, - cache_read_input_tokens: None, - cache_creation_input_tokens: None, + ) -> Result { + Ok(crate::providers::provider_stream_from_response( + ChatCompletionResponse { + id: "mock".into(), + model: "mock".into(), + content: "[summarized]".into(), + reasoning_content: None, + provider_state: None, + tool_calls: vec![], + usage: Usage { + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + cached_tokens: None, + cache_read_input_tokens: None, + cache_creation_input_tokens: None, + }, }, - }) + )) } fn ptype(&self) -> &str { diff --git a/src/providers/anthropic.rs b/src/providers/anthropic.rs index 975d68d..131c979 100644 --- a/src/providers/anthropic.rs +++ b/src/providers/anthropic.rs @@ -5,7 +5,10 @@ use std::collections::HashMap; use std::time::Duration; use super::traits::Usage; -use super::{ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, Tool, ToolCall}; +use super::{ + ChatCompletionRequest, ChatCompletionResponse, DynProviderError, LLMProvider, Message, + ProviderStream, Tool, ToolCall, provider_stream_from_response, +}; use crate::bus::message::ContentBlock; use crate::storage::Storage; use std::sync::Arc; @@ -232,10 +235,10 @@ struct AnthropicUsage { #[async_trait] impl LLMProvider for AnthropicProvider { - async fn chat( + async fn stream( &self, request: ChatCompletionRequest, - ) -> Result> { + ) -> Result { let start = std::time::Instant::now(); let url = format!("{}/v1/messages", self.base_url); let max_tokens = request.max_tokens.or(self.max_tokens).unwrap_or(1024); @@ -382,6 +385,7 @@ impl LLMProvider for AnthropicProvider { model: anthropic_resp.model.unwrap_or_default(), content, reasoning_content: reasoning, + provider_state: None, tool_calls, usage: Usage { prompt_tokens: anthropic_resp @@ -427,7 +431,7 @@ impl LLMProvider for AnthropicProvider { .await; } - Ok(response) + Ok(provider_stream_from_response(response)) } fn ptype(&self) -> &str { diff --git a/src/providers/mod.rs b/src/providers/mod.rs index 57f70c7..c86abf4 100644 --- a/src/providers/mod.rs +++ b/src/providers/mod.rs @@ -1,11 +1,16 @@ pub mod anthropic; pub mod openai; +pub mod stream; pub mod traits; pub use self::anthropic::AnthropicProvider; pub use self::openai::OpenAIProvider; use crate::config::LLMProviderConfig; +pub use stream::{ + DynProviderError, FinishReason, ProviderChunk, ProviderStream, ProviderStreamItem, + collect_provider_stream, provider_stream_from_response, +}; pub use traits::{ ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, Tool, ToolCall, ToolFunction, Usage, diff --git a/src/providers/openai.rs b/src/providers/openai.rs index ceab97a..c898844 100644 --- a/src/providers/openai.rs +++ b/src/providers/openai.rs @@ -1,12 +1,15 @@ use async_trait::async_trait; +use futures_util::stream; use reqwest::Client; -use serde::Deserialize; use serde_json::{Value, json}; -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; use std::time::Duration; +use thiserror::Error; -use super::traits::Usage; -use super::{ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, ToolCall}; +use super::{ + ChatCompletionRequest, DynProviderError, FinishReason, LLMProvider, Message, ProviderChunk, + ProviderStream, Usage, +}; use crate::bus::message::ContentBlock; use crate::storage::Storage; use std::sync::Arc; @@ -208,82 +211,411 @@ impl OpenAIProvider { body } + + fn build_stream_request_body(&self, request: &ChatCompletionRequest) -> Value { + let mut body = self.build_request_body(request); + body["stream"] = Value::Bool(true); + body["stream_options"] = json!({ "include_usage": true }); + body + } } -#[derive(Deserialize)] -struct OpenAIResponse { - id: String, - model: String, - choices: Vec, - #[serde(default)] - usage: OpenAIUsage, +#[derive(Debug, Error)] +enum OpenAIStreamError { + #[error("invalid UTF-8 in SSE event: {0}")] + Utf8(#[from] std::string::FromUtf8Error), + #[error("invalid OpenAI-compatible SSE payload: {0}")] + Json(#[from] serde_json::Error), + #[error("OpenAI-compatible stream ended without a finish marker")] + MissingFinish, } -#[derive(Deserialize)] -struct OpenAIChoice { - message: OpenAIMessage, +#[derive(Default)] +struct SseFramer { + buffer: Vec, } -fn null_or_missing_tool_calls<'de, D>(deserializer: D) -> Result, D::Error> -where - D: serde::Deserializer<'de>, -{ - Ok(Option::>::deserialize(deserializer)?.unwrap_or_default()) +impl SseFramer { + fn push(&mut self, bytes: &[u8]) -> Result, OpenAIStreamError> { + self.buffer.extend_from_slice(bytes); + self.drain_frames(false) + } + + fn finish(&mut self) -> Result, OpenAIStreamError> { + self.drain_frames(true) + } + + fn drain_frames(&mut self, finish: bool) -> Result, OpenAIStreamError> { + let mut frames = Vec::new(); + loop { + let delimiter = find_sse_delimiter(&self.buffer); + let Some((position, delimiter_len)) = delimiter else { + break; + }; + let frame = self.buffer.drain(..position).collect::>(); + self.buffer.drain(..delimiter_len); + if let Some(data) = sse_data(frame)? { + frames.push(data); + } + } + if finish && !self.buffer.is_empty() { + let frame = std::mem::take(&mut self.buffer); + if let Some(data) = sse_data(frame)? { + frames.push(data); + } + } + Ok(frames) + } } -#[derive(Deserialize)] -struct OpenAIMessage { - #[serde(default)] - content: Option, - #[serde(default)] - reasoning_content: Option, - #[serde(default, deserialize_with = "null_or_missing_tool_calls")] - tool_calls: Vec, +fn find_sse_delimiter(buffer: &[u8]) -> Option<(usize, usize)> { + let lf = buffer.windows(2).position(|window| window == b"\n\n"); + let crlf = buffer.windows(4).position(|window| window == b"\r\n\r\n"); + match (lf, crlf) { + (Some(left), Some(right)) if left <= right => Some((left, 2)), + (Some(_), Some(right)) => Some((right, 4)), + (Some(position), None) => Some((position, 2)), + (None, Some(position)) => Some((position, 4)), + (None, None) => None, + } } -#[derive(Deserialize)] -struct OpenAIToolCall { - id: String, - #[serde(rename = "function")] - function: OAIFunction, +fn sse_data(frame: Vec) -> Result, OpenAIStreamError> { + let frame = String::from_utf8(frame)?; + let data = frame + .lines() + .filter_map(|line| { + line.strip_prefix("data:") + .map(|value| value.strip_prefix(' ').unwrap_or(value)) + }) + .collect::>() + .join("\n"); + Ok((!data.is_empty()).then_some(data)) } -#[derive(Deserialize)] -struct OAIFunction { - name: String, - arguments: String, +#[derive(Clone, Copy, PartialEq, Eq)] +enum InlineMode { + Text, + Reasoning, } -#[derive(Deserialize, Default)] -struct OpenAIUsage { - #[serde(default)] - prompt_tokens: u32, - #[serde(default)] - completion_tokens: u32, - #[serde(default)] - total_tokens: u32, - #[serde(default)] - cached_tokens: Option, - #[serde(default)] - prompt_tokens_details: Option, +struct InlineReasoningParser { + mode: InlineMode, + pending: String, } -#[derive(Deserialize, Default)] -struct OpenAIPromptTokensDetails { - #[serde(default)] - cached_tokens: Option, +impl Default for InlineReasoningParser { + fn default() -> Self { + Self { + mode: InlineMode::Text, + pending: String::new(), + } + } +} + +impl InlineReasoningParser { + const TAGS: [(&'static str, InlineMode); 4] = [ + ("", InlineMode::Reasoning), + ("", InlineMode::Reasoning), + ("", InlineMode::Text), + ("", InlineMode::Text), + ]; + + fn push(&mut self, delta: &str) -> Vec { + self.pending.push_str(delta); + self.drain(false) + } + + fn finish(&mut self) -> Vec { + self.drain(true) + } + + fn drain(&mut self, finish: bool) -> Vec { + let mut chunks = Vec::new(); + loop { + let next_tag = Self::TAGS + .iter() + .filter_map(|(tag, mode)| self.pending.find(tag).map(|index| (index, *tag, *mode))) + .min_by_key(|(index, _, _)| *index); + if let Some((index, tag, mode)) = next_tag { + let text = self.pending[..index].to_string(); + self.emit_text(text, &mut chunks); + self.pending.drain(..index + tag.len()); + self.mode = mode; + continue; + } + + let retained = if finish { + 0 + } else { + longest_tag_prefix_suffix(&self.pending, &Self::TAGS) + }; + let emit_len = self.pending.len() - retained; + if emit_len > 0 { + let text = self.pending[..emit_len].to_string(); + self.pending.drain(..emit_len); + self.emit_text(text, &mut chunks); + } + break; + } + chunks + } + + fn emit_text(&self, text: String, chunks: &mut Vec) { + if text.is_empty() { + return; + } + chunks.push(match self.mode { + InlineMode::Text => ProviderChunk::Text(text), + InlineMode::Reasoning => ProviderChunk::Reasoning(text), + }); + } +} + +fn longest_tag_prefix_suffix(value: &str, tags: &[(&str, InlineMode)]) -> usize { + let mut best = 0; + for boundary in value + .char_indices() + .map(|(index, _)| index) + .chain([value.len()]) + { + let suffix = &value[boundary..]; + if tags.iter().any(|(tag, _)| tag.starts_with(suffix)) { + best = best.max(suffix.len()); + } + } + best +} + +#[derive(Default)] +struct PartialStreamTool { + id: Option, + name: Option, + started: bool, +} + +#[derive(Default)] +struct OpenAISseDecoder { + framer: SseFramer, + inline_reasoning: InlineReasoningParser, + tools: HashMap, + metadata_emitted: bool, + done_emitted: bool, +} + +impl OpenAISseDecoder { + fn push(&mut self, bytes: &[u8]) -> Result, OpenAIStreamError> { + let frames = self.framer.push(bytes)?; + self.decode_frames(frames) + } + + fn finish(&mut self) -> Result, OpenAIStreamError> { + let frames = self.framer.finish()?; + let mut chunks = self.decode_frames(frames)?; + chunks.extend(self.inline_reasoning.finish()); + if !self.done_emitted { + return Err(OpenAIStreamError::MissingFinish); + } + Ok(chunks) + } + + fn decode_frames( + &mut self, + frames: Vec, + ) -> Result, OpenAIStreamError> { + let mut chunks = Vec::new(); + for data in frames { + if data == "[DONE]" { + chunks.extend(self.inline_reasoning.finish()); + if !self.done_emitted { + chunks.push(ProviderChunk::Done(FinishReason::Stop)); + self.done_emitted = true; + } + continue; + } + let payload: Value = serde_json::from_str(&data)?; + if !self.metadata_emitted { + let id = payload + .get("id") + .and_then(Value::as_str) + .unwrap_or_default(); + let model = payload + .get("model") + .and_then(Value::as_str) + .unwrap_or_default(); + if !id.is_empty() || !model.is_empty() { + chunks.push(ProviderChunk::Metadata { + id: id.to_string(), + model: model.to_string(), + }); + self.metadata_emitted = true; + } + } + if let Some(usage) = payload.get("usage").filter(|value| !value.is_null()) { + chunks.push(ProviderChunk::Usage(parse_openai_usage(usage))); + } + let Some(choice) = payload + .get("choices") + .and_then(Value::as_array) + .and_then(|choices| choices.first()) + else { + continue; + }; + if let Some(delta) = choice.get("delta") { + if let Some(reasoning) = delta + .get("reasoning_content") + .and_then(Value::as_str) + .or_else(|| delta.get("reasoning").and_then(Value::as_str)) + .filter(|value| !value.is_empty()) + { + chunks.push(ProviderChunk::Reasoning(reasoning.to_string())); + } + if let Some(content) = delta + .get("content") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + { + chunks.extend(self.inline_reasoning.push(content)); + } + if let Some(tool_calls) = delta.get("tool_calls").and_then(Value::as_array) { + self.decode_tool_calls(tool_calls, &mut chunks); + } + } + if let Some(reason) = choice.get("finish_reason").and_then(Value::as_str) { + chunks.extend(self.inline_reasoning.finish()); + self.flush_unstarted_tools(&mut chunks); + if !self.done_emitted { + chunks.push(ProviderChunk::Done(FinishReason::from_provider(reason))); + self.done_emitted = true; + } + } + } + Ok(chunks) + } + + fn decode_tool_calls(&mut self, calls: &[Value], chunks: &mut Vec) { + for (fallback_index, call) in calls.iter().enumerate() { + let index = call + .get("index") + .and_then(Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .unwrap_or(fallback_index); + let tool = self.tools.entry(index).or_default(); + if let Some(id) = call.get("id").and_then(Value::as_str) { + tool.id = Some(id.to_string()); + } + let function = call.get("function"); + if let Some(name) = function + .and_then(|value| value.get("name")) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + { + tool.name = Some(name.to_string()); + } + if !tool.started && tool.name.is_some() { + chunks.push(ProviderChunk::ToolCallStart { + index, + id: tool.id.clone(), + name: tool.name.clone(), + }); + tool.started = true; + } + if let Some(arguments) = function + .and_then(|value| value.get("arguments")) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + { + chunks.push(ProviderChunk::ToolCallArguments { + index, + delta: arguments.to_string(), + }); + } + } + } + + fn flush_unstarted_tools(&mut self, chunks: &mut Vec) { + let mut indexes = self.tools.keys().copied().collect::>(); + indexes.sort_unstable(); + for index in indexes { + let tool = self + .tools + .get_mut(&index) + .expect("tool index came from map"); + if !tool.started { + chunks.push(ProviderChunk::ToolCallStart { + index, + id: tool.id.clone(), + name: tool.name.clone(), + }); + tool.started = true; + } + } + } +} + +fn parse_openai_usage(value: &Value) -> Usage { + let direct_cached = value.get("cached_tokens").and_then(Value::as_u64); + let nested_cached = value + .get("prompt_tokens_details") + .and_then(|details| details.get("cached_tokens")) + .and_then(Value::as_u64); + Usage { + prompt_tokens: json_u32(value, "prompt_tokens"), + completion_tokens: json_u32(value, "completion_tokens"), + total_tokens: json_u32(value, "total_tokens"), + cached_tokens: nested_cached + .or(direct_cached) + .and_then(|tokens| u32::try_from(tokens).ok()), + cache_read_input_tokens: None, + cache_creation_input_tokens: None, + } +} + +fn json_u32(value: &Value, key: &str) -> u32 { + value + .get(key) + .and_then(Value::as_u64) + .and_then(|number| u32::try_from(number).ok()) + .unwrap_or_default() +} + +struct OpenAIHttpStream { + response: reqwest::Response, + decoder: OpenAISseDecoder, + pending: VecDeque, + reached_eof: bool, +} + +async fn next_openai_chunk( + mut state: OpenAIHttpStream, +) -> Result, DynProviderError> { + loop { + if let Some(chunk) = state.pending.pop_front() { + return Ok(Some((chunk, state))); + } + if state.reached_eof { + return Ok(None); + } + match state.response.chunk().await? { + Some(bytes) => state.pending.extend(state.decoder.push(&bytes)?), + None => { + state.pending.extend(state.decoder.finish()?); + state.reached_eof = true; + } + } + } } #[async_trait] impl LLMProvider for OpenAIProvider { - async fn chat( + async fn stream( &self, request: ChatCompletionRequest, - ) -> Result> { + ) -> Result { let start = std::time::Instant::now(); let url = format!("{}/chat/completions", self.base_url); - let body = self.build_request_body(&request); + let body = self.build_stream_request_body(&request); // Debug: Log LLM request summary (only in debug builds) #[cfg(debug_assertions)] @@ -339,10 +671,8 @@ impl LLMProvider for OpenAIProvider { })?; let status = resp.status(); - let text = resp.text().await?; - tracing::debug!(status = %status, resp_body = %text, "LLM response"); - if !status.is_success() { + let text = resp.text().await?; let error = format!("API error {}: {}", status, text); tracing::error!( provider = %self.name, @@ -369,93 +699,13 @@ impl LLMProvider for OpenAIProvider { return Err(error.into()); } - let openai_resp: OpenAIResponse = match serde_json::from_str(&text) { - Ok(response) => response, - Err(e) => { - let err_msg = format!("decode error: {} | body: {}", e, &text); - if let Some(ref storage) = self.storage { - let dur = start.elapsed().as_millis() as u64; - if let Err(error) = storage - .append_llm_call( - &self.name, - &self.model_id, - &req_body_str, - Some(&text), - Some(&err_msg), - dur, - ) - .await - { - tracing::warn!("failed to persist LLM call (decode error): {}", error); - } - } - return Err(err_msg.into()); - } + let state = OpenAIHttpStream { + response: resp, + decoder: OpenAISseDecoder::default(), + pending: VecDeque::new(), + reached_eof: false, }; - - let first_choice = openai_resp - .choices - .into_iter() - .next() - .ok_or("no choices in response")?; - - let content = first_choice - .message - .content - .as_ref() - .unwrap_or(&String::new()) - .clone(); - - let tool_calls: Vec = first_choice - .message - .tool_calls - .iter() - .map(|tc| ToolCall { - id: tc.id.clone(), - name: tc.function.name.clone(), - arguments: serde_json::from_str(&tc.function.arguments) - .unwrap_or(serde_json::Value::Null), - }) - .collect(); - - let usage = openai_resp.usage; - let nested_cached_tokens = usage - .prompt_tokens_details - .as_ref() - .and_then(|d| d.cached_tokens); - let cached_tokens = nested_cached_tokens.or(usage.cached_tokens); - let response = ChatCompletionResponse { - id: openai_resp.id, - model: openai_resp.model, - content, - reasoning_content: first_choice.message.reasoning_content, - tool_calls, - usage: Usage { - prompt_tokens: usage.prompt_tokens, - completion_tokens: usage.completion_tokens, - total_tokens: usage.total_tokens, - cached_tokens, - cache_read_input_tokens: None, - cache_creation_input_tokens: None, - }, - }; - - if let Some(ref storage) = self.storage - && let Err(e) = storage - .append_llm_call( - &self.name, - &self.model_id, - &req_body_str, - Some(&text), - None, - start.elapsed().as_millis() as u64, - ) - .await - { - tracing::warn!("failed to persist LLM call: {}", e); - } - - Ok(response) + Ok(Box::pin(stream::try_unfold(state, next_openai_chunk))) } fn ptype(&self) -> &str { @@ -474,7 +724,7 @@ impl LLMProvider for OpenAIProvider { #[cfg(test)] mod tests { use super::*; - use crate::providers::Message; + use crate::providers::{Message, ToolCall}; #[test] fn test_build_request_body_includes_assistant_tool_calls() { @@ -519,6 +769,10 @@ mod tests { tool_calls[0]["function"]["arguments"], "{\"expression\":\"1+1\"}" ); + + let stream_body = provider.build_stream_request_body(&request); + assert_eq!(stream_body["stream"], true); + assert_eq!(stream_body["stream_options"]["include_usage"], true); } #[test] @@ -552,76 +806,113 @@ mod tests { assert_eq!(converted[1]["content"], "second image"); } - #[test] - fn test_decode_response_accepts_null_tool_calls() { - let text = r#"{ - "id": "d21abaa6552741949e2aba76bde59359", - "choices": [{ - "finish_reason": "stop", - "index": 0, - "message": { - "content": "你好！", - "role": "assistant", - "tool_calls": null, - "reasoning_content": "The user sent a greeting." - } - }], - "created": 1781622889, - "model": "mimo-v2.5", - "object": "chat.completion", - "usage": { - "completion_tokens": 65, - "prompt_tokens": 11741, - "total_tokens": 11806, - "completion_tokens_details": {"reasoning_tokens": 40}, - "prompt_tokens_details": {} - } - }"#; - - let response: OpenAIResponse = serde_json::from_str(text).unwrap(); - let message = &response.choices[0].message; - - assert_eq!(message.content.as_deref(), Some("你好！")); - assert_eq!( - message.reasoning_content.as_deref(), - Some("The user sent a greeting.") + #[tokio::test] + async fn sse_decoder_handles_byte_boundaries_reasoning_content_and_usage() { + let input = concat!( + "data: {\"id\":\"r1\",\"model\":\"m1\",\"choices\":[{\"delta\":{\"reasoning_content\":\"why\",\"content\":\"你好\"},\"finish_reason\":null}]}\r\n\r\n", + "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15,\"prompt_tokens_details\":{\"cached_tokens\":3}}}\n\n", + "data: [DONE]\n\n" ); - assert!(message.tool_calls.is_empty()); - assert_eq!(response.usage.total_tokens, 11806); + let mut decoder = OpenAISseDecoder::default(); + let mut chunks = Vec::new(); + for byte in input.as_bytes() { + chunks.extend(decoder.push(std::slice::from_ref(byte)).unwrap()); + } + chunks.extend(decoder.finish().unwrap()); + + let response = crate::providers::collect_provider_stream(Box::pin( + futures_util::stream::iter(chunks.into_iter().map(Ok)), + )) + .await + .unwrap(); + assert_eq!(response.id, "r1"); + assert_eq!(response.model, "m1"); + assert_eq!(response.reasoning_content.as_deref(), Some("why")); + assert_eq!(response.content, "你好"); + assert_eq!(response.usage.total_tokens, 15); + assert_eq!(response.usage.cached_tokens, Some(3)); } #[test] - fn test_decode_response_exposes_cached_tokens() { - let text = r#"{ - "id": "d21abaa6552741949e2aba76bde59359", - "choices": [{ - "finish_reason": "stop", - "index": 0, - "message": { - "content": "你好！", - "role": "assistant", - "tool_calls": null - } - }], - "created": 1781622889, - "model": "mimo-v2.5", - "object": "chat.completion", - "usage": { - "completion_tokens": 65, - "prompt_tokens": 11741, - "total_tokens": 11806, - "prompt_tokens_details": {"cached_tokens": 1200} - } - }"#; - - let response: OpenAIResponse = serde_json::from_str(text).unwrap(); - assert_eq!( - response - .usage - .prompt_tokens_details - .as_ref() - .and_then(|d| d.cached_tokens), - Some(1200) + fn inline_reasoning_tags_may_span_sse_chunks() { + let input = concat!( + "data: {\"choices\":[{\"delta\":{\"content\":\"secretanswer\"},\"finish_reason\":\"stop\"}]}\n\n", + "data: [DONE]\n\n" ); + let mut decoder = OpenAISseDecoder::default(); + let chunks = decoder.push(input.as_bytes()).unwrap(); + assert_eq!( + chunks + .iter() + .filter_map(|chunk| match chunk { + ProviderChunk::Reasoning(value) => Some(value.as_str()), + _ => None, + }) + .collect::(), + "secret" + ); + assert_eq!( + chunks + .iter() + .filter_map(|chunk| match chunk { + ProviderChunk::Text(value) => Some(value.as_str()), + _ => None, + }) + .collect::(), + "answer" + ); + } + + #[test] + fn reasoning_alias_is_used_when_reasoning_content_is_null() { + let input = concat!( + "data: {\"choices\":[{\"delta\":{\"reasoning_content\":null,\"reasoning\":\"alias\"},\"finish_reason\":\"stop\"}]}\n\n", + "data: [DONE]\n\n" + ); + let mut decoder = OpenAISseDecoder::default(); + let chunks = decoder.push(input.as_bytes()).unwrap(); + assert!( + chunks + .iter() + .any(|chunk| matches!(chunk, ProviderChunk::Reasoning(value) if value == "alias")) + ); + } + + #[tokio::test] + async fn tool_call_arguments_are_assembled_across_sse_events() { + let input = concat!( + "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"calculator\",\"arguments\":\"{\\\"expression\\\":\"}}]},\"finish_reason\":null}]}\n\n", + "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"\\\"1+1\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n", + "data: [DONE]\n\n" + ); + let mut decoder = OpenAISseDecoder::default(); + let chunks = decoder.push(input.as_bytes()).unwrap(); + let response = crate::providers::collect_provider_stream(Box::pin( + futures_util::stream::iter(chunks.into_iter().map(Ok)), + )) + .await + .unwrap(); + assert_eq!(response.tool_calls.len(), 1); + assert_eq!(response.tool_calls[0].id, "call_1"); + assert_eq!(response.tool_calls[0].name, "calculator"); + assert_eq!( + response.tool_calls[0].arguments, + serde_json::json!({"expression":"1+1"}) + ); + } + + #[test] + fn missing_finish_marker_is_an_error() { + let mut decoder = OpenAISseDecoder::default(); + decoder + .push(b"data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n") + .unwrap(); + assert!(matches!( + decoder.finish(), + Err(OpenAIStreamError::MissingFinish) + )); } } diff --git a/src/providers/stream.rs b/src/providers/stream.rs new file mode 100644 index 0000000..6de3ed0 --- /dev/null +++ b/src/providers/stream.rs @@ -0,0 +1,225 @@ +use std::collections::BTreeMap; +use std::error::Error; +use std::pin::Pin; + +use futures_util::{Stream, StreamExt, stream}; +use serde::{Deserialize, Serialize}; + +use crate::bus::ProviderReasoningState; + +use super::{ChatCompletionResponse, ToolCall, Usage}; + +pub type DynProviderError = Box; +pub type ProviderStreamItem = Result; +pub type ProviderStream = Pin + Send>>; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FinishReason { + Stop, + ToolCalls, + Length, + ContentFilter, + Other(String), +} + +impl FinishReason { + pub fn from_provider(value: &str) -> Self { + match value { + "stop" | "end_turn" | "stop_sequence" => Self::Stop, + "tool_calls" | "tool_use" => Self::ToolCalls, + "length" | "max_tokens" => Self::Length, + "content_filter" => Self::ContentFilter, + other => Self::Other(other.to_string()), + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ProviderChunk { + Metadata { + id: String, + model: String, + }, + Text(String), + Reasoning(String), + ToolCallStart { + index: usize, + id: Option, + name: Option, + }, + ToolCallArguments { + index: usize, + delta: String, + }, + ProviderState(ProviderReasoningState), + Usage(Usage), + Done(FinishReason), +} + +#[derive(Default)] +struct PartialToolCall { + id: Option, + name: Option, + arguments: String, +} + +pub async fn collect_provider_stream( + mut provider_stream: ProviderStream, +) -> Result { + let mut id = String::new(); + let mut model = String::new(); + let mut content = String::new(); + let mut reasoning_content = String::new(); + let mut provider_state = None; + let mut usage = Usage::default(); + let mut tool_calls = BTreeMap::::new(); + + while let Some(chunk) = provider_stream.next().await { + match chunk? { + ProviderChunk::Metadata { + id: response_id, + model: response_model, + } => { + id = response_id; + model = response_model; + } + ProviderChunk::Text(delta) => content.push_str(&delta), + ProviderChunk::Reasoning(delta) => reasoning_content.push_str(&delta), + ProviderChunk::ToolCallStart { + index, + id: call_id, + name, + } => { + let partial = tool_calls.entry(index).or_default(); + if call_id.is_some() { + partial.id = call_id; + } + if name.is_some() { + partial.name = name; + } + } + ProviderChunk::ToolCallArguments { index, delta } => { + tool_calls + .entry(index) + .or_default() + .arguments + .push_str(&delta); + } + ProviderChunk::ProviderState(state) => provider_state = Some(state), + ProviderChunk::Usage(value) => usage = value, + ProviderChunk::Done(_) => {} + } + } + + let tool_calls = tool_calls + .into_iter() + .map(|(index, partial)| ToolCall { + id: partial.id.unwrap_or_else(|| format!("tool_call_{index}")), + name: partial.name.unwrap_or_default(), + arguments: serde_json::from_str(&partial.arguments).unwrap_or(serde_json::Value::Null), + }) + .collect(); + + Ok(ChatCompletionResponse { + id, + model, + content, + reasoning_content: (!reasoning_content.is_empty()).then_some(reasoning_content), + provider_state, + tool_calls, + usage, + }) +} + +/// Adapt a complete response to the stream-first provider contract. +/// +/// This is intentionally a compatibility bridge for providers while their +/// native streaming parser is implemented; consumers still use one interface. +pub fn provider_stream_from_response(response: ChatCompletionResponse) -> ProviderStream { + let finish_reason = if response.tool_calls.is_empty() { + FinishReason::Stop + } else { + FinishReason::ToolCalls + }; + let mut chunks = vec![ProviderChunk::Metadata { + id: response.id, + model: response.model, + }]; + if let Some(reasoning) = response.reasoning_content { + chunks.push(ProviderChunk::Reasoning(reasoning)); + } + if !response.content.is_empty() { + chunks.push(ProviderChunk::Text(response.content)); + } + for (index, call) in response.tool_calls.into_iter().enumerate() { + chunks.push(ProviderChunk::ToolCallStart { + index, + id: Some(call.id), + name: Some(call.name), + }); + chunks.push(ProviderChunk::ToolCallArguments { + index, + delta: serde_json::to_string(&call.arguments).unwrap_or_else(|_| "null".to_string()), + }); + } + if let Some(state) = response.provider_state { + chunks.push(ProviderChunk::ProviderState(state)); + } + chunks.push(ProviderChunk::Usage(response.usage)); + chunks.push(ProviderChunk::Done(finish_reason)); + Box::pin(stream::iter(chunks.into_iter().map(Ok))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn collector_assembles_interleaved_tool_argument_fragments() { + let chunks = vec![ + ProviderChunk::Metadata { + id: "response".into(), + model: "model".into(), + }, + ProviderChunk::Reasoning("why".into()), + ProviderChunk::Text("answer".into()), + ProviderChunk::ToolCallStart { + index: 1, + id: Some("second".into()), + name: Some("b".into()), + }, + ProviderChunk::ToolCallArguments { + index: 1, + delta: "{\"n\":".into(), + }, + ProviderChunk::ToolCallStart { + index: 0, + id: Some("first".into()), + name: Some("a".into()), + }, + ProviderChunk::ToolCallArguments { + index: 0, + delta: "{}".into(), + }, + ProviderChunk::ToolCallArguments { + index: 1, + delta: "2}".into(), + }, + ProviderChunk::Done(FinishReason::ToolCalls), + ]; + + let response = collect_provider_stream(Box::pin(stream::iter(chunks.into_iter().map(Ok)))) + .await + .unwrap(); + + assert_eq!(response.content, "answer"); + assert_eq!(response.reasoning_content.as_deref(), Some("why")); + assert_eq!(response.tool_calls.len(), 2); + assert_eq!(response.tool_calls[0].id, "first"); + assert_eq!( + response.tool_calls[1].arguments, + serde_json::json!({"n": 2}) + ); + } +} diff --git a/src/providers/traits.rs b/src/providers/traits.rs index f547b65..55d82d0 100644 --- a/src/providers/traits.rs +++ b/src/providers/traits.rs @@ -2,6 +2,8 @@ use crate::bus::message::ContentBlock; use async_trait::async_trait; use serde::{Deserialize, Serialize}; +use super::stream::{DynProviderError, ProviderStream, collect_provider_stream}; + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Message { pub role: String, @@ -101,11 +103,12 @@ pub struct ChatCompletionResponse { pub model: String, pub content: String, pub reasoning_content: Option, + pub provider_state: Option, pub tool_calls: Vec, pub usage: Usage, } -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct Usage { pub prompt_tokens: u32, pub completion_tokens: u32, @@ -120,10 +123,17 @@ pub struct Usage { #[async_trait] pub trait LLMProvider: Send + Sync { + async fn stream( + &self, + request: ChatCompletionRequest, + ) -> Result; + async fn chat( &self, request: ChatCompletionRequest, - ) -> Result>; + ) -> Result { + collect_provider_stream(self.stream(request).await?).await + } fn ptype(&self) -> &str;