diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e79463..17a6ab7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,6 +40,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. `None`) for parallel-dispatch resource-conflict detection, plus the `FnTool::with_resource_key` builder. - `MockTool::with_delay(Duration)` builder for timing-sensitive tests. +- `StructuredOutput` trait, `ResponseFormat`, `RequestOptions`, and + `StructuredError` (`structured` module): request guaranteed-schema JSON + responses from the model. Includes a lenient JSON extraction helper (handles + markdown fences/prose prefixes) and a `request_structured::()` + convenience function. +- `ApiClient::stream_messages_with_options` and + `create_message_with_options` default methods (additive — existing impls + compile unchanged). `OpenAiClient`, `AnthropicClient`, and `GeminiClient` + override both to inject the schema (OpenAI via native `response_format`, + Anthropic via forced-tool tool-forcing, Gemini via `generationConfig` + `responseMimeType` + `responseSchema`). +- `ToolOutput::structured`, `structured_value()`, and + `structured_as::()` for typed tool results that round-trip through JSON. ### Changed diff --git a/src/api.rs b/src/api.rs index 0604c4b..a2a7865 100644 --- a/src/api.rs +++ b/src/api.rs @@ -186,6 +186,74 @@ pub trait ApiClient: Send + Sync { system: Option, tools: Option>, ) -> Pin> + Send + '_>>; + + /// Streaming variant that honors [`RequestOptions`](crate::structured::RequestOptions). + /// + /// When `options` is empty (the default), delegates to + /// [`stream_messages`](ApiClient::stream_messages). When `options` + /// contains fields the client does not support (e.g. `response_format` + /// on a client that has not overridden this method), yields an + /// [`ApiError::config`] error as the first stream item. + /// `OpenAiClient`, `AnthropicClient`, and `GeminiClient` override + /// this to inject the schema. + fn stream_messages_with_options( + &self, + messages: Vec, + system: Option, + tools: Option>, + options: crate::structured::RequestOptions, + ) -> Pin> + Send + 'static>> { + if options.response_format.is_none() { + return self.stream_messages(messages, system, tools); + } + Box::pin(futures::stream::once(async { + Err(ApiError::config( + "this client does not support structured output (response_format)", + )) + })) + } + + /// Non-streaming variant that honors [`RequestOptions`](crate::structured::RequestOptions). + /// + /// When `options` is empty (the default), delegates to + /// [`create_message`](ApiClient::create_message). When `options` + /// contains fields the client does not support (e.g. `response_format` + /// on a client that has not overridden this method), returns an + /// [`ApiError::config`] error. This is the primary path for structured + /// output — a complete JSON document must be present before + /// deserialization, so callers that need a typed `T` should use this + /// method and then + /// [`StructuredOutput::from_value`](crate::structured::StructuredOutput::from_value) + /// on the extracted payload. + fn create_message_with_options( + &self, + messages: Vec, + system: Option, + tools: Option>, + options: crate::structured::RequestOptions, + ) -> Pin> + Send + '_>> { + if options.response_format.is_none() { + return self.create_message(messages, system, tools); + } + Box::pin(async { + Err(ApiError::config( + "this client does not support structured output (response_format)", + )) + }) + } + + /// Extract the structured-output payload from a provider response. + /// + /// Each provider knows its own response envelope shape. This method pulls + /// the inner JSON value that should be fed to + /// [`StructuredOutput::from_value`](crate::structured::StructuredOutput::from_value). + /// The default implementation returns the raw value as-is (for mock + /// clients and custom providers that already return the structured value + /// without an envelope). `OpenAiClient`, `AnthropicClient`, and + /// `GeminiClient` override this to navigate their respective envelopes. + fn extract_structured(&self, raw: &serde_json::Value) -> serde_json::Value { + raw.clone() + } } /// Owned, single-threaded API client handle. diff --git a/src/lib.rs b/src/lib.rs index 6b3c7a4..092c053 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -83,6 +83,7 @@ pub mod provider; pub mod reflection; pub mod runtime; pub mod stream; +pub mod structured; #[cfg(feature = "testing")] pub mod testing; pub mod tool; diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index ee799b0..656a27f 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -38,6 +38,7 @@ use crate::tool::ToolSchema; // ================================================== // Constants // ================================================== + const DEFAULT_BASE_URL: &str = "https://api.anthropic.com"; const DEFAULT_MODEL: &str = "claude-sonnet-4-20250514"; const ANTHROPIC_VERSION: &str = "2023-06-01"; @@ -48,8 +49,6 @@ const DEFAULT_MAX_TOKENS: u32 = 8192; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb -/// Maximum bytes to read from an error response body. Prevents OOM when a -/// misconfigured or malicious server returns a multi-GB body on a 4xx/5xx. const MAX_ERROR_BODY: usize = 8 * 1024; // 8 Kb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); @@ -65,30 +64,79 @@ const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); /// Also works with Anthropic-compatible endpoints such as `Z.ai` /// — use a custom `base_url` via [`AnthropicClientBuilder::base_url`]. pub struct AnthropicClient { + /// The underlying HTTP client (connection-pooled `reqwest::Client`). + /// + /// Created once at build time with the configured timeouts; reused + /// across all requests for connection pooling. http: reqwest::Client, + + /// The Anthropic API key used for authentication. + /// + /// Sent as the `x-api-key` header on every request. Set via + /// [`AnthropicClientBuilder::api_key`]. api_key: String, + + /// The base URL for API requests. + /// + /// The Messages API endpoint is `{base_url}/v1/messages`. Defaults to + /// `https://api.anthropic.com`; override for proxies or + /// Anthropic-compatible endpoints. base_url: String, + + /// The current model identifier, stored behind a mutex for runtime + /// hot-swapping. + /// + /// Changed via [`ApiClient::set_model`] when the + /// [`FallbackManager`](crate::fallback::FallbackManager) trips to a + /// fallback model. model: parking_lot::Mutex, + + /// The maximum output tokens per response. + /// + /// Anthropic requires this field on every request. Defaults to 8192. + /// Set via [`AnthropicClientBuilder::max_tokens`]. max_tokens: u32, } impl AnthropicClient { /// Create a builder for configuring an [`AnthropicClient`]. + /// + /// Returns an [`AnthropicClientBuilder`] with sensible defaults. The only + /// required field is `api_key`; everything else has a production-ready + /// default. Call `.api_key(...).build()` to finish, or chain additional + /// setters for custom configuration. + /// + /// # Example + /// + /// ```rust,no_run + /// use loopctl::provider::AnthropicClient; + /// + /// let client = AnthropicClient::builder() + /// .api_key("sk-ant-...") + /// .model("claude-sonnet-4-20250514") + /// .build() + /// .unwrap(); + /// ``` #[must_use] pub fn builder() -> AnthropicClientBuilder { AnthropicClientBuilder::default() } - /// Create from environment variables. + /// Create a client from environment variables. /// - /// Reads: - /// - `ANTHROPIC_API_KEY` — required. - /// - `ANTHROPIC_BASE_URL` — optional, defaults to `https://api.anthropic.com`. - /// - `ANTHROPIC_MODEL` — optional, defaults to `claude-sonnet-4-20250514`. + /// Reads the following variables: + /// + /// - `ANTHROPIC_API_KEY` — **required**. The API key for authentication. + /// - `ANTHROPIC_BASE_URL` — optional. Defaults to `https://api.anthropic.com`. + /// Override when targeting a proxy or Anthropic-compatible endpoint. + /// - `ANTHROPIC_MODEL` — optional. Defaults to `claude-sonnet-4-20250514`. + /// + /// This is a convenience constructor that delegates to + /// [`builder`](Self::builder) with the env vars as setter arguments. /// /// # Errors /// - /// Returns [`ApiError`] if no API key is found. + /// Returns [`ApiError`] if `ANTHROPIC_API_KEY` is not set or is empty. pub fn from_env() -> Result { let api_key = std::env::var("ANTHROPIC_API_KEY") .map_err(|_| ApiError::auth_invalid_key("ANTHROPIC_API_KEY not set"))?; @@ -103,15 +151,21 @@ impl AnthropicClient { .build() } - /// Send a POST request to the Messages endpoint. + /// Send a POST request to the Anthropic Messages API endpoint. /// - /// Shared by both [`ApiClient::stream_messages`] and - /// [`ApiClient::create_message`]. + /// Shared by [`stream_messages`](ApiClient::stream_messages), + /// [`create_message`](ApiClient::create_message), + /// [`stream_messages_with_options`](ApiClient::stream_messages_with_options), + /// and [`create_message_with_options`](ApiClient::create_message_with_options). + /// Sends the JSON body with `x-api-key` and `anthropic-version` headers. + /// On a non-success status, reads the error body (capped at + /// `MAX_ERROR_BODY` bytes) and returns it as an [`ApiError`]. /// /// # Errors /// - /// Returns [`ApiError`] if the request fails or the server - /// responds with a non-success status code. + /// Returns [`ApiError::http`] if the HTTP request fails, or + /// [`ApiError::http_with_status`] if the server responds with a + /// non-success status code. async fn post_messages( http: &reqwest::Client, url: &str, @@ -140,7 +194,11 @@ impl AnthropicClient { } } - /// Build the Messages API URL for this client. + /// Build the full URL for the Anthropic Messages API endpoint. + /// + /// Appends `/v1/messages` to the client's `base_url`. All four + /// `ApiClient` methods (`stream_messages`, `create_message`, and their + /// `*_with_options` variants) POST to this URL. fn messages_url(&self) -> String { format!("{}/v1/messages", self.base_url) } @@ -168,8 +226,38 @@ impl ApiClient for AnthropicClient { messages: Vec, system: Option, tools: Option>, + ) -> Pin> + Send + 'static>> { + self.stream_messages_with_options( + messages, + system, + tools, + crate::structured::RequestOptions::default(), + ) + } + + fn create_message( + &self, + messages: Vec, + system: Option, + tools: Option>, + ) -> Pin> + Send + '_>> { + self.create_message_with_options( + messages, + system, + tools, + crate::structured::RequestOptions::default(), + ) + } + + fn stream_messages_with_options( + &self, + messages: Vec, + system: Option, + tools: Option>, + options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { let model = self.model.lock().clone(); + let rf = options.response_format.as_ref(); let body = build_request_body( &model, &messages, @@ -177,6 +265,7 @@ impl ApiClient for AnthropicClient { tools.as_deref(), true, self.max_tokens, + rf, ); let url = self.messages_url(); let api_key = self.api_key.clone(); @@ -200,13 +289,15 @@ impl ApiClient for AnthropicClient { }) } - fn create_message( + fn create_message_with_options( &self, messages: Vec, system: Option, tools: Option>, + options: crate::structured::RequestOptions, ) -> Pin> + Send + '_>> { let model = self.model.lock().clone(); + let rf = options.response_format.as_ref(); let body = build_request_body( &model, &messages, @@ -214,6 +305,7 @@ impl ApiClient for AnthropicClient { tools.as_deref(), false, self.max_tokens, + rf, ); let url = self.messages_url(); Box::pin(async move { @@ -232,6 +324,21 @@ impl ApiClient for AnthropicClient { serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) }) } + + fn extract_structured(&self, raw: &Value) -> Value { + raw.get("content") + .and_then(|c| c.as_array()) + .and_then(|blocks| { + blocks.iter().find_map(|block| { + if block.get("type").and_then(|t| t.as_str()) == Some("tool_use") { + block.get("input").cloned() + } else { + None + } + }) + }) + .unwrap_or_else(|| raw.clone()) + } } // ================================================== @@ -239,12 +346,44 @@ impl ApiClient for AnthropicClient { // ================================================== /// Builder for [`AnthropicClient`]. +/// +/// Created via [`AnthropicClientBuilder::default`] or +/// [`AnthropicClient::builder`]. All fields have sensible defaults except +/// `api_key`, which must be set before [`build`](Self::build). pub struct AnthropicClientBuilder { + /// The Anthropic API key used for authentication. + /// + /// Required — [`build`](Self::build) returns an error if this is `None`. + /// The key is sent as the `x-api-key` header on every request. Set via + /// [`api_key`](Self::api_key) on the builder, or read from + /// `ANTHROPIC_API_KEY` via [`AnthropicClient::from_env`]. api_key: Option, + + /// The base URL for API requests. + /// + /// Defaults to `https://api.anthropic.com`. Override when targeting a + /// proxy, gateway, or Anthropic-compatible endpoint (e.g. Z.AI). base_url: String, + + /// The default model identifier (e.g. `claude-sonnet-4-20250514`). + /// + /// Can be changed at runtime via [`AnthropicClient::set_model`]. model: String, + + /// The maximum number of output tokens per response. + /// + /// Anthropic requires this field on every request. Defaults to 8192. max_tokens: u32, + + /// The total HTTP request timeout (connect + response + body). + /// + /// Bounds the entire request lifecycle. Defaults to 120 seconds. timeout: Duration, + + /// The TCP connection establishment timeout (including TLS handshake). + /// + /// Separate from the total timeout so a slow-connecting server can be + /// detected faster than a slow-responding one. Defaults to 10 seconds. connect_timeout: Duration, } @@ -262,30 +401,44 @@ impl Default for AnthropicClientBuilder { } impl AnthropicClientBuilder { - /// Set the API key. + /// Set the API key for authentication. + /// + /// Required — [`build`](Self::build) returns an error if this is not set. + /// The key is sent as the `x-api-key` header on every request. #[must_use] pub fn api_key(mut self, key: impl Into) -> Self { self.api_key = Some(key.into()); self } - /// Set the base URL (e.g. `https://api.z.ai/api/anthropic`). + /// Set the base URL for API requests. + /// + /// Defaults to `https://api.anthropic.com`. Override when targeting a + /// proxy, gateway, or Anthropic-compatible endpoint (e.g. Z.AI at + /// `https://api.z.ai/api/anthropic`). #[must_use] pub fn base_url(mut self, url: impl Into) -> Self { self.base_url = url.into(); self } - /// Set the model name (e.g. `claude-sonnet-4-20250514`). + /// Set the default model identifier. + /// + /// The model string is sent as the `model` field on every request. Can be + /// changed at runtime via [`AnthropicClient::set_model`] (e.g. when the + /// [`FallbackManager`](crate::fallback::FallbackManager) trips). #[must_use] pub fn model(mut self, model: impl Into) -> Self { self.model = model.into(); self } - /// Set the maximum output tokens per response. + /// Set the maximum number of output tokens per response. /// - /// Anthropic requires this field. Defaults to 8192. + /// Anthropic requires this field on every request — unlike OpenAI, which + /// defaults it server-side. It bounds the length of a single model + /// response. Defaults to 8192. Increase for long-form generation; + /// decrease to cap cost on simple queries. #[must_use] pub fn max_tokens(mut self, tokens: u32) -> Self { self.max_tokens = tokens; @@ -313,11 +466,15 @@ impl AnthropicClientBuilder { self } - /// Build the client. + /// Construct the [`AnthropicClient`] from the builder's configuration. + /// + /// Creates the internal `reqwest::Client` with the configured timeouts, + /// and validates that an API key was provided. /// /// # Errors /// - /// Returns [`ApiError`] if no API key was set. + /// Returns [`ApiError`] if no API key was set via + /// [`api_key`](Self::api_key). pub fn build(self) -> Result { let api_key = self .api_key @@ -353,8 +510,23 @@ fn build_request_body( tools: Option<&[ToolSchema]>, stream: bool, max_tokens: u32, + response_format: Option<&crate::structured::ResponseFormat>, ) -> Value { let msgs: Vec = messages.iter().map(convert_message).collect(); + let (tools_val, tool_choice) = if let Some(rf) = response_format { + let forced_tool = serde_json::json!({ + "name": rf.name, + "description": "Return the result via this tool", + "input_schema": rf.schema, + }); + let choice = serde_json::json!({ + "type": "tool", + "name": rf.name, + }); + (Some(vec![forced_tool]), Some(choice)) + } else { + (tools.map(convert_tools), None) + }; let mut body = serde_json::json!({ "model": model, @@ -362,16 +534,23 @@ fn build_request_body( "messages": msgs, "system": system.unwrap_or(""), "stream": stream, - "tools": tools.map(convert_tools), + "tools": tools_val, }); // Remove `tools` if None so we don't send a null field. - if tools.is_none() { + if tools_val.is_none() { if let Some(obj) = body.as_object_mut() { obj.remove("tools"); } } + // Inject tool_choice when forcing a structured-output tool. + if let Some(choice) = tool_choice { + if let Some(obj) = body.as_object_mut() { + obj.insert("tool_choice".to_string(), choice); + } + } + body } @@ -443,7 +622,16 @@ fn convert_message(m: &Message) -> Value { } } -/// Convert tool schemas into the Anthropic `tools` array shape. +/// Convert framework tool schemas into the Anthropic `tools` array shape. +/// +/// Each [`ToolSchema`] becomes a JSON object with `name`, `description`, and +/// `input_schema` — the three fields Anthropic's tool-use API expects. The +/// `input_schema` is passed through verbatim from the framework's schema (a +/// JSON Schema Draft 07 object), since Anthropic validates it server-side. +/// +/// When structured output is active (`response_format` set), this function is +/// not called — instead, [`build_request_body`] synthesizes a single forced +/// tool whose `input_schema` is the target `ResponseFormat::schema`. fn convert_tools(tools: &[ToolSchema]) -> Vec { tools .iter() @@ -824,7 +1012,15 @@ mod tests { #[test] fn request_body_user_text_single_string() { let msgs = vec![Message::user("hello")]; - let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + None, + ); let messages = body["messages"].as_array().unwrap(); assert_eq!(messages.len(), 1); @@ -842,6 +1038,7 @@ mod tests { None, false, DEFAULT_MAX_TOKENS, + None, ); assert_eq!(body["system"], "be brief"); } @@ -849,7 +1046,15 @@ mod tests { #[test] fn request_body_system_empty_when_none() { let msgs = vec![Message::user("hi")]; - let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + None, + ); assert_eq!(body["system"], ""); } @@ -863,6 +1068,7 @@ mod tests { None, false, DEFAULT_MAX_TOKENS, + None, ); assert_eq!(body["model"], "claude-sonnet-4"); } @@ -870,14 +1076,30 @@ mod tests { #[test] fn request_body_max_tokens() { let msgs = vec![Message::user("hi")]; - let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + None, + ); assert_eq!(body["max_tokens"], DEFAULT_MAX_TOKENS); } #[test] fn request_body_user_role() { let msgs = vec![Message::user("hi")]; - let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + None, + ); assert_eq!(body["messages"][0]["role"], "user"); } @@ -887,7 +1109,15 @@ mod tests { Role::Assistant, vec![MessagePart::text("hello")], )]; - let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + None, + ); assert_eq!(body["messages"][0]["role"], "assistant"); assert_eq!(body["messages"][0]["content"], "hello"); } @@ -902,7 +1132,15 @@ mod tests { input: serde_json::json!({"msg": "hi"}), }], )]; - let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + None, + ); let msg = &body["messages"][0]; assert_eq!(msg["role"], "assistant"); @@ -923,7 +1161,15 @@ mod tests { is_error: None, }], )]; - let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + None, + ); let msg = &body["messages"][0]; assert_eq!(msg["role"], "user"); @@ -948,6 +1194,7 @@ mod tests { Some(&tools), false, DEFAULT_MAX_TOKENS, + None, ); let tools_arr = body["tools"].as_array().unwrap(); @@ -960,7 +1207,15 @@ mod tests { #[test] fn request_body_tools_absent_when_none() { let msgs = vec![Message::user("hi")]; - let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + None, + ); assert!(body.get("tools").is_none()); } @@ -977,7 +1232,15 @@ mod tests { }, ], )]; - let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + None, + ); let msg = &body["messages"][0]; assert_eq!(msg["role"], "assistant"); @@ -994,7 +1257,15 @@ mod tests { Message::new(Role::Assistant, vec![MessagePart::text("hi")]), Message::user("bye"), ]; - let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + None, + ); let messages = body["messages"].as_array().unwrap(); assert_eq!(messages.len(), 3); @@ -1338,4 +1609,106 @@ mod tests { fn max_response_body_is_ten_mb() { assert_eq!(MAX_RESPONSE_BODY, 10 * 1024 * 1024); } + + #[test] + fn request_body_response_format_forces_tool() { + let msgs = vec![Message::user("classify this")]; + let rf = crate::structured::ResponseFormat::new( + "action", + serde_json::json!({"type": "object", "properties": {"x": {"type": "string"}}}), + ); + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + Some(&rf), + ); + + // Exactly one forced tool with the schema's name + input_schema. + let tools = body["tools"].as_array().expect("tools should be an array"); + assert_eq!(tools.len(), 1, "should have exactly one forced tool"); + assert_eq!(tools[0]["name"], "action"); + assert_eq!(tools[0]["input_schema"], rf.schema); + + // tool_choice forces the named tool. + assert_eq!(body["tool_choice"]["type"], "tool"); + assert_eq!(body["tool_choice"]["name"], "action"); + } + + #[test] + fn request_body_response_format_suppresses_caller_tools() { + let msgs = vec![Message::user("hi")]; + let caller_tool = ToolSchema { + tool: "read".into(), + description: "Read a file".into(), + input_schema: serde_json::json!({"type": "object"}), + }; + let rf = + crate::structured::ResponseFormat::new("result", serde_json::json!({"type": "object"})); + let body = build_request_body( + "claude-3", + &msgs, + None, + Some(&[caller_tool]), + false, + DEFAULT_MAX_TOKENS, + Some(&rf), + ); + + // The forced tool replaces the caller's tools — not appended. + let tools = body["tools"].as_array().expect("tools should be an array"); + assert_eq!(tools.len(), 1); + assert_eq!( + tools[0]["name"], "result", + "caller tools should be suppressed" + ); + } + + #[test] + fn request_body_no_response_format_has_no_tool_choice() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body( + "claude-3", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + None, + ); + assert!( + body.get("tool_choice").is_none(), + "tool_choice should only appear with response_format" + ); + } + + #[test] + fn extract_structured_from_tool_use_input() { + let client = AnthropicClient::builder().api_key("test").build().unwrap(); + let raw = serde_json::json!({ + "content": [{ + "type": "tool_use", + "name": "action", + "input": {"tool": "write", "args": {}} + }] + }); + let value = client.extract_structured(&raw); + assert_eq!(value["tool"], "write"); + } + + #[test] + fn extract_structured_text_only_falls_back_to_raw() { + let client = AnthropicClient::builder().api_key("test").build().unwrap(); + let raw = serde_json::json!({ + "id": "msg_1", + "model": "claude-3", + "content": [{"type": "text", "text": "I cannot do that."}] + }); + let value = client.extract_structured(&raw); + // No tool_use block → returns the raw envelope; T::from_value fails. + assert_eq!(value["id"], "msg_1"); + } } diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index 52abf99..32e17fc 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -47,8 +47,6 @@ const TEXT_PART_INDEX: usize = 0; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb -/// Maximum bytes to read from an error response body. Prevents OOM when a -/// misconfigured or malicious server returns a multi-GB body on a 4xx/5xx. const MAX_ERROR_BODY: usize = 8 * 1024; // 8 Kb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); @@ -61,26 +59,70 @@ const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); /// Implements [`ApiClient`] by translating between the framework's /// [`StreamEvent`] protocol and the Gemini Streaming Generate Content API. pub struct GeminiClient { + /// The underlying HTTP client (connection-pooled `reqwest::Client`). + /// + /// Created once at build time with the configured timeouts; reused + /// across all requests for connection pooling. http: reqwest::Client, + + /// The Gemini API key used for authentication. + /// + /// Sent as the `x-goog-api-key` header on every request. Set via + /// [`GeminiClientBuilder::api_key`]. api_key: String, + + /// The base URL for API requests. + /// + /// The streaming endpoint is `{base_url}/models/{model}:streamGenerateContent` + /// and the non-streaming endpoint is `{base_url}/models/{model}:generateContent`. + /// Defaults to `https://generativelanguage.googleapis.com/v1beta`. base_url: String, + + /// The current model identifier, stored behind a mutex for runtime + /// hot-swapping. + /// + /// Changed via [`ApiClient::set_model`] when the + /// [`FallbackManager`](crate::fallback::FallbackManager) trips to a + /// fallback model. model: parking_lot::Mutex, } impl GeminiClient { /// Create a builder for configuring a [`GeminiClient`]. + /// + /// Returns a [`GeminiClientBuilder`] with sensible defaults. The only + /// required field is `api_key`; everything else has a production-ready + /// default. Call `.api_key(...).build()` to finish, or chain additional + /// setters for custom configuration. + /// + /// # Example + /// + /// ```rust,no_run + /// use loopctl::provider::GeminiClient; + /// + /// let client = GeminiClient::builder() + /// .api_key("AI...") + /// .model("gemini-2.0-flash") + /// .build() + /// .unwrap(); + /// ``` #[must_use] pub fn builder() -> GeminiClientBuilder { GeminiClientBuilder::default() } - /// Create from environment variables. + /// Create a client from environment variables. /// - /// Reads: - /// - `GEMINI_API_KEY` (or `GOOGLE_API_KEY`) — required. - /// - `GEMINI_BASE_URL` — optional, defaults to + /// Reads the following variables: + /// + /// - `GEMINI_API_KEY` (or `GOOGLE_API_KEY`) — **required**. The API key + /// for authentication. + /// - `GEMINI_BASE_URL` — optional. Defaults to /// `https://generativelanguage.googleapis.com/v1beta`. - /// - `GEMINI_MODEL` — optional, defaults to `gemini-2.0-flash`. + /// - `GEMINI_MODEL` — optional. Defaults to `gemini-2.0-flash`. + /// + /// This is a convenience constructor that delegates to + /// [`builder`](Self::builder) with the env vars as setter arguments. /// /// # Errors /// @@ -101,8 +143,8 @@ impl GeminiClient { /// Build the streaming Generate Content URL. /// - /// Gemini puts the model in the URL path and the API key as a query - /// parameter rather than using headers. + /// Gemini puts the model in the URL path. The API key is sent via + /// the `x-goog-api-key` header, not as a query parameter. fn stream_url(&self) -> String { let model = self.model.lock().clone(); format!( @@ -111,7 +153,13 @@ impl GeminiClient { ) } - /// Build the non-streaming Generate Content URL. + /// Build the full URL for the Gemini non-streaming Generate Content + /// endpoint. + /// + /// Constructs `{base_url}/models/{model}:generateContent`. The API key + /// is sent via the `x-goog-api-key` header, not in the URL. + /// Used by [`ApiClient::create_message`] and its `*_with_options` + /// variant. fn generate_url(&self) -> String { let model = self.model.lock().clone(); format!("{}/models/{}:generateContent", self.base_url, model) @@ -177,7 +225,7 @@ impl ApiClient for GeminiClient { system: Option, tools: Option>, ) -> Pin> + Send + 'static>> { - let body = build_request_body(&messages, system.as_deref(), tools.as_deref()); + let body = build_request_body(&messages, system.as_deref(), tools.as_deref(), None); let url = self.stream_url(); let http = self.http.clone(); let api_key = self.api_key.clone(); @@ -206,7 +254,7 @@ impl ApiClient for GeminiClient { system: Option, tools: Option>, ) -> Pin> + Send + '_>> { - let body = build_request_body(&messages, system.as_deref(), tools.as_deref()); + let body = build_request_body(&messages, system.as_deref(), tools.as_deref(), None); let url = self.generate_url(); Box::pin(async move { @@ -225,18 +273,116 @@ impl ApiClient for GeminiClient { serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) }) } -} + fn stream_messages_with_options( + &self, + messages: Vec, + system: Option, + tools: Option>, + options: crate::structured::RequestOptions, + ) -> Pin> + Send + 'static>> { + let rf = options.response_format.as_ref(); + let body = build_request_body(&messages, system.as_deref(), tools.as_deref(), rf); + let url = self.stream_url(); + let http = self.http.clone(); + let api_key = self.api_key.clone(); + + Box::pin(async_stream::try_stream! { + let resp = Self::post_content(&http, &url, &api_key, &body).await?; + let mut sse = SseReader::from_response(resp); + let mut emitter = StreamEmitter::default(); + + while let Some(data) = sse.next_data().await? { + emitter.process_chunk(&data); + for ev in emitter.drain() { + yield ev; + } + } + + for ev in emitter.finish() { + yield ev; + } + }) + } + + fn create_message_with_options( + &self, + messages: Vec, + system: Option, + tools: Option>, + options: crate::structured::RequestOptions, + ) -> Pin> + Send + '_>> { + let response_format = options.response_format.as_ref(); + let body = build_request_body( + &messages, + system.as_deref(), + tools.as_deref(), + response_format, + ); + let url = self.generate_url(); + Box::pin(async move { + let resp = Self::post_content(&self.http, &url, &self.api_key, &body).await?; + let resp = resp + .bytes() + .await + .map_err(|e| ApiError::http(e.to_string()))?; + if resp.len() > MAX_RESPONSE_BODY { + return Err(ApiError::http(format!( + "response body too large: {} bytes (max {})", + resp.len(), + MAX_RESPONSE_BODY + ))); + } + serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) + }) + } + + fn extract_structured(&self, raw: &Value) -> Value { + let Some(text) = raw + .pointer("/candidates/0/content/parts/0/text") + .and_then(serde_json::Value::as_str) + else { + return raw.clone(); + }; + crate::structured::parse_json_lenient(text) + .unwrap_or_else(|| Value::String(text.to_string())) + } +} // ================================================== // Builder // ================================================== /// Builder for [`GeminiClient`]. +/// +/// Created via [`GeminiClientBuilder::default`] or +/// [`GeminiClient::builder`]. All fields have sensible defaults except +/// `api_key`, which must be set before [`build`](Self::build). pub struct GeminiClientBuilder { + /// The Gemini API key for authentication (required). + /// + /// Must be set before building. Sent as the `x-goog-api-key` header on + /// every request. api_key: Option, + + /// The base URL for API requests. + /// + /// Defaults to `https://generativelanguage.googleapis.com/v1beta`. base_url: String, + + /// The default model identifier. + /// + /// Can be changed at runtime via [`GeminiClient::set_model`]. model: String, + + /// The total HTTP request timeout (connect + response + body). + /// + /// Bounds the entire request lifecycle. Defaults to 120 seconds. timeout: Duration, + + /// The TCP connection establishment timeout (including TLS handshake). + /// + /// Separate from the total timeout so a slow-connecting server can be + /// detected faster. Defaults to 10 seconds. connect_timeout: Duration, } @@ -253,21 +399,32 @@ impl Default for GeminiClientBuilder { } impl GeminiClientBuilder { - /// Set the API key. + /// Set the API key for authentication. + /// + /// Required — [`build`](Self::build) returns an error if this is not set. + /// The key is sent as the `x-goog-api-key` header on every request. #[must_use] pub fn api_key(mut self, key: impl Into) -> Self { self.api_key = Some(key.into()); self } - /// Set the base URL. + /// Set the base URL for API requests. + /// + /// Defaults to `https://generativelanguage.googleapis.com/v1beta`. + /// Override when targeting a proxy or Google AI-compatible endpoint. #[must_use] pub fn base_url(mut self, url: impl Into) -> Self { self.base_url = url.into(); self } - /// Set the model name (e.g. `gemini-2.0-flash`). + /// Set the default model identifier. + /// + /// The model string is embedded in the request URL (e.g. + /// `/models/{model}:generateContent`). Can be changed at runtime via + /// [`GeminiClient::set_model`] (e.g. when the + /// [`FallbackManager`](crate::fallback::FallbackManager) trips). #[must_use] pub fn model(mut self, model: impl Into) -> Self { self.model = model.into(); @@ -331,6 +488,7 @@ fn build_request_body( messages: &[Message], system: Option<&str>, tools: Option<&[ToolSchema]>, + response_format: Option<&crate::structured::ResponseFormat>, ) -> Value { let contents: Vec = messages.iter().map(convert_message).collect(); let mut body = serde_json::json!({ "contents": contents }); @@ -341,10 +499,23 @@ fn build_request_body( serde_json::json!({"parts": [{"text": sys}]}), ); } - if let Some(tool_list) = tools { + + if response_format.is_none() { + if let Some(tool_list) = tools { + obj.insert( + "tools".into(), + serde_json::json!([{"functionDeclarations": convert_tools(tool_list)}]), + ); + } + } + + if let Some(rf) = response_format { obj.insert( - "tools".into(), - serde_json::json!([{"functionDeclarations": convert_tools(tool_list)}]), + "generationConfig".into(), + serde_json::json!({ + "responseMimeType": "application/json", + "responseJsonSchema": rf.schema, + }), ); } } @@ -389,7 +560,14 @@ fn convert_part(p: &MessagePart) -> Option { } } -/// Convert tool schemas into the Gemini `functionDeclarations` array. +/// Convert framework tool schemas into the Gemini `functionDeclarations` +/// array shape. +/// +/// Each [`ToolSchema`] becomes a JSON object with `name`, `description`, and +/// `parameters` — the fields Gemini's function-calling API expects. When +/// structured output is active (`response_format` set), this function is not +/// called — [`build_request_body`] injects `generationConfig.responseJsonSchema` +/// instead, and `tools` is suppressed. fn convert_tools(tools: &[ToolSchema]) -> Vec { tools .iter() @@ -418,7 +596,13 @@ struct SseReader { } impl SseReader { - /// Wrap a streaming HTTP response. + /// Wrap a streaming HTTP response into an SSE reader. + /// + /// Takes the response body's byte stream and converts it into a + /// line-oriented reader that yields SSE event data. Used by + /// [`stream_messages`](crate::api::ApiClient::stream_messages) and its + /// `*_with_options` variant to parse Gemini's streaming + /// `streamGenerateContent` responses. fn from_response(resp: Response) -> Self { let bytes = resp.bytes_stream().map(|res| { res.map(|b| String::from_utf8_lossy(&b).into_owned()) @@ -474,7 +658,12 @@ impl SseReader { } } - /// Pop the first `\n`-terminated line from the buffer, if present. + /// Pop the first `\n`-terminated line from the internal buffer. + /// + /// Returns the line (trimmed) if a newline is present, and removes it + /// (plus the newline) from the buffer. Returns `None` if the buffer + /// does not yet contain a complete line — the caller should wait for + /// more bytes from the HTTP stream. fn take_line(&mut self) -> Option { let pos = self.buf.find('\n')?; let line = self.buf[..pos].trim().to_string(); @@ -504,7 +693,14 @@ struct StreamEmitter { } impl StreamEmitter { - /// Process a single Gemini SSE chunk, appending events to the queue. + /// Process a single Gemini SSE JSON chunk into stream events. + /// + /// On the first call, emits [`MessageStart`](StreamEvent::MessageStart). + /// Then delegates to the three extractors: [`extract_text`](Self::extract_text) + /// for text deltas, [`extract_function_call`](Self::extract_function_call) + /// for tool calls, and [`extract_finish_reason`](Self::extract_finish_reason) + /// for the terminal stop signal. Events accumulate in the internal queue + /// until [`drain`](Self::drain) is called. fn process_chunk(&mut self, json: &Value) { if !self.started { self.started = true; @@ -522,7 +718,11 @@ impl StreamEmitter { self.extract_finish_reason(json); } - /// Extract text delta from `candidates[0].content.parts[0].text`. + /// Extract a text delta from the chunk and emit an `IndexedDelta` event. + /// + /// Reads `candidates[0].content.parts[0].text`. If the text is non-empty, + /// pushes a [`DeltaPart::Text`](crate::stream::DeltaPart::Text) event at + /// index 0. Does nothing if the path is absent or the text is empty. fn extract_text(&mut self, json: &Value) { if let Some(text) = json .pointer("/candidates/0/content/parts/0/text") @@ -539,7 +739,15 @@ impl StreamEmitter { } } - /// Extract function call from `candidates[0].content.parts[0].functionCall`. + /// Extract a function (tool) call from the chunk and emit the + /// corresponding part-start and input-json events. + /// + /// Reads `candidates[0].content.parts[0].functionCall` for the tool + /// `name` and `args`. Emits a [`PartStart`](StreamEvent::PartStart) + /// with a [`ToolCall`](crate::message::MessagePart::ToolCall) part, + /// followed by an [`InputJson`](crate::stream::DeltaPart::InputJson) + /// delta carrying the serialized arguments. Does nothing if no function + /// call is present in the chunk. fn extract_function_call(&mut self, json: &Value) { if let Some(func_call) = json.pointer("/candidates/0/content/parts/0/functionCall") { let name = func_call @@ -567,7 +775,12 @@ impl StreamEmitter { } } - /// Extract finish reason and emit stop events. + /// Extract the finish reason from the chunk and emit stop events. + /// + /// Reads `candidates[0].finishReason` (e.g. `"STOP"`, `"MAX_TOKENS"`). + /// When present, emits [`PartStop`](StreamEvent::PartStop) followed by a + /// [`MessageDelta`](StreamEvent::MessageDelta) carrying the mapped + /// [`StreamStopReason`]. Does nothing if the chunk has no finish reason. fn extract_finish_reason(&mut self, json: &Value) { let Some(reason) = json .pointer("/candidates/0/finishReason") @@ -589,12 +802,22 @@ impl StreamEmitter { })); } - /// Drain all pending events. + /// Drain all pending events from the internal queue. + /// + /// Returns the accumulated [`StreamEvent`]s and clears the queue. + /// Called by the stream loop after each chunk is processed, so events + /// are yielded to the consumer promptly rather than buffered until the + /// end of the stream. fn drain(&mut self) -> Vec { std::mem::take(&mut self.pending) } - /// Emit the terminal [`MessageStop`] if the stream was started. + /// Finalize the stream, emitting the terminal + /// [`MessageStop`](StreamEvent::MessageStop) if one was started. + /// + /// Drains any remaining pending events and appends the stop event. + /// Safe to call exactly once at the end of the stream; subsequent calls + /// return an empty vec (the `finished` flag guards against double-stop). fn finish(&mut self) -> Vec { let mut out = self.drain(); if self.started && !self.finished { @@ -604,6 +827,12 @@ impl StreamEmitter { out } + /// Push an event onto the internal pending queue. + /// + /// Events are held until [`drain`](Self::drain) is called. This is the + /// single write point — all extractors (`extract_text`, + /// `extract_function_call`, `extract_finish_reason`) and the lifecycle + /// methods (`process_chunk`, `finish`) funnel through here. fn push(&mut self, ev: StreamEvent) { self.pending.push(ev); } @@ -621,7 +850,7 @@ mod tests { #[test] fn request_body_user_text() { let msgs = vec![Message::user("hello")]; - let body = build_request_body(&msgs, None, None); + let body = build_request_body(&msgs, None, None, None); let contents = body["contents"].as_array().unwrap(); assert_eq!(contents.len(), 1); @@ -633,7 +862,7 @@ mod tests { #[test] fn request_body_includes_system_instruction() { let msgs = vec![Message::user("hi")]; - let body = build_request_body(&msgs, Some("be brief"), None); + let body = build_request_body(&msgs, Some("be brief"), None, None); let sys = &body["systemInstruction"]; assert!(sys.is_object()); @@ -643,7 +872,7 @@ mod tests { #[test] fn request_body_no_system_instruction_when_none() { let msgs = vec![Message::user("hi")]; - let body = build_request_body(&msgs, None, None); + let body = build_request_body(&msgs, None, None, None); assert!(body.get("systemInstruction").is_none()); } @@ -653,14 +882,14 @@ mod tests { Role::Assistant, vec![MessagePart::text("hello")], )]; - let body = build_request_body(&msgs, None, None); + let body = build_request_body(&msgs, None, None, None); assert_eq!(body["contents"][0]["role"], "model"); } #[test] fn request_body_user_role() { let msgs = vec![Message::user("hi")]; - let body = build_request_body(&msgs, None, None); + let body = build_request_body(&msgs, None, None, None); assert_eq!(body["contents"][0]["role"], "user"); } @@ -674,7 +903,7 @@ mod tests { input: serde_json::json!({"msg": "hi"}), }], )]; - let body = build_request_body(&msgs, None, None); + let body = build_request_body(&msgs, None, None, None); let parts = body["contents"][0]["parts"].as_array().unwrap(); assert_eq!(parts[0]["functionCall"]["name"], "echo"); @@ -691,7 +920,7 @@ mod tests { is_error: None, }], )]; - let body = build_request_body(&msgs, None, None); + let body = build_request_body(&msgs, None, None, None); let parts = body["contents"][0]["parts"].as_array().unwrap(); assert_eq!(parts[0]["functionResponse"]["name"], "call_1"); @@ -709,7 +938,7 @@ mod tests { description: "Search the web".into(), input_schema: serde_json::json!({"type": "object"}), }]; - let body = build_request_body(&msgs, None, Some(&tools)); + let body = build_request_body(&msgs, None, Some(&tools), None); let tools_arr = body["tools"].as_array().unwrap(); assert_eq!(tools_arr.len(), 1); @@ -722,7 +951,7 @@ mod tests { #[test] fn request_body_no_tools_when_none() { let msgs = vec![Message::user("hi")]; - let body = build_request_body(&msgs, None, None); + let body = build_request_body(&msgs, None, None, None); assert!(body.get("tools").is_none()); } @@ -733,7 +962,7 @@ mod tests { Message::new(Role::Assistant, vec![MessagePart::text("hi")]), Message::user("bye"), ]; - let body = build_request_body(&msgs, None, None); + let body = build_request_body(&msgs, None, None, None); let contents = body["contents"].as_array().unwrap(); assert_eq!(contents.len(), 3); @@ -1135,4 +1364,83 @@ mod tests { fn max_response_body_is_ten_mb() { assert_eq!(MAX_RESPONSE_BODY, 10 * 1024 * 1024); } + + #[test] + fn request_body_response_format_injects_generation_config() { + let msgs = vec![Message::user("hi")]; + let rf = crate::structured::ResponseFormat::new( + "result", + serde_json::json!({"type": "object", "properties": {"x": {"type": "string"}}}), + ); + let body = build_request_body(&msgs, None, None, Some(&rf)); + + assert_eq!( + body["generationConfig"]["responseMimeType"], + "application/json" + ); + assert_eq!( + body["generationConfig"]["responseJsonSchema"], + serde_json::json!({"type": "object", "properties": {"x": {"type": "string"}}}) + ); + } + + #[test] + fn request_body_response_format_absent_when_none() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body(&msgs, None, None, None); + assert!( + body.get("generationConfig").is_none(), + "generationConfig should be absent when no response_format" + ); + } + + #[test] + fn request_body_response_format_suppresses_tools() { + let msgs = vec![Message::user("hi")]; + let caller_tool = ToolSchema { + tool: "read".into(), + description: "Read".into(), + input_schema: serde_json::json!({"type": "object"}), + }; + let rf = + crate::structured::ResponseFormat::new("result", serde_json::json!({"type": "object"})); + let body = build_request_body(&msgs, None, Some(&[caller_tool]), Some(&rf)); + + assert!( + body.get("tools").is_none(), + "tools should be suppressed when response_format is set" + ); + assert!(body.get("generationConfig").is_some()); + } + + #[test] + fn extract_structured_from_text_field() { + let client = GeminiClient::builder().api_key("test").build().unwrap(); + let raw = serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{ + "text": r#"{"tool": "write", "args": {}}"# + }] + } + }] + }); + let value = client.extract_structured(&raw); + assert_eq!(value["tool"], "write"); + } + + #[test] + fn extract_structured_prose_falls_back_to_raw() { + let client = GeminiClient::builder().api_key("test").build().unwrap(); + let raw = serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{"text": "I cannot produce that."}] + } + }] + }); + let value = client.extract_structured(&raw); + // Prose text not parseable as JSON → falls back to the string value. + assert_eq!(value, serde_json::json!("I cannot produce that.")); + } } diff --git a/src/provider/openai.rs b/src/provider/openai.rs index 8a63ef5..e71d337 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -50,8 +50,6 @@ const TEXT_PART_INDEX: usize = 0; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb -/// Maximum bytes to read from an error response body. Prevents OOM when a -/// misconfigured or malicious server returns a multi-GB body on a 4xx/5xx. const MAX_ERROR_BODY: usize = 8 * 1024; // 8 Kb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); @@ -67,26 +65,70 @@ const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); /// Works with any OpenAI-compatible endpoint. Use a custom `base_url` /// to target `DeepSeek`, `Grok`, Ollama, `vLLM`, or other compatible APIs. pub struct OpenAiClient { + /// The underlying HTTP client (connection-pooled `reqwest::Client`). + /// + /// Created once at build time with the configured timeouts; reused + /// across all requests for connection pooling. http: reqwest::Client, + + /// The API key used for authentication. + /// + /// Sent as the `Authorization: Bearer ` header on every request. + /// Set via [`OpenAiClientBuilder::api_key`]. api_key: String, + + /// The base URL for API requests. + /// + /// The chat-completions endpoint is `{base_url}/chat/completions`. + /// Defaults to `https://api.openai.com/v1`; override for + /// `DeepSeek`, `Grok`, `Ollama`, `vLLM`, or other compatible endpoints. base_url: String, + + /// The current model identifier, stored behind a mutex for runtime + /// hot-swapping. + /// + /// Changed via [`ApiClient::set_model`] when the + /// [`FallbackManager`](crate::fallback::FallbackManager) trips to a + /// fallback model. model: parking_lot::Mutex, } impl OpenAiClient { /// Create a builder for configuring an [`OpenAiClient`]. + /// + /// Returns an [`OpenAiClientBuilder`] with sensible defaults. The only + /// required field is `api_key`; everything else has a production-ready + /// default. Call `.api_key(...).build()` to finish, or chain additional + /// setters for custom configuration. + /// + /// # Example + /// + /// ```rust,no_run + /// use loopctl::provider::OpenAiClient; + /// + /// let client = OpenAiClient::builder() + /// .api_key("sk-...") + /// .model("gpt-4o") + /// .build() + /// .unwrap(); + /// ``` #[must_use] pub fn builder() -> OpenAiClientBuilder { OpenAiClientBuilder::default() } - /// Create from environment variables. + /// Create a client from environment variables. /// - /// Reads: - /// - `OPENAI_API_KEY` (or `API_KEY`) — required. - /// - `OPENAI_BASE_URL` (or `BASE_URL`) — optional, defaults to - /// `https://api.openai.com/v1`. - /// - `OPENAI_MODEL` (or `MODEL`) — optional, defaults to `gpt-4o`. + /// Reads the following variables: + /// + /// - `OPENAI_API_KEY` (or `API_KEY`) — **required**. The API key for + /// authentication. + /// - `OPENAI_BASE_URL` (or `BASE_URL`) — optional. Defaults to + /// `https://api.openai.com/v1`. Override for OpenAI-compatible endpoints. + /// - `OPENAI_MODEL` (or `MODEL`) — optional. Defaults to `gpt-4o`. + /// + /// This is a convenience constructor that delegates to + /// [`builder`](Self::builder) with the env vars as setter arguments. /// /// # Errors /// @@ -111,7 +153,11 @@ impl OpenAiClient { .build() } - /// Build the chat-completions URL for this client. + /// Build the full URL for the OpenAI chat-completions endpoint. + /// + /// Appends `/chat/completions` to the client's `base_url`. All four + /// `ApiClient` methods (`stream_messages`, `create_message`, and their + /// `*_with_options` variants) POST to this URL. fn completions_url(&self) -> String { format!("{}/chat/completions", self.base_url) } @@ -179,7 +225,7 @@ impl ApiClient for OpenAiClient { tools: Option>, ) -> Pin> + Send + 'static>> { let model = self.model.lock().clone(); - let body = RequestBody::build(&model, &messages, system.as_deref(), tools.as_deref()); + let body = RequestBody::build(&model, &messages, system.as_deref(), tools.as_deref(), None); let url = self.completions_url(); let api_key = self.api_key.clone(); let http = self.http.clone(); @@ -212,7 +258,7 @@ impl ApiClient for OpenAiClient { tools: Option>, ) -> Pin> + Send + '_>> { let model = self.model.lock().clone(); - let body = RequestBody::build(&model, &messages, system.as_deref(), tools.as_deref()); + let body = RequestBody::build(&model, &messages, system.as_deref(), tools.as_deref(), None); let url = self.completions_url(); Box::pin(async move { @@ -233,18 +279,126 @@ impl ApiClient for OpenAiClient { serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) }) } -} + fn stream_messages_with_options( + &self, + messages: Vec, + system: Option, + tools: Option>, + options: crate::structured::RequestOptions, + ) -> Pin> + Send + 'static>> { + let model = self.model.lock().clone(); + let rf = options.response_format.as_ref(); + let body = RequestBody::build(&model, &messages, system.as_deref(), tools.as_deref(), rf); + let url = self.completions_url(); + let api_key = self.api_key.clone(); + let http = self.http.clone(); + + Box::pin(async_stream::try_stream! { + let resp = Self::post_completions(&http, &url, &api_key, &body.to_json(true)).await?; + let mut sse = SseReader::from_response(resp); + let mut emitter = StreamEmitter::default(); + + while let Some(data) = sse.next_data().await? { + let Some(chunk) = OpenAiChunk::parse(&data) else { + continue; + }; + emitter.process_chunk(&chunk); + for ev in emitter.drain() { + yield ev; + } + } + + for ev in emitter.finish() { + yield ev; + } + }) + } + + fn create_message_with_options( + &self, + messages: Vec, + system: Option, + tools: Option>, + options: crate::structured::RequestOptions, + ) -> Pin> + Send + '_>> { + let model = self.model.lock().clone(); + let rf = options.response_format.as_ref(); + let body = RequestBody::build(&model, &messages, system.as_deref(), tools.as_deref(), rf); + let url = self.completions_url(); + + Box::pin(async move { + let resp = + Self::post_completions(&self.http, &url, &self.api_key, &body.to_json(false)) + .await?; + let resp = resp + .bytes() + .await + .map_err(|e| ApiError::http(e.to_string()))?; + if resp.len() > MAX_RESPONSE_BODY { + return Err(ApiError::http(format!( + "response body too large: {} bytes (max {})", + resp.len(), + MAX_RESPONSE_BODY + ))); + } + serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) + }) + } + + fn extract_structured(&self, raw: &Value) -> Value { + let Some(content) = raw + .get("choices") + .and_then(|c| c.get(0)) + .and_then(|c| c.get("message")) + .and_then(|m| m.get("content")) + else { + return raw.clone(); + }; + if let Some(text) = content.as_str() { + crate::structured::parse_json_lenient(text) + .unwrap_or_else(|| Value::String(text.to_string())) + } else { + content.clone() + } + } +} // ================================================== // Builder // ================================================== /// Builder for [`OpenAiClient`]. +/// +/// Created via [`OpenAiClientBuilder::default`] or +/// [`OpenAiClient::builder`]. All fields have sensible defaults except +/// `api_key`, which must be set before [`build`](Self::build). pub struct OpenAiClientBuilder { + /// The API key for authentication (required). + /// + /// Must be set before building. Sent as the `Authorization: Bearer` + /// header on every request. api_key: Option, + + /// The base URL for API requests. + /// + /// Defaults to `https://api.openai.com/v1`. Override for + /// `DeepSeek`, `Grok`, `Ollama`, `vLLM`, or other OpenAI-compatible endpoints. base_url: String, + + /// The default model identifier. + /// + /// Can be changed at runtime via [`OpenAiClient::set_model`]. model: String, + + /// The total HTTP request timeout (connect + response + body). + /// + /// Bounds the entire request lifecycle. Defaults to 120 seconds. timeout: Duration, + + /// The TCP connection establishment timeout (including TLS handshake). + /// + /// Separate from the total timeout so a slow-connecting server can be + /// detected faster. Defaults to 10 seconds. connect_timeout: Duration, } @@ -261,21 +415,33 @@ impl Default for OpenAiClientBuilder { } impl OpenAiClientBuilder { - /// Set the API key. + /// Set the API key for authentication. + /// + /// Required — [`build`](Self::build) returns an error if this is not set. + /// The key is sent as the `Authorization: Bearer ` header on every + /// request. #[must_use] pub fn api_key(mut self, key: impl Into) -> Self { self.api_key = Some(key.into()); self } - /// Set the base URL (e.g. `https://api.deepseek.com/v1`). + /// Set the base URL for API requests. + /// + /// Defaults to `https://api.openai.com/v1`. Override when targeting an + /// OpenAI-compatible endpoint (e.g. `https://api.deepseek.com/v1`, + /// `http://localhost:11434/v1` for Ollama, or a vLLM server). #[must_use] pub fn base_url(mut self, url: impl Into) -> Self { self.base_url = url.into(); self } - /// Set the model name (e.g. `gpt-4o`, `deepseek-chat`). + /// Set the default model identifier. + /// + /// The model string is sent as the `model` field on every request. Can be + /// changed at runtime via [`OpenAiClient::set_model`] (e.g. when the + /// [`FallbackManager`](crate::fallback::FallbackManager) trips). #[must_use] pub fn model(mut self, model: impl Into) -> Self { self.model = model.into(); @@ -338,19 +504,48 @@ impl OpenAiClientBuilder { /// body for both streaming and non-streaming requests, toggling only /// the `stream` flag via [`to_json`](Self::to_json). struct RequestBody { + /// The model identifier sent as the `model` field on every request. + /// + /// Copied from [`OpenAiClient`]'s model field (which is mutable at + /// runtime via [`ApiClient::set_model`]). Each request carries it so + /// the provider knows which model to invoke. model: String, + + /// The conversation messages converted to OpenAI's JSON format. + /// + /// Each message is an object with `role` (`"system"`, `"user"`, or + /// `"assistant"`) and `content` (text string, tool-call array, or tool + /// result). Built by [`RequestBody::build`] from the framework's + /// [`Message`] list via [`convert_message`]. messages: Vec, + + /// The registered tools in OpenAI function-calling format, or `None` + /// when no tools are registered or when `response_format` is set + /// (mutual exclusion). tools: Option>, + + /// The structured-output `response_format` JSON object, or `None` + /// when structured output is not requested. When `Some`, the field is + /// emitted as `response_format: { type: "json_schema", ... }` and + /// `tools` is suppressed. + response_format: Option, } impl RequestBody { - /// Translate the framework's [`Message`] list into the OpenAI - /// Chat Completions request shape. + /// Translate the framework's [`Message`] list into the OpenAI Chat + /// Completions request shape. + /// + /// Converts messages to OpenAI's `role`/`content` JSON format, wraps + /// tool schemas in the `function` envelope, and — when + /// `response_format` is set — suppresses `tools` (OpenAI's structured + /// output and free-form tool-calling are mutually exclusive) and emits + /// the `response_format: json_schema` object. fn build( model: &str, messages: &[Message], system: Option<&str>, tools: Option<&[ToolSchema]>, + response_format: Option<&crate::structured::ResponseFormat>, ) -> Self { let mut msgs = Vec::with_capacity(messages.len().saturating_add(1)); @@ -362,22 +557,51 @@ impl RequestBody { msgs.push(convert_message(m)); } + let tools = if response_format.is_some() { + None + } else { + tools.map(convert_tools) + }; + + let rf = response_format.map(|rf| { + serde_json::json!({ + "type": "json_schema", + "json_schema": { + "name": rf.name, + "schema": rf.schema, + "strict": rf.strict + } + }) + }); + Self { model: model.into(), messages: msgs, - tools: tools.map(convert_tools), + tools, + response_format: rf, } } - /// Serialize to a [`serde_json::Value`] with the `stream` flag - /// set as requested. + /// Serialize to a [`serde_json::Value`] for the HTTP request body. + /// + /// Emits `model`, `messages`, `stream` (toggled by the parameter), + /// and `tools`. When `response_format` is set, appends the + /// `response_format` key; otherwise omits it entirely (not `null`). fn to_json(&self, stream: bool) -> Value { - serde_json::json!({ + let mut body = serde_json::json!({ "model": self.model, "messages": self.messages, "stream": stream, - "tools": self.tools, - }) + }); + if let Some(obj) = body.as_object_mut() { + if let Some(tools) = &self.tools { + obj.insert("tools".to_string(), Value::Array(tools.clone())); + } + if let Some(rf) = &self.response_format { + obj.insert("response_format".to_string(), rf.clone()); + } + } + body } } @@ -432,7 +656,13 @@ fn convert_message(m: &Message) -> Value { } } -/// Build an assistant message JSON that includes `tool_calls`. +/// Build an assistant message JSON object that includes `tool_calls`. +/// +/// Constructs the OpenAI-shaped `{ role, content, tool_calls }` object from +/// the accumulated text parts and tool-call entries. When there is no text +/// (pure tool-call turn), `content` is set to `null` — OpenAI's convention +/// for tool-call-only assistant messages. The `tool_calls` array carries the +/// converted tool-call entries produced by [`convert_message`]. fn build_assistant_message(role: &str, tool_calls: &[Value], text_parts: &[&str]) -> Value { let text = text_parts.join(""); let content = if text.is_empty() { @@ -459,7 +689,13 @@ fn merge_tool_results(results: &[Value]) -> Value { } } -/// Convert tool schemas into the OpenAI `tools` array shape. +/// Convert framework tool schemas into the OpenAI `tools` array shape. +/// +/// Each [`ToolSchema`] becomes a JSON object with `type: "function"` and a +/// nested `function` object carrying `name`, `description`, and `parameters` +/// (the framework's `input_schema`). When structured output is active +/// (`response_format` set), this function is not called — `tools` is +/// suppressed entirely. fn convert_tools(tools: &[ToolSchema]) -> Vec { tools .iter() @@ -555,7 +791,12 @@ impl SseReader { } } - /// Pop the first `\n`-terminated line from the buffer, if present. + /// Pop the first `\n`-terminated line from the internal buffer. + /// + /// Returns the trimmed line if a newline is present, and removes it + /// (plus the newline) from the buffer. Returns `None` if the buffer + /// does not yet contain a complete line — the caller should wait for + /// more bytes from the HTTP stream. fn take_line(&mut self) -> Option { let pos = self.buf.find('\n')?; let line = self.buf[..pos].trim().to_string(); @@ -647,8 +888,14 @@ struct StreamEmitter { } impl StreamEmitter { - /// Process a single parsed chunk, appending events to the - /// internal pending queue. + /// Process a single parsed OpenAI SSE chunk into stream events. + /// + /// On the first call, emits [`MessageStart`](StreamEvent::MessageStart) + /// with the chunk's message ID and model. Then delegates to + /// [`process_delta`](Self::process_delta) for text/tool-call deltas and + /// [`process_finish`](Self::process_finish) for the terminal finish + /// reason. Events accumulate in the internal queue until + /// [`drain`](Self::drain) is called. fn process_chunk(&mut self, chunk: &OpenAiChunk) { if !self.started { self.started = true; @@ -674,7 +921,12 @@ impl StreamEmitter { } } - /// Translate a delta into text/tool-call events. + /// Translate a single delta object into text and/or tool-call events. + /// + /// If the delta carries non-empty `content`, emits a `PartStart` (on the + /// first text delta) followed by `IndexedDelta(Text)` events. If it + /// carries `tool_calls`, delegates each to + /// [`process_tool_call`](Self::process_tool_call). fn process_delta(&mut self, delta: &OpenAiDelta) { if let Some(text) = &delta.content && !text.is_empty() @@ -699,7 +951,14 @@ impl StreamEmitter { } } - /// Handle a single tool-call delta. + /// Handle a single tool-call delta from the stream. + /// + /// On the first delta for a tool call (when `function` is present), + /// emits a [`PartStart`](StreamEvent::PartStart) with a + /// [`ToolCall`](crate::message::MessagePart::ToolCall) part carrying the + /// tool ID and name. Subsequent deltas carrying `function.arguments` + /// fragments emit [`InputJson`](crate::stream::DeltaPart::InputJson) + /// events so the caller can accumulate the full JSON input. fn process_tool_call(&mut self, tc: &OpenAiToolCallDelta) { if tc.function.is_some() { // New tool call — emit PartStart. @@ -732,6 +991,14 @@ impl StreamEmitter { } /// Handle a finish reason, emitting the appropriate stop events. + /// + /// Closes any open text parts and tool-call parts with + /// [`PartStop`](StreamEvent::PartStop), then emits a + /// [`MessageDelta`](StreamEvent::MessageDelta) carrying the mapped + /// [`StreamStopReason`]. Maps `"tool_calls"` → + /// [`ToolCall`](StreamStopReason::ToolCall), `"length"` → + /// [`MaxTokens`](StreamStopReason::MaxTokens), and anything else via + /// [`StreamStopReason::from_api_str`]. No-ops if already finished. fn process_finish(&mut self, reason: &str) { if self.finished { return; @@ -762,8 +1029,11 @@ impl StreamEmitter { })); } - /// Emit the terminal [`MessageStop`] if the stream was started, - /// returning all remaining events. + /// Finalize the stream, emitting the terminal + /// [`MessageStop`](StreamEvent::MessageStop) if one was started. + /// + /// Drains any remaining pending events and appends the stop event. + /// Called exactly once at the end of the SSE stream. fn finish(&mut self) -> Vec { let mut out = self.drain(); if self.started { @@ -772,11 +1042,20 @@ impl StreamEmitter { out } - /// Drain all pending events. + /// Drain all pending events from the internal queue. + /// + /// Returns the accumulated [`StreamEvent`]s and clears the queue. + /// Called by the stream loop after each chunk is processed so events + /// are yielded promptly rather than buffered until stream end. fn drain(&mut self) -> Vec { std::mem::take(&mut self.pending) } + /// Push an event onto the internal pending queue. + /// + /// The single write point — all methods (`process_delta`, + /// `process_tool_call`, `process_finish`, `process_chunk`) funnel + /// through here. Events are held until [`drain`](Self::drain) is called. fn push(&mut self, ev: StreamEvent) { self.pending.push(ev); } @@ -795,7 +1074,7 @@ mod tests { #[test] fn request_body_includes_system_message_first() { let msgs = vec![Message::user("hello")]; - let body = RequestBody::build("gpt-4o", &msgs, Some("be brief"), None); + let body = RequestBody::build("gpt-4o", &msgs, Some("be brief"), None, None); let json = body.to_json(true); let messages = json["messages"].as_array().unwrap(); @@ -808,7 +1087,7 @@ mod tests { #[test] fn request_body_without_system() { let msgs = vec![Message::user("hi")]; - let body = RequestBody::build("gpt-4o", &msgs, None, None); + let body = RequestBody::build("gpt-4o", &msgs, None, None, None); let json = body.to_json(false); let messages = json["messages"].as_array().unwrap(); @@ -819,7 +1098,7 @@ mod tests { #[test] fn request_body_stream_flag_toggles() { let msgs = vec![Message::user("hi")]; - let body = RequestBody::build("gpt-4o", &msgs, None, None); + let body = RequestBody::build("gpt-4o", &msgs, None, None, None); assert_eq!(body.to_json(true)["stream"], true); assert_eq!(body.to_json(false)["stream"], false); @@ -833,7 +1112,7 @@ mod tests { description: "Echo".into(), input_schema: serde_json::json!({"type": "object"}), }]; - let body = RequestBody::build("my-model", &msgs, None, Some(&tools)); + let body = RequestBody::build("my-model", &msgs, None, Some(&tools), None); let json = body.to_json(true); assert_eq!(json["model"], "my-model"); @@ -844,11 +1123,14 @@ mod tests { } #[test] - fn request_body_tools_null_when_none() { + fn request_body_tools_absent_when_none() { let msgs = vec![Message::user("hi")]; - let body = RequestBody::build("gpt-4o", &msgs, None, None); + let body = RequestBody::build("gpt-4o", &msgs, None, None, None); let json = body.to_json(false); - assert!(json["tools"].is_null()); + assert!( + json.get("tools").is_none(), + "tools key should be absent when no tools are set" + ); } #[test] @@ -1400,4 +1682,96 @@ mod tests { let within = MAX_RESPONSE_BODY; assert!(within <= MAX_RESPONSE_BODY); } + + #[test] + fn request_body_response_format_emitted() { + let msgs = vec![Message::user("hi")]; + let rf = + crate::structured::ResponseFormat::new("action", serde_json::json!({"type": "object"})); + let body = RequestBody::build("gpt-4o", &msgs, None, None, Some(&rf)); + let json = body.to_json(false); + + assert_eq!(json["response_format"]["type"], "json_schema"); + assert_eq!(json["response_format"]["json_schema"]["name"], "action"); + assert_eq!( + json["response_format"]["json_schema"]["schema"], + serde_json::json!({"type": "object"}) + ); + assert_eq!(json["response_format"]["json_schema"]["strict"], true); + } + + #[test] + fn request_body_response_format_absent_when_none() { + let msgs = vec![Message::user("hi")]; + let body = RequestBody::build("gpt-4o", &msgs, None, None, None); + let json = body.to_json(false); + assert!( + json.get("response_format").is_none(), + "response_format should be absent (not null) when not set" + ); + } + + #[test] + fn request_body_response_format_suppresses_tools() { + let msgs = vec![Message::user("hi")]; + let caller_tool = ToolSchema { + tool: "read".into(), + description: "Read".into(), + input_schema: serde_json::json!({"type": "object"}), + }; + let rf = + crate::structured::ResponseFormat::new("result", serde_json::json!({"type": "object"})); + let body = RequestBody::build("gpt-4o", &msgs, None, Some(&[caller_tool]), Some(&rf)); + let json = body.to_json(false); + + assert!( + json.get("tools").is_none(), + "tools key should be absent when response_format is set" + ); + assert!(json.get("response_format").is_some()); + } + + #[test] + fn extract_structured_from_string_content() { + let client = OpenAiClient::builder().api_key("test").build().unwrap(); + let raw = serde_json::json!({ + "choices": [{ + "message": { + "content": r#"{"tool": "write", "args": {}}"# + } + }] + }); + let value = client.extract_structured(&raw); + assert_eq!(value["tool"], "write"); + } + + #[test] + fn extract_structured_from_object_content() { + let client = OpenAiClient::builder().api_key("test").build().unwrap(); + let raw = serde_json::json!({ + "choices": [{ + "message": { + "content": {"tool": "read", "args": {}} + } + }] + }); + let value = client.extract_structured(&raw); + assert_eq!(value["tool"], "read"); + } + + #[test] + fn extract_structured_prose_falls_back_to_raw() { + let client = OpenAiClient::builder().api_key("test").build().unwrap(); + let raw = serde_json::json!({ + "choices": [{ + "message": { + "content": "I cannot produce that." + } + }] + }); + let value = client.extract_structured(&raw); + // When content is prose (not parseable JSON), falls back to the + // string value; T::from_value will then fail with Deserialize. + assert_eq!(value, serde_json::json!("I cannot produce that.")); + } } diff --git a/src/structured.rs b/src/structured.rs new file mode 100644 index 0000000..a8fe299 --- /dev/null +++ b/src/structured.rs @@ -0,0 +1,687 @@ +//! Structured output — request guaranteed-schema JSON responses from the model. +//! +//! This module provides: +//! +//! - [`StructuredOutput`] — a type-level trait that names a type, exposes its +//! JSON Schema, and deserializes from a `serde_json::Value`. +//! - [`ResponseFormat`] + [`RequestOptions`] — the request-side carrier that +//! tells the provider to constrain output to the schema. +//! - [`StructuredError`] — errors raised by the structured-output machinery. +//! - [`request_structured`] — a convenience helper that hides the +//! options/extraction dance behind a single generic call. +//! +//! # Quick Start +//! +//! ```rust,ignore +//! use loopctl::structured::{StructuredOutput, request_structured}; +//! use serde::{Deserialize, Serialize}; +//! use serde_json::json; +//! +//! #[derive(Debug, Serialize, Deserialize, PartialEq)] +//! pub struct Action { +//! pub tool: String, +//! pub args: serde_json::Value, +//! } +//! +//! impl StructuredOutput for Action { +//! fn name() -> &'static str { "action" } +//! fn schema() -> serde_json::Value { +//! json!({ +//! "type": "object", +//! "properties": { +//! "tool": { "type": "string" }, +//! "args": {} +//! }, +//! "required": ["tool", "args"], +//! "additionalProperties": false +//! }) +//! } +//! } +//! +//! // let action: Action = request_structured(&client, messages, system).await?; +//! ``` + +use crate::api::ApiClient; +use crate::message::Message; + +/// A type that can be requested from the model as a JSON-schema-conformant +/// response, and deserialized from the model's output. +/// +/// This is a *type-level* trait (like `serde::Serialize`), not a provider +/// trait — it is never used as a trait object. Implement it on any `Sized + +/// Send + 'static` type that also implements `serde::de::DeserializeOwned`. +/// +/// The schema returned by [`schema`](Self::schema) is injected into the +/// provider request (OpenAI `response_format` / Anthropic forced tool); the +/// model's output is parsed back via [`from_value`](Self::from_value). +/// +/// # Manual schema vs derive +/// +/// By default, implement `schema()` by returning a `serde_json::json!` +/// literal — no extra dependency, matching how +/// [`ToolSchema::input_schema`](crate::tool::ToolSchema::input_schema) is +/// authored today. +/// +/// # Example +/// +/// ```rust,ignore +/// use loopctl::structured::StructuredOutput; +/// use serde::{Deserialize, Serialize}; +/// use serde_json::json; +/// +/// #[derive(Debug, Serialize, Deserialize)] +/// pub struct Action { +/// pub tool: String, +/// pub args: serde_json::Value, +/// } +/// +/// impl StructuredOutput for Action { +/// fn name() -> &'static str { "action" } +/// fn schema() -> serde_json::Value { +/// json!({ +/// "type": "object", +/// "properties": { +/// "tool": { "type": "string" }, +/// "args": {} +/// }, +/// "required": ["tool", "args"], +/// "additionalProperties": false +/// }) +/// } +/// } +/// ``` +pub trait StructuredOutput: Sized + Send + 'static { + /// Logical name for the schema. + /// + /// Used verbatim as the OpenAI `json_schema.name` field and as the + /// synthesized Anthropic forced-tool name. Must match + /// `^[a-zA-Z0-9_-]+$` (alphanumeric, underscore, hyphen only) because + /// both providers validate this identifier. + fn name() -> &'static str; + + /// The JSON Schema (Draft 07) describing the desired output object. + /// + /// The schema is injected into the provider request: OpenAI emits it as + /// `response_format.json_schema.schema`; Anthropic uses it as the + /// forced tool's `input_schema`. The model's output is expected to + /// conform to this schema — the [`from_value`](Self::from_value) method + /// then deserializes it into `Self`. + /// + /// Implement this by returning a `serde_json::json!({ … })` literal + /// (matching the pattern used by + /// [`ToolSchema::input_schema`](crate::tool::ToolSchema::input_schema)). + fn schema() -> serde_json::Value; + + /// Deserialize an instance from the model's JSON output. + /// + /// The default implementation is `serde_json::from_value::(v)`, + /// which is correct for any `Self: DeserializeOwned`. Override only for + /// post-processing (e.g. trimming, defaults, cross-field validation). + /// + /// # Errors + /// + /// Returns [`StructuredError::Deserialize`] if the value does not match + /// the type (and, by construction, the schema). + fn from_value(v: serde_json::Value) -> Result + where + Self: serde::de::DeserializeOwned, + { + serde_json::from_value(v).map_err(StructuredError::Deserialize) + } +} + +/// A request to constrain the model's output to a named JSON schema. +/// +/// Construct with [`ResponseFormat::from_type`] from any [`StructuredOutput`] +/// type, or manually from a raw schema + name via [`ResponseFormat::new`]. +/// Passed to the provider via [`RequestOptions`]. +#[derive(Debug, Clone)] +pub struct ResponseFormat { + /// Logical name for the schema. + /// + /// Copied from [`StructuredOutput::name`] when constructed via + /// [`from_type`](Self::from_type). Used as the OpenAI `json_schema.name` + /// and the Anthropic forced-tool name. Must match `^[a-zA-Z0-9_-]+$`. + pub name: String, + + /// The JSON Schema the model's output must satisfy. + /// + /// Injected into the provider request verbatim: OpenAI emits it as + /// `response_format.json_schema.schema`; Anthropic uses it as the + /// forced tool's `input_schema`. + pub schema: serde_json::Value, + + /// Whether to enforce the schema server-side ("strict" mode). + /// + /// When `true` (the default), OpenAI guarantees the output conforms to + /// the schema. Providers that lack strict mode (Anthropic tool-forcing) + /// ignore this flag. + pub strict: bool, +} + +impl ResponseFormat { + /// Build a [`ResponseFormat`] from a [`StructuredOutput`] type. + /// + /// # Example + /// + /// ```rust,ignore + /// use loopctl::structured::{ResponseFormat, StructuredOutput}; + /// + /// let rf = ResponseFormat::from_type::(); + /// assert_eq!(rf.name, MyOutput::name()); + /// assert_eq!(rf.schema, MyOutput::schema()); + /// assert!(rf.strict); + /// ``` + #[must_use] + pub fn from_type() -> Self { + Self { + name: T::name().to_string(), + schema: T::schema(), + strict: true, + } + } + + /// Build a [`ResponseFormat`] from raw parts. + /// + /// Use this when the schema is constructed dynamically (e.g. at runtime + /// from config or a database) rather than from a static type. For the + /// common case of deriving the format from a `StructuredOutput` type, + /// prefer [`from_type`](Self::from_type). + /// + /// Sets `strict: true` by default — the model is constrained server-side + /// where supported (OpenAI strict mode, Anthropic ignores the flag). + #[must_use] + pub fn new(name: impl Into, schema: serde_json::Value) -> Self { + Self { + name: name.into(), + schema, + strict: true, + } + } +} + +/// Optional per-request knobs layered on top of a `stream_messages` / +/// `create_message` call. +/// +/// Additive and forward-compatible: every field is optional, and providers +/// that don't understand a field ignore it. Today carries only +/// [`response_format`](Self::response_format); future tasks (max-tokens, +/// temperature, seed) extend it without touching the trait again. +#[derive(Debug, Clone, Default)] +pub struct RequestOptions { + /// If set, ask the model to return JSON conforming to this schema. + /// + /// When `Some`, the provider injects the schema into the request + /// (OpenAI `response_format` / Anthropic forced tool). When `None`, + /// the model's output is unconstrained — the default behaviour. + pub response_format: Option, +} + +impl RequestOptions { + /// Create empty options with no response format set. + /// + /// Equivalent to [`RequestOptions::default`]. Use + /// [`response_format`](Self::response_format) to chain a format + /// builder-style. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Set the response format, builder-style. + /// + /// When set, the provider constrains the model's output to the schema. + /// When left `None` (the default), the model's output is unconstrained. + #[must_use] + pub fn response_format(mut self, rf: ResponseFormat) -> Self { + self.response_format = Some(rf); + self + } +} + +/// Errors raised by the structured-output machinery. +#[derive(Debug, thiserror::Error)] +pub enum StructuredError { + /// The model's output did not deserialize into the target type. + /// + /// Carries the underlying `serde_json::Error` with its exact location + /// (line/column within the JSON). This typically means the model returned + /// valid JSON but with missing fields, wrong types, or unexpected + /// structure relative to `T`'s schema. + #[error("structured output did not match the expected schema: {0}")] + Deserialize(#[from] serde_json::Error), + + /// The provider API call failed (HTTP error, auth failure, timeout, rate + /// limit). + /// + /// Carries the underlying [`ApiError`](crate::api::error::ApiError). + #[error("API error during structured output request: {0}")] + Api(crate::api::error::ApiError), +} + +/// Parse a string as JSON, with a lenient fallback that finds the outermost +/// `{ ... }` or `[ ... ]` substring. +/// +/// This is the single biggest lever for hitting the ≥95% schema-valid bar on +/// real-world providers that wrap JSON in markdown fences or prefix it with +/// prose. +/// +/// Returns `None` if the content cannot be parsed as JSON (even after the +/// lenient rescue). +pub(crate) fn parse_json_lenient(text: &str) -> Option { + if let Ok(v) = serde_json::from_str(text) { + return Some(v); + } + // Lenient rescue: find the outermost { ... } or [ ... ]. + extract_json_substring(text) +} + +/// Find and parse the outermost JSON object or array in a string. +/// +/// Scans for the first `{` or `[`, tracks brace/bracket depth, and extracts +/// the substring up to the matching close. String-aware: braces/brackets +/// inside JSON string literals (`"..."`) do not affect depth, and `\"` +/// escapes are honored. +pub(crate) fn extract_json_substring(text: &str) -> Option { + let bytes = text.as_bytes(); + let mut start = None; + let mut depth: i32 = 0; + let mut close = b'\0'; + let mut in_string = false; + let mut escaped = false; + + for (i, &byte) in bytes.iter().enumerate() { + if in_string { + if escaped { + escaped = false; + } else if byte == b'\\' { + escaped = true; + } else if byte == b'"' { + in_string = false; + } + continue; + } + + match byte { + b'"' => { + in_string = true; + } + b'{' | b'[' => { + if start.is_none() { + start = Some(i); + close = if byte == b'{' { b'}' } else { b']' }; + } + depth = depth.saturating_add(1); + } + b'}' | b']' => { + if let Some(s) = start { + if byte == close { + depth = depth.saturating_sub(1); + if depth == 0 { + let slice = text.get(s..=i).unwrap_or(text); + if let Ok(v) = serde_json::from_str(slice) { + return Some(v); + } + start = None; + } + } else { + // Mismatched closing delimiter — not valid JSON. + start = None; + depth = 0; + } + } + } + _ => {} + } + } + None +} + +/// Request a typed, schema-conformant value from the model. +/// +/// This is the ergonomic entry point callers use. It: +/// 1. Builds [`RequestOptions`] with the [`ResponseFormat`] for `T`. +/// 2. Calls [`create_message_with_options`](ApiClient::create_message_with_options) +/// on the client. +/// 3. Extracts the structured value from the provider's response. +/// 4. Deserializes it into `T` via [`StructuredOutput::from_value`]. +/// +/// # Errors +/// +/// Returns [`StructuredError`] if the provider call fails, the response +/// cannot be parsed as JSON, or the JSON does not match `T`'s schema. +pub async fn request_structured( + client: &dyn ApiClient, + messages: Vec, + system: Option, +) -> Result { + let opts = RequestOptions::new().response_format(ResponseFormat::from_type::()); + let raw = client + .create_message_with_options(messages, system, None, opts) + .await + .map_err(StructuredError::Api)?; + let value = client.extract_structured(&raw); + T::from_value(value) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::future::Future; + use std::pin::Pin; + + #[derive(Debug, serde::Serialize, serde::Deserialize, PartialEq)] + struct Action { + tool: String, + args: serde_json::Value, + } + + impl StructuredOutput for Action { + fn name() -> &'static str { + "action" + } + fn schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "tool": { "type": "string" }, + "args": {} + }, + "required": ["tool", "args"], + "additionalProperties": false + }) + } + } + + fn fixture_action() -> serde_json::Value { + serde_json::json!({ + "tool": "write", + "args": { "path": "/tmp/test.txt" } + }) + } + + #[test] + fn structured_output_round_trip() { + let v = fixture_action(); + let action: Action = Action::from_value(v).expect("should deserialize"); + assert_eq!(action.tool, "write"); + assert_eq!(action.args, serde_json::json!({ "path": "/tmp/test.txt" })); + } + + #[test] + fn response_format_from_type() { + let rf = ResponseFormat::from_type::(); + assert_eq!(rf.name, "action"); + assert_eq!(rf.schema, Action::schema()); + assert!(rf.strict); + } + + #[test] + fn request_options_builder() { + let opts = RequestOptions::new(); + assert!(opts.response_format.is_none()); + + let rf = ResponseFormat::from_type::(); + let opts = RequestOptions::new().response_format(rf); + assert!(opts.response_format.is_some()); + assert_eq!(opts.response_format.as_ref().unwrap().name, "action"); + } + + #[test] + fn parse_json_lenient_plain_json() { + let v = parse_json_lenient(r#"{"a": 1}"#).unwrap(); + assert_eq!(v["a"], 1); + } + + #[test] + fn parse_json_lenient_with_prefix() { + let v = parse_json_lenient(r#"Here is the JSON: {"a": 1}"#).unwrap(); + assert_eq!(v["a"], 1); + } + + #[test] + fn parse_json_lenient_markdown_fences() { + let v = parse_json_lenient("```json\n{\"a\": 1}\n```").unwrap(); + assert_eq!(v["a"], 1); + } + + #[test] + fn parse_json_lenient_array() { + let v = parse_json_lenient(r#"prefix [1, 2, 3] suffix"#).unwrap(); + assert_eq!(v[0], 1); + } + + #[test] + fn parse_json_lenient_no_json() { + let result = parse_json_lenient("just prose, nothing here"); + assert!(result.is_none()); + } + + #[test] + fn structured_error_displays() { + let json_err = serde_json::from_str::("bad").unwrap_err(); + let err = StructuredError::Deserialize(json_err); + assert!(err.to_string().contains("schema")); + } + + #[test] + fn parse_json_lenient_brace_inside_string() { + let v = parse_json_lenient(r#"prefix {"a": "}"} suffix"#).unwrap(); + assert_eq!(v["a"], "}"); + } + + #[test] + fn parse_json_lenient_bracket_inside_string() { + let v = parse_json_lenient(r#"before {"x": "]"} after"#).unwrap(); + assert_eq!(v["x"], "]"); + } + + #[test] + fn parse_json_lenient_escaped_quote_in_string() { + let v = parse_json_lenient(r#"here {"a": "he said \"hi\""} there"#).unwrap(); + assert_eq!(v["a"], "he said \"hi\""); + } + + #[test] + fn parse_json_lenient_nested_objects_in_prose() { + let v = parse_json_lenient(r#"result: {"outer": {"inner": 42}}"#).unwrap(); + assert_eq!(v["outer"]["inner"], 42); + } + + #[test] + fn parse_json_lenient_mismatched_delimiter_then_valid() { + let v = parse_json_lenient(r#"{oops] then {"a":1}"#).unwrap(); + assert_eq!(v["a"], 1); + } + + struct PlainMockClient; + impl crate::api::ApiClient for PlainMockClient { + fn model(&self) -> String { + "test".to_string() + } + fn stream_messages( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin< + Box< + dyn futures::Stream< + Item = Result, + > + Send + + 'static, + >, + > { + Box::pin(futures::stream::empty()) + } + fn create_message( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin< + Box< + dyn Future> + + Send + + '_, + >, + > { + Box::pin(async { Ok(serde_json::json!({})) }) + } + } + + #[tokio::test] + async fn default_client_rejects_response_format() { + let client = PlainMockClient; + let opts = RequestOptions::new().response_format(ResponseFormat::from_type::()); + let result = client + .create_message_with_options(vec![], None, None, opts) + .await; + assert!( + result.is_err(), + "client without structured-output support should reject response_format" + ); + let err_msg = result.unwrap_err().to_string(); + assert!( + err_msg.contains("does not support structured output"), + "error should explain why: {err_msg}" + ); + } + + #[tokio::test] + async fn default_client_delegates_empty_options() { + let client = PlainMockClient; + let opts = RequestOptions::new(); + let result = client + .create_message_with_options(vec![], None, None, opts) + .await; + assert!(result.is_ok(), "empty options should delegate normally"); + } + + struct StructuredMockClient; + impl crate::api::ApiClient for StructuredMockClient { + fn model(&self) -> String { + "test".to_string() + } + fn stream_messages( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin< + Box< + dyn futures::Stream< + Item = Result, + > + Send + + 'static, + >, + > { + Box::pin(futures::stream::empty()) + } + fn create_message( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin< + Box< + dyn Future> + + Send + + '_, + >, + > { + Box::pin(async { Ok(serde_json::json!({})) }) + } + fn create_message_with_options( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + _options: RequestOptions, + ) -> Pin< + Box< + dyn Future> + + Send + + '_, + >, + > { + Box::pin(async { + Ok(serde_json::json!({ + "tool": "write", + "args": {"path": "/test"} + })) + }) + } + } + + #[tokio::test] + async fn request_structured_end_to_end() { + let client = StructuredMockClient; + let action: Action = request_structured(&client, vec![], None) + .await + .expect("should succeed"); + assert_eq!(action.tool, "write"); + } + + struct ProseMockClient; + impl crate::api::ApiClient for ProseMockClient { + fn model(&self) -> String { + "test".to_string() + } + fn stream_messages( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin< + Box< + dyn futures::Stream< + Item = Result, + > + Send + + 'static, + >, + > { + Box::pin(futures::stream::empty()) + } + fn create_message( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin< + Box< + dyn Future> + + Send + + '_, + >, + > { + Box::pin(async { Ok(serde_json::json!({})) }) + } + fn create_message_with_options( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + _options: RequestOptions, + ) -> Pin< + Box< + dyn Future> + + Send + + '_, + >, + > { + Box::pin(async { Ok(serde_json::json!("I cannot produce that.")) }) + } + } + + #[tokio::test] + async fn request_structured_prose_returns_deserialize_error() { + let client = ProseMockClient; + let err = request_structured::(&client, vec![], None) + .await + .expect_err("should fail"); + // Prose is a valid JSON string but doesn't match Action's schema, + // so deserialization fails. + assert!(matches!(err, StructuredError::Deserialize(_))); + } +} diff --git a/src/tool.rs b/src/tool.rs index 93264ae..5b9d7b5 100644 --- a/src/tool.rs +++ b/src/tool.rs @@ -65,7 +65,7 @@ pub mod health; #[cfg(feature = "tool_shield")] pub mod shield; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Serialize, de::DeserializeOwned}; use serde_json::Value; use std::collections::HashMap; use std::fmt; @@ -360,6 +360,74 @@ impl ToolOutput { } } } + + /// Construct a successful result from any serializable value. + /// + /// The value is JSON-serialized into the + /// [`ToolContent::Text`](crate::message::ToolContent::Text) payload and + /// flagged so [`structured_value`](Self::structured_value) can recover it. + /// `T` need not implement [`StructuredOutput`](crate::structured::StructuredOutput) + /// — any `Serialize` works — + /// but types that *do* get a round-trippable accessor via + /// [`structured_as`](Self::structured_as). + /// + /// If serialization fails (should not happen for normal structs), returns + /// an error result rather than panicking. + /// + /// # Example + /// + /// ```rust + /// use loopctl::tool::ToolOutput; + /// use serde::Serialize; + /// + /// #[derive(Serialize)] + /// struct Data { count: u32 } + /// + /// let out = ToolOutput::structured(&Data { count: 42 }); + /// assert!(!out.is_error); + /// assert!(out.structured_value().is_some()); + /// ``` + pub fn structured(value: &T) -> Self { + match serde_json::to_string(value) { + Ok(json) => Self::success(json), + Err(e) => Self::error_text(format!("structured serialization failed: {e}")), + } + } + + /// Parse the payload back into a [`serde_json::Value`], if it is valid JSON. + /// + /// Returns `None` for multipart payloads or non-JSON text. Use this when + /// the consumer does not know the concrete type at compile time (e.g. an + /// observer that re-serializes for a TUI). + #[must_use] + pub fn structured_value(&self) -> Option { + let MessageToolContent::Text(s) = &self.payload else { + return None; + }; + serde_json::from_str(s).ok() + } + + /// Parse the payload into a concrete `T: StructuredOutput`. + /// + /// Round-trips a value produced by [`structured`](Self::structured) or any + /// tool whose text output happens to conform to `T`'s schema. Returns + /// `None` if the payload is not valid JSON for `T`. + /// + /// # Example + /// + /// ```rust,ignore + /// use loopctl::tool::ToolOutput; + /// use loopctl::structured::StructuredOutput; + /// + /// let out = ToolOutput::structured(&action); + /// let back: Action = out.structured_as().unwrap(); + /// ``` + #[must_use] + pub fn structured_as( + &self, + ) -> Option { + self.structured_value().and_then(|v| T::from_value(v).ok()) + } } impl From for ToolOutput { @@ -1523,4 +1591,106 @@ mod tests { let tool = EchoTool; assert!(tool.system_prompt().is_none()); } + + #[test] + fn tool_output_structured_round_trip() { + use crate::structured::StructuredOutput; + + #[derive(Debug, serde::Serialize, serde::Deserialize, PartialEq)] + struct Action { + tool: String, + args: serde_json::Value, + } + + impl StructuredOutput for Action { + fn name() -> &'static str { + "action" + } + fn schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "tool": { "type": "string" }, + "args": {} + }, + "required": ["tool", "args"] + }) + } + } + + let action = Action { + tool: "write".to_string(), + args: serde_json::json!({"path": "/a"}), + }; + let out = ToolOutput::structured(&action); + assert!(!out.is_error); + assert!(out.structured_value().is_some()); + let back: Action = out.structured_as().expect("should round-trip"); + assert_eq!(back, action); + } + + #[test] + fn tool_output_structured_value_none_for_non_json() { + let out = ToolOutput::text("not json"); + assert!(out.structured_value().is_none()); + } + + #[test] + fn tool_output_structured_with_plain_serialize() { + // structured works with any Serialize type, not just StructuredOutput. + #[derive(serde::Serialize)] + struct Count { + n: u32, + } + + let out = ToolOutput::structured(&Count { n: 7 }); + assert!(!out.is_error); + let v = out.structured_value().expect("should be valid JSON"); + assert_eq!(v["n"], 7); + } + + #[test] + fn tool_output_structured_primitive() { + let out = ToolOutput::structured(&42u32); + assert!(!out.is_error); + let v = out.structured_value().expect("should parse"); + assert_eq!(v, 42); + } + + #[test] + fn tool_output_structured_value_for_multipart_is_none() { + use crate::message::{ToolContent, ToolContentPart}; + let out = ToolOutput::success(ToolContent::Multipart(vec![ToolContentPart::Text { + text: "a".into(), + }])); + assert!(out.structured_value().is_none()); + } + + #[test] + fn tool_output_structured_as_returns_none_when_type_mismatches() { + use crate::structured::StructuredOutput; + + #[derive(Debug, serde::Serialize, serde::Deserialize, PartialEq)] + struct Target { + name: String, + } + + impl StructuredOutput for Target { + fn name() -> &'static str { + "target" + } + fn schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"] + }) + } + } + + // Feed JSON that has a different shape than Target expects. + let out = ToolOutput::text(r#"{"count": 5}"#); + let result: Option = out.structured_as(); + assert!(result.is_none(), "mismatched shape should not deserialize"); + } }