1use async_trait::async_trait;
5use reqwest::header::{HeaderMap, HeaderValue, CONTENT_TYPE};
6use serde_json::{json, Value};
7use tokenmiser_config::ProviderConfig;
8
9use crate::{ChatChoice, ChatMessage, ChatRequest, ChatResponse, Provider, ProviderError, Usage};
10
11const ANTHROPIC_VERSION: &str = "2023-06-01";
12
13pub struct AnthropicProvider {
14 cfg: ProviderConfig,
15 client: reqwest::Client,
16 api_key: Option<String>,
17}
18
19impl AnthropicProvider {
20 pub fn new(cfg: ProviderConfig) -> Self {
21 let api_key = cfg.api_key_env.as_ref().and_then(|k| std::env::var(k).ok());
22
23 let client = reqwest::Client::builder()
24 .pool_max_idle_per_host(64)
25 .build()
26 .expect("reqwest client construction");
27
28 Self {
29 cfg,
30 client,
31 api_key,
32 }
33 }
34
35 fn headers(&self) -> Result<HeaderMap, ProviderError> {
36 let mut h = HeaderMap::new();
37 h.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
38 h.insert(
39 "anthropic-version",
40 HeaderValue::from_static(ANTHROPIC_VERSION),
41 );
42 if let Some(k) = &self.api_key {
43 let v = HeaderValue::from_str(k)
44 .map_err(|e| ProviderError::Malformed(format!("invalid api key: {e}")))?;
45 h.insert("x-api-key", v);
46 }
47 Ok(h)
48 }
49
50 fn to_anthropic_body(req: &ChatRequest) -> Value {
52 let mut system: Option<String> = None;
53 let mut messages: Vec<Value> = Vec::with_capacity(req.messages.len());
54
55 for m in &req.messages {
56 match m.role.as_str() {
57 "system" => {
58 if let Some(s) = content_to_string(&m.content) {
60 system = Some(match system {
61 Some(prev) => format!("{prev}\n\n{s}"),
62 None => s,
63 });
64 }
65 }
66 role => {
67 messages.push(json!({
68 "role": role,
69 "content": m.content,
70 }));
71 }
72 }
73 }
74
75 let mut body = json!({
76 "model": req.model,
77 "messages": messages,
78 "max_tokens": req.max_tokens.unwrap_or(4096),
80 });
81
82 if let Some(s) = system {
83 body["system"] = Value::String(s);
84 }
85 if let Some(t) = req.temperature {
86 body["temperature"] = json!(t);
87 }
88 if let Some(p) = req.top_p {
89 body["top_p"] = json!(p);
90 }
91
92 body
93 }
94
95 fn from_anthropic_body(v: Value, requested_model: &str) -> Result<ChatResponse, ProviderError> {
97 let id = v
98 .get("id")
99 .and_then(Value::as_str)
100 .unwrap_or("anthropic")
101 .to_string();
102 let model = v
103 .get("model")
104 .and_then(Value::as_str)
105 .unwrap_or(requested_model)
106 .to_string();
107
108 let mut text = String::new();
110 if let Some(blocks) = v.get("content").and_then(Value::as_array) {
111 for b in blocks {
112 if b.get("type").and_then(Value::as_str) == Some("text") {
113 if let Some(t) = b.get("text").and_then(Value::as_str) {
114 text.push_str(t);
115 }
116 }
117 }
118 }
119
120 let usage = v.get("usage");
121 let prompt_tokens = usage
122 .and_then(|u| u.get("input_tokens"))
123 .and_then(Value::as_u64)
124 .unwrap_or(0);
125 let completion_tokens = usage
126 .and_then(|u| u.get("output_tokens"))
127 .and_then(Value::as_u64)
128 .unwrap_or(0);
129
130 let finish_reason = v
131 .get("stop_reason")
132 .and_then(Value::as_str)
133 .map(|s| match s {
134 "end_turn" => "stop".to_string(),
135 "max_tokens" => "length".to_string(),
136 other => other.to_string(),
137 });
138
139 Ok(ChatResponse {
140 id,
141 object: "chat.completion".into(),
142 created: chrono::Utc::now().timestamp() as u64,
143 model,
144 choices: vec![ChatChoice {
145 index: 0,
146 message: ChatMessage {
147 role: "assistant".into(),
148 content: Value::String(text),
149 extra: Default::default(),
150 },
151 finish_reason,
152 logprobs: None,
153 }],
154 usage: Usage {
155 prompt_tokens,
156 completion_tokens,
157 total_tokens: prompt_tokens + completion_tokens,
158 },
159 extra: Default::default(),
160 })
161 }
162}
163
164#[async_trait]
165impl Provider for AnthropicProvider {
166 fn name(&self) -> &str {
167 &self.cfg.name
168 }
169
170 fn config(&self) -> &ProviderConfig {
171 &self.cfg
172 }
173
174 async fn complete(&self, req: &ChatRequest) -> Result<ChatResponse, ProviderError> {
175 if self.api_key.is_none() {
176 return Err(ProviderError::MissingApiKey(
177 self.cfg.api_key_env.clone().unwrap_or_default(),
178 ));
179 }
180
181 let url = format!("{}/messages", self.cfg.base_url.trim_end_matches('/'));
182 let body = Self::to_anthropic_body(req);
183
184 let res = self
185 .client
186 .post(&url)
187 .headers(self.headers()?)
188 .json(&body)
189 .send()
190 .await?;
191
192 let status = res.status();
193 let text = res.text().await?;
194
195 if !status.is_success() {
196 return Err(ProviderError::Upstream {
197 status: status.as_u16(),
198 body: text,
199 });
200 }
201
202 let v: Value = serde_json::from_str(&text)?;
203 Self::from_anthropic_body(v, &req.model)
204 }
205}
206
207fn content_to_string(v: &Value) -> Option<String> {
208 match v {
209 Value::String(s) => Some(s.clone()),
210 Value::Array(arr) => {
211 let mut buf = String::new();
212 for item in arr {
213 if let Some(t) = item.get("text").and_then(Value::as_str) {
214 buf.push_str(t);
215 }
216 }
217 if buf.is_empty() {
218 None
219 } else {
220 Some(buf)
221 }
222 }
223 _ => None,
224 }
225}
226
227#[cfg(test)]
228mod tests {
229 use super::*;
230
231 #[test]
232 fn system_message_moves_to_top_level() {
233 let req = ChatRequest {
234 model: "claude-sonnet-4-6".into(),
235 messages: vec![
236 ChatMessage {
237 role: "system".into(),
238 content: Value::String("you are concise".into()),
239 extra: Default::default(),
240 },
241 ChatMessage {
242 role: "user".into(),
243 content: Value::String("hi".into()),
244 extra: Default::default(),
245 },
246 ],
247 temperature: Some(0.2),
248 max_tokens: Some(100),
249 top_p: None,
250 stream: None,
251 extra: Default::default(),
252 };
253 let body = AnthropicProvider::to_anthropic_body(&req);
254 assert_eq!(body["system"], "you are concise");
255 assert_eq!(body["messages"].as_array().unwrap().len(), 1);
256 assert_eq!(body["messages"][0]["role"], "user");
257 assert_eq!(body["max_tokens"], 100);
258 }
259
260 #[test]
261 fn parses_anthropic_response_to_openai_shape() {
262 let body = json!({
263 "id": "msg_01",
264 "model": "claude-sonnet-4-6",
265 "content": [{"type": "text", "text": "hello"}],
266 "stop_reason": "end_turn",
267 "usage": {"input_tokens": 5, "output_tokens": 2}
268 });
269 let r = AnthropicProvider::from_anthropic_body(body, "claude-sonnet-4-6").unwrap();
270 assert_eq!(r.choices[0].message.content, Value::String("hello".into()));
271 assert_eq!(r.choices[0].finish_reason.as_deref(), Some("stop"));
272 assert_eq!(r.usage.prompt_tokens, 5);
273 assert_eq!(r.usage.completion_tokens, 2);
274 assert_eq!(r.usage.total_tokens, 7);
275 }
276}