stakpak_shared/models/integrations/
anthropic.rs

1use crate::models::llm::{
2    LLMChoice, LLMCompletionResponse, LLMMessage, LLMMessageContent, LLMTokenUsage, LLMTool,
3    PromptTokensDetails,
4};
5use crate::models::model_pricing::{ContextAware, ContextPricingTier, ModelContextInfo};
6use serde::{Deserialize, Serialize};
7
8#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
9pub enum AnthropicModel {
10    #[serde(rename = "claude-haiku-4-5-20251001")]
11    Claude45Haiku,
12    #[serde(rename = "claude-sonnet-4-5-20250929")]
13    Claude45Sonnet,
14    #[serde(rename = "claude-opus-4-5-20251101")]
15    Claude45Opus,
16}
17impl std::fmt::Display for AnthropicModel {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        match self {
20            AnthropicModel::Claude45Haiku => write!(f, "claude-haiku-4-5-20251001"),
21            AnthropicModel::Claude45Sonnet => write!(f, "claude-sonnet-4-5-20250929"),
22            AnthropicModel::Claude45Opus => write!(f, "claude-opus-4-5-20251101"),
23        }
24    }
25}
26
27impl AnthropicModel {
28    pub fn from_string(s: &str) -> Result<Self, String> {
29        serde_json::from_value(serde_json::Value::String(s.to_string()))
30            .map_err(|_| "Failed to deserialize Anthropic model".to_string())
31    }
32
33    /// Default smart model for Anthropic
34    pub const DEFAULT_SMART_MODEL: AnthropicModel = AnthropicModel::Claude45Opus;
35
36    /// Default eco model for Anthropic
37    pub const DEFAULT_ECO_MODEL: AnthropicModel = AnthropicModel::Claude45Haiku;
38
39    /// Default recovery model for Anthropic
40    pub const DEFAULT_RECOVERY_MODEL: AnthropicModel = AnthropicModel::Claude45Haiku;
41
42    /// Get default smart model as string
43    pub fn default_smart_model() -> String {
44        Self::DEFAULT_SMART_MODEL.to_string()
45    }
46
47    /// Get default eco model as string
48    pub fn default_eco_model() -> String {
49        Self::DEFAULT_ECO_MODEL.to_string()
50    }
51
52    /// Get default recovery model as string
53    pub fn default_recovery_model() -> String {
54        Self::DEFAULT_RECOVERY_MODEL.to_string()
55    }
56}
57
58impl ContextAware for AnthropicModel {
59    fn context_info(&self) -> ModelContextInfo {
60        let model_name = self.to_string();
61
62        if model_name.starts_with("claude-haiku") {
63            return ModelContextInfo {
64                max_tokens: 200_000,
65                pricing_tiers: vec![ContextPricingTier {
66                    label: "Standard".to_string(),
67                    input_cost_per_million: 1.0,
68                    output_cost_per_million: 5.0,
69                    upper_bound: None,
70                }],
71                approach_warning_threshold: 0.8,
72            };
73        }
74
75        if model_name.starts_with("claude-sonnet") {
76            return ModelContextInfo {
77                max_tokens: 1_000_000,
78                pricing_tiers: vec![
79                    ContextPricingTier {
80                        label: "<200K tokens".to_string(),
81                        input_cost_per_million: 3.0,
82                        output_cost_per_million: 15.0,
83                        upper_bound: Some(200_000),
84                    },
85                    ContextPricingTier {
86                        label: ">200K tokens".to_string(),
87                        input_cost_per_million: 6.0,
88                        output_cost_per_million: 22.5,
89                        upper_bound: None,
90                    },
91                ],
92                approach_warning_threshold: 0.8,
93            };
94        }
95
96        if model_name.starts_with("claude-opus") {
97            return ModelContextInfo {
98                max_tokens: 200_000,
99                pricing_tiers: vec![ContextPricingTier {
100                    label: "Standard".to_string(),
101                    input_cost_per_million: 5.0,
102                    output_cost_per_million: 25.0,
103                    upper_bound: None,
104                }],
105                approach_warning_threshold: 0.8,
106            };
107        }
108
109        panic!("Unknown model: {}", model_name);
110    }
111
112    fn model_name(&self) -> String {
113        match self {
114            AnthropicModel::Claude45Sonnet => "Claude Sonnet 4.5".to_string(),
115            AnthropicModel::Claude45Haiku => "Claude Haiku 4.5".to_string(),
116            AnthropicModel::Claude45Opus => "Claude Opus 4.5".to_string(),
117        }
118    }
119}
120
121#[derive(Serialize, Deserialize, Debug)]
122pub struct AnthropicInput {
123    pub model: AnthropicModel,
124    pub messages: Vec<LLMMessage>,
125    pub grammar: Option<String>,
126    pub max_tokens: u32,
127    pub stop_sequences: Option<Vec<String>>,
128    pub tools: Option<Vec<LLMTool>>,
129    pub thinking: ThinkingInput,
130}
131
132#[derive(Serialize, Deserialize, Debug)]
133pub struct ThinkingInput {
134    pub r#type: ThinkingType,
135    // Must be ≥1024 and less than max_tokens
136    pub budget_tokens: u32,
137}
138
139impl Default for ThinkingInput {
140    fn default() -> Self {
141        Self {
142            r#type: ThinkingType::default(),
143            budget_tokens: 1024,
144        }
145    }
146}
147
148#[derive(Serialize, Deserialize, Debug, Default)]
149#[serde(rename_all = "lowercase")]
150pub enum ThinkingType {
151    Enabled,
152    #[default]
153    Disabled,
154}
155
156#[derive(Serialize, Deserialize, Debug)]
157pub struct AnthropicOutputUsage {
158    pub input_tokens: u32,
159    pub output_tokens: u32,
160    #[serde(default)]
161    pub cache_creation_input_tokens: Option<u32>,
162    #[serde(default)]
163    pub cache_read_input_tokens: Option<u32>,
164}
165
166#[derive(Serialize, Deserialize, Debug)]
167pub struct AnthropicOutput {
168    pub id: String,
169    pub r#type: String,
170    pub role: String,
171    pub content: LLMMessageContent,
172    pub model: String,
173    pub stop_reason: String,
174    pub usage: AnthropicOutputUsage,
175}
176
177#[derive(Serialize, Deserialize, Debug)]
178pub struct AnthropicErrorOutput {
179    pub r#type: String,
180    pub error: AnthropicError,
181}
182
183#[derive(Serialize, Deserialize, Debug)]
184pub struct AnthropicError {
185    pub message: String,
186    pub r#type: String,
187}
188
189impl From<AnthropicOutput> for LLMCompletionResponse {
190    fn from(val: AnthropicOutput) -> Self {
191        let choices = vec![LLMChoice {
192            finish_reason: Some(val.stop_reason.clone()),
193            index: 0,
194            message: LLMMessage {
195                role: val.role.clone(),
196                content: val.content,
197            },
198        }];
199
200        LLMCompletionResponse {
201            id: val.id,
202            model: val.model,
203            object: val.r#type,
204            choices,
205            created: chrono::Utc::now().timestamp_millis() as u64,
206            usage: Some(val.usage.into()),
207        }
208    }
209}
210
211#[derive(Serialize, Deserialize, Debug)]
212pub struct AnthropicStreamEvent {
213    #[serde(rename = "type")]
214    pub event: String,
215    #[serde(flatten)]
216    pub data: AnthropicStreamEventData,
217}
218
219impl From<AnthropicOutputUsage> for LLMTokenUsage {
220    fn from(usage: AnthropicOutputUsage) -> Self {
221        let input_tokens = usage.input_tokens
222            + usage.cache_creation_input_tokens.unwrap_or(0)
223            + usage.cache_read_input_tokens.unwrap_or(0);
224        let output_tokens = usage.output_tokens;
225        Self {
226            completion_tokens: output_tokens,
227            prompt_tokens: input_tokens,
228            total_tokens: input_tokens + output_tokens,
229            prompt_tokens_details: Some(PromptTokensDetails {
230                input_tokens: Some(input_tokens),
231                output_tokens: Some(output_tokens),
232                cache_read_input_tokens: usage.cache_read_input_tokens,
233                cache_write_input_tokens: usage.cache_creation_input_tokens,
234            }),
235        }
236    }
237}
238
239#[derive(Serialize, Deserialize, Debug)]
240pub struct AnthropicStreamOutput {
241    pub id: String,
242    pub r#type: String,
243    pub role: String,
244    pub content: LLMMessageContent,
245    pub model: String,
246    pub stop_reason: Option<String>,
247    pub usage: AnthropicOutputUsage,
248}
249
250#[derive(Serialize, Deserialize, Debug)]
251#[serde(rename_all = "snake_case", tag = "type")]
252pub enum AnthropicStreamEventData {
253    MessageStart {
254        message: AnthropicStreamOutput,
255    },
256    ContentBlockStart {
257        index: usize,
258        content_block: ContentBlock,
259    },
260    ContentBlockDelta {
261        index: usize,
262        delta: ContentDelta,
263    },
264    ContentBlockStop {
265        index: usize,
266    },
267    MessageDelta {
268        delta: MessageDelta,
269        usage: Option<AnthropicOutputUsage>,
270    },
271    MessageStop {},
272    Ping {},
273}
274
275#[derive(Serialize, Deserialize, Debug)]
276#[serde(tag = "type")]
277pub enum ContentBlock {
278    #[serde(rename = "text")]
279    Text { text: String },
280    #[serde(rename = "thinking")]
281    Thinking { thinking: String },
282    #[serde(rename = "tool_use")]
283    ToolUse {
284        id: String,
285        name: String,
286        input: serde_json::Value,
287    },
288}
289
290#[derive(Serialize, Deserialize, Debug)]
291#[serde(tag = "type")]
292pub enum ContentDelta {
293    #[serde(rename = "text_delta")]
294    TextDelta { text: String },
295    #[serde(rename = "thinking_delta")]
296    ThinkingDelta { thinking: String },
297    #[serde(rename = "input_json_delta")]
298    InputJsonDelta { partial_json: String },
299}
300
301#[derive(Serialize, Deserialize, Debug)]
302pub struct MessageDelta {
303    pub stop_reason: Option<String>,
304    pub stop_sequence: Option<String>,
305}
306
307#[derive(Serialize, Deserialize, Clone, Debug, Default, PartialEq)]
308pub struct AnthropicConfig {
309    pub api_endpoint: Option<String>,
310    pub api_key: Option<String>,
311}