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 — the caller treats an
58    /// unknown model as zero cost (no spend accrued) rather than guessing.
59    pub fn infer_from_name(name: &str) -> Option<(Provider, Self)> {
60        let n = name.to_ascii_lowercase();
61        // Most-specific patterns first; substrings of others must come earlier.
62        if n.contains("gpt-4o") || n.contains("gpt4o") {
63            Some((Provider::OpenAi, Model::Gpt4o))
64        } else if n.contains("gpt-3.5") || n.contains("gpt35") || n.contains("gpt-35") {
65            Some((Provider::OpenAi, Model::Gpt35Turbo))
66        } else if n.contains("gpt-4") || n.contains("gpt4") {
67            Some((Provider::OpenAi, Model::Gpt4))
68        } else if n.contains("opus") {
69            Some((Provider::Anthropic, Model::Claude3Opus))
70        } else if n.contains("sonnet") {
71            Some((Provider::Anthropic, Model::Claude3Sonnet))
72        } else if n.contains("haiku") {
73            Some((Provider::Anthropic, Model::Claude3Haiku))
74        } else if n.contains("command-r-plus") || n.contains("command-r+") {
75            Some((Provider::Cohere, Model::CommandRPlus))
76        } else if n.contains("command-r") {
77            Some((Provider::Cohere, Model::CommandR))
78        } else {
79            None
80        }
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    #[test]
89    fn provider_variants_are_distinct() {
90        assert_eq!(Provider::OpenAi, Provider::OpenAi);
91        assert_ne!(Provider::OpenAi, Provider::Anthropic);
92        assert_ne!(Provider::OpenAi, Provider::Cohere);
93        assert_ne!(Provider::Anthropic, Provider::Cohere);
94    }
95
96    #[test]
97    fn model_variants_are_distinct() {
98        assert_eq!(Model::Gpt4o, Model::Gpt4o);
99        assert_ne!(Model::Gpt4o, Model::Gpt4);
100        assert_ne!(Model::Claude3Opus, Model::Claude3Haiku);
101        assert_ne!(Model::CommandRPlus, Model::CommandR);
102    }
103
104    #[test]
105    fn infer_openai_gpt4o_most_specific() {
106        assert_eq!(
107            Model::infer_from_name("gpt-4o-2024-08-06"),
108            Some((Provider::OpenAi, Model::Gpt4o))
109        );
110        assert_eq!(Model::infer_from_name("GPT4O"), Some((Provider::OpenAi, Model::Gpt4o)));
111    }
112
113    #[test]
114    fn infer_openai_gpt35_before_gpt4() {
115        assert_eq!(
116            Model::infer_from_name("gpt-3.5-turbo"),
117            Some((Provider::OpenAi, Model::Gpt35Turbo))
118        );
119    }
120
121    #[test]
122    fn infer_openai_gpt4() {
123        assert_eq!(Model::infer_from_name("gpt-4"), Some((Provider::OpenAi, Model::Gpt4)));
124    }
125
126    #[test]
127    fn infer_anthropic_family() {
128        assert_eq!(
129            Model::infer_from_name("claude-3-opus-20240229"),
130            Some((Provider::Anthropic, Model::Claude3Opus))
131        );
132        assert_eq!(
133            Model::infer_from_name("claude-3-5-sonnet"),
134            Some((Provider::Anthropic, Model::Claude3Sonnet))
135        );
136        assert_eq!(
137            Model::infer_from_name("claude-3-haiku"),
138            Some((Provider::Anthropic, Model::Claude3Haiku))
139        );
140    }
141
142    #[test]
143    fn infer_cohere_command_r_plus_before_command_r() {
144        assert_eq!(
145            Model::infer_from_name("command-r-plus"),
146            Some((Provider::Cohere, Model::CommandRPlus))
147        );
148        assert_eq!(
149            Model::infer_from_name("command-r"),
150            Some((Provider::Cohere, Model::CommandR))
151        );
152    }
153
154    #[test]
155    fn infer_unknown_returns_none() {
156        assert_eq!(Model::infer_from_name("mystery-model-v9"), None);
157        assert_eq!(Model::infer_from_name(""), None);
158    }
159}
aa_core/llm.rs

aa_core/
llm.rs
