Skip to main content

uni_xervo/provider/
anthropic.rs

1use crate::api::{ModelAliasSpec, ModelTask};
2use crate::error::{Result, RuntimeError};
3use crate::provider::remote_common::{RemoteProviderBase, check_http_status, resolve_api_key};
4use crate::traits::{
5    GenerationOptions, GenerationResult, GeneratorModel, LoadedModelHandle, Message, MessageRole,
6    ModelProvider, ProviderCapabilities, ProviderHealth, TokenUsage,
7};
8use async_trait::async_trait;
9use reqwest::Client;
10use serde_json::json;
11use std::sync::Arc;
12
13/// Remote provider that calls the [Anthropic Messages API](https://docs.anthropic.com/en/api/messages)
14/// for text generation. Does not support embedding or reranking.
15///
16/// Requires the `ANTHROPIC_API_KEY` environment variable (or a custom env var
17/// name via the `api_key_env` option).
18pub struct RemoteAnthropicProvider {
19    base: RemoteProviderBase,
20}
21
22impl Default for RemoteAnthropicProvider {
23    fn default() -> Self {
24        Self {
25            base: RemoteProviderBase::new(),
26        }
27    }
28}
29
30impl RemoteAnthropicProvider {
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    #[cfg(test)]
36    fn insert_test_breaker(&self, key: crate::api::ModelRuntimeKey, age: std::time::Duration) {
37        self.base.insert_test_breaker(key, age);
38    }
39
40    #[cfg(test)]
41    fn breaker_count(&self) -> usize {
42        self.base.breaker_count()
43    }
44
45    #[cfg(test)]
46    fn force_cleanup_now_for_test(&self) {
47        self.base.force_cleanup_now_for_test();
48    }
49}
50
51#[async_trait]
52impl ModelProvider for RemoteAnthropicProvider {
53    fn provider_id(&self) -> &'static str {
54        "remote/anthropic"
55    }
56
57    fn capabilities(&self) -> ProviderCapabilities {
58        ProviderCapabilities {
59            supported_tasks: vec![ModelTask::Generate],
60        }
61    }
62
63    async fn load(&self, spec: &ModelAliasSpec) -> Result<LoadedModelHandle> {
64        let cb = self.base.circuit_breaker_for(spec);
65        let api_key = resolve_api_key(&spec.options, "api_key_env", "ANTHROPIC_API_KEY")?;
66
67        let anthropic_version = spec
68            .options
69            .get("anthropic_version")
70            .and_then(|v| v.as_str())
71            .unwrap_or("2023-06-01")
72            .to_string();
73
74        match spec.task {
75            ModelTask::Generate => {
76                let model = AnthropicGeneratorModel {
77                    client: self.base.client.clone(),
78                    cb,
79                    model_id: spec.model_id.clone(),
80                    api_key,
81                    anthropic_version,
82                };
83                let handle: Arc<dyn GeneratorModel> = Arc::new(model);
84                Ok(Arc::new(handle) as LoadedModelHandle)
85            }
86            ModelTask::Raw => Err(RuntimeError::CapabilityMismatch(
87                "Anthropic provider does not support task Raw".to_string(),
88            )),
89            _ => Err(RuntimeError::CapabilityMismatch(format!(
90                "Anthropic provider does not support task {:?}",
91                spec.task
92            ))),
93        }
94    }
95
96    async fn health(&self) -> ProviderHealth {
97        ProviderHealth::Healthy
98    }
99}
100
101struct AnthropicGeneratorModel {
102    client: Client,
103    cb: crate::reliability::CircuitBreakerWrapper,
104    model_id: String,
105    api_key: String,
106    anthropic_version: String,
107}
108
109impl crate::traits::ModelInfo for AnthropicGeneratorModel {
110    fn model_id(&self) -> &str {
111        &self.model_id
112    }
113}
114
115fn build_anthropic_payload(
116    model_id: &str,
117    messages: &[serde_json::Value],
118    options: &GenerationOptions,
119    system: Option<&str>,
120) -> serde_json::Value {
121    let max_tokens = options.max_tokens.unwrap_or(1024);
122
123    let mut body = json!({
124        "model": model_id,
125        "max_tokens": max_tokens,
126        "messages": messages,
127    });
128
129    if let Some(system_text) = system {
130        body["system"] = json!(system_text);
131    }
132    if let Some(temperature) = options.temperature {
133        body["temperature"] = json!(temperature);
134    }
135    if let Some(top_p) = options.top_p {
136        body["top_p"] = json!(top_p);
137    }
138
139    body
140}
141
142#[async_trait]
143impl GeneratorModel for AnthropicGeneratorModel {
144    async fn generate(
145        &self,
146        messages: &[Message],
147        options: GenerationOptions,
148    ) -> Result<GenerationResult> {
149        // Extract system messages into a single combined string
150        let system_parts: Vec<String> = messages
151            .iter()
152            .filter(|m| m.role == MessageRole::System)
153            .map(|m| m.text())
154            .collect();
155        let system_text = if system_parts.is_empty() {
156            None
157        } else {
158            Some(system_parts.join("\n"))
159        };
160
161        let messages: Vec<serde_json::Value> = messages
162            .iter()
163            .filter(|msg| msg.role != MessageRole::System)
164            .map(|msg| {
165                let role = match msg.role {
166                    MessageRole::User => "user",
167                    MessageRole::Assistant => "assistant",
168                    MessageRole::System => unreachable!("system messages filtered above"),
169                };
170                json!({ "role": role, "content": msg.text() })
171            })
172            .collect();
173
174        self.cb
175            .call(move || async move {
176                let body = build_anthropic_payload(
177                    &self.model_id,
178                    &messages,
179                    &options,
180                    system_text.as_deref(),
181                );
182
183                let response = self
184                    .client
185                    .post("https://api.anthropic.com/v1/messages")
186                    .header("x-api-key", &self.api_key)
187                    .header("anthropic-version", &self.anthropic_version)
188                    .header("content-type", "application/json")
189                    .json(&body)
190                    .send()
191                    .await
192                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
193
194                let body: serde_json::Value = check_http_status("Anthropic", response)?
195                    .json()
196                    .await
197                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
198
199                let text = body
200                    .get("content")
201                    .and_then(|c| c.as_array())
202                    .and_then(|arr| arr.first())
203                    .and_then(|item| item.get("text"))
204                    .and_then(|t| t.as_str())
205                    .unwrap_or("")
206                    .to_string();
207
208                let usage = body.get("usage").map(|u| TokenUsage {
209                    prompt_tokens: u["input_tokens"].as_u64().unwrap_or(0) as usize,
210                    completion_tokens: u["output_tokens"].as_u64().unwrap_or(0) as usize,
211                    total_tokens: (u["input_tokens"].as_u64().unwrap_or(0)
212                        + u["output_tokens"].as_u64().unwrap_or(0))
213                        as usize,
214                });
215
216                Ok(GenerationResult {
217                    text,
218                    usage,
219                    images: vec![],
220                    audio: None,
221                })
222            })
223            .await
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230    use crate::api::ModelRuntimeKey;
231    use crate::provider::remote_common::RemoteProviderBase;
232    use crate::traits::ModelProvider;
233    use std::time::Duration;
234
235    static ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
236
237    fn spec(alias: &str, task: ModelTask, model_id: &str) -> ModelAliasSpec {
238        ModelAliasSpec {
239            alias: alias.to_string(),
240            task,
241            provider_id: "remote/anthropic".to_string(),
242            model_id: model_id.to_string(),
243            revision: None,
244            warmup: crate::api::WarmupPolicy::Lazy,
245            required: false,
246            timeout: None,
247            load_timeout: None,
248            retry: None,
249            options: serde_json::Value::Null,
250        }
251    }
252
253    #[tokio::test]
254    async fn breaker_reused_for_same_runtime_key() {
255        let _lock = ENV_LOCK.lock().await;
256        unsafe { std::env::set_var("ANTHROPIC_API_KEY", "test-key") };
257
258        let provider = RemoteAnthropicProvider::new();
259        let s1 = spec("gen/a", ModelTask::Generate, "claude-sonnet-4-5-20250929");
260        let s2 = spec("gen/b", ModelTask::Generate, "claude-sonnet-4-5-20250929");
261
262        let _ = provider.load(&s1).await.unwrap();
263        let _ = provider.load(&s2).await.unwrap();
264
265        assert_eq!(provider.breaker_count(), 1);
266
267        unsafe { std::env::remove_var("ANTHROPIC_API_KEY") };
268    }
269
270    #[tokio::test]
271    async fn breaker_cleanup_evicts_stale_entries() {
272        let _lock = ENV_LOCK.lock().await;
273        unsafe { std::env::set_var("ANTHROPIC_API_KEY", "test-key") };
274
275        let provider = RemoteAnthropicProvider::new();
276        let stale = spec(
277            "gen/stale",
278            ModelTask::Generate,
279            "claude-sonnet-4-5-20250929",
280        );
281        let fresh = spec(
282            "gen/fresh",
283            ModelTask::Generate,
284            "claude-haiku-3-5-20241022",
285        );
286        provider.insert_test_breaker(
287            ModelRuntimeKey::new(&stale),
288            RemoteProviderBase::BREAKER_TTL + Duration::from_secs(5),
289        );
290        provider.insert_test_breaker(ModelRuntimeKey::new(&fresh), Duration::from_secs(1));
291        assert_eq!(provider.breaker_count(), 2);
292
293        provider.force_cleanup_now_for_test();
294        let _ = provider.load(&fresh).await.unwrap();
295
296        assert_eq!(provider.breaker_count(), 1);
297
298        unsafe { std::env::remove_var("ANTHROPIC_API_KEY") };
299    }
300
301    #[tokio::test]
302    async fn embed_capability_mismatch() {
303        let _lock = ENV_LOCK.lock().await;
304        unsafe { std::env::set_var("ANTHROPIC_API_KEY", "test-key") };
305
306        let provider = RemoteAnthropicProvider::new();
307        let s = spec("embed/a", ModelTask::Embed, "claude-sonnet-4-5-20250929");
308        let result = provider.load(&s).await;
309        assert!(result.is_err());
310        assert!(
311            result
312                .unwrap_err()
313                .to_string()
314                .contains("does not support task")
315        );
316
317        unsafe { std::env::remove_var("ANTHROPIC_API_KEY") };
318    }
319
320    #[tokio::test]
321    async fn rerank_capability_mismatch() {
322        let _lock = ENV_LOCK.lock().await;
323        unsafe { std::env::set_var("ANTHROPIC_API_KEY", "test-key") };
324
325        let provider = RemoteAnthropicProvider::new();
326        let s = spec("rerank/a", ModelTask::Rerank, "claude-sonnet-4-5-20250929");
327        let result = provider.load(&s).await;
328        assert!(result.is_err());
329        assert!(
330            result
331                .unwrap_err()
332                .to_string()
333                .contains("does not support task")
334        );
335
336        unsafe { std::env::remove_var("ANTHROPIC_API_KEY") };
337    }
338
339    #[test]
340    fn payload_defaults_max_tokens_to_1024() {
341        let messages = vec![json!({"role": "user", "content": "hello"})];
342        let payload = build_anthropic_payload(
343            "claude-sonnet-4-5-20250929",
344            &messages,
345            &GenerationOptions::default(),
346            None,
347        );
348        assert_eq!(payload["max_tokens"], 1024);
349    }
350
351    #[test]
352    fn payload_uses_explicit_max_tokens() {
353        let messages = vec![json!({"role": "user", "content": "hello"})];
354        let payload = build_anthropic_payload(
355            "claude-sonnet-4-5-20250929",
356            &messages,
357            &GenerationOptions {
358                max_tokens: Some(512),
359                ..Default::default()
360            },
361            None,
362        );
363        assert_eq!(payload["max_tokens"], 512);
364    }
365
366    #[test]
367    fn payload_includes_system_field() {
368        let messages = vec![json!({"role": "user", "content": "hello"})];
369        let payload = build_anthropic_payload(
370            "claude-sonnet-4-5-20250929",
371            &messages,
372            &GenerationOptions::default(),
373            Some("you are helpful"),
374        );
375        assert_eq!(payload["system"], "you are helpful");
376    }
377
378    #[test]
379    fn payload_omits_system_field_when_none() {
380        let messages = vec![json!({"role": "user", "content": "hello"})];
381        let payload = build_anthropic_payload(
382            "claude-sonnet-4-5-20250929",
383            &messages,
384            &GenerationOptions::default(),
385            None,
386        );
387        assert!(payload.get("system").is_none());
388    }
389}