use anyhow::{bail, Context, Result};
use async_trait::async_trait;
pub mod client;
pub mod streaming;
pub mod tool_calling;
pub mod types;
pub use client::{ApiClient, RetryConfig};
pub use streaming::{StreamChunk, StreamingResponse};
pub use tool_calling::{
attach_tools, extract_tool_calls, extract_tool_calls_from_text, parsed_to_tool_call,
};
pub use types::*;
const DISABLED_THINKING_SYSTEM_MESSAGE: &str =
"CRITICAL INSTRUCTION: DO NOT use <think> blocks or any thinking process in your response. Output your final response directly and immediately.";
fn sanitize_assistant_content(messages: &mut [Message]) {
for message in messages.iter_mut() {
if message.role != "assistant" {
continue;
}
let has_tool_calls = message
.tool_calls
.as_ref()
.is_some_and(|calls| !calls.is_empty());
if has_tool_calls || !message.content.text().trim().is_empty() {
continue;
}
let filler = message
.reasoning_content
.as_deref()
.map(str::trim)
.filter(|s| !s.is_empty())
.unwrap_or("(no textual output)")
.to_string();
message.content = MessageContent::Text(filler);
}
}
fn canonicalize_message_order(messages: &mut Vec<Message>) {
sanitize_assistant_content(messages);
let sys_indices: Vec<usize> = messages
.iter()
.enumerate()
.filter(|(_, m)| m.role == "system")
.map(|(i, _)| i)
.collect();
if sys_indices.len() > 1 {
let merged_content: String = sys_indices
.iter()
.map(|&i| messages[i].content.to_string())
.collect::<Vec<_>>()
.join("\n\n");
for &i in sys_indices.iter().rev() {
messages.remove(i);
}
messages.insert(0, Message::system(merged_content));
}
let has_user = messages.iter().any(|m| m.role == "user");
if !has_user {
let insert_pos = if messages.first().map(|m| m.role.as_str()) == Some("system") {
1
} else {
0
};
messages.insert(insert_pos, Message::user("Continue with the task."));
}
}
fn maybe_prepend_disabled_thinking_instruction(
messages: &mut Vec<Message>,
thinking: &ThinkingMode,
) {
if matches!(thinking, ThinkingMode::Disabled) {
messages.insert(0, Message::system(DISABLED_THINKING_SYSTEM_MESSAGE));
}
}
fn attach_tools_to_body(
body: &mut serde_json::Value,
tools: &Option<Vec<ToolDefinition>>,
native_tool_choice: bool,
) {
tool_calling::attach_tools(body, tools, native_tool_choice);
}
const RESERVED_EXTRA_BODY_KEYS: &[&str] = &[
"model",
"messages",
"tools",
"tool_choice",
"stream",
"max_tokens",
"temperature",
"thinking",
];
const ALLOWED_EXTRA_BODY_KEYS: &[&str] = &[
"top_p",
"top_k",
"min_p",
"repetition_penalty",
"frequency_penalty",
"presence_penalty",
"seed",
"stop",
"chat_template_kwargs",
"guided_json",
"guided_regex",
"guided_choice",
"skip_special_tokens",
"spaces_between_special_tokens",
"add_generation_prompt",
"best_of",
"use_beam_search",
"length_penalty",
"early_stopping",
"provider",
"models",
"route",
"transforms",
];
pub(crate) fn merge_extra_body(
body: &mut serde_json::Value,
extra_body: Option<&serde_json::Map<String, serde_json::Value>>,
context: &str,
) -> Result<()> {
let Some(extra_body) = extra_body else {
return Ok(());
};
let body_obj = body
.as_object_mut()
.context("request body must be a JSON object")?;
for key in extra_body.keys() {
let k = key.as_str();
if RESERVED_EXTRA_BODY_KEYS.contains(&k) {
bail!("{} extra_body cannot override reserved key: {}", context, k);
}
if !ALLOWED_EXTRA_BODY_KEYS.contains(&k) {
bail!(
"{} extra_body contains disallowed key '{}'. \
Only sampling and backend-specific parameters are permitted. \
Allowed keys: {}",
context,
k,
ALLOWED_EXTRA_BODY_KEYS.join(", ")
);
}
}
for (key, value) in extra_body {
body_obj.insert(key.clone(), value.clone());
}
Ok(())
}
#[async_trait]
pub trait LlmClient: Send + Sync {
async fn chat(
&self,
messages: Vec<Message>,
tools: Option<Vec<ToolDefinition>>,
thinking: ThinkingMode,
) -> Result<ChatResponse>;
async fn chat_stream(
&self,
messages: Vec<Message>,
tools: Option<Vec<ToolDefinition>>,
thinking: ThinkingMode,
) -> Result<StreamingResponse>;
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum ThinkingMode {
Enabled,
Disabled,
Budget(usize),
}
#[cfg(test)]
#[path = "tests.rs"]
mod tests;
#[cfg(test)]
#[path = "../../tests/unit/api/message_compat_test.rs"]
mod message_compat_test;
#[cfg(test)]
#[path = "mock.rs"]
pub mod mock;