oxi_ai/
types.rs

1//! Core domain types for oxi-ai
2
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::hash::Hash;
6
7/// Provider API identifier.
8///
9/// Selects the wire-format / protocol dialect spoken to a particular LLM provider.
10/// Moved to the `oxi-catalog` crate (omp aligns `KnownApi` with `pi-catalog`).
11/// Re-exported here for backward compatibility; new code should use
12/// `oxi_catalog::Api` directly.
13pub use oxi_catalog::Api;
14
15/// Cache retention preference
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
17#[serde(rename_all = "lowercase")]
18pub enum CacheRetention {
19    /// No caching (default).
20    #[default]
21    None,
22    /// Short-lived cache.
23    Short,
24    /// Long-lived cache.
25    Long,
26}
27
28/// Model thinking/reasoning level
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
30#[serde(rename_all = "lowercase")]
31#[non_exhaustive]
32pub enum ThinkingLevel {
33    /// Extended reasoning disabled (default).
34    #[default]
35    Off,
36    /// Minimal reasoning.
37    Minimal,
38    /// Low reasoning.
39    Low,
40    /// Medium reasoning.
41    Medium,
42    /// High reasoning.
43    High,
44    /// Very high reasoning.
45    XHigh,
46}
47
48impl ThinkingLevel {
49    /// Returns the reasoning level as a string. Returns `None` for `Off`.
50    pub fn as_str(&self) -> Option<&str> {
51        match self {
52            ThinkingLevel::Off => None,
53            ThinkingLevel::Minimal => Some("minimal"),
54            ThinkingLevel::Low => Some("low"),
55            ThinkingLevel::Medium => Some("medium"),
56            ThinkingLevel::High => Some("high"),
57            ThinkingLevel::XHigh => Some("xhigh"),
58        }
59    }
60}
61
62/// Input modalities
63#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
64#[serde(rename_all = "lowercase")]
65#[non_exhaustive]
66pub enum InputModality {
67    /// Text input.
68    Text,
69    /// Image input.
70    Image,
71}
72
73/// Cost structure – prices per million tokens.
74#[derive(Debug, Clone, Default, Serialize, Deserialize)]
75#[serde(default)]
76pub struct Cost {
77    /// Input token cost ($/M tokens).
78    #[serde(default)]
79    pub input: f64,
80    /// Output token cost ($/M tokens).
81    #[serde(default)]
82    pub output: f64,
83    /// Cached-input read cost ($/M tokens).
84    #[serde(default)]
85    pub cache_read: f64,
86    /// Cache write cost ($/M tokens).
87    #[serde(default)]
88    pub cache_write: f64,
89}
90
91impl Cost {
92    /// Sum of all cost components.
93    pub fn total(&self) -> f64 {
94        self.input + self.output + self.cache_read + self.cache_write
95    }
96}
97
98/// Stop reason – why the model finished generating.
99#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
100#[serde(rename_all = "camelCase")]
101#[non_exhaustive]
102pub enum StopReason {
103    /// Normal stop – the model finished its response.
104    Stop,
105    /// Hit the maximum output token limit.
106    Length,
107    /// Stopped to invoke a tool.
108    ToolUse,
109    /// An error occurred during generation.
110    Error,
111    /// Generation was aborted by the client.
112    Aborted,
113}
114
115/// Token usage statistics.
116#[derive(Debug, Clone, Default, Serialize, Deserialize)]
117pub struct Usage {
118    /// Number of input (prompt) tokens.
119    #[serde(default)]
120    pub input: usize,
121    /// Number of output (completion) tokens.
122    #[serde(default)]
123    pub output: usize,
124    /// Number of tokens read from cache.
125    #[serde(default)]
126    pub cache_read: usize,
127    /// Number of tokens written to cache.
128    #[serde(default)]
129    pub cache_write: usize,
130    /// Total tokens (input + output + cache).
131    #[serde(default)]
132    pub total_tokens: usize,
133    /// Computed cost in dollars.
134    #[serde(default)]
135    pub cost: Cost,
136}
137
138impl Usage {
139    /// Recalculate `total_tokens` and per-component costs from raw token counts.
140    ///
141    /// If pricing parameters are provided, they override the default $1/M rate.
142    pub fn calculate_cost(
143        &mut self,
144        input_cost_per_million: Option<f64>,
145        output_cost_per_million: Option<f64>,
146    ) {
147        self.total_tokens = self.input + self.output + self.cache_read + self.cache_write;
148        self.cost.input = input_cost_per_million.unwrap_or(1.0) * self.input as f64 / 1_000_000.0;
149        self.cost.output =
150            output_cost_per_million.unwrap_or(1.0) * self.output as f64 / 1_000_000.0;
151        self.cost.cache_read = (self.cache_read as f64) / 1_000_000.0;
152        self.cost.cache_write = (self.cache_write as f64) / 1_000_000.0;
153    }
154}
155
156/// Compatibility settings for OpenAI-compatible APIs.
157///
158/// Not every OpenAI-compatible provider supports every feature.
159/// These flags let the streaming layer adapt its request shape.
160#[derive(Debug, Clone, Default, Serialize, Deserialize)]
161#[serde(default)]
162pub struct CompatSettings {
163    /// Whether the provider supports the `store` parameter.
164    #[serde(default = "default_true")]
165    pub supports_store: bool,
166    /// Whether the provider recognises the `developer` role.
167    #[serde(default = "default_true")]
168    pub supports_developer_role: bool,
169    /// Whether the provider supports `reasoning_effort`.
170    #[serde(default = "default_true")]
171    pub supports_reasoning_effort: bool,
172    /// Whether the provider returns usage data in streaming responses.
173    #[serde(default = "default_true")]
174    pub supports_usage_in_streaming: bool,
175    /// Which JSON field name to use for the max-tokens parameter.
176    #[serde(default)]
177    pub max_tokens_field: Option<MaxTokensField>,
178    /// Whether tool results must include the tool name.
179    #[serde(default = "default_false")]
180    pub requires_tool_result_name: bool,
181    /// Whether an assistant message must follow every tool result.
182    #[serde(default = "default_false")]
183    pub requires_assistant_after_tool_result: bool,
184    /// Whether thinking should be sent as plain text.
185    #[serde(default = "default_false")]
186    pub requires_thinking_as_text: bool,
187    /// Provider-specific thinking wire-format.
188    #[serde(default)]
189    pub thinking_format: Option<ThinkingFormat>,
190}
191
192fn default_true() -> bool {
193    true
194}
195fn default_false() -> bool {
196    false
197}
198
199/// Which JSON field to use for the maximum output token count.
200#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
201#[serde(rename_all = "kebab-case")]
202pub enum MaxTokensField {
203    /// Use `max_completion_tokens`.
204    MaxCompletionTokens,
205    /// Use `max_tokens`.
206    MaxTokens,
207}
208
209/// Provider-specific wire format for extended thinking.
210#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
211#[serde(rename_all = "lowercase")]
212pub enum ThinkingFormat {
213    /// OpenAI native thinking format.
214    OpenAI,
215    /// OpenRouter thinking format.
216    OpenRouter,
217    /// DeepSeek thinking format.
218    DeepSeek,
219    /// Zai thinking format.
220    Zai,
221    /// Qwen API thinking format.
222    Qwen,
223    /// Qwen chat-template thinking format.
224    QwenChatTemplate,
225}
226
227/// Task complexity level for routing decisions.
228#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
229pub enum Complexity {
230    /// Simple, single-step tasks (e.g., "translate this text")
231    Trivial,
232    /// Routine tasks needing moderate reasoning (e.g., "write a function")
233    Simple,
234    /// Tasks requiring multi-step reasoning (e.g., "architect a service")
235    Moderate,
236    /// Complex tasks needing deep analysis (e.g., "write a full codebase")
237    #[default]
238    Complex,
239    /// Research-grade tasks needing the best models
240    Research,
241}
242
243impl Complexity {
244    /// Returns the relative cost tier (0=cheapest, 4=most expensive) for routing
245    pub fn cost_tier(&self) -> u8 {
246        match self {
247            Self::Trivial => 0,
248            Self::Simple => 1,
249            Self::Moderate => 2,
250            Self::Complex => 3,
251            Self::Research => 4,
252        }
253    }
254}
255
256/// Tool result returned by agent tool execution.
257#[derive(Debug, Clone, Serialize, Deserialize)]
258pub struct ToolResult {
259    /// ID of the tool call this result corresponds to.
260    pub tool_call_id: String,
261    /// Human-readable result or error text.
262    pub content: String,
263    /// `"success"` or `"error"`.
264    pub status: String,
265}
266
267impl ToolResult {
268    /// Create a successful tool result.
269    pub fn success(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
270        Self {
271            tool_call_id: tool_call_id.into(),
272            content: content.into(),
273            status: "success".to_string(),
274        }
275    }
276
277    /// Create an error tool result.
278    pub fn error(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
279        Self {
280            tool_call_id: tool_call_id.into(),
281            content: content.into(),
282            status: "error".to_string(),
283        }
284    }
285
286    /// Returns `true` if this result represents an error.
287    pub fn is_error(&self) -> bool {
288        self.status == "error"
289    }
290}
291
292/// Images API provider type.
293#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
294#[non_exhaustive]
295pub enum ImagesApi {
296    /// OpenRouter API (supports multiple image generation models).
297    OpenRouter,
298}
299
300impl std::fmt::Display for ImagesApi {
301    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
302        match self {
303            ImagesApi::OpenRouter => write!(f, "openrouter"),
304        }
305    }
306}
307
308/// Request for image generation via an Images API provider.
309#[derive(Debug, Clone, Serialize, Deserialize)]
310#[serde(default)]
311pub struct ImageGenerationRequest {
312    /// The text prompt describing the desired image.
313    pub prompt: String,
314    /// Model identifier (e.g. `"openai/dall-e-3"`, `"black-forest-labs/flux-1-dev"`).
315    pub model: Option<String>,
316    /// Output size. Provider-dependent. Examples: `"1024x1024"`, `"1024x1792"`.
317    pub size: Option<String>,
318    /// Number of images to generate (default 1).
319    pub n: Option<u32>,
320    /// Output format: `"url"` (default) or `"b64_json"`.
321    pub response_format: Option<String>,
322}
323
324impl Default for ImageGenerationRequest {
325    fn default() -> Self {
326        Self {
327            prompt: String::new(),
328            model: None,
329            size: None,
330            n: Some(1),
331            response_format: Some("b64_json".to_string()),
332        }
333    }
334}
335
336/// Response from an image generation API call.
337#[derive(Debug, Clone, Serialize, Deserialize, Default)]
338#[serde(default)]
339pub struct ImageGenerationResponse {
340    /// Vector of generated image bytes (one per `n`). Raw PNG/JPEG data.
341    pub images: Vec<Vec<u8>>,
342    /// Revised prompt from the model (providers may rewrite prompts).
343    pub revised_prompt: Option<String>,
344}
345
346/// LLM model definition.
347///
348/// Describes a model's capabilities, endpoint, and cost structure.
349#[derive(Debug, Clone, Serialize, Deserialize)]
350pub struct Model {
351    /// Unique model identifier (e.g. `"gpt-4o"`, `"claude-3-5-sonnet"`).
352    pub id: String,
353    /// Human-readable display name.
354    pub name: String,
355    /// Which API dialect this model speaks.
356    pub api: Api,
357    /// Provider name (e.g. `"openai"`, `"anthropic"`).
358    pub provider: String,
359    /// Base URL for the provider API.
360    pub base_url: String,
361    /// Whether this model supports extended reasoning / thinking.
362    #[serde(default)]
363    pub reasoning: bool,
364    /// Supported input modalities.
365    #[serde(default)]
366    pub input: Vec<InputModality>,
367    /// Pricing information.
368    #[serde(default)]
369    pub cost: Cost,
370    /// Maximum context window in tokens.
371    pub context_window: usize,
372    /// Maximum output tokens per request.
373    pub max_tokens: usize,
374    /// Extra HTTP headers to send with every request.
375    #[serde(default)]
376    pub headers: HashMap<String, String>,
377    /// Compatibility tweaks for non-standard providers.
378    #[serde(default)]
379    pub compat: Option<CompatSettings>,
380}
381
382impl Model {
383    /// Create a new model with sensible defaults.
384    pub fn new(
385        id: impl Into<String>,
386        name: impl Into<String>,
387        api: Api,
388        provider: impl Into<String>,
389        base_url: impl Into<String>,
390    ) -> Self {
391        Self {
392            id: id.into(),
393            name: name.into(),
394            api,
395            provider: provider.into(),
396            base_url: base_url.into(),
397            reasoning: false,
398            input: vec![InputModality::Text],
399            cost: Cost::default(),
400            context_window: 128_000,
401            max_tokens: 32_000,
402            headers: HashMap::new(),
403            compat: None,
404        }
405    }
406
407    /// Returns `true` if the model accepts image inputs.
408    pub fn supports_vision(&self) -> bool {
409        self.input.contains(&InputModality::Image)
410    }
411
412    /// Returns `true` if the model supports extended reasoning.
413    pub fn supports_reasoning(&self) -> bool {
414        self.reasoning
415    }
416}
417
418#[cfg(test)]
419mod tests {
420    use super::*;
421
422    #[test]
423    fn model_roundtrip() {
424        let mut model = Model::new(
425            "gpt-4o",
426            "GPT-4o",
427            Api::OpenAiCompletions,
428            "openai",
429            "https://api.openai.com/v1",
430        );
431        model.reasoning = true;
432        model.input.push(InputModality::Image);
433        model.cost = Cost {
434            input: 5.0,
435            output: 15.0,
436            cache_read: 2.5,
437            cache_write: 0.0,
438        };
439        model.compat = Some(CompatSettings::default());
440
441        let json = serde_json::to_string(&model).unwrap();
442        let deserialized: Model = serde_json::from_str(&json).unwrap();
443
444        assert_eq!(deserialized.id, "gpt-4o");
445        assert_eq!(deserialized.name, "GPT-4o");
446        assert_eq!(deserialized.api, Api::OpenAiCompletions);
447        assert_eq!(deserialized.provider, "openai");
448        assert!(deserialized.reasoning);
449        assert!(deserialized.supports_vision());
450        assert!(deserialized.supports_reasoning());
451        assert_eq!(deserialized.cost.input, 5.0);
452        assert_eq!(deserialized.cost.output, 15.0);
453    }
454
455    #[test]
456    fn usage_calculate_cost() {
457        let mut usage = Usage {
458            input: 1_000_000,
459            output: 500_000,
460            cache_read: 200_000,
461            cache_write: 100_000,
462            ..Default::default()
463        };
464        usage.calculate_cost(None, None);
465
466        assert_eq!(usage.total_tokens, 1_800_000);
467        assert_eq!(usage.cost.input, 1.0);
468        assert_eq!(usage.cost.output, 0.5);
469        assert_eq!(usage.cost.cache_read, 0.2);
470        assert_eq!(usage.cost.cache_write, 0.1);
471    }
472
473    #[test]
474    fn cost_total() {
475        let cost = Cost {
476            input: 3.0,
477            output: 6.0,
478            cache_read: 1.0,
479            cache_write: 0.5,
480        };
481        assert!((cost.total() - 10.5).abs() < f64::EPSILON);
482
483        let default_cost = Cost::default();
484        assert_eq!(default_cost.total(), 0.0);
485    }
486
487    #[test]
488    fn api_display() {
489        assert_eq!(Api::OpenAiCompletions.to_string(), "openai-completions");
490        assert_eq!(Api::OpenAiResponses.to_string(), "openai-responses");
491        assert_eq!(Api::AnthropicMessages.to_string(), "anthropic-messages");
492        assert_eq!(Api::GoogleGenerativeAi.to_string(), "google-generative-ai");
493        assert_eq!(Api::GoogleVertex.to_string(), "google-vertex");
494        assert_eq!(
495            Api::AzureOpenAiResponses.to_string(),
496            "azure-openai-responses"
497        );
498        assert_eq!(
499            Api::BedrockConverseStream.to_string(),
500            "bedrock-converse-stream"
501        );
502    }
503
504    #[test]
505    fn api_serde_roundtrip() {
506        for api in [
507            Api::OpenAiCompletions,
508            Api::OpenAiResponses,
509            Api::AnthropicMessages,
510            Api::GoogleGenerativeAi,
511            Api::GoogleVertex,
512            Api::AzureOpenAiResponses,
513            Api::BedrockConverseStream,
514        ] {
515            let json = serde_json::to_string(&api).unwrap();
516            let back: Api = serde_json::from_str(&json).unwrap();
517            assert_eq!(api, back);
518        }
519    }
520
521    #[test]
522    fn thinking_level_serde() {
523        for level in [
524            ThinkingLevel::Off,
525            ThinkingLevel::Minimal,
526            ThinkingLevel::Low,
527            ThinkingLevel::Medium,
528            ThinkingLevel::High,
529            ThinkingLevel::XHigh,
530        ] {
531            let json = serde_json::to_string(&level).unwrap();
532            let back: ThinkingLevel = serde_json::from_str(&json).unwrap();
533            assert_eq!(level, back);
534        }
535        // Verify default
536        assert_eq!(ThinkingLevel::default(), ThinkingLevel::Off);
537        // Verify rename values
538        assert_eq!(
539            serde_json::to_string(&ThinkingLevel::High).unwrap(),
540            "\"high\""
541        );
542        assert_eq!(
543            serde_json::to_string(&ThinkingLevel::Off).unwrap(),
544            "\"off\""
545        );
546        // as_str
547        assert!(ThinkingLevel::Off.as_str().is_none());
548        assert_eq!(ThinkingLevel::High.as_str(), Some("high"));
549        assert_eq!(ThinkingLevel::XHigh.as_str(), Some("xhigh"));
550    }
551
552    #[test]
553    fn stop_reason_serde() {
554        assert_eq!(
555            serde_json::to_string(&StopReason::ToolUse).unwrap(),
556            "\"toolUse\""
557        );
558        let back: StopReason = serde_json::from_str("\"toolUse\"").unwrap();
559        assert_eq!(back, StopReason::ToolUse);
560    }
561
562    #[test]
563    fn tool_result_helpers() {
564        let success = ToolResult::success("call_1", "result text");
565        assert_eq!(success.tool_call_id, "call_1");
566        assert_eq!(success.content, "result text");
567        assert_eq!(success.status, "success");
568        assert!(!success.is_error());
569
570        let error = ToolResult::error("call_2", "something failed");
571        assert!(error.is_error());
572        assert_eq!(error.status, "error");
573    }
574
575    #[test]
576    fn cache_retention_default() {
577        assert_eq!(CacheRetention::default(), CacheRetention::None);
578    }
579}
oxi_ai/types.rs

oxi_ai/
types.rs
