diff --git a/Cargo.lock b/Cargo.lock index cc4b5b6a76a7..7d5bfb6af8df 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4677,6 +4677,7 @@ name = "goose-providers" version = "1.37.0" dependencies = [ "anyhow", + "async-stream", "async-trait", "base64 0.22.1", "chrono", diff --git a/crates/goose-providers/Cargo.toml b/crates/goose-providers/Cargo.toml index 28cc3a9e9cb2..a30ed41ca690 100644 --- a/crates/goose-providers/Cargo.toml +++ b/crates/goose-providers/Cargo.toml @@ -42,6 +42,7 @@ uuid = { workspace = true, features = ["v4"] } utoipa.workspace = true base64.workspace = true rmcp.workspace = true +async-stream.workspace = true [dev-dependencies] tempfile.workspace = true diff --git a/crates/goose-providers/src/openai.rs b/crates/goose-providers/src/openai.rs index 7fb67b9a2e6e..e875fa048b04 100644 --- a/crates/goose-providers/src/openai.rs +++ b/crates/goose-providers/src/openai.rs @@ -1,7 +1,18 @@ +use crate::errors::ProviderError; +use crate::json::safely_parse_json; +use crate::text::ThinkFilter; +use anyhow::anyhow; +use async_stream::try_stream; +use futures::Stream; +use goose_types::{Message, MessageContent, MessageProviderMetadata, ProviderUsage, Usage}; +use rmcp::model::{object, CallToolRequestParams, ErrorCode, ErrorData, Role}; use serde::{Deserialize, Serialize}; -use serde_json::Value; +use serde_json::{json, Value}; +use std::borrow::Cow; use std::collections::HashMap; +pub const THOUGHT_SIGNATURE_KEY: &str = "thoughtSignature"; + pub type ToolCallData = HashMap< i32, ( @@ -111,6 +122,459 @@ pub fn merge_reasoning_text(prefix: &str, suffix: &str) -> String { format!("{prefix}{suffix}") } +pub fn extract_content_and_signature( + delta_content: Option<&DeltaContent>, +) -> (Option, Option) { + match delta_content { + Some(DeltaContent::String(s)) => (Some(s.clone()), None), + Some(DeltaContent::Array(parts)) => { + let text_parts: Vec<_> = parts.iter().filter(|p| p.r#type == "text").collect(); + + let text = text_parts + .iter() + .filter_map(|p| p.text.as_deref()) + .collect::(); + + let signature = text_parts + .iter() + .find_map(|p| p.thought_signature.as_ref()) + .cloned(); + + let text = if text.is_empty() { None } else { Some(text) }; + + (text, signature) + } + None => (None, None), + } +} + +pub fn get_usage(usage: &Value) -> Usage { + let usage = usage + .get("usage") + .filter(|nested| nested.is_object()) + .unwrap_or(usage); + + // Try standard OpenAI fields first, then fall back to Ollama-native fields + // (prompt_eval_count / eval_count) for compatibility with older Ollama builds + // that don't translate to OpenAI field names. + // Parse the value before falling back so that present-but-null keys + // (e.g. "completion_tokens": null) don't block the fallback. + let input_tokens = usage + .get("prompt_tokens") + .and_then(|v| v.as_i64()) + .or_else(|| usage.get("prompt_eval_count").and_then(|v| v.as_i64())) + .map(|v| v as i32); + + let output_tokens = usage + .get("completion_tokens") + .and_then(|v| v.as_i64()) + .or_else(|| usage.get("eval_count").and_then(|v| v.as_i64())) + .map(|v| v as i32); + + let cache_read_input_tokens = usage + .get("cache_read_input_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32); + + let cache_write_input_tokens = usage + .get("cache_creation_input_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32); + + let total_tokens = usage + .get("total_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32) + .or_else(|| match (input_tokens, output_tokens) { + (Some(input), Some(output)) => Some(input.saturating_add(output)), + _ => None, + }); + + Usage::new(input_tokens, output_tokens, total_tokens) + .with_cache_tokens(cache_read_input_tokens, cache_write_input_tokens) +} + +pub fn extract_usage_with_output_tokens( + chunk: &StreamingChunk, + fallback_model: Option<&str>, +) -> Option { + chunk + .usage + .as_ref() + .and_then(|u| { + chunk + .model + .as_deref() + .or(fallback_model) + .map(|model| ProviderUsage { + usage: get_usage(u), + model: model.to_string(), + }) + }) + .filter(|u| u.usage.output_tokens.is_some()) +} + +pub fn strip_data_prefix(line: &str) -> Option<&str> { + // SSE spec allows both "data: value" and "data:value" (space after colon is optional) + line.strip_prefix("data: ") + .or_else(|| line.strip_prefix("data:")) + .map(|s| s.trim()) +} + +pub fn parse_streaming_chunk(line: &str) -> Result { + let value: Value = serde_json::from_str(line).map_err(|e| { + ProviderError::RequestFailed(format!("Failed to parse streaming chunk: {e}: {line:?}")) + })?; + + if let Some(error) = value.get("error") { + let message = error + .get("message") + .and_then(|m| m.as_str()) + .unwrap_or("Unknown server error"); + return Err(ProviderError::ServerError(message.to_string())); + } + + if value.get("object").and_then(|o| o.as_str()) == Some("error") { + let message = value + .get("message") + .and_then(|m| m.as_str()) + .unwrap_or("Unknown server error"); + return Err(ProviderError::ServerError(message.to_string())); + } + + serde_json::from_value(value).map_err(|e| { + ProviderError::RequestFailed(format!("Failed to parse streaming chunk: {e}: {line:?}")) + }) +} + +pub fn response_to_streaming_message( + mut stream: S, +) -> impl Stream, Option)>> + 'static +where + S: Stream> + Unpin + Send + 'static, +{ + try_stream! { + use futures::StreamExt; + + let mut accumulated_reasoning: Vec = Vec::new(); + let mut accumulated_reasoning_content = String::new(); + let mut think_filter = ThinkFilter::new(); + let mut saw_structured_reasoning = false; + let mut yielded_reasoning_content_len = 0usize; + let mut last_signature: Option = None; + // Buffer inline ... content until we know whether structured + // reasoning will arrive. Emitting it immediately and then receiving + // reasoning_content in a later chunk would produce duplicated reasoning. + let mut pending_inline_thinking = String::new(); + let mut last_seen_model: Option = None; + + 'outer: while let Some(response) = stream.next().await { + let response_str = response?; + let line = strip_data_prefix(&response_str); + + if line.is_some_and(|l| l == "[DONE]") { + break 'outer; + } + + if line.is_none() || line.is_some_and(|l| l.is_empty()) { + continue + } + + let chunk: StreamingChunk = parse_streaming_chunk( + line.ok_or_else(|| anyhow!("unexpected stream format"))? + )?; + if let Some(model) = &chunk.model { + last_seen_model = Some(model.clone()); + } + + if !chunk.choices.is_empty() { + if let Some(details) = &chunk.choices[0].delta.reasoning_details { + accumulated_reasoning.extend(details.iter().cloned()); + } + if let Some(rc) = chunk.choices[0].delta.reasoning_text() { + accumulated_reasoning_content.push_str(rc); + if !rc.is_empty() { + saw_structured_reasoning = true; + pending_inline_thinking.clear(); + } + } + } + + let mut usage = extract_usage_with_output_tokens(&chunk, last_seen_model.as_deref()); + + if chunk.choices.is_empty() { + yield (None, usage) + } else if chunk.choices[0].delta.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty()) { + let mut tool_call_data: ToolCallData = HashMap::new(); + + if let Some(tool_calls) = &chunk.choices[0].delta.tool_calls { + for tool_call in tool_calls { + if let (Some(index), Some(id), Some(name)) = (tool_call.index, &tool_call.id, &tool_call.function.name) { + tool_call_data.insert(index, (id.clone(), name.clone(), tool_call.function.arguments.clone(), tool_call.extra.clone())); + } + } + } + + let is_complete = chunk.choices[0].finish_reason == Some("tool_calls".to_string()); + + if !is_complete { + let mut done = false; + while !done { + if let Some(response_chunk) = stream.next().await { + let response_str = response_chunk?; + if let Some(line) = strip_data_prefix(&response_str) { + if line == "[DONE]" { + break 'outer; + } + + let tool_chunk: StreamingChunk = parse_streaming_chunk(line)?; + if let Some(model) = &tool_chunk.model { + last_seen_model = Some(model.clone()); + } + + if let Some(chunk_usage) = extract_usage_with_output_tokens(&tool_chunk, last_seen_model.as_deref()) { + usage = Some(chunk_usage); + } + + if !tool_chunk.choices.is_empty() { + if let Some(details) = &tool_chunk.choices[0].delta.reasoning_details { + accumulated_reasoning.extend(details.iter().cloned()); + } + if let Some(rc) = tool_chunk.choices[0].delta.reasoning_text() { + accumulated_reasoning_content.push_str(rc); + if !rc.is_empty() { + saw_structured_reasoning = true; + pending_inline_thinking.clear(); + } + } + if let Some(delta_tool_calls) = &tool_chunk.choices[0].delta.tool_calls { + for delta_call in delta_tool_calls { + if let Some(index) = delta_call.index { + if let Some((_, _, ref mut args, ref mut extra)) = tool_call_data.get_mut(&index) { + args.push_str(&delta_call.function.arguments); + if extra.is_none() && delta_call.extra.is_some() { + *extra = delta_call.extra.clone(); + } else if let (Some(existing), Some(new_extra)) = (extra.as_mut(), &delta_call.extra) { + for (key, value) in new_extra { + existing.entry(key.clone()).or_insert(value.clone()); + } + } + } else if let (Some(id), Some(name)) = (&delta_call.id, &delta_call.function.name) { + tool_call_data.insert(index, (id.clone(), name.clone(), delta_call.function.arguments.clone(), delta_call.extra.clone())); + } + } + } + } + if tool_chunk.choices[0].finish_reason.is_some() { + done = true; + } + } else { + done = true; + } + } + } else { + break; + } + } + } + + let _metadata: Option = if !accumulated_reasoning.is_empty() { + let mut map = MessageProviderMetadata::new(); + map.insert("reasoning_details".to_string(), json!(accumulated_reasoning)); + Some(map) + } else { + None + }; + + let filtered = think_filter.push(""); + let mut flush_thinking = String::new(); + if !saw_structured_reasoning { + flush_thinking.push_str(&pending_inline_thinking); + flush_thinking.push_str(&filtered.thinking); + } + pending_inline_thinking.clear(); + if !filtered.content.is_empty() || !flush_thinking.is_empty() { + let mut filtered_contents = Vec::new(); + if !filtered.content.is_empty() { + filtered_contents.push(MessageContent::text(filtered.content)); + } + if !flush_thinking.is_empty() { + filtered_contents.push(MessageContent::thinking(flush_thinking, "")); + } + + if !filtered_contents.is_empty() { + let mut msg = Message::new( + Role::Assistant, + chrono::Utc::now().timestamp(), + filtered_contents, + ); + + if let Some(id) = chunk.id.clone() { + msg = msg.with_id(id); + } + + yield (Some(msg), None); + } + } + + let mut contents = Vec::new(); + if yielded_reasoning_content_len < accumulated_reasoning_content.len() { + if let Some(unyielded_reasoning) = + accumulated_reasoning_content.get(yielded_reasoning_content_len..) + { + if !unyielded_reasoning.is_empty() { + contents.push(MessageContent::thinking(unyielded_reasoning, "")); + } + } + } + accumulated_reasoning_content.clear(); + yielded_reasoning_content_len = 0; + let mut sorted_indices: Vec<_> = tool_call_data.keys().cloned().collect(); + sorted_indices.sort(); + + for index in sorted_indices { + if let Some((id, function_name, arguments, extra_fields)) = tool_call_data.get(&index) { + let parsed = if arguments.is_empty() { + Ok(json!({})) + } else { + safely_parse_json(arguments) + }; + + let metadata = if let Some(sig) = &last_signature { + let mut combined = extra_fields.clone().unwrap_or_default(); + combined.insert( + THOUGHT_SIGNATURE_KEY.to_string(), + json!(sig) + ); + Some(combined) + } else { + extra_fields.as_ref().filter(|m| !m.is_empty()).cloned() + }; + + let content = match parsed { + Ok(params) => { + MessageContent::tool_request_with_metadata( + id.clone(), + Ok(CallToolRequestParams::new(function_name.clone()).with_arguments(object(params))), + metadata.as_ref(), + ) + }, + Err(e) => { + let error = ErrorData { + code: ErrorCode::INVALID_PARAMS, + message: Cow::from(format!( + "Could not interpret tool use parameters for id {}: {}", + id, e + )), + data: None, + }; + MessageContent::tool_request_with_metadata(id.clone(), Err(error), metadata.as_ref()) + } + }; + contents.push(content); + } + } + + let mut msg = Message::new( + Role::Assistant, + chrono::Utc::now().timestamp(), + contents, + ); + + // Add ID if present + if let Some(id) = chunk.id { + msg = msg.with_id(id); + } + + yield ( + Some(msg), + usage, + ) + } else if chunk.choices[0].delta.content.is_some() || chunk.choices[0].delta.reasoning_text().is_some() { + let mut content = Vec::new(); + + if let Some(reasoning) = chunk.choices[0].delta.reasoning_text() { + let signature = last_signature.as_deref().unwrap_or(""); + content.push(MessageContent::thinking(reasoning, signature)); + yielded_reasoning_content_len = accumulated_reasoning_content.len(); + } + + let (text_content, thought_signature) = extract_content_and_signature(chunk.choices[0].delta.content.as_ref()); + + if let Some(sig) = thought_signature { + last_signature = Some(sig); + } + + if let Some(text) = text_content { + let filtered = think_filter.push(&text); + + if !saw_structured_reasoning && !filtered.thinking.is_empty() { + pending_inline_thinking.push_str(&filtered.thinking); + } + + if !filtered.content.is_empty() { + content.push(MessageContent::text(filtered.content)); + } + } + + if !content.is_empty() { + let mut msg = Message::new( + Role::Assistant, + chrono::Utc::now().timestamp(), + content, + ); + + if let Some(id) = chunk.id { + msg = msg.with_id(id); + } + + yield ( + Some(msg), + if chunk.choices[0].finish_reason.is_some() { + usage + } else { + None + }, + ) + } else if usage.is_some() { + yield (None, usage) + } + } else if usage.is_some() { + yield (None, usage) + } + } + + let filtered = think_filter.finish(); + let mut trailing_thinking = String::new(); + if !saw_structured_reasoning { + trailing_thinking.push_str(&pending_inline_thinking); + trailing_thinking.push_str(&filtered.thinking); + } + pending_inline_thinking.clear(); + + if !filtered.content.is_empty() || !trailing_thinking.is_empty() { + let mut content = Vec::new(); + + if !filtered.content.is_empty() { + content.push(MessageContent::text(filtered.content)); + } + + if !trailing_thinking.is_empty() { + content.push(MessageContent::thinking(trailing_thinking, "")); + } + + yield ( + Some(Message::new( + Role::Assistant, + chrono::Utc::now().timestamp(), + content, + )), + None, + ) + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -131,4 +595,91 @@ mod tests { assert_eq!(delta.reasoning_text(), Some("preferred")); } + + #[test] + fn get_usage_reads_nested_usage_object() { + let usage = get_usage(&json!({ + "id": "chatcmpl_test", + "usage": { + "prompt_tokens": 84, + "completion_tokens": 21, + "total_tokens": 105, + "cache_read_input_tokens": 60, + "cache_creation_input_tokens": 10 + } + })); + + assert_eq!(usage.input_tokens, Some(84)); + assert_eq!(usage.output_tokens, Some(21)); + assert_eq!(usage.total_tokens, Some(105)); + assert_eq!(usage.cache_read_input_tokens, Some(60)); + assert_eq!(usage.cache_write_input_tokens, Some(10)); + } + + #[tokio::test] + async fn streaming_response_extracts_inline_think_blocks() -> anyhow::Result<()> { + use futures::{pin_mut, stream, StreamExt}; + + let lines = [ + r#"data: {"id":"chunk-1","choices":[{"delta":{"content":"xy"},"index":0,"finish_reason":"stop"}]}"#, + "data: [DONE]", + ]; + let response_stream = stream::iter(lines.into_iter().map(|line| Ok(line.to_string()))); + let messages = response_to_streaming_message(response_stream); + pin_mut!(messages); + + let mut text = String::new(); + let mut thinking = String::new(); + + while let Some(result) = messages.next().await { + let (message, _) = result?; + if let Some(message) = message { + for item in message.content { + match item { + MessageContent::Text(text_content) => text.push_str(&text_content.text), + MessageContent::Thinking(thinking_content) => { + thinking.push_str(&thinking_content.thinking) + } + _ => {} + } + } + } + } + + assert_eq!(text, "y"); + assert_eq!(thinking, "x"); + + Ok(()) + } + + #[tokio::test] + async fn streaming_response_preserves_nested_tool_call_metadata() -> anyhow::Result<()> { + use futures::{pin_mut, stream, StreamExt}; + + let lines = [ + r#"data: {"model":"test-model","choices":[{"delta":{"role":"assistant","tool_calls":[{"extra_content":{"google":{"thought_signature":"nested_stream_sig"}},"id":"call_nested","function":{"name":"test_tool","arguments":"{}"},"type":"function","index":0}]},"index":0,"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110},"object":"chat.completion.chunk","id":"test-id","created":1234567890}"#, + "data: [DONE]", + ]; + let response_stream = stream::iter(lines.into_iter().map(|line| Ok(line.to_string()))); + let messages = response_to_streaming_message(response_stream); + pin_mut!(messages); + + while let Some(result) = messages.next().await { + let (message, _) = result?; + if let Some(message) = message { + if let MessageContent::ToolRequest(request) = &message.content[0] { + let metadata = request.metadata.as_ref().expect("metadata should exist"); + assert_eq!( + metadata["extra_content"]["google"]["thought_signature"], + "nested_stream_sig" + ); + return Ok(()); + } + } + } + + panic!("expected tool call message with nested metadata"); + } } diff --git a/crates/goose/src/providers/formats/openai.rs b/crates/goose/src/providers/formats/openai.rs index 8992534f2342..da9464ceffb7 100644 --- a/crates/goose/src/providers/formats/openai.rs +++ b/crates/goose/src/providers/formats/openai.rs @@ -1,61 +1,33 @@ -use crate::conversation::message::{Message, MessageContent, ProviderMetadata}; +use crate::conversation::message::{Message, MessageContent}; use crate::mcp_utils::extract_text_from_resource; use crate::model::GooseModelConfigExt; -use crate::providers::base::{split_think_blocks, ProviderUsage, ThinkFilter, Usage}; -use crate::providers::errors::ProviderError; use crate::providers::utils::{ convert_image, detect_image_path, extract_reasoning_effort, is_openai_responses_model, is_valid_function_name, load_image_file, openai_reasoning_effort_for_thinking, safely_parse_json, sanitize_function_name, ImageFormat, }; use anyhow::{anyhow, Error}; -use async_stream::try_stream; use chrono; -use futures::Stream; +use goose_providers::openai::merge_reasoning_text; #[cfg(test)] use goose_providers::openai::DeltaToolCallFunction; -pub use goose_providers::openai::OpenAiFormatOptions; -use goose_providers::openai::{merge_reasoning_text, DeltaContent, StreamingChunk, ToolCallData}; +pub use goose_providers::openai::{get_usage, response_to_streaming_message, OpenAiFormatOptions}; +use goose_providers::text::split_think_blocks; use goose_types::ModelConfig; +#[cfg(test)] +use goose_types::ProviderUsage; use rmcp::model::{ object, AnnotateAble, CallToolRequestParams, Content, ErrorCode, ErrorData, RawContent, Role, Tool, }; use serde_json::{json, Value}; use std::borrow::Cow; -use std::collections::HashMap; use std::ops::Deref; fn is_reserved_request_param_key(key: &str) -> bool { matches!(key, "messages" | "model" | "stream" | "stream_options") } -fn extract_content_and_signature( - delta_content: Option<&DeltaContent>, -) -> (Option, Option) { - match delta_content { - Some(DeltaContent::String(s)) => (Some(s.clone()), None), - Some(DeltaContent::Array(parts)) => { - let text_parts: Vec<_> = parts.iter().filter(|p| p.r#type == "text").collect(); - - let text = text_parts - .iter() - .filter_map(|p| p.text.as_deref()) - .collect::(); - - let signature = text_parts - .iter() - .find_map(|p| p.thought_signature.as_ref()) - .cloned(); - - let text = if text.is_empty() { None } else { Some(text) }; - - (text, signature) - } - None => (None, None), - } -} - pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec { format_messages_with_options( messages, @@ -581,72 +553,6 @@ pub fn response_to_message(response: &Value) -> anyhow::Result { )) } -pub fn get_usage(usage: &Value) -> Usage { - let usage = usage - .get("usage") - .filter(|nested| nested.is_object()) - .unwrap_or(usage); - - // Try standard OpenAI fields first, then fall back to Ollama-native fields - // (prompt_eval_count / eval_count) for compatibility with older Ollama builds - // that don't translate to OpenAI field names. - // Parse the value before falling back so that present-but-null keys - // (e.g. "completion_tokens": null) don't block the fallback. - let input_tokens = usage - .get("prompt_tokens") - .and_then(|v| v.as_i64()) - .or_else(|| usage.get("prompt_eval_count").and_then(|v| v.as_i64())) - .map(|v| v as i32); - - let output_tokens = usage - .get("completion_tokens") - .and_then(|v| v.as_i64()) - .or_else(|| usage.get("eval_count").and_then(|v| v.as_i64())) - .map(|v| v as i32); - - let cache_read_input_tokens = usage - .get("cache_read_input_tokens") - .and_then(|v| v.as_i64()) - .map(|v| v as i32); - - let cache_write_input_tokens = usage - .get("cache_creation_input_tokens") - .and_then(|v| v.as_i64()) - .map(|v| v as i32); - - let total_tokens = usage - .get("total_tokens") - .and_then(|v| v.as_i64()) - .map(|v| v as i32) - .or_else(|| match (input_tokens, output_tokens) { - (Some(input), Some(output)) => Some(input.saturating_add(output)), - _ => None, - }); - - Usage::new(input_tokens, output_tokens, total_tokens) - .with_cache_tokens(cache_read_input_tokens, cache_write_input_tokens) -} - -fn extract_usage_with_output_tokens( - chunk: &StreamingChunk, - fallback_model: Option<&str>, -) -> Option { - chunk - .usage - .as_ref() - .and_then(|u| { - chunk - .model - .as_deref() - .or(fallback_model) - .map(|model| ProviderUsage { - usage: get_usage(u), - model: model.to_string(), - }) - }) - .filter(|u| u.usage.output_tokens.is_some()) -} - /// Validates and fixes tool schemas to ensure they have proper parameter structure. /// If parameters exist, ensures they have properties and required fields, or removes parameters entirely. pub fn validate_tool_schemas(tools: &mut [Value]) { @@ -749,367 +655,6 @@ fn normalize_nullable(schema: &mut Value) { } } -fn strip_data_prefix(line: &str) -> Option<&str> { - // SSE spec allows both "data: value" and "data:value" (space after colon is optional) - line.strip_prefix("data: ") - .or_else(|| line.strip_prefix("data:")) - .map(|s| s.trim()) -} - -fn parse_streaming_chunk(line: &str) -> Result { - let value: Value = serde_json::from_str(line).map_err(|e| { - ProviderError::RequestFailed(format!("Failed to parse streaming chunk: {e}: {line:?}")) - })?; - - if let Some(error) = value.get("error") { - let message = error - .get("message") - .and_then(|m| m.as_str()) - .unwrap_or("Unknown server error"); - return Err(ProviderError::ServerError(message.to_string())); - } - - if value.get("object").and_then(|o| o.as_str()) == Some("error") { - let message = value - .get("message") - .and_then(|m| m.as_str()) - .unwrap_or("Unknown server error"); - return Err(ProviderError::ServerError(message.to_string())); - } - - serde_json::from_value(value).map_err(|e| { - ProviderError::RequestFailed(format!("Failed to parse streaming chunk: {e}: {line:?}")) - }) -} - -pub fn response_to_streaming_message( - mut stream: S, -) -> impl Stream, Option)>> + 'static -where - S: Stream> + Unpin + Send + 'static, -{ - try_stream! { - use futures::StreamExt; - - let mut accumulated_reasoning: Vec = Vec::new(); - let mut accumulated_reasoning_content = String::new(); - let mut think_filter = ThinkFilter::new(); - let mut saw_structured_reasoning = false; - let mut yielded_reasoning_content_len = 0usize; - let mut last_signature: Option = None; - // Buffer inline ... content until we know whether structured - // reasoning will arrive. Emitting it immediately and then receiving - // reasoning_content in a later chunk would produce duplicated reasoning. - let mut pending_inline_thinking = String::new(); - let mut last_seen_model: Option = None; - - 'outer: while let Some(response) = stream.next().await { - let response_str = response?; - let line = strip_data_prefix(&response_str); - - if line.is_some_and(|l| l == "[DONE]") { - break 'outer; - } - - if line.is_none() || line.is_some_and(|l| l.is_empty()) { - continue - } - - let chunk: StreamingChunk = parse_streaming_chunk( - line.ok_or_else(|| anyhow!("unexpected stream format"))? - )?; - if let Some(model) = &chunk.model { - last_seen_model = Some(model.clone()); - } - - if !chunk.choices.is_empty() { - if let Some(details) = &chunk.choices[0].delta.reasoning_details { - accumulated_reasoning.extend(details.iter().cloned()); - } - if let Some(rc) = chunk.choices[0].delta.reasoning_text() { - accumulated_reasoning_content.push_str(rc); - if !rc.is_empty() { - saw_structured_reasoning = true; - pending_inline_thinking.clear(); - } - } - } - - let mut usage = extract_usage_with_output_tokens(&chunk, last_seen_model.as_deref()); - - if chunk.choices.is_empty() { - yield (None, usage) - } else if chunk.choices[0].delta.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty()) { - let mut tool_call_data: ToolCallData = HashMap::new(); - - if let Some(tool_calls) = &chunk.choices[0].delta.tool_calls { - for tool_call in tool_calls { - if let (Some(index), Some(id), Some(name)) = (tool_call.index, &tool_call.id, &tool_call.function.name) { - tool_call_data.insert(index, (id.clone(), name.clone(), tool_call.function.arguments.clone(), tool_call.extra.clone())); - } - } - } - - let is_complete = chunk.choices[0].finish_reason == Some("tool_calls".to_string()); - - if !is_complete { - let mut done = false; - while !done { - if let Some(response_chunk) = stream.next().await { - let response_str = response_chunk?; - if let Some(line) = strip_data_prefix(&response_str) { - if line == "[DONE]" { - break 'outer; - } - - let tool_chunk: StreamingChunk = parse_streaming_chunk(line)?; - if let Some(model) = &tool_chunk.model { - last_seen_model = Some(model.clone()); - } - - if let Some(chunk_usage) = extract_usage_with_output_tokens(&tool_chunk, last_seen_model.as_deref()) { - usage = Some(chunk_usage); - } - - if !tool_chunk.choices.is_empty() { - if let Some(details) = &tool_chunk.choices[0].delta.reasoning_details { - accumulated_reasoning.extend(details.iter().cloned()); - } - if let Some(rc) = tool_chunk.choices[0].delta.reasoning_text() { - accumulated_reasoning_content.push_str(rc); - if !rc.is_empty() { - saw_structured_reasoning = true; - pending_inline_thinking.clear(); - } - } - if let Some(delta_tool_calls) = &tool_chunk.choices[0].delta.tool_calls { - for delta_call in delta_tool_calls { - if let Some(index) = delta_call.index { - if let Some((_, _, ref mut args, ref mut extra)) = tool_call_data.get_mut(&index) { - args.push_str(&delta_call.function.arguments); - if extra.is_none() && delta_call.extra.is_some() { - *extra = delta_call.extra.clone(); - } else if let (Some(existing), Some(new_extra)) = (extra.as_mut(), &delta_call.extra) { - for (key, value) in new_extra { - existing.entry(key.clone()).or_insert(value.clone()); - } - } - } else if let (Some(id), Some(name)) = (&delta_call.id, &delta_call.function.name) { - tool_call_data.insert(index, (id.clone(), name.clone(), delta_call.function.arguments.clone(), delta_call.extra.clone())); - } - } - } - } - if tool_chunk.choices[0].finish_reason.is_some() { - done = true; - } - } else { - done = true; - } - } - } else { - break; - } - } - } - - let _metadata: Option = if !accumulated_reasoning.is_empty() { - let mut map = ProviderMetadata::new(); - map.insert("reasoning_details".to_string(), json!(accumulated_reasoning)); - Some(map) - } else { - None - }; - - let filtered = think_filter.push(""); - let mut flush_thinking = String::new(); - if !saw_structured_reasoning { - flush_thinking.push_str(&pending_inline_thinking); - flush_thinking.push_str(&filtered.thinking); - } - pending_inline_thinking.clear(); - if !filtered.content.is_empty() || !flush_thinking.is_empty() { - let mut filtered_contents = Vec::new(); - if !filtered.content.is_empty() { - filtered_contents.push(MessageContent::text(filtered.content)); - } - if !flush_thinking.is_empty() { - filtered_contents.push(MessageContent::thinking(flush_thinking, "")); - } - - if !filtered_contents.is_empty() { - let mut msg = Message::new( - Role::Assistant, - chrono::Utc::now().timestamp(), - filtered_contents, - ); - - if let Some(id) = chunk.id.clone() { - msg = msg.with_id(id); - } - - yield (Some(msg), None); - } - } - - let mut contents = Vec::new(); - if yielded_reasoning_content_len < accumulated_reasoning_content.len() { - if let Some(unyielded_reasoning) = - accumulated_reasoning_content.get(yielded_reasoning_content_len..) - { - if !unyielded_reasoning.is_empty() { - contents.push(MessageContent::thinking(unyielded_reasoning, "")); - } - } - } - accumulated_reasoning_content.clear(); - yielded_reasoning_content_len = 0; - let mut sorted_indices: Vec<_> = tool_call_data.keys().cloned().collect(); - sorted_indices.sort(); - - for index in sorted_indices { - if let Some((id, function_name, arguments, extra_fields)) = tool_call_data.get(&index) { - let parsed = if arguments.is_empty() { - Ok(json!({})) - } else { - safely_parse_json(arguments) - }; - - let metadata = if let Some(sig) = &last_signature { - let mut combined = extra_fields.clone().unwrap_or_default(); - combined.insert( - crate::providers::formats::google::THOUGHT_SIGNATURE_KEY.to_string(), - json!(sig) - ); - Some(combined) - } else { - extra_fields.as_ref().filter(|m| !m.is_empty()).cloned() - }; - - let content = match parsed { - Ok(params) => { - MessageContent::tool_request_with_metadata( - id.clone(), - Ok(CallToolRequestParams::new(function_name.clone()).with_arguments(object(params))), - metadata.as_ref(), - ) - }, - Err(e) => { - let error = ErrorData { - code: ErrorCode::INVALID_PARAMS, - message: Cow::from(format!( - "Could not interpret tool use parameters for id {}: {}", - id, e - )), - data: None, - }; - MessageContent::tool_request_with_metadata(id.clone(), Err(error), metadata.as_ref()) - } - }; - contents.push(content); - } - } - - let mut msg = Message::new( - Role::Assistant, - chrono::Utc::now().timestamp(), - contents, - ); - - // Add ID if present - if let Some(id) = chunk.id { - msg = msg.with_id(id); - } - - yield ( - Some(msg), - usage, - ) - } else if chunk.choices[0].delta.content.is_some() || chunk.choices[0].delta.reasoning_text().is_some() { - let mut content = Vec::new(); - - if let Some(reasoning) = chunk.choices[0].delta.reasoning_text() { - let signature = last_signature.as_deref().unwrap_or(""); - content.push(MessageContent::thinking(reasoning, signature)); - yielded_reasoning_content_len = accumulated_reasoning_content.len(); - } - - let (text_content, thought_signature) = extract_content_and_signature(chunk.choices[0].delta.content.as_ref()); - - if let Some(sig) = thought_signature { - last_signature = Some(sig); - } - - if let Some(text) = text_content { - let filtered = think_filter.push(&text); - - if !saw_structured_reasoning && !filtered.thinking.is_empty() { - pending_inline_thinking.push_str(&filtered.thinking); - } - - if !filtered.content.is_empty() { - content.push(MessageContent::text(filtered.content)); - } - } - - if !content.is_empty() { - let mut msg = Message::new( - Role::Assistant, - chrono::Utc::now().timestamp(), - content, - ); - - if let Some(id) = chunk.id { - msg = msg.with_id(id); - } - - yield ( - Some(msg), - if chunk.choices[0].finish_reason.is_some() { - usage - } else { - None - }, - ) - } else if usage.is_some() { - yield (None, usage) - } - } else if usage.is_some() { - yield (None, usage) - } - } - - let filtered = think_filter.finish(); - let mut trailing_thinking = String::new(); - if !saw_structured_reasoning { - trailing_thinking.push_str(&pending_inline_thinking); - trailing_thinking.push_str(&filtered.thinking); - } - pending_inline_thinking.clear(); - - if !filtered.content.is_empty() || !trailing_thinking.is_empty() { - let mut content = Vec::new(); - - if !filtered.content.is_empty() { - content.push(MessageContent::text(filtered.content)); - } - - if !trailing_thinking.is_empty() { - content.push(MessageContent::thinking(trailing_thinking, "")); - } - - yield ( - Some(Message::new( - Role::Assistant, - chrono::Utc::now().timestamp(), - content, - )), - None, - ) - } - } -} - pub fn create_request( model_config: &ModelConfig, system: &str,