Skip to main content

cairn_models/
openai.rs

1use async_trait::async_trait;
2use cairn_runtime::{
3    Content, Model, ModelError, ModelRequest, ModelResponse, ModelStream, Role, StopReason, StreamEvent, Usage,
4};
5use serde_json::{Value, json};
6
7use crate::sse::{self, SseEvent};
8use crate::{Pricing, merge_options, result_string, send_json};
9
10/// Chat Completions adapter for any OpenAI-compatible endpoint.
11pub struct OpenAiCompatModel {
12    provider: String,
13    model: String,
14    base_url: String,
15    api_key: Option<String>,
16    max_tokens_field: &'static str,
17    pricing: Option<Pricing>,
18    client: reqwest::Client,
19}
20
21impl OpenAiCompatModel {
22    /// Generic endpoint. `provider` names the key used for
23    /// `provider_options` and the prefix of [`Model::id`].
24    pub fn new(
25        provider: impl Into<String>,
26        model: impl Into<String>,
27        base_url: impl Into<String>,
28        api_key: Option<String>,
29    ) -> Self {
30        Self {
31            provider: provider.into(),
32            model: model.into(),
33            base_url: base_url.into().trim_end_matches('/').to_owned(),
34            api_key,
35            max_tokens_field: "max_tokens",
36            pricing: None,
37            client: reqwest::Client::new(),
38        }
39    }
40
41    pub fn openai(model: impl Into<String>, api_key: impl Into<String>) -> Self {
42        let mut m = Self::new("openai", model, "https://api.openai.com/v1", Some(api_key.into()));
43        m.max_tokens_field = "max_completion_tokens";
44        m
45    }
46
47    pub fn gemini(model: impl Into<String>, api_key: impl Into<String>) -> Self {
48        Self::new("gemini", model, "https://generativelanguage.googleapis.com/v1beta/openai", Some(api_key.into()))
49    }
50
51    pub fn openrouter(model: impl Into<String>, api_key: impl Into<String>) -> Self {
52        Self::new("openrouter", model, "https://openrouter.ai/api/v1", Some(api_key.into()))
53    }
54
55    pub fn with_pricing(mut self, pricing: Pricing) -> Self {
56        self.pricing = Some(pricing);
57        self
58    }
59
60    /// Reasoning this adapter's own provider produced, as opposed to another
61    /// provider's earlier in a fallback chain.
62    fn own_reasoning<'a>(&self, c: &'a Content) -> Option<&'a str> {
63        match c {
64            Content::Thinking { text, provider_data }
65                if provider_data.as_ref().and_then(|d| d["provider"].as_str()) == Some(self.provider.as_str()) =>
66            {
67                Some(text.as_str())
68            }
69            _ => None,
70        }
71    }
72
73    pub(crate) fn build_body(&self, req: &ModelRequest) -> Value {
74        // Reasoning models (IFM's K2, DeepSeek and others) return
75        // `reasoning_content` and refuse a later turn unless every assistant
76        // message carries it back verbatim. Once this provider has returned
77        // any, replay it on every assistant message: its own text where there
78        // is some, "" where another model wrote the turn.
79        let replay_reasoning = req.messages.iter().flat_map(|m| &m.content).any(|c| self.own_reasoning(c).is_some());
80        let mut messages = Vec::new();
81        if let Some(system) = &req.system {
82            messages.push(json!({ "role": "system", "content": system }));
83        }
84        for m in &req.messages {
85            match m.role {
86                Role::User => {
87                    let mut text = String::new();
88                    for c in &m.content {
89                        match c {
90                            Content::ToolResult { tool_use_id, content, .. } => messages.push(json!({
91                                "role": "tool",
92                                "tool_call_id": tool_use_id,
93                                "content": result_string(content),
94                            })),
95                            Content::Text { text: t } => text.push_str(t),
96                            _ => {}
97                        }
98                    }
99                    if !text.is_empty() {
100                        messages.push(json!({ "role": "user", "content": text }));
101                    }
102                }
103                Role::Assistant => {
104                    let text = cairn_runtime::message::content_text(&m.content);
105                    let calls: Vec<Value> = m
106                        .content
107                        .iter()
108                        .filter_map(|c| match c {
109                            Content::ToolUse { id, name, input } => Some(json!({
110                                "id": id,
111                                "type": "function",
112                                "function": { "name": name, "arguments": input.to_string() },
113                            })),
114                            _ => None,
115                        })
116                        .collect();
117                    let mut msg = json!({
118                        "role": "assistant",
119                        "content": if text.is_empty() { Value::Null } else { json!(text) },
120                    });
121                    if !calls.is_empty() {
122                        msg["tool_calls"] = json!(calls);
123                    }
124                    if replay_reasoning {
125                        let r: String = m.content.iter().filter_map(|c| self.own_reasoning(c)).collect();
126                        msg["reasoning_content"] = json!(r);
127                    }
128                    messages.push(msg);
129                }
130            }
131        }
132
133        let mut body = json!({ "model": self.model, "messages": messages });
134        body[self.max_tokens_field] = json!(req.max_tokens);
135        if !req.tools.is_empty() {
136            body["tools"] = json!(
137                req.tools
138                    .iter()
139                    .map(|t| json!({
140                        "type": "function",
141                        "function": { "name": t.name, "description": t.description, "parameters": t.input_schema },
142                    }))
143                    .collect::<Vec<_>>()
144            );
145        }
146        if let Some(t) = req.temperature {
147            body["temperature"] = json!(t);
148        }
149        merge_options(&mut body, req.provider_options.get(&self.provider));
150        body
151    }
152
153    pub(crate) fn parse_response(&self, v: &Value) -> Result<ModelResponse, ModelError> {
154        let choice = &v["choices"][0];
155        let msg = &choice["message"];
156        if msg.is_null() {
157            return Err(ModelError::Decode("missing `choices[0].message`".into()));
158        }
159
160        let mut content = Vec::new();
161        // Kept even when empty: "" is a value such providers expect back.
162        if let Some(r) = msg["reasoning_content"].as_str().or_else(|| msg["reasoning"].as_str()) {
163            content.push(Content::Thinking {
164                text: r.to_owned(),
165                provider_data: Some(json!({ "provider": self.provider })),
166            });
167        }
168        if let Some(t) = msg["content"].as_str().filter(|s| !s.is_empty()) {
169            content.push(Content::Text { text: t.to_owned() });
170        }
171        for call in msg["tool_calls"].as_array().into_iter().flatten() {
172            let args = call["function"]["arguments"].as_str().unwrap_or("{}");
173            content.push(Content::ToolUse {
174                id: call["id"].as_str().unwrap_or_default().to_owned(),
175                name: call["function"]["name"].as_str().unwrap_or_default().to_owned(),
176                input: serde_json::from_str(args).unwrap_or_else(|_| json!({ "_unparsed_arguments": args })),
177            });
178        }
179
180        let stop_reason = match choice["finish_reason"].as_str() {
181            Some("stop") => StopReason::EndTurn,
182            Some("tool_calls") | Some("function_call") => StopReason::ToolUse,
183            Some("length") => StopReason::MaxTokens,
184            Some(other) => StopReason::Other(other.to_owned()),
185            None => StopReason::Other("unknown".into()),
186        };
187        let u = &v["usage"];
188        let prompt = u["prompt_tokens"].as_u64().unwrap_or(0);
189        let cached = u["prompt_tokens_details"]["cached_tokens"].as_u64().unwrap_or(0);
190        Ok(ModelResponse {
191            model: v["model"].as_str().unwrap_or(&self.model).to_owned(),
192            content,
193            stop_reason,
194            usage: Usage {
195                input_tokens: prompt.saturating_sub(cached),
196                output_tokens: u["completion_tokens"].as_u64().unwrap_or(0),
197                cache_read_tokens: cached,
198                cache_write_tokens: 0,
199            },
200            ..Default::default()
201        })
202    }
203}
204
205#[async_trait]
206impl Model for OpenAiCompatModel {
207    fn id(&self) -> String {
208        format!("{}:{}", self.provider, self.model)
209    }
210
211    async fn complete(&self, request: &ModelRequest) -> Result<ModelResponse, ModelError> {
212        let body = self.build_body(request);
213        let mut http = self.client.post(format!("{}/chat/completions", self.base_url));
214        if let Some(key) = &self.api_key {
215            http = http.bearer_auth(key);
216        }
217        let v = send_json(http, &body).await?;
218        self.parse_response(&v)
219    }
220
221    async fn stream(&self, request: &ModelRequest) -> Result<ModelStream, ModelError> {
222        let mut body = self.build_body(request);
223        body["stream"] = json!(true);
224        body["stream_options"] = json!({ "include_usage": true });
225        let mut http = self.client.post(format!("{}/chat/completions", self.base_url));
226        if let Some(key) = &self.api_key {
227            http = http.bearer_auth(key);
228        }
229        let resp = http.json(&body).send().await.map_err(|e| ModelError::Transport(e.to_string()))?;
230        let resp = sse::check_status(resp).await?;
231        let acc = StreamAcc { provider: self.provider.clone(), model: self.model.clone(), ..Default::default() };
232        Ok(sse::translate(sse::events(resp), acc, StreamAcc::step, StreamAcc::end))
233    }
234
235    fn pricing(&self) -> Option<Pricing> {
236        self.pricing
237    }
238}
239
240#[derive(Default)]
241pub(crate) struct StreamAcc {
242    provider: String,
243    model: String,
244    text: String,
245    reasoning: String,
246    /// Whether the stream carried a reasoning field at all, even an empty one.
247    saw_reasoning: bool,
248    /// (id, name, arguments) per tool-call index.
249    tools: Vec<(String, String, String)>,
250    finish: Option<String>,
251    usage: Usage,
252    done: bool,
253}
254
255impl StreamAcc {
256    pub(crate) fn step(&mut self, ev: SseEvent) -> Result<Vec<StreamEvent>, ModelError> {
257        if ev.data.trim() == "[DONE]" {
258            return self.complete();
259        }
260        let v: Value = serde_json::from_str(&ev.data).map_err(|e| ModelError::Decode(e.to_string()))?;
261        if !v["error"].is_null() {
262            let status = v["error"]["code"].as_u64().filter(|c| (400..600).contains(c)).unwrap_or(500) as u16;
263            return Err(ModelError::Api { status, body: v["error"].to_string() });
264        }
265        if let Some(m) = v["model"].as_str() {
266            self.model = m.to_owned();
267        }
268        if v["usage"].is_object() {
269            let u = &v["usage"];
270            let prompt = u["prompt_tokens"].as_u64().unwrap_or(0);
271            let cached = u["prompt_tokens_details"]["cached_tokens"].as_u64().unwrap_or(0);
272            self.usage = Usage {
273                input_tokens: prompt.saturating_sub(cached),
274                output_tokens: u["completion_tokens"].as_u64().unwrap_or(0),
275                cache_read_tokens: cached,
276                cache_write_tokens: 0,
277            };
278        }
279
280        let mut out = Vec::new();
281        let choice = &v["choices"][0];
282        let delta = &choice["delta"];
283        if let Some(r) = delta["reasoning_content"].as_str().or_else(|| delta["reasoning"].as_str()) {
284            self.saw_reasoning = true;
285            if !r.is_empty() {
286                self.reasoning.push_str(r);
287                out.push(StreamEvent::ThinkingDelta(r.to_owned()));
288            }
289        }
290        if let Some(t) = delta["content"].as_str().filter(|s| !s.is_empty()) {
291            self.text.push_str(t);
292            out.push(StreamEvent::TextDelta(t.to_owned()));
293        }
294        for call in delta["tool_calls"].as_array().into_iter().flatten() {
295            let index = call["index"].as_u64().unwrap_or(self.tools.len() as u64) as usize;
296            if index >= self.tools.len() {
297                self.tools.resize(index + 1, Default::default());
298            }
299            let slot = &mut self.tools[index];
300            if let Some(id) = call["id"].as_str().filter(|s| !s.is_empty()) {
301                slot.0 = id.to_owned();
302                slot.1 = call["function"]["name"].as_str().unwrap_or_default().to_owned();
303                out.push(StreamEvent::ToolCallStart { index, id: slot.0.clone(), name: slot.1.clone() });
304            }
305            if let Some(args) = call["function"]["arguments"].as_str().filter(|s| !s.is_empty()) {
306                slot.2.push_str(args);
307                out.push(StreamEvent::ToolInputDelta { index, partial_json: args.to_owned() });
308            }
309        }
310        if let Some(f) = choice["finish_reason"].as_str() {
311            self.finish = Some(f.to_owned());
312        }
313        Ok(out)
314    }
315
316    /// Some servers close the stream without `[DONE]`; complete if they
317    /// at least sent a finish reason.
318    pub(crate) fn end(&mut self) -> Result<Vec<StreamEvent>, ModelError> {
319        if self.done {
320            Ok(Vec::new())
321        } else if self.finish.is_some() {
322            self.complete()
323        } else {
324            Err(ModelError::Transport("stream closed before completion".into()))
325        }
326    }
327
328    fn complete(&mut self) -> Result<Vec<StreamEvent>, ModelError> {
329        if self.done {
330            return Ok(Vec::new());
331        }
332        self.done = true;
333        let mut content = Vec::new();
334        if self.saw_reasoning || !self.reasoning.is_empty() {
335            content.push(Content::Thinking {
336                text: std::mem::take(&mut self.reasoning),
337                provider_data: Some(json!({ "provider": self.provider })),
338            });
339        }
340        if !self.text.is_empty() {
341            content.push(Content::Text { text: std::mem::take(&mut self.text) });
342        }
343        for (id, name, args) in std::mem::take(&mut self.tools) {
344            let input = if args.trim().is_empty() {
345                json!({})
346            } else {
347                serde_json::from_str(&args).unwrap_or_else(|_| json!({ "_unparsed_arguments": args }))
348            };
349            content.push(Content::ToolUse { id, name, input });
350        }
351        let stop_reason = match self.finish.as_deref() {
352            Some("stop") => StopReason::EndTurn,
353            Some("tool_calls") | Some("function_call") => StopReason::ToolUse,
354            Some("length") => StopReason::MaxTokens,
355            Some(other) => StopReason::Other(other.to_owned()),
356            None => StopReason::Other("unknown".into()),
357        };
358        Ok(vec![StreamEvent::Completed(ModelResponse {
359            model: self.model.clone(),
360            content,
361            stop_reason,
362            usage: self.usage,
363            ..Default::default()
364        })])
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371    use cairn_runtime::Message;
372
373    #[test]
374    fn reasoning_is_replayed_verbatim_once_the_provider_returns_it() {
375        let m = OpenAiCompatModel::new("k2", "IFM/K2", "https://example.invalid/v1", None::<String>);
376        let first = m
377            .parse_response(&json!({"model": "IFM/K2", "choices": [{"finish_reason": "stop",
378                "message": {"content": "Hello", "reasoning_content": "think think"}}]}))
379            .unwrap();
380        let empty = m
381            .parse_response(&json!({"model": "IFM/K2", "choices": [{"finish_reason": "stop",
382                "message": {"content": "Again", "reasoning_content": ""}}]}))
383            .unwrap();
384        let req = ModelRequest {
385            messages: vec![
386                Message::user("a"),
387                Message::assistant(first.content),
388                Message::user("b"),
389                // A turn another model wrote, after a fallback.
390                Message::assistant(vec![Content::Text { text: "from haiku".into() }]),
391                Message::user("c"),
392                Message::assistant(empty.content),
393                Message::user("d"),
394            ],
395            max_tokens: 10,
396            ..Default::default()
397        };
398        let msgs = m.build_body(&req)["messages"].as_array().unwrap().clone();
399        assert_eq!(msgs[1]["reasoning_content"], "think think");
400        assert_eq!(msgs[1]["content"], "Hello", "reasoning never leaks into the content");
401        assert_eq!(msgs[3]["reasoning_content"], "", "another model's turn carries an empty one");
402        assert_eq!(msgs[5]["reasoning_content"], "", "an empty one is kept, not dropped");
403    }
404
405    #[test]
406    fn providers_without_reasoning_get_no_reasoning_field() {
407        let m = OpenAiCompatModel::openai("gpt-x", "k");
408        let req = ModelRequest {
409            messages: vec![
410                Message::user("a"),
411                Message::assistant(vec![
412                    // Another provider's thinking is not this one's to replay.
413                    Content::Thinking { text: "x".into(), provider_data: Some(json!({"provider": "anthropic"})) },
414                    Content::Text { text: "hi".into() },
415                ]),
416                Message::user("b"),
417            ],
418            max_tokens: 10,
419            ..Default::default()
420        };
421        let msgs = m.build_body(&req)["messages"].as_array().unwrap().clone();
422        assert!(msgs[1].get("reasoning_content").is_none());
423    }
424
425    #[test]
426    fn round_trips_tool_calls() {
427        let m = OpenAiCompatModel::openai("gpt-x", "k");
428        let req = ModelRequest {
429            system: Some("sys".into()),
430            messages: vec![
431                Message::user("hi"),
432                Message::assistant(vec![
433                    Content::Thinking { text: "x".into(), provider_data: Some(json!({"provider": "anthropic"})) },
434                    Content::ToolUse { id: "c1".into(), name: "f".into(), input: json!({"a": 1}) },
435                ]),
436                Message {
437                    role: Role::User,
438                    content: vec![Content::ToolResult {
439                        tool_use_id: "c1".into(),
440                        content: json!("ok"),
441                        is_error: false,
442                    }],
443                },
444            ],
445            max_tokens: 50,
446            ..Default::default()
447        };
448        let body = m.build_body(&req);
449        assert_eq!(body["max_completion_tokens"], 50);
450        let msgs = body["messages"].as_array().unwrap();
451        assert_eq!(msgs.len(), 4);
452        assert_eq!(msgs[2]["tool_calls"][0]["function"]["arguments"], r#"{"a":1}"#);
453        assert!(msgs[2]["content"].is_null());
454        assert_eq!(msgs[3]["role"], "tool");
455
456        let resp = m
457            .parse_response(&json!({
458                "model": "gpt-x",
459                "choices": [{"finish_reason": "tool_calls", "message": {"content": null, "tool_calls": [
460                    {"id": "c2", "type": "function", "function": {"name": "f", "arguments": "{\"a\":2}"}}
461                ]}}],
462                "usage": {"prompt_tokens": 100, "completion_tokens": 7, "prompt_tokens_details": {"cached_tokens": 60}}
463            }))
464            .unwrap();
465        assert_eq!(resp.stop_reason, StopReason::ToolUse);
466        assert_eq!(resp.content[0], Content::ToolUse { id: "c2".into(), name: "f".into(), input: json!({"a": 2}) });
467        assert_eq!(resp.usage.input_tokens, 40);
468        assert_eq!(resp.usage.cache_read_tokens, 60);
469    }
470
471    #[test]
472    fn accumulates_stream_with_parallel_tool_calls() {
473        let raw = r#"data: {"model":"gpt-x","choices":[{"index":0,"delta":{"role":"assistant","content":""}}]}
474
475data: {"choices":[{"index":0,"delta":{"content":"Check"}}]}
476
477data: {"choices":[{"index":0,"delta":{"content":"ing"}}]}
478
479data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"c1","type":"function","function":{"name":"a","arguments":""}}]}}]}
480
481data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"id":"c2","type":"function","function":{"name":"b","arguments":"{}"}}]}}]}
482
483data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"x\":"}}]}}]}
484
485data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"1}"}}]}}]}
486
487data: {"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}
488
489data: {"choices":[],"usage":{"prompt_tokens":50,"completion_tokens":9,"prompt_tokens_details":{"cached_tokens":20}}}
490
491data: [DONE]
492"#;
493        let mut acc = StreamAcc { provider: "openai".into(), ..Default::default() };
494        let events: Vec<StreamEvent> = sse::parse_text(raw).into_iter().flat_map(|e| acc.step(e).unwrap()).collect();
495        assert_eq!(events[0], StreamEvent::TextDelta("Check".into()));
496        assert!(events.contains(&StreamEvent::ToolCallStart { index: 1, id: "c2".into(), name: "b".into() }));
497        let StreamEvent::Completed(r) = events.last().unwrap() else { panic!("no completion") };
498        assert_eq!(r.stop_reason, StopReason::ToolUse);
499        assert_eq!(r.content[0], Content::Text { text: "Checking".into() });
500        assert_eq!(r.content[1], Content::ToolUse { id: "c1".into(), name: "a".into(), input: json!({"x": 1}) });
501        assert_eq!(r.content[2], Content::ToolUse { id: "c2".into(), name: "b".into(), input: json!({}) });
502        assert_eq!(r.usage.input_tokens, 30);
503        assert_eq!(r.usage.cache_read_tokens, 20);
504        assert!(acc.end().unwrap().is_empty(), "no duplicate completion");
505    }
506}