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
10pub 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 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 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 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 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 saw_reasoning: bool,
248 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 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 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 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}