use crate::error::LlmError;
use async_trait::async_trait;
use futures_core::Stream;
use serde::{Deserialize, Serialize};
use std::pin::Pin;
use std::time::Duration;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
pub role: Role,
pub content: String,
#[serde(default)]
pub tool_calls: Vec<ToolCall>,
#[serde(default)]
pub tool_call_id: Option<String>,
}
impl Message {
pub fn system(content: impl Into<String>) -> Self {
Self::text(Role::System, content)
}
pub fn user(content: impl Into<String>) -> Self {
Self::text(Role::User, content)
}
pub fn assistant(content: impl Into<String>) -> Self {
Self::text(Role::Assistant, content)
}
fn text(role: Role, content: impl Into<String>) -> Self {
Self {
role,
content: content.into(),
tool_calls: Vec::new(),
tool_call_id: None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
#[non_exhaustive]
pub enum Role {
System,
User,
Assistant,
Tool,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ToolCall {
pub id: String,
pub name: String,
pub args: serde_json::Value,
}
impl ToolCall {
pub fn new(id: impl Into<String>, name: impl Into<String>, args: serde_json::Value) -> Self {
Self {
id: id.into(),
name: name.into(),
args,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ToolDef {
pub name: String,
pub description: String,
pub json_schema: serde_json::Value,
}
impl ToolDef {
pub fn new(
name: impl Into<String>,
description: impl Into<String>,
json_schema: serde_json::Value,
) -> Self {
Self {
name: name.into(),
description: description.into(),
json_schema,
}
}
}
#[derive(Debug, Clone)]
pub struct ChatRequest {
pub messages: Vec<Message>,
pub tools: Vec<ToolDef>,
pub temperature: Option<f32>,
pub max_tokens: Option<u32>,
pub response_format: ResponseFormat,
pub stop: Vec<String>,
pub timeout: Option<Duration>,
}
impl ChatRequest {
pub fn new(messages: Vec<Message>) -> Self {
Self {
messages,
tools: Vec::new(),
temperature: None,
max_tokens: None,
response_format: ResponseFormat::Text,
stop: Vec::new(),
timeout: None,
}
}
pub fn builder() -> ChatRequestBuilder {
ChatRequestBuilder::default()
}
}
#[derive(Default)]
pub struct ChatRequestBuilder {
messages: Vec<Message>,
tools: Vec<ToolDef>,
temperature: Option<f32>,
max_tokens: Option<u32>,
response_format: Option<ResponseFormat>,
timeout: Option<Duration>,
}
impl ChatRequestBuilder {
pub fn system(mut self, content: impl Into<String>) -> Self {
self.messages.push(Message::system(content));
self
}
pub fn user(mut self, content: impl Into<String>) -> Self {
self.messages.push(Message::user(content));
self
}
pub fn assistant(mut self, content: impl Into<String>) -> Self {
self.messages.push(Message::assistant(content));
self
}
pub fn message(mut self, message: Message) -> Self {
self.messages.push(message);
self
}
pub fn tools(mut self, tools: Vec<ToolDef>) -> Self {
self.tools = tools;
self
}
pub fn temperature(mut self, temperature: f32) -> Self {
self.temperature = Some(temperature);
self
}
pub fn max_tokens(mut self, max_tokens: u32) -> Self {
self.max_tokens = Some(max_tokens);
self
}
pub fn response_format(mut self, response_format: ResponseFormat) -> Self {
self.response_format = Some(response_format);
self
}
pub fn timeout(mut self, timeout: Duration) -> Self {
self.timeout = Some(timeout);
self
}
pub fn build(self) -> ChatRequest {
ChatRequest {
messages: self.messages,
tools: self.tools,
temperature: self.temperature,
max_tokens: self.max_tokens,
response_format: self.response_format.unwrap_or(ResponseFormat::Text),
stop: Vec::new(),
timeout: self.timeout,
}
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum ResponseFormat {
Text,
Json {
schema: serde_json::Value,
},
StructuredOutput {
schema: serde_json::Value,
},
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct ChatResponse {
pub message: Message,
pub usage: Usage,
pub finish_reason: FinishReason,
pub model: String,
}
impl ChatResponse {
pub fn new(
message: Message,
usage: Usage,
finish_reason: FinishReason,
model: impl Into<String>,
) -> Self {
Self {
message,
usage,
finish_reason,
model: model.into(),
}
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[non_exhaustive]
pub struct Usage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
#[serde(default)]
pub cache_read_tokens: u32,
#[serde(default)]
pub cache_creation_tokens: u32,
}
impl Usage {
pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
Self {
prompt_tokens,
completion_tokens,
cache_read_tokens: 0,
cache_creation_tokens: 0,
}
}
#[must_use]
pub fn with_cache_tokens(mut self, cache_read_tokens: u32, cache_creation_tokens: u32) -> Self {
self.cache_read_tokens = cache_read_tokens;
self.cache_creation_tokens = cache_creation_tokens;
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum FinishReason {
Stop,
ToolCalls,
Length,
ContentFilter,
Error,
}
#[derive(Debug, Clone, Default)]
pub struct Capabilities {
pub tool_calling: bool,
pub streaming: bool,
pub structured_output: bool,
pub embeddings: bool,
pub max_context_tokens: u32,
pub vision: bool,
}
impl Capabilities {
pub fn builder() -> CapabilitiesBuilder {
CapabilitiesBuilder(Capabilities::default())
}
}
#[derive(Debug, Clone, Default)]
pub struct CapabilitiesBuilder(Capabilities);
impl CapabilitiesBuilder {
pub fn tool_calling(mut self, v: bool) -> Self {
self.0.tool_calling = v;
self
}
pub fn streaming(mut self, v: bool) -> Self {
self.0.streaming = v;
self
}
pub fn structured_output(mut self, v: bool) -> Self {
self.0.structured_output = v;
self
}
pub fn embeddings(mut self, v: bool) -> Self {
self.0.embeddings = v;
self
}
pub fn max_context_tokens(mut self, v: u32) -> Self {
self.0.max_context_tokens = v;
self
}
pub fn vision(mut self, v: bool) -> Self {
self.0.vision = v;
self
}
pub fn build(self) -> Capabilities {
self.0
}
}
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct ChatChunk {
pub delta: String,
pub tool_calls: Vec<ToolCall>,
pub finish_reason: Option<FinishReason>,
pub usage: Option<Usage>,
}
impl ChatChunk {
pub fn new(
delta: String,
tool_calls: Vec<ToolCall>,
finish_reason: Option<FinishReason>,
usage: Option<Usage>,
) -> Self {
Self {
delta,
tool_calls,
finish_reason,
usage,
}
}
}
pub type ChunkStream = Pin<Box<dyn Stream<Item = Result<ChatChunk, LlmError>> + Send + 'static>>;
pub type Embedding = Vec<f32>;
#[async_trait]
pub trait LlmClient: Send + Sync {
fn name(&self) -> &str;
fn capabilities(&self) -> &Capabilities;
async fn complete(&self, req: ChatRequest) -> Result<ChatResponse, LlmError>;
async fn stream(&self, req: ChatRequest) -> Result<ChunkStream, LlmError>;
async fn embed(&self, texts: &[String]) -> Result<Vec<Embedding>, LlmError>;
}
#[cfg(test)]
mod tests {
use super::*;
#[allow(dead_code)]
fn _assert_dyn_compatible(_: &dyn LlmClient) {}
#[test]
fn chat_request_default_has_no_tools() {
let req = ChatRequest::new(vec![]);
assert!(req.tools.is_empty());
assert!(matches!(req.response_format, ResponseFormat::Text));
}
#[test]
fn message_constructors_set_role_and_clear_tool_fields() {
let m = Message::system("be terse");
assert_eq!(m.role, Role::System);
assert_eq!(m.content, "be terse");
assert!(m.tool_calls.is_empty());
assert!(m.tool_call_id.is_none());
let user = Message::user("hi");
assert_eq!((user.role, user.content.as_str()), (Role::User, "hi"));
let assistant = Message::assistant("ok");
assert_eq!(
(assistant.role, assistant.content.as_str()),
(Role::Assistant, "ok")
);
}
#[test]
fn chat_request_builder_sets_every_option_and_appends_raw_messages() {
let timeout = Duration::from_secs(7);
let schema = serde_json::json!({"type": "object"});
let req = ChatRequest::builder()
.assistant("a")
.message(Message::user("raw"))
.tools(vec![ToolDef::new("t", "desc", serde_json::json!({}))])
.max_tokens(42)
.response_format(ResponseFormat::Json {
schema: schema.clone(),
})
.timeout(timeout)
.build();
let roles: Vec<Role> = req.messages.iter().map(|m| m.role).collect();
assert_eq!(roles, vec![Role::Assistant, Role::User]);
assert_eq!(req.messages[1].content, "raw");
assert_eq!(req.tools.len(), 1);
assert_eq!(req.max_tokens, Some(42));
assert_eq!(req.timeout, Some(timeout));
assert!(matches!(req.response_format, ResponseFormat::Json { .. }));
}
#[test]
fn chat_request_builder_orders_messages_and_keeps_defaults() {
let req = ChatRequest::builder()
.system("sys")
.user("u")
.temperature(0.5)
.build();
let roles: Vec<Role> = req.messages.iter().map(|m| m.role).collect();
assert_eq!(roles, vec![Role::System, Role::User]);
assert_eq!(req.temperature, Some(0.5));
assert!(req.tools.is_empty());
assert!(matches!(req.response_format, ResponseFormat::Text));
}
#[test]
fn capabilities_builder_default_matches_struct_default() {
let built = Capabilities::builder().build();
let direct = Capabilities::default();
assert_eq!(built.tool_calling, direct.tool_calling);
assert_eq!(built.streaming, direct.streaming);
assert_eq!(built.structured_output, direct.structured_output);
assert_eq!(built.embeddings, direct.embeddings);
assert_eq!(built.max_context_tokens, direct.max_context_tokens);
assert_eq!(built.vision, direct.vision);
}
#[test]
fn capabilities_builder_sets_tool_calling() {
let c = Capabilities::builder().tool_calling(true).build();
assert!(c.tool_calling);
}
#[test]
fn capabilities_builder_sets_streaming() {
let c = Capabilities::builder().streaming(true).build();
assert!(c.streaming);
}
#[test]
fn capabilities_builder_sets_structured_output() {
let c = Capabilities::builder().structured_output(true).build();
assert!(c.structured_output);
}
#[test]
fn capabilities_builder_sets_embeddings() {
let c = Capabilities::builder().embeddings(true).build();
assert!(c.embeddings);
}
#[test]
fn capabilities_builder_sets_max_context_tokens() {
let c = Capabilities::builder().max_context_tokens(8000).build();
assert_eq!(c.max_context_tokens, 8000);
}
#[test]
fn capabilities_builder_sets_vision() {
let c = Capabilities::builder().vision(true).build();
assert!(c.vision);
}
#[test]
fn capabilities_builder_chains_all_setters() {
let c = Capabilities::builder()
.tool_calling(true)
.streaming(true)
.structured_output(true)
.embeddings(true)
.max_context_tokens(32_000)
.vision(false)
.build();
assert!(c.tool_calling && c.streaming && c.structured_output && c.embeddings);
assert_eq!(c.max_context_tokens, 32_000);
assert!(!c.vision);
}
#[test]
fn chat_chunk_usage_defaults_to_none_in_struct_literal() {
let chunk = ChatChunk {
delta: String::new(),
tool_calls: vec![],
finish_reason: None,
usage: None,
};
assert!(chunk.usage.is_none());
}
#[test]
fn chat_chunk_with_usage_round_trips() {
let chunk = ChatChunk {
delta: "done".into(),
tool_calls: vec![],
finish_reason: Some(FinishReason::Stop),
usage: Some(Usage::new(10, 32)),
};
let u = chunk.usage.as_ref().expect("usage set");
assert_eq!(u.prompt_tokens, 10);
assert_eq!(u.completion_tokens, 32);
}
#[test]
fn usage_new_sets_token_counts() {
let u = Usage::new(12, 34);
assert_eq!(u.prompt_tokens, 12);
assert_eq!(u.completion_tokens, 34);
}
#[test]
fn usage_cache_tokens_default_to_zero() {
let u = Usage::new(12, 34);
assert_eq!(u.cache_read_tokens, 0);
assert_eq!(u.cache_creation_tokens, 0);
}
#[test]
fn usage_deserializes_legacy_json_without_cache_fields() {
let legacy = r#"{"prompt_tokens":10,"completion_tokens":5}"#;
let u: Usage = serde_json::from_str(legacy).unwrap();
assert_eq!(u.cache_read_tokens, 0);
assert_eq!(u.cache_creation_tokens, 0);
}
#[test]
fn usage_with_cache_tokens_sets_cache_fields() {
let u = Usage::new(12, 34).with_cache_tokens(100, 8);
assert_eq!(u.prompt_tokens, 12);
assert_eq!(u.completion_tokens, 34);
assert_eq!(u.cache_read_tokens, 100);
assert_eq!(u.cache_creation_tokens, 8);
}
#[test]
fn tooldef_new_sets_fields() {
let d = ToolDef::new("echo", "echoes input", serde_json::json!({"type":"object"}));
assert_eq!(d.name, "echo");
assert_eq!(d.description, "echoes input");
assert_eq!(d.json_schema, serde_json::json!({"type":"object"}));
}
#[test]
fn toolcall_new_sets_fields() {
let c = ToolCall::new("id-1", "search", serde_json::json!({"q":"x"}));
assert_eq!(c.id, "id-1");
assert_eq!(c.name, "search");
assert_eq!(c.args, serde_json::json!({"q":"x"}));
}
#[test]
fn chatresponse_new_carries_model() {
let msg = Message {
role: Role::Assistant,
content: "hi".into(),
tool_calls: vec![],
tool_call_id: None,
};
let r = ChatResponse::new(
msg,
Usage::new(1, 2),
FinishReason::Stop,
"anthropic:claude-sonnet-4-6",
);
assert_eq!(r.model, "anthropic:claude-sonnet-4-6");
assert_eq!(r.finish_reason, FinishReason::Stop);
assert_eq!(r.message.content, "hi");
assert_eq!(r.usage.prompt_tokens, 1);
}
}