1#[cfg(test)]
2#[path = "openai_test.rs"]
3mod tests;
4
5use crate::backend::mcp::{Tool, ToolInputSchema};
6use crate::backend::utils::context_truncation;
7use crate::backend::{ArcBackend, Backend, TITLE_PROMPT};
8use crate::config::{self, ModelSetting, user_agent};
9use crate::models::{
10 ArcEventTx, BackendConnection, BackendPrompt, BackendResponse, BackendUsage, Event, Message,
11 Model,
12};
13use crate::{info_event, warn_event};
14use async_trait::async_trait;
15use eyre::{Context, Result, bail};
16use futures::TryStreamExt;
17use serde::{Deserialize, Serialize};
18use serde_json::Value;
19use std::collections::{BTreeMap, HashMap};
20use std::ops::Deref;
21use std::sync::Arc;
22use std::{fmt::Display, time};
23use thiserror::Error;
24use tokio::io::AsyncBufReadExt;
25use tokio_util::io::StreamReader;
26
27use super::mcp;
28
29pub struct OpenAI {
30 alias: String,
31 endpoint: String,
32 api_key: Option<String>,
33 timeout: Option<time::Duration>,
34 mcp: Option<Arc<dyn mcp::McpClient>>,
35
36 want_models: Vec<String>,
37 model_settings: HashMap<String, ModelSetting>,
38
39 max_output_tokens: Option<usize>,
40}
41
42#[async_trait]
43impl Backend for OpenAI {
44 fn name(&self) -> &str {
45 &self.alias
46 }
47
48 async fn list_models(&self) -> Result<Vec<Model>> {
49 let mut req = reqwest::Client::new()
50 .get(format!("{}/v1/models", self.endpoint))
51 .header("User-Agent", user_agent());
52
53 if let Some(timeout) = self.timeout {
54 req = req.timeout(timeout);
55 }
56
57 if let Some(token) = &self.api_key {
58 req = req.bearer_auth(token);
59 }
60
61 let res = req.send().await.wrap_err("listing models")?;
62
63 if !res.status().is_success() {
64 let http_code = res.status().as_u16();
65 let err: ErrorResponse = res.json().await.wrap_err("parsing error response")?;
66 let mut err = err.error;
67 err.http_code = http_code;
68 return Err(err.into());
69 }
70
71 let res = res
72 .json::<ModelListResponse>()
73 .await
74 .wrap_err("parsing model list response")?;
75
76 let all = self.want_models.is_empty();
77
78 let mut models = res
79 .data
80 .into_iter()
81 .filter(|m| all || self.want_models.contains(&m.id))
82 .map(|m| Model::new(m.id).with_provider(&self.alias))
83 .collect::<Vec<_>>();
84
85 models.sort_by(|a, b| a.id().cmp(b.id()));
86
87 Ok(models)
88 }
89
90 async fn get_completion(&self, prompt: BackendPrompt, event_tx: ArcEventTx) -> Result<()> {
91 if prompt.model().is_empty() {
92 bail!("no model is set");
93 }
94
95 let init_conversation = prompt.context().is_empty();
96 let content = if init_conversation && !prompt.no_generate_title() {
97 format!("{}\n{}", prompt.text(), TITLE_PROMPT)
98 } else {
99 prompt.text().to_string()
100 };
101
102 let mut messages = prompt.context().to_vec();
103 messages.push(Message::new_user("user", content));
104
105 if let Some(max_output_tokens) = self.max_output_tokens {
106 context_truncation(&mut messages, max_output_tokens);
107 }
108
109 let messages = messages
110 .into_iter()
111 .map(|m| MessageRequest::from(&m))
112 .collect::<Vec<_>>();
113
114 self.chat_completion(None, init_conversation, prompt.model(), &messages, event_tx)
115 .await?;
116 Ok(())
117 }
118}
119
120impl From<OpenAI> for ArcBackend {
121 fn from(value: OpenAI) -> Self {
122 Arc::new(value)
123 }
124}
125
126impl From<&BackendConnection> for OpenAI {
127 fn from(value: &BackendConnection) -> Self {
128 let mut openai = OpenAI::default().with_endpoint(value.endpoint());
129
130 if let Some(api_key) = value.api_key() {
131 openai.api_key = Some(api_key.to_string());
132 }
133
134 if let Some(timeout) = value.timeout() {
135 openai.timeout = Some(timeout);
136 }
137
138 if let Some(alias) = value.alias() {
139 openai.alias = alias.to_string();
140 }
141
142 openai.max_output_tokens = value.max_output_tokens();
143
144 openai.want_models = value.models().to_vec();
145 openai
146 }
147}
148
149impl OpenAI {
150 pub fn new() -> Self {
151 Self::default()
152 }
153
154 pub async fn init(&mut self) -> Result<()> {
155 let models = self.list_models().await.wrap_err("listing models")?;
156 for settings in &config::instance().backend.model_settings {
157 let re = settings.model.build().wrap_err("building model filter")?;
158 if let Some(model) = models.iter().find(|m| re.is_match(m.id())) {
159 self.model_settings
160 .insert(model.id().to_string(), settings.clone());
161 }
162 }
163 Ok(())
164 }
165
166 pub fn with_want_models(mut self, models: Vec<String>) -> Self {
167 self.want_models = models;
168 self
169 }
170
171 pub fn with_endpoint(mut self, endpoint: &str) -> Self {
172 self.endpoint = endpoint.to_string();
173 self
174 }
175
176 pub fn with_api_key(mut self, api_key: &str) -> Self {
177 self.api_key = Some(api_key.to_string());
178 self
179 }
180
181 pub fn with_timeout(mut self, timeout: time::Duration) -> Self {
182 self.timeout = Some(timeout);
183 self
184 }
185
186 pub fn endpoint(&self) -> &str {
187 &self.endpoint
188 }
189
190 pub fn api_key(&self) -> Option<&str> {
191 self.api_key.as_deref()
192 }
193
194 pub fn timeout(&self) -> Option<time::Duration> {
195 self.timeout
196 }
197
198 pub fn want_models(&self) -> &[String] {
199 &self.want_models
200 }
201
202 pub fn with_mcp(mut self, mcp: Arc<dyn mcp::McpClient>) -> Self {
203 self.mcp = Some(mcp);
204 self
205 }
206
207 async fn get_mcp_tools(&self, event_tx: ArcEventTx) -> Vec<Tool> {
208 if let Some(mcp) = &self.mcp {
209 let tools = match mcp.list_tools().await {
210 Ok(tools) => tools,
211 Err(e) => {
212 let _ = event_tx
213 .send(warn_event!(format!("Unable to list tools: {}", e)))
214 .await;
215 return vec![];
216 }
217 };
218 return tools;
219 }
220 vec![]
221 }
222
223 async fn chat_completion(
224 &self,
225 override_id: Option<String>,
226 init_conversation: bool,
227 model: &str,
228 messages: &[MessageRequest],
229 event_tx: ArcEventTx,
230 ) -> Result<()> {
231 let settings = self.model_settings.get(model);
232
233 let enable_mcp = if let Some(settings) = settings {
234 settings.enable_mcp.unwrap_or(true)
235 } else {
236 true
237 };
238
239 let tools = if enable_mcp {
240 self.get_mcp_tools(event_tx.clone()).await
241 } else {
242 vec![]
243 };
244
245 let completion_req = CompletionRequest {
246 model: model.to_string(),
247 messages: messages.to_vec(),
248 stream: true,
249 max_completion_tokens: self.max_output_tokens,
250 tool_choice: if !tools.is_empty() {
251 Some("auto".to_string())
252 } else {
253 None
254 },
255 tools: tools.iter().map(ToolRequest::from).collect(),
256 };
257
258 let mut req = reqwest::Client::new()
259 .post(format!("{}/v1/chat/completions", self.endpoint))
260 .header("Content-Type", "application/json")
261 .header("User-Agent", user_agent());
262
263 if let Some(timeout) = self.timeout {
264 req = req.timeout(timeout);
265 }
266
267 if let Some(token) = &self.api_key {
268 req = req.bearer_auth(token);
269 }
270
271 log::trace!("Sending completion request: {:?}", completion_req);
272
273 let res = req
274 .json(&completion_req)
275 .send()
276 .await
277 .wrap_err("sending completion request")?;
278
279 if !res.status().is_success() {
280 let http_code = res.status().as_u16();
281 let resp = res.text().await.wrap_err("parsing error response")?;
282 log::error!("Error response: {}", resp);
283 let err = serde_json::from_str::<ErrorResponse>(&resp)
284 .wrap_err(format!("parsing error response: {}", resp))?;
285 let mut err = err.error;
286 err.http_code = http_code;
287 return Err(err.into());
288 }
289
290 let stream = res.bytes_stream().map_err(|e| {
291 let err_msg = e.to_string();
292 std::io::Error::new(std::io::ErrorKind::Interrupted, err_msg)
293 });
294
295 let mut line_readers = StreamReader::new(stream).lines();
296
297 let mut message_id = override_id.unwrap_or_default();
298 let mut usage: Option<BackendUsage> = None;
299
300 let mut call_tools: BTreeMap<usize, ToolCallResponse> = BTreeMap::new();
301
302 let mut current_message = MessageRequest {
303 role: "assistant".to_string(),
304 content: String::new(),
305 tool_call_id: None,
306 ..Default::default()
307 };
308
309 while let Ok(line) = line_readers.next_line().await {
310 if line.is_none() {
311 break;
312 }
313
314 let mut line = line.unwrap().trim().to_string();
315 log::trace!("streaming response: {}", line);
316 if !line.starts_with("data: ") {
317 continue;
318 }
319
320 line = line[6..].to_string();
321 if line == "[DONE]" {
322 break;
323 }
324
325 let data = serde_json::from_str::<CompletionResponse>(&line)
326 .wrap_err(format!("parsing completion response line: {}", line))?;
327
328 let c = match data.choices.first() {
329 Some(c) => c,
330 None => continue,
331 };
332
333 if message_id.is_empty() {
334 message_id = data.id;
335 }
336
337 c.delta.tool_calls.iter().for_each(|e| {
338 if let Some(tool) = call_tools.get_mut(&e.index) {
339 tool.function
340 .arguments
341 .as_mut()
342 .unwrap()
343 .push_str(e.function.arguments.as_deref().unwrap_or_default());
344 return;
345 }
346 call_tools.insert(e.index, e.clone());
347 });
348
349 let text = match c.delta.content {
350 Some(ref text) => text.deref().to_string(),
351 None => continue,
352 };
353
354 current_message.content.push_str(&text);
355
356 event_tx
357 .send(Event::ChatCompletionResponse(
358 BackendResponse::new(&message_id, model)
359 .with_text(&text)
360 .with_init_conversation(init_conversation),
361 ))
362 .await?;
363
364 if call_tools.is_empty() {
365 if let Some(usage_data) = data.usage {
366 usage = Some(BackendUsage {
367 prompt_tokens: usage_data.prompt_tokens,
368 completion_tokens: usage_data.completion_tokens,
369 total_tokens: usage_data.total_tokens,
370 });
371 }
372 }
373 }
374
375 if call_tools.is_empty() {
376 let mut msg = BackendResponse::new(&message_id, model)
377 .with_done()
378 .with_init_conversation(init_conversation);
379 if let Some(usage) = usage {
380 msg = msg.with_usage(usage);
381 }
382 event_tx.send(Event::ChatCompletionResponse(msg)).await?;
383 return Ok(());
384 }
385
386 event_tx
387 .send(Event::ChatCompletionResponse(
388 BackendResponse::new(&message_id, model)
389 .with_text("\n")
390 .with_init_conversation(init_conversation),
391 ))
392 .await?;
393
394 let call_tools = call_tools.into_values().collect::<Vec<_>>();
395 let tool_call_messages = self
398 .call_tool(&call_tools, &tools, event_tx.clone())
399 .await
400 .wrap_err("calling tools")?;
401 let mut messages = messages.to_vec();
402 current_message.tool_calls = call_tools;
403 messages.push(current_message);
404 messages.extend(tool_call_messages);
405
406 Box::pin(self.chat_completion(
407 Some(message_id),
408 init_conversation,
409 model,
410 &messages,
411 event_tx,
412 ))
413 .await?;
414 Ok(())
415 }
416
417 async fn call_tool(
418 &self,
419 calls: &[ToolCallResponse],
420 tools: &[Tool],
421 event_tx: ArcEventTx,
422 ) -> Result<Vec<MessageRequest>> {
423 if self.mcp.is_none() {
424 bail!("MCP is not set");
425 }
426 let mut results = vec![];
427
428 let notice_on_call = config::instance()
429 .backend
430 .mcp
431 .notice_on_call_tool
432 .unwrap_or_default();
433
434 for call in calls {
435 let args = match call.function.arguments.as_ref() {
436 Some(args) => Some(
437 serde_json::from_str::<Value>(args).wrap_err("parsing tool call arguments")?,
438 ),
439 _ => None,
440 };
441
442 let tool_name = call
443 .function
444 .name
445 .as_ref()
446 .ok_or_else(|| eyre::eyre!("missing tool name"))?;
447 if notice_on_call {
448 let provider = match tools.iter().find(|t| &t.name == tool_name) {
449 Some(t) => t.provider.clone(),
450 None => "unknown".to_string(),
451 };
452
453 event_tx
454 .send(info_event!(format!(
455 "Calling tool \"{}\" (provider: {})",
456 tool_name, provider
457 )))
458 .await?;
459 }
460
461 log::debug!("Calling tool {} with args: {:?}", tool_name, args);
463
464 let resp = self
465 .mcp
466 .as_ref()
467 .unwrap()
468 .clone()
469 .call_tool(call.function.name.as_ref().unwrap(), args)
470 .await
471 .wrap_err("calling tool")?;
472 let result =
473 serde_json::to_string(&resp.content).wrap_err("serializing tool result")?;
474 results.push(MessageRequest {
475 role: "tool".to_string(),
476 content: result,
477 tool_call_id: call.id.clone(),
478 ..Default::default()
479 });
480 }
481 Ok(results)
482 }
483}
484
485impl Default for OpenAI {
486 fn default() -> Self {
487 Self {
488 max_output_tokens: None,
489 alias: "OpenAI".to_string(),
490 endpoint: "https://api.openai.com".to_string(),
491 api_key: None,
492 timeout: None,
493 want_models: vec![],
494 mcp: None,
495 model_settings: HashMap::new(),
496 }
497 }
498}
499
500#[derive(Default, Debug, Serialize, Deserialize)]
501struct ModelResponse {
502 id: String,
503}
504
505#[derive(Default, Debug, Serialize, Deserialize)]
506struct ModelListResponse {
507 data: Vec<ModelResponse>,
508}
509
510#[derive(Default, Debug, Clone, Serialize, Deserialize)]
511struct MessageRequest {
512 role: String,
513 content: String,
514 #[serde(skip_serializing_if = "Option::is_none")]
515 tool_call_id: Option<String>,
516 #[serde(skip_serializing_if = "Vec::is_empty", default)]
517 tool_calls: Vec<ToolCallResponse>,
518}
519
520#[derive(Default, Debug, Serialize, Deserialize)]
521struct CompletionRequest {
522 model: String,
523 messages: Vec<MessageRequest>,
524 stream: bool,
525 #[serde(skip_serializing_if = "Option::is_none")]
526 max_completion_tokens: Option<usize>,
527 #[serde(skip_serializing_if = "Option::is_none")]
528 tool_choice: Option<String>,
529 #[serde(skip_serializing_if = "Vec::is_empty")]
530 tools: Vec<ToolRequest>,
531}
532
533#[derive(Default, Debug, Serialize, Deserialize)]
534struct ToolRequest {
535 #[serde(rename = "type")]
536 tool_type: String,
537 function: FunctionRequest,
538}
539
540#[derive(Default, Debug, Serialize, Deserialize)]
541struct FunctionRequest {
542 name: String,
543 #[serde(skip_serializing_if = "Option::is_none")]
544 description: Option<String>,
545 parameters: ToolInputSchema,
546}
547
548#[derive(Default, Debug, Serialize, Deserialize)]
549struct CompletionDeltaResponse {
550 content: Option<String>,
551 #[serde(skip_serializing_if = "Vec::is_empty", default)]
552 tool_calls: Vec<ToolCallResponse>,
553}
554
555#[derive(Default, Debug, Serialize, Deserialize)]
556struct CompletionChoiceResponse {
557 delta: CompletionDeltaResponse,
558 finish_reason: Option<String>,
559}
560
561#[derive(Default, Debug, Serialize, Deserialize, Clone)]
562struct ToolCallResponse {
563 index: usize,
564 #[serde(skip_serializing_if = "Option::is_none")]
565 id: Option<String>,
566 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
567 tool_type: Option<String>,
568 function: FunctionResponse,
569}
570
571#[derive(Default, Debug, Serialize, Deserialize, Clone)]
572struct FunctionResponse {
573 #[serde(skip_serializing_if = "Option::is_none")]
574 name: Option<String>,
575 #[serde(skip_serializing_if = "Option::is_none")]
576 arguments: Option<String>,
577}
578
579#[derive(Default, Debug, Serialize, Deserialize)]
580struct CompletionResponse {
581 id: String,
582 choices: Vec<CompletionChoiceResponse>,
583 usage: Option<CompletionUsageResponse>,
584}
585
586#[derive(Default, Debug, Serialize, Deserialize)]
587struct CompletionUsageResponse {
588 prompt_tokens: usize,
589 completion_tokens: usize,
590 total_tokens: usize,
591}
592
593#[derive(Default, Debug, Serialize, Deserialize)]
594struct ErrorResponse {
595 error: OpenAIError,
596}
597
598#[derive(Default, Error, Debug, Serialize, Deserialize)]
599pub struct OpenAIError {
600 #[serde(skip)]
601 pub http_code: u16,
602 pub message: String,
603 #[serde(rename = "type")]
604 pub err_type: String,
605 pub param: Option<String>,
606 pub code: Option<String>,
607}
608
609impl Display for OpenAIError {
610 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
611 write!(f, "OpenAI error ({}): {}", self.http_code, self.message)
612 }
613}
614
615impl From<&Message> for MessageRequest {
616 fn from(msg: &Message) -> Self {
617 Self {
618 role: if msg.is_context() {
619 "system".to_string()
620 } else if msg.is_system() {
621 "assistant".to_string()
622 } else {
623 "user".to_string()
624 },
625 content: msg.text().to_string(),
626 tool_call_id: None,
627 tool_calls: vec![],
628 }
629 }
630}
631
632impl From<&Tool> for ToolRequest {
633 fn from(tool: &Tool) -> Self {
634 Self {
635 tool_type: "function".to_string(),
636 function: FunctionRequest {
637 name: tool.name.clone(),
638 description: tool.description.clone(),
639 parameters: tool.input_schema.clone(),
640 },
641 }
642 }
643}