Skip to main content

llm/
provider.rs

1use crate::LlmModel;
2use crate::Result as LlmResult;
3use std::future::Future;
4use std::pin::Pin;
5use tokio_stream::Stream;
6
7use super::{Context, LlmResponse};
8
9/// A stream of [`LlmResponse`] events from an LLM provider.
10///
11/// This is a pinned, boxed, `Send` stream used as the return type of
12/// [`StreamingModelProvider::stream_response`]. Boxing is required to support
13/// trait objects (`Vec<Box<dyn StreamingModelProvider>>`) in types like
14/// [`AlloyedModelProvider`](crate::alloyed::AlloyedModelProvider).
15pub type LlmResponseStream = Pin<Box<dyn Stream<Item = LlmResult<LlmResponse>> + Send>>;
16
17#[doc = include_str!("docs/provider_factory.md")]
18pub trait ProviderFactory: Sized {
19    /// Create provider from environment variables and default configuration
20    fn from_env() -> impl Future<Output = LlmResult<Self>> + Send;
21
22    /// Set or update the model for this provider (builder pattern)
23    fn with_model(self, model: &str) -> Self;
24}
25
26#[doc = include_str!("docs/streaming_model_provider.md")]
27pub trait StreamingModelProvider: Send + Sync {
28    fn stream_response(&self, context: &Context) -> LlmResponseStream;
29    fn display_name(&self) -> String;
30
31    /// Context window size in tokens for the current model.
32    /// Returns `None` for unknown models (e.g. Ollama, `LlamaCpp`).
33    fn context_window(&self) -> Option<u32>;
34
35    /// The `LlmModel` this provider is currently configured to use.
36    /// Returns `None` for providers where the model is unknown at compile time
37    /// (e.g. test fakes).
38    fn model(&self) -> Option<LlmModel> {
39        None
40    }
41}
42
43/// Look up context window for a known provider + model ID combo via the catalog.
44///
45/// Returns `None` if the model is not in the catalog.
46pub fn get_context_window(provider: &str, model_id: &str) -> Option<u32> {
47    let key = format!("{provider}:{model_id}");
48    key.parse::<LlmModel>().ok().and_then(|m| m.context_window())
49}
50
51impl StreamingModelProvider for Box<dyn StreamingModelProvider> {
52    fn stream_response(&self, context: &Context) -> LlmResponseStream {
53        (**self).stream_response(context)
54    }
55
56    fn display_name(&self) -> String {
57        (**self).display_name()
58    }
59
60    fn context_window(&self) -> Option<u32> {
61        (**self).context_window()
62    }
63
64    fn model(&self) -> Option<LlmModel> {
65        (**self).model()
66    }
67}
68
69impl<T: StreamingModelProvider> StreamingModelProvider for std::sync::Arc<T> {
70    fn stream_response(&self, context: &Context) -> LlmResponseStream {
71        (**self).stream_response(context)
72    }
73
74    fn display_name(&self) -> String {
75        (**self).display_name()
76    }
77
78    fn context_window(&self) -> Option<u32> {
79        (**self).context_window()
80    }
81
82    fn model(&self) -> Option<LlmModel> {
83        (**self).model()
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn lookup_context_window_known_model() {
93        assert_eq!(get_context_window("anthropic", "claude-opus-4-6"), Some(1_000_000));
94    }
95
96    #[test]
97    fn lookup_context_window_openrouter_model() {
98        // OpenRouter Qwen models should resolve from catalog
99        let result = get_context_window("openrouter", "anthropic/claude-opus-4");
100        assert_eq!(result, Some(200_000));
101    }
102
103    #[test]
104    fn lookup_context_window_unknown_model() {
105        assert_eq!(get_context_window("anthropic", "unknown-model-xyz"), None);
106    }
107
108    #[test]
109    fn lookup_context_window_unknown_provider() {
110        assert_eq!(get_context_window("unknown-provider", "some-model"), None);
111    }
112}