1use serde::{Deserialize, Serialize};
2use tokio::sync::broadcast;
3use tokio_util::sync::CancellationToken;
4
5use crate::error::RuntimeError;
6use crate::event::{NodeEvent, Observable};
7use crate::message::{ImageData, Message, MessagePart, MessageRole};
8use crate::provider::{
9 AssistantMessage, CallTiming, DEFAULT_STREAM_BUFFER, LlmRequest, Provider, StopReason,
10 TokenUsage, estimate_tokens,
11};
12use crate::providers::classify_attachment_error;
13use crate::tool::BoxFut;
14
15pub struct OpenAiProvider {
16 name: String,
17 api_key: String,
18 base_url: String,
19 client: reqwest::Client,
20 max_tokens: Option<u32>,
21}
22
23impl OpenAiProvider {
24 pub fn new(name: impl Into<String>, api_key: impl Into<String>) -> Self {
25 Self {
26 name: name.into(),
27 api_key: api_key.into(),
28 base_url: "https://api.openai.com/v1".into(),
29 client: reqwest::Client::new(),
30 max_tokens: None,
31 }
32 }
33
34 pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
35 self.base_url = url.into();
36 self
37 }
38
39 pub fn with_max_tokens(mut self, n: u32) -> Self {
40 self.max_tokens = Some(n);
41 self
42 }
43
44 fn build_body(&self, req: &LlmRequest, stream: bool) -> ChatCompletionsRequest {
45 let mut wire_messages: Vec<ChatMessage> = Vec::new();
46 if let Some(sys) = &req.system {
47 wire_messages.push(ChatMessage {
48 role: "system",
49 content: Some(ChatContent::Text(sys.clone())),
50 tool_calls: None,
51 tool_call_id: None,
52 });
53 }
54 for m in &req.messages {
55 wire_messages.push(build_wire_message(m));
56 }
57 let tools: Vec<WireToolSpec> = req
58 .tools
59 .iter()
60 .map(|t| WireToolSpec {
61 kind: "function",
62 function: WireToolFunction {
63 name: t.name.clone(),
64 description: t.description.clone(),
65 parameters: t.input_schema.clone(),
66 },
67 })
68 .collect();
69 ChatCompletionsRequest {
70 model: req.model.clone(),
71 stream,
72 max_tokens: self.max_tokens,
73 messages: wire_messages,
74 tools,
75 stream_options: if stream {
76 Some(StreamOptions {
77 include_usage: true,
78 })
79 } else {
80 None
81 },
82 }
83 }
84
85 fn build_request(&self, req: &LlmRequest, stream: bool) -> reqwest::RequestBuilder {
86 let body = self.build_body(req, stream);
87 self.client
88 .post(format!("{}/chat/completions", self.base_url))
89 .bearer_auth(&self.api_key)
90 .json(&body)
91 }
92
93 #[doc(hidden)]
94 pub fn wire_body_bytes(&self, req: &LlmRequest, stream: bool) -> Vec<u8> {
95 serde_json::to_vec(&self.build_body(req, stream)).expect("serialize wire body")
96 }
97}
98
99fn build_wire_message(m: &Message) -> ChatMessage {
100 match m.role {
101 MessageRole::System => ChatMessage {
102 role: "system",
103 content: Some(ChatContent::Text(m.text_concat())),
104 tool_calls: None,
105 tool_call_id: None,
106 },
107 MessageRole::Tool => {
108 let (id, content) = extract_tool_result(m);
109 ChatMessage {
110 role: "tool",
111 content: Some(ChatContent::Text(content)),
112 tool_calls: None,
113 tool_call_id: Some(id),
114 }
115 }
116 MessageRole::Assistant => {
117 let (text_parts, tool_uses) = split_assistant_parts(&m.parts);
118 let content = if text_parts.is_empty() {
119 None
120 } else {
121 Some(ChatContent::Text(text_parts.join("")))
122 };
123 let tool_calls = if tool_uses.is_empty() {
124 None
125 } else {
126 Some(tool_uses)
127 };
128 ChatMessage {
129 role: "assistant",
130 content,
131 tool_calls,
132 tool_call_id: None,
133 }
134 }
135 MessageRole::User => {
136 let parts = build_user_parts(&m.parts);
137 let content = if parts.iter().all(|p| matches!(p, ChatPart::Text { .. })) {
138 let joined: String = parts
139 .iter()
140 .filter_map(|p| match p {
141 ChatPart::Text { text } => Some(text.as_str()),
142 _ => None,
143 })
144 .collect();
145 Some(ChatContent::Text(joined))
146 } else {
147 Some(ChatContent::Parts(parts))
148 };
149 ChatMessage {
150 role: "user",
151 content,
152 tool_calls: None,
153 tool_call_id: None,
154 }
155 }
156 }
157}
158
159fn build_user_parts(parts: &[MessagePart]) -> Vec<ChatPart> {
160 let mut out = Vec::with_capacity(parts.len());
161 for p in parts {
162 match p {
163 MessagePart::CompactSummary { summary, .. } => out.push(ChatPart::Text {
164 text: summary.clone(),
165 }),
166 MessagePart::Text { text } => out.push(ChatPart::Text { text: text.clone() }),
167 MessagePart::Image { source } => {
168 let url = match &source.data {
169 ImageData::Base64 { data } => {
170 format!("data:{};base64,{}", source.media_type, data)
171 }
172 ImageData::Path { path } => {
173 let bytes = std::fs::read(path).unwrap_or_default();
174 use base64::Engine;
175 let data = base64::engine::general_purpose::STANDARD.encode(&bytes);
176 format!("data:{};base64,{}", source.media_type, data)
177 }
178 };
179 out.push(ChatPart::ImageUrl {
180 image_url: ImageUrl { url },
181 });
182 }
183 _ => {}
184 }
185 }
186 out
187}
188
189fn extract_tool_result(m: &Message) -> (String, String) {
190 for p in &m.parts {
191 if let MessagePart::ToolResult {
192 tool_use_id,
193 content,
194 ..
195 } = p
196 {
197 return (tool_use_id.clone(), content.clone());
198 }
199 }
200 (String::new(), m.text_concat())
201}
202
203fn split_assistant_parts(parts: &[MessagePart]) -> (Vec<String>, Vec<WireToolCall>) {
204 let mut text = Vec::new();
205 let mut tools = Vec::new();
206 for p in parts {
207 match p {
208 MessagePart::Text { text: t } => text.push(t.clone()),
209 MessagePart::ToolUse { id, name, input } => tools.push(WireToolCall {
210 id: id.clone(),
211 kind: "function",
212 function: WireFunctionCall {
213 name: name.clone(),
214 arguments: input.to_string(),
215 },
216 }),
217 _ => {}
218 }
219 }
220 (text, tools)
221}
222
223impl Provider for OpenAiProvider {
224 fn name(&self) -> &str {
225 &self.name
226 }
227
228 fn call<'a>(&'a self, req: LlmRequest) -> BoxFut<'a, Result<AssistantMessage, RuntimeError>> {
229 let request = self.build_request(&req, false);
230 let turn_id = next_turn_id_from_req(&req);
231 Box::pin(async move {
232 let resp = request.send().await.map_err(net_err)?;
233 let status = resp.status();
234 let body: ChatCompletionsResponse = if status.is_success() {
235 resp.json().await.map_err(net_err)?
236 } else {
237 let body_text = resp.text().await.unwrap_or_default();
238 if let Some(reason) = classify_attachment_error(status.as_u16(), &body_text) {
239 return Err(RuntimeError::AttachmentError { reason });
240 }
241 return Err(RuntimeError::ToolFailed(format!(
242 "openai http {status}: {body_text}"
243 )));
244 };
245 Ok(response_to_assistant(body, turn_id))
246 })
247 }
248
249 fn call_streaming(&self, req: LlmRequest) -> Observable<AssistantMessage> {
250 let request = self.build_request(&req, true);
251 let turn_id = next_turn_id_from_req(&req);
252 let (tx, events) = broadcast::channel(DEFAULT_STREAM_BUFFER);
253 let cancel = CancellationToken::new();
254 let cancel_for_task = cancel.clone();
255 let output: BoxFut<'static, Result<AssistantMessage, RuntimeError>> = Box::pin(
256 async move {
257 use eventsource_stream::Eventsource;
258 use futures::StreamExt;
259
260 let resp = tokio::select! {
261 biased;
262 _ = cancel_for_task.cancelled() => return Err(RuntimeError::Cancelled("openai cancelled before send".into())),
263 r = request.send() => r.map_err(net_err)?,
264 };
265 let status = resp.status();
266 if !status.is_success() {
267 let body = resp.text().await.unwrap_or_default();
268 if let Some(reason) = classify_attachment_error(status.as_u16(), &body) {
269 return Err(RuntimeError::AttachmentError { reason });
270 }
271 return Err(RuntimeError::ToolFailed(format!(
272 "openai http {status}: {body}"
273 )));
274 }
275
276 let mut stream = resp.bytes_stream().eventsource();
277 let mut acc_text = String::new();
278 let mut acc_thinking = String::new();
279 let mut cumulative = 0u64;
280 let mut final_usage: Option<OpenAiUsage> = None;
281 let mut resp_model: Option<String> = None;
282 let mut resp_id: Option<String> = None;
283 let mut partial_tool_calls: Vec<PartialToolCall> = Vec::new();
284 let mut stop_reason = StopReason::End;
285 while let Some(event) = tokio::select! {
286 biased;
287 _ = cancel_for_task.cancelled() => None,
288 next = stream.next() => next,
289 } {
290 let event = event.map_err(|e| RuntimeError::ToolFailed(format!("sse: {e}")))?;
291 if event.data == "[DONE]" {
292 break;
293 }
294 if event.data.is_empty() {
295 continue;
296 }
297 let parsed: serde_json::Value = match serde_json::from_str(&event.data) {
298 Ok(v) => v,
299 Err(_) => continue,
300 };
301 if let Some(content) = parsed
302 .pointer("/choices/0/delta/content")
303 .and_then(|v| v.as_str())
304 && !content.is_empty()
305 {
306 acc_text.push_str(content);
307 cumulative += estimate_tokens(content);
308 let _ = tx.send(NodeEvent::LlmChunk {
309 text: content.to_string(),
310 cumulative_tokens: cumulative,
311 });
312 }
313 if let Some(reasoning) = parsed
314 .pointer("/choices/0/delta/reasoning_content")
315 .and_then(|v| v.as_str())
316 && !reasoning.is_empty()
317 {
318 acc_thinking.push_str(reasoning);
319 let _ = tx.send(NodeEvent::ThinkingChunk {
320 text: reasoning.to_string(),
321 });
322 }
323 if let Some(m) = parsed.get("model").and_then(|v| v.as_str()) {
324 resp_model = Some(m.to_string());
325 }
326 if let Some(id) = parsed.get("id").and_then(|v| v.as_str()) {
327 resp_id = Some(id.to_string());
328 }
329 if let Some(tcs) = parsed
330 .pointer("/choices/0/delta/tool_calls")
331 .and_then(|v| v.as_array())
332 {
333 for tc in tcs {
334 let idx =
335 tc.get("index").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
336 while partial_tool_calls.len() <= idx {
337 partial_tool_calls.push(PartialToolCall::default());
338 }
339 let slot = &mut partial_tool_calls[idx];
340 if let Some(id) = tc.get("id").and_then(|v| v.as_str()) {
341 slot.id = id.to_string();
342 }
343 if let Some(name) =
344 tc.pointer("/function/name").and_then(|v| v.as_str())
345 {
346 slot.name = name.to_string();
347 }
348 if let Some(args) =
349 tc.pointer("/function/arguments").and_then(|v| v.as_str())
350 {
351 slot.arguments.push_str(args);
352 }
353 }
354 }
355 if let Some(reason) = parsed
356 .pointer("/choices/0/finish_reason")
357 .and_then(|v| v.as_str())
358 {
359 stop_reason = parse_stop_reason(reason);
360 }
361 if let Some(usage_obj) = parsed.get("usage") {
362 if !usage_obj.is_null() {
363 final_usage =
364 serde_json::from_value::<OpenAiUsage>(usage_obj.clone()).ok();
365 }
366 }
367 }
368 if cancel_for_task.is_cancelled() {
369 let _ = tx.send(NodeEvent::LlmDone {
370 total_tokens: cumulative,
371 });
372 return Err(RuntimeError::Cancelled(
373 "openai cancelled mid-stream".into(),
374 ));
375 }
376 let total = final_usage
377 .as_ref()
378 .and_then(|u| u.completion_tokens)
379 .unwrap_or(cumulative);
380 let _ = tx.send(NodeEvent::LlmDone {
381 total_tokens: total,
382 });
383
384 let mut parts: Vec<MessagePart> = Vec::new();
385 if !acc_thinking.is_empty() {
386 parts.push(MessagePart::Thinking {
387 thinking: acc_thinking,
388 signature: None,
389 });
390 }
391 if !acc_text.is_empty() {
392 parts.push(MessagePart::Text { text: acc_text });
393 }
394 for tc in partial_tool_calls {
395 if tc.id.is_empty() && tc.name.is_empty() && tc.arguments.is_empty() {
396 continue;
397 }
398 let input: serde_json::Value = if tc.arguments.is_empty() {
399 serde_json::Value::Object(Default::default())
400 } else {
401 serde_json::from_str(&tc.arguments).unwrap_or(serde_json::Value::Null)
402 };
403 parts.push(MessagePart::ToolUse {
404 id: tc.id,
405 name: tc.name,
406 input,
407 });
408 }
409
410 let token_usage = if let Some(u) = &final_usage {
411 let cached = u
412 .prompt_tokens_details
413 .as_ref()
414 .and_then(|d| d.cached_tokens)
415 .unwrap_or(0);
416 let cache_write = u
417 .prompt_tokens_details
418 .as_ref()
419 .and_then(|d| d.cache_write_tokens)
420 .unwrap_or(0);
421 let cache_read = cached.max(u.prompt_cache_hit_tokens.unwrap_or(0));
422 let reasoning = u
423 .completion_tokens_details
424 .as_ref()
425 .and_then(|d| d.reasoning_tokens)
426 .unwrap_or(0);
427 TokenUsage {
428 input: u.prompt_tokens.unwrap_or(0).saturating_sub(cache_read),
429 cached_input: cache_read,
430 output: u.completion_tokens.unwrap_or(0),
431 cache_write,
432 reasoning_tokens: reasoning,
433 }
434 } else {
435 TokenUsage {
436 output: total,
437 ..Default::default()
438 }
439 };
440
441 Ok(AssistantMessage {
442 message: Message {
443 role: MessageRole::Assistant,
444 parts,
445 turn_id,
446 },
447 stop_reason,
448 token_usage,
449 timing: CallTiming::default(),
450 model: resp_model.unwrap_or_default(),
451 response_id: resp_id,
452 })
453 },
454 );
455 Observable {
456 output,
457 events,
458 cancel,
459 }
460 }
461}
462
463#[derive(Default)]
464struct PartialToolCall {
465 id: String,
466 name: String,
467 arguments: String,
468}
469
470fn response_to_assistant(
471 body: ChatCompletionsResponse,
472 turn_id: crate::event::TurnId,
473) -> AssistantMessage {
474 let mut parts: Vec<MessagePart> = Vec::new();
475 let mut stop_reason = StopReason::End;
476 if let Some(choice) = body.choices.into_iter().next() {
477 if let Some(msg) = choice.message {
478 if let Some(content) = msg.content {
479 parts.push(MessagePart::Text { text: content });
480 }
481 if let Some(tool_calls) = msg.tool_calls {
482 for tc in tool_calls {
483 let input: serde_json::Value = if tc.function.arguments.is_empty() {
484 serde_json::Value::Object(Default::default())
485 } else {
486 serde_json::from_str(&tc.function.arguments)
487 .unwrap_or(serde_json::Value::Null)
488 };
489 parts.push(MessagePart::ToolUse {
490 id: tc.id,
491 name: tc.function.name,
492 input,
493 });
494 }
495 }
496 }
497 if let Some(reason) = choice.finish_reason {
498 stop_reason = parse_stop_reason(&reason);
499 }
500 }
501 let usage = body.usage.map(|u| {
502 let cached = u
503 .prompt_tokens_details
504 .as_ref()
505 .and_then(|d| d.cached_tokens)
506 .unwrap_or(0);
507 let cache_write = u
508 .prompt_tokens_details
509 .as_ref()
510 .and_then(|d| d.cache_write_tokens)
511 .unwrap_or(0);
512 let cache_read = cached.max(u.prompt_cache_hit_tokens.unwrap_or(0));
513 let reasoning = u
514 .completion_tokens_details
515 .as_ref()
516 .and_then(|d| d.reasoning_tokens)
517 .unwrap_or(0);
518 TokenUsage {
519 input: u.prompt_tokens.unwrap_or(0).saturating_sub(cache_read),
520 cached_input: cache_read,
521 output: u.completion_tokens.unwrap_or(0),
522 cache_write,
523 reasoning_tokens: reasoning,
524 }
525 });
526 AssistantMessage {
527 message: Message {
528 role: MessageRole::Assistant,
529 parts,
530 turn_id,
531 },
532 stop_reason,
533 token_usage: usage.unwrap_or_default(),
534 timing: CallTiming::default(),
535 model: body.model.unwrap_or_default(),
536 response_id: body.id,
537 }
538}
539
540fn parse_stop_reason(s: &str) -> StopReason {
541 match s {
542 "tool_calls" | "function_call" => StopReason::ToolUse,
543 "length" => StopReason::Length,
544 _ => StopReason::End,
545 }
546}
547
548fn next_turn_id_from_req(req: &LlmRequest) -> crate::event::TurnId {
549 req.messages
550 .first()
551 .map(|m| m.turn_id.clone())
552 .unwrap_or_else(crate::event::TurnId::now)
553}
554
555fn net_err(e: reqwest::Error) -> RuntimeError {
556 RuntimeError::ToolFailed(format!("openai net: {e}"))
557}
558
559#[derive(Serialize)]
560struct ChatCompletionsRequest {
561 model: String,
562 stream: bool,
563 #[serde(skip_serializing_if = "Option::is_none")]
564 max_tokens: Option<u32>,
565 messages: Vec<ChatMessage>,
566 #[serde(skip_serializing_if = "Vec::is_empty")]
567 tools: Vec<WireToolSpec>,
568 #[serde(skip_serializing_if = "Option::is_none")]
569 stream_options: Option<StreamOptions>,
570}
571
572#[derive(Serialize)]
573struct StreamOptions {
574 include_usage: bool,
575}
576
577#[derive(Serialize)]
578struct WireToolSpec {
579 #[serde(rename = "type")]
580 kind: &'static str,
581 function: WireToolFunction,
582}
583
584#[derive(Serialize)]
585struct WireToolFunction {
586 name: String,
587 #[serde(skip_serializing_if = "Option::is_none")]
588 description: Option<String>,
589 parameters: serde_json::Value,
590}
591
592#[derive(Serialize)]
593struct ChatMessage {
594 role: &'static str,
595 #[serde(skip_serializing_if = "Option::is_none")]
596 content: Option<ChatContent>,
597 #[serde(skip_serializing_if = "Option::is_none")]
598 tool_calls: Option<Vec<WireToolCall>>,
599 #[serde(skip_serializing_if = "Option::is_none")]
600 tool_call_id: Option<String>,
601}
602
603#[derive(Serialize)]
604#[serde(untagged)]
605enum ChatContent {
606 Text(String),
607 Parts(Vec<ChatPart>),
608}
609
610#[derive(Serialize)]
611#[serde(tag = "type", rename_all = "snake_case")]
612enum ChatPart {
613 Text { text: String },
614 ImageUrl { image_url: ImageUrl },
615}
616
617#[derive(Serialize)]
618struct ImageUrl {
619 url: String,
620}
621
622#[derive(Serialize)]
623struct WireToolCall {
624 id: String,
625 #[serde(rename = "type")]
626 kind: &'static str,
627 function: WireFunctionCall,
628}
629
630#[derive(Serialize)]
631struct WireFunctionCall {
632 name: String,
633 arguments: String,
634}
635
636#[derive(Deserialize)]
637struct ChatCompletionsResponse {
638 choices: Vec<ChatChoice>,
639 #[serde(default)]
640 usage: Option<OpenAiUsage>,
641 #[serde(default)]
642 model: Option<String>,
643 #[serde(default)]
644 id: Option<String>,
645}
646
647#[derive(Deserialize, Default)]
648struct OpenAiUsage {
649 #[serde(default)]
650 prompt_tokens: Option<u64>,
651 #[serde(default)]
652 completion_tokens: Option<u64>,
653 #[serde(default)]
654 prompt_tokens_details: Option<PromptTokensDetails>,
655 #[serde(default)]
656 completion_tokens_details: Option<CompletionTokensDetails>,
657 #[serde(default)]
658 prompt_cache_hit_tokens: Option<u64>,
659 #[serde(default)]
660 #[allow(dead_code)]
661 prompt_cache_miss_tokens: Option<u64>,
662}
663
664#[derive(Deserialize, Default)]
665struct PromptTokensDetails {
666 #[serde(default)]
667 cached_tokens: Option<u64>,
668 #[serde(default)]
669 cache_write_tokens: Option<u64>,
670}
671
672#[derive(Deserialize, Default)]
673struct CompletionTokensDetails {
674 #[serde(default)]
675 reasoning_tokens: Option<u64>,
676}
677
678#[derive(Deserialize)]
679struct ChatChoice {
680 #[serde(default)]
681 message: Option<ChatChoiceMessage>,
682 #[serde(default)]
683 finish_reason: Option<String>,
684}
685
686#[derive(Deserialize)]
687struct ChatChoiceMessage {
688 #[serde(default)]
689 content: Option<String>,
690 #[serde(default)]
691 tool_calls: Option<Vec<RespToolCall>>,
692}
693
694#[derive(Deserialize)]
695struct RespToolCall {
696 id: String,
697 function: RespFunctionCall,
698}
699
700#[derive(Deserialize)]
701struct RespFunctionCall {
702 name: String,
703 #[serde(default)]
704 arguments: String,
705}