diff --git a/Cargo.lock b/Cargo.lock index 638abc4b9ad5e0..1eea7f35c8e3ce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9630,6 +9630,7 @@ dependencies = [ "anyhow", "async-lock", "cloud_llm_client", + "collections", "futures 0.3.32", "gpui_shared_string", "http_client", @@ -10213,6 +10214,7 @@ dependencies = [ "anyhow", "futures 0.3.32", "http_client", + "language_model_core", "schemars 1.0.4", "serde", "serde_json", @@ -10237,6 +10239,7 @@ dependencies = [ "anyhow", "futures 0.3.32", "http_client", + "language_model_core", "schemars 1.0.4", "serde", "serde_json", diff --git a/crates/language_model_core/Cargo.toml b/crates/language_model_core/Cargo.toml index f4a59a3e08dc30..825818f4322aaa 100644 --- a/crates/language_model_core/Cargo.toml +++ b/crates/language_model_core/Cargo.toml @@ -10,12 +10,12 @@ workspace = true [lib] path = "src/language_model_core.rs" -doctest = false [dependencies] anyhow.workspace = true async-lock.workspace = true cloud_llm_client.workspace = true +collections.workspace = true futures.workspace = true gpui_shared_string.workspace = true http_client.workspace = true diff --git a/crates/language_model_core/src/chat_completion.rs b/crates/language_model_core/src/chat_completion.rs new file mode 100644 index 00000000000000..b5c0219647a4ec --- /dev/null +++ b/crates/language_model_core/src/chat_completion.rs @@ -0,0 +1,1632 @@ +//! Wire types and event mapping for streaming responses of OpenAI-compatible +//! Chat Completions APIs. +//! +//! Multiple providers (OpenAI, OpenRouter, LM Studio, llama.cpp, and various +//! OpenAI-compatible proxies) share this format, so the types are deliberately +//! lenient: every field a consumer does not strictly require is optional or +//! defaulted, because real-world providers routinely omit fields or send +//! explicit `null`s where the OpenAI reference implementation would not. + +use crate::util::{fix_streamed_json, parse_tool_arguments}; +use crate::{ + LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelToolUse, + LanguageModelToolUseInput, StopReason, TokenUsage, +}; +use collections::HashMap; +use futures::{Stream, StreamExt}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::pin::Pin; + +/// A single decoded Chat Completions stream chunk: either an event or a +/// provider error envelope (`{"error": {...}}`). +/// +/// The error payload is generic so providers with richer error envelopes +/// (e.g. OpenRouter) can preserve their extra fields. +#[derive(Serialize, Debug)] +#[serde(untagged)] +pub enum ResponseStreamResult { + Ok(ResponseStreamEvent), + Err { error: E }, +} + +/// `#[derive(Deserialize)]` with `#[serde(untagged)]` is avoided here because +/// untagged enums report unhelpful errors ("data did not match any variant") +/// and, with the lenient event type below, would silently swallow error +/// envelopes as empty events. Instead: +/// +/// * A non-null top-level `error` field decodes as an error, taking +/// precedence over any event payload in the same chunk. +/// * A chunk with a `choices` or `usage` field decodes as an event, +/// preserving the underlying deserialization error message on failure. +/// * Anything else (e.g. non-standard error envelopes like `{"detail": ...}` +/// or `{"object": "error", ...}`) is rejected so it surfaces as a stream +/// error instead of being silently dropped as an empty event. +impl<'de, E> Deserialize<'de> for ResponseStreamResult +where + E: serde::de::DeserializeOwned, +{ + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let value = Value::deserialize(deserializer)?; + match value.get("error") { + Some(error) if !error.is_null() => { + let error = E::deserialize(error).map_err(serde::de::Error::custom)?; + Ok(ResponseStreamResult::Err { error }) + } + _ => { + if value.get("choices").is_none() && value.get("usage").is_none() { + return Err(serde::de::Error::custom(format!( + "unrecognized chat completion stream chunk: {value}" + ))); + } + let event = + ResponseStreamEvent::deserialize(&value).map_err(serde::de::Error::custom)?; + Ok(ResponseStreamResult::Ok(event)) + } + } + } +} + +/// The error payload most OpenAI-compatible providers send inside an +/// `{"error": {...}}` envelope. +#[derive(Serialize, Deserialize, Debug)] +pub struct ResponseStreamError { + pub message: String, +} + +#[derive(Serialize, Deserialize, Debug, Default)] +pub struct ResponseStreamEvent { + /// Usage-only chunks from some providers omit `choices` entirely or send + /// an explicit `null` instead of an empty array. + #[serde(default, deserialize_with = "null_as_default")] + pub choices: Vec, + pub usage: Option, +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct ChoiceDelta { + #[serde(default, deserialize_with = "null_as_default")] + pub index: u32, + pub delta: Option, + pub finish_reason: Option, +} + +/// Deserializes a missing field or an explicit `null` as the type's default, +/// honoring this module's leniency guarantee for fields that are not +/// themselves optional. +fn null_as_default<'de, D, T>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, + T: Default + Deserialize<'de>, +{ + Ok(Option::::deserialize(deserializer)?.unwrap_or_default()) +} + +#[derive(Serialize, Deserialize, Debug, Default, Eq, PartialEq)] +pub struct ResponseMessageDelta { + pub content: Option, + /// Reasoning text as sent by OpenRouter and compatible providers. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reasoning: Option, + /// Reasoning text as sent by DeepSeek-style providers and `llama-server` + /// (when started with a reasoning format, e.g. `--reasoning-format deepseek`). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reasoning_content: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, + /// Provider-defined structured reasoning metadata. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reasoning_details: Option, +} + +#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] +pub struct ToolCallChunk { + /// `None` when the provider omits the index, sends `null`, or sends a + /// negative sentinel such as `-1` (#42584). + #[serde( + default, + deserialize_with = "lenient_index", + skip_serializing_if = "Option::is_none" + )] + pub index: Option, + pub id: Option, + + // There is also an optional `type` field that would determine if a + // function is there. Sometimes this streams in with the `function` before + // it streams in the `type` + pub function: Option, +} + +/// Treats an explicit `null` and negative sentinels like `-1` (#42584) as an +/// absent index rather than a malformed chunk. +fn lenient_index<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + Ok(Option::::deserialize(deserializer)?.and_then(|index| usize::try_from(index).ok())) +} + +#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] +pub struct FunctionChunk { + pub name: Option, + pub arguments: Option, + /// Provider-defined metadata required to replay a reasoning tool call. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub thought_signature: Option, +} + +#[derive(Clone, Serialize, Deserialize, Debug, Default)] +pub struct Usage { + pub prompt_tokens: Option, + pub completion_tokens: Option, + pub total_tokens: Option, + /// Prompt-cache usage when reported by the provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prompt_tokens_details: Option, +} + +/// Reports prompt-cache token usage from compatible providers. +#[derive(Clone, Serialize, Deserialize, Debug, Default)] +pub struct PromptTokensDetails { + /// Tokens read from a prompt cache. + pub cached_tokens: Option, + /// Tokens written to a prompt cache. + pub cache_write_tokens: Option, +} + +impl Usage { + /// Converts to a [`TokenUsage`] update, splitting cache reads and writes + /// out of `prompt_tokens` when the provider reports them. + /// + /// Returns `None` unless both `prompt_tokens` and `completion_tokens` are + /// present, because a partial usage object carries no usable totals. + pub fn token_usage(&self) -> Option { + let prompt_tokens = self.prompt_tokens?; + let completion_tokens = self.completion_tokens?; + let details = self.prompt_tokens_details.as_ref(); + let cache_creation_input_tokens = details + .and_then(|details| details.cache_write_tokens) + .unwrap_or(0); + let cache_read_input_tokens = details + .and_then(|details| details.cached_tokens) + .unwrap_or(0); + Some(TokenUsage { + input_tokens: prompt_tokens + .saturating_sub(cache_creation_input_tokens) + .saturating_sub(cache_read_input_tokens), + output_tokens: completion_tokens, + cache_creation_input_tokens, + cache_read_input_tokens, + }) + } +} + +/// Accumulates structured reasoning metadata from compatible providers. +/// +/// Array entries are matched by `index` and then `id`. Fragmented `text`, +/// `summary`, and `data` fields are concatenated while other non-null fields +/// replace their previous values. +/// +/// # Examples +/// +/// ``` +/// use language_model_core::chat_completion::ReasoningDetailsAccumulator; +/// use serde_json::json; +/// +/// let mut accumulator = ReasoningDetailsAccumulator::default(); +/// accumulator.push(json!([{"index": 0, "text": "first "}])); +/// let details = accumulator +/// .push(json!([{"index": 0, "text": "second"}])) +/// .expect("non-empty reasoning details"); +/// +/// assert_eq!(details[0]["text"], "first second"); +/// ``` +#[derive(Debug, Default)] +pub struct ReasoningDetailsAccumulator { + accumulated: Option, +} + +impl ReasoningDetailsAccumulator { + /// Merges `chunk` and returns the updated metadata snapshot. + /// + /// `null` and empty arrays do not replace previously accumulated metadata + /// and return `None`. + pub fn push(&mut self, chunk: Value) -> Option { + match chunk { + Value::Null => None, + Value::Array(chunks) if chunks.is_empty() => None, + Value::Array(chunks) => { + let mut details = match self.accumulated.take() { + Some(Value::Array(details)) => details, + _ => Vec::new(), + }; + for chunk in chunks { + merge_reasoning_detail(&mut details, chunk); + } + let accumulated = Value::Array(details); + self.accumulated = Some(accumulated.clone()); + Some(accumulated) + } + chunk => { + self.accumulated = Some(chunk.clone()); + Some(chunk) + } + } + } +} + +fn merge_reasoning_detail(details: &mut Vec, chunk: Value) { + let index = chunk.get("index").and_then(Value::as_u64); + let target_index = index + .and_then(|index| { + details + .iter() + .position(|detail| detail.get("index").and_then(Value::as_u64) == Some(index)) + }) + .or_else(|| { + let id = chunk.get("id").and_then(Value::as_str)?; + details + .iter() + .position(|detail| detail.get("id").and_then(Value::as_str) == Some(id)) + }); + let Some(target_index) = target_index else { + details.push(chunk); + return; + }; + let (Some(target), Some(chunk)) = (details[target_index].as_object_mut(), chunk.as_object()) + else { + return; + }; + for (key, value) in chunk { + if matches!(key.as_str(), "text" | "summary" | "data") + && let Some(fragment) = value.as_str() + && let Some(existing) = target.get(key).and_then(Value::as_str) + { + target.insert(key.clone(), Value::String(format!("{existing}{fragment}"))); + } else if !value.is_null() { + target.insert(key.clone(), value.clone()); + } + } +} + +#[derive(Default)] +struct RawToolCall { + index: Option, + id: String, + name: String, + arguments: String, + thought_signature: Option, +} + +/// Accumulates streamed tool calls, matching chunks by their ID and index when +/// both are present. Either value alone is unreliable because some providers +/// omit indices (#42584), reuse indices, or repeat IDs across parallel calls. +#[derive(Default)] +struct ToolCallAccumulator { + calls: Vec, + calls_by_index: HashMap, +} + +#[derive(Debug, thiserror::Error)] +#[error( + "cannot unambiguously attribute a tool call chunk to one of the \ + {calls_in_flight} tool calls in progress" +)] +struct AmbiguousToolCallChunk { + calls_in_flight: usize, +} + +impl ToolCallAccumulator { + /// Returns the call `chunk` belongs to, creating it if the chunk starts + /// a new call. Fails rather than guessing (and corrupting arguments) when + /// the available identity matches multiple calls in flight. + fn entry(&mut self, chunk: &ToolCallChunk) -> Result<&mut RawToolCall, AmbiguousToolCallChunk> { + let id = chunk.id.as_deref().filter(|id| !id.is_empty()); + let call = match (id, chunk.index) { + (Some(id), index) => self.call_for_id(id, index)?, + (None, Some(index)) => self.call_for_index(index), + (None, None) => match self.calls.len() { + 0 => self.new_call(), + 1 => 0, + calls_in_flight => return Err(AmbiguousToolCallChunk { calls_in_flight }), + }, + }; + Ok(&mut self.calls[call]) + } + + fn call_for_id( + &mut self, + id: &str, + index: Option, + ) -> Result { + let known_call = match index { + Some(index) => self + .calls + .iter() + .rposition(|call| call.id == id && call.index == Some(index)) + .or_else(|| { + // A new ID continues the call at its index only while that + // call is anonymous; a different ID means the index was + // reused for another parallel call. + let call = self.calls_by_index.get(&index).copied()?; + self.calls[call].id.is_empty().then_some(call) + }) + .or_else(|| { + self.calls + .iter() + .rposition(|call| call.id == id && call.index.is_none()) + }), + None => { + let matching_calls = self.calls.iter().filter(|call| call.id == id).count(); + if matching_calls > 1 { + return Err(AmbiguousToolCallChunk { + calls_in_flight: matching_calls, + }); + } + self.calls + .iter() + .position(|call| call.id == id) + .or_else(|| { + // Some providers send the ID only partway through a call. + (self.calls.len() == 1 && self.calls[0].id.is_empty()).then_some(0) + }) + } + }; + let call = known_call.unwrap_or_else(|| self.new_call()); + if self.calls[call].id.is_empty() { + self.calls[call].id = id.to_string(); + } + if let Some(index) = index { + self.calls[call].index = Some(index); + self.calls_by_index.insert(index, call); + } + Ok(call) + } + + fn call_for_index(&mut self, index: usize) -> usize { + match self.calls_by_index.get(&index) { + Some(&call) => call, + None => { + let call = self.new_call(); + self.calls[call].index = Some(index); + self.calls_by_index.insert(index, call); + call + } + } + } + + fn new_call(&mut self) -> usize { + self.calls.push(RawToolCall::default()); + self.calls.len() - 1 + } + + /// Removes and returns every call, in the order they appeared in the stream. + fn drain(&mut self) -> impl Iterator + '_ { + self.calls_by_index.clear(); + self.calls.drain(..) + } +} + +/// Maps a stream of Chat Completions chunks to [`LanguageModelCompletionEvent`]s. +/// +/// This is shared by every provider that speaks an OpenAI-compatible Chat +/// Completions dialect (OpenAI, OpenAI-compatible endpoints, OpenRouter, +/// LM Studio, llama.cpp, Bedrock/Mantle, and others), so its behavior must +/// stay provider-neutral: any provider-specific interpretation belongs in the +/// provider's own adapter before or after this mapping. +pub struct ChatCompletionEventMapper { + tool_calls: ToolCallAccumulator, + tool_call_accumulation_failed: bool, + reasoning_details: ReasoningDetailsAccumulator, +} + +impl ChatCompletionEventMapper { + pub fn new() -> Self { + Self { + tool_calls: ToolCallAccumulator::default(), + tool_call_accumulation_failed: false, + reasoning_details: ReasoningDetailsAccumulator::default(), + } + } + + pub fn map_stream( + mut self, + events: Pin>>>, + ) -> impl Stream> + where + E: Into, + { + events.flat_map(move |event| { + futures::stream::iter(match event { + Ok(event) => self.map_event(event), + Err(error) => vec![Err(error.into())], + }) + }) + } + + pub fn map_event( + &mut self, + event: ResponseStreamEvent, + ) -> Vec> { + let mut events = Vec::new(); + if let Some(token_usage) = event.usage.as_ref().and_then(|usage| usage.token_usage()) { + events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage))); + } + + let Some(choice) = event.choices.first() else { + return events; + }; + + if let Some(delta) = choice.delta.as_ref() { + if let Some(reasoning_details) = delta.reasoning_details.clone() + && let Some(reasoning_details) = self.reasoning_details.push(reasoning_details) + { + events.push(Ok(LanguageModelCompletionEvent::ReasoningDetails( + reasoning_details, + ))); + } + if let Some(reasoning) = delta.reasoning.clone() { + push_thinking_event(reasoning, &mut events); + } + if let Some(reasoning_content) = delta.reasoning_content.clone() { + push_thinking_event(reasoning_content, &mut events); + } + if let Some(content) = delta.content.clone() { + if !content.is_empty() { + events.push(Ok(LanguageModelCompletionEvent::Text(content))); + } + } + + if !self.tool_call_accumulation_failed + && let Some(tool_calls) = delta.tool_calls.as_ref() + { + for tool_call in tool_calls { + let entry = match self.tool_calls.entry(tool_call) { + Ok(entry) => entry, + Err(error) => { + self.tool_call_accumulation_failed = true; + self.tool_calls = ToolCallAccumulator::default(); + events.push(Err(anyhow::Error::new(error).into())); + break; + } + }; + + if let Some(function) = tool_call.function.as_ref() { + if let Some(name) = function.name.clone() + && !name.is_empty() + { + entry.name = name; + } + + if let Some(arguments) = function.arguments.clone() { + entry.arguments.push_str(&arguments); + } + + if let Some(thought_signature) = function.thought_signature.clone() { + entry.thought_signature = Some(thought_signature); + } + } + + if !entry.id.is_empty() && !entry.name.is_empty() { + if let Ok(input) = + serde_json::from_str::(&fix_streamed_json(&entry.arguments)) + { + events.push(Ok(LanguageModelCompletionEvent::ToolUse( + LanguageModelToolUse { + id: entry.id.clone().into(), + name: entry.name.as_str().into(), + is_input_complete: false, + input: LanguageModelToolUseInput::Json(input), + raw_input: entry.arguments.clone(), + thought_signature: entry.thought_signature.clone(), + }, + ))); + } + } + } + } + } + + match choice.finish_reason.as_deref() { + Some("stop") => { + events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))); + } + Some("tool_calls") => { + if !self.tool_call_accumulation_failed { + events.extend(self.tool_calls.drain().map( + |tool_call| match parse_tool_arguments(&tool_call.arguments) { + Ok(input) => Ok(LanguageModelCompletionEvent::ToolUse( + LanguageModelToolUse { + id: tool_call.id.clone().into(), + name: tool_call.name.as_str().into(), + is_input_complete: true, + input: LanguageModelToolUseInput::Json(input), + raw_input: tool_call.arguments.clone(), + thought_signature: tool_call.thought_signature.clone(), + }, + )), + Err(error) => Ok(LanguageModelCompletionEvent::ToolUseJsonParseError { + id: tool_call.id.into(), + tool_name: tool_call.name.into(), + raw_input: tool_call.arguments.clone().into(), + json_parse_error: error.to_string(), + }), + }, + )); + } + + events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse))); + } + Some("length") => { + events.push(Ok(LanguageModelCompletionEvent::Stop( + StopReason::MaxTokens, + ))); + } + Some(stop_reason) => { + log::error!("Unexpected chat completion stop_reason: {stop_reason:?}",); + events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))); + } + None => {} + } + + events + } +} + +fn push_thinking_event( + text: String, + events: &mut Vec>, +) { + if !text.is_empty() { + events.push(Ok(LanguageModelCompletionEvent::Thinking { + text, + signature: None, + })); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::LanguageModelToolUse; + use pretty_assertions::assert_eq; + use serde_json::json; + + fn parse(chunk: &str) -> ResponseStreamResult { + serde_json::from_str::(chunk).unwrap() + } + + fn expect_event(chunk: &str) -> ResponseStreamEvent { + match parse(chunk) { + ResponseStreamResult::Ok(event) => event, + ResponseStreamResult::Err { .. } => panic!("expected an event, got an error: {chunk}"), + } + } + + #[test] + fn parses_usage_only_chunk_with_null_prompt_cache_tokens() { + let event = expect_event( + r#"{"choices":[],"usage":{"prompt_tokens":5,"completion_tokens":3,"total_tokens":8,"prompt_tokens_details":{"cached_tokens":0,"cache_write_tokens":null}}}"#, + ); + let details = event.usage.unwrap().prompt_tokens_details.unwrap(); + assert_eq!(details.cached_tokens, Some(0)); + assert_eq!(details.cache_write_tokens, None); + } + + #[test] + fn parses_chunk_without_choices() { + let event = expect_event(r#"{"usage":{"total_tokens":8}}"#); + assert!(event.choices.is_empty()); + assert_eq!(event.usage.unwrap().total_tokens, Some(8)); + } + + #[test] + fn parses_explicit_nulls_for_non_optional_fields() { + let event = expect_event( + r#"{"choices":[{"index":null,"delta":{"tool_calls":[{"index":null,"id":"call_1","function":null}]},"finish_reason":null}],"usage":null}"#, + ); + let choice = &event.choices[0]; + assert_eq!(choice.index, 0); + let tool_call = &choice.delta.as_ref().unwrap().tool_calls.as_ref().unwrap()[0]; + assert_eq!(tool_call.index, None); + assert_eq!(tool_call.id.as_deref(), Some("call_1")); + + let event = expect_event(r#"{"choices":null,"usage":{"total_tokens":8}}"#); + assert!(event.choices.is_empty()); + assert_eq!(event.usage.unwrap().total_tokens, Some(8)); + } + + #[test] + fn parses_empty_usage_object() { + let event = expect_event(r#"{"choices":[],"usage":{}}"#); + let usage = event.usage.unwrap(); + assert_eq!(usage.prompt_tokens, None); + assert_eq!(usage.total_tokens, None); + } + + #[test] + fn parses_tool_call_chunk_without_index() { + let event = expect_event( + r#"{"choices":[{"index":0,"delta":{"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"edit_file","arguments":""}}]}}]}"#, + ); + let mut choices = event.choices; + let delta = choices.remove(0).delta.unwrap(); + let tool_call = delta.tool_calls.unwrap().remove(0); + assert_eq!(tool_call.index, None); + assert_eq!(tool_call.id.as_deref(), Some("call_1")); + assert_eq!( + tool_call.function.unwrap().name.as_deref(), + Some("edit_file") + ); + } + + // MiniMax has sent `index: -1` on every tool call chunk (#42584). + #[test] + fn parses_negative_tool_call_index_as_absent() { + let event = expect_event( + r#"{"choices":[{"index":0,"delta":{"tool_calls":[{"index":-1,"id":"call_1","function":{"name":"edit_file","arguments":"{}"}}]}}]}"#, + ); + let tool_call = &event.choices[0] + .delta + .as_ref() + .unwrap() + .tool_calls + .as_ref() + .unwrap()[0]; + assert_eq!(tool_call.index, None); + assert_eq!(tool_call.id.as_deref(), Some("call_1")); + } + + #[test] + fn parses_error_envelope() { + match parse(r#"{"error":{"message":"quota exceeded"}}"#) { + ResponseStreamResult::Err { error } => assert_eq!(error.message, "quota exceeded"), + ResponseStreamResult::Ok(_) => panic!("expected an error"), + } + } + + #[test] + fn parses_custom_error_payload() { + #[derive(Deserialize)] + struct CustomError { + code: u16, + message: String, + } + + let chunk = r#"{"error":{"code":429,"message":"slow down"}}"#; + match serde_json::from_str::>(chunk).unwrap() { + ResponseStreamResult::Err { error } => { + assert_eq!(error.code, 429); + assert_eq!(error.message, "slow down"); + } + ResponseStreamResult::Ok(_) => panic!("expected an error"), + } + } + + #[test] + fn error_takes_precedence_over_event_payload() { + match parse(r#"{"error":{"message":"boom"},"choices":[{"delta":{"content":"hi"}}]}"#) { + ResponseStreamResult::Err { error } => assert_eq!(error.message, "boom"), + ResponseStreamResult::Ok(_) => panic!("expected an error"), + } + } + + #[test] + fn rejects_unrecognized_chunks() { + for chunk in [ + r#"{"detail":"Internal Server Error"}"#, + r#"{"object":"error","message":"engine overloaded","code":50302}"#, + r#""catastrophe""#, + ] { + let error = serde_json::from_str::(chunk).unwrap_err(); + assert!( + error.to_string().contains("unrecognized"), + "expected {chunk} to be rejected, got: {error}" + ); + } + } + + #[test] + fn null_error_field_is_not_an_error() { + let event = expect_event(r#"{"error":null,"choices":[{"delta":{"content":"hi"}}]}"#); + assert_eq!( + event.choices[0].delta.as_ref().unwrap().content.as_deref(), + Some("hi") + ); + } + + #[test] + fn token_usage_requires_both_totals() { + let usage = Usage { + prompt_tokens: Some(10), + completion_tokens: None, + ..Default::default() + }; + assert!(usage.token_usage().is_none()); + } + + #[test] + fn token_usage_splits_cache_tokens_out_of_prompt_tokens() { + let usage = Usage { + prompt_tokens: Some(12), + completion_tokens: Some(7), + total_tokens: Some(19), + prompt_tokens_details: Some(PromptTokensDetails { + cached_tokens: Some(5), + cache_write_tokens: Some(3), + }), + }; + let token_usage = usage.token_usage().unwrap(); + assert_eq!(token_usage.input_tokens, 4); + assert_eq!(token_usage.output_tokens, 7); + assert_eq!(token_usage.cache_creation_input_tokens, 3); + assert_eq!(token_usage.cache_read_input_tokens, 5); + } + + #[test] + fn reports_the_underlying_field_error() { + let error = serde_json::from_str::( + r#"{"choices":[{"index":0,"delta":{"content":42}}]}"#, + ) + .unwrap_err(); + let message = error.to_string(); + assert!( + !message.contains("did not match any variant"), + "expected a specific error, got: {message}" + ); + assert!( + message.contains("invalid type"), + "expected the underlying field error, got: {message}" + ); + } + + fn map_completion_events( + events: Vec, + ) -> Vec { + let mut mapper = ChatCompletionEventMapper::new(); + let mut all_events = Vec::new(); + for event in events { + all_events.extend(mapper.map_event(event)); + } + all_events.into_iter().filter_map(|e| e.ok()).collect() + } + + #[test] + fn stream_maps_reasoning() { + let events = map_completion_events(vec![ResponseStreamEvent { + choices: vec![ChoiceDelta { + index: 0, + delta: Some(ResponseMessageDelta { + content: None, + reasoning: Some("thinking".into()), + tool_calls: None, + reasoning_content: None, + reasoning_details: None, + }), + finish_reason: None, + }], + usage: None, + }]); + + assert_eq!( + events, + vec![LanguageModelCompletionEvent::Thinking { + text: "thinking".into(), + signature: None, + }] + ); + } + + #[test] + fn stream_maps_length_finish_reason_to_max_tokens_stop() { + let events = map_completion_events(vec![ResponseStreamEvent { + choices: vec![ChoiceDelta { + index: 0, + delta: None, + finish_reason: Some("length".into()), + }], + usage: None, + }]); + + assert_eq!( + events, + vec![LanguageModelCompletionEvent::Stop(StopReason::MaxTokens)] + ); + } + + #[test] + fn chunk_without_choices_or_usage_maps_to_no_events() { + let mut mapper = ChatCompletionEventMapper::new(); + assert!(mapper.map_event(ResponseStreamEvent::default()).is_empty()); + } + + #[test] + fn usage_update_precedes_text_and_stop_events_from_the_same_chunk() { + let mut mapper = ChatCompletionEventMapper::new(); + let events = mapper.map_event(ResponseStreamEvent { + choices: vec![ChoiceDelta { + index: 0, + delta: Some(ResponseMessageDelta { + content: Some("Hello!".to_string()), + ..Default::default() + }), + finish_reason: Some("stop".to_string()), + }], + usage: Some(Usage { + prompt_tokens: Some(11), + completion_tokens: Some(7), + total_tokens: Some(18), + prompt_tokens_details: None, + }), + }); + + assert!(matches!( + events.as_slice(), + [ + Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage { + input_tokens: 11, + output_tokens: 7, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + })), + Ok(LanguageModelCompletionEvent::Text(text)), + Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)), + ] if text == "Hello!" + )); + } + + #[test] + fn usage_update_precedes_tool_use_and_stop_events_from_the_same_chunk() { + let mut mapper = ChatCompletionEventMapper::new(); + let events = mapper.map_event(ResponseStreamEvent { + choices: vec![ChoiceDelta { + index: 0, + delta: Some(ResponseMessageDelta { + tool_calls: Some(vec![ToolCallChunk { + index: Some(0), + id: Some("tool-call-id".to_string()), + function: Some(FunctionChunk { + name: Some("test_tool".to_string()), + arguments: Some(r#"{"value":1}"#.to_string()), + thought_signature: None, + }), + }]), + ..Default::default() + }), + finish_reason: Some("tool_calls".to_string()), + }], + usage: Some(Usage { + prompt_tokens: Some(13), + completion_tokens: Some(5), + total_tokens: Some(18), + prompt_tokens_details: None, + }), + }); + + assert!(matches!( + events.as_slice(), + [ + Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage { + input_tokens: 13, + output_tokens: 5, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + })), + Ok(LanguageModelCompletionEvent::ToolUse(LanguageModelToolUse { + is_input_complete: false, + .. + })), + Ok(LanguageModelCompletionEvent::ToolUse(LanguageModelToolUse { + id, + name, + is_input_complete: true, + .. + })), + Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse)), + ] if id.to_string() == "tool-call-id" && name.as_ref() == "test_tool" + )); + } + + #[test] + fn stream_merges_reasoning_details_and_maps_compatible_usage_and_signatures() { + let response_events = serde_json::from_value(json!([ + { + "choices": [{ + "index": 0, + "delta": { + "reasoning_details": [{ + "id": "reasoning-1", + "index": 0, + "type": "reasoning.text", + "text": "first " + }], + "tool_calls": [{ + "index": 0, + "id": "call-1", + "function": { + "name": "search", + "arguments": "{", + "thought_signature": "signature" + } + }] + }, + "finish_reason": null + }], + "usage": null + }, + { + "choices": [{ + "index": 0, + "delta": { + "reasoning_details": [{ + "id": "reasoning-1", + "index": 0, + "text": "second" + }], + "tool_calls": [{ + "index": 0, + "function": { + "arguments": "}" + } + }] + }, + "finish_reason": "tool_calls" + }], + "usage": null + }, + { + "choices": [], + "usage": { + "prompt_tokens": 10000, + "completion_tokens": 500, + "total_tokens": 10500, + "prompt_tokens_details": { + "cached_tokens": 6000, + "cache_write_tokens": 1000 + } + } + } + ])) + .expect("valid compatible Chat Completions events"); + let events = map_completion_events(response_events); + + assert!(events.iter().any(|event| { + matches!( + event, + LanguageModelCompletionEvent::ReasoningDetails(details) + if details[0]["text"] == "first second" + ) + })); + assert!(events.iter().any(|event| { + matches!( + event, + LanguageModelCompletionEvent::ToolUse(tool_use) + if tool_use.is_input_complete + && tool_use.thought_signature.as_deref() == Some("signature") + ) + })); + assert!(events.iter().any(|event| { + matches!( + event, + LanguageModelCompletionEvent::UsageUpdate(TokenUsage { + input_tokens: 3_000, + output_tokens: 500, + cache_creation_input_tokens: 1_000, + cache_read_input_tokens: 6_000, + }) + ) + })); + } + + #[test] + fn reasoning_details_accumulator_replaces_an_incompatible_previous_shape() { + let mut accumulator = ReasoningDetailsAccumulator::default(); + assert_eq!( + accumulator.push(json!({"summary": "provider-defined"})), + Some(json!({"summary": "provider-defined"})) + ); + assert_eq!( + accumulator.push(json!([{"index": 0, "text": "reasoning"}])), + Some(json!([{"index": 0, "text": "reasoning"}])) + ); + } + + // OpenRouter sends an empty `reasoning_details` array in the finish chunk; + // it must not wipe out details accumulated from earlier chunks. + #[test] + fn reasoning_details_accumulator_ignores_null_and_empty_array_chunks() { + let mut accumulator = ReasoningDetailsAccumulator::default(); + assert!( + accumulator + .push(json!([{"index": 0, "type": "reasoning.text", "text": "thinking"}])) + .is_some() + ); + assert!( + accumulator + .push( + json!([{"index": 0, "type": "reasoning.encrypted", "data": "encrypted-blob"}]) + ) + .is_some() + ); + + assert_eq!(accumulator.push(json!([])), None); + assert_eq!(accumulator.push(serde_json::Value::Null), None); + + let details = accumulator + .push(json!([{"index": 0, "text": " more"}])) + .expect("accumulated reasoning details"); + assert_eq!(details[0]["text"], "thinking more"); + assert_eq!(details[0]["data"], "encrypted-blob"); + } + + #[test] + fn stream_maps_preserves_tool_id_and_name_across_empty_deltas() { + // DashScope sends id="" and name="" in subsequent tool_calls delta + // chunks after the first chunk. ChatCompletionEventMapper must not overwrite + // the accumulated id and name with these empty strings. + + let events = vec![ + // First chunk: id and name are present + ResponseStreamEvent { + choices: vec![ChoiceDelta { + index: 0, + delta: Some(ResponseMessageDelta { + content: None, + reasoning: None, + tool_calls: Some(vec![ToolCallChunk { + index: Some(0), + id: Some("call_dashscope_test".into()), + function: Some(FunctionChunk { + name: Some("list_directory".into()), + arguments: Some("".into()), + thought_signature: None, + }), + }]), + reasoning_content: None, + reasoning_details: None, + }), + finish_reason: None, + }], + usage: None, + }, + // Subsequent chunks: DashScope sends id="" and name="" + ResponseStreamEvent { + choices: vec![ChoiceDelta { + index: 0, + delta: Some(ResponseMessageDelta { + content: None, + reasoning: None, + tool_calls: Some(vec![ToolCallChunk { + index: Some(0), + id: Some("".into()), + function: Some(FunctionChunk { + name: Some("".into()), + arguments: Some("{\"path\": \"".into()), + thought_signature: None, + }), + }]), + reasoning_content: None, + reasoning_details: None, + }), + finish_reason: None, + }], + usage: None, + }, + ResponseStreamEvent { + choices: vec![ChoiceDelta { + index: 0, + delta: Some(ResponseMessageDelta { + content: None, + reasoning: None, + tool_calls: Some(vec![ToolCallChunk { + index: Some(0), + id: Some("".into()), + function: Some(FunctionChunk { + name: Some("".into()), + arguments: Some("blog-scraper\"}".into()), + thought_signature: None, + }), + }]), + reasoning_content: None, + reasoning_details: None, + }), + finish_reason: None, + }], + usage: None, + }, + // Final chunk: finish_reason = "tool_calls" + ResponseStreamEvent { + choices: vec![ChoiceDelta { + index: 0, + delta: None, + finish_reason: Some("tool_calls".into()), + }], + usage: None, + }, + ]; + + let mapped = map_completion_events(events); + + // Events emitted: + // 1. Partial ToolUse from chunk 1 (fix_json("") → "{}", parseable) + // 2. Partial ToolUse from chunk 3 (arguments fully assembled) + // 3. Complete ToolUse from finish_reason="tool_calls" drain + // 4. Stop(ToolUse) + assert_eq!(mapped.len(), 4); + + // Verify the complete ToolUse event (from finish_reason drain) + // has the correct id, name, and accumulated arguments. + let complete_tool_use = mapped.iter().find_map(|event| { + if let LanguageModelCompletionEvent::ToolUse(tool_use) = event { + if tool_use.is_input_complete { + return Some(tool_use); + } + } + None + }); + assert!( + complete_tool_use.is_some(), + "expected a completed ToolUse event" + ); + let tool_use = complete_tool_use.unwrap(); + assert_eq!( + tool_use.id.to_string(), + "call_dashscope_test", + "id must survive empty-string overwrites" + ); + assert_eq!( + tool_use.name.as_ref(), + "list_directory", + "name must survive empty-string overwrites" + ); + assert_eq!( + tool_use.raw_input, "{\"path\": \"blog-scraper\"}", + "arguments should accumulate across chunks" + ); + + // Verify the Stop event + assert!(mapped.iter().any(|event| { + matches!( + event, + LanguageModelCompletionEvent::Stop(StopReason::ToolUse) + ) + })); + } + + fn tool_call_chunk_event(tool_calls: Vec) -> ResponseStreamEvent { + ResponseStreamEvent { + choices: vec![ChoiceDelta { + index: 0, + delta: Some(ResponseMessageDelta { + tool_calls: Some(tool_calls), + ..Default::default() + }), + finish_reason: None, + }], + usage: None, + } + } + + fn finish_tool_calls_event() -> ResponseStreamEvent { + ResponseStreamEvent { + choices: vec![ChoiceDelta { + index: 0, + delta: None, + finish_reason: Some("tool_calls".into()), + }], + usage: None, + } + } + + fn completed_tool_calls(events: &[LanguageModelCompletionEvent]) -> Vec<&LanguageModelToolUse> { + events + .iter() + .filter_map(|event| match event { + LanguageModelCompletionEvent::ToolUse(tool_use) if tool_use.is_input_complete => { + Some(tool_use) + } + _ => None, + }) + .collect() + } + + // Ollama has emitted parallel calls with distinct IDs that both use `index: 0`. + #[test] + fn separates_parallel_tool_calls_sharing_an_index() { + let mapped = map_completion_events(vec![ + tool_call_chunk_event(vec![ToolCallChunk { + index: Some(0), + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: Some("get_weather".into()), + arguments: Some(r#"{"city":"Berlin"}"#.into()), + thought_signature: None, + }), + }]), + tool_call_chunk_event(vec![ToolCallChunk { + index: Some(0), + id: Some("call_b".into()), + function: Some(FunctionChunk { + name: Some("get_weather".into()), + arguments: Some(r#"{"city":"Tokyo"}"#.into()), + thought_signature: None, + }), + }]), + finish_tool_calls_event(), + ]); + + let tool_uses = completed_tool_calls(&mapped); + assert_eq!(tool_uses.len(), 2); + assert_eq!(tool_uses[0].id.to_string(), "call_a"); + assert_eq!(tool_uses[0].raw_input, r#"{"city":"Berlin"}"#); + assert_eq!(tool_uses[1].id.to_string(), "call_b"); + assert_eq!(tool_uses[1].raw_input, r#"{"city":"Tokyo"}"#); + } + + #[test] + fn separates_parallel_tool_calls_sharing_an_id() { + let mapped = map_completion_events(vec![ + tool_call_chunk_event(vec![ToolCallChunk { + index: Some(0), + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: Some("get_weather".into()), + arguments: Some(r#"{"city":"#.into()), + thought_signature: None, + }), + }]), + tool_call_chunk_event(vec![ToolCallChunk { + index: Some(1), + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: Some("get_weather".into()), + arguments: Some(r#"{"city":"#.into()), + thought_signature: None, + }), + }]), + tool_call_chunk_event(vec![ToolCallChunk { + index: Some(0), + id: None, + function: Some(FunctionChunk { + name: None, + arguments: Some(r#""Berlin"}"#.into()), + thought_signature: None, + }), + }]), + tool_call_chunk_event(vec![ToolCallChunk { + index: Some(1), + id: None, + function: Some(FunctionChunk { + name: None, + arguments: Some(r#""Tokyo"}"#.into()), + thought_signature: None, + }), + }]), + finish_tool_calls_event(), + ]); + + let tool_uses = completed_tool_calls(&mapped); + assert_eq!(tool_uses.len(), 2); + assert_eq!(tool_uses[0].raw_input, r#"{"city":"Berlin"}"#); + assert_eq!(tool_uses[1].raw_input, r#"{"city":"Tokyo"}"#); + } + + #[test] + fn rejects_id_only_chunks_for_calls_sharing_an_id() { + let mut mapper = ChatCompletionEventMapper::new(); + for index in [0, 1] { + mapper.map_event(tool_call_chunk_event(vec![ToolCallChunk { + index: Some(index), + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: Some("get_weather".into()), + arguments: Some("{".into()), + thought_signature: None, + }), + }])); + } + + let events = mapper.map_event(tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: None, + arguments: Some("}".into()), + thought_signature: None, + }), + }])); + + match events.as_slice() { + [Err(error)] => { + let message = error.to_string(); + assert!( + message.contains("2 tool calls in progress"), + "expected an attribution error, got: {message}" + ); + } + events => panic!("expected a single error event, got: {events:?}"), + } + } + + #[test] + fn continues_the_latest_call_at_a_reused_index() { + let mapped = map_completion_events(vec![ + tool_call_chunk_event(vec![ToolCallChunk { + index: Some(0), + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: Some("get_weather".into()), + arguments: Some(r#"{"city":"Berlin"}"#.into()), + thought_signature: None, + }), + }]), + tool_call_chunk_event(vec![ToolCallChunk { + index: Some(0), + id: Some("call_b".into()), + function: Some(FunctionChunk { + name: Some("get_weather".into()), + arguments: Some(r#"{"city":"#.into()), + thought_signature: None, + }), + }]), + tool_call_chunk_event(vec![ToolCallChunk { + index: Some(0), + id: None, + function: Some(FunctionChunk { + name: None, + arguments: Some(r#""Tokyo"}"#.into()), + thought_signature: None, + }), + }]), + finish_tool_calls_event(), + ]); + + let tool_uses = completed_tool_calls(&mapped); + assert_eq!(tool_uses.len(), 2); + assert_eq!(tool_uses[0].raw_input, r#"{"city":"Berlin"}"#); + assert_eq!(tool_uses[1].raw_input, r#"{"city":"Tokyo"}"#); + } + + // MiniMax identifies calls by ID alone, without usable indices (#42584). + #[test] + fn accumulates_tool_calls_without_indices_by_id() { + let mapped = map_completion_events(vec![ + tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: Some("list_directory".into()), + arguments: Some(r#"{"path":"#.into()), + thought_signature: None, + }), + }]), + tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: None, + arguments: Some(r#""src"}"#.into()), + thought_signature: None, + }), + }]), + finish_tool_calls_event(), + ]); + + let tool_uses = completed_tool_calls(&mapped); + assert_eq!(tool_uses.len(), 1); + assert_eq!(tool_uses[0].id.to_string(), "call_a"); + assert_eq!(tool_uses[0].raw_input, r#"{"path":"src"}"#); + } + + #[test] + fn continues_the_only_call_for_chunks_without_id_or_index() { + let mapped = map_completion_events(vec![ + tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: Some("list_directory".into()), + arguments: Some(r#"{"path":"#.into()), + thought_signature: None, + }), + }]), + tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: None, + function: Some(FunctionChunk { + name: None, + arguments: Some(r#""src"}"#.into()), + thought_signature: None, + }), + }]), + finish_tool_calls_event(), + ]); + + let tool_uses = completed_tool_calls(&mapped); + assert_eq!(tool_uses.len(), 1); + assert_eq!(tool_uses[0].raw_input, r#"{"path":"src"}"#); + } + + #[test] + fn adopts_a_late_id_for_the_call_at_the_same_index() { + let mapped = map_completion_events(vec![ + tool_call_chunk_event(vec![ToolCallChunk { + index: Some(0), + id: None, + function: Some(FunctionChunk { + name: Some("list_directory".into()), + arguments: Some(r#"{"path":"#.into()), + thought_signature: None, + }), + }]), + tool_call_chunk_event(vec![ToolCallChunk { + index: Some(0), + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: None, + arguments: Some(r#""src"}"#.into()), + thought_signature: None, + }), + }]), + finish_tool_calls_event(), + ]); + + let tool_uses = completed_tool_calls(&mapped); + assert_eq!(tool_uses.len(), 1); + assert_eq!(tool_uses[0].id.to_string(), "call_a"); + assert_eq!(tool_uses[0].raw_input, r#"{"path":"src"}"#); + } + + #[test] + fn adopts_a_late_id_for_the_only_anonymous_call_without_an_index() { + let mapped = map_completion_events(vec![ + tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: None, + function: Some(FunctionChunk { + name: Some("list_directory".into()), + arguments: Some(r#"{"path":"#.into()), + thought_signature: None, + }), + }]), + tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: None, + arguments: Some(r#""src"}"#.into()), + thought_signature: None, + }), + }]), + finish_tool_calls_event(), + ]); + + let tool_uses = completed_tool_calls(&mapped); + assert_eq!(tool_uses.len(), 1); + assert_eq!(tool_uses[0].id.to_string(), "call_a"); + assert_eq!(tool_uses[0].raw_input, r#"{"path":"src"}"#); + } + + #[test] + fn rejects_unattributable_tool_call_chunks() { + let mut mapper = ChatCompletionEventMapper::new(); + for id in ["call_a", "call_b"] { + mapper.map_event(tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: Some(id.into()), + function: Some(FunctionChunk { + name: Some("list_directory".into()), + arguments: Some("{".into()), + thought_signature: None, + }), + }])); + } + + let events = mapper.map_event(tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: None, + function: Some(FunctionChunk { + name: None, + arguments: Some(r#""path":"src"}"#.into()), + thought_signature: None, + }), + }])); + + match events.as_slice() { + [Err(error)] => { + let message = error.to_string(); + assert!( + message.contains("2 tool calls in progress"), + "expected an attribution error, got: {message}" + ); + } + events => panic!("expected a single error event, got: {events:?}"), + } + } + + #[test] + fn does_not_emit_tool_calls_after_an_attribution_error() { + let mut mapper = ChatCompletionEventMapper::new(); + for id in ["call_a", "call_b"] { + mapper.map_event(tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: Some(id.into()), + function: Some(FunctionChunk { + name: Some("list_directory".into()), + arguments: Some("{".into()), + thought_signature: None, + }), + }])); + } + + let error_events = mapper.map_event(tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: None, + function: Some(FunctionChunk { + name: None, + arguments: Some(r#""path":"src"}"#.into()), + thought_signature: None, + }), + }])); + assert!(matches!(error_events.as_slice(), [Err(_)])); + + let continuation_events = mapper.map_event(tool_call_chunk_event(vec![ToolCallChunk { + index: None, + id: Some("call_a".into()), + function: Some(FunctionChunk { + name: None, + arguments: Some("}".into()), + thought_signature: None, + }), + }])); + assert!(continuation_events.is_empty()); + + let finish_events = mapper.map_event(finish_tool_calls_event()); + + assert!(matches!( + finish_events.as_slice(), + [Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse))] + )); + } + + // The `index: -1` stream shape from #42584, end to end. + #[test] + fn maps_parallel_tool_calls_with_negative_indices() { + let response_events = serde_json::from_value(json!([ + { + "choices": [{ + "index": 0, + "delta": { + "tool_calls": [ + { + "index": -1, + "id": "call_a", + "type": "function", + "function": {"name": "get_weather", "arguments": "{\"city\":\"Berlin\"}"} + }, + { + "index": -1, + "id": "call_b", + "type": "function", + "function": {"name": "get_weather", "arguments": "{\"city\":\"Tokyo\"}"} + } + ] + }, + "finish_reason": "tool_calls" + }] + } + ])) + .expect("valid compatible Chat Completions events"); + + let mapped = map_completion_events(response_events); + let tool_uses = completed_tool_calls(&mapped); + assert_eq!(tool_uses.len(), 2); + assert_eq!(tool_uses[0].id.to_string(), "call_a"); + assert_eq!(tool_uses[0].raw_input, r#"{"city":"Berlin"}"#); + assert_eq!(tool_uses[1].id.to_string(), "call_b"); + assert_eq!(tool_uses[1].raw_input, r#"{"city":"Tokyo"}"#); + } +} diff --git a/crates/language_model_core/src/language_model_core.rs b/crates/language_model_core/src/language_model_core.rs index 470844f8f0f472..86f51950ce1f5f 100644 --- a/crates/language_model_core/src/language_model_core.rs +++ b/crates/language_model_core/src/language_model_core.rs @@ -1,3 +1,4 @@ +pub mod chat_completion; mod provider; mod rate_limiter; mod request; diff --git a/crates/language_models/src/provider/bedrock.rs b/crates/language_models/src/provider/bedrock.rs index 5710b7ef92cf2e..d3b5c41efbee48 100644 --- a/crates/language_models/src/provider/bedrock.rs +++ b/crates/language_models/src/provider/bedrock.rs @@ -66,11 +66,14 @@ use util::ResultExt; use crate::AllLanguageModelSettings; use crate::provider::open_ai::{ - ChatCompletionMaxTokensParameter, OpenAiEventMapper, OpenAiResponseEventMapper, into_open_ai, + ChatCompletionMaxTokensParameter, OpenAiResponseEventMapper, into_open_ai, into_open_ai_response, }; +use language_model::chat_completion::{ + ChatCompletionEventMapper, ResponseStreamEvent, ResponseStreamResult, +}; use language_model::util::{fix_streamed_json, parse_tool_arguments}; -use open_ai::{ReasoningEffort, RequestError, ResponseStreamEvent}; +use open_ai::{ReasoningEffort, RequestError}; actions!(bedrock, [Tab, TabPrev]); @@ -1210,22 +1213,10 @@ async fn resolve_mantle_auth( } } -#[derive(Deserialize)] -#[serde(untagged)] -enum MantleChatStreamResult { - Ok(ResponseStreamEvent), - Err { error: MantleChatStreamError }, -} - -#[derive(Deserialize)] -struct MantleChatStreamError { - message: String, -} - fn parse_mantle_chat_stream_line(line: &str) -> Result { - match serde_json::from_str(line) { - Ok(MantleChatStreamResult::Ok(response)) => Ok(response), - Ok(MantleChatStreamResult::Err { error }) => Err(anyhow!(error.message)), + match serde_json::from_str::(line) { + Ok(ResponseStreamResult::Ok(response)) => Ok(response), + Ok(ResponseStreamResult::Err { error }) => Err(anyhow!(error.message)), Err(error) => { log::error!( "Failed to parse Mantle chat completion stream event: `{}`\nResponse: `{}`", @@ -1973,7 +1964,7 @@ impl LanguageModel for BedrockMantleModel { let completions = self.stream_completion(request, cx); let executor = cx.background_executor().clone(); async move { - let mapper = OpenAiEventMapper::new(); + let mapper = ChatCompletionEventMapper::new(); Ok(language_model::stream_in_background( mapper.map_stream(completions.await?).boxed(), executor, diff --git a/crates/language_models/src/provider/llama_cpp.rs b/crates/language_models/src/provider/llama_cpp.rs index 297c1e4284c3dd..3b18b732253b55 100644 --- a/crates/language_models/src/provider/llama_cpp.rs +++ b/crates/language_models/src/provider/llama_cpp.rs @@ -2,18 +2,16 @@ use anyhow::Result; use collections::{HashMap, HashSet}; use credentials_provider::CredentialsProvider; use fs::Fs; -use futures::Stream; use futures::{FutureExt, StreamExt, future::BoxFuture, stream::BoxStream}; use gpui::{App, AsyncApp, Context, Entity, Task, TaskExt}; use http_client::{CustomHeaders, HttpClient}; -use language_model::util::parse_tool_arguments; use language_model::{ ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, InlineDescription, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice, - LanguageModelToolResultContent, LanguageModelToolUse, MessageContent, ProviderSettingsView, - RateLimiter, Role, StopReason, SubPageProviderSettings, TokenUsage, env_var, + LanguageModelToolResultContent, MessageContent, ProviderSettingsView, RateLimiter, Role, + SubPageProviderSettings, env_var, }; use llama_cpp::{ LLAMA_CPP_API_URL, ModelEntry, Props, get_models, get_props, stream_chat_completion, @@ -21,7 +19,6 @@ use llama_cpp::{ }; pub use settings::LlamaCppAvailableModel as AvailableModel; use settings::{Settings, SettingsStore, update_settings_file}; -use std::pin::Pin; use std::sync::LazyLock; use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; use std::time::Duration; @@ -32,6 +29,7 @@ use ui_input::InputField; use util::ResultExt; use crate::AllLanguageModelSettings; +use language_model::chat_completion::ChatCompletionEventMapper; const LLAMA_CPP_DOWNLOAD_URL: &str = "https://llama.app"; const LLAMA_CPP_MODELS_URL: &str = "https://huggingface.co/models?library=gguf&sort=trending"; @@ -942,147 +940,13 @@ impl LanguageModel for LlamaCppLanguageModel { }; let completions = self.stream_completion(request, cx); async move { - let mapper = LlamaCppEventMapper::new(); + let mapper = ChatCompletionEventMapper::new(); Ok(mapper.map_stream(completions.await?).boxed()) } .boxed() } } -struct LlamaCppEventMapper { - tool_calls_by_index: HashMap, -} - -impl LlamaCppEventMapper { - fn new() -> Self { - Self { - tool_calls_by_index: HashMap::default(), - } - } - - pub fn map_stream( - mut self, - events: Pin>>>, - ) -> impl Stream> - { - events.flat_map(move |event| { - futures::stream::iter(match event { - Ok(event) => self.map_event(event), - Err(error) => vec![Err(LanguageModelCompletionError::from(error))], - }) - }) - } - - pub fn map_event( - &mut self, - event: llama_cpp::ResponseStreamEvent, - ) -> Vec> { - let mut events = Vec::new(); - - if let Some(usage) = event.usage { - events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage { - input_tokens: usage.prompt_tokens, - output_tokens: usage.completion_tokens, - cache_creation_input_tokens: 0, - cache_read_input_tokens: 0, - }))); - } - - if let Some(choice) = event.choices.into_iter().next() { - if let Some(reasoning_content) = choice.delta.reasoning_content { - events.push(Ok(LanguageModelCompletionEvent::Thinking { - text: reasoning_content, - signature: None, - })); - } - - if let Some(content) = choice.delta.content { - if !content.is_empty() { - events.push(Ok(LanguageModelCompletionEvent::Text(content))); - } - } - - if let Some(tool_calls) = choice.delta.tool_calls { - for tool_call in tool_calls { - let entry = self.tool_calls_by_index.entry(tool_call.index).or_default(); - - if let Some(tool_id) = tool_call.id { - entry.id = tool_id; - } - - if let Some(function) = tool_call.function { - if let Some(name) = function.name { - // Only the first chunk carries the function name; - // later chunks send an empty name with arguments. - if !name.is_empty() { - entry.name = name; - } - } - - if let Some(arguments) = function.arguments { - entry.arguments.push_str(&arguments); - } - } - } - } - - if let Some(finish_reason) = choice.finish_reason.as_deref() { - match finish_reason { - "stop" => { - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))); - } - "tool_calls" => { - events.extend(self.tool_calls_by_index.drain().map(|(_, tool_call)| { - match parse_tool_arguments(&tool_call.arguments) { - Ok(input) => Ok(LanguageModelCompletionEvent::ToolUse( - LanguageModelToolUse { - id: tool_call.id.into(), - name: tool_call.name.into(), - is_input_complete: true, - input: language_model::LanguageModelToolUseInput::Json( - input, - ), - raw_input: tool_call.arguments, - thought_signature: None, - }, - )), - Err(error) => { - Ok(LanguageModelCompletionEvent::ToolUseJsonParseError { - id: tool_call.id.into(), - tool_name: tool_call.name.into(), - raw_input: tool_call.arguments.into(), - json_parse_error: error.to_string(), - }) - } - } - })); - - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse))); - } - "length" => { - events.push(Ok(LanguageModelCompletionEvent::Stop( - StopReason::MaxTokens, - ))); - } - unexpected => { - log::warn!("Unexpected llama.cpp finish_reason: {unexpected:?}"); - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))); - } - } - } - } - - events - } -} - -#[derive(Default)] -struct RawToolCall { - id: String, - name: String, - arguments: String, -} - fn add_message_content_part( new_part: llama_cpp::MessagePart, role: Role, @@ -1638,6 +1502,7 @@ mod tests { use super::*; use gpui::TestAppContext; use http_client::FakeHttpClient; + use language_model::LanguageModelToolUse; use parking_lot::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -1964,90 +1829,6 @@ mod tests { } } - #[test] - fn usage_event_precedes_stop_event() { - let mut mapper = LlamaCppEventMapper::new(); - let events = mapper.map_event(llama_cpp::ResponseStreamEvent { - model: "test-model".to_string(), - object: "chat.completion.chunk".to_string(), - choices: vec![llama_cpp::ChoiceDelta { - index: 0, - delta: llama_cpp::ResponseMessageDelta { - content: None, - reasoning_content: None, - tool_calls: None, - }, - finish_reason: Some("stop".to_string()), - }], - usage: Some(llama_cpp::Usage { - prompt_tokens: 11, - completion_tokens: 7, - total_tokens: 18, - }), - }); - - assert!(matches!( - events.as_slice(), - [ - Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage { - input_tokens: 11, - output_tokens: 7, - cache_creation_input_tokens: 0, - cache_read_input_tokens: 0, - })), - Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)), - ] - )); - } - - #[test] - fn usage_event_precedes_tool_use_stop_event() { - let mut mapper = LlamaCppEventMapper::new(); - let events = mapper.map_event(llama_cpp::ResponseStreamEvent { - model: "test-model".to_string(), - object: "chat.completion.chunk".to_string(), - choices: vec![llama_cpp::ChoiceDelta { - index: 0, - delta: llama_cpp::ResponseMessageDelta { - content: None, - reasoning_content: None, - tool_calls: Some(vec![llama_cpp::ToolCallChunk { - index: 0, - id: Some("tool-call-id".to_string()), - function: Some(llama_cpp::FunctionChunk { - name: Some("test_tool".to_string()), - arguments: Some(r#"{"value":1}"#.to_string()), - }), - }]), - }, - finish_reason: Some("tool_calls".to_string()), - }], - usage: Some(llama_cpp::Usage { - prompt_tokens: 13, - completion_tokens: 5, - total_tokens: 18, - }), - }); - - assert!(matches!( - events.as_slice(), - [ - Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage { - input_tokens: 13, - output_tokens: 5, - cache_creation_input_tokens: 0, - cache_read_input_tokens: 0, - })), - Ok(LanguageModelCompletionEvent::ToolUse(LanguageModelToolUse { - id, - name, - .. - })), - Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse)), - ] if id.to_string() == "tool-call-id" && name.as_ref() == "test_tool" - )); - } - #[gpui::test] async fn authenticate_fetches_models_after_loading_api_key(cx: &mut TestAppContext) { cx.update(|cx| { diff --git a/crates/language_models/src/provider/lmstudio.rs b/crates/language_models/src/provider/lmstudio.rs index 518b4a26155021..a2de082c1b403d 100644 --- a/crates/language_models/src/provider/lmstudio.rs +++ b/crates/language_models/src/provider/lmstudio.rs @@ -1,14 +1,13 @@ use anyhow::{Result, anyhow}; use credentials_provider::CredentialsProvider; use fs::Fs; -use futures::Stream; use futures::{FutureExt, StreamExt, future::BoxFuture, stream::BoxStream}; use gpui::{App, AsyncApp, Context, Entity, Subscription, Task, TaskExt}; use http_client::{CustomHeaders, HttpClient}; use language_model::{ ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelToolChoice, LanguageModelToolResultContent, - LanguageModelToolUse, MessageContent, StopReason, TokenUsage, env_var, + MessageContent, env_var, }; use language_model::{ InlineDescription, LanguageModelId, LanguageModelName, LanguageModelProvider, @@ -19,17 +18,13 @@ use lmstudio::{LMSTUDIO_API_URL, ModelType, get_models}; pub use settings::LmStudioAvailableModel as AvailableModel; use settings::{Settings, SettingsStore, update_settings_file}; -use std::pin::Pin; use std::sync::LazyLock; -use std::{ - collections::{BTreeMap, HashMap}, - sync::Arc, -}; +use std::{collections::BTreeMap, sync::Arc}; use ui::{ButtonLike, ConfiguredApiCard, Divider, List, ListBulletItem, Tooltip, prelude::*}; use ui_input::InputField; use crate::AllLanguageModelSettings; -use language_model::util::parse_tool_arguments; +use language_model::chat_completion::ChatCompletionEventMapper; const LMSTUDIO_DOWNLOAD_URL: &str = "https://lmstudio.ai/download"; const LMSTUDIO_CATALOG_URL: &str = "https://lmstudio.ai/models"; @@ -559,278 +554,13 @@ impl LanguageModel for LmStudioLanguageModel { }; let completions = self.stream_completion(request, cx); async move { - let mapper = LmStudioEventMapper::new(); + let mapper = ChatCompletionEventMapper::new(); Ok(mapper.map_stream(completions.await?).boxed()) } .boxed() } } -struct LmStudioEventMapper { - tool_calls_by_index: HashMap, -} - -impl LmStudioEventMapper { - fn new() -> Self { - Self { - tool_calls_by_index: HashMap::default(), - } - } - - pub fn map_stream( - mut self, - events: Pin>>>, - ) -> impl Stream> - { - events.flat_map(move |event| { - futures::stream::iter(match event { - Ok(event) => self.map_event(event), - Err(error) => vec![Err(LanguageModelCompletionError::from(error))], - }) - }) - } - - pub fn map_event( - &mut self, - event: lmstudio::ResponseStreamEvent, - ) -> Vec> { - let mut events = Vec::new(); - - if let Some(usage) = event.usage { - events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage { - input_tokens: usage.prompt_tokens, - output_tokens: usage.completion_tokens, - cache_creation_input_tokens: 0, - cache_read_input_tokens: 0, - }))); - } - - // The final usage summary chunk from OpenAI-compatible servers has an empty choices array. - // Return accumulated events instead of treating it as an error. - let Some(choice) = event.choices.into_iter().next() else { - return events; - }; - - if let Some(content) = choice.delta.content { - events.push(Ok(LanguageModelCompletionEvent::Text(content))); - } - - if let Some(reasoning_content) = choice.delta.reasoning_content { - events.push(Ok(LanguageModelCompletionEvent::Thinking { - text: reasoning_content, - signature: None, - })); - } - - if let Some(tool_calls) = choice.delta.tool_calls { - for tool_call in tool_calls { - let entry = self.tool_calls_by_index.entry(tool_call.index).or_default(); - - if let Some(tool_id) = tool_call.id { - entry.id = tool_id; - } - - if let Some(function) = tool_call.function { - if let Some(name) = function.name { - // At the time of writing this code LM Studio (0.3.15) is incompatible with the OpenAI API: - // 1. It sends function name in the first chunk - // 2. It sends empty string in the function name field in all subsequent chunks for arguments - // According to https://platform.openai.com/docs/guides/function-calling?api-mode=responses#streaming - // function name field should be sent only inside the first chunk. - if !name.is_empty() { - entry.name = name; - } - } - - if let Some(arguments) = function.arguments { - entry.arguments.push_str(&arguments); - } - } - } - } - - match choice.finish_reason.as_deref() { - Some("stop") => { - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))); - } - Some("tool_calls") => { - events.extend(self.tool_calls_by_index.drain().map(|(_, tool_call)| { - match parse_tool_arguments(&tool_call.arguments) { - Ok(input) => Ok(LanguageModelCompletionEvent::ToolUse( - LanguageModelToolUse { - id: tool_call.id.into(), - name: tool_call.name.into(), - is_input_complete: true, - input: language_model::LanguageModelToolUseInput::Json(input), - raw_input: tool_call.arguments, - thought_signature: None, - }, - )), - Err(error) => Ok(LanguageModelCompletionEvent::ToolUseJsonParseError { - id: tool_call.id.into(), - tool_name: tool_call.name.into(), - raw_input: tool_call.arguments.into(), - json_parse_error: error.to_string(), - }), - } - })); - - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse))); - } - Some(stop_reason) => { - log::error!("Unexpected LMStudio stop_reason: {stop_reason:?}",); - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))); - } - None => {} - } - - events - } -} - -#[derive(Default)] -struct RawToolCall { - id: String, - name: String, - arguments: String, -} - -#[cfg(test)] -mod tests { - use super::*; - use lmstudio::{ChoiceDelta, ResponseMessageDelta, ResponseStreamEvent, Usage}; - - fn make_event(choices: Vec, usage: Option) -> ResponseStreamEvent { - ResponseStreamEvent { - created: 0, - model: "test-model".to_string(), - object: "chat.completion.chunk".to_string(), - choices, - usage, - } - } - - fn make_content_choice(content: &str) -> ChoiceDelta { - ChoiceDelta { - index: 0, - delta: ResponseMessageDelta { - role: None, - content: Some(content.to_string()), - reasoning_content: None, - tool_calls: None, - }, - finish_reason: None, - } - } - - fn make_stop_choice() -> ChoiceDelta { - ChoiceDelta { - index: 0, - delta: ResponseMessageDelta { - role: None, - content: None, - reasoning_content: None, - tool_calls: None, - }, - finish_reason: Some("stop".to_string()), - } - } - - // OpenAI-compatible servers send a final chunk with usage data and an empty - // choices array. Before this fix, the mapper returned an error for empty - // choices, discarding usage entirely. - #[test] - fn test_usage_in_final_empty_choices_chunk() { - let mut mapper = LmStudioEventMapper::new(); - let event = make_event( - vec![], - Some(Usage { - prompt_tokens: 10, - completion_tokens: 20, - total_tokens: 30, - }), - ); - - let results: Vec<_> = mapper - .map_event(event) - .into_iter() - .map(|r| r.unwrap()) - .collect(); - - assert_eq!( - results, - vec![LanguageModelCompletionEvent::UsageUpdate(TokenUsage { - input_tokens: 10, - output_tokens: 20, - cache_creation_input_tokens: 0, - cache_read_input_tokens: 0, - })] - ); - } - - #[test] - fn test_empty_choices_without_usage_returns_empty() { - let mut mapper = LmStudioEventMapper::new(); - let event = make_event(vec![], None); - - let results = mapper.map_event(event); - - assert!(results.is_empty()); - } - - // Usage data can also arrive in a regular chunk that also contains content. - // Both events must be emitted, with UsageUpdate first. - #[test] - fn test_usage_emitted_alongside_content() { - let mut mapper = LmStudioEventMapper::new(); - let event = make_event( - vec![make_content_choice("Hello!")], - Some(Usage { - prompt_tokens: 5, - completion_tokens: 3, - total_tokens: 8, - }), - ); - - let results: Vec<_> = mapper - .map_event(event) - .into_iter() - .map(|r| r.unwrap()) - .collect(); - - assert_eq!( - results[0], - LanguageModelCompletionEvent::UsageUpdate(TokenUsage { - input_tokens: 5, - output_tokens: 3, - cache_creation_input_tokens: 0, - cache_read_input_tokens: 0, - }) - ); - assert_eq!( - results[1], - LanguageModelCompletionEvent::Text("Hello!".to_string()) - ); - } - - #[test] - fn test_stop_event_emitted_on_finish_reason() { - let mut mapper = LmStudioEventMapper::new(); - let event = make_event(vec![make_stop_choice()], None); - - let results: Vec<_> = mapper - .map_event(event) - .into_iter() - .map(|r| r.unwrap()) - .collect(); - - assert_eq!( - results, - vec![LanguageModelCompletionEvent::Stop(StopReason::EndTurn)] - ); - } -} - fn add_message_content_part( new_part: lmstudio::MessagePart, role: Role, diff --git a/crates/language_models/src/provider/open_ai.rs b/crates/language_models/src/provider/open_ai.rs index 8c123999afe885..328f8005c7b1bf 100644 --- a/crates/language_models/src/provider/open_ai.rs +++ b/crates/language_models/src/provider/open_ai.rs @@ -25,9 +25,10 @@ use std::sync::{Arc, LazyLock}; use strum::IntoEnumIterator; use ui::IconName; +use language_model::chat_completion::ChatCompletionEventMapper; use open_ai::completion::token_usage_from_response_usage; pub use open_ai::completion::{ - ChatCompletionMaxTokensParameter, OpenAiEventMapper, OpenAiResponseEventMapper, into_open_ai, + ChatCompletionMaxTokensParameter, OpenAiResponseEventMapper, into_open_ai, into_open_ai_response, }; @@ -669,7 +670,7 @@ impl LanguageModel for OpenAiLanguageModel { let completions = self.stream_completion(request, cx); let executor = cx.background_executor().clone(); async move { - let mapper = OpenAiEventMapper::new(); + let mapper = ChatCompletionEventMapper::new(); Ok(stream_in_background( mapper.map_stream(completions.await?).boxed(), executor, diff --git a/crates/language_models/src/provider/open_ai_compatible.rs b/crates/language_models/src/provider/open_ai_compatible.rs index 47a109b04db42b..090c348f167d0b 100644 --- a/crates/language_models/src/provider/open_ai_compatible.rs +++ b/crates/language_models/src/provider/open_ai_compatible.rs @@ -3,6 +3,7 @@ use credentials_provider::CredentialsProvider; use futures::{FutureExt, StreamExt, future::BoxFuture}; use gpui::{App, AppContext, AsyncApp, Entity, Task}; use http_client::{CustomHeaders, HttpClient}; +use language_model::chat_completion::ChatCompletionEventMapper; use language_model::{ AuthenticateError, IconOrSvg, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelEffortLevel, LanguageModelId, LanguageModelName, @@ -23,9 +24,7 @@ use crate::provider::api_compatible::{ ApiCompatibleProviderConfigurationView, ApiCompatibleProviderSettings, ApiCompatibleProviderState, }; -use crate::provider::open_ai::{ - OpenAiEventMapper, OpenAiResponseEventMapper, into_open_ai, into_open_ai_response, -}; +use crate::provider::open_ai::{OpenAiResponseEventMapper, into_open_ai, into_open_ai_response}; pub use settings::OpenAiCompatibleAvailableModel as AvailableModel; pub use settings::OpenAiCompatibleModelCapabilities as ModelCapabilities; @@ -436,7 +435,7 @@ impl LanguageModel for OpenAiCompatibleLanguageModel { let completions = self.stream_completion(request, cx); let executor = cx.background_executor().clone(); async move { - let mapper = OpenAiEventMapper::new(); + let mapper = ChatCompletionEventMapper::new(); Ok(language_model::stream_in_background( mapper.map_stream(completions.await?).boxed(), executor, diff --git a/crates/language_models/src/provider/open_router.rs b/crates/language_models/src/provider/open_router.rs index 381ff36bd1d0ab..84b4f7b898d42a 100644 --- a/crates/language_models/src/provider/open_router.rs +++ b/crates/language_models/src/provider/open_router.rs @@ -1,30 +1,27 @@ use anyhow::Result; -use collections::HashMap; + use credentials_provider::CredentialsProvider; -use futures::{FutureExt, Stream, StreamExt, future::BoxFuture}; +use futures::{FutureExt, StreamExt, future::BoxFuture}; use gpui::{App, AppContext, AsyncApp, Context, Entity, SharedString, Task}; use http_client::{CustomHeaders, HttpClient}; +use language_model::chat_completion::ChatCompletionEventMapper; use language_model::{ ApiKeyConfiguration, ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelEffortLevel, LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, - LanguageModelToolChoice, LanguageModelToolResultContent, LanguageModelToolUse, MessageContent, - ProviderSettingsView, RateLimiter, Role, StopReason, TokenUsage, env_var, + LanguageModelToolChoice, LanguageModelToolResultContent, MessageContent, ProviderSettingsView, + RateLimiter, Role, env_var, }; -use open_ai::completion::ReasoningDetailsAccumulator; use open_router::{ Model, ModelMode as OpenRouterModelMode, OPEN_ROUTER_API_URL, ReasoningEffort, ResponseStreamEvent, list_models, }; use settings::{OpenRouterAvailableModel as AvailableModel, Settings, SettingsStore}; use sha2::{Digest as _, Sha256}; -use std::pin::Pin; use std::sync::{Arc, LazyLock}; use ui::IconName; -use language_model::util::{fix_streamed_json, parse_tool_arguments}; - const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("openrouter"); const PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("OpenRouter"); @@ -446,7 +443,7 @@ impl LanguageModel for OpenRouterLanguageModel { let executor = cx.background_executor().clone(); let future = self.request_limiter.stream(async move { let response = request.await?; - let events = OpenRouterEventMapper::new().map_stream(response); + let events = ChatCompletionEventMapper::new().map_stream(response); Ok(language_model::stream_in_background( events.boxed(), executor, @@ -774,373 +771,10 @@ fn add_message_content_part( } } -pub struct OpenRouterEventMapper { - tool_calls_by_index: HashMap, - reasoning_details: ReasoningDetailsAccumulator, -} - -impl OpenRouterEventMapper { - pub fn new() -> Self { - Self { - tool_calls_by_index: HashMap::default(), - reasoning_details: ReasoningDetailsAccumulator::default(), - } - } - - pub fn map_stream( - mut self, - events: Pin< - Box< - dyn Send + Stream>, - >, - >, - ) -> impl Stream> - { - events.flat_map(move |event| { - futures::stream::iter(match event { - Ok(event) => self.map_event(event), - Err(error) => vec![Err(error.into())], - }) - }) - } - - pub fn map_event( - &mut self, - event: ResponseStreamEvent, - ) -> Vec> { - let mut events = Vec::new(); - - if let Some(usage) = event.usage { - let cache_creation_input_tokens = usage - .prompt_tokens_details - .as_ref() - .map_or(0, |details| details.cache_write_tokens); - let cache_read_input_tokens = usage - .prompt_tokens_details - .as_ref() - .map_or(0, |details| details.cached_tokens); - let input_tokens = usage.prompt_tokens.saturating_sub( - cache_creation_input_tokens.saturating_add(cache_read_input_tokens), - ); - - events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage { - input_tokens, - output_tokens: usage.completion_tokens, - cache_creation_input_tokens, - cache_read_input_tokens, - }))); - } - - let Some(choice) = event.choices.first() else { - return events; - }; - - if let Some(details) = choice.delta.reasoning_details.clone() - && let Some(details) = self.reasoning_details.push(details) - { - events.push(Ok(LanguageModelCompletionEvent::ReasoningDetails(details))); - } - - if let Some(reasoning) = choice.delta.reasoning.clone() { - events.push(Ok(LanguageModelCompletionEvent::Thinking { - text: reasoning, - signature: None, - })); - } - - if let Some(content) = choice.delta.content.clone() { - // OpenRouter send empty content string with the reasoning content - // This is a workaround for the OpenRouter API bug - if !content.is_empty() { - events.push(Ok(LanguageModelCompletionEvent::Text(content))); - } - } - - if let Some(tool_calls) = choice.delta.tool_calls.as_ref() { - for tool_call in tool_calls { - let entry = self.tool_calls_by_index.entry(tool_call.index).or_default(); - - if let Some(tool_id) = tool_call.id.clone() { - entry.id = tool_id; - } - - if let Some(function) = tool_call.function.as_ref() { - if let Some(name) = function.name.clone() { - entry.name = name; - } - - if let Some(arguments) = function.arguments.clone() { - entry.arguments.push_str(&arguments); - } - - if let Some(signature) = function.thought_signature.clone() { - entry.thought_signature = Some(signature); - } - } - - if !entry.id.is_empty() && !entry.name.is_empty() { - if let Ok(input) = serde_json::from_str::( - &fix_streamed_json(&entry.arguments), - ) { - events.push(Ok(LanguageModelCompletionEvent::ToolUse( - LanguageModelToolUse { - id: entry.id.clone().into(), - name: entry.name.as_str().into(), - is_input_complete: false, - input: language_model::LanguageModelToolUseInput::Json(input), - raw_input: entry.arguments.clone(), - thought_signature: entry.thought_signature.clone(), - }, - ))); - } - } - } - } - - match choice.finish_reason.as_deref() { - Some("stop") => { - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))); - } - Some("tool_calls") => { - events.extend(self.tool_calls_by_index.drain().map(|(_, tool_call)| { - match parse_tool_arguments(&tool_call.arguments) { - Ok(input) => Ok(LanguageModelCompletionEvent::ToolUse( - LanguageModelToolUse { - id: tool_call.id.clone().into(), - name: tool_call.name.as_str().into(), - is_input_complete: true, - input: language_model::LanguageModelToolUseInput::Json(input), - raw_input: tool_call.arguments.clone(), - thought_signature: tool_call.thought_signature.clone(), - }, - )), - Err(error) => Ok(LanguageModelCompletionEvent::ToolUseJsonParseError { - id: tool_call.id.clone().into(), - tool_name: tool_call.name.as_str().into(), - raw_input: tool_call.arguments.clone().into(), - json_parse_error: error.to_string(), - }), - } - })); - - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse))); - } - Some(stop_reason) => { - log::error!("Unexpected OpenRouter stop_reason: {stop_reason:?}",); - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))); - } - None => {} - } - - events - } -} - -#[derive(Default)] -struct RawToolCall { - id: String, - name: String, - arguments: String, - thought_signature: Option, -} - #[cfg(test)] mod tests { use super::*; - use open_router::{ChoiceDelta, FunctionChunk, ResponseMessageDelta, ToolCallChunk}; - - #[gpui::test] - async fn test_reasoning_details_preservation_with_tool_calls() { - // This test verifies that reasoning_details are properly captured and preserved - // when a model uses tool calling with reasoning/thinking tokens. - // - // The key regression this prevents: - // - OpenRouter sends multiple reasoning_details updates during streaming - // - First with actual content (encrypted reasoning data) - // - Then with empty array on completion - // - We must NOT overwrite the real data with the empty array - - let mut mapper = OpenRouterEventMapper::new(); - - // Simulate the streaming events as they come from OpenRouter/Gemini - let events = vec![ - // Event 1: Initial reasoning details with text - ResponseStreamEvent { - id: Some("response_123".into()), - created: 1234567890, - model: "google/gemini-3.1-pro-preview".into(), - choices: vec![ChoiceDelta { - index: 0, - delta: ResponseMessageDelta { - role: None, - content: None, - reasoning: None, - tool_calls: None, - reasoning_details: Some(serde_json::json!([ - { - "type": "reasoning.text", - "text": "Let me analyze this request...", - "format": "google-gemini-v1", - "index": 0 - } - ])), - }, - finish_reason: None, - }], - usage: None, - }, - // Event 2: More reasoning details - ResponseStreamEvent { - id: Some("response_123".into()), - created: 1234567890, - model: "google/gemini-3.1-pro-preview".into(), - choices: vec![ChoiceDelta { - index: 0, - delta: ResponseMessageDelta { - role: None, - content: None, - reasoning: None, - tool_calls: None, - reasoning_details: Some(serde_json::json!([ - { - "type": "reasoning.encrypted", - "data": "EtgDCtUDAdHtim9OF5jm4aeZSBAtl/randomized123", - "format": "google-gemini-v1", - "index": 0, - "id": "tool_call_abc123" - } - ])), - }, - finish_reason: None, - }], - usage: None, - }, - // Event 3: Tool call starts - ResponseStreamEvent { - id: Some("response_123".into()), - created: 1234567890, - model: "google/gemini-3.1-pro-preview".into(), - choices: vec![ChoiceDelta { - index: 0, - delta: ResponseMessageDelta { - role: None, - content: None, - reasoning: None, - tool_calls: Some(vec![ToolCallChunk { - index: 0, - id: Some("tool_call_abc123".into()), - function: Some(FunctionChunk { - name: Some("list_directory".into()), - arguments: Some("{\"path\":\"test\"}".into()), - thought_signature: Some("sha256:test_signature_xyz789".into()), - }), - }]), - reasoning_details: None, - }, - finish_reason: None, - }], - usage: None, - }, - // Event 4: Empty reasoning_details on tool_calls finish - // This is the critical event - we must not overwrite with this empty array! - ResponseStreamEvent { - id: Some("response_123".into()), - created: 1234567890, - model: "google/gemini-3.1-pro-preview".into(), - choices: vec![ChoiceDelta { - index: 0, - delta: ResponseMessageDelta { - role: None, - content: None, - reasoning: None, - tool_calls: None, - reasoning_details: Some(serde_json::json!([])), - }, - finish_reason: Some("tool_calls".into()), - }], - usage: None, - }, - ]; - - // Process all events - let mut collected_events = Vec::new(); - for event in events { - let mapped = mapper.map_event(event); - collected_events.extend(mapped); - } - - // Verify we got the expected events - let mut has_tool_use = false; - let mut reasoning_details_events = Vec::new(); - let mut thought_signature_value = None; - - for event_result in collected_events { - match event_result { - Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => { - has_tool_use = true; - assert_eq!(tool_use.id.to_string(), "tool_call_abc123"); - assert_eq!(tool_use.name.as_ref(), "list_directory"); - thought_signature_value = tool_use.thought_signature.clone(); - } - Ok(LanguageModelCompletionEvent::ReasoningDetails(details)) => { - reasoning_details_events.push(details); - } - _ => {} - } - } - - assert!(has_tool_use, "Should have emitted ToolUse event"); - assert_eq!(reasoning_details_events.len(), 2); - let final_details = reasoning_details_events - .last() - .and_then(serde_json::Value::as_array) - .and_then(|details| details.first()) - .expect("accumulated reasoning details"); - assert_eq!(final_details["text"], "Let me analyze this request..."); - assert_eq!( - final_details["data"], - "EtgDCtUDAdHtim9OF5jm4aeZSBAtl/randomized123" - ); - assert_eq!( - thought_signature_value.as_deref(), - Some("sha256:test_signature_xyz789") - ); - } - - #[gpui::test] - async fn test_usage_only_chunk_with_empty_choices_does_not_error() { - let mut mapper = OpenRouterEventMapper::new(); - - let events = mapper.map_event(ResponseStreamEvent { - id: Some("response_123".into()), - created: 1234567890, - model: "google/gemini-3-flash-preview".into(), - choices: Vec::new(), - usage: Some(open_router::Usage { - prompt_tokens: 12, - completion_tokens: 7, - total_tokens: 19, - prompt_tokens_details: Some(open_router::PromptTokensDetails { - cached_tokens: 5, - cache_write_tokens: 3, - }), - }), - }); - - assert_eq!(events.len(), 1); - match events.into_iter().next() { - Some(Ok(LanguageModelCompletionEvent::UsageUpdate(usage))) => { - assert_eq!(usage.input_tokens, 4); - assert_eq!(usage.output_tokens, 7); - assert_eq!(usage.cache_creation_input_tokens, 3); - assert_eq!(usage.cache_read_input_tokens, 5); - assert_eq!(usage.total_tokens(), 19); - } - other => panic!("Expected usage update event, got: {other:?}"), - } - } - #[gpui::test] async fn test_session_id_is_stable_without_exposing_thread_id() { let model = open_router::Model::new( diff --git a/crates/language_models/src/provider/opencode.rs b/crates/language_models/src/provider/opencode.rs index ec0c5a0a90b960..79e5d7fdbd5286 100644 --- a/crates/language_models/src/provider/opencode.rs +++ b/crates/language_models/src/provider/opencode.rs @@ -29,9 +29,10 @@ use util::ResultExt; use crate::provider::anthropic::{AnthropicEventMapper, into_anthropic}; use crate::provider::google::{GoogleEventMapper, into_google}; use crate::provider::open_ai::{ - ChatCompletionMaxTokensParameter, OpenAiEventMapper, OpenAiResponseEventMapper, into_open_ai, + ChatCompletionMaxTokensParameter, OpenAiResponseEventMapper, into_open_ai, into_open_ai_response, }; +use language_model::chat_completion::{ChatCompletionEventMapper, ResponseStreamEvent}; fn normalize_reasoning_effort(effort: &str) -> Option { match effort.trim().to_ascii_lowercase().as_str() { @@ -423,10 +424,8 @@ impl OpenCodeLanguageModel { http_client: Arc, extra_headers: CustomHeaders, cx: &AsyncApp, - ) -> BoxFuture< - 'static, - Result>>, - > { + ) -> BoxFuture<'static, Result>>> + { // OpenAI crate appends /chat/completions to api_url, so we pass base + "/v1" let base_url = self.base_api_url(cx); let api_url: SharedString = format!("{base_url}/v1").into(); @@ -706,7 +705,7 @@ impl LanguageModel for OpenCodeLanguageModel { self.stream_openai_chat(openai_request, http_client, extra_headers, cx); let executor = cx.background_executor().clone(); async move { - let mapper = OpenAiEventMapper::new(); + let mapper = ChatCompletionEventMapper::new(); Ok(language_model::stream_in_background( mapper.map_stream(stream.await?).boxed(), executor, diff --git a/crates/language_models/src/provider/vercel_ai_gateway.rs b/crates/language_models/src/provider/vercel_ai_gateway.rs index ede1e45a341887..b22c7d3b4a494d 100644 --- a/crates/language_models/src/provider/vercel_ai_gateway.rs +++ b/crates/language_models/src/provider/vercel_ai_gateway.rs @@ -6,6 +6,7 @@ use gpui::{App, AppContext, AsyncApp, Context, Entity, SharedString, Task}; use http_client::{ AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt, http, }; +use language_model::chat_completion::{ChatCompletionEventMapper, ResponseStreamEvent}; use language_model::{ ApiKeyConfiguration, ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, @@ -13,7 +14,6 @@ use language_model::{ LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice, ProviderSettingsView, RateLimiter, env_var, }; -use open_ai::ResponseStreamEvent; use serde::Deserialize; pub use settings::OpenAiCompatibleModelCapabilities as ModelCapabilities; pub use settings::VercelAiGatewayAvailableModel as AvailableModel; @@ -463,7 +463,7 @@ impl LanguageModel for VercelAiGatewayLanguageModel { let completions = self.stream_open_ai(request, cx); let executor = cx.background_executor().clone(); async move { - let mapper = crate::provider::open_ai::OpenAiEventMapper::new(); + let mapper = ChatCompletionEventMapper::new(); Ok(language_model::stream_in_background( mapper.map_stream(completions.await?).boxed(), executor, diff --git a/crates/language_models/src/provider/x_ai.rs b/crates/language_models/src/provider/x_ai.rs index 711f71afa0ffe9..8bc9e8562d8582 100644 --- a/crates/language_models/src/provider/x_ai.rs +++ b/crates/language_models/src/provider/x_ai.rs @@ -4,6 +4,7 @@ use credentials_provider::CredentialsProvider; use futures::{FutureExt, StreamExt, future::BoxFuture}; use gpui::{App, AppContext, AsyncApp, Context, Entity, SharedString, Task}; use http_client::{CustomHeaders, HttpClient}; +use language_model::chat_completion::{ChatCompletionEventMapper, ResponseStreamEvent}; use language_model::{ ApiKeyConfiguration, ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelEffortLevel, @@ -11,7 +12,6 @@ use language_model::{ LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice, ProviderSettingsView, RateLimiter, env_var, }; -use open_ai::ResponseStreamEvent; pub use settings::XaiAvailableModel as AvailableModel; use settings::{Settings, SettingsStore}; use std::sync::{Arc, LazyLock}; @@ -430,7 +430,7 @@ impl LanguageModel for XAiLanguageModel { let completions = self.stream_completion(request, cx); let executor = cx.background_executor().clone(); async move { - let mapper = crate::provider::open_ai::OpenAiEventMapper::new(); + let mapper = ChatCompletionEventMapper::new(); Ok(language_model::stream_in_background( mapper.map_stream(completions.await?).boxed(), executor, diff --git a/crates/language_models_cloud/src/language_models_cloud.rs b/crates/language_models_cloud/src/language_models_cloud.rs index 50a69171e0fbf6..a365f3f966c500 100644 --- a/crates/language_models_cloud/src/language_models_cloud.rs +++ b/crates/language_models_cloud/src/language_models_cloud.rs @@ -43,8 +43,9 @@ use anthropic::completion::{ AnthropicEventMapper, AnthropicPromptCacheMode, collect_compaction_result, into_anthropic, }; use google_ai::completion::{GoogleEventMapper, into_google}; +use language_model::chat_completion::ChatCompletionEventMapper; use open_ai::completion::{ - ChatCompletionMaxTokensParameter, OpenAiEventMapper, OpenAiResponseEventMapper, into_open_ai, + ChatCompletionMaxTokensParameter, OpenAiResponseEventMapper, into_open_ai, into_open_ai_response, token_usage_from_response_usage, }; @@ -912,7 +913,7 @@ impl LanguageModel for CloudLanguageModel, - pub usage: Option, -} - -#[derive(Serialize, Deserialize, Debug)] -pub struct ChoiceDelta { - pub index: u32, - pub delta: ResponseMessageDelta, - pub finish_reason: Option, -} - -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct ResponseMessageDelta { - pub content: Option, - /// `llama-server` emits reasoning as a dedicated `reasoning_content` field - /// when started with a reasoning format (e.g. `--reasoning-format deepseek`). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub reasoning_content: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tool_calls: Option>, -} - -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct ToolCallChunk { - pub index: usize, - pub id: Option, - pub function: Option, -} - -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct FunctionChunk { - pub name: Option, - pub arguments: Option, -} - /// Response of `GET /v1/models`. /// /// In single-model mode `data` has exactly one entry describing the loaded @@ -522,7 +468,7 @@ pub async fn stream_chat_completion( if line == "[DONE]" { None } else { - match serde_json::from_str(line) { + match serde_json::from_str::(line) { Ok(ResponseStreamResult::Ok(response)) => Some(Ok(response)), Ok(ResponseStreamResult::Err { error }) => { Some(Err(anyhow!(error.message))) @@ -889,7 +835,7 @@ mod tests { ] }); let event: ResponseStreamEvent = serde_json::from_value(event).unwrap(); - let delta = &event.choices[0].delta; + let delta = event.choices[0].delta.as_ref().unwrap(); assert_eq!(delta.reasoning_content.as_deref(), Some("thinking...")); assert_eq!(delta.tool_calls.as_ref().unwrap().len(), 1); } diff --git a/crates/lmstudio/Cargo.toml b/crates/lmstudio/Cargo.toml index 825507b9152bc0..2ef8b96d32e6db 100644 --- a/crates/lmstudio/Cargo.toml +++ b/crates/lmstudio/Cargo.toml @@ -19,6 +19,7 @@ schemars = ["dep:schemars"] anyhow.workspace = true futures.workspace = true http_client.workspace = true +language_model_core.workspace = true schemars = { workspace = true, optional = true } serde.workspace = true serde_json.workspace = true diff --git a/crates/lmstudio/src/lmstudio.rs b/crates/lmstudio/src/lmstudio.rs index 4f5c977f1296ab..498cbde985e0e4 100644 --- a/crates/lmstudio/src/lmstudio.rs +++ b/crates/lmstudio/src/lmstudio.rs @@ -3,46 +3,16 @@ use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::B use http_client::{ AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt, http, }; +pub use language_model_core::chat_completion::{ + ChoiceDelta, FunctionChunk, ResponseMessageDelta, ResponseStreamError, ResponseStreamEvent, + ResponseStreamResult, ToolCallChunk, Usage, +}; use serde::{Deserialize, Serialize}; use serde_json::Value; -use std::{convert::TryFrom, time::Duration}; +use std::time::Duration; pub const LMSTUDIO_API_URL: &str = "http://localhost:1234/api/v0"; -#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)] -#[serde(rename_all = "lowercase")] -pub enum Role { - User, - Assistant, - System, - Tool, -} - -impl TryFrom for Role { - type Error = anyhow::Error; - - fn try_from(value: String) -> Result { - match value.as_str() { - "user" => Ok(Self::User), - "assistant" => Ok(Self::Assistant), - "system" => Ok(Self::System), - "tool" => Ok(Self::Tool), - _ => anyhow::bail!("invalid role '{value}'"), - } - } -} - -impl From for String { - fn from(val: Role) -> Self { - match val { - Role::User => "user".to_owned(), - Role::Assistant => "assistant".to_owned(), - Role::System => "system".to_owned(), - Role::Tool => "tool".to_owned(), - } - } -} - #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)] pub struct Model { @@ -231,46 +201,6 @@ pub struct ChatCompletionRequest { pub tool_choice: Option, } -#[derive(Serialize, Deserialize, Debug)] -pub struct ChatResponse { - pub id: String, - pub object: String, - pub created: u64, - pub model: String, - pub choices: Vec, -} - -#[derive(Serialize, Deserialize, Debug)] -pub struct ChoiceDelta { - pub index: u32, - pub delta: ResponseMessageDelta, - pub finish_reason: Option, -} - -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct ToolCallChunk { - pub index: usize, - pub id: Option, - - // There is also an optional `type` field that would determine if a - // function is there. Sometimes this streams in with the `function` before - // it streams in the `type` - pub function: Option, -} - -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct FunctionChunk { - pub name: Option, - pub arguments: Option, -} - -#[derive(Serialize, Deserialize, Debug)] -pub struct Usage { - pub prompt_tokens: u64, - pub completion_tokens: u64, - pub total_tokens: u64, -} - #[derive(Debug, Default, Clone, Deserialize, PartialEq)] #[serde(transparent)] pub struct Capabilities(Vec); @@ -285,27 +215,6 @@ impl Capabilities { } } -#[derive(Serialize, Deserialize, Debug)] -pub struct LmStudioError { - pub message: String, -} - -#[derive(Serialize, Deserialize, Debug)] -#[serde(untagged)] -pub enum ResponseStreamResult { - Ok(ResponseStreamEvent), - Err { error: LmStudioError }, -} - -#[derive(Serialize, Deserialize, Debug)] -pub struct ResponseStreamEvent { - pub created: u32, - pub model: String, - pub object: String, - pub choices: Vec, - pub usage: Option, -} - #[derive(Deserialize)] pub struct ListModelsResponse { pub data: Vec, @@ -350,56 +259,6 @@ pub enum CompatibilityType { Mlx, } -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct ResponseMessageDelta { - pub role: Option, - pub content: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub reasoning_content: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tool_calls: Option>, -} - -pub async fn complete( - client: &dyn HttpClient, - api_url: &str, - api_key: Option<&str>, - request: ChatCompletionRequest, - extra_headers: &CustomHeaders, -) -> Result { - let uri = format!("{api_url}/chat/completions"); - let mut request_builder = HttpRequest::builder() - .method(Method::POST) - .uri(uri) - .header("Content-Type", "application/json"); - - if let Some(api_key) = api_key { - request_builder = request_builder.header("Authorization", format!("Bearer {}", api_key)); - } - - let serialized_request = serde_json::to_string(&request)?; - let request = request_builder - .extra_headers(extra_headers) - .body(AsyncBody::from(serialized_request))?; - - let mut response = client.send(request).await?; - if response.status().is_success() { - let mut body = Vec::new(); - response.body_mut().read_to_end(&mut body).await?; - let response_message: ChatResponse = serde_json::from_slice(&body)?; - Ok(response_message) - } else { - let mut body = Vec::new(); - response.body_mut().read_to_end(&mut body).await?; - let body_str = std::str::from_utf8(&body)?; - anyhow::bail!( - "Failed to connect to API: {} {}", - response.status(), - body_str - ); - } -} - pub async fn stream_chat_completion( client: &dyn HttpClient, api_url: &str, @@ -432,9 +291,9 @@ pub async fn stream_chat_completion( if line == "[DONE]" { None } else { - match serde_json::from_str(line) { + match serde_json::from_str::(line) { Ok(ResponseStreamResult::Ok(response)) => Some(Ok(response)), - Ok(ResponseStreamResult::Err { error, .. }) => { + Ok(ResponseStreamResult::Err { error }) => { Some(Err(anyhow!(error.message))) } Err(error) => Some(Err(anyhow!(error))), diff --git a/crates/open_ai/src/chat_completion_transport_tests.rs b/crates/open_ai/src/chat_completion_transport_tests.rs index d13a6212644bb8..8251b3a6f64d17 100644 --- a/crates/open_ai/src/chat_completion_transport_tests.rs +++ b/crates/open_ai/src/chat_completion_transport_tests.rs @@ -62,13 +62,13 @@ fn streaming_transport_serializes_custom_requests_and_reports_done() { .expect("stream events") }); - assert_eq!( - events, - vec![ - ChatCompletionStreamEvent::Data(json!({"chunk": 1})), + match events.as_slice() { + [ + ChatCompletionStreamEvent::Data(data), ChatCompletionStreamEvent::Done, - ] - ); + ] => assert_eq!(data.get(), r#"{"chunk":1}"#), + events => panic!("unexpected events: {events:?}"), + } assert_eq!( captured_request .lock() @@ -109,10 +109,10 @@ fn streaming_transport_does_not_synthesize_done_at_eof() { .expect("stream events") }); - assert_eq!( - events, - vec![ChatCompletionStreamEvent::Data(json!({"chunk": 1}))] - ); + match events.as_slice() { + [ChatCompletionStreamEvent::Data(data)] => assert_eq!(data.get(), r#"{"chunk":1}"#), + events => panic!("unexpected events: {events:?}"), + } } #[test] diff --git a/crates/open_ai/src/completion.rs b/crates/open_ai/src/completion.rs index e82d9695ca90cc..6abd9faf079190 100644 --- a/crates/open_ai/src/completion.rs +++ b/crates/open_ai/src/completion.rs @@ -25,8 +25,8 @@ use crate::responses::{ provider_compaction_items, provider_compaction_state_from_items, }; use crate::{ - FunctionContent, FunctionDefinition, ImageUrl, MessagePart, ReasoningEffort, - ResponseStreamEvent, ServiceTier, ToolCall, ToolCallContent, + FunctionContent, FunctionDefinition, ImageUrl, MessagePart, ReasoningEffort, ServiceTier, + ToolCall, ToolCallContent, }; const RESPONSE_MESSAGE_PHASE_COMMENTARY: &str = "commentary"; @@ -729,281 +729,6 @@ fn add_message_content_part( } } -/// Accumulates structured reasoning metadata from compatible providers. -/// -/// Array entries are matched by `index` and then `id`. Fragmented `text`, -/// `summary`, and `data` fields are concatenated while other non-null fields -/// replace their previous values. -/// -/// # Examples -/// -/// ``` -/// use open_ai::completion::ReasoningDetailsAccumulator; -/// use serde_json::json; -/// -/// let mut accumulator = ReasoningDetailsAccumulator::default(); -/// accumulator.push(json!([{"index": 0, "text": "first "}])); -/// let details = accumulator -/// .push(json!([{"index": 0, "text": "second"}])) -/// .expect("non-empty reasoning details"); -/// -/// assert_eq!(details[0]["text"], "first second"); -/// ``` -#[derive(Debug, Default)] -pub struct ReasoningDetailsAccumulator { - accumulated: Option, -} - -impl ReasoningDetailsAccumulator { - /// Merges `chunk` and returns the updated metadata snapshot. - /// - /// `null` and empty arrays do not replace previously accumulated metadata - /// and return `None`. - pub fn push(&mut self, chunk: serde_json::Value) -> Option { - match chunk { - serde_json::Value::Null => None, - serde_json::Value::Array(chunks) if chunks.is_empty() => None, - serde_json::Value::Array(chunks) => { - let mut details = match self.accumulated.take() { - Some(serde_json::Value::Array(details)) => details, - _ => Vec::new(), - }; - for chunk in chunks { - merge_reasoning_detail(&mut details, chunk); - } - let accumulated = serde_json::Value::Array(details); - self.accumulated = Some(accumulated.clone()); - Some(accumulated) - } - chunk => { - self.accumulated = Some(chunk.clone()); - Some(chunk) - } - } - } -} - -pub struct OpenAiEventMapper { - tool_calls_by_index: HashMap, - reasoning_details: ReasoningDetailsAccumulator, -} - -impl OpenAiEventMapper { - pub fn new() -> Self { - Self { - tool_calls_by_index: HashMap::default(), - reasoning_details: ReasoningDetailsAccumulator::default(), - } - } - - pub fn map_stream( - mut self, - events: Pin>>>, - ) -> impl Stream> - { - events.flat_map(move |event| { - futures::stream::iter(match event { - Ok(event) => self.map_event(event), - Err(error) => vec![Err(LanguageModelCompletionError::from(anyhow!(error)))], - }) - }) - } - - pub fn map_event( - &mut self, - event: ResponseStreamEvent, - ) -> Vec> { - let mut events = Vec::new(); - if let Some(usage) = event.usage - && let Some(prompt_tokens) = usage.prompt_tokens - && let Some(completion_tokens) = usage.completion_tokens - { - let cache_creation_input_tokens = usage - .prompt_tokens_details - .as_ref() - .and_then(|details| details.cache_write_tokens) - .unwrap_or(0); - let cache_read_input_tokens = usage - .prompt_tokens_details - .as_ref() - .and_then(|details| details.cached_tokens) - .unwrap_or(0); - events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage { - input_tokens: prompt_tokens - .saturating_sub(cache_creation_input_tokens) - .saturating_sub(cache_read_input_tokens), - output_tokens: completion_tokens, - cache_creation_input_tokens, - cache_read_input_tokens, - }))); - } - - let Some(choice) = event.choices.first() else { - return events; - }; - - if let Some(delta) = choice.delta.as_ref() { - if let Some(reasoning_details) = delta.reasoning_details.clone() - && let Some(reasoning_details) = self.reasoning_details.push(reasoning_details) - { - events.push(Ok(LanguageModelCompletionEvent::ReasoningDetails( - reasoning_details, - ))); - } - if let Some(reasoning) = delta.reasoning.clone() { - push_thinking_event(reasoning, &mut events); - } - if let Some(reasoning_content) = delta.reasoning_content.clone() { - push_thinking_event(reasoning_content, &mut events); - } - if let Some(content) = delta.content.clone() { - if !content.is_empty() { - events.push(Ok(LanguageModelCompletionEvent::Text(content))); - } - } - - if let Some(tool_calls) = delta.tool_calls.as_ref() { - for tool_call in tool_calls { - let entry = self.tool_calls_by_index.entry(tool_call.index).or_default(); - - if let Some(tool_id) = tool_call.id.clone() - && !tool_id.is_empty() - { - entry.id = tool_id; - } - - if let Some(function) = tool_call.function.as_ref() { - if let Some(name) = function.name.clone() - && !name.is_empty() - { - entry.name = name; - } - - if let Some(arguments) = function.arguments.clone() { - entry.arguments.push_str(&arguments); - } - - if let Some(thought_signature) = function.thought_signature.clone() { - entry.thought_signature = Some(thought_signature); - } - } - - if !entry.id.is_empty() && !entry.name.is_empty() { - if let Ok(input) = serde_json::from_str::( - &fix_streamed_json(&entry.arguments), - ) { - events.push(Ok(LanguageModelCompletionEvent::ToolUse( - LanguageModelToolUse { - id: entry.id.clone().into(), - name: entry.name.as_str().into(), - is_input_complete: false, - input: LanguageModelToolUseInput::Json(input), - raw_input: entry.arguments.clone(), - thought_signature: entry.thought_signature.clone(), - }, - ))); - } - } - } - } - } - - match choice.finish_reason.as_deref() { - Some("stop") => { - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))); - } - Some("tool_calls") => { - events.extend(self.tool_calls_by_index.drain().map(|(_, tool_call)| { - match parse_tool_arguments(&tool_call.arguments) { - Ok(input) => Ok(LanguageModelCompletionEvent::ToolUse( - LanguageModelToolUse { - id: tool_call.id.clone().into(), - name: tool_call.name.as_str().into(), - is_input_complete: true, - input: LanguageModelToolUseInput::Json(input), - raw_input: tool_call.arguments.clone(), - thought_signature: tool_call.thought_signature.clone(), - }, - )), - Err(error) => Ok(LanguageModelCompletionEvent::ToolUseJsonParseError { - id: tool_call.id.into(), - tool_name: tool_call.name.into(), - raw_input: tool_call.arguments.clone().into(), - json_parse_error: error.to_string(), - }), - } - })); - - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse))); - } - Some(stop_reason) => { - log::error!("Unexpected OpenAI stop_reason: {stop_reason:?}",); - events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))); - } - None => {} - } - - events - } -} - -fn push_thinking_event( - text: String, - events: &mut Vec>, -) { - if !text.is_empty() { - events.push(Ok(LanguageModelCompletionEvent::Thinking { - text, - signature: None, - })); - } -} - -fn merge_reasoning_detail(details: &mut Vec, chunk: serde_json::Value) { - let index = chunk.get("index").and_then(serde_json::Value::as_u64); - let target_index = index - .and_then(|index| { - details.iter().position(|detail| { - detail.get("index").and_then(serde_json::Value::as_u64) == Some(index) - }) - }) - .or_else(|| { - let id = chunk.get("id").and_then(serde_json::Value::as_str)?; - details - .iter() - .position(|detail| detail.get("id").and_then(serde_json::Value::as_str) == Some(id)) - }); - let Some(target_index) = target_index else { - details.push(chunk); - return; - }; - let (Some(target), Some(chunk)) = (details[target_index].as_object_mut(), chunk.as_object()) - else { - return; - }; - for (key, value) in chunk { - if matches!(key.as_str(), "text" | "summary" | "data") - && let Some(fragment) = value.as_str() - && let Some(existing) = target.get(key).and_then(serde_json::Value::as_str) - { - target.insert( - key.clone(), - serde_json::Value::String(format!("{existing}{fragment}")), - ); - } else if !value.is_null() { - target.insert(key.clone(), value.clone()); - } - } -} - -#[derive(Default)] -struct RawToolCall { - id: String, - name: String, - arguments: String, - thought_signature: Option, -} - pub struct OpenAiResponseEventMapper { /// The backend whose infrastructure produced this stream; stamped on any /// compaction state it emits so replay is limited to the same backend. @@ -1774,9 +1499,6 @@ mod tests { use serde_json::json; use super::*; - use crate::{ - ChoiceDelta, FunctionChunk, ResponseMessageDelta, ResponseStreamEvent, ToolCallChunk, - }; fn map_response_events(events: Vec) -> Vec { block_on(async { @@ -1790,17 +1512,6 @@ mod tests { }) } - fn map_completion_events( - events: Vec, - ) -> Vec { - let mut mapper = OpenAiEventMapper::new(); - let mut all_events = Vec::new(); - for event in events { - all_events.extend(mapper.map_event(event)); - } - all_events.into_iter().filter_map(|e| e.ok()).collect() - } - fn response_item_message(id: &str) -> ResponseOutputItem { ResponseOutputItem::Message(ResponseOutputMessage { id: Some(id.to_string()), @@ -4323,274 +4034,6 @@ mod tests { ); } - #[test] - fn stream_maps_reasoning() { - let events = map_completion_events(vec![ResponseStreamEvent { - choices: vec![ChoiceDelta { - index: 0, - delta: Some(ResponseMessageDelta { - role: None, - content: None, - reasoning: Some("thinking".into()), - tool_calls: None, - reasoning_content: None, - reasoning_details: None, - }), - finish_reason: None, - }], - usage: None, - }]); - - assert_eq!( - events, - vec![LanguageModelCompletionEvent::Thinking { - text: "thinking".into(), - signature: None, - }] - ); - } - - #[test] - fn stream_merges_reasoning_details_and_maps_compatible_usage_and_signatures() { - let response_events = serde_json::from_value(json!([ - { - "choices": [{ - "index": 0, - "delta": { - "reasoning_details": [{ - "id": "reasoning-1", - "index": 0, - "type": "reasoning.text", - "text": "first " - }], - "tool_calls": [{ - "index": 0, - "id": "call-1", - "function": { - "name": "search", - "arguments": "{", - "thought_signature": "signature" - } - }] - }, - "finish_reason": null - }], - "usage": null - }, - { - "choices": [{ - "index": 0, - "delta": { - "reasoning_details": [{ - "id": "reasoning-1", - "index": 0, - "text": "second" - }], - "tool_calls": [{ - "index": 0, - "function": { - "arguments": "}" - } - }] - }, - "finish_reason": "tool_calls" - }], - "usage": null - }, - { - "choices": [], - "usage": { - "prompt_tokens": 10000, - "completion_tokens": 500, - "total_tokens": 10500, - "prompt_tokens_details": { - "cached_tokens": 6000, - "cache_write_tokens": 1000 - } - } - } - ])) - .expect("valid compatible Chat Completions events"); - let events = map_completion_events(response_events); - - assert!(events.iter().any(|event| { - matches!( - event, - LanguageModelCompletionEvent::ReasoningDetails(details) - if details[0]["text"] == "first second" - ) - })); - assert!(events.iter().any(|event| { - matches!( - event, - LanguageModelCompletionEvent::ToolUse(tool_use) - if tool_use.is_input_complete - && tool_use.thought_signature.as_deref() == Some("signature") - ) - })); - assert!(events.iter().any(|event| { - matches!( - event, - LanguageModelCompletionEvent::UsageUpdate(TokenUsage { - input_tokens: 3_000, - output_tokens: 500, - cache_creation_input_tokens: 1_000, - cache_read_input_tokens: 6_000, - }) - ) - })); - } - - #[test] - fn reasoning_details_accumulator_replaces_an_incompatible_previous_shape() { - let mut accumulator = ReasoningDetailsAccumulator::default(); - assert_eq!( - accumulator.push(json!({"summary": "provider-defined"})), - Some(json!({"summary": "provider-defined"})) - ); - assert_eq!( - accumulator.push(json!([{"index": 0, "text": "reasoning"}])), - Some(json!([{"index": 0, "text": "reasoning"}])) - ); - } - - #[test] - fn stream_maps_preserves_tool_id_and_name_across_empty_deltas() { - // DashScope sends id="" and name="" in subsequent tool_calls delta - // chunks after the first chunk. OpenAiEventMapper must not overwrite - // the accumulated id and name with these empty strings. - - let events = vec![ - // First chunk: id and name are present - ResponseStreamEvent { - choices: vec![ChoiceDelta { - index: 0, - delta: Some(ResponseMessageDelta { - role: None, - content: None, - reasoning: None, - tool_calls: Some(vec![ToolCallChunk { - index: 0, - id: Some("call_dashscope_test".into()), - function: Some(FunctionChunk { - name: Some("list_directory".into()), - arguments: Some("".into()), - thought_signature: None, - }), - }]), - reasoning_content: None, - reasoning_details: None, - }), - finish_reason: None, - }], - usage: None, - }, - // Subsequent chunks: DashScope sends id="" and name="" - ResponseStreamEvent { - choices: vec![ChoiceDelta { - index: 0, - delta: Some(ResponseMessageDelta { - role: None, - content: None, - reasoning: None, - tool_calls: Some(vec![ToolCallChunk { - index: 0, - id: Some("".into()), - function: Some(FunctionChunk { - name: Some("".into()), - arguments: Some("{\"path\": \"".into()), - thought_signature: None, - }), - }]), - reasoning_content: None, - reasoning_details: None, - }), - finish_reason: None, - }], - usage: None, - }, - ResponseStreamEvent { - choices: vec![ChoiceDelta { - index: 0, - delta: Some(ResponseMessageDelta { - role: None, - content: None, - reasoning: None, - tool_calls: Some(vec![ToolCallChunk { - index: 0, - id: Some("".into()), - function: Some(FunctionChunk { - name: Some("".into()), - arguments: Some("blog-scraper\"}".into()), - thought_signature: None, - }), - }]), - reasoning_content: None, - reasoning_details: None, - }), - finish_reason: None, - }], - usage: None, - }, - // Final chunk: finish_reason = "tool_calls" - ResponseStreamEvent { - choices: vec![ChoiceDelta { - index: 0, - delta: None, - finish_reason: Some("tool_calls".into()), - }], - usage: None, - }, - ]; - - let mapped = map_completion_events(events); - - // Events emitted: - // 1. Partial ToolUse from chunk 1 (fix_json("") → "{}", parseable) - // 2. Partial ToolUse from chunk 3 (arguments fully assembled) - // 3. Complete ToolUse from finish_reason="tool_calls" drain - // 4. Stop(ToolUse) - assert_eq!(mapped.len(), 4); - - // Verify the complete ToolUse event (from finish_reason drain) - // has the correct id, name, and accumulated arguments. - let complete_tool_use = mapped.iter().find_map(|event| { - if let LanguageModelCompletionEvent::ToolUse(tool_use) = event { - if tool_use.is_input_complete { - return Some(tool_use); - } - } - None - }); - assert!( - complete_tool_use.is_some(), - "expected a completed ToolUse event" - ); - let tool_use = complete_tool_use.unwrap(); - assert_eq!( - tool_use.id.to_string(), - "call_dashscope_test", - "id must survive empty-string overwrites" - ); - assert_eq!( - tool_use.name.as_ref(), - "list_directory", - "name must survive empty-string overwrites" - ); - assert_eq!( - tool_use.raw_input, "{\"path\": \"blog-scraper\"}", - "arguments should accumulate across chunks" - ); - - // Verify the Stop event - assert!(mapped.iter().any(|event| { - matches!( - event, - LanguageModelCompletionEvent::Stop(StopReason::ToolUse) - ) - })); - } - #[test] fn into_open_ai_response_prepends_provider_input_unchanged() { let provider_input = json!([ diff --git a/crates/open_ai/src/open_ai.rs b/crates/open_ai/src/open_ai.rs index f25c3249660a0c..e2a0cda6a46d60 100644 --- a/crates/open_ai/src/open_ai.rs +++ b/crates/open_ai/src/open_ai.rs @@ -13,18 +13,19 @@ use http_client::{ http::{HeaderMap, HeaderValue}, }; pub use language_model_core::ReasoningEffort; +pub use language_model_core::chat_completion::{ + ChoiceDelta, FunctionChunk, PromptTokensDetails, ResponseMessageDelta, ResponseStreamError, + ResponseStreamEvent, ResponseStreamResult, ToolCallChunk, Usage, +}; use serde::{Deserialize, Serialize}; use serde_json::Value; +use serde_json::value::RawValue; use std::{convert::TryFrom, future::Future, io}; use strum::EnumIter; use thiserror::Error; pub const OPEN_AI_API_URL: &str = "https://api.openai.com/v1"; -fn is_none_or_empty, U>(opt: &Option) -> bool { - opt.as_ref().is_none_or(|v| v.as_ref().is_empty()) -} - #[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)] #[serde(rename_all = "lowercase")] pub enum Role { @@ -800,66 +801,6 @@ pub struct Choice { pub finish_reason: Option, } -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct ResponseMessageDelta { - pub role: Option, - pub content: Option, - pub reasoning: Option, - #[serde(default, skip_serializing_if = "is_none_or_empty")] - pub tool_calls: Option>, - #[serde(default, skip_serializing_if = "is_none_or_empty")] - pub reasoning_content: Option, - /// Provider-defined structured reasoning metadata. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub reasoning_details: Option, -} - -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct ToolCallChunk { - pub index: usize, - pub id: Option, - - // There is also an optional `type` field that would determine if a - // function is there. Sometimes this streams in with the `function` before - // it streams in the `type` - pub function: Option, -} - -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct FunctionChunk { - pub name: Option, - pub arguments: Option, - /// Provider-defined metadata required to replay a reasoning tool call. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub thought_signature: Option, -} - -/// Reports prompt-cache token usage from compatible providers. -#[derive(Clone, Serialize, Deserialize, Debug, Default)] -pub struct PromptTokensDetails { - /// Tokens read from a prompt cache. - pub cached_tokens: Option, - /// Tokens written to a prompt cache. - pub cache_write_tokens: Option, -} - -#[derive(Clone, Serialize, Deserialize, Debug)] -pub struct Usage { - pub prompt_tokens: Option, - pub completion_tokens: Option, - pub total_tokens: Option, - /// Prompt-cache usage when reported by the provider. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub prompt_tokens_details: Option, -} - -#[derive(Serialize, Deserialize, Debug)] -pub struct ChoiceDelta { - pub index: u32, - pub delta: Option, - pub finish_reason: Option, -} - /// An error produced while sending an OpenAI-compatible request. /// /// Transport and wire-format failures retain their category so callers can @@ -908,37 +849,20 @@ pub enum RequestError { Other(#[from] anyhow::Error), } -#[derive(Serialize, Deserialize, Debug)] -pub struct ResponseStreamError { - message: String, -} - -#[derive(Serialize, Deserialize, Debug)] -#[serde(untagged)] -pub enum ResponseStreamResult { - Ok(ResponseStreamEvent), - Err { error: ResponseStreamError }, -} - #[derive(Deserialize)] struct ResponseErrorEnvelope { error: responses::ResponseError, } -#[derive(Serialize, Deserialize, Debug)] -pub struct ResponseStreamEvent { - pub choices: Vec, - pub usage: Option, -} - /// A framed Chat Completions server-sent event. /// /// `Done` is distinct from the underlying response body ending so callers can /// tell whether the server completed the stream according to the protocol. -#[derive(Debug, PartialEq)] +#[derive(Debug)] pub enum ChatCompletionStreamEvent { - /// A JSON payload from a `data` field. - Data(Value), + /// A JSON payload from a `data` field, kept as raw text so consumers can + /// deserialize it into their own wire types in a single pass. + Data(Box), /// The protocol terminator `data: [DONE]`. Done, } @@ -1117,7 +1041,7 @@ pub async fn stream_completion( Ok(ChatCompletionStreamEvent::Done) => return None, Err(error) => return Some(Err(anyhow!(error))), }; - match ResponseStreamResult::deserialize(&value) { + match serde_json::from_str::(value.get()) { Ok(ResponseStreamResult::Ok(response)) => Some(Ok(response)), Ok(ResponseStreamResult::Err { error }) => Some(Err(anyhow!(error.message))), Err(error) => { diff --git a/crates/open_router/src/open_router.rs b/crates/open_router/src/open_router.rs index 14cb90079e0241..6272e063a3d624 100644 --- a/crates/open_router/src/open_router.rs +++ b/crates/open_router/src/open_router.rs @@ -4,6 +4,11 @@ use http_client::{ AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt, http, }; pub use language_model_core::ReasoningEffort; +use language_model_core::chat_completion::ResponseStreamResult; +pub use language_model_core::chat_completion::{ + ChoiceDelta, FunctionChunk, PromptTokensDetails, ResponseMessageDelta, ResponseStreamEvent, + ToolCallChunk, Usage, +}; use open_ai::ChatCompletionStreamEvent; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -34,10 +39,6 @@ fn extract_retry_after(headers: &http::HeaderMap) -> Option None } -fn is_none_or_empty, U>(opt: &Option) -> bool { - opt.as_ref().is_none_or(|v| v.as_ref().is_empty()) -} - #[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)] #[serde(rename_all = "lowercase")] pub enum Role { @@ -397,73 +398,6 @@ pub struct FunctionContent { pub thought_signature: Option, } -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct ResponseMessageDelta { - pub role: Option, - pub content: Option, - pub reasoning: Option, - #[serde(default, skip_serializing_if = "is_none_or_empty")] - pub tool_calls: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub reasoning_details: Option, -} - -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct ToolCallChunk { - pub index: usize, - pub id: Option, - pub function: Option, -} - -#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)] -pub struct FunctionChunk { - pub name: Option, - pub arguments: Option, - #[serde(default)] - pub thought_signature: Option, -} - -#[derive(Serialize, Deserialize, Debug, Default)] -pub struct PromptTokensDetails { - #[serde(default)] - pub cached_tokens: u64, - #[serde(default)] - pub cache_write_tokens: u64, -} - -#[derive(Serialize, Deserialize, Debug)] -pub struct Usage { - pub prompt_tokens: u64, - pub completion_tokens: u64, - pub total_tokens: u64, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub prompt_tokens_details: Option, -} - -#[derive(Serialize, Deserialize, Debug)] -pub struct ChoiceDelta { - pub index: u32, - pub delta: ResponseMessageDelta, - pub finish_reason: Option, -} - -#[derive(Serialize, Deserialize, Debug)] -pub struct ResponseStreamEvent { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub id: Option, - pub created: u32, - pub model: String, - pub choices: Vec, - pub usage: Option, -} - -#[derive(Deserialize)] -#[serde(untagged)] -enum ResponseStreamResult { - Response(ResponseStreamEvent), - Error(OpenRouterErrorResponse), -} - #[derive(Serialize, Deserialize, Debug)] pub struct Response { pub id: String, @@ -543,9 +477,9 @@ pub async fn stream_completion( return Some(Err(OpenRouterError::ChatCompletion(error))); } }; - match serde_json::from_value(value) { - Ok(ResponseStreamResult::Response(response)) => Some(Ok(response)), - Ok(ResponseStreamResult::Error(OpenRouterErrorResponse { error })) => { + match serde_json::from_str::>(value.get()) { + Ok(ResponseStreamResult::Ok(response)) => Some(Ok(response)), + Ok(ResponseStreamResult::Err { error }) => { Some(Err(OpenRouterError::ApiError(ApiError { status: None, code: error.code, @@ -895,7 +829,7 @@ mod tests { }); assert_eq!(responses.len(), 1); - assert_eq!(responses[0].model, "vendor/model"); + assert!(responses[0].choices.is_empty()); let headers = captured_headers.lock().expect("captured headers lock"); let headers = headers.as_ref().expect("captured headers"); assert_eq!(headers["http-referer"], "https://zed.dev");