use async_trait::async_trait;
use futures::stream::BoxStream;
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub mod gemini;
pub mod openai;
pub mod provider;
pub mod tool_call;
pub use provider::{
LlmConfigError, LlmSettings, PROVIDER_ENV_VARS, ReasoningPlan, SUPPORTED_PROVIDERS,
SelectedLlm, provider_from_env, provider_from_settings,
};
pub use tool_call::{PartialToolCall, ToolCallAccumulator};
#[derive(Clone, Copy)]
pub(crate) struct Env<'a>(&'a dyn Fn(&str) -> Option<String>);
impl<'a> Env<'a> {
#[cfg(test)]
pub(crate) fn new(lookup: &'a dyn Fn(&str) -> Option<String>) -> Self {
Self(lookup)
}
pub(crate) fn os() -> Env<'static> {
const LOOKUP: &dyn Fn(&str) -> Option<String> = &os_lookup;
Env(LOOKUP)
}
pub(crate) fn get(&self, key: &str) -> Option<String> {
(self.0)(key)
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
}
fn os_lookup(key: &str) -> Option<String> {
std::env::var(key).ok()
}
#[derive(Debug, thiserror::Error)]
pub enum LlmError {
#[error("API error: {0}")]
ApiError(String),
#[error("context length exceeded: {detail}")]
ContextLengthExceeded {
detail: String,
prompt_tokens: Option<u32>,
context_window: Option<u32>,
},
#[error("Network error: {0}")]
NetworkError(String),
#[error("Serialization error: {0}")]
SerializationError(String),
#[error("Provider error: {0}")]
ProviderError(String),
}
pub(crate) fn describe_transport_error(error: &dyn std::error::Error) -> String {
let mut message = error.to_string();
let mut source = error.source();
while let Some(cause) = source {
message.push_str(": ");
message.push_str(&cause.to_string());
source = cause.source();
}
message
}
const CONTEXT_LENGTH_MARKERS: [&str; 8] = [
"context_length_exceeded",
"maximum context length",
"context length",
"context size",
"too many tokens",
"exceeds the maximum",
"exceed_context_size_error",
"input token count",
];
pub(crate) fn classify_api_error(label: &str, status: reqwest::StatusCode, body: &str) -> LlmError {
let message = format!("{label} ({status}): {body}");
let haystack = message.to_lowercase();
if CONTEXT_LENGTH_MARKERS
.iter()
.any(|marker| haystack.contains(marker))
{
let (prompt_tokens, context_window) = context_refusal_numbers(body);
return LlmError::ContextLengthExceeded {
detail: message,
prompt_tokens,
context_window,
};
}
LlmError::ApiError(message)
}
fn context_refusal_numbers(body: &str) -> (Option<u32>, Option<u32>) {
let json: Option<Value> = serde_json::from_str(body).ok();
let field = |name: &str| {
let json = json.as_ref()?;
json.get(name)
.or_else(|| json.get("error").and_then(|error| error.get(name)))
.and_then(Value::as_u64)
.and_then(|n| u32::try_from(n).ok())
};
let prose = body.to_lowercase();
let prompt_tokens =
field("n_prompt_tokens").or_else(|| number_after(&prose, "input token count ("));
let context_window =
field("n_ctx").or_else(|| number_after(&prose, "maximum context length is "));
(prompt_tokens, context_window)
}
fn number_after(haystack: &str, pattern: &str) -> Option<u32> {
let start = haystack.find(pattern)? + pattern.len();
let digits: &str = haystack[start..]
.split(|c: char| !c.is_ascii_digit())
.next()?;
digits.parse().ok()
}
#[derive(Clone, Default)]
pub(crate) struct ReasoningSupport(std::sync::Arc<std::sync::atomic::AtomicBool>);
impl ReasoningSupport {
pub(crate) fn refused(&self) -> bool {
self.0.load(std::sync::atomic::Ordering::Relaxed)
}
pub(crate) fn record_refusal(&self) {
self.0.store(true, std::sync::atomic::Ordering::Relaxed);
}
}
pub(crate) fn refuses_reasoning(status: u16, body: &str, field: &str) -> bool {
status == 400 && body.to_lowercase().contains(field)
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TokenUsage {
pub prompt_tokens: Option<u32>,
pub completion_tokens: Option<u32>,
pub reasoning_tokens: Option<u32>,
pub total_tokens: Option<u32>,
}
impl TokenUsage {
pub fn is_empty(&self) -> bool {
*self == Self::default()
}
}
impl std::fmt::Display for TokenUsage {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let field = |value: Option<u32>| match value {
Some(count) => count.to_string(),
None => "?".to_string(),
};
write!(
f,
"prompt={} completion={} reasoning={} total={}",
field(self.prompt_tokens),
field(self.completion_tokens),
field(self.reasoning_tokens),
field(self.total_tokens)
)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum MessageRole {
System,
User,
Assistant,
Tool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
pub name: String,
pub description: String,
pub parameters: Value, }
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
pub id: String, pub name: String,
pub arguments: String, }
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
pub role: MessageRole,
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCall>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
}
impl ChatMessage {
pub fn system(content: impl Into<String>) -> Self {
Self {
role: MessageRole::System,
content: Some(content.into()),
tool_calls: None,
tool_call_id: None,
name: None,
}
}
pub fn user(content: impl Into<String>) -> Self {
Self {
role: MessageRole::User,
content: Some(content.into()),
tool_calls: None,
tool_call_id: None,
name: None,
}
}
pub fn assistant(content: impl Into<String>) -> Self {
Self {
role: MessageRole::Assistant,
content: Some(content.into()),
tool_calls: None,
tool_call_id: None,
name: None,
}
}
pub fn tool_result(
tool_call_id: impl Into<String>,
name: impl Into<String>,
content: impl Into<String>,
) -> Self {
Self {
role: MessageRole::Tool,
content: Some(content.into()),
tool_calls: None,
tool_call_id: Some(tool_call_id.into()),
name: Some(name.into()),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReasoningEffort {
Low,
Medium,
High,
}
impl ReasoningEffort {
pub fn as_str(self) -> &'static str {
match self {
ReasoningEffort::Low => "low",
ReasoningEffort::Medium => "medium",
ReasoningEffort::High => "high",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Reasoning {
Off,
Effort(ReasoningEffort),
Budget(u32),
}
const REASONING_EXPECTED: &str =
r#""off", "low", "medium", "high", or a number of reasoning tokens"#;
impl std::str::FromStr for Reasoning {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.trim() {
"off" => Ok(Reasoning::Off),
"low" => Ok(Reasoning::Effort(ReasoningEffort::Low)),
"medium" => Ok(Reasoning::Effort(ReasoningEffort::Medium)),
"high" => Ok(Reasoning::Effort(ReasoningEffort::High)),
budget => budget
.parse()
.map(Reasoning::Budget)
.map_err(|_| format!("expected {REASONING_EXPECTED}; got {s:?}")),
}
}
}
impl std::fmt::Display for Reasoning {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Reasoning::Off => f.write_str("off"),
Reasoning::Effort(effort) => f.write_str(effort.as_str()),
Reasoning::Budget(tokens) => write!(f, "{tokens}"),
}
}
}
impl Serialize for Reasoning {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
match self {
Reasoning::Budget(tokens) => serializer.serialize_u32(*tokens),
level => serializer.serialize_str(&level.to_string()),
}
}
}
impl<'de> Deserialize<'de> for Reasoning {
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
use serde::de::{Error, Unexpected, Visitor};
struct ReasoningVisitor;
impl Visitor<'_> for ReasoningVisitor {
type Value = Reasoning;
fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(REASONING_EXPECTED)
}
fn visit_str<E: Error>(self, value: &str) -> Result<Reasoning, E> {
value
.parse()
.map_err(|_| E::invalid_value(Unexpected::Str(value), &self))
}
fn visit_u64<E: Error>(self, value: u64) -> Result<Reasoning, E> {
u32::try_from(value)
.map(Reasoning::Budget)
.map_err(|_| E::invalid_value(Unexpected::Unsigned(value), &self))
}
fn visit_i64<E: Error>(self, value: i64) -> Result<Reasoning, E> {
u32::try_from(value)
.map(Reasoning::Budget)
.map_err(|_| E::invalid_value(Unexpected::Signed(value), &self))
}
}
deserializer.deserialize_any(ReasoningVisitor)
}
}
#[derive(Debug, Clone)]
pub struct LlmRequest {
pub messages: Vec<ChatMessage>,
pub tools: Option<Vec<ToolDefinition>>,
pub temperature: Option<f32>,
pub max_tokens: Option<u32>,
pub force_json: bool,
pub reasoning: Option<Reasoning>,
}
impl LlmRequest {
pub fn new(messages: Vec<ChatMessage>) -> Self {
Self {
messages,
tools: None,
temperature: None,
max_tokens: None,
force_json: false,
reasoning: None,
}
}
pub fn reasoning(mut self, reasoning: Reasoning) -> Self {
self.reasoning = Some(reasoning);
self
}
pub fn temperature(mut self, temp: f32) -> Self {
self.temperature = Some(temp);
self
}
pub fn max_tokens(mut self, tokens: u32) -> Self {
self.max_tokens = Some(tokens);
self
}
pub fn tools(mut self, tools: Vec<ToolDefinition>) -> Self {
self.tools = Some(tools);
self
}
pub fn force_json(mut self, force: bool) -> Self {
self.force_json = force;
self
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FinishReason {
Stop,
Length,
ToolCalls,
ContentFilter,
Other(String),
}
impl FinishReason {
pub fn from_wire(value: &str) -> Self {
match value {
"stop" | "STOP" => Self::Stop,
"length" | "MAX_TOKENS" => Self::Length,
"tool_calls" | "function_call" => Self::ToolCalls,
"content_filter" | "SAFETY" | "RECITATION" | "PROHIBITED_CONTENT" | "BLOCKLIST"
| "SPII" => Self::ContentFilter,
other => Self::Other(other.to_string()),
}
}
}
impl std::fmt::Display for FinishReason {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
FinishReason::Stop => f.write_str("stop"),
FinishReason::Length => f.write_str("length"),
FinishReason::ToolCalls => f.write_str("tool_calls"),
FinishReason::ContentFilter => f.write_str("content_filter"),
FinishReason::Other(reason) => f.write_str(reason),
}
}
}
#[derive(Debug, Clone)]
pub struct LlmResponse {
pub content: Option<String>,
pub tool_calls: Option<Vec<ToolCall>>,
pub reasoning: Option<String>,
pub usage: Option<TokenUsage>,
pub finish: Option<FinishReason>,
}
#[derive(Debug, Clone)]
pub enum LlmStreamEvent {
ContentChunk(String),
Reasoning(String),
ToolCallChunk {
id: String,
name: Option<String>,
arguments: String,
},
ToolCall(ToolCall),
Usage(TokenUsage),
Finish(FinishReason),
}
#[async_trait]
pub trait LlmProvider: Send + Sync {
async fn chat_completion(&self, request: LlmRequest) -> Result<LlmResponse, LlmError>;
async fn chat_completion_stream(
&self,
request: LlmRequest,
) -> Result<BoxStream<'static, Result<LlmStreamEvent, LlmError>>, LlmError>;
}
#[cfg(test)]
mod error_tests {
use super::*;
fn status(code: u16) -> reqwest::StatusCode {
reqwest::StatusCode::from_u16(code).unwrap()
}
#[test]
fn an_over_long_request_is_classified_as_a_context_length_failure() {
let cases = [
(
"OpenAI API error",
r#"{"error":{"message":"This model's maximum context length is 128000 tokens","code":"context_length_exceeded"}}"#,
None,
Some(128000),
),
(
"OpenAI stream error",
"Requested 200000 tokens, exceeds the maximum for this model",
None,
None,
),
(
"Gemini API error",
r#"{"error":{"status":"INVALID_ARGUMENT","message":"The input token count (1200000) exceeds the maximum"}}"#,
Some(1200000),
None,
),
("OpenAI API error", "too many tokens in prompt", None, None),
(
"OpenAI stream error",
r#"{"error":{"code":400,"message":"request (40089 tokens) exceeds the available context size (32768 tokens), try increasing it","type":"exceed_context_size_error","n_prompt_tokens":40089,"n_ctx":32768}}"#,
Some(40089),
Some(32768),
),
];
for (label, body, prompt, window) in cases {
match classify_api_error(label, status(400), body) {
LlmError::ContextLengthExceeded {
detail,
prompt_tokens,
context_window,
} => {
assert!(detail.contains(label) && detail.contains(body), "{detail}");
assert_eq!(prompt_tokens, prompt, "prompt tokens for: {body}");
assert_eq!(context_window, window, "window for: {body}");
}
other => panic!("should classify as context length: {body} → {other:?}"),
}
}
}
#[test]
fn other_failures_stay_api_errors() {
let cases = [
(
"OpenAI API error",
401,
r#"{"error":{"message":"Incorrect API key provided"}}"#,
),
("OpenAI API error", 429, "Rate limit reached for requests"),
("Gemini API error", 503, "The model is overloaded"),
];
for (label, code, body) in cases {
assert!(
matches!(
classify_api_error(label, status(code), body),
LlmError::ApiError(_)
),
"should stay an API error: {body}"
);
}
}
#[test]
fn a_refused_reasoning_parameter_is_recognized_from_the_body() {
let openai = [
r#"{"error":{"message":"Unsupported parameter: 'reasoning_effort' is not supported with this model.","type":"invalid_request_error","param":"reasoning_effort","code":"unsupported_parameter"}}"#,
r#"{"error":{"message":"Invalid value: 'none'. Supported values are: 'low', 'medium' and 'high'.","type":"invalid_request_error","param":"reasoning_effort","code":"invalid_value"}}"#,
r#"{"error":{"message":"Unrecognized request argument supplied: reasoning_effort"}}"#,
];
for body in openai {
assert!(refuses_reasoning(400, body, "reasoning"), "{body}");
}
let gemini = [
r#"{"error":{"code":400,"message":"Invalid JSON payload received. Unknown name \"thinkingLevel\" at 'generation_config': Cannot find field.","status":"INVALID_ARGUMENT"}}"#,
r#"{"error":{"code":400,"message":"Budget 128 is invalid. thinkingBudget must be 0 or in the range [128, 32768]","status":"INVALID_ARGUMENT"}}"#,
];
for body in gemini {
assert!(refuses_reasoning(400, body, "thinking"), "{body}");
}
}
#[test]
fn other_failures_are_not_read_as_a_refusal() {
assert!(!refuses_reasoning(
400,
r#"{"error":{"message":"Incorrect API key provided"}}"#,
"reasoning"
));
assert!(!refuses_reasoning(
429,
r#"{"error":{"message":"Rate limit reached"}}"#,
"reasoning"
));
assert!(!refuses_reasoning(
503,
r#"{"error":{"message":"reasoning_effort backend unavailable"}}"#,
"reasoning"
));
}
#[test]
fn a_recorded_refusal_is_shared() {
let support = ReasoningSupport::default();
let clone = support.clone();
assert!(!support.refused());
clone.record_refusal();
assert!(support.refused());
}
#[test]
fn usage_with_nothing_reported_reads_as_empty() {
assert!(TokenUsage::default().is_empty());
assert!(
!TokenUsage {
prompt_tokens: Some(0),
..Default::default()
}
.is_empty(),
"a reported zero is a report, not an absence"
);
}
#[test]
fn finish_reasons_normalize_across_provider_spellings() {
let cases = [
("stop", FinishReason::Stop),
("STOP", FinishReason::Stop),
("length", FinishReason::Length),
("MAX_TOKENS", FinishReason::Length),
("tool_calls", FinishReason::ToolCalls),
("function_call", FinishReason::ToolCalls),
("content_filter", FinishReason::ContentFilter),
("SAFETY", FinishReason::ContentFilter),
("RECITATION", FinishReason::ContentFilter),
];
for (wire, expected) in cases {
assert_eq!(FinishReason::from_wire(wire), expected, "for {wire:?}");
}
assert_eq!(
FinishReason::from_wire("MALFORMED_FUNCTION_CALL"),
FinishReason::Other("MALFORMED_FUNCTION_CALL".to_string())
);
}
}