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
62/// Higher-level LLM capabilities for agent operations
63#[async_trait]
64pub trait LLMCapability: Send + Sync {
65    async fn select_tool(
66        &self,
67        context: &TaskContext,
68        user_input: &str,
69    ) -> Result<ToolSelection, LLMError>;
70
71    async fn generate_tool_args(
72        &self,
73        tool_id: &str,
74        user_input: &str,
75        schema: &serde_json::Value,
76    ) -> Result<serde_json::Value, LLMError>;
77
78    async fn evaluate_yesno(
79        &self,
80        question: &str,
81        context: &TaskContext,
82    ) -> Result<(bool, String), LLMError>;
83
84    async fn classify(&self, input: &str, categories: &[String])
85    -> Result<(String, f32), LLMError>;
86
87    async fn process_task(
88        &self,
89        context: &TaskContext,
90        system_prompt: &str,
91    ) -> Result<LLMResponse, LLMError>;
92}
93
94/// Task context for LLM operations
95#[derive(Debug, Clone, Serialize, Deserialize)]
96pub struct TaskContext {
97    pub current_state: Option<String>,
98    pub available_tools: Vec<String>,
99    pub memory_slots: HashMap<String, serde_json::Value>,
100    pub recent_messages: Vec<ChatMessage>,
101}
102
103/// Tool selection result
104#[derive(Debug, Clone, Serialize, Deserialize)]
105pub struct ToolSelection {
106    pub tool_id: String,
107    pub confidence: f32,
108    pub reasoning: Option<String>,
109}
110
111/// LLM error types
112#[derive(Debug, Error)]
113pub enum LLMError {
114    #[error("API error: {message}")]
115    API {
116        message: String,
117        status: Option<u16>,
118    },
119
120    #[error("Network error: {0}")]
121    Network(String),
122
123    #[error("Rate limit exceeded: {retry_after:?}")]
124    RateLimit {
125        retry_after: Option<std::time::Duration>,
126    },
127
128    #[error("Configuration error: {0}")]
129    Config(String),
130
131    #[error("Model not found: {0}")]
132    ModelNotFound(String),
133
134    #[error("Content filtered: {0}")]
135    ContentFiltered(String),
136
137    #[error("Serialization error: {0}")]
138    Serialization(String),
139
140    #[error("Other error: {0}")]
141    Other(String),
142}
143
144impl From<serde_json::Error> for LLMError {
145    fn from(err: serde_json::Error) -> Self {
146        LLMError::Serialization(err.to_string())
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use crate::types::FinishReason;
154
155    struct LegacyProvider;
156
157    #[async_trait]
158    impl LLMProvider for LegacyProvider {
159        async fn complete(
160            &self,
161            _messages: &[ChatMessage],
162            _config: Option<&LLMConfig>,
163        ) -> Result<LLMResponse, LLMError> {
164            Ok(LLMResponse::new("legacy", FinishReason::Stop))
165        }
166
167        async fn complete_stream(
168            &self,
169            _messages: &[ChatMessage],
170            _config: Option<&LLMConfig>,
171        ) -> Result<
172            Box<dyn futures::Stream<Item = Result<LLMChunk, LLMError>> + Unpin + Send>,
173            LLMError,
174        > {
175            Ok(Box::new(futures::stream::empty()))
176        }
177
178        fn provider_name(&self) -> &str {
179            "legacy"
180        }
181
182        fn supports(&self, _feature: LLMFeature) -> bool {
183            false
184        }
185    }
186
187    #[test]
188    fn additive_tool_methods_preserve_legacy_implementations() {
189        let provider = LegacyProvider;
190        let request = LLMToolRequest {
191            tools: Vec::new(),
192            choice: ToolChoice::Auto,
193        };
194
195        assert!(provider.configured_tool_choice().is_none());
196        assert!(!provider.supports_tool_choice(&ToolChoice::Auto));
197        let error = futures::executor::block_on(provider.complete_with_tools(&[], None, &request))
198            .unwrap_err();
199        assert!(
200            error
201                .to_string()
202                .contains("does not support native tool completion")
203        );
204    }
205}
ai_agents_core/traits/llm.rs

ai_agents_core/traits/
llm.rs
