use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::time::Duration;
use super::StructuredRequest;
use super::{Completion, LlmProvider, RetryPolicy, StreamSender, TokenUsage, ToolCall, ToolSpec};
use crate::error::{Error, Result};
use crate::message::{Message, Role};
#[derive(Debug, Clone)]
pub struct OpenAiProvider {
base_url: String,
api_key: String,
model: String,
client: Client,
temperature: f64,
max_tokens: u32,
retry: RetryPolicy,
stream_tx: Option<StreamSender>,
}
impl OpenAiProvider {
pub fn new(
base_url: impl Into<String>,
api_key: impl Into<String>,
model: impl Into<String>,
) -> Self {
Self {
base_url: base_url.into().trim_end_matches('/').to_string(),
api_key: api_key.into(),
model: model.into(),
client: Client::builder()
.timeout(Duration::from_secs(180))
.build()
.expect("reqwest client build"),
temperature: 0.2,
max_tokens: 16384,
retry: RetryPolicy::default(),
stream_tx: None,
}
}
pub fn with_stream_tx(mut self, tx: StreamSender) -> Self {
self.stream_tx = Some(tx);
self
}
fn make_err(&self, ctx: impl Into<String>) -> Error {
Error::Llm {
provider: self.model.clone(),
message: ctx.into(),
}
}
pub fn with_temperature(mut self, t: f64) -> Self {
self.temperature = t;
self
}
pub fn with_max_tokens(mut self, n: u32) -> Self {
self.max_tokens = n;
self
}
pub fn with_retry_policy(mut self, policy: RetryPolicy) -> Self {
self.retry = policy;
self
}
async fn post_json_with_retry(&self, url: &str, body: &Value, label: &str) -> Result<String> {
let mut attempt = 0;
loop {
tracing::debug!(target: "recursive::llm", request = %body, "POST {url} ({label})");
let result = self
.client
.post(url)
.bearer_auth(&self.api_key)
.json(body)
.send()
.await;
match result {
Ok(resp) => {
let status = resp.status();
if status.is_success() {
let text = resp.text().await?;
if text.trim().is_empty() {
if let Some(backoff) = self.retry.backoff_for(attempt, None, true) {
tracing::warn!(
target: "recursive::llm",
attempt,
backoff_ms = backoff.as_millis(),
"HTTP 200 but empty body ({label}), retrying"
);
tokio::time::sleep(backoff).await;
attempt += 1;
continue;
}
return Err(self.make_err("HTTP 200 but response body is empty"));
}
return Ok(text);
}
let text = resp.text().await?;
tracing::debug!(target: "recursive::llm", body = %text, "error response ({label})");
if let Some(backoff) =
self.retry
.backoff_for(attempt, Some(status.as_u16()), false)
{
tracing::warn!(
target: "recursive::llm",
attempt,
backoff_ms = backoff.as_millis(),
status = status.as_u16(),
"transient HTTP error, retrying ({label})"
);
tokio::time::sleep(backoff).await;
attempt += 1;
continue;
}
return Err(self.make_err(format!("HTTP {status}: {text}")));
}
Err(e) => {
if let Some(backoff) = self.retry.backoff_for(attempt, None, true) {
tracing::warn!(
target: "recursive::llm",
attempt,
backoff_ms = backoff.as_millis(),
error = %e,
"network error, retrying ({label})"
);
tokio::time::sleep(backoff).await;
attempt += 1;
continue;
}
return Err(self.make_err(format!("request failed: {e}")));
}
}
}
}
}
#[async_trait]
impl LlmProvider for OpenAiProvider {
#[tracing::instrument(skip(self, messages, tools), fields(
provider = %self.base_url.split('/').next_back().unwrap_or("unknown"),
model = %self.model
))]
async fn complete(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Completion> {
let body = build_request(
&self.model,
self.temperature,
self.max_tokens,
messages,
tools,
);
let url = format!("{}/chat/completions", self.base_url);
let text = self.post_json_with_retry(&url, &body, "complete").await?;
let parsed: ChatResponse = serde_json::from_str(&text)
.map_err(|e| self.make_err(format!("failed to parse response: {e}; body: {text}")))?;
let choice = parsed
.choices
.into_iter()
.next()
.ok_or_else(|| self.make_err("response had no choices"))?;
Ok(parse_completion(choice, parsed.usage))
}
async fn complete_structured(&self, req: StructuredRequest) -> Result<Value> {
let mut body = build_request(
&self.model,
self.temperature,
self.max_tokens,
&req.messages,
&[],
);
body["response_format"] = serde_json::json!({
"type": "json_schema",
"json_schema": {
"name": req.schema_name,
"strict": true,
"schema": req.schema,
}
});
let url = format!("{}/chat/completions", self.base_url);
let text = self.post_json_with_retry(&url, &body, "structured").await?;
let parsed: ChatResponse = serde_json::from_str(&text)
.map_err(|e| self.make_err(format!("failed to parse response: {e}; body: {text}")))?;
let choice = parsed
.choices
.into_iter()
.next()
.ok_or_else(|| self.make_err("response had no choices"))?;
let completion = parse_completion(choice, parsed.usage);
if completion.content.trim().is_empty() {
return Err(self.make_err("structured response had empty content"));
}
serde_json::from_str(&completion.content).map_err(|e| {
self.make_err(format!(
"failed to parse structured response as JSON: {e}; content: {}",
completion.content
))
})
}
async fn stream(
&self,
messages: &[Message],
tools: &[ToolSpec],
stream_tx: Option<StreamSender>,
) -> Result<Completion> {
let tx = stream_tx.or_else(|| self.stream_tx.clone());
self.stream_inner(messages, tools, tx).await
}
}
impl OpenAiProvider {
async fn stream_inner(
&self,
messages: &[Message],
tools: &[ToolSpec],
stream_tx: Option<StreamSender>,
) -> Result<Completion> {
let mut body = build_request(
&self.model,
self.temperature,
self.max_tokens,
messages,
tools,
);
body["stream"] = Value::Bool(true);
let url = format!("{}/chat/completions", self.base_url);
let mut attempt = 0;
loop {
tracing::debug!(target: "recursive::llm", request = %body, "POST {url} (stream)");
let result = self
.client
.post(&url)
.bearer_auth(&self.api_key)
.json(&body)
.send()
.await;
match result {
Ok(resp) => {
let status = resp.status();
if status.is_success() {
return self.parse_sse_stream(resp, stream_tx.clone()).await;
}
let text = resp.text().await?;
tracing::debug!(target: "recursive::llm", body = %text, "error response (stream)");
if let Some(backoff) =
self.retry
.backoff_for(attempt, Some(status.as_u16()), false)
{
tracing::warn!(
target: "recursive::llm",
attempt,
backoff_ms = backoff.as_millis(),
status = status.as_u16(),
"transient HTTP error, retrying (stream)"
);
tokio::time::sleep(backoff).await;
attempt += 1;
continue;
}
return Err(self.make_err(format!("HTTP {status}: {text}")));
}
Err(e) => {
if let Some(backoff) = self.retry.backoff_for(attempt, None, true) {
tracing::warn!(
target: "recursive::llm",
attempt,
backoff_ms = backoff.as_millis(),
error = %e,
"network error, retrying (stream)"
);
tokio::time::sleep(backoff).await;
attempt += 1;
continue;
}
return Err(self.make_err(format!("request failed: {e}")));
}
}
}
}
async fn parse_sse_stream(
&self,
resp: reqwest::Response,
stream_tx: Option<StreamSender>,
) -> Result<Completion> {
let mut content = String::new();
let mut reasoning_content = String::new();
let tool_calls: Vec<ToolCall> = Vec::new();
let mut finish_reason: Option<String> = None;
let mut usage: Option<TokenUsage> = None;
let reader = resp.text().await?;
for line in reader.lines() {
if let Some(data) = line.strip_prefix("data: ") {
if data.trim() == "[DONE]" {
break;
}
let chunk: Value = serde_json::from_str(data)
.map_err(|e| self.make_err(format!("SSE parse error: {e}; data: {data}")))?;
if let Some(choices) = chunk.get("choices").and_then(|c| c.as_array()) {
if let Some(choice) = choices.first() {
if let Some(delta) = choice.get("delta") {
if let Some(delta_content) =
delta.get("content").and_then(|c| c.as_str())
{
if !delta_content.is_empty() {
content.push_str(delta_content);
if let Some(ref tx) = stream_tx {
let _ = tx.send(delta_content.to_string());
}
}
}
if let Some(delta_reasoning) =
delta.get("reasoning_content").and_then(|c| c.as_str())
{
if !delta_reasoning.is_empty() {
reasoning_content.push_str(delta_reasoning);
}
}
}
if let Some(fr) = choice.get("finish_reason").and_then(|f| f.as_str()) {
if !fr.is_empty() {
finish_reason = Some(fr.to_string());
}
}
}
}
if let Some(u) = chunk.get("usage") {
let prompt =
u.get("prompt_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32;
let completion = u
.get("completion_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32;
let total = u.get("total_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32;
let cache_hit = u
.get("prompt_cache_hit_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32;
let cache_miss = u
.get("prompt_cache_miss_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32;
usage = Some(TokenUsage {
prompt_tokens: prompt,
completion_tokens: completion,
total_tokens: total,
cache_hit_tokens: cache_hit,
cache_miss_tokens: cache_miss,
});
}
}
}
Ok(Completion {
content,
tool_calls,
finish_reason,
usage,
reasoning_content: if reasoning_content.is_empty() {
None
} else {
Some(reasoning_content)
},
})
}
}
fn build_request(
model: &str,
temperature: f64,
max_tokens: u32,
messages: &[Message],
tools: &[ToolSpec],
) -> Value {
let mut req = serde_json::json!({
"model": model,
"temperature": temperature,
"max_tokens": max_tokens,
"messages": messages.iter().map(serialize_message).collect::<Vec<_>>(),
});
if !tools.is_empty() {
let tools_json: Vec<Value> = tools
.iter()
.map(|t| {
serde_json::json!({
"type": "function",
"function": {
"name": t.name,
"description": t.description,
"parameters": t.parameters,
}
})
})
.collect();
req["tools"] = Value::Array(tools_json);
req["tool_choice"] = Value::String("auto".into());
}
req
}
fn serialize_message(m: &Message) -> Value {
let role = match m.role {
Role::System => "system",
Role::User => "user",
Role::Assistant => "assistant",
Role::Tool => "tool",
};
let mut obj = serde_json::Map::new();
obj.insert("role".into(), Value::String(role.into()));
obj.insert("content".into(), Value::String(m.content.clone()));
if let Some(id) = &m.tool_call_id {
obj.insert("tool_call_id".into(), Value::String(id.clone()));
}
if let Some(ref reasoning) = m.reasoning_content {
obj.insert("reasoning_content".into(), Value::String(reasoning.clone()));
}
if !m.tool_calls.is_empty() {
let calls: Vec<Value> = m
.tool_calls
.iter()
.map(|c| {
serde_json::json!({
"id": c.id,
"type": "function",
"function": {
"name": c.name,
"arguments": serde_json::to_string(&c.arguments).unwrap_or_else(|_| "{}".into()),
}
})
})
.collect();
obj.insert("tool_calls".into(), Value::Array(calls));
}
Value::Object(obj)
}
#[derive(Debug, Deserialize)]
struct ChatResponse {
choices: Vec<ChatChoice>,
#[serde(default)]
usage: Option<ResponseUsage>,
}
#[derive(Debug, Deserialize)]
struct ResponseUsage {
#[serde(default)]
prompt_tokens: Option<u32>,
#[serde(default)]
completion_tokens: Option<u32>,
#[serde(default)]
total_tokens: Option<u32>,
#[serde(default)]
prompt_cache_hit_tokens: Option<u32>,
#[serde(default)]
prompt_cache_miss_tokens: Option<u32>,
}
impl ResponseUsage {
fn to_token_usage(&self) -> TokenUsage {
TokenUsage {
prompt_tokens: self.prompt_tokens.unwrap_or(0),
completion_tokens: self.completion_tokens.unwrap_or(0),
total_tokens: self.total_tokens.unwrap_or(0),
cache_hit_tokens: self.prompt_cache_hit_tokens.unwrap_or(0),
cache_miss_tokens: self.prompt_cache_miss_tokens.unwrap_or(0),
}
}
}
#[derive(Debug, Deserialize)]
struct ChatChoice {
message: ChatChoiceMessage,
#[serde(default)]
finish_reason: Option<String>,
}
#[derive(Debug, Deserialize)]
struct ChatChoiceMessage {
#[serde(default)]
content: Option<String>,
#[serde(default)]
reasoning_content: Option<String>,
#[serde(default)]
tool_calls: Vec<RawToolCall>,
}
#[derive(Debug, Deserialize, Serialize)]
struct RawToolCall {
id: String,
#[serde(default)]
function: RawFunction,
}
#[derive(Debug, Default, Deserialize, Serialize)]
struct RawFunction {
#[serde(default)]
name: String,
#[serde(default)]
arguments: String,
}
fn parse_completion(choice: ChatChoice, usage: Option<ResponseUsage>) -> Completion {
let content = choice.message.content.unwrap_or_default();
let reasoning_content = choice.message.reasoning_content;
let tool_calls = choice
.message
.tool_calls
.into_iter()
.map(|c| {
let args: Value = if c.function.arguments.trim().is_empty() {
Value::Object(Default::default())
} else {
serde_json::from_str(&c.function.arguments)
.unwrap_or_else(|_| Value::String(c.function.arguments.clone()))
};
ToolCall {
id: c.id,
name: c.function.name,
arguments: args,
}
})
.collect();
Completion {
content,
tool_calls,
finish_reason: choice.finish_reason,
usage: usage.map(|u| u.to_token_usage()),
reasoning_content,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_plain_text_choice() {
let raw = r#"{"choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}]}"#;
let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
assert_eq!(c.content, "hi");
assert!(c.tool_calls.is_empty());
assert_eq!(c.finish_reason.as_deref(), Some("stop"));
assert!(c.usage.is_none());
}
#[test]
fn parses_tool_call_choice() {
let raw = r#"{
"choices":[{
"message":{
"role":"assistant",
"content":null,
"tool_calls":[{
"id":"call_1",
"type":"function",
"function":{"name":"read_file","arguments":"{\"path\":\"src/lib.rs\"}"}
}]
},
"finish_reason":"tool_calls"
}]
}"#;
let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
assert_eq!(c.tool_calls.len(), 1);
assert_eq!(c.tool_calls[0].name, "read_file");
assert_eq!(c.tool_calls[0].arguments["path"], "src/lib.rs");
}
#[test]
fn parses_usage_from_response() {
let raw = r#"{
"choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],
"usage":{"prompt_tokens":41,"completion_tokens":7,"total_tokens":48}
}"#;
let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
assert!(c.usage.is_some());
let u = c.usage.unwrap();
assert_eq!(u.prompt_tokens, 41);
assert_eq!(u.completion_tokens, 7);
assert_eq!(u.total_tokens, 48);
}
#[test]
fn parses_missing_usage_as_none() {
let raw = r#"{"choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}]}"#;
let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
assert!(c.usage.is_none());
}
#[test]
fn parses_partial_usage_fills_zeros() {
let raw = r#"{
"choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],
"usage":{"total_tokens":50}
}"#;
let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
assert!(c.usage.is_some());
let u = c.usage.unwrap();
assert_eq!(u.prompt_tokens, 0);
assert_eq!(u.completion_tokens, 0);
assert_eq!(u.total_tokens, 50);
assert_eq!(u.cache_hit_tokens, 0);
assert_eq!(u.cache_miss_tokens, 0);
}
#[test]
fn parses_cache_fields_from_deepseek_usage() {
let raw = r#"{
"choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],
"usage":{
"prompt_tokens":100,
"completion_tokens":50,
"total_tokens":150,
"prompt_cache_hit_tokens":60,
"prompt_cache_miss_tokens":40
}
}"#;
let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
assert!(c.usage.is_some());
let u = c.usage.unwrap();
assert_eq!(u.prompt_tokens, 100);
assert_eq!(u.completion_tokens, 50);
assert_eq!(u.total_tokens, 150);
assert_eq!(u.cache_hit_tokens, 60);
assert_eq!(u.cache_miss_tokens, 40);
}
#[test]
fn parses_cache_fields_as_zero_when_absent() {
let raw = r#"{
"choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],
"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}
}"#;
let parsed: ChatResponse = serde_json::from_str(raw).unwrap();
let c = parse_completion(parsed.choices.into_iter().next().unwrap(), parsed.usage);
assert!(c.usage.is_some());
let u = c.usage.unwrap();
assert_eq!(u.cache_hit_tokens, 0);
assert_eq!(u.cache_miss_tokens, 0);
}
#[test]
fn serialises_assistant_with_tool_calls() {
let msg = Message::assistant_with_tool_calls(
"",
vec![ToolCall {
id: "abc".into(),
name: "write_file".into(),
arguments: serde_json::json!({"path":"a","contents":"b"}),
}],
);
let v = serialize_message(&msg);
assert_eq!(v["tool_calls"][0]["function"]["name"], "write_file");
let args = v["tool_calls"][0]["function"]["arguments"]
.as_str()
.unwrap();
let decoded: Value = serde_json::from_str(args).unwrap();
assert_eq!(decoded["contents"], "b");
}
#[test]
fn builds_request_without_tools_omits_field() {
let req = build_request("m", 0.2, 16384, &[Message::user("hi")], &[]);
assert!(req.get("tools").is_none());
assert_eq!(req["messages"][0]["role"], "user");
assert_eq!(req["max_tokens"], 16384);
}
#[test]
fn builds_request_includes_max_tokens() {
let req = build_request("m", 0.2, 1024, &[Message::user("hi")], &[]);
assert_eq!(req["max_tokens"], 1024);
}
#[tokio::test]
async fn error_includes_model_name_on_network_failure() {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
drop(listener);
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
let provider = OpenAiProvider::new(format!("http://{addr}"), "sk-noop", "test-model");
let err = provider
.complete(&[Message::user("hi")], &[])
.await
.unwrap_err();
let msg = err.to_string();
assert!(
msg.contains("test-model"),
"error should contain model name: {msg}"
);
}
#[tokio::test]
async fn error_includes_model_name_and_status_on_http_error() {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
std::thread::spawn(move || {
let (mut stream, _) = listener.accept().unwrap();
use std::io::{Read, Write};
let mut buf = [0u8; 4096];
let _ = stream.read(&mut buf);
write!(
stream,
"HTTP/1.1 400 Bad Request\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{{}}"
)
.unwrap();
stream.flush().unwrap();
});
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let provider = OpenAiProvider::new(format!("http://{addr}"), "sk-noop", "test-model-http");
let err = provider
.complete(&[Message::user("hi")], &[])
.await
.unwrap_err();
let msg = err.to_string();
assert!(
msg.contains("test-model-http"),
"error should contain model name: {msg}"
);
assert!(
msg.contains("400"),
"error should contain HTTP status: {msg}"
);
}
#[tokio::test]
async fn stream_concatenates_sse_chunks() {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
std::thread::spawn(move || {
let (mut stream, _) = listener.accept().unwrap();
use std::io::{Read, Write};
let mut buf = [0u8; 4096];
let _ = stream.read(&mut buf);
let body = "\
data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n\
data: {\"choices\":[{\"delta\":{\"content\":\" \"},\"finish_reason\":null}]}\n\n\
data: {\"choices\":[{\"delta\":{\"content\":\"World\"},\"finish_reason\":\"stop\"}]}\n\n\
data: [DONE]\n\n";
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
body.len(),
body,
)
.unwrap();
stream.flush().unwrap();
});
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let provider =
OpenAiProvider::new(format!("http://{addr}"), "sk-noop", "test-stream-model");
let completion = provider
.stream(&[Message::user("hi")], &[], None)
.await
.unwrap();
assert_eq!(completion.content, "Hello World");
assert_eq!(completion.finish_reason.as_deref(), Some("stop"));
}
#[tokio::test]
async fn stream_fallback_delegates_to_complete() {
use crate::llm::MockProvider;
let provider = MockProvider::new(vec![super::Completion {
content: "fallback works".to_string(),
tool_calls: vec![],
finish_reason: Some("stop".to_string()),
usage: None,
reasoning_content: None,
}]);
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let completion = provider
.stream(&[Message::user("hi")], &[], Some(tx))
.await
.unwrap();
assert_eq!(completion.content, "fallback works");
let delta = rx.recv().await.unwrap();
assert_eq!(delta, "fallback works");
}
#[test]
fn policy_retries_5xx_with_exponential_backoff() {
let policy = RetryPolicy::default(); assert_eq!(
policy.backoff_for(0, Some(503), false),
Some(Duration::from_secs(1))
);
assert_eq!(
policy.backoff_for(1, Some(500), false),
Some(Duration::from_secs(2))
);
assert_eq!(policy.backoff_for(2, Some(500), false), None);
}
#[test]
fn policy_retries_network_errors() {
let policy = RetryPolicy::default();
assert_eq!(
policy.backoff_for(0, None, true),
Some(Duration::from_secs(1))
);
}
#[test]
fn policy_does_not_retry_4xx() {
let policy = RetryPolicy::default();
assert_eq!(policy.backoff_for(0, Some(400), false), None);
assert_eq!(policy.backoff_for(0, Some(401), false), None);
assert_eq!(policy.backoff_for(0, Some(404), false), None);
assert_eq!(policy.backoff_for(0, Some(429), false), None);
}
#[test]
fn policy_caps_backoff_at_max() {
let policy = RetryPolicy {
max_retries: 10,
initial_backoff: Duration::from_secs(1),
max_backoff: Duration::from_secs(3),
};
assert_eq!(
policy.backoff_for(5, Some(500), false),
Some(Duration::from_secs(3))
);
}
#[tokio::test]
async fn openai_structured_includes_schema_in_request_body() {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let captured = std::sync::Arc::new(std::sync::Mutex::new(String::new()));
let captured_clone = captured.clone();
std::thread::spawn(move || {
let (mut stream, _) = listener.accept().unwrap();
use std::io::{Read, Write};
let mut buf = [0u8; 8192];
let n = stream.read(&mut buf).unwrap();
let request = String::from_utf8_lossy(&buf[..n]).to_string();
if let Some(body_start) = request.find("\r\n\r\n") {
let body = request[body_start + 4..].trim().to_string();
*captured_clone.lock().unwrap() = body.clone();
}
let response_body = r#"{"choices":[{"message":{"role":"assistant","content":"{\"answer\":42}"},"finish_reason":"stop"}],"usage":null}"#;
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
response_body.len(),
response_body,
)
.unwrap();
stream.flush().unwrap();
});
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let provider =
OpenAiProvider::new(format!("http://{addr}"), "sk-noop", "test-structured-model");
let schema = serde_json::json!({
"type": "object",
"properties": {
"answer": {"type": "integer"}
},
"required": ["answer"]
});
let req = StructuredRequest {
messages: vec![Message::user("what is 6 times 7?".to_string())],
schema: schema.clone(),
schema_name: "math_answer".to_string(),
};
let _ = provider.complete_structured(req).await;
let body_str = captured.lock().unwrap().clone();
let body: serde_json::Value = serde_json::from_str(&body_str).unwrap();
let rf = body.get("response_format").unwrap();
assert_eq!(rf["type"], "json_schema");
assert_eq!(rf["json_schema"]["name"], "math_answer");
assert_eq!(rf["json_schema"]["strict"], true);
assert_eq!(rf["json_schema"]["schema"], schema);
}
#[tokio::test]
async fn openai_structured_parses_response_json() {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
std::thread::spawn(move || {
let (mut stream, _) = listener.accept().unwrap();
use std::io::{Read, Write};
let mut buf = [0u8; 8192];
let _ = stream.read(&mut buf);
let response_body = r#"{"choices":[{"message":{"role":"assistant","content":"{\"summary\":\"test\",\"kept_facts\":[\"a\",\"b\"]}"},"finish_reason":"stop"}],"usage":null}"#;
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
response_body.len(),
response_body,
)
.unwrap();
stream.flush().unwrap();
});
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let provider =
OpenAiProvider::new(format!("http://{addr}"), "sk-noop", "test-structured-model");
let schema = serde_json::json!({
"type": "object",
"properties": {
"summary": {"type": "string"},
"kept_facts": {"type": "array", "items": {"type": "string"}}
},
"required": ["summary", "kept_facts"]
});
let req = StructuredRequest {
messages: vec![Message::user("summarize".to_string())],
schema,
schema_name: "summary".to_string(),
};
let result = provider.complete_structured(req).await;
assert!(result.is_ok(), "got error: {:?}", result.err());
let value = result.unwrap();
assert_eq!(value["summary"], "test");
assert_eq!(value["kept_facts"][0], "a");
assert_eq!(value["kept_facts"][1], "b");
}
}