omni_dev/claude/
ai.rs

1//! AI client trait and metadata definitions.
2
3pub mod bedrock;
4pub mod claude;
5pub mod claude_cli;
6pub mod openai;
7
8use std::future::Future;
9use std::pin::Pin;
10use std::time::{Duration, Instant};
11
12use anyhow::{Context, Result};
13use reqwest::Client;
14use serde_json::Value;
15
16use crate::claude::error::ClaudeError;
17use crate::claude::model_config::get_model_registry;
18use crate::request_log;
19
20/// HTTP request timeout for AI API calls.
21///
22/// Set to 5 minutes to accommodate large prompts and long model responses
23/// (up to 64k output tokens) while preventing indefinite hangs.
24pub(crate) const REQUEST_TIMEOUT: Duration = Duration::from_secs(300);
25
26/// Metadata about an AI client implementation.
27#[derive(Clone, Debug)]
28pub struct AiClientMetadata {
29    /// Service provider name.
30    pub provider: String,
31    /// Model identifier.
32    pub model: String,
33    /// Maximum context length supported.
34    pub max_context_length: usize,
35    /// Maximum token response length supported.
36    pub max_response_length: usize,
37    /// Active beta header, if any: (key, value).
38    pub active_beta: Option<(String, String)>,
39}
40
41/// Prompt formatting families for AI providers.
42///
43/// Determines provider-specific prompt behaviour (e.g., how template
44/// instructions are phrased). Parse once at the boundary via
45/// [`AiClientMetadata::prompt_style`] and match on the enum downstream.
46#[derive(Clone, Copy, Debug, PartialEq, Eq)]
47pub enum PromptStyle {
48    /// Claude models handle "literal template" instructions correctly.
49    Claude,
50    /// OpenAI-compatible models (OpenAI, Ollama) need different formatting.
51    OpenAi,
52}
53
54impl AiClientMetadata {
55    /// Derives the prompt style from the provider name.
56    ///
57    /// Matches against the exact strings set by each [`AiClient`] implementation:
58    /// - `"OpenAI"` and `"Ollama"` → [`PromptStyle::OpenAi`]
59    /// - `"Anthropic"` and `"Anthropic Bedrock"` → [`PromptStyle::Claude`]
60    ///
61    /// Unrecognised provider strings default to [`PromptStyle::Claude`].
62    #[must_use]
63    pub fn prompt_style(&self) -> PromptStyle {
64        match self.provider.as_str() {
65            "OpenAI" | "Ollama" => PromptStyle::OpenAi,
66            _ => PromptStyle::Claude,
67        }
68    }
69}
70
71// ── Shared helpers for AI client implementations ────────────────────
72
73/// Builds an HTTP client with the standard request timeout.
74pub(crate) fn build_http_client() -> Result<Client> {
75    Client::builder()
76        .timeout(REQUEST_TIMEOUT)
77        .build()
78        .context("Failed to build HTTP client")
79}
80
81/// Appends a best-effort `service = claude` HTTP record for one AI-backend
82/// `POST` attempt (direct Anthropic, Bedrock, or OpenAI-compatible).
83pub(crate) fn record_ai_http(
84    url: &str,
85    started: Instant,
86    result: &reqwest::Result<reqwest::Response>,
87) {
88    match result {
89        Ok(r) => {
90            request_log::record_http(
91                "claude",
92                "POST",
93                url,
94                started,
95                Some(r.status().as_u16()),
96                None,
97            );
98        }
99        Err(e) => {
100            request_log::record_http("claude", "POST", url, started, None, Some(&e.to_string()));
101        }
102    }
103}
104
105/// Returns the maximum output tokens for a model from the registry,
106/// respecting beta overrides.
107#[must_use]
108pub(crate) fn registry_max_output_tokens(
109    model: &str,
110    active_beta: &Option<(String, String)>,
111) -> i32 {
112    let registry = get_model_registry();
113    if let Some((_, value)) = active_beta {
114        registry.get_max_output_tokens_with_beta(model, value) as i32
115    } else {
116        registry.get_max_output_tokens(model) as i32
117    }
118}
119
120/// Returns the (input context length, max response length) for a model
121/// from the registry, respecting beta overrides.
122#[must_use]
123pub(crate) fn registry_model_limits(
124    model: &str,
125    active_beta: &Option<(String, String)>,
126) -> (usize, usize) {
127    let registry = get_model_registry();
128    match active_beta {
129        Some((_, value)) => (
130            registry.get_input_context_with_beta(model, value),
131            registry.get_max_output_tokens_with_beta(model, value),
132        ),
133        None => (
134            registry.get_input_context(model),
135            registry.get_max_output_tokens(model),
136        ),
137    }
138}
139
140/// Checks an HTTP response for error status and returns a structured error
141/// if non-success.
142///
143/// On success, returns the response unchanged for further processing.
144/// On failure, reads the error body and returns a
145/// [`ClaudeError::ApiRequestFailed`].
146pub(crate) async fn check_error_response(response: reqwest::Response) -> Result<reqwest::Response> {
147    if response.status().is_success() {
148        return Ok(response);
149    }
150    let status = response.status();
151    let error_text = response.text().await.unwrap_or_else(|e| {
152        tracing::debug!("Failed to read error response body: {e}");
153        String::new()
154    });
155    Err(ClaudeError::ApiRequestFailed(format!("HTTP {status}: {error_text}")).into())
156}
157
158/// Logs successful text extraction from an AI API response.
159pub(crate) fn log_response_success(provider: &str, result: &Result<String>) {
160    if let Ok(text) = result {
161        tracing::debug!(
162            response_len = text.len(),
163            "Successfully extracted text content from {} API response",
164            provider
165        );
166        tracing::debug!(
167            response_content = %text,
168            "{} API response content",
169            provider
170        );
171    }
172}
173
174/// Capabilities advertised by an [`AiClient`] implementation.
175///
176/// Used by call sites to decide whether to attach a structured-response
177/// schema (or other backend-specific request options) before dispatching.
178/// The default value is the conservative ''nothing supported'' baseline so
179/// new fields can be added without forcing existing implementations to
180/// update.
181#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
182pub struct AiClientCapabilities {
183    /// Whether the backend can enforce a JSON Schema on its response.
184    ///
185    /// When `true`, the call site may set
186    /// [`RequestOptions::response_schema`]; the backend will hand the schema
187    /// to its underlying API (e.g. `claude -p --json-schema <file>`) and the
188    /// API re-prompts until the model produces a validating response.
189    pub supports_response_schema: bool,
190}
191
192/// Whether the response should be formatted as YAML (default) or JSON
193/// matching a schema.
194///
195/// Used by the prompts module to swap the format-specific portion of a
196/// structured prompt without rewriting the semantic instructions.
197#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
198pub enum ResponseFormat {
199    /// Plain YAML, with the prompt asking the model to emit a fenced or
200    /// bare YAML document.
201    #[default]
202    Yaml,
203    /// JSON object that matches a schema attached via
204    /// [`RequestOptions::response_schema`]. The prompt drops the YAML
205    /// structure literal and tells the model to return only the JSON
206    /// object.
207    JsonSchema,
208}
209
210impl ResponseFormat {
211    /// Returns the response format that should be used given a backend's
212    /// capabilities.
213    #[must_use]
214    pub fn from_capabilities(caps: &AiClientCapabilities) -> Self {
215        if caps.supports_response_schema {
216            Self::JsonSchema
217        } else {
218            Self::Yaml
219        }
220    }
221}
222
223/// Per-request options passed to [`AiClient::send_request_with_options`].
224///
225/// Schema and other knobs live on the request, not the client, so a shared
226/// client cannot leak settings between concurrent calls. Backends that do
227/// not support an option are expected to ignore it (and the call site is
228/// expected to consult [`AiClient::capabilities`] before setting it).
229#[derive(Clone, Debug, Default)]
230pub struct RequestOptions {
231    /// Optional JSON Schema (as a `serde_json::Value`) constraining the
232    /// model's response. Only honoured by backends whose
233    /// [`AiClientCapabilities::supports_response_schema`] is `true`.
234    pub response_schema: Option<Value>,
235}
236
237impl RequestOptions {
238    /// Returns a new [`RequestOptions`] with [`Self::response_schema`] set.
239    #[must_use]
240    pub fn with_response_schema(mut self, schema: Value) -> Self {
241        self.response_schema = Some(schema);
242        self
243    }
244}
245
246/// Trait for AI service clients.
247pub trait AiClient: Send + Sync {
248    /// Sends a request to the AI service and returns the raw response.
249    fn send_request<'a>(
250        &'a self,
251        system_prompt: &'a str,
252        user_prompt: &'a str,
253    ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>>;
254
255    /// Returns metadata about the AI client implementation.
256    fn get_metadata(&self) -> AiClientMetadata;
257
258    /// Returns the optional capabilities advertised by this backend.
259    ///
260    /// The default implementation returns the all-disabled baseline so
261    /// existing backends remain source-compatible. Backends that gain new
262    /// capabilities (e.g. structured-output enforcement) should override
263    /// this method.
264    fn capabilities(&self) -> AiClientCapabilities {
265        AiClientCapabilities::default()
266    }
267
268    /// Sends a request with optional per-request settings.
269    ///
270    /// The default implementation drops `options` and dispatches via
271    /// [`Self::send_request`]. Backends that honour any field in
272    /// [`RequestOptions`] (e.g. `response_schema`) override this method.
273    /// Backends that don't honour an option must ignore it; call sites
274    /// should consult [`capabilities`](Self::capabilities) before setting
275    /// options that not all backends support.
276    fn send_request_with_options<'a>(
277        &'a self,
278        system_prompt: &'a str,
279        user_prompt: &'a str,
280        _options: RequestOptions,
281    ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>> {
282        self.send_request(system_prompt, user_prompt)
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289
290    fn meta(provider: &str) -> AiClientMetadata {
291        AiClientMetadata {
292            provider: provider.to_string(),
293            model: "test-model".to_string(),
294            max_context_length: 1024,
295            max_response_length: 1024,
296            active_beta: None,
297        }
298    }
299
300    #[test]
301    fn prompt_style_openai() {
302        assert_eq!(meta("OpenAI").prompt_style(), PromptStyle::OpenAi);
303    }
304
305    #[test]
306    fn prompt_style_ollama() {
307        assert_eq!(meta("Ollama").prompt_style(), PromptStyle::OpenAi);
308    }
309
310    #[test]
311    fn prompt_style_anthropic() {
312        assert_eq!(meta("Anthropic").prompt_style(), PromptStyle::Claude);
313    }
314
315    #[test]
316    fn prompt_style_bedrock() {
317        assert_eq!(
318            meta("Anthropic Bedrock").prompt_style(),
319            PromptStyle::Claude
320        );
321    }
322
323    #[test]
324    fn prompt_style_unknown_defaults_to_claude() {
325        assert_eq!(meta("SomeNewProvider").prompt_style(), PromptStyle::Claude);
326    }
327
328    /// Ensure case-sensitive matching: "openai" (lowercase) is not a known provider
329    /// string and must not silently match as OpenAI.
330    #[test]
331    fn prompt_style_case_sensitive() {
332        assert_eq!(meta("openai").prompt_style(), PromptStyle::Claude);
333        assert_eq!(meta("ollama").prompt_style(), PromptStyle::Claude);
334    }
335
336    #[test]
337    fn capabilities_default_is_all_disabled() {
338        let caps = AiClientCapabilities::default();
339        assert!(!caps.supports_response_schema);
340    }
341
342    #[test]
343    fn response_format_default_is_yaml() {
344        assert_eq!(ResponseFormat::default(), ResponseFormat::Yaml);
345    }
346
347    #[test]
348    fn response_format_from_capabilities_disabled_picks_yaml() {
349        let caps = AiClientCapabilities::default();
350        assert_eq!(
351            ResponseFormat::from_capabilities(&caps),
352            ResponseFormat::Yaml
353        );
354    }
355
356    #[test]
357    fn response_format_from_capabilities_enabled_picks_json_schema() {
358        let caps = AiClientCapabilities {
359            supports_response_schema: true,
360        };
361        assert_eq!(
362            ResponseFormat::from_capabilities(&caps),
363            ResponseFormat::JsonSchema
364        );
365    }
366
367    #[test]
368    fn request_options_with_response_schema_sets_field() {
369        let value = serde_json::json!({"type": "object"});
370        let opts = RequestOptions::default().with_response_schema(value.clone());
371        assert_eq!(opts.response_schema, Some(value));
372    }
373
374    #[test]
375    fn request_options_default_has_no_schema() {
376        let opts = RequestOptions::default();
377        assert!(opts.response_schema.is_none());
378    }
379}
omni_dev/claude/ai.rs

omni_dev/claude/
ai.rs
