ai_agents_core/traits/
llm.rs

1//! LLM provider traits
2
3use async_trait::async_trait;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use thiserror::Error;
7
8use crate::message::ChatMessage;
9use crate::types::{LLMChunk, LLMConfig, LLMFeature, LLMResponse, LLMToolRequest, ToolChoice};
10
11/// Core LLM provider trait.
12///
13/// Implement this to integrate a custom LLM backend. Most users can use
14/// `UnifiedLLMProvider` which supports OpenAI, Anthropic, and other providers
15/// out of the box.
16#[async_trait]
17pub trait LLMProvider: Send + Sync {
18    /// Send messages and get a complete response.
19    async fn complete(
20        &self,
21        messages: &[ChatMessage],
22        config: Option<&LLMConfig>,
23    ) -> Result<LLMResponse, LLMError>;
24
25    /// Send messages with provider-native tool definitions.
26    async fn complete_with_tools(
27        &self,
28        _messages: &[ChatMessage],
29        _config: Option<&LLMConfig>,
30        _request: &LLMToolRequest,
31    ) -> Result<LLMResponse, LLMError> {
32        Err(LLMError::Other(format!(
33            "provider '{}' does not support native tool completion",
34            self.provider_name()
35        )))
36    }
37
38    /// Returns a provider-level tool choice override when configured.
39    fn configured_tool_choice(&self) -> Option<ToolChoice> {
40        None
41    }
42
43    /// Reports whether this provider can enforce a native tool choice.
44    fn supports_tool_choice(&self, _choice: &ToolChoice) -> bool {
45        false
46    }
47
48    /// Send messages and get a streaming response.
49    async fn complete_stream(
50        &self,
51        messages: &[ChatMessage],
52        config: Option<&LLMConfig>,
53    ) -> Result<Box<dyn futures::Stream<Item = Result<LLMChunk, LLMError>> + Unpin + Send>, LLMError>;
54
55    /// Provider identifier (e.g. `"openai"`, `"anthropic"`).
56    fn provider_name(&self) -> &str;
57
58    /// Check if this provider supports a given feature.
59    fn supports(&self, feature: LLMFeature) -> bool;
60
61    /// Reports whether a provider error is terminal for this exact attempt.
62    ///
63    /// Terminal errors represent local configuration, protocol, or history failures that retry,
64    /// provider fallback, and static success responses must not hide. The default preserves the
65    /// behavior of existing custom providers by classifying every error as recoverable.
66    fn is_terminal_error(&self, _error: &LLMError) -> bool {
67        false
68    }
69}
70
71/// Higher-level LLM capabilities for agent operations
72#[async_trait]
73pub trait LLMCapability: Send + Sync {
74    async fn select_tool(
75        &self,
76        context: &TaskContext,
77        user_input: &str,
78    ) -> Result<ToolSelection, LLMError>;
79
80    async fn generate_tool_args(
81        &self,
82        tool_id: &str,
83        user_input: &str,
84        schema: &serde_json::Value,
85    ) -> Result<serde_json::Value, LLMError>;
86
87    async fn evaluate_yesno(
88        &self,
89        question: &str,
90        context: &TaskContext,
91    ) -> Result<(bool, String), LLMError>;
92
93    async fn classify(&self, input: &str, categories: &[String])
94    -> Result<(String, f32), LLMError>;
95
96    async fn process_task(
97        &self,
98        context: &TaskContext,
99        system_prompt: &str,
100    ) -> Result<LLMResponse, LLMError>;
101}
102
103/// Task context for LLM operations
104#[derive(Debug, Clone, Serialize, Deserialize)]
105pub struct TaskContext {
106    pub current_state: Option<String>,
107    pub available_tools: Vec<String>,
108    pub memory_slots: HashMap<String, serde_json::Value>,
109    pub recent_messages: Vec<ChatMessage>,
110}
111
112/// Tool selection result
113#[derive(Debug, Clone, Serialize, Deserialize)]
114pub struct ToolSelection {
115    pub tool_id: String,
116    pub confidence: f32,
117    pub reasoning: Option<String>,
118}
119
120/// LLM error types
121#[derive(Debug, Error)]
122pub enum LLMError {
123    #[error("API error: {message}")]
124    API {
125        message: String,
126        status: Option<u16>,
127    },
128
129    #[error("Network error: {0}")]
130    Network(String),
131
132    #[error("Rate limit exceeded: {retry_after:?}")]
133    RateLimit {
134        retry_after: Option<std::time::Duration>,
135    },
136
137    #[error("Configuration error: {0}")]
138    Config(String),
139
140    #[error("Model not found: {0}")]
141    ModelNotFound(String),
142
143    #[error("Content filtered: {0}")]
144    ContentFiltered(String),
145
146    #[error("Serialization error: {0}")]
147    Serialization(String),
148
149    #[error("Other error: {0}")]
150    Other(String),
151}
152
153impl From<serde_json::Error> for LLMError {
154    fn from(err: serde_json::Error) -> Self {
155        LLMError::Serialization(err.to_string())
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162    use crate::types::FinishReason;
163
164    struct LegacyProvider;
165
166    struct TerminalProvider;
167
168    #[async_trait]
169    impl LLMProvider for LegacyProvider {
170        async fn complete(
171            &self,
172            _messages: &[ChatMessage],
173            _config: Option<&LLMConfig>,
174        ) -> Result<LLMResponse, LLMError> {
175            Ok(LLMResponse::new("legacy", FinishReason::Stop))
176        }
177
178        async fn complete_stream(
179            &self,
180            _messages: &[ChatMessage],
181            _config: Option<&LLMConfig>,
182        ) -> Result<
183            Box<dyn futures::Stream<Item = Result<LLMChunk, LLMError>> + Unpin + Send>,
184            LLMError,
185        > {
186            Ok(Box::new(futures::stream::empty()))
187        }
188
189        fn provider_name(&self) -> &str {
190            "legacy"
191        }
192
193        fn supports(&self, _feature: LLMFeature) -> bool {
194            false
195        }
196    }
197
198    #[async_trait]
199    impl LLMProvider for TerminalProvider {
200        async fn complete(
201            &self,
202            messages: &[ChatMessage],
203            config: Option<&LLMConfig>,
204        ) -> Result<LLMResponse, LLMError> {
205            LegacyProvider.complete(messages, config).await
206        }
207
208        async fn complete_stream(
209            &self,
210            messages: &[ChatMessage],
211            config: Option<&LLMConfig>,
212        ) -> Result<
213            Box<dyn futures::Stream<Item = Result<LLMChunk, LLMError>> + Unpin + Send>,
214            LLMError,
215        > {
216            LegacyProvider.complete_stream(messages, config).await
217        }
218
219        fn provider_name(&self) -> &str {
220            "terminal"
221        }
222
223        fn supports(&self, _feature: LLMFeature) -> bool {
224            false
225        }
226
227        fn is_terminal_error(&self, error: &LLMError) -> bool {
228            matches!(error, LLMError::Config(_) | LLMError::Serialization(_))
229        }
230    }
231
232    #[test]
233    fn additive_tool_methods_preserve_legacy_implementations() {
234        let provider = LegacyProvider;
235        let request = LLMToolRequest {
236            tools: Vec::new(),
237            choice: ToolChoice::Auto,
238        };
239
240        assert!(provider.configured_tool_choice().is_none());
241        assert!(!provider.supports_tool_choice(&ToolChoice::Auto));
242        assert!(!provider.is_terminal_error(&LLMError::Config("legacy".to_string())));
243        let error = futures::executor::block_on(provider.complete_with_tools(&[], None, &request))
244            .unwrap_err();
245        assert!(
246            error
247                .to_string()
248                .contains("does not support native tool completion")
249        );
250    }
251
252    #[test]
253    fn terminal_error_classification_is_additive_and_provider_owned() {
254        let provider = TerminalProvider;
255
256        assert!(provider.is_terminal_error(&LLMError::Config("bad config".to_string())));
257        assert!(provider.is_terminal_error(&LLMError::Serialization("bad history".to_string())));
258        assert!(!provider.is_terminal_error(&LLMError::Network("temporary".to_string())));
259    }
260}
ai_agents_core/traits/llm.rs

ai_agents_core/traits/
llm.rs
