use async_trait::async_trait;
use reqwest::Client;
use serde_json::Value;
use crate::inference::adapter::InferenceAdapter;
use crate::inference::error::InferenceError;
use crate::inference::registry::ProviderCapabilities;
use crate::inference::streaming::{ChatStream, buffered_stream, decode_event_stream};
use crate::inference::types::{ChatRequest, ChatResponse, RequestUsageConfig, SecretString};
pub struct OpenAiCompatConfig {
pub name: String,
pub base_url: String,
pub api_key: SecretString,
pub extra_headers: Vec<(String, String)>,
pub capabilities: ProviderCapabilities,
}
pub struct OpenAiCompatAdapter {
name: String,
endpoint: String,
api_key: SecretString,
extra_headers: Vec<(String, String)>,
capabilities: ProviderCapabilities,
http: Client,
}
impl OpenAiCompatAdapter {
pub fn new(config: OpenAiCompatConfig) -> Result<Self, InferenceError> {
let endpoint = format!("{}/chat/completions", config.base_url.trim_end_matches('/'));
let http = Client::builder()
.use_rustls_tls()
.build()
.map_err(|e| InferenceError::Transport(e.to_string()))?;
Ok(Self {
name: config.name,
endpoint,
api_key: config.api_key,
extra_headers: config.extra_headers,
capabilities: config.capabilities,
http,
})
}
pub fn endpoint(&self) -> &str {
&self.endpoint
}
fn prepare_body(&self, request: &ChatRequest) -> ChatRequest {
let mut req = request.clone();
if self.wants_detailed_usage() && req.usage.is_none() {
req.usage = Some(RequestUsageConfig::detailed());
}
req
}
fn stream_body(&self, request: &ChatRequest) -> Result<Value, InferenceError> {
let prepared = self.prepare_body(request);
let mut body =
serde_json::to_value(&prepared).map_err(|e| InferenceError::Deserialise {
message: e.to_string(),
body: String::new(),
})?;
if let Value::Object(map) = &mut body {
map.insert("stream".to_string(), Value::Bool(true));
map.insert(
"stream_options".to_string(),
serde_json::json!({ "include_usage": true }),
);
}
Ok(body)
}
}
#[async_trait]
impl InferenceAdapter for OpenAiCompatAdapter {
fn name(&self) -> &str {
&self.name
}
fn capabilities(&self) -> &ProviderCapabilities {
&self.capabilities
}
async fn chat(&self, request: &ChatRequest) -> Result<ChatResponse, InferenceError> {
let body = self.prepare_body(request);
let mut builder = self
.http
.post(&self.endpoint)
.header(
reqwest::header::AUTHORIZATION,
format!("Bearer {}", self.api_key.expose()),
)
.header(reqwest::header::CONTENT_TYPE, "application/json");
for (name, value) in &self.extra_headers {
builder = builder.header(name, value);
}
let resp = builder
.json(&body)
.send()
.await
.map_err(|e| InferenceError::Transport(e.to_string()))?;
let status = resp.status();
if !status.is_success() {
let body = resp.text().await.unwrap_or_default();
return Err(InferenceError::Api {
status: status.as_u16(),
body,
});
}
let text = resp
.text()
.await
.map_err(|e| InferenceError::Transport(e.to_string()))?;
serde_json::from_str::<ChatResponse>(&text).map_err(|e| InferenceError::Deserialise {
message: e.to_string(),
body: text,
})
}
async fn chat_stream(&self, request: &ChatRequest) -> Result<ChatStream, InferenceError> {
let body = self.stream_body(request)?;
let mut builder = self
.http
.post(&self.endpoint)
.header(
reqwest::header::AUTHORIZATION,
format!("Bearer {}", self.api_key.expose()),
)
.header(reqwest::header::CONTENT_TYPE, "application/json")
.header(reqwest::header::ACCEPT, "text/event-stream");
for (name, value) in &self.extra_headers {
builder = builder.header(name, value);
}
let resp = builder
.json(&body)
.send()
.await
.map_err(|e| InferenceError::Transport(e.to_string()))?;
let status = resp.status();
if !status.is_success() {
let body = resp.text().await.unwrap_or_default();
return Err(InferenceError::Api {
status: status.as_u16(),
body,
});
}
let is_sse = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.map(|v| v.to_ascii_lowercase().contains("text/event-stream"))
.unwrap_or(false);
if is_sse {
return Ok(decode_event_stream(resp.bytes_stream()));
}
let text = resp
.text()
.await
.map_err(|e| InferenceError::Transport(e.to_string()))?;
let response = serde_json::from_str::<ChatResponse>(&text).map_err(|e| {
InferenceError::Deserialise {
message: e.to_string(),
body: text,
}
})?;
Ok(buffered_stream(response))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::inference::registry::{ProviderId, capabilities};
use crate::inference::types::ChatMessage;
fn config_for(id: ProviderId, base_url: &str) -> OpenAiCompatConfig {
OpenAiCompatConfig {
name: id.as_str().to_string(),
base_url: base_url.to_string(),
api_key: SecretString::new("sk-test"), extra_headers: Vec::new(),
capabilities: *capabilities(id),
}
}
#[test]
fn endpoint_appends_path() {
let a =
OpenAiCompatAdapter::new(config_for(ProviderId::OpenAI, "https://api.openai.com/v1"))
.expect("build");
assert_eq!(a.endpoint(), "https://api.openai.com/v1/chat/completions");
let b =
OpenAiCompatAdapter::new(config_for(ProviderId::OpenAI, "https://api.openai.com/v1/"))
.expect("build");
assert_eq!(b.endpoint(), "https://api.openai.com/v1/chat/completions");
}
#[test]
fn usage_directive_injected() {
let a = OpenAiCompatAdapter::new(config_for(
ProviderId::OpenRouter,
"https://openrouter.ai/api/v1",
))
.expect("build");
let req = ChatRequest::new("openai/gpt-4o-mini", vec![ChatMessage::user("hi")]);
let prepared = a.prepare_body(&req);
assert_eq!(prepared.usage, Some(RequestUsageConfig::detailed()));
}
#[test]
fn usage_directive_absent_without_detailed() {
for id in [ProviderId::Fireworks, ProviderId::OpenAI] {
let a =
OpenAiCompatAdapter::new(config_for(id, "https://example.test/v1")).expect("build");
let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
assert!(
a.prepare_body(&req).usage.is_none(),
"{} must not inject usage",
id.as_str()
);
}
}
#[test]
fn usage_directive_not_overwritten() {
let a = OpenAiCompatAdapter::new(config_for(
ProviderId::OpenRouter,
"https://openrouter.ai/api/v1",
))
.expect("build");
let mut req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
req.usage = Some(RequestUsageConfig { include: false });
assert_eq!(
a.prepare_body(&req).usage,
Some(RequestUsageConfig { include: false })
);
}
#[test]
fn stream_body_sets_stream_and_usage_options() {
let a = OpenAiCompatAdapter::new(config_for(ProviderId::OpenAI, "https://example.test/v1"))
.expect("build");
let req = ChatRequest::new("m", vec![ChatMessage::user("hi")]);
let body = a.stream_body(&req).expect("body");
assert_eq!(body["stream"], serde_json::json!(true));
assert_eq!(
body["stream_options"]["include_usage"],
serde_json::json!(true)
);
assert_eq!(body["model"], "m");
}
#[test]
fn stream_body_preserves_openrouter_usage_directive() {
let a = OpenAiCompatAdapter::new(config_for(
ProviderId::OpenRouter,
"https://openrouter.ai/api/v1",
))
.expect("build");
let req = ChatRequest::new("openai/gpt-4o-mini", vec![ChatMessage::user("hi")]);
let body = a.stream_body(&req).expect("body");
assert_eq!(body["usage"], serde_json::json!({ "include": true }));
assert_eq!(body["stream"], serde_json::json!(true));
}
}