aa_core/
llm.rs

1//! LLM provider and model identifiers shared across the workspace.
2//!
3//! These types and the [`Model::infer_from_name`] helper were originally
4//! introduced inside `aa-gateway::budget::types` (AAASM-3353). They are
5//! relocated here (AAASM-3362) so that SDKs and other crates can reuse the
6//! provider/model taxonomy and the model-name → `(Provider, Model)` inference
7//! without taking a dependency on `aa-gateway`. Pricing tables remain in
8//! `aa-gateway` — only the enums and the inference logic are core-appropriate.
9
10/// LLM provider identifier.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
13#[cfg_attr(feature = "serde", serde(rename_all = "snake_case"))]
14pub enum Provider {
15    /// OpenAI (GPT-* models).
16    OpenAi,
17    /// Anthropic (Claude models).
18    Anthropic,
19    /// Cohere (Command models).
20    Cohere,
21}
22
23/// LLM model identifier.
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
25#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
26#[cfg_attr(feature = "serde", serde(rename_all = "snake_case"))]
27pub enum Model {
28    /// OpenAI GPT-4o.
29    Gpt4o,
30    /// OpenAI GPT-4.
31    Gpt4,
32    /// OpenAI GPT-3.5 Turbo.
33    Gpt35Turbo,
34    /// Anthropic Claude 3 Opus.
35    Claude3Opus,
36    /// Anthropic Claude 3 Sonnet.
37    Claude3Sonnet,
38    /// Anthropic Claude 3 Haiku.
39    Claude3Haiku,
40    /// Cohere Command R+.
41    CommandRPlus,
42    /// Cohere Command R.
43    CommandR,
44}
45
46impl Model {
47    /// Infer the `(Provider, Model)` pair from a free-form model name string.
48    ///
49    /// AAASM-3353 — the live `CheckAction` proto (`LlmCallContext`) carries only
50    /// the model name string; it does NOT carry a provider field. Rather than
51    /// change `proto/` (out of scope), the provider is inferred here from the
52    /// model name. The match is a case-insensitive substring test against the
53    /// known model families, ordered most-specific-first so that e.g.
54    /// `gpt-4o-2024-08-06` maps to `Gpt4o` (not `Gpt4`) and
55    /// `command-r-plus` maps to `CommandRPlus` (not `CommandR`).
56    ///
57    /// Returns `None` for an unrecognised model name. Callers must decide how
58    /// to price an unknown model — the gateway budget stage fails closed at a
59    /// conservative fallback rate (AAASM-4069) rather than treating it as free.
60    pub fn infer_from_name(name: &str) -> Option<(Provider, Self)> {
61        let n = name.to_ascii_lowercase();
62        // Most-specific patterns first; substrings of others must come earlier.
63        if n.contains("gpt-4o") || n.contains("gpt4o") {
64            Some((Provider::OpenAi, Model::Gpt4o))
65        } else if n.contains("gpt-3.5") || n.contains("gpt35") || n.contains("gpt-35") {
66            Some((Provider::OpenAi, Model::Gpt35Turbo))
67        } else if n.contains("gpt-4") || n.contains("gpt4") {
68            Some((Provider::OpenAi, Model::Gpt4))
69        } else if n.contains("opus") {
70            Some((Provider::Anthropic, Model::Claude3Opus))
71        } else if n.contains("sonnet") {
72            Some((Provider::Anthropic, Model::Claude3Sonnet))
73        } else if n.contains("haiku") {
74            Some((Provider::Anthropic, Model::Claude3Haiku))
75        } else if n.contains("command-r-plus") || n.contains("command-r+") {
76            Some((Provider::Cohere, Model::CommandRPlus))
77        } else if n.contains("command-r") {
78            Some((Provider::Cohere, Model::CommandR))
79        } else {
80            None
81        }
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    fn provider_variants_are_distinct() {
91        assert_eq!(Provider::OpenAi, Provider::OpenAi);
92        assert_ne!(Provider::OpenAi, Provider::Anthropic);
93        assert_ne!(Provider::OpenAi, Provider::Cohere);
94        assert_ne!(Provider::Anthropic, Provider::Cohere);
95    }
96
97    #[test]
98    fn model_variants_are_distinct() {
99        assert_eq!(Model::Gpt4o, Model::Gpt4o);
100        assert_ne!(Model::Gpt4o, Model::Gpt4);
101        assert_ne!(Model::Claude3Opus, Model::Claude3Haiku);
102        assert_ne!(Model::CommandRPlus, Model::CommandR);
103    }
104
105    #[test]
106    fn infer_openai_gpt4o_most_specific() {
107        assert_eq!(
108            Model::infer_from_name("gpt-4o-2024-08-06"),
109            Some((Provider::OpenAi, Model::Gpt4o))
110        );
111        assert_eq!(Model::infer_from_name("GPT4O"), Some((Provider::OpenAi, Model::Gpt4o)));
112    }
113
114    #[test]
115    fn infer_openai_gpt35_before_gpt4() {
116        assert_eq!(
117            Model::infer_from_name("gpt-3.5-turbo"),
118            Some((Provider::OpenAi, Model::Gpt35Turbo))
119        );
120    }
121
122    #[test]
123    fn infer_openai_gpt4() {
124        assert_eq!(Model::infer_from_name("gpt-4"), Some((Provider::OpenAi, Model::Gpt4)));
125    }
126
127    #[test]
128    fn infer_anthropic_family() {
129        assert_eq!(
130            Model::infer_from_name("claude-3-opus-20240229"),
131            Some((Provider::Anthropic, Model::Claude3Opus))
132        );
133        assert_eq!(
134            Model::infer_from_name("claude-3-5-sonnet"),
135            Some((Provider::Anthropic, Model::Claude3Sonnet))
136        );
137        assert_eq!(
138            Model::infer_from_name("claude-3-haiku"),
139            Some((Provider::Anthropic, Model::Claude3Haiku))
140        );
141    }
142
143    #[test]
144    fn infer_cohere_command_r_plus_before_command_r() {
145        assert_eq!(
146            Model::infer_from_name("command-r-plus"),
147            Some((Provider::Cohere, Model::CommandRPlus))
148        );
149        assert_eq!(
150            Model::infer_from_name("command-r"),
151            Some((Provider::Cohere, Model::CommandR))
152        );
153    }
154
155    #[test]
156    fn infer_unknown_returns_none() {
157        assert_eq!(Model::infer_from_name("mystery-model-v9"), None);
158        assert_eq!(Model::infer_from_name(""), None);
159    }
160}
aa_core/llm.rs

aa_core/
llm.rs
