1use 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
27pub enum Flavour {
28 OpenAi,
31 Compatible,
33}
34
35pub struct OpenAi(pub Flavour);
37
38fn made_up_id() -> String {
40 uuid::Uuid::new_v4()
41 .simple()
42 .to_string()
43 .chars()
44 .take(9)
45 .collect()
46}
47
48fn 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
67fn 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 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 "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 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 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 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 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 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 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 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 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 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 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 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 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}