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 HTTP record for one AI-backend request attempt. The
82/// `service` tag distinguishes the backend that issued it (`anthropic`,
83/// `bedrock`, `openai`, or `ollama`) so their traffic is filterable apart;
84/// `method` covers both the chat `POST` and the metadata probes (GET/POST).
85pub(crate) fn record_ai_http(
86    service: &str,
87    method: &str,
88    url: &str,
89    started: Instant,
90    result: &reqwest::Result<reqwest::Response>,
91) {
92    request_log::record_http_result(service, method, url, started, result);
93}
94
95/// Returns the maximum output tokens for a model from the registry,
96/// respecting beta overrides.
97#[must_use]
98pub(crate) fn registry_max_output_tokens(
99    model: &str,
100    active_beta: &Option<(String, String)>,
101) -> i32 {
102    let registry = get_model_registry();
103    if let Some((_, value)) = active_beta {
104        registry.get_max_output_tokens_with_beta(model, value) as i32
105    } else {
106        registry.get_max_output_tokens(model) as i32
107    }
108}
109
110/// Returns the (input context length, max response length) for a model
111/// from the registry, respecting beta overrides.
112#[must_use]
113pub(crate) fn registry_model_limits(
114    model: &str,
115    active_beta: &Option<(String, String)>,
116) -> (usize, usize) {
117    let registry = get_model_registry();
118    match active_beta {
119        Some((_, value)) => (
120            registry.get_input_context_with_beta(model, value),
121            registry.get_max_output_tokens_with_beta(model, value),
122        ),
123        None => (
124            registry.get_input_context(model),
125            registry.get_max_output_tokens(model),
126        ),
127    }
128}
129
130/// Checks an HTTP response for error status and returns a structured error
131/// if non-success.
132///
133/// On success, returns the response unchanged for further processing.
134/// On failure, reads the error body and returns a
135/// [`ClaudeError::ApiRequestFailed`].
136pub(crate) async fn check_error_response(response: reqwest::Response) -> Result<reqwest::Response> {
137    if response.status().is_success() {
138        return Ok(response);
139    }
140    let status = response.status();
141    let error_text = response.text().await.unwrap_or_else(|e| {
142        tracing::debug!("Failed to read error response body: {e}");
143        String::new()
144    });
145    Err(ClaudeError::ApiRequestFailed(format!("HTTP {status}: {error_text}")).into())
146}
147
148/// Logs successful text extraction from an AI API response.
149pub(crate) fn log_response_success(provider: &str, result: &Result<String>) {
150    if let Ok(text) = result {
151        tracing::debug!(
152            response_len = text.len(),
153            "Successfully extracted text content from {} API response",
154            provider
155        );
156        tracing::debug!(
157            response_content = %text,
158            "{} API response content",
159            provider
160        );
161    }
162}
163
164/// Capabilities advertised by an [`AiClient`] implementation.
165///
166/// Used by call sites to decide whether to attach a structured-response
167/// schema (or other backend-specific request options) before dispatching.
168/// The default value is the conservative ''nothing supported'' baseline so
169/// new fields can be added without forcing existing implementations to
170/// update.
171#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
172pub struct AiClientCapabilities {
173    /// Whether the backend can enforce a JSON Schema on its response.
174    ///
175    /// When `true`, the call site may set
176    /// [`RequestOptions::response_schema`]; the backend will hand the schema
177    /// to its underlying API (e.g. `claude -p --json-schema <file>`) and the
178    /// API re-prompts until the model produces a validating response.
179    pub supports_response_schema: bool,
180}
181
182/// Whether the response should be formatted as YAML (default) or JSON
183/// matching a schema.
184///
185/// Used by the prompts module to swap the format-specific portion of a
186/// structured prompt without rewriting the semantic instructions.
187#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
188pub enum ResponseFormat {
189    /// Plain YAML, with the prompt asking the model to emit a fenced or
190    /// bare YAML document.
191    #[default]
192    Yaml,
193    /// JSON object that matches a schema attached via
194    /// [`RequestOptions::response_schema`]. The prompt drops the YAML
195    /// structure literal and tells the model to return only the JSON
196    /// object.
197    JsonSchema,
198}
199
200impl ResponseFormat {
201    /// Returns the response format that should be used given a backend's
202    /// capabilities.
203    #[must_use]
204    pub fn from_capabilities(caps: &AiClientCapabilities) -> Self {
205        if caps.supports_response_schema {
206            Self::JsonSchema
207        } else {
208            Self::Yaml
209        }
210    }
211}
212
213/// Per-request options passed to [`AiClient::send_request_with_options`].
214///
215/// Schema and other knobs live on the request, not the client, so a shared
216/// client cannot leak settings between concurrent calls. Backends that do
217/// not support an option are expected to ignore it (and the call site is
218/// expected to consult [`AiClient::capabilities`] before setting it).
219#[derive(Clone, Debug, Default)]
220pub struct RequestOptions {
221    /// Optional JSON Schema (as a `serde_json::Value`) constraining the
222    /// model's response. Only honoured by backends whose
223    /// [`AiClientCapabilities::supports_response_schema`] is `true`.
224    pub response_schema: Option<Value>,
225}
226
227impl RequestOptions {
228    /// Returns a new [`RequestOptions`] with [`Self::response_schema`] set.
229    #[must_use]
230    pub fn with_response_schema(mut self, schema: Value) -> Self {
231        self.response_schema = Some(schema);
232        self
233    }
234}
235
236/// Trait for AI service clients.
237pub trait AiClient: Send + Sync {
238    /// Sends a request to the AI service and returns the raw response.
239    fn send_request<'a>(
240        &'a self,
241        system_prompt: &'a str,
242        user_prompt: &'a str,
243    ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>>;
244
245    /// Returns metadata about the AI client implementation.
246    fn get_metadata(&self) -> AiClientMetadata;
247
248    /// Returns the optional capabilities advertised by this backend.
249    ///
250    /// The default implementation returns the all-disabled baseline so
251    /// existing backends remain source-compatible. Backends that gain new
252    /// capabilities (e.g. structured-output enforcement) should override
253    /// this method.
254    fn capabilities(&self) -> AiClientCapabilities {
255        AiClientCapabilities::default()
256    }
257
258    /// Sends a request with optional per-request settings.
259    ///
260    /// The default implementation drops `options` and dispatches via
261    /// [`Self::send_request`]. Backends that honour any field in
262    /// [`RequestOptions`] (e.g. `response_schema`) override this method.
263    /// Backends that don't honour an option must ignore it; call sites
264    /// should consult [`capabilities`](Self::capabilities) before setting
265    /// options that not all backends support.
266    fn send_request_with_options<'a>(
267        &'a self,
268        system_prompt: &'a str,
269        user_prompt: &'a str,
270        _options: RequestOptions,
271    ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>> {
272        self.send_request(system_prompt, user_prompt)
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279
280    fn meta(provider: &str) -> AiClientMetadata {
281        AiClientMetadata {
282            provider: provider.to_string(),
283            model: "test-model".to_string(),
284            max_context_length: 1024,
285            max_response_length: 1024,
286            active_beta: None,
287        }
288    }
289
290    #[test]
291    fn prompt_style_openai() {
292        assert_eq!(meta("OpenAI").prompt_style(), PromptStyle::OpenAi);
293    }
294
295    #[test]
296    fn prompt_style_ollama() {
297        assert_eq!(meta("Ollama").prompt_style(), PromptStyle::OpenAi);
298    }
299
300    #[test]
301    fn prompt_style_anthropic() {
302        assert_eq!(meta("Anthropic").prompt_style(), PromptStyle::Claude);
303    }
304
305    #[test]
306    fn prompt_style_bedrock() {
307        assert_eq!(
308            meta("Anthropic Bedrock").prompt_style(),
309            PromptStyle::Claude
310        );
311    }
312
313    #[test]
314    fn prompt_style_unknown_defaults_to_claude() {
315        assert_eq!(meta("SomeNewProvider").prompt_style(), PromptStyle::Claude);
316    }
317
318    /// Ensure case-sensitive matching: "openai" (lowercase) is not a known provider
319    /// string and must not silently match as OpenAI.
320    #[test]
321    fn prompt_style_case_sensitive() {
322        assert_eq!(meta("openai").prompt_style(), PromptStyle::Claude);
323        assert_eq!(meta("ollama").prompt_style(), PromptStyle::Claude);
324    }
325
326    #[test]
327    fn capabilities_default_is_all_disabled() {
328        let caps = AiClientCapabilities::default();
329        assert!(!caps.supports_response_schema);
330    }
331
332    #[test]
333    fn response_format_default_is_yaml() {
334        assert_eq!(ResponseFormat::default(), ResponseFormat::Yaml);
335    }
336
337    #[test]
338    fn response_format_from_capabilities_disabled_picks_yaml() {
339        let caps = AiClientCapabilities::default();
340        assert_eq!(
341            ResponseFormat::from_capabilities(&caps),
342            ResponseFormat::Yaml
343        );
344    }
345
346    #[test]
347    fn response_format_from_capabilities_enabled_picks_json_schema() {
348        let caps = AiClientCapabilities {
349            supports_response_schema: true,
350        };
351        assert_eq!(
352            ResponseFormat::from_capabilities(&caps),
353            ResponseFormat::JsonSchema
354        );
355    }
356
357    #[test]
358    fn request_options_with_response_schema_sets_field() {
359        let value = serde_json::json!({"type": "object"});
360        let opts = RequestOptions::default().with_response_schema(value.clone());
361        assert_eq!(opts.response_schema, Some(value));
362    }
363
364    #[test]
365    fn request_options_default_has_no_schema() {
366        let opts = RequestOptions::default();
367        assert!(opts.response_schema.is_none());
368    }
369}
omni_dev/claude/ai.rs

omni_dev/claude/
ai.rs
