telos_agent/model/provider/
traits.rs

1//! The [`ModelProvider`] trait and the [`ErasedProvider`] helper.
2
3use async_stream::try_stream;
4use async_trait::async_trait;
5use futures_core::stream::Stream;
6
7use crate::error::AgentError;
8
9use super::types::{CompletionRequest, CompletionResponse, ProviderEvent};
10
11/// Abstract LLM backend. Implement this to plug in a new model provider.
12#[async_trait]
13pub trait ModelProvider: Send + Sync {
14    /// Issue a single non-streaming completion request.
15    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, AgentError>;
16
17    /// Stream a completion as a sequence of [`ProviderEvent`]s.
18    ///
19    /// The default implementation calls [`complete`](Self::complete) and
20    /// re-emits the result as one synthetic stream — providers that genuinely
21    /// stream should override this for incremental delivery.
22    fn stream_complete<'a>(
23        &'a self,
24        request: CompletionRequest,
25    ) -> std::pin::Pin<Box<dyn Stream<Item = Result<ProviderEvent, AgentError>> + Send + 'a>> {
26        Box::pin(try_stream! {
27            let response = self.complete(request).await?;
28            yield ProviderEvent::MessageStart;
29            for block in &response.message.blocks {
30                match block {
31                    crate::model::message::ContentBlock::Text(text) => {
32                        yield ProviderEvent::TextDelta(text.text.clone());
33                    }
34                    crate::model::message::ContentBlock::Thinking(thinking) => {
35                        yield ProviderEvent::ThinkingDelta(thinking.text.clone());
36                    }
37                    crate::model::message::ContentBlock::ToolCall(call) => {
38                        yield ProviderEvent::ToolCall(call.clone());
39                    }
40                    crate::model::message::ContentBlock::ToolResult(_) => {}
41                }
42            }
43            yield ProviderEvent::MessageStop {
44                stop_reason: response.stop_reason,
45                usage: response.usage,
46                model: response.model,
47            };
48        })
49    }
50
51    /// Return the maximum number of tokens that can be requested from this provider.
52    ///
53    /// The default implementation returns 128,000 tokens, which is the maximum
54    /// supported by the `cl100k_base` tokenizer used by most providers. Providers
55    /// with a different tokenizer (e.g. Gemini's SentencePiece) can override this.
56    fn max_tokens(&self) -> u32 {
57        128_000
58    }
59
60    /// Estimate the number of tokens for the given text.
61    ///
62    /// The default implementation uses the `cl100k_base` tokenizer (via
63    /// `tiktoken-rs`). Since DeepSeek, Kimi, and most OpenAI-compatible
64    /// providers all use cl100k_base-compatible BPE tokenizers, this gives
65    /// ±5% accuracy across all built-in providers.
66    ///
67    /// Providers with a different tokenizer (e.g. Gemini's SentencePiece)
68    /// can override this.
69    fn estimate_tokens(&self, text: &str) -> usize {
70        crate::model::tokens::count_tokens(text)
71    }
72}
73
74/// A newtype that implements [`ModelProvider`] by delegating to an erased
75/// `&dyn ModelProvider` reference. This is useful when adapting an existing
76/// borrowed provider to APIs that require a concrete [`ModelProvider`] type.
77pub struct ErasedProvider<'a>(pub &'a (dyn ModelProvider + 'a));
78
79#[async_trait]
80impl ModelProvider for ErasedProvider<'_> {
81    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, AgentError> {
82        self.0.complete(request).await
83    }
84
85    fn stream_complete<'a>(
86        &'a self,
87        request: CompletionRequest,
88    ) -> std::pin::Pin<Box<dyn Stream<Item = Result<ProviderEvent, AgentError>> + Send + 'a>> {
89        self.0.stream_complete(request)
90    }
91
92    fn max_tokens(&self) -> u32 {
93        self.0.max_tokens()
94    }
95
96    fn estimate_tokens(&self, text: &str) -> usize {
97        self.0.estimate_tokens(text)
98    }
99}
telos_agent/model/provider/traits.rs

telos_agent/model/provider/
traits.rs
