llm_universal/
openai.rs

1use async_trait::async_trait;
2use serde_json::{json, Value};
3
4use crate::{
5    send_json, ChatRequest, ChatResponse, LlmError, LlmProvider, Message, Part, Result, Role,
6    StopReason, Usage,
7};
8
9/// Client for the OpenAI Chat Completions API (`POST {base_url}/chat/completions`).
10///
11/// The same wire format is served by many other providers; use the presets
12/// ([`groq`](Self::groq), [`mistral`](Self::mistral), [`deepseek`](Self::deepseek), [`xai`](Self::xai),
13/// [`openrouter`](Self::openrouter), [`together`](Self::together), [`fireworks`](Self::fireworks),
14/// [`perplexity`](Self::perplexity), [`azure`](Self::azure), [`ollama`](Self::ollama),
15/// [`lm_studio`](Self::lm_studio)) or [`with_base_url`](Self::with_base_url) for anything else (vLLM, LiteLLM, …).
16#[derive(Clone)]
17pub struct OpenAiProvider {
18    client: reqwest::Client,
19    api_key: String,
20    base_url: String,
21    auth: Auth,
22}
23
24/// How the key is sent.
25#[derive(Clone, Copy, Debug)]
26enum Auth {
27    /// `Authorization: Bearer <key>` (OpenAI and most compatible servers).
28    Bearer,
29    /// `api-key: <key>` (Azure OpenAI).
30    ApiKeyHeader,
31}
32
33impl std::fmt::Debug for OpenAiProvider {
34    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
35        f.debug_struct("OpenAiProvider")
36            .field("api_key", &"<redacted>")
37            .field("base_url", &self.base_url)
38            .finish()
39    }
40}
41
42impl OpenAiProvider {
43    pub fn new(api_key: impl Into<String>) -> Self {
44        Self {
45            client: crate::http_client(),
46            api_key: api_key.into(),
47            base_url: "https://api.openai.com/v1".to_string(),
48            auth: Auth::Bearer,
49        }
50    }
51
52    fn preset(api_key: impl Into<String>, base_url: &str) -> Self {
53        Self::new(api_key).with_base_url(base_url)
54    }
55
56    /// Groq: `https://api.groq.com/openai/v1`.
57    pub fn groq(api_key: impl Into<String>) -> Self {
58        Self::preset(api_key, "https://api.groq.com/openai/v1")
59    }
60
61    /// Mistral AI: `https://api.mistral.ai/v1`.
62    pub fn mistral(api_key: impl Into<String>) -> Self {
63        Self::preset(api_key, "https://api.mistral.ai/v1")
64    }
65
66    /// DeepSeek: `https://api.deepseek.com/v1`.
67    pub fn deepseek(api_key: impl Into<String>) -> Self {
68        Self::preset(api_key, "https://api.deepseek.com/v1")
69    }
70
71    /// xAI (Grok): `https://api.x.ai/v1`.
72    pub fn xai(api_key: impl Into<String>) -> Self {
73        Self::preset(api_key, "https://api.x.ai/v1")
74    }
75
76    /// OpenRouter: `https://openrouter.ai/api/v1`.
77    pub fn openrouter(api_key: impl Into<String>) -> Self {
78        Self::preset(api_key, "https://openrouter.ai/api/v1")
79    }
80
81    /// Together AI: `https://api.together.xyz/v1`.
82    pub fn together(api_key: impl Into<String>) -> Self {
83        Self::preset(api_key, "https://api.together.xyz/v1")
84    }
85
86    /// Fireworks AI: `https://api.fireworks.ai/inference/v1`.
87    pub fn fireworks(api_key: impl Into<String>) -> Self {
88        Self::preset(api_key, "https://api.fireworks.ai/inference/v1")
89    }
90
91    /// Perplexity: `https://api.perplexity.ai`.
92    pub fn perplexity(api_key: impl Into<String>) -> Self {
93        Self::preset(api_key, "https://api.perplexity.ai")
94    }
95
96    /// Azure OpenAI v1 API: `https://{resource}.openai.azure.com/openai/v1`, key in the `api-key` header.
97    /// Pass the deployment name as the model.
98    pub fn azure(resource: &str, api_key: impl Into<String>) -> Self {
99        let mut provider = Self::preset(
100            api_key,
101            &format!("https://{resource}.openai.azure.com/openai/v1"),
102        );
103        provider.auth = Auth::ApiKeyHeader;
104        provider
105    }
106
107    /// Ollama on this machine: `http://localhost:11434/v1`, no key.
108    pub fn ollama() -> Self {
109        Self::preset("", "http://localhost:11434/v1")
110    }
111
112    /// LM Studio on this machine: `http://localhost:1234/v1`, no key.
113    pub fn lm_studio() -> Self {
114        Self::preset("", "http://localhost:1234/v1")
115    }
116
117    /// Points the client at another server. The key may be empty if the server does not need one.
118    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
119        self.base_url = base_url.into().trim_end_matches('/').to_string();
120        self
121    }
122
123    /// Uses your own `reqwest::Client`, for example to share a connection pool or set a proxy.
124    pub fn with_client(mut self, client: reqwest::Client) -> Self {
125        self.client = client;
126        self
127    }
128
129    pub(crate) fn build_body(req: &ChatRequest) -> Value {
130        let mut messages = Vec::new();
131        if let Some(system) = &req.system {
132            messages.push(json!({ "role": "system", "content": system }));
133        }
134        for message in &req.messages {
135            messages.extend(to_openai_messages(message));
136        }
137
138        let mut body = json!({ "model": req.model, "messages": messages });
139
140        if !req.tools.is_empty() {
141            let tools: Vec<Value> = req
142                .tools
143                .iter()
144                .map(|t| {
145                    json!({
146                        "type": "function",
147                        "function": {
148                            "name": t.name,
149                            "description": t.description,
150                            "parameters": t.parameters,
151                        }
152                    })
153                })
154                .collect();
155            body["tools"] = json!(tools);
156        }
157        if let Some(max_tokens) = req.max_tokens {
158            body["max_completion_tokens"] = json!(max_tokens);
159        }
160        if let Some(temperature) = req.temperature {
161            body["temperature"] = json!(temperature);
162        }
163        body
164    }
165}
166
167pub(crate) fn to_openai_messages(message: &Message) -> Vec<Value> {
168    if message.role == Role::Tool {
169        // OpenAI expects one `tool` message per result.
170        return message
171            .parts
172            .iter()
173            .filter_map(|p| match p {
174                Part::ToolResult {
175                    call_id, content, ..
176                } => Some(json!({
177                    "role": "tool",
178                    "tool_call_id": call_id,
179                    "content": content,
180                })),
181                _ => None,
182            })
183            .collect();
184    }
185
186    let role = match message.role {
187        Role::System => "system",
188        Role::User => "user",
189        Role::Assistant | Role::Tool => "assistant",
190    };
191    let text = message.text_content();
192    let tool_calls: Vec<Value> = message
193        .parts
194        .iter()
195        .filter_map(|p| match p {
196            Part::ToolCall {
197                id,
198                name,
199                arguments,
200            } => Some(json!({
201                "id": id,
202                "type": "function",
203                "function": { "name": name, "arguments": arguments.to_string() },
204            })),
205            _ => None,
206        })
207        .collect();
208
209    let content = if text.is_empty() && !tool_calls.is_empty() {
210        Value::Null
211    } else {
212        json!(text)
213    };
214    let mut msg = json!({ "role": role, "content": content });
215    if !tool_calls.is_empty() {
216        msg["tool_calls"] = json!(tool_calls);
217    }
218    vec![msg]
219}
220
221/// Parses an OpenAI-style `tool_calls` array (also used by Cohere).
222pub(crate) fn parse_tool_calls(calls: &Value) -> Vec<Part> {
223    let Some(calls) = calls.as_array() else {
224        return Vec::new();
225    };
226    calls
227        .iter()
228        .map(|call| {
229            // Arguments are a JSON string. An empty string means a tool without arguments;
230            // invalid JSON (e.g. cut off by max tokens) is kept raw instead of failing the call.
231            let raw = call["function"]["arguments"].as_str().unwrap_or_default();
232            let arguments = if raw.trim().is_empty() {
233                json!({})
234            } else {
235                serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.to_string()))
236            };
237            Part::ToolCall {
238                id: call["id"].as_str().unwrap_or_default().to_string(),
239                name: call["function"]["name"]
240                    .as_str()
241                    .unwrap_or_default()
242                    .to_string(),
243                arguments,
244            }
245        })
246        .collect()
247}
248
249fn parse_response(v: &Value) -> Result<ChatResponse> {
250    let choice = v["choices"]
251        .get(0)
252        .ok_or_else(|| LlmError::InvalidResponse("missing field `choices[0]`".into()))?;
253    let message = &choice["message"];
254
255    let mut parts = Vec::new();
256    if let Some(text) = message["content"].as_str() {
257        if !text.is_empty() {
258            parts.push(Part::Text {
259                text: text.to_string(),
260            });
261        }
262    }
263    // A refusal arrives in its own `refusal` field while `content` is null.
264    let refusal = message["refusal"].as_str();
265    if let Some(text) = refusal {
266        parts.push(Part::Text {
267            text: text.to_string(),
268        });
269    }
270    parts.extend(parse_tool_calls(&message["tool_calls"]));
271
272    let stop_reason = match choice["finish_reason"].as_str() {
273        _ if refusal.is_some() => StopReason::Refusal,
274        Some("stop") => StopReason::EndTurn,
275        Some("tool_calls") | Some("function_call") => StopReason::ToolUse,
276        Some("length") => StopReason::MaxTokens,
277        Some("content_filter") => StopReason::Refusal,
278        other => StopReason::Other(other.unwrap_or("unknown").to_string()),
279    };
280
281    let tokens =
282        |field: &str| u32::try_from(v["usage"][field].as_u64().unwrap_or(0)).unwrap_or(u32::MAX);
283    let usage = Usage {
284        input_tokens: tokens("prompt_tokens"),
285        output_tokens: tokens("completion_tokens"),
286    };
287
288    Ok(ChatResponse {
289        model: v["model"].as_str().unwrap_or_default().to_string(),
290        message: Message::new(Role::Assistant, parts),
291        stop_reason,
292        usage,
293    })
294}
295
296#[async_trait]
297impl LlmProvider for OpenAiProvider {
298    async fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
299        // OpenAI allows tool names up to 64 characters.
300        crate::validate(
301            "OpenAiProvider",
302            &self.api_key,
303            false,
304            &self.base_url,
305            64,
306            req,
307        )?;
308        let url = format!("{}/chat/completions", self.base_url);
309        let mut builder = self.client.post(url).json(&Self::build_body(req));
310        if !self.api_key.is_empty() {
311            builder = match self.auth {
312                Auth::Bearer => builder.header(
313                    reqwest::header::AUTHORIZATION,
314                    crate::secret_header("OpenAiProvider", &format!("Bearer {}", self.api_key))?,
315                ),
316                Auth::ApiKeyHeader => builder.header(
317                    "api-key",
318                    crate::secret_header("OpenAiProvider", &self.api_key)?,
319                ),
320            };
321        }
322        let body = send_json(builder).await?;
323        parse_response(&body)
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use super::*;
330    use crate::Tool;
331
332    #[test]
333    fn builds_tool_request() {
334        let req = ChatRequest::new("gpt-4o")
335            .system("sys")
336            .max_tokens(100)
337            .tool(Tool::new("f", "d", json!({"type": "object"})))
338            .message(Message::user("hi"));
339        let body = OpenAiProvider::build_body(&req);
340        assert_eq!(body["messages"][0]["role"], "system");
341        assert_eq!(body["messages"][1]["content"], "hi");
342        assert_eq!(body["tools"][0]["function"]["name"], "f");
343        assert_eq!(body["max_completion_tokens"], 100);
344    }
345
346    #[test]
347    fn parses_tool_call_response() {
348        let v = json!({
349            "model": "gpt-4o",
350            "choices": [{
351                "finish_reason": "tool_calls",
352                "message": {
353                    "content": null,
354                    "tool_calls": [{
355                        "id": "call_1",
356                        "type": "function",
357                        "function": { "name": "f", "arguments": "{\"x\":1}" }
358                    }]
359                }
360            }],
361            "usage": { "prompt_tokens": 5, "completion_tokens": 7 }
362        });
363        let resp = parse_response(&v).unwrap();
364        assert_eq!(resp.stop_reason, StopReason::ToolUse);
365        assert_eq!(resp.tool_calls()[0].2["x"], 1);
366        assert_eq!(resp.usage.output_tokens, 7);
367    }
368
369    #[test]
370    fn tolerates_empty_and_truncated_arguments() {
371        let v = json!({
372            "choices": [{
373                "finish_reason": "length",
374                "message": {
375                    "content": null,
376                    "tool_calls": [
377                        { "id": "a", "function": { "name": "f", "arguments": "" } },
378                        { "id": "b", "function": { "name": "f", "arguments": "{\"x\":" } }
379                    ]
380                }
381            }]
382        });
383        let resp = parse_response(&v).unwrap();
384        let calls = resp.tool_calls();
385        assert_eq!(calls[0].2, json!({}));
386        assert_eq!(calls[1].2, json!("{\"x\":"));
387        assert_eq!(resp.stop_reason, StopReason::MaxTokens);
388    }
389
390    #[test]
391    fn maps_refusal_field() {
392        let v = json!({
393            "choices": [{
394                "finish_reason": "stop",
395                "message": { "content": null, "refusal": "I can't help with that." }
396            }]
397        });
398        let resp = parse_response(&v).unwrap();
399        assert_eq!(resp.stop_reason, StopReason::Refusal);
400        assert_eq!(resp.text(), "I can't help with that.");
401    }
402
403    #[test]
404    fn debug_redacts_api_key() {
405        let debug = format!("{:?}", OpenAiProvider::new("sk-secret"));
406        assert!(!debug.contains("sk-secret"));
407    }
408}
llm_universal/openai.rs

llm_universal/
openai.rs
