Skip to main content

recursive/llm/
openai.rs

1//! OpenAI-compatible chat-completions adapter.
2//!
3//! Targets the `/chat/completions` shape that OpenAI, Azure (via gateway),
4//! GLM (Zhipu), DeepSeek, Moonshot, Together, Ollama and many others speak.
5//! The only thing that varies is the base URL + model name + API key, which
6//! is all driven by config.
7
8use async_trait::async_trait;
9use reqwest::Client;
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12use std::time::Duration;
13
14use super::{Completion, LlmProvider, ToolCall, ToolSpec};
15use crate::error::{Error, Result};
16use crate::message::{Message, Role};
17
18/// Retry policy for transient failures (network timeouts, 5xx errors).
19#[derive(Debug, Clone)]
20pub struct RetryPolicy {
21    pub max_retries: usize,
22    pub initial_backoff: Duration,
23    pub max_backoff: Duration,
24}
25
26impl Default for RetryPolicy {
27    fn default() -> Self {
28        Self {
29            max_retries: 2,
30            initial_backoff: Duration::from_secs(1),
31            max_backoff: Duration::from_secs(8),
32        }
33    }
34}
35
36impl RetryPolicy {
37    /// Decide whether the caller should wait and try again.
38    /// `attempt` is 0-indexed (0 = the first retry decision after the
39    /// initial try has failed). Returns `Some(backoff)` to retry,
40    /// `None` to give up and propagate the error.
41    pub fn backoff_for(
42        &self,
43        attempt: usize,
44        status: Option<u16>,
45        is_network_error: bool,
46    ) -> Option<Duration> {
47        if attempt >= self.max_retries {
48            return None;
49        }
50
51        let is_transient = is_network_error || status.is_some_and(|s| (500..600).contains(&s));
52
53        if !is_transient {
54            return None;
55        }
56
57        let backoff = self.initial_backoff * 2u32.pow(attempt as u32);
58        Some(backoff.min(self.max_backoff))
59    }
60}
61
62#[derive(Debug, Clone)]
63pub struct OpenAiProvider {
64    base_url: String,
65    api_key: String,
66    model: String,
67    client: Client,
68    temperature: f64,
69    retry: RetryPolicy,
70}
71
72impl OpenAiProvider {
73    pub fn new(
74        base_url: impl Into<String>,
75        api_key: impl Into<String>,
76        model: impl Into<String>,
77    ) -> Self {
78        Self {
79            base_url: base_url.into().trim_end_matches('/').to_string(),
80            api_key: api_key.into(),
81            model: model.into(),
82            client: Client::builder()
83                .timeout(Duration::from_secs(180))
84                .build()
85                .expect("reqwest client build"),
86            temperature: 0.2,
87            retry: RetryPolicy::default(),
88        }
89    }
90
91    pub fn with_temperature(mut self, t: f64) -> Self {
92        self.temperature = t;
93        self
94    }
95
96    pub fn with_retry_policy(mut self, policy: RetryPolicy) -> Self {
97        self.retry = policy;
98        self
99    }
100}
101
102#[async_trait]
103impl LlmProvider for OpenAiProvider {
104    async fn complete(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Completion> {
105        let body = build_request(&self.model, self.temperature, messages, tools);
106        let url = format!("{}/chat/completions", self.base_url);
107
108        let mut attempt = 0;
109        loop {
110            tracing::debug!(target: "recursive::llm", request = %body, "POST {}", url);
111            let result = self
112                .client
113                .post(&url)
114                .bearer_auth(&self.api_key)
115                .json(&body)
116                .send()
117                .await;
118
119            match result {
120                Ok(resp) => {
121                    let status = resp.status();
122                    let is_network_error = false;
123
124                    if status.is_success() {
125                        let text = resp.text().await?;
126                        let parsed: ChatResponse = serde_json::from_str(&text).map_err(|e| {
127                            Error::Llm(format!("failed to parse response: {e}; body: {text}"))
128                        })?;
129                        let choice = parsed
130                            .choices
131                            .into_iter()
132                            .next()
133                            .ok_or_else(|| Error::Llm("response had no choices".into()))?;
134                        return Ok(parse_completion(choice));
135                    }
136
137                    // Non-2xx response: check if it's transient (5xx)
138                    let text = resp.text().await?;
139                    tracing::debug!(target: "recursive::llm", body = %text, "error response");
140
141                    if let Some(backoff) =
142                        self.retry
143                            .backoff_for(attempt, Some(status.as_u16()), is_network_error)
144                    {
145                        tracing::warn!(
146                            target: "recursive::llm",
147                            attempt,
148                            backoff_ms = backoff.as_millis(),
149                            status = status.as_u16(),
150                            "transient HTTP error, retrying"
151                        );
152                        tokio::time::sleep(backoff).await;
153                        attempt += 1;
154                        continue;
155                    }
156
157                    // Non-transient (4xx or other)
158                    return Err(Error::Llm(format!("HTTP {}: {}", status, text)));
159                }
160                Err(e) => {
161                    // Network error
162                    if let Some(backoff) = self.retry.backoff_for(attempt, None, true) {
163                        tracing::warn!(
164                            target: "recursive::llm",
165                            attempt,
166                            backoff_ms = backoff.as_millis(),
167                            error = %e,
168                            "network error, retrying"
169                        );
170                        tokio::time::sleep(backoff).await;
171                        attempt += 1;
172                        continue;
173                    }
174
175                    return Err(Error::Llm(format!("request failed: {e}")));
176                }
177            }
178        }
179    }
180}
181
182fn build_request(model: &str, temperature: f64, messages: &[Message], tools: &[ToolSpec]) -> Value {
183    let mut req = serde_json::json!({
184        "model": model,
185        "temperature": temperature,
186        "messages": messages.iter().map(serialize_message).collect::<Vec<_>>(),
187    });
188    if !tools.is_empty() {
189        let tools_json: Vec<Value> = tools
190            .iter()
191            .map(|t| {
192                serde_json::json!({
193                    "type": "function",
194                    "function": {
195                        "name": t.name,
196                        "description": t.description,
197                        "parameters": t.parameters,
198                    }
199                })
200            })
201            .collect();
202        req["tools"] = Value::Array(tools_json);
203        req["tool_choice"] = Value::String("auto".into());
204    }
205    req
206}
207
208fn serialize_message(m: &Message) -> Value {
209    let role = match m.role {
210        Role::System => "system",
211        Role::User => "user",
212        Role::Assistant => "assistant",
213        Role::Tool => "tool",
214    };
215    let mut obj = serde_json::Map::new();
216    obj.insert("role".into(), Value::String(role.into()));
217    obj.insert("content".into(), Value::String(m.content.clone()));
218    if let Some(id) = &m.tool_call_id {
219        obj.insert("tool_call_id".into(), Value::String(id.clone()));
220    }
221    if !m.tool_calls.is_empty() {
222        let calls: Vec<Value> = m
223            .tool_calls
224            .iter()
225            .map(|c| {
226                serde_json::json!({
227                    "id": c.id,
228                    "type": "function",
229                    "function": {
230                        "name": c.name,
231                        "arguments": serde_json::to_string(&c.arguments).unwrap_or_else(|_| "{}".into()),
232                    }
233                })
234            })
235            .collect();
236        obj.insert("tool_calls".into(), Value::Array(calls));
237    }
238    Value::Object(obj)
239}
240
241#[derive(Debug, Deserialize)]
242struct ChatResponse {
243    choices: Vec<ChatChoice>,
244}
245
246#[derive(Debug, Deserialize)]
247struct ChatChoice {
248    message: ChatChoiceMessage,
249    #[serde(default)]
250    finish_reason: Option<String>,
251}
252
253#[derive(Debug, Deserialize)]
254struct ChatChoiceMessage {
255    #[serde(default)]
256    content: Option<String>,
257    #[serde(default)]
258    tool_calls: Vec<RawToolCall>,
259}
260
261#[derive(Debug, Deserialize, Serialize)]
262struct RawToolCall {
263    id: String,
264    #[serde(default)]
265    function: RawFunction,
266}
267
268#[derive(Debug, Default, Deserialize, Serialize)]
269struct RawFunction {
270    #[serde(default)]
271    name: String,
272    #[serde(default)]
273    arguments: String,
274}
275
276fn parse_completion(choice: ChatChoice) -> Completion {
277    let content = choice.message.content.unwrap_or_default();
278    let tool_calls = choice
279        .message
280        .tool_calls
281        .into_iter()
282        .map(|c| {
283            let args: Value = if c.function.arguments.trim().is_empty() {
284                Value::Object(Default::default())
285            } else {
286                serde_json::from_str(&c.function.arguments)
287                    .unwrap_or_else(|_| Value::String(c.function.arguments.clone()))
288            };
289            ToolCall {
290                id: c.id,
291                name: c.function.name,
292                arguments: args,
293            }
294        })
295        .collect();
296    Completion {
297        content,
298        tool_calls,
299        finish_reason: choice.finish_reason,
300    }
301}
302
303#[cfg(test)]
304mod tests {
305    use super::*;
306
307    #[test]
308    fn parses_plain_text_choice() {
309        let raw = r#"{"choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}]}"#;
310        let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
311        let c = parse_completion(parsed.choices.into_iter().next().unwrap());
312        assert_eq!(c.content, "hi");
313        assert!(c.tool_calls.is_empty());
314        assert_eq!(c.finish_reason.as_deref(), Some("stop"));
315    }
316
317    #[test]
318    fn parses_tool_call_choice() {
319        let raw = r#"{
320            "choices":[{
321                "message":{
322                    "role":"assistant",
323                    "content":null,
324                    "tool_calls":[{
325                        "id":"call_1",
326                        "type":"function",
327                        "function":{"name":"read_file","arguments":"{\"path\":\"src/lib.rs\"}"}
328                    }]
329                },
330                "finish_reason":"tool_calls"
331            }]
332        }"#;
333        let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
334        let c = parse_completion(parsed.choices.into_iter().next().unwrap());
335        assert_eq!(c.tool_calls.len(), 1);
336        assert_eq!(c.tool_calls[0].name, "read_file");
337        assert_eq!(c.tool_calls[0].arguments["path"], "src/lib.rs");
338    }
339
340    #[test]
341    fn serialises_assistant_with_tool_calls() {
342        let msg = Message::assistant_with_tool_calls(
343            "",
344            vec![ToolCall {
345                id: "abc".into(),
346                name: "write_file".into(),
347                arguments: serde_json::json!({"path":"a","contents":"b"}),
348            }],
349        );
350        let v = serialize_message(&msg);
351        assert_eq!(v["tool_calls"][0]["function"]["name"], "write_file");
352        let args = v["tool_calls"][0]["function"]["arguments"]
353            .as_str()
354            .unwrap();
355        let decoded: Value = serde_json::from_str(args).unwrap();
356        assert_eq!(decoded["contents"], "b");
357    }
358
359    #[test]
360    fn builds_request_without_tools_omits_field() {
361        let req = build_request("m", 0.2, &[Message::user("hi")], &[]);
362        assert!(req.get("tools").is_none());
363        assert_eq!(req["messages"][0]["role"], "user");
364    }
365
366    #[test]
367    fn policy_retries_5xx_with_exponential_backoff() {
368        let policy = RetryPolicy::default(); // max_retries = 2
369                                             // Attempt 0: should return initial_backoff (1s)
370        assert_eq!(
371            policy.backoff_for(0, Some(503), false),
372            Some(Duration::from_secs(1))
373        );
374        // Attempt 1: should return 2s (1s * 2^1)
375        assert_eq!(
376            policy.backoff_for(1, Some(500), false),
377            Some(Duration::from_secs(2))
378        );
379        // Attempt 2: should return None (exceeds max_retries=2)
380        assert_eq!(policy.backoff_for(2, Some(500), false), None);
381    }
382
383    #[test]
384    fn policy_retries_network_errors() {
385        let policy = RetryPolicy::default();
386        // Network error at attempt 0 should return initial_backoff
387        assert_eq!(
388            policy.backoff_for(0, None, true),
389            Some(Duration::from_secs(1))
390        );
391    }
392
393    #[test]
394    fn policy_does_not_retry_4xx() {
395        let policy = RetryPolicy::default();
396        // 4xx errors should not be retried
397        assert_eq!(policy.backoff_for(0, Some(400), false), None);
398        assert_eq!(policy.backoff_for(0, Some(401), false), None);
399        assert_eq!(policy.backoff_for(0, Some(404), false), None);
400        assert_eq!(policy.backoff_for(0, Some(429), false), None);
401    }
402
403    #[test]
404    fn policy_caps_backoff_at_max() {
405        let policy = RetryPolicy {
406            max_retries: 10,
407            initial_backoff: Duration::from_secs(1),
408            max_backoff: Duration::from_secs(3),
409        };
410        // At attempt 5, exponential backoff would be 32s but capped to 3s
411        assert_eq!(
412            policy.backoff_for(5, Some(500), false),
413            Some(Duration::from_secs(3))
414        );
415    }
416}