use crate::errors::AppError;
use crate::retry::AttemptOutcome;
use secrecy::{ExposeSecret, SecretBox};
use serde::{Deserialize, Serialize};
use std::time::Duration;
use crate::constants::DEFAULT_OPENROUTER_CHAT_URL;
const DEFAULT_TIMEOUT_SECS: u64 = 600;
const DEFAULT_CONNECT_TIMEOUT_SECS: u64 = 10;
const SCHEMA_NAME: &str = "enrich_output";
#[derive(Serialize)]
struct ChatRequest<'a> {
model: &'a str,
messages: Vec<ChatMessage<'a>>,
response_format: ResponseFormat,
provider: ProviderPrefs,
#[serde(skip_serializing_if = "Option::is_none")]
reasoning: Option<ReasoningPrefs>,
#[serde(skip_serializing_if = "Option::is_none")]
max_tokens: Option<u32>,
}
#[derive(Serialize)]
struct ChatMessage<'a> {
role: &'a str,
content: String,
}
#[derive(Serialize)]
struct ResponseFormat {
#[serde(rename = "type")]
format_type: &'static str,
json_schema: JsonSchemaSpec,
}
#[derive(Serialize)]
struct JsonSchemaSpec {
name: &'static str,
strict: bool,
schema: serde_json::Value,
}
#[derive(Serialize)]
struct ProviderPrefs {
require_parameters: bool,
}
#[derive(Serialize)]
struct ReasoningPrefs {
enabled: bool,
}
#[derive(Deserialize)]
struct ChatResponse {
#[serde(default)]
choices: Vec<Choice>,
#[serde(default)]
usage: Option<Usage>,
#[serde(default)]
error: Option<crate::openrouter_http::ApiError>,
}
#[derive(Deserialize)]
struct Choice {
message: RespMessage,
#[serde(default)]
finish_reason: Option<String>,
}
#[derive(Deserialize)]
struct RespMessage {
#[serde(default)]
content: Option<String>,
}
#[derive(Deserialize)]
struct Usage {
#[serde(default)]
cost: Option<f64>,
#[serde(default)]
prompt_tokens: Option<u32>,
#[serde(default)]
completion_tokens: Option<u32>,
}
#[derive(Debug)]
pub struct ChatCompletion {
pub value: serde_json::Value,
pub cost_usd: f64,
pub finish_reason: Option<String>,
pub prompt_tokens: Option<u32>,
pub completion_tokens: Option<u32>,
}
#[derive(Debug)]
pub struct ChatError {
pub source: AppError,
pub finish_reason: Option<String>,
pub prompt_tokens: Option<u32>,
pub completion_tokens: Option<u32>,
pub retry_class: AttemptOutcome,
}
impl ChatError {
fn new(source: AppError, retry_class: AttemptOutcome) -> Self {
Self {
source,
finish_reason: None,
prompt_tokens: None,
completion_tokens: None,
retry_class,
}
}
fn with_diagnostics(
source: AppError,
finish_reason: Option<String>,
prompt_tokens: Option<u32>,
completion_tokens: Option<u32>,
retry_class: AttemptOutcome,
) -> Self {
Self {
source,
finish_reason,
prompt_tokens,
completion_tokens,
retry_class,
}
}
}
impl std::fmt::Display for ChatError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
std::fmt::Display::fmt(&self.source, f)
}
}
impl std::error::Error for ChatError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
Some(&self.source)
}
}
pub struct OpenRouterChatClient {
client: reqwest::Client,
api_key: SecretBox<String>,
model: String,
base_url: String,
}
impl OpenRouterChatClient {
pub fn new(
api_key: SecretBox<String>,
model: String,
timeout_secs: u64,
) -> Result<Self, AppError> {
let base_url = crate::runtime_config::openrouter_chat_url(DEFAULT_OPENROUTER_CHAT_URL);
Self::new_with_base_url(api_key, model, timeout_secs, base_url)
}
pub fn new_with_base_url(
api_key: SecretBox<String>,
model: String,
timeout_secs: u64,
base_url: String,
) -> Result<Self, AppError> {
let timeout_secs = if timeout_secs == 0 {
DEFAULT_TIMEOUT_SECS
} else {
timeout_secs
};
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(timeout_secs))
.connect_timeout(Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SECS))
.user_agent(concat!("sqlite-graphrag/", env!("CARGO_PKG_VERSION")))
.build()
.map_err(|e| {
AppError::Validation(crate::i18n::validation::http_client_build_failed(&e))
})?;
Ok(Self {
client,
api_key,
model,
base_url,
})
}
#[cfg(test)]
pub fn new_with_url(
api_key: SecretBox<String>,
model: String,
base_url: String,
timeout_secs: u64,
) -> Result<Self, AppError> {
Self::new_with_base_url(api_key, model, timeout_secs, base_url)
}
pub fn model(&self) -> &str {
&self.model
}
pub async fn complete(
&self,
system_prompt: &str,
input_text: &str,
schema_str: &str,
max_tokens: Option<u32>,
) -> Result<ChatCompletion, ChatError> {
let schema: serde_json::Value = serde_json::from_str(schema_str).map_err(|e| {
ChatError::new(
AppError::Validation(crate::i18n::validation::invalid_json_schema_for_request(&e)),
AttemptOutcome::HardFailure,
)
})?;
let mut current_max_tokens = max_tokens;
for length_attempt in 0..=crate::constants::ENRICH_MAX_LENGTH_RETRIES {
let response = self
.complete_one_attempt(&schema, system_prompt, input_text, current_max_tokens)
.await?;
let finish_reason = response
.choices
.first()
.and_then(|c| c.finish_reason.clone());
let prompt_tokens = response.usage.as_ref().and_then(|u| u.prompt_tokens);
let completion_tokens = response.usage.as_ref().and_then(|u| u.completion_tokens);
let truncated = finish_reason.as_deref() == Some("length");
let retries_left = length_attempt < crate::constants::ENRICH_MAX_LENGTH_RETRIES;
if truncated && retries_left {
let next_max_tokens = grow_max_tokens(current_max_tokens);
tracing::warn!(
model = %self.model,
attempt = length_attempt,
previous_max_tokens = ?current_max_tokens,
next_max_tokens,
"OpenRouter completion truncated (finish_reason=length); \
retrying with a larger max_tokens budget"
);
current_max_tokens = Some(next_max_tokens);
continue;
}
if truncated {
tracing::warn!(
model = %self.model,
max_length_retries = crate::constants::ENRICH_MAX_LENGTH_RETRIES,
max_tokens = ?current_max_tokens,
"OpenRouter completion still truncated after exhausting \
max_tokens growth"
);
}
return self.finish_completion(
response,
finish_reason,
prompt_tokens,
completion_tokens,
);
}
unreachable!("loop always returns within ENRICH_MAX_LENGTH_RETRIES + 1 iterations")
}
async fn complete_one_attempt(
&self,
schema: &serde_json::Value,
system_prompt: &str,
input_text: &str,
max_tokens: Option<u32>,
) -> Result<ChatResponse, ChatError> {
let primary = self.build_request(
schema.clone(),
system_prompt,
input_text,
max_tokens,
Some(ReasoningPrefs { enabled: false }),
);
match self.execute_with_retry(&primary).await {
Ok(r) => Ok(r),
Err(first_err) => {
if reasoning_disable_rejected(&first_err) {
tracing::warn!(
model = %self.model,
"model rejected reasoning.enabled=false (mandatory); \
retrying once with reasoning omitted"
);
let fallback = self.build_request(
schema.clone(),
system_prompt,
input_text,
max_tokens,
None,
);
match self.execute_with_retry(&fallback).await {
Ok(r) => Ok(r),
Err(_) => Err(first_err),
}
} else {
Err(first_err)
}
}
}
}
fn finish_completion(
&self,
response: ChatResponse,
finish_reason: Option<String>,
prompt_tokens: Option<u32>,
completion_tokens: Option<u32>,
) -> Result<ChatCompletion, ChatError> {
let content = response
.choices
.into_iter()
.next()
.and_then(|c| c.message.content)
.filter(|c| !c.trim().is_empty())
.ok_or_else(|| {
AppError::Validation(crate::i18n::validation::model_no_structured_content(
&self.model,
))
})
.map_err(|e| {
ChatError::with_diagnostics(
e,
finish_reason.clone(),
prompt_tokens,
completion_tokens,
AttemptOutcome::Transient,
)
})?;
let value = crate::json_repair::repair_to_value(&content).map_err(|e| {
ChatError::with_diagnostics(
AppError::Validation(crate::i18n::validation::model_json_parse_failed(
&self.model,
&e,
)),
finish_reason.clone(),
prompt_tokens,
completion_tokens,
AttemptOutcome::Transient,
)
})?;
if !value.is_object() {
return Err(ChatError::with_diagnostics(
AppError::Validation(crate::i18n::validation::model_non_object_json(
&self.model,
json_shape_name(&value),
)),
finish_reason,
prompt_tokens,
completion_tokens,
AttemptOutcome::Transient,
));
}
let cost = response.usage.and_then(|u| u.cost).unwrap_or(0.0);
Ok(ChatCompletion {
value,
cost_usd: cost,
finish_reason,
prompt_tokens,
completion_tokens,
})
}
fn build_request<'a>(
&'a self,
schema: serde_json::Value,
system_prompt: &str,
input_text: &str,
max_tokens: Option<u32>,
reasoning: Option<ReasoningPrefs>,
) -> ChatRequest<'a> {
let mut messages = Vec::with_capacity(2);
messages.push(ChatMessage {
role: "system",
content: system_prompt.to_string(),
});
if !input_text.is_empty() {
messages.push(ChatMessage {
role: "user",
content: input_text.to_string(),
});
}
ChatRequest {
model: &self.model,
messages,
response_format: ResponseFormat {
format_type: "json_schema",
json_schema: JsonSchemaSpec {
name: SCHEMA_NAME,
strict: true,
schema,
},
},
provider: ProviderPrefs {
require_parameters: true,
},
reasoning,
max_tokens,
}
}
async fn execute_with_retry(
&self,
request: &ChatRequest<'_>,
) -> Result<ChatResponse, ChatError> {
let mut last_err: Option<ChatError> = None;
for attempt in 0..crate::openrouter_http::MAX_RETRIES {
let result = self
.client
.post(&self.base_url)
.header(
"Authorization",
format!("Bearer {}", self.api_key.expose_secret()),
)
.json(request)
.send()
.await;
let resp = match result {
Ok(r) => r,
Err(e) if e.is_timeout() => {
return Err(ChatError::new(
AppError::Validation(crate::i18n::validation::openrouter_chat_timed_out()),
AttemptOutcome::Transient,
));
}
Err(e) => {
last_err = Some(ChatError::new(
AppError::Validation(crate::i18n::validation::http_request_failed(&e)),
AttemptOutcome::Transient,
));
crate::openrouter_http::backoff(attempt).await;
continue;
}
};
let status = resp.status();
if status.is_success() {
let body = resp.text().await.map_err(|e| {
ChatError::new(
AppError::Validation(
crate::i18n::validation::failed_to_read_response_body(&e),
),
AttemptOutcome::Transient,
)
})?;
match serde_json::from_str::<ChatResponse>(&body) {
Ok(parsed) => {
if let Some(api_err) = parsed.error {
let retry_class =
crate::openrouter_http::provider_error_retry_class(&api_err);
return Err(ChatError::new(
AppError::ProviderError {
code: api_err.code_string(),
message: api_err.message,
},
retry_class,
));
}
return Ok(parsed);
}
Err(e) => {
tracing::warn!(
attempt,
body_len = body.len(),
"HTTP 200 but parse failed (retrying): {e}"
);
last_err = Some(ChatError::new(
AppError::Validation(
crate::i18n::validation::failed_to_parse_chat_response(&e),
),
AttemptOutcome::Transient,
));
crate::openrouter_http::backoff(attempt).await;
continue;
}
}
}
if status.as_u16() == 401 {
return Err(ChatError::new(
AppError::Validation(crate::i18n::validation::openrouter_invalid_api_key_401()),
AttemptOutcome::HardFailure,
));
}
if status.as_u16() == 400 || status.as_u16() == 404 {
let body = resp.text().await.unwrap_or_default();
return Err(ChatError::new(
AppError::Validation(crate::i18n::validation::openrouter_status_error(
&status,
&self.model,
&body,
)),
AttemptOutcome::HardFailure,
));
}
if status.as_u16() == 429 {
let retry_after = resp
.headers()
.get("retry-after")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(2);
tracing::warn!(
attempt,
retry_after_secs = retry_after,
"OpenRouter rate limited, waiting"
);
last_err = Some(ChatError::new(
AppError::RateLimited {
detail: format!("OpenRouter HTTP 429 (retry-after {retry_after}s)"),
},
AttemptOutcome::Transient,
));
tokio::time::sleep(Duration::from_secs(retry_after)).await;
continue;
}
if status.is_server_error() {
tracing::warn!(attempt, status = %status, "OpenRouter server error, retrying");
last_err = Some(ChatError::new(
AppError::Validation(crate::i18n::validation::openrouter_server_error(&status)),
AttemptOutcome::Transient,
));
crate::openrouter_http::backoff(attempt).await;
continue;
}
let body = resp.text().await.unwrap_or_default();
return Err(ChatError::new(
AppError::Validation(crate::i18n::validation::unexpected_http_status(
&status, &body,
)),
crate::openrouter_http::status_retry_class(status),
));
}
Err(last_err.unwrap_or_else(|| {
ChatError::new(
AppError::Validation(crate::i18n::validation::openrouter_chat_max_retries()),
AttemptOutcome::Transient,
)
}))
}
}
fn grow_max_tokens(current: Option<u32>) -> u32 {
let base = current.unwrap_or(crate::constants::ENRICH_INITIAL_MAX_TOKENS);
base.saturating_mul(crate::constants::ENRICH_MAX_TOKENS_GROWTH_FACTOR)
.min(crate::constants::ENRICH_MAX_TOKENS_CEILING)
}
fn reasoning_disable_rejected(err: &ChatError) -> bool {
let msg = err.source.to_string().to_lowercase();
msg.contains("400") && msg.contains("reasoning")
}
fn json_shape_name(value: &serde_json::Value) -> &'static str {
match value {
serde_json::Value::Null => "null",
serde_json::Value::Bool(_) => "boolean",
serde_json::Value::Number(_) => "number",
serde_json::Value::String(_) => "string",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => "object",
}
}
#[cfg(test)]
#[path = "chat_api_tests.rs"]
mod tests;