1use std::{future::Future, time::Duration};
14
15use axum::http::HeaderMap;
16use serde_json::Value;
17
18pub mod anthropic;
19pub mod http;
20pub mod openai;
21
22#[cfg(test)]
23pub(crate) mod testing {
24 use axum::http::{HeaderMap, HeaderName};
25 use serde_json::Value;
26
27 use super::HttpResponse;
28
29 pub fn recorded(text: &str) -> HttpResponse {
32 let v: Value = serde_json::from_str(text).unwrap();
33 if v.get("status").is_none() {
34 return HttpResponse {
35 status: 200,
36 headers: HeaderMap::new(),
37 body: text.to_string(),
38 };
39 }
40 let mut headers = HeaderMap::new();
41 for (k, val) in v["headers"].as_object().unwrap() {
42 headers.insert(
43 HeaderName::from_bytes(k.as_bytes()).unwrap(),
44 val.as_str().unwrap().parse().unwrap(),
45 );
46 }
47 HttpResponse {
48 status: v["status"].as_u64().unwrap() as u16,
49 headers,
50 body: v["body"].to_string(),
51 }
52 }
53}
54
55#[derive(Clone, Debug, PartialEq)]
57pub struct ToolSpec {
58 pub name: String,
59 pub description: String,
60 pub schema: Value,
61}
62
63#[derive(Clone, Debug, PartialEq)]
65pub struct ToolCall {
66 pub id: String,
67 pub name: String,
68 pub input: Result<Value, String>,
71}
72
73#[derive(Clone, Debug, PartialEq)]
75pub struct ToolResult {
76 pub call_id: String,
77 pub name: String,
78 pub text: String,
79 pub is_error: bool,
80}
81
82#[derive(Clone, Debug, PartialEq)]
84pub enum Turn {
85 User(String),
86 Assistant {
89 text: String,
90 calls: Vec<ToolCall>,
91 raw: Value,
92 },
93 ToolResults(Vec<ToolResult>),
95}
96
97pub struct Request<'a> {
99 pub model: &'a str,
100 pub system: &'a str,
101 pub tools: &'a [ToolSpec],
102 pub turns: &'a [Turn],
103 pub max_tokens: u32,
104}
105
106#[derive(Clone, Copy, Debug, PartialEq, Eq)]
108pub enum Stop {
109 Done,
111 ToolCalls,
113 Truncated,
115 Refused,
117 Paused,
119}
120
121#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
124pub struct Usage {
125 pub input: u64,
126 pub cached_input: u64,
127 pub output: u64,
128 pub reported: bool,
131}
132
133#[derive(Clone, Debug, PartialEq)]
135pub struct Reply {
136 pub text: String,
137 pub calls: Vec<ToolCall>,
138 pub raw: Value,
140 pub stop: Stop,
141 pub usage: Usage,
142}
143
144#[derive(Clone, Debug, PartialEq)]
147pub struct HttpCall {
148 pub path: &'static str,
149 pub headers: Vec<(&'static str, String)>,
150 pub body: Value,
151}
152
153#[derive(Clone, Debug)]
155pub struct HttpResponse {
156 pub status: u16,
157 pub headers: HeaderMap,
158 pub body: String,
159}
160
161#[derive(Clone, Copy, Debug, PartialEq, Eq)]
163pub enum AuthStyle {
164 XApiKey,
166 Bearer,
168}
169
170#[derive(Clone, Debug, PartialEq, thiserror::Error)]
172pub enum LlmError {
173 #[error("the AI service rejected the API key")]
174 Auth,
175 #[error("rate limited by the AI service")]
176 RateLimited { retry_after: Option<Duration> },
177 #[error("the AI service is overloaded or failed ({0})")]
178 Overloaded(u16),
179 #[error("the AI service rejected the request ({status}): {message}")]
180 Rejected { status: u16, message: String },
181 #[error("the AI service refused: its usage or spending limit is reached")]
182 Quota,
183 #[error("the input is longer than the model accepts")]
184 ContextLength,
185 #[error("network error: {0}")]
186 Network(String),
187 #[error("time limit reached")]
188 Timeout,
189 #[error("unexpected API response: {0}")]
190 Malformed(String),
191}
192
193impl LlmError {
194 pub fn retryable(&self) -> bool {
196 matches!(
197 self,
198 LlmError::RateLimited { .. } | LlmError::Overloaded(_) | LlmError::Network(_)
199 )
200 }
201}
202
203pub trait Provider: Send + Sync {
205 fn name(&self) -> &'static str;
207 fn auth(&self) -> AuthStyle;
209 fn encode(&self, request: &Request) -> HttpCall;
211 fn decode(&self, body: &Value) -> Result<Reply, LlmError>;
213 fn error(&self, status: u16, headers: &HeaderMap, body: &str) -> LlmError;
215}
216
217pub trait Transport: Sync {
219 fn send(&self, call: &HttpCall) -> impl Future<Output = Result<HttpResponse, LlmError>> + Send;
220}
221
222pub async fn complete<T: Transport>(
224 provider: &dyn Provider,
225 transport: &T,
226 request: &Request<'_>,
227) -> Result<Reply, LlmError> {
228 let response = transport.send(&provider.encode(request)).await?;
229 if !(200..300).contains(&response.status) {
230 return Err(provider.error(response.status, &response.headers, &response.body));
231 }
232 let body: Value =
233 serde_json::from_str(&response.body).map_err(|e| LlmError::Malformed(e.to_string()))?;
234 provider.decode(&body)
235}
236
237pub fn retry_after(headers: &HeaderMap) -> Option<Duration> {
239 headers
240 .get("retry-after")?
241 .to_str()
242 .ok()?
243 .trim()
244 .parse::<u64>()
245 .ok()
246 .map(Duration::from_secs)
247}
248
249pub fn error_message(body: &str) -> String {
251 let v = serde_json::from_str::<Value>(body).unwrap_or_default();
252 let message = v["error"]["message"]
254 .as_str()
255 .or_else(|| v["error"].as_str())
256 .or_else(|| v["message"].as_str())
257 .unwrap_or_default();
258 cut(message, 300)
259}
260
261pub fn status_error(status: u16, headers: &HeaderMap, body: &str) -> LlmError {
267 let lower = body.to_ascii_lowercase();
268 match status {
269 401 | 403 => LlmError::Auth,
270 402 => LlmError::Quota,
271 429 if lower.contains("insufficient_quota") || lower.contains("billing") => LlmError::Quota,
272 429 => LlmError::RateLimited {
273 retry_after: retry_after(headers),
274 },
275 500..=599 => LlmError::Overloaded(status),
276 400 | 413 | 422
277 if lower.contains("context_length_exceeded")
278 || lower.contains("prompt is too long")
279 || lower.contains("maximum context length")
280 || lower.contains("context length")
281 || lower.contains("context window") =>
282 {
283 LlmError::ContextLength
284 }
285 _ => LlmError::Rejected {
286 status,
287 message: error_message(body),
288 },
289 }
290}
291
292pub fn cut(s: &str, max: usize) -> String {
294 if s.len() <= max {
295 return s.to_string();
296 }
297 let mut end = max;
298 while !s.is_char_boundary(end) {
299 end -= 1;
300 }
301 format!("{}… [cut – narrow the request]", &s[..end])
302}