Skip to main content

teamcraft/llm/
mod.rs

1//! One provider-neutral API layer for the assistant (design:
2//! `docs/src/research/agent-interface.md`).
3//!
4//! This module knows nothing about TeamCraft (no `Actor`, no `mcp`): it
5//! holds the neutral conversation types, the pure [`Provider`] trait (encode a
6//! request, decode a reply, map an error – no I/O, so recorded provider
7//! bodies test it directly) and the [`Transport`] trait with its one HTTP
8//! implementation in [`http`]. The wire formats live in one file each
9//! ([`anthropic`]).
10//!
11//! The tool loop in `crate::assistant` only sees these types.
12
13use 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    /// A recorded API answer: a bare body (status 200) or
30    /// `{status, headers, body}` for an error.
31    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/// A tool the model may call. `schema` is the JSON Schema of the input.
56#[derive(Clone, Debug, PartialEq)]
57pub struct ToolSpec {
58    pub name: String,
59    pub description: String,
60    pub schema: Value,
61}
62
63/// One tool call of the model.
64#[derive(Clone, Debug, PartialEq)]
65pub struct ToolCall {
66    pub id: String,
67    pub name: String,
68    /// `Err` holds arguments that could not be parsed (the loop answers the
69    /// call with a tool error and the run goes on).
70    pub input: Result<Value, String>,
71}
72
73/// The result of one tool call, handed back to the model.
74#[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/// One entry of the conversation.
83#[derive(Clone, Debug, PartialEq)]
84pub enum Turn {
85    User(String),
86    /// What the model said. `raw` is the provider's own form and is sent back
87    /// unchanged (Anthropic thinking blocks); the loop never edits it.
88    Assistant {
89        text: String,
90        calls: Vec<ToolCall>,
91        raw: Value,
92    },
93    /// Anthropic: one user message; OpenAI: one tool message per result.
94    ToolResults(Vec<ToolResult>),
95}
96
97/// One model call.
98pub 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/// Why the model stopped.
107#[derive(Clone, Copy, Debug, PartialEq, Eq)]
108pub enum Stop {
109    /// The answer is complete.
110    Done,
111    /// The model asks for tool calls.
112    ToolCalls,
113    /// Cut at the output limit.
114    Truncated,
115    /// The model declined.
116    Refused,
117    /// Paused by the provider; send the conversation again unchanged.
118    Paused,
119}
120
121/// Token counts of one call. `input` means all prompt tokens (cached
122/// ones included).
123#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
124pub struct Usage {
125    pub input: u64,
126    pub cached_input: u64,
127    pub output: u64,
128    /// The server reported token counts. Without them (some local servers)
129    /// the token limit of a run cannot count the call.
130    pub reported: bool,
131}
132
133/// A decoded model reply.
134#[derive(Clone, Debug, PartialEq)]
135pub struct Reply {
136    pub text: String,
137    pub calls: Vec<ToolCall>,
138    /// The provider's own form of the turn (see [`Turn::Assistant`]).
139    pub raw: Value,
140    pub stop: Stop,
141    pub usage: Usage,
142}
143
144/// A request ready to send: the path below the base URL, headers (without
145/// the API key) and the JSON body.
146#[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/// An HTTP answer.
154#[derive(Clone, Debug)]
155pub struct HttpResponse {
156    pub status: u16,
157    pub headers: HeaderMap,
158    pub body: String,
159}
160
161/// How the API key is sent.
162#[derive(Clone, Copy, Debug, PartialEq, Eq)]
163pub enum AuthStyle {
164    /// `x-api-key: <key>`
165    XApiKey,
166    /// `Authorization: Bearer <key>`
167    Bearer,
168}
169
170/// Why a model call failed.
171#[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    /// Whether a single retry may help.
195    pub fn retryable(&self) -> bool {
196        matches!(
197            self,
198            LlmError::RateLimited { .. } | LlmError::Overloaded(_) | LlmError::Network(_)
199        )
200    }
201}
202
203/// One wire format. Pure: no network, no clock, no database.
204pub trait Provider: Send + Sync {
205    /// Short name for logs and the audit record (`anthropic`).
206    fn name(&self) -> &'static str;
207    /// How the API key goes into the request.
208    fn auth(&self) -> AuthStyle;
209    /// The HTTP request of one model call.
210    fn encode(&self, request: &Request) -> HttpCall;
211    /// The reply from a successful response body.
212    fn decode(&self, body: &Value) -> Result<Reply, LlmError>;
213    /// The error for a response with a status outside 2xx.
214    fn error(&self, status: u16, headers: &HeaderMap, body: &str) -> LlmError;
215}
216
217/// Sends one HTTP call and returns the response as it is (any status).
218pub trait Transport: Sync {
219    fn send(&self, call: &HttpCall) -> impl Future<Output = Result<HttpResponse, LlmError>> + Send;
220}
221
222/// One model call, one attempt: encode, send, map errors, decode.
223pub 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
237/// `retry-after` in seconds (the HTTP-date form is not used by the APIs).
238pub 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
249/// The `error.message` of a JSON error body, cut short.
250pub fn error_message(body: &str) -> String {
251    let v = serde_json::from_str::<Value>(body).unwrap_or_default();
252    // `{"error": {"message": …}}`, or a plain string (Ollama, some proxies)
253    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
261/// The one error table of all providers (design page): 401/403 are an
262/// authentication error, 402 and a billing-quota 429 a quota error (no
263/// retry), other 429s are rate limits with their `retry-after`, 5xx (also
264/// 529) an overloaded service, a context-length rejection its own error,
265/// everything else a rejected request.
266pub 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
292/// Cuts `s` to at most `max` bytes at a character boundary.
293pub 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}