Skip to main content

teamcraft/llm/
openai.rs

1//! Chat Completions (`POST {base}/chat/completions`), the format OpenAI,
2//! Azure OpenAI (v1), Mistral, OpenRouter, Groq, Ollama, vLLM and LM Studio
3//! speak. Stateless like the Messages API, so the loop stays the same.
4//!
5//! Rules (design page):
6//!
7//! * `strict`, `tool_choice` and `parallel_tool_calls` are never sent.
8//! * Tool arguments are a JSON string that may be invalid: a call whose
9//!   arguments do not parse as an object gets a tool error, the run goes on.
10//! * A missing call id is made up (9 alphanumeric characters – Mistral
11//!   rejects longer ids); an id from the server is sent back unchanged.
12//! * Tool calls win over `finish_reason` (some servers say `stop`).
13//! * There is no error flag on a tool message: a failed result starts with
14//!   `Error: `.
15//! * Without `usage` in the reply, `Usage::reported` is false.
16
17use axum::http::HeaderMap;
18use serde_json::{json, Value};
19
20use super::{
21    status_error, AuthStyle, HttpCall, LlmError, Provider, Reply, Request, Stop, ToolCall, Turn,
22    Usage,
23};
24
25/// Which flavour of Chat Completions.
26#[derive(Clone, Copy, Debug, PartialEq, Eq)]
27pub enum Flavour {
28    /// OpenAI and Azure OpenAI v1: `max_completion_tokens` (`max_tokens`
29    /// does not work with reasoning models).
30    OpenAi,
31    /// Everything else: `max_tokens`, which every server knows.
32    Compatible,
33}
34
35/// The Chat Completions wire format.
36pub struct OpenAi(pub Flavour);
37
38/// A call id of 9 random alphanumeric characters.
39fn made_up_id() -> String {
40    uuid::Uuid::new_v4()
41        .simple()
42        .to_string()
43        .chars()
44        .take(9)
45        .collect()
46}
47
48/// The message of one turn (a tool-result turn is several messages).
49fn messages(turn: &Turn) -> Vec<Value> {
50    match turn {
51        Turn::User(text) => vec![json!({ "role": "user", "content": text })],
52        Turn::Assistant { raw, .. } => vec![raw.clone()],
53        Turn::ToolResults(results) => results
54            .iter()
55            .map(|r| {
56                json!({
57                    "role": "tool",
58                    "tool_call_id": r.call_id,
59                    "name": r.name,
60                    "content": if r.is_error { format!("Error: {}", r.text) } else { r.text.clone() },
61                })
62            })
63            .collect(),
64    }
65}
66
67/// The text of a message: a string, or an array of `{type: "text"}` parts.
68fn content_text(content: &Value) -> String {
69    match content {
70        Value::String(s) => s.trim().to_string(),
71        Value::Array(parts) => parts
72            .iter()
73            .filter_map(|p| p["text"].as_str())
74            .collect::<Vec<_>>()
75            .join("\n")
76            .trim()
77            .to_string(),
78        _ => String::new(),
79    }
80}
81
82impl Provider for OpenAi {
83    fn name(&self) -> &'static str {
84        match self.0 {
85            Flavour::OpenAi => "openai",
86            Flavour::Compatible => "openai-compatible",
87        }
88    }
89
90    fn auth(&self) -> AuthStyle {
91        AuthStyle::Bearer
92    }
93
94    fn encode(&self, request: &Request) -> HttpCall {
95        let mut all = vec![json!({ "role": "system", "content": request.system })];
96        all.extend(request.turns.iter().flat_map(messages));
97        let mut body = json!({
98            "model": request.model,
99            "messages": all,
100        });
101        let limit = match self.0 {
102            Flavour::OpenAi => "max_completion_tokens",
103            Flavour::Compatible => "max_tokens",
104        };
105        body[limit] = json!(request.max_tokens);
106        if !request.tools.is_empty() {
107            body["tools"] = request
108                .tools
109                .iter()
110                .map(|t| {
111                    json!({
112                        "type": "function",
113                        "function": {
114                            "name": t.name,
115                            "description": t.description,
116                            "parameters": t.schema,
117                        },
118                    })
119                })
120                .collect();
121        }
122        HttpCall {
123            path: "/chat/completions",
124            headers: Vec::new(),
125            body,
126        }
127    }
128
129    fn decode(&self, body: &Value) -> Result<Reply, LlmError> {
130        let malformed = |what: &str| LlmError::Malformed(what.to_string());
131        let choice = body["choices"]
132            .as_array()
133            .and_then(|c| c.first())
134            .ok_or_else(|| malformed("response without choices"))?;
135        let message = &choice["message"];
136        if !message.is_object() {
137            return Err(malformed("response without a message"));
138        }
139        let text = content_text(&message["content"]);
140        let refusal = message["refusal"]
141            .as_str()
142            .is_some_and(|r| !r.trim().is_empty());
143
144        let mut calls = Vec::new();
145        let mut raw_calls = Vec::new();
146        for call in message["tool_calls"].as_array().into_iter().flatten() {
147            let name = call["function"]["name"]
148                .as_str()
149                .filter(|n| !n.is_empty())
150                .ok_or_else(|| malformed("tool call without a function name"))?;
151            let id = call["id"]
152                .as_str()
153                .filter(|i| !i.is_empty())
154                .map(str::to_string)
155                .unwrap_or_else(made_up_id);
156            // a JSON string; some servers send the object itself
157            let arguments = &call["function"]["arguments"];
158            let (input, arguments_text) = match arguments {
159                Value::String(s) if s.trim().is_empty() => (Ok(json!({})), "{}".to_string()),
160                Value::String(s) => (
161                    match serde_json::from_str::<Value>(s) {
162                        Ok(v) if v.is_object() => Ok(v),
163                        Ok(_) => Err("the arguments are not a JSON object".to_string()),
164                        Err(e) => Err(format!("the arguments are not valid JSON: {e}")),
165                    },
166                    s.clone(),
167                ),
168                Value::Object(_) => (Ok(arguments.clone()), arguments.to_string()),
169                Value::Null => (Ok(json!({})), "{}".to_string()),
170                _ => (
171                    Err("the arguments are not a JSON object".to_string()),
172                    "{}".to_string(),
173                ),
174            };
175            raw_calls.push(json!({
176                "id": id,
177                "type": "function",
178                "function": {
179                    "name": name,
180                    // an unparsable string goes back as it came; the model
181                    // sees its tool error next to it
182                    "arguments": arguments_text,
183                },
184            }));
185            calls.push(ToolCall {
186                id,
187                name: name.to_string(),
188                input,
189            });
190        }
191
192        let finish = choice["finish_reason"].as_str().unwrap_or_default();
193        let stop = if !calls.is_empty() {
194            // tool calls win over `finish_reason`
195            Stop::ToolCalls
196        } else if refusal || finish == "content_filter" {
197            Stop::Refused
198        } else if finish == "length" {
199            Stop::Truncated
200        } else if finish == "tool_calls" || finish == "function_call" {
201            return Err(malformed("finish_reason tool_calls without tool calls"));
202        } else if finish.is_empty() && text.is_empty() {
203            return Err(malformed("response without finish_reason or text"));
204        } else {
205            Stop::Done
206        };
207
208        let mut raw = json!({ "role": "assistant", "content": text });
209        if !raw_calls.is_empty() {
210            raw["tool_calls"] = Value::Array(raw_calls);
211        }
212        let usage = &body["usage"];
213        let count = |v: &Value| v.as_u64();
214        let prompt = count(&usage["prompt_tokens"]);
215        let completion = count(&usage["completion_tokens"]);
216        Ok(Reply {
217            text,
218            calls,
219            raw,
220            stop,
221            usage: Usage {
222                input: prompt.unwrap_or(0),
223                cached_input: count(&usage["prompt_tokens_details"]["cached_tokens"]).unwrap_or(0),
224                output: completion.unwrap_or(0),
225                reported: prompt.is_some() || completion.is_some(),
226            },
227        })
228    }
229
230    fn error(&self, status: u16, headers: &HeaderMap, body: &str) -> LlmError {
231        status_error(status, headers, body)
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use std::time::Duration;
238
239    use super::*;
240    use crate::llm::{complete, testing::recorded, HttpResponse, ToolResult, ToolSpec, Transport};
241
242    macro_rules! fixture {
243        ($name:literal) => {
244            recorded(include_str!(concat!(
245                "../../tests/fixtures/llm/openai/",
246                $name,
247                ".json"
248            )))
249        };
250    }
251
252    struct Fixed(HttpResponse);
253    impl Transport for Fixed {
254        async fn send(&self, _: &HttpCall) -> Result<HttpResponse, LlmError> {
255            Ok(self.0.clone())
256        }
257    }
258
259    async fn reply(resp: HttpResponse) -> Result<Reply, LlmError> {
260        let req = Request {
261            model: "m",
262            system: "s",
263            tools: &[],
264            turns: &[Turn::User("q".into())],
265            max_tokens: 10,
266        };
267        complete(&OpenAi(Flavour::Compatible), &Fixed(resp), &req).await
268    }
269
270    fn spec() -> ToolSpec {
271        ToolSpec {
272            name: "get_item".into(),
273            description: "d".into(),
274            schema: json!({ "type": "object", "properties": { "item": { "type": "string" } } }),
275        }
276    }
277
278    #[test]
279    fn requests_have_the_chat_completions_shape() {
280        let tools = [spec()];
281        let turns = [
282            Turn::User("hi".into()),
283            Turn::Assistant {
284                text: String::new(),
285                calls: vec![],
286                raw: json!({ "role": "assistant", "content": "", "tool_calls": [
287                    { "id": "c1", "type": "function", "function": { "name": "get_item", "arguments": "{}" } },
288                    { "id": "c2", "type": "function", "function": { "name": "get_item", "arguments": "{" } } ] }),
289            },
290            Turn::ToolResults(vec![
291                ToolResult {
292                    call_id: "c1".into(),
293                    name: "get_item".into(),
294                    text: "ok".into(),
295                    is_error: false,
296                },
297                ToolResult {
298                    call_id: "c2".into(),
299                    name: "get_item".into(),
300                    text: "bad input".into(),
301                    is_error: true,
302                },
303            ]),
304        ];
305        let request = Request {
306            model: "gpt-x",
307            system: "sys",
308            tools: &tools,
309            turns: &turns,
310            max_tokens: 777,
311        };
312        let call = OpenAi(Flavour::OpenAi).encode(&request);
313        assert_eq!(call.path, "/chat/completions");
314        assert!(call.headers.is_empty(), "the key is added by the transport");
315        let b = &call.body;
316        assert_eq!(b["model"], "gpt-x");
317        assert_eq!(b["max_completion_tokens"], 777);
318        assert!(b.get("max_tokens").is_none());
319        for never in [
320            "tool_choice",
321            "parallel_tool_calls",
322            "stream",
323            "temperature",
324        ] {
325            assert!(b.get(never).is_none(), "{never} must not be sent");
326        }
327        assert!(b["tools"][0]["function"].get("strict").is_none());
328        assert_eq!(b["tools"][0]["type"], "function");
329        assert_eq!(b["tools"][0]["function"]["name"], "get_item");
330        assert_eq!(b["tools"][0]["function"]["parameters"]["type"], "object");
331        let m = b["messages"].as_array().unwrap();
332        assert_eq!(m[0], json!({ "role": "system", "content": "sys" }));
333        assert_eq!(m[1]["role"], "user");
334        assert_eq!(m[2]["tool_calls"][0]["id"], "c1");
335        // one tool message per result, in order, with the error marker
336        assert_eq!(m.len(), 5);
337        assert_eq!(m[3]["role"], "tool");
338        assert_eq!(m[3]["tool_call_id"], "c1");
339        assert_eq!(m[3]["name"], "get_item");
340        assert_eq!(m[3]["content"], "ok");
341        assert_eq!(m[4]["tool_call_id"], "c2");
342        assert_eq!(m[4]["content"], "Error: bad input");
343
344        let call = OpenAi(Flavour::Compatible).encode(&request);
345        assert_eq!(call.body["max_tokens"], 777);
346        assert!(call.body.get("max_completion_tokens").is_none());
347    }
348
349    #[tokio::test]
350    async fn openai_replies_are_decoded() {
351        let r = reply(fixture!("openai_answer")).await.unwrap();
352        assert_eq!(r.stop, Stop::Done);
353        assert_eq!(r.text, "PRT-000001 is released.");
354        assert_eq!(
355            (
356                r.usage.input,
357                r.usage.cached_input,
358                r.usage.output,
359                r.usage.reported
360            ),
361            (120, 100, 12, true),
362            "prompt_tokens already include the cached ones"
363        );
364
365        let r = reply(fixture!("openai_tool_calls")).await.unwrap();
366        assert_eq!(r.stop, Stop::ToolCalls);
367        assert_eq!(r.calls.len(), 2);
368        assert_eq!(r.calls[0].id, "call_abc123");
369        assert_eq!(r.calls[0].name, "get_item");
370        assert_eq!(r.calls[0].input, Ok(json!({ "item": "PRT-000001" })));
371        assert_eq!(r.calls[1].id, "call_def456");
372        // what goes back is the assistant message with its tool calls
373        assert_eq!(r.raw["role"], "assistant");
374        assert_eq!(r.raw["tool_calls"][1]["id"], "call_def456");
375        assert_eq!(
376            r.raw["tool_calls"][0]["function"]["arguments"],
377            "{\"item\":\"PRT-000001\"}"
378        );
379
380        let r = reply(fixture!("openai_length")).await.unwrap();
381        assert_eq!((r.stop, r.text.as_str()), (Stop::Truncated, "partial"));
382    }
383
384    #[tokio::test]
385    async fn azure_replies_are_decoded() {
386        // Azure adds content-filter annotations and a leading prompt_filter_results
387        let r = reply(fixture!("azure_tool_calls")).await.unwrap();
388        assert_eq!(r.stop, Stop::ToolCalls);
389        assert_eq!(r.calls[0].name, "bom_health");
390        assert_eq!(r.calls[0].input, Ok(json!({ "item": "ASM-000002" })));
391        assert!(r.usage.reported);
392
393        let r = reply(fixture!("azure_content_filter")).await.unwrap();
394        assert_eq!(r.stop, Stop::Refused);
395    }
396
397    #[tokio::test]
398    async fn mistral_replies_are_decoded() {
399        let r = reply(fixture!("mistral_tool_calls")).await.unwrap();
400        assert_eq!(r.stop, Stop::ToolCalls);
401        assert_eq!(r.calls[0].id, "D681PevKs", "9-character id kept");
402        assert_eq!(r.calls[0].input, Ok(json!({ "item": "PRT-000001" })));
403        assert_eq!(r.usage.input, 90);
404    }
405
406    #[tokio::test]
407    async fn ollama_replies_are_decoded() {
408        // `finish_reason: stop` together with tool calls: the calls win
409        let r = reply(fixture!("ollama_tool_call_stop")).await.unwrap();
410        assert_eq!(r.stop, Stop::ToolCalls);
411        assert_eq!(r.calls[0].input, Ok(json!({ "item": "PRT-000001" })));
412        // arguments sent as an object instead of a string
413        let r = reply(fixture!("ollama_object_arguments")).await.unwrap();
414        assert_eq!(r.calls[0].input, Ok(json!({ "item": "PRT-000002" })));
415        assert_eq!(
416            r.raw["tool_calls"][0]["function"]["arguments"], "{\"item\":\"PRT-000002\"}",
417            "sent back as the string the API expects"
418        );
419        // a plain-string error and a missing model
420        assert_eq!(
421            reply(fixture!("ollama_model_not_found")).await.unwrap_err(),
422            LlmError::Rejected {
423                status: 404,
424                message: "model \"llama9\" not found, try pulling it first".into()
425            }
426        );
427    }
428
429    #[tokio::test]
430    async fn vllm_replies_are_decoded() {
431        // no call id, no usage, empty arguments for a tool without input
432        let r = reply(fixture!("vllm_no_id_no_usage")).await.unwrap();
433        assert_eq!(r.stop, Stop::ToolCalls);
434        let id = &r.calls[0].id;
435        assert!(
436            id.len() == 9 && id.chars().all(|c| c.is_ascii_alphanumeric()),
437            "made-up id: {id}"
438        );
439        assert_eq!(
440            r.raw["tool_calls"][0]["id"],
441            id.as_str(),
442            "same id sent back"
443        );
444        assert_eq!(r.calls[0].name, "my_worklist");
445        assert_eq!(r.calls[0].input, Ok(json!({})));
446        assert!(!r.usage.reported);
447        assert_eq!((r.usage.input, r.usage.output), (0, 0));
448
449        // arguments that are not JSON: the call is kept, the input is an error
450        let r = reply(fixture!("vllm_bad_arguments")).await.unwrap();
451        assert_eq!(r.stop, Stop::ToolCalls);
452        assert!(r.calls[0].input.is_err());
453        assert_eq!(
454            r.raw["tool_calls"][0]["function"]["arguments"], "{item: PRT-1",
455            "the model sees what it sent"
456        );
457    }
458
459    #[test]
460    fn unusual_bodies() {
461        let d = |v: Value| OpenAi(Flavour::Compatible).decode(&v);
462        // arguments that are valid JSON but not an object
463        let r = d(
464            json!({ "choices": [{ "finish_reason": "tool_calls", "message": { "content": null,
465            "tool_calls": [{ "id": "a", "function": { "name": "x", "arguments": "[1]" } }] } }] }),
466        )
467        .unwrap();
468        assert!(r.calls[0].input.is_err());
469        // text in parts, and a refusal
470        let r = d(
471            json!({ "choices": [{ "finish_reason": "stop", "message": { "content":
472            [{ "type": "text", "text": "A" }, { "type": "text", "text": "B" }] } }] }),
473        )
474        .unwrap();
475        assert_eq!(r.text, "A\nB");
476        let r = d(json!({ "choices": [{ "finish_reason": "stop", "message": {
477            "content": null, "refusal": "I cannot help with that." } }] }))
478        .unwrap();
479        assert_eq!(r.stop, Stop::Refused);
480        // malformed bodies
481        for bad in [
482            json!({}),
483            json!({ "choices": [] }),
484            json!({ "choices": [{ "message": {} }] }),
485            json!({ "choices": [{ "finish_reason": "tool_calls", "message": { "content": "" } }] }),
486            json!({ "choices": [{ "finish_reason": "stop", "message": { "tool_calls": [{ "function": {} }] } }] }),
487        ] {
488            assert!(
489                matches!(d(bad.clone()), Err(LlmError::Malformed(_))),
490                "{bad}"
491            );
492        }
493    }
494
495    #[tokio::test]
496    async fn the_error_table() {
497        assert_eq!(
498            reply(fixture!("openai_401")).await.unwrap_err(),
499            LlmError::Auth
500        );
501        assert_eq!(
502            reply(fixture!("openai_rate_limit")).await.unwrap_err(),
503            LlmError::RateLimited {
504                retry_after: Some(Duration::from_secs(2))
505            }
506        );
507        assert_eq!(
508            reply(fixture!("openai_insufficient_quota"))
509                .await
510                .unwrap_err(),
511            LlmError::Quota
512        );
513        assert_eq!(
514            reply(fixture!("openai_context_length")).await.unwrap_err(),
515            LlmError::ContextLength
516        );
517        assert_eq!(
518            reply(fixture!("openai_503")).await.unwrap_err(),
519            LlmError::Overloaded(503)
520        );
521        let p = OpenAi(Flavour::OpenAi);
522        let h = HeaderMap::new();
523        assert_eq!(p.error(402, &h, ""), LlmError::Quota);
524        assert_eq!(p.error(403, &h, ""), LlmError::Auth);
525        assert_eq!(p.error(529, &h, ""), LlmError::Overloaded(529));
526        assert_eq!(
527            p.error(429, &h, ""),
528            LlmError::RateLimited { retry_after: None }
529        );
530        assert!(matches!(
531            p.error(302, &h, ""),
532            LlmError::Rejected { status: 302, .. }
533        ));
534        assert_eq!(
535            p.error(400, &h, "{\"error\":{\"message\":\"x\"}}"),
536            LlmError::Rejected {
537                status: 400,
538                message: "x".into()
539            }
540        );
541        assert_eq!(
542            p.error(
543                400,
544                &h,
545                "This model's maximum context length is 8192 tokens"
546            ),
547            LlmError::ContextLength
548        );
549    }
550}