trae_agent_rs_core/llm/
client.rs

1//! LLM client trait and response structures
2
3use crate::error::{ LlmError, Result };
4use crate::tools::ToolCall;
5use async_trait::async_trait;
6use serde::{ Deserialize, Serialize };
7use std::collections::HashMap;
8
9use super::message::LlmMessage;
10
11/// Trait for LLM clients
12#[async_trait]
13pub trait LlmClient: Send + Sync {
14    /// Send a chat completion request
15    async fn chat_completion(
16        &self,
17        messages: Vec<LlmMessage>,
18        tools: Option<Vec<ToolDefinition>>,
19        options: Option<ChatOptions>
20    ) -> Result<LlmResponse>;
21
22    /// Get the model name
23    fn model_name(&self) -> &str;
24
25    /// Get the provider name
26    fn provider_name(&self) -> &str;
27
28    /// Check if the client supports streaming
29    fn supports_streaming(&self) -> bool {
30        false
31    }
32
33    /// Send a streaming chat completion request
34    async fn chat_completion_stream(
35        &self,
36        _messages: Vec<LlmMessage>,
37        _tools: Option<Vec<ToolDefinition>>,
38        _options: Option<ChatOptions>
39    ) -> Result<Box<dyn futures::Stream<Item = Result<LlmStreamChunk>> + Send + Unpin + '_>> {
40        Err(
41            (LlmError::InvalidRequest {
42                message: "Streaming not supported by this client".to_string(),
43            }).into()
44        )
45    }
46}
47
48/// Response from an LLM
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct LlmResponse {
51    /// The generated message
52    pub message: LlmMessage,
53
54    /// Usage statistics
55    pub usage: Option<Usage>,
56
57    /// Model used for generation
58    pub model: String,
59
60    /// Finish reason
61    pub finish_reason: Option<FinishReason>,
62
63    /// Additional metadata
64    pub metadata: Option<HashMap<String, serde_json::Value>>,
65}
66
67/// Streaming chunk from an LLM
68#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct LlmStreamChunk {
70    /// Delta content
71    pub delta: Option<String>,
72
73    /// Tool calls in this chunk
74    pub tool_calls: Option<Vec<ToolCall>>,
75
76    /// Finish reason if this is the last chunk
77    pub finish_reason: Option<FinishReason>,
78
79    /// Usage statistics (usually only in the last chunk)
80    pub usage: Option<Usage>,
81}
82
83/// Usage statistics for a request
84#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct Usage {
86    /// Number of tokens in the prompt
87    pub prompt_tokens: u32,
88
89    /// Number of tokens in the completion
90    pub completion_tokens: u32,
91
92    /// Total number of tokens
93    pub total_tokens: u32,
94}
95
96/// Reason why generation finished
97#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
98#[serde(rename_all = "snake_case")]
99pub enum FinishReason {
100    /// Generation completed naturally
101    Stop,
102
103    /// Hit the maximum token limit
104    Length,
105
106    /// Model decided to call a tool
107    ToolCalls,
108
109    /// Content was filtered
110    ContentFilter,
111
112    /// Other reason
113    Other(String),
114}
115
116/// Tool definition for function calling
117#[derive(Debug, Clone, Serialize, Deserialize)]
118pub struct ToolDefinition {
119    /// Type of tool (usually "function")
120    #[serde(rename = "type")]
121    pub tool_type: String,
122
123    /// Function definition
124    pub function: FunctionDefinition,
125}
126
127/// Function definition for tool calling
128#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct FunctionDefinition {
130    /// Name of the function
131    pub name: String,
132
133    /// Description of what the function does
134    pub description: String,
135
136    /// JSON schema for the function parameters
137    pub parameters: serde_json::Value,
138}
139
140/// Options for chat completion
141#[derive(Debug, Clone, Serialize, Deserialize)]
142pub struct ChatOptions {
143    /// Maximum tokens to generate
144    pub max_tokens: Option<u32>,
145
146    /// Temperature for generation
147    pub temperature: Option<f32>,
148
149    /// Top-p sampling parameter
150    pub top_p: Option<f32>,
151
152    /// Top-k sampling parameter
153    pub top_k: Option<u32>,
154
155    /// Stop sequences
156    pub stop: Option<Vec<String>>,
157
158    /// Whether to stream the response
159    pub stream: Option<bool>,
160
161    /// Tool choice strategy
162    pub tool_choice: Option<ToolChoice>,
163}
164
165/// Tool choice strategy
166#[derive(Debug, Clone, Serialize, Deserialize)]
167#[serde(untagged)]
168pub enum ToolChoice {
169    /// Let the model decide
170    Auto,
171
172    /// Never use tools
173    None,
174
175    /// Force use of a specific tool
176    Required {
177        name: String,
178    },
179}
180
181impl Default for ChatOptions {
182    fn default() -> Self {
183        Self {
184            max_tokens: Some(8192),
185            temperature: Some(0.7),
186            top_p: Some(1.0),
187            top_k: None,
188            stop: None,
189            stream: Some(false),
190            tool_choice: Some(ToolChoice::Auto),
191        }
192    }
193}
trae_agent_rs_core/llm/client.rs

trae_agent_rs_core/llm/
client.rs
