//! LLM provider trait and types. use async_trait::async_trait; use rust_decimal::Decimal; use serde::{Deserialize, Serialize}; use crate::error::LlmError; /// Role in a conversation. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum Role { System, User, Assistant, Tool, } /// A message in a conversation. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ChatMessage { pub role: Role, pub content: String, /// Tool call ID if this is a tool result message. #[serde(skip_serializing_if = "Option::is_none")] pub tool_call_id: Option, /// Name of the tool for tool results. #[serde(skip_serializing_if = "Option::is_none")] pub name: Option, } impl ChatMessage { /// Create a system message. pub fn system(content: impl Into) -> Self { Self { role: Role::System, content: content.into(), tool_call_id: None, name: None, } } /// Create a user message. pub fn user(content: impl Into) -> Self { Self { role: Role::User, content: content.into(), tool_call_id: None, name: None, } } /// Create an assistant message. pub fn assistant(content: impl Into) -> Self { Self { role: Role::Assistant, content: content.into(), tool_call_id: None, name: None, } } /// Create a tool result message. pub fn tool_result( tool_call_id: impl Into, name: impl Into, content: impl Into, ) -> Self { Self { role: Role::Tool, content: content.into(), tool_call_id: Some(tool_call_id.into()), name: Some(name.into()), } } } /// Request for a chat completion. #[derive(Debug, Clone)] pub struct CompletionRequest { pub messages: Vec, pub max_tokens: Option, pub temperature: Option, pub stop_sequences: Option>, } impl CompletionRequest { /// Create a new completion request. pub fn new(messages: Vec) -> Self { Self { messages, max_tokens: None, temperature: None, stop_sequences: None, } } /// Set max tokens. pub fn with_max_tokens(mut self, max_tokens: u32) -> Self { self.max_tokens = Some(max_tokens); self } /// Set temperature. pub fn with_temperature(mut self, temperature: f32) -> Self { self.temperature = Some(temperature); self } } /// Response from a chat completion. #[derive(Debug, Clone)] pub struct CompletionResponse { pub content: String, pub input_tokens: u32, pub output_tokens: u32, pub finish_reason: FinishReason, } /// Why the completion finished. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum FinishReason { Stop, Length, ToolUse, ContentFilter, Unknown, } /// Definition of a tool for the LLM. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolDefinition { pub name: String, pub description: String, pub parameters: serde_json::Value, } /// A tool call requested by the LLM. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolCall { pub id: String, pub name: String, pub arguments: serde_json::Value, } /// Result of a tool execution to send back to the LLM. #[derive(Debug, Clone)] pub struct ToolResult { pub tool_call_id: String, pub name: String, pub content: String, pub is_error: bool, } /// Request for a completion with tool use. #[derive(Debug, Clone)] pub struct ToolCompletionRequest { pub messages: Vec, pub tools: Vec, pub max_tokens: Option, pub temperature: Option, /// How to handle tool use: "auto", "required", or "none". pub tool_choice: Option, } impl ToolCompletionRequest { /// Create a new tool completion request. pub fn new(messages: Vec, tools: Vec) -> Self { Self { messages, tools, max_tokens: None, temperature: None, tool_choice: None, } } /// Set max tokens. pub fn with_max_tokens(mut self, max_tokens: u32) -> Self { self.max_tokens = Some(max_tokens); self } /// Set temperature. pub fn with_temperature(mut self, temperature: f32) -> Self { self.temperature = Some(temperature); self } /// Set tool choice mode. pub fn with_tool_choice(mut self, choice: impl Into) -> Self { self.tool_choice = Some(choice.into()); self } } /// Response from a completion with potential tool calls. #[derive(Debug, Clone)] pub struct ToolCompletionResponse { /// Text content (may be empty if tool calls are present). pub content: Option, /// Tool calls requested by the model. pub tool_calls: Vec, pub input_tokens: u32, pub output_tokens: u32, pub finish_reason: FinishReason, } /// Trait for LLM providers. #[async_trait] pub trait LlmProvider: Send + Sync { /// Get the model name. fn model_name(&self) -> &str; /// Get cost per token (input, output). fn cost_per_token(&self) -> (Decimal, Decimal); /// Complete a chat conversation. async fn complete(&self, request: CompletionRequest) -> Result; /// Complete with tool use support. async fn complete_with_tools( &self, request: ToolCompletionRequest, ) -> Result; /// List available models from the provider. /// Default implementation returns empty list. async fn list_models(&self) -> Result, LlmError> { Ok(Vec::new()) } /// Calculate cost for a completion. fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal { let (input_cost, output_cost) = self.cost_per_token(); input_cost * Decimal::from(input_tokens) + output_cost * Decimal::from(output_tokens) } }