diff --git a/crates/unified-llm/src/error.rs b/crates/unified-llm/src/error.rs index cd878c95c..7aac6453c 100644 --- a/crates/unified-llm/src/error.rs +++ b/crates/unified-llm/src/error.rs @@ -85,15 +85,23 @@ pub enum SdkError { impl SdkError { #[must_use] pub const fn retryable(&self) -> bool { - matches!( - self, - Self::Provider { - kind: ProviderErrorKind::RateLimit | ProviderErrorKind::Server, - .. - } | Self::RequestTimeout { .. } - | Self::Network { .. } - | Self::Stream { .. } - ) + match self { + Self::Provider { kind, .. } => match kind { + ProviderErrorKind::Authentication + | ProviderErrorKind::AccessDenied + | ProviderErrorKind::NotFound + | ProviderErrorKind::InvalidRequest + | ProviderErrorKind::ContextLength + | ProviderErrorKind::QuotaExceeded + | ProviderErrorKind::ContentFilter => false, + _ => true, + }, + Self::InvalidToolCall { .. } + | Self::NoObjectGenerated { .. } + | Self::Abort { .. } + | Self::Configuration { .. } => false, + _ => true, + } } #[must_use] @@ -279,22 +287,38 @@ mod tests { } #[test] - fn non_retryable_errors() { - let kinds = [ - ProviderErrorKind::AccessDenied, - ProviderErrorKind::NotFound, - ProviderErrorKind::InvalidRequest, - ProviderErrorKind::ContextLength, - ProviderErrorKind::QuotaExceeded, - ProviderErrorKind::ContentFilter, - ]; - for kind in &kinds { - let err = SdkError::Provider { - kind: *kind, - detail: Box::new(ProviderErrorDetail::new("error", "openai")), - }; - assert!(!err.retryable(), "Expected non-retryable: {err}"); - } + fn non_retryable_provider_errors() { + let detail = || Box::new(ProviderErrorDetail::new("error", "openai")); + + let access_denied = SdkError::Provider { kind: ProviderErrorKind::AccessDenied, detail: detail() }; + assert!(!access_denied.retryable()); + + let not_found = SdkError::Provider { kind: ProviderErrorKind::NotFound, detail: detail() }; + assert!(!not_found.retryable()); + + let invalid_req = SdkError::Provider { kind: ProviderErrorKind::InvalidRequest, detail: detail() }; + assert!(!invalid_req.retryable()); + + let ctx_length = SdkError::Provider { kind: ProviderErrorKind::ContextLength, detail: detail() }; + assert!(!ctx_length.retryable()); + + let quota = SdkError::Provider { kind: ProviderErrorKind::QuotaExceeded, detail: detail() }; + assert!(!quota.retryable()); + + let content_filter = SdkError::Provider { kind: ProviderErrorKind::ContentFilter, detail: detail() }; + assert!(!content_filter.retryable()); + } + + #[test] + fn non_retryable_sdk_errors() { + let invalid_tool = SdkError::InvalidToolCall { message: "bad tool".into() }; + assert!(!invalid_tool.retryable()); + + let no_object = SdkError::NoObjectGenerated { message: "no output".into() }; + assert!(!no_object.retryable()); + + let abort = SdkError::Abort { message: "aborted".into() }; + assert!(!abort.retryable()); } #[test] diff --git a/crates/unified-llm/src/generate.rs b/crates/unified-llm/src/generate.rs index 1e4f8603e..b34574ed7 100644 --- a/crates/unified-llm/src/generate.rs +++ b/crates/unified-llm/src/generate.rs @@ -2,7 +2,7 @@ use crate::client::Client; use crate::error::SdkError; use crate::provider::StreamEventStream; use crate::retry::retry; -use crate::tools::{execute_all_tools, Tool}; +use crate::tools::{execute_all_tools_with_repair, RepairToolCallFn, Tool}; use crate::types::{ FinishReason, GenerateResult, Message, ObjectStreamEvent, Request, Response, ResponseFormat, ResponseFormatType, RetryPolicy, StepResult, StreamEvent, TimeoutConfig, ToolCall, ToolChoice, @@ -172,7 +172,7 @@ pub async fn generate(params: GenerateParams) -> Result = tools.iter().map(std::convert::AsRef::as_ref).collect(); - tool_results = execute_all_tools(&tool_refs, &tool_calls, &messages, abort_signal.as_ref()).await; + tool_results = execute_all_tools_with_repair(&tool_refs, &tool_calls, &messages, abort_signal.as_ref(), params.repair_tool_call.as_ref()).await; } } @@ -207,7 +207,7 @@ pub async fn generate(params: GenerateParams) -> Result, /// Custom stop condition checked after each tool round (Section 4.3). pub stop_when: Option, + /// Callback to repair invalid tool call arguments (Section 5.8). + pub repair_tool_call: Option, } impl GenerateParams { @@ -285,6 +287,7 @@ impl GenerateParams { client: None, abort_signal: None, stop_when: None, + repair_tool_call: None, } } @@ -414,6 +417,12 @@ impl GenerateParams { self.stop_when = Some(Arc::new(f)); self } + + #[must_use] + pub fn repair_tool_call(mut self, repair: RepairToolCallFn) -> Self { + self.repair_tool_call = Some(repair); + self + } } /// `StreamAccumulator` collects stream events into a complete Response (Section 4.4). @@ -586,6 +595,7 @@ async fn stream_with_tool_loop(params: GenerateParams) -> Result 0 && params @@ -703,7 +713,7 @@ async fn stream_with_tool_loop(params: GenerateParams) -> Result = tool_list.iter().map(std::convert::AsRef::as_ref).collect(); - let tool_results = execute_all_tools(&tool_refs, &tool_calls, &messages, abort_signal.as_ref()).await; + let tool_results = execute_all_tools_with_repair(&tool_refs, &tool_calls, &messages, abort_signal.as_ref(), repair_tool_call.as_ref()).await; if tool_results.is_empty() { return; @@ -746,7 +756,7 @@ async fn stream_with_tool_loop(params: GenerateParams) -> Result> + Send>>; +/// Wraps an `ObjectStream` with an `object()` accessor for the final parsed value. +/// +/// Implements `Stream>` so it can be used +/// as a drop-in replacement for `ObjectStream`. Tracks the last `Complete` event's +/// object internally so callers can retrieve it after the stream ends. +pub struct ObjectStreamResult { + inner: ObjectStream, + object: Option, +} + +impl ObjectStreamResult { + fn new(inner: ObjectStream) -> Self { + Self { + inner, + object: None, + } + } + + /// Returns the final parsed object after the stream has yielded a `Complete` event. + #[must_use] + pub fn object(&self) -> Option<&serde_json::Value> { + self.object.as_ref() + } +} + +impl Stream for ObjectStreamResult { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let inner = self.inner.as_mut(); + match inner.poll_next(cx) { + Poll::Ready(Some(Ok(event))) => { + if let ObjectStreamEvent::Complete { ref object, .. } = event { + self.object = Some(object.clone()); + } + Poll::Ready(Some(Ok(event))) + } + other => other, + } + } +} + /// Streaming structured output with incremental JSON parsing (Section 4.6). /// /// Combines streaming with structured output: sets `response_format` to `json_schema`, @@ -878,7 +930,7 @@ pub type ObjectStream = pub async fn stream_object( params: GenerateParams, schema: serde_json::Value, -) -> Result { +) -> Result { let params = GenerateParams { response_format: Some(ResponseFormat { kind: ResponseFormatType::JsonSchema, @@ -943,7 +995,7 @@ pub async fn stream_object( }, ); - Ok(Box::pin(mapped.flatten())) + Ok(ObjectStreamResult::new(Box::pin(mapped.flatten()))) } #[cfg(test)] diff --git a/crates/unified-llm/src/providers/anthropic.rs b/crates/unified-llm/src/providers/anthropic.rs index 721c25e3c..d87008cfc 100644 --- a/crates/unified-llm/src/providers/anthropic.rs +++ b/crates/unified-llm/src/providers/anthropic.rs @@ -3,7 +3,8 @@ use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; use crate::error::SdkError; use crate::provider::{ProviderAdapter, StreamEventStream}; use crate::providers::common::{ - extract_system_prompt, parse_error_body, parse_rate_limit_headers, send_and_read_response, + extract_system_prompt, parse_error_body, parse_rate_limit_headers, parse_retry_after, + send_and_read_response, }; use crate::types::{ ContentPart, FinishReason, Message, Request, Response, ResponseFormatType, Role, StreamEvent, @@ -1148,6 +1149,7 @@ impl ProviderAdapter for Adapter { let status = http_resp.status(); if !status.is_success() { + let retry_after = parse_retry_after(http_resp.headers()); let body = http_resp.text().await.map_err(|e| SdkError::Network { message: e.to_string(), })?; @@ -1158,7 +1160,7 @@ impl ProviderAdapter for Adapter { "anthropic".to_string(), code, raw, - None, + retry_after, )); } diff --git a/crates/unified-llm/src/providers/openai.rs b/crates/unified-llm/src/providers/openai.rs index 6a0c4d096..29163a376 100644 --- a/crates/unified-llm/src/providers/openai.rs +++ b/crates/unified-llm/src/providers/openai.rs @@ -253,11 +253,15 @@ fn translate_input(messages: &[Message]) -> (Option, Vec, - content: impl Into, + content: serde_json::Value, is_error: bool, ) -> Self { let id = tool_call_id.into(); @@ -243,7 +243,7 @@ impl Message { role: Role::Tool, content: vec![ContentPart::ToolResult(ToolResult { tool_call_id: id.clone(), - content: serde_json::Value::String(content.into()), + content, is_error, image_data: None, image_media_type: None, @@ -603,6 +603,15 @@ pub struct TimeoutConfig { pub per_step: Option, } +impl From for TimeoutConfig { + fn from(total: f64) -> Self { + Self { + total: Some(total), + per_step: None, + } + } +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AdapterTimeout { pub connect: f64, @@ -797,7 +806,7 @@ mod tests { #[test] fn message_tool_result_constructor() { - let msg = Message::tool_result("call_123", "72F and sunny", false); + let msg = Message::tool_result("call_123", serde_json::Value::String("72F and sunny".into()), false); assert_eq!(msg.role, Role::Tool); assert_eq!(msg.tool_call_id, Some("call_123".to_string())); match &msg.content[0] {