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//!
8//! ## Deferred tool loading
9//!
10//! This provider does **not** support `Tool::should_defer()` or
11//! `complete_with_search`. The default `complete_with_search` impl
12//! in `LlmProvider` concatenates the eager and deferred lists and
13//! calls `complete()` — so the model sees every tool on every turn.
14//! To enable deferred loading, use the Anthropic provider.
15
16use async_trait::async_trait;
17use reqwest::Client;
18use serde::{Deserialize, Serialize};
19use serde_json::Value;
20use std::time::Duration;
21
22use super::StructuredRequest;
23use super::{Completion, LlmProvider, RetryPolicy, StreamSender, TokenUsage, ToolCall, ToolSpec};
24use crate::error::{Error, Result};
25use crate::message::{Message, Role};
26
27#[derive(Debug, Clone)]
28pub struct OpenAiProvider {
29    base_url: String,
30    api_key: String,
31    model: String,
32    client: Client,
33    temperature: f64,
34    max_tokens: u32,
35    retry: RetryPolicy,
36    stream_tx: Option<StreamSender>,
37}
38
39impl OpenAiProvider {
40    pub fn new(
41        base_url: impl Into<String>,
42        api_key: impl Into<String>,
43        model: impl Into<String>,
44    ) -> Self {
45        Self {
46            base_url: base_url.into().trim_end_matches('/').to_string(),
47            api_key: api_key.into(),
48            model: model.into(),
49            client: Client::builder()
50                .timeout(Duration::from_secs(180))
51                .build()
52                .expect("reqwest client build"),
53            temperature: 0.2,
54            // DeepSeek defaults to a per-response cap of 4096 tokens; any
55            // tool call whose `arguments` string holds more than that — e.g.
56            // a `write_file` with a multi-kilobyte `contents` field — gets
57            // truncated server-side and arrives as malformed JSON. 16384 is
58            // both within DeepSeek's hard ceiling (8192 for v3, 32K-64K for
59            // newer models) and big enough for whole-file writes. Callers
60            // can override with `with_max_tokens` if their provider supports
61            // more or needs less.
62            max_tokens: 16384,
63            retry: RetryPolicy::default(),
64            stream_tx: None,
65        }
66    }
67
68    /// Enable streaming by providing a channel sender for partial tokens.
69    pub fn with_stream_tx(mut self, tx: StreamSender) -> Self {
70        self.stream_tx = Some(tx);
71        self
72    }
73
74    /// Build an `Error::Llm` with the model name prefixed.
75    fn make_err(&self, ctx: impl Into<String>) -> Error {
76        Error::Llm {
77            provider: self.model.clone(),
78            message: ctx.into(),
79        }
80    }
81
82    pub fn with_temperature(mut self, t: f64) -> Self {
83        self.temperature = t;
84        self
85    }
86
87    pub fn with_max_tokens(mut self, n: u32) -> Self {
88        self.max_tokens = n;
89        self
90    }
91
92    pub fn with_retry_policy(mut self, policy: RetryPolicy) -> Self {
93        self.retry = policy;
94        self
95    }
96
97    /// POST `body` to `url` with retry, returning the raw response text on success.
98    ///
99    /// Handles: network errors, 5xx transient errors, and HTTP 200 with empty body
100    /// (MiniMax transient failure). All three are retried with exponential back-off
101    /// according to `self.retry`. Non-transient 4xx errors are returned immediately.
102    async fn post_json_with_retry(&self, url: &str, body: &Value, label: &str) -> Result<String> {
103        let mut attempt = 0;
104        loop {
105            tracing::debug!(target: "recursive::llm", request = %body, "POST {url} ({label})");
106            let result = self
107                .client
108                .post(url)
109                .bearer_auth(&self.api_key)
110                .json(body)
111                .send()
112                .await;
113
114            match result {
115                Ok(resp) => {
116                    let status = resp.status();
117                    if status.is_success() {
118                        let text = resp.text().await?;
119                        if text.trim().is_empty() {
120                            if let Some(backoff) = self.retry.backoff_for(attempt, None, true) {
121                                tracing::warn!(
122                                    target: "recursive::llm",
123                                    attempt,
124                                    backoff_ms = backoff.as_millis(),
125                                    "HTTP 200 but empty body ({label}), retrying"
126                                );
127                                tokio::time::sleep(backoff).await;
128                                attempt += 1;
129                                continue;
130                            }
131                            return Err(self.make_err("HTTP 200 but response body is empty"));
132                        }
133                        return Ok(text);
134                    }
135
136                    let text = resp.text().await?;
137                    tracing::debug!(target: "recursive::llm", body = %text, "error response ({label})");
138
139                    if let Some(backoff) =
140                        self.retry
141                            .backoff_for(attempt, Some(status.as_u16()), false)
142                    {
143                        tracing::warn!(
144                            target: "recursive::llm",
145                            attempt,
146                            backoff_ms = backoff.as_millis(),
147                            status = status.as_u16(),
148                            "transient HTTP error, retrying ({label})"
149                        );
150                        tokio::time::sleep(backoff).await;
151                        attempt += 1;
152                        continue;
153                    }
154
155                    return Err(self.make_err(format!("HTTP {status}: {text}")));
156                }
157                Err(e) => {
158                    if let Some(backoff) = self.retry.backoff_for(attempt, None, true) {
159                        tracing::warn!(
160                            target: "recursive::llm",
161                            attempt,
162                            backoff_ms = backoff.as_millis(),
163                            error = %e,
164                            "network error, retrying ({label})"
165                        );
166                        tokio::time::sleep(backoff).await;
167                        attempt += 1;
168                        continue;
169                    }
170                    return Err(self.make_err(format!("request failed: {e}")));
171                }
172            }
173        }
174    }
175}
176
177#[async_trait]
178impl LlmProvider for OpenAiProvider {
179    #[tracing::instrument(skip(self, messages, tools), fields(
180        provider = %self.base_url.split('/').next_back().unwrap_or("unknown"),
181        model = %self.model
182    ))]
183    async fn complete(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Completion> {
184        let body = build_request(
185            &self.model,
186            self.temperature,
187            self.max_tokens,
188            messages,
189            tools,
190        );
191        let url = format!("{}/chat/completions", self.base_url);
192        let text = self.post_json_with_retry(&url, &body, "complete").await?;
193        let parsed: ChatResponse = serde_json::from_str(&text)
194            .map_err(|e| self.make_err(format!("failed to parse response: {e}; body: {text}")))?;
195        let choice = parsed
196            .choices
197            .into_iter()
198            .next()
199            .ok_or_else(|| self.make_err("response had no choices"))?;
200        Ok(parse_completion(choice, parsed.usage))
201    }
202
203    async fn complete_structured(&self, req: StructuredRequest) -> Result<Value> {
204        let mut body = build_request(
205            &self.model,
206            self.temperature,
207            self.max_tokens,
208            &req.messages,
209            &[],
210        );
211        body["response_format"] = serde_json::json!({
212            "type": "json_schema",
213            "json_schema": {
214                "name": req.schema_name,
215                "strict": true,
216                "schema": req.schema,
217            }
218        });
219        let url = format!("{}/chat/completions", self.base_url);
220        let text = self.post_json_with_retry(&url, &body, "structured").await?;
221        let parsed: ChatResponse = serde_json::from_str(&text)
222            .map_err(|e| self.make_err(format!("failed to parse response: {e}; body: {text}")))?;
223        let choice = parsed
224            .choices
225            .into_iter()
226            .next()
227            .ok_or_else(|| self.make_err("response had no choices"))?;
228        let completion = parse_completion(choice, parsed.usage);
229        if completion.content.trim().is_empty() {
230            return Err(self.make_err("structured response had empty content"));
231        }
232        serde_json::from_str(&completion.content).map_err(|e| {
233            self.make_err(format!(
234                "failed to parse structured response as JSON: {e}; content: {}",
235                completion.content
236            ))
237        })
238    }
239
240    async fn stream(
241        &self,
242        messages: &[Message],
243        tools: &[ToolSpec],
244        stream_tx: Option<StreamSender>,
245    ) -> Result<Completion> {
246        // If no stream_tx provided, fall back to the instance-level one
247        let tx = stream_tx.or_else(|| self.stream_tx.clone());
248        self.stream_inner(messages, tools, tx).await
249    }
250}
251
252impl OpenAiProvider {
253    async fn stream_inner(
254        &self,
255        messages: &[Message],
256        tools: &[ToolSpec],
257        stream_tx: Option<StreamSender>,
258    ) -> Result<Completion> {
259        let mut body = build_request(
260            &self.model,
261            self.temperature,
262            self.max_tokens,
263            messages,
264            tools,
265        );
266        body["stream"] = Value::Bool(true);
267        let url = format!("{}/chat/completions", self.base_url);
268
269        // Stream requests need the raw Response object for SSE parsing, so we
270        // can't use post_json_with_retry (which returns text). Retry only on
271        // non-2xx and network errors; a successful 2xx hands off to parse_sse_stream.
272        let mut attempt = 0;
273        loop {
274            tracing::debug!(target: "recursive::llm", request = %body, "POST {url} (stream)");
275            let result = self
276                .client
277                .post(&url)
278                .bearer_auth(&self.api_key)
279                .json(&body)
280                .send()
281                .await;
282
283            match result {
284                Ok(resp) => {
285                    let status = resp.status();
286                    if status.is_success() {
287                        return self.parse_sse_stream(resp, stream_tx.clone()).await;
288                    }
289                    let text = resp.text().await?;
290                    tracing::debug!(target: "recursive::llm", body = %text, "error response (stream)");
291                    if let Some(backoff) =
292                        self.retry
293                            .backoff_for(attempt, Some(status.as_u16()), false)
294                    {
295                        tracing::warn!(
296                            target: "recursive::llm",
297                            attempt,
298                            backoff_ms = backoff.as_millis(),
299                            status = status.as_u16(),
300                            "transient HTTP error, retrying (stream)"
301                        );
302                        tokio::time::sleep(backoff).await;
303                        attempt += 1;
304                        continue;
305                    }
306                    return Err(self.make_err(format!("HTTP {status}: {text}")));
307                }
308                Err(e) => {
309                    if let Some(backoff) = self.retry.backoff_for(attempt, None, true) {
310                        tracing::warn!(
311                            target: "recursive::llm",
312                            attempt,
313                            backoff_ms = backoff.as_millis(),
314                            error = %e,
315                            "network error, retrying (stream)"
316                        );
317                        tokio::time::sleep(backoff).await;
318                        attempt += 1;
319                        continue;
320                    }
321                    return Err(self.make_err(format!("request failed: {e}")));
322                }
323            }
324        }
325    }
326
327    /// Parse an SSE stream from a successful HTTP response.
328    ///
329    /// Reads `data: {...}\n\n` chunks line-by-line, extracts
330    /// `choices[0].delta.content` deltas, accumulates them, and emits
331    /// each delta through `stream_tx` if configured. Returns the final
332    /// `Completion` matching the non-streaming shape.
333    async fn parse_sse_stream(
334        &self,
335        resp: reqwest::Response,
336        stream_tx: Option<StreamSender>,
337    ) -> Result<Completion> {
338        let mut content = String::new();
339        let mut reasoning_content = String::new();
340        let tool_calls: Vec<ToolCall> = Vec::new();
341        let mut finish_reason: Option<String> = None;
342        let mut usage: Option<TokenUsage> = None;
343
344        // Read the byte stream line by line
345        let reader = resp.text().await?;
346        for line in reader.lines() {
347            if let Some(data) = line.strip_prefix("data: ") {
348                // Skip the final "[DONE]" marker
349                if data.trim() == "[DONE]" {
350                    break;
351                }
352
353                // Parse the JSON chunk
354                let chunk: Value = serde_json::from_str(data)
355                    .map_err(|e| self.make_err(format!("SSE parse error: {e}; data: {data}")))?;
356
357                // Extract delta content
358                if let Some(choices) = chunk.get("choices").and_then(|c| c.as_array()) {
359                    if let Some(choice) = choices.first() {
360                        // Delta content
361                        if let Some(delta) = choice.get("delta") {
362                            if let Some(delta_content) =
363                                delta.get("content").and_then(|c| c.as_str())
364                            {
365                                if !delta_content.is_empty() {
366                                    content.push_str(delta_content);
367                                    if let Some(ref tx) = stream_tx {
368                                        let _ = tx.send(delta_content.to_string());
369                                    }
370                                }
371                            }
372                            // Accumulate reasoning_content deltas (DeepSeek thinking mode)
373                            if let Some(delta_reasoning) =
374                                delta.get("reasoning_content").and_then(|c| c.as_str())
375                            {
376                                if !delta_reasoning.is_empty() {
377                                    reasoning_content.push_str(delta_reasoning);
378                                }
379                            }
380                        }
381
382                        // Finish reason (only on the last chunk)
383                        if let Some(fr) = choice.get("finish_reason").and_then(|f| f.as_str()) {
384                            if !fr.is_empty() {
385                                finish_reason = Some(fr.to_string());
386                            }
387                        }
388                    }
389                }
390
391                // Usage (only on the last chunk for some providers)
392                if let Some(u) = chunk.get("usage") {
393                    let prompt =
394                        u.get("prompt_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32;
395                    let completion = u
396                        .get("completion_tokens")
397                        .and_then(|v| v.as_u64())
398                        .unwrap_or(0) as u32;
399                    let total = u.get("total_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32;
400                    let cache_hit = u
401                        .get("prompt_cache_hit_tokens")
402                        .and_then(|v| v.as_u64())
403                        .unwrap_or(0) as u32;
404                    let cache_miss = u
405                        .get("prompt_cache_miss_tokens")
406                        .and_then(|v| v.as_u64())
407                        .unwrap_or(0) as u32;
408                    usage = Some(TokenUsage {
409                        prompt_tokens: prompt,
410                        completion_tokens: completion,
411                        total_tokens: total,
412                        cache_hit_tokens: cache_hit,
413                        cache_miss_tokens: cache_miss,
414                    });
415                }
416            }
417        }
418
419        Ok(Completion {
420            content,
421            tool_calls,
422            finish_reason,
423            usage,
424            reasoning_content: if reasoning_content.is_empty() {
425                None
426            } else {
427                Some(reasoning_content)
428            },
429        })
430    }
431}
432
433fn build_request(
434    model: &str,
435    temperature: f64,
436    max_tokens: u32,
437    messages: &[Message],
438    tools: &[ToolSpec],
439) -> Value {
440    let mut req = serde_json::json!({
441        "model": model,
442        "temperature": temperature,
443        "max_tokens": max_tokens,
444        "messages": messages.iter().map(serialize_message).collect::<Vec<_>>(),
445    });
446    if !tools.is_empty() {
447        let tools_json: Vec<Value> = tools
448            .iter()
449            .map(|t| {
450                serde_json::json!({
451                    "type": "function",
452                    "function": {
453                        "name": t.name,
454                        "description": t.description,
455                        "parameters": t.parameters,
456                    }
457                })
458            })
459            .collect();
460        req["tools"] = Value::Array(tools_json);
461        req["tool_choice"] = Value::String("auto".into());
462    }
463    req
464}
465
466fn serialize_message(m: &Message) -> Value {
467    let role = match m.role {
468        Role::System => "system",
469        Role::User => "user",
470        Role::Assistant => "assistant",
471        Role::Tool => "tool",
472    };
473    let mut obj = serde_json::Map::new();
474    obj.insert("role".into(), Value::String(role.into()));
475    obj.insert("content".into(), Value::String(m.content.clone()));
476    if let Some(id) = &m.tool_call_id {
477        obj.insert("tool_call_id".into(), Value::String(id.clone()));
478    }
479    // Echo reasoning_content back to the API (required by DeepSeek thinking mode)
480    if let Some(ref reasoning) = m.reasoning_content {
481        obj.insert("reasoning_content".into(), Value::String(reasoning.clone()));
482    }
483    if !m.tool_calls.is_empty() {
484        let calls: Vec<Value> = m
485            .tool_calls
486            .iter()
487            .map(|c| {
488                serde_json::json!({
489                    "id": c.id,
490                    "type": "function",
491                    "function": {
492                        "name": c.name,
493                        "arguments": serde_json::to_string(&c.arguments).unwrap_or_else(|_| "{}".into()),
494                    }
495                })
496            })
497            .collect();
498        obj.insert("tool_calls".into(), Value::Array(calls));
499    }
500    Value::Object(obj)
501}
502
503#[derive(Debug, Deserialize)]
504struct ChatResponse {
505    choices: Vec<ChatChoice>,
506    #[serde(default)]
507    usage: Option<ResponseUsage>,
508}
509
510#[derive(Debug, Deserialize)]
511struct ResponseUsage {
512    #[serde(default)]
513    prompt_tokens: Option<u32>,
514    #[serde(default)]
515    completion_tokens: Option<u32>,
516    #[serde(default)]
517    total_tokens: Option<u32>,
518    #[serde(default)]
519    prompt_cache_hit_tokens: Option<u32>,
520    #[serde(default)]
521    prompt_cache_miss_tokens: Option<u32>,
522}
523
524impl ResponseUsage {
525    fn to_token_usage(&self) -> TokenUsage {
526        TokenUsage {
527            prompt_tokens: self.prompt_tokens.unwrap_or(0),
528            completion_tokens: self.completion_tokens.unwrap_or(0),
529            total_tokens: self.total_tokens.unwrap_or(0),
530            cache_hit_tokens: self.prompt_cache_hit_tokens.unwrap_or(0),
531            cache_miss_tokens: self.prompt_cache_miss_tokens.unwrap_or(0),
532        }
533    }
534}
535
536#[derive(Debug, Deserialize)]
537struct ChatChoice {
538    message: ChatChoiceMessage,
539    #[serde(default)]
540    finish_reason: Option<String>,
541}
542
543#[derive(Debug, Deserialize)]
544struct ChatChoiceMessage {
545    #[serde(default)]
546    content: Option<String>,
547    #[serde(default)]
548    reasoning_content: Option<String>,
549    #[serde(default)]
550    tool_calls: Vec<RawToolCall>,
551}
552
553#[derive(Debug, Deserialize, Serialize)]
554struct RawToolCall {
555    id: String,
556    #[serde(default)]
557    function: RawFunction,
558}
559
560#[derive(Debug, Default, Deserialize, Serialize)]
561struct RawFunction {
562    #[serde(default)]
563    name: String,
564    #[serde(default)]
565    arguments: String,
566}
567
568fn parse_completion(choice: ChatChoice, usage: Option<ResponseUsage>) -> Completion {
569    let content = choice.message.content.unwrap_or_default();
570    let reasoning_content = choice.message.reasoning_content;
571    let tool_calls = choice
572        .message
573        .tool_calls
574        .into_iter()
575        .map(|c| {
576            let args: Value = if c.function.arguments.trim().is_empty() {
577                Value::Object(Default::default())
578            } else {
579                serde_json::from_str(&c.function.arguments)
580                    .unwrap_or_else(|_| Value::String(c.function.arguments.clone()))
581            };
582            ToolCall {
583                id: c.id,
584                name: c.function.name,
585                arguments: args,
586            }
587        })
588        .collect();
589    Completion {
590        content,
591        tool_calls,
592        finish_reason: choice.finish_reason,
593        usage: usage.map(|u| u.to_token_usage()),
594        reasoning_content,
595    }
596}
597
598#[cfg(test)]
599mod tests {
600    use super::*;
601
602    #[test]
603    fn parses_plain_text_choice() {
604        let raw = r#"{"choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}]}"#;
605        let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
606        let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
607        assert_eq!(c.content, "hi");
608        assert!(c.tool_calls.is_empty());
609        assert_eq!(c.finish_reason.as_deref(), Some("stop"));
610        assert!(c.usage.is_none());
611    }
612
613    #[test]
614    fn parses_tool_call_choice() {
615        let raw = r#"{
616            "choices":[{
617                "message":{
618                    "role":"assistant",
619                    "content":null,
620                    "tool_calls":[{
621                        "id":"call_1",
622                        "type":"function",
623                        "function":{"name":"read_file","arguments":"{\"path\":\"src/lib.rs\"}"}
624                    }]
625                },
626                "finish_reason":"tool_calls"
627            }]
628        }"#;
629        let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
630        let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
631        assert_eq!(c.tool_calls.len(), 1);
632        assert_eq!(c.tool_calls[0].name, "read_file");
633        assert_eq!(c.tool_calls[0].arguments["path"], "src/lib.rs");
634    }
635
636    #[test]
637    fn parses_usage_from_response() {
638        let raw = r#"{
639            "choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],
640            "usage":{"prompt_tokens":41,"completion_tokens":7,"total_tokens":48}
641        }"#;
642        let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
643        let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
644        assert!(c.usage.is_some());
645        let u = c.usage.unwrap();
646        assert_eq!(u.prompt_tokens, 41);
647        assert_eq!(u.completion_tokens, 7);
648        assert_eq!(u.total_tokens, 48);
649    }
650
651    #[test]
652    fn parses_missing_usage_as_none() {
653        let raw = r#"{"choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}]}"#;
654        let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
655        let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
656        assert!(c.usage.is_none());
657    }
658
659    #[test]
660    fn parses_partial_usage_fills_zeros() {
661        let raw = r#"{
662            "choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],
663            "usage":{"total_tokens":50}
664        }"#;
665        let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
666        let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
667        assert!(c.usage.is_some());
668        let u = c.usage.unwrap();
669        assert_eq!(u.prompt_tokens, 0);
670        assert_eq!(u.completion_tokens, 0);
671        assert_eq!(u.total_tokens, 50);
672        assert_eq!(u.cache_hit_tokens, 0);
673        assert_eq!(u.cache_miss_tokens, 0);
674    }
675
676    #[test]
677    fn parses_cache_fields_from_deepseek_usage() {
678        let raw = r#"{
679            "choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],
680            "usage":{
681                "prompt_tokens":100,
682                "completion_tokens":50,
683                "total_tokens":150,
684                "prompt_cache_hit_tokens":60,
685                "prompt_cache_miss_tokens":40
686            }
687        }"#;
688        let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
689        let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
690        assert!(c.usage.is_some());
691        let u = c.usage.unwrap();
692        assert_eq!(u.prompt_tokens, 100);
693        assert_eq!(u.completion_tokens, 50);
694        assert_eq!(u.total_tokens, 150);
695        assert_eq!(u.cache_hit_tokens, 60);
696        assert_eq!(u.cache_miss_tokens, 40);
697    }
698
699    #[test]
700    fn parses_cache_fields_as_zero_when_absent() {
701        let raw = r#"{
702            "choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],
703            "usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}
704        }"#;
705        let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
706        let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
707        assert!(c.usage.is_some());
708        let u = c.usage.unwrap();
709        assert_eq!(u.cache_hit_tokens, 0);
710        assert_eq!(u.cache_miss_tokens, 0);
711    }
712
713    #[test]
714    fn serialises_assistant_with_tool_calls() {
715        let msg = Message::assistant_with_tool_calls(
716            "",
717            vec![ToolCall {
718                id: "abc".into(),
719                name: "write_file".into(),
720                arguments: serde_json::json!({"path":"a","contents":"b"}),
721            }],
722        );
723        let v = serialize_message(&msg);
724        assert_eq!(v["tool_calls"][0]["function"]["name"], "write_file");
725        let args = v["tool_calls"][0]["function"]["arguments"]
726            .as_str()
727            .unwrap();
728        let decoded: Value = serde_json::from_str(args).unwrap();
729        assert_eq!(decoded["contents"], "b");
730    }
731
732    #[test]
733    fn builds_request_without_tools_omits_field() {
734        let req = build_request("m", 0.2, 16384, &[Message::user("hi")], &[]);
735        assert!(req.get("tools").is_none());
736        assert_eq!(req["messages"][0]["role"], "user");
737        assert_eq!(req["max_tokens"], 16384);
738    }
739
740    #[test]
741    fn builds_request_includes_max_tokens() {
742        let req = build_request("m", 0.2, 1024, &[Message::user("hi")], &[]);
743        assert_eq!(req["max_tokens"], 1024);
744    }
745
746    #[tokio::test]
747    async fn error_includes_model_name_on_network_failure() {
748        // Bind a listener on an ephemeral port, then drop it so the port is freed.
749        // The next connect attempt gets ECONNREFUSED (a network error).
750        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
751        let addr = listener.local_addr().unwrap();
752        drop(listener);
753
754        // Small sleep so the OS releases the port.
755        tokio::time::sleep(std::time::Duration::from_millis(20)).await;
756
757        let provider = OpenAiProvider::new(format!("http://{addr}"), "sk-noop", "test-model");
758        let err = provider
759            .complete(&[Message::user("hi")], &[])
760            .await
761            .unwrap_err();
762        let msg = err.to_string();
763        assert!(
764            msg.contains("test-model"),
765            "error should contain model name: {msg}"
766        );
767    }
768
769    #[tokio::test]
770    async fn error_includes_model_name_and_status_on_http_error() {
771        // Spawn a one-shot TCP listener that sends a 400 response.
772        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
773        let addr = listener.local_addr().unwrap();
774
775        // Use a dedicated thread with a blocking server to avoid tokio issues.
776        std::thread::spawn(move || {
777            let (mut stream, _) = listener.accept().unwrap();
778            use std::io::{Read, Write};
779            let mut buf = [0u8; 4096];
780            let _ = stream.read(&mut buf);
781            write!(
782                stream,
783                "HTTP/1.1 400 Bad Request\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{{}}"
784            )
785            .unwrap();
786            stream.flush().unwrap();
787        });
788
789        // Give the server a moment to start
790        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
791
792        let provider = OpenAiProvider::new(format!("http://{addr}"), "sk-noop", "test-model-http");
793        let err = provider
794            .complete(&[Message::user("hi")], &[])
795            .await
796            .unwrap_err();
797        let msg = err.to_string();
798        assert!(
799            msg.contains("test-model-http"),
800            "error should contain model name: {msg}"
801        );
802        assert!(
803            msg.contains("400"),
804            "error should contain HTTP status: {msg}"
805        );
806    }
807
808    #[tokio::test]
809    async fn stream_concatenates_sse_chunks() {
810        // Spawn a one-shot TCP server that serves a canned SSE response
811        // with 3 chunks. Assert the returned Completion's content is the
812        // concatenation of the deltas.
813        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
814        let addr = listener.local_addr().unwrap();
815
816        std::thread::spawn(move || {
817            let (mut stream, _) = listener.accept().unwrap();
818            use std::io::{Read, Write};
819            let mut buf = [0u8; 4096];
820            let _ = stream.read(&mut buf);
821            let body = "\
822data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n\
823data: {\"choices\":[{\"delta\":{\"content\":\" \"},\"finish_reason\":null}]}\n\n\
824data: {\"choices\":[{\"delta\":{\"content\":\"World\"},\"finish_reason\":\"stop\"}]}\n\n\
825data: [DONE]\n\n";
826            write!(
827                stream,
828                "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
829                body.len(),
830                body,
831            )
832            .unwrap();
833            stream.flush().unwrap();
834        });
835
836        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
837
838        let provider =
839            OpenAiProvider::new(format!("http://{addr}"), "sk-noop", "test-stream-model");
840        let completion = provider
841            .stream(&[Message::user("hi")], &[], None)
842            .await
843            .unwrap();
844        assert_eq!(completion.content, "Hello World");
845        assert_eq!(completion.finish_reason.as_deref(), Some("stop"));
846    }
847
848    #[tokio::test]
849    async fn stream_fallback_delegates_to_complete() {
850        // MockProvider doesn't override stream, so it falls back to complete.
851        // Verify the fallback path works by using a MockProvider.
852        use crate::llm::MockProvider;
853        let provider = MockProvider::new(vec![super::Completion {
854            content: "fallback works".to_string(),
855            tool_calls: vec![],
856            finish_reason: Some("stop".to_string()),
857            usage: None,
858            reasoning_content: None,
859        }]);
860        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
861        let completion = provider
862            .stream(&[Message::user("hi")], &[], Some(tx))
863            .await
864            .unwrap();
865        assert_eq!(completion.content, "fallback works");
866        // Should have received the full content as a single delta
867        let delta = rx.recv().await.unwrap();
868        assert_eq!(delta, "fallback works");
869    }
870
871    #[test]
872    fn policy_retries_5xx_with_exponential_backoff() {
873        let policy = RetryPolicy::default(); // max_retries = 2
874                                             // Attempt 0: should return initial_backoff (1s)
875        assert_eq!(
876            policy.backoff_for(0, Some(503), false),
877            Some(Duration::from_secs(1))
878        );
879        // Attempt 1: should return 2s (1s * 2^1)
880        assert_eq!(
881            policy.backoff_for(1, Some(500), false),
882            Some(Duration::from_secs(2))
883        );
884        // Attempt 2: should return None (exceeds max_retries=2)
885        assert_eq!(policy.backoff_for(2, Some(500), false), None);
886    }
887
888    #[test]
889    fn policy_retries_network_errors() {
890        let policy = RetryPolicy::default();
891        // Network error at attempt 0 should return initial_backoff
892        assert_eq!(
893            policy.backoff_for(0, None, true),
894            Some(Duration::from_secs(1))
895        );
896    }
897
898    #[test]
899    fn policy_does_not_retry_4xx() {
900        let policy = RetryPolicy::default();
901        // 4xx errors should not be retried
902        assert_eq!(policy.backoff_for(0, Some(400), false), None);
903        assert_eq!(policy.backoff_for(0, Some(401), false), None);
904        assert_eq!(policy.backoff_for(0, Some(404), false), None);
905        assert_eq!(policy.backoff_for(0, Some(429), false), None);
906    }
907
908    #[test]
909    fn policy_caps_backoff_at_max() {
910        let policy = RetryPolicy {
911            max_retries: 10,
912            initial_backoff: Duration::from_secs(1),
913            max_backoff: Duration::from_secs(3),
914        };
915        // At attempt 5, exponential backoff would be 32s but capped to 3s
916        assert_eq!(
917            policy.backoff_for(5, Some(500), false),
918            Some(Duration::from_secs(3))
919        );
920    }
921
922    #[tokio::test]
923    async fn openai_structured_includes_schema_in_request_body() {
924        // Spawn a mock server that captures the request body and returns
925        // a valid JSON response. Assert the request body contains the
926        // response_format block with the schema.
927        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
928        let addr = listener.local_addr().unwrap();
929        let captured = std::sync::Arc::new(std::sync::Mutex::new(String::new()));
930        let captured_clone = captured.clone();
931
932        std::thread::spawn(move || {
933            let (mut stream, _) = listener.accept().unwrap();
934            use std::io::{Read, Write};
935            let mut buf = [0u8; 8192];
936            let n = stream.read(&mut buf).unwrap();
937            let request = String::from_utf8_lossy(&buf[..n]).to_string();
938            // Extract the JSON body (after the blank line)
939            if let Some(body_start) = request.find("\r\n\r\n") {
940                let body = request[body_start + 4..].trim().to_string();
941                *captured_clone.lock().unwrap() = body.clone();
942            }
943            let response_body = r#"{"choices":[{"message":{"role":"assistant","content":"{\"answer\":42}"},"finish_reason":"stop"}],"usage":null}"#;
944            write!(
945                stream,
946                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
947                response_body.len(),
948                response_body,
949            )
950            .unwrap();
951            stream.flush().unwrap();
952        });
953
954        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
955
956        let provider =
957            OpenAiProvider::new(format!("http://{addr}"), "sk-noop", "test-structured-model");
958
959        let schema = serde_json::json!({
960            "type": "object",
961            "properties": {
962                "answer": {"type": "integer"}
963            },
964            "required": ["answer"]
965        });
966
967        let req = StructuredRequest {
968            messages: vec![Message::user("what is 6 times 7?".to_string())],
969            schema: schema.clone(),
970            schema_name: "math_answer".to_string(),
971        };
972
973        let _ = provider.complete_structured(req).await;
974
975        // Check the captured request body
976        let body_str = captured.lock().unwrap().clone();
977        let body: serde_json::Value = serde_json::from_str(&body_str).unwrap();
978
979        // Assert response_format is present with the right shape
980        let rf = body.get("response_format").unwrap();
981        assert_eq!(rf["type"], "json_schema");
982        assert_eq!(rf["json_schema"]["name"], "math_answer");
983        assert_eq!(rf["json_schema"]["strict"], true);
984        assert_eq!(rf["json_schema"]["schema"], schema);
985    }
986
987    #[tokio::test]
988    async fn openai_structured_parses_response_json() {
989        // Spawn a mock server that returns a known JSON response.
990        // Assert the parsed value matches.
991        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
992        let addr = listener.local_addr().unwrap();
993
994        std::thread::spawn(move || {
995            let (mut stream, _) = listener.accept().unwrap();
996            use std::io::{Read, Write};
997            let mut buf = [0u8; 8192];
998            let _ = stream.read(&mut buf);
999            // Return a JSON object with summary and kept_facts
1000            let response_body = r#"{"choices":[{"message":{"role":"assistant","content":"{\"summary\":\"test\",\"kept_facts\":[\"a\",\"b\"]}"},"finish_reason":"stop"}],"usage":null}"#;
1001            write!(
1002                stream,
1003                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1004                response_body.len(),
1005                response_body,
1006            )
1007            .unwrap();
1008            stream.flush().unwrap();
1009        });
1010
1011        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1012
1013        let provider =
1014            OpenAiProvider::new(format!("http://{addr}"), "sk-noop", "test-structured-model");
1015
1016        let schema = serde_json::json!({
1017            "type": "object",
1018            "properties": {
1019                "summary": {"type": "string"},
1020                "kept_facts": {"type": "array", "items": {"type": "string"}}
1021            },
1022            "required": ["summary", "kept_facts"]
1023        });
1024
1025        let req = StructuredRequest {
1026            messages: vec![Message::user("summarize".to_string())],
1027            schema,
1028            schema_name: "summary".to_string(),
1029        };
1030
1031        let result = provider.complete_structured(req).await;
1032        assert!(result.is_ok(), "got error: {:?}", result.err());
1033        let value = result.unwrap();
1034        assert_eq!(value["summary"], "test");
1035        assert_eq!(value["kept_facts"][0], "a");
1036        assert_eq!(value["kept_facts"][1], "b");
1037    }
1038}