1use 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 max_tokens: 16384,
63 retry: RetryPolicy::default(),
64 stream_tx: None,
65 }
66 }
67
68 pub fn with_stream_tx(mut self, tx: StreamSender) -> Self {
70 self.stream_tx = Some(tx);
71 self
72 }
73
74 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 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 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 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 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 let reader = resp.text().await?;
346 for line in reader.lines() {
347 if let Some(data) = line.strip_prefix("data: ") {
348 if data.trim() == "[DONE]" {
350 break;
351 }
352
353 let chunk: Value = serde_json::from_str(data)
355 .map_err(|e| self.make_err(format!("SSE parse error: {e}; data: {data}")))?;
356
357 if let Some(choices) = chunk.get("choices").and_then(|c| c.as_array()) {
359 if let Some(choice) = choices.first() {
360 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 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 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 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 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 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
751 let addr = listener.local_addr().unwrap();
752 drop(listener);
753
754 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 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
773 let addr = listener.local_addr().unwrap();
774
775 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 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 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 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 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(); assert_eq!(
876 policy.backoff_for(0, Some(503), false),
877 Some(Duration::from_secs(1))
878 );
879 assert_eq!(
881 policy.backoff_for(1, Some(500), false),
882 Some(Duration::from_secs(2))
883 );
884 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 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 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 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 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 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 let body_str = captured.lock().unwrap().clone();
977 let body: serde_json::Value = serde_json::from_str(&body_str).unwrap();
978
979 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 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 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}