llm/
provider.rs

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

llm/
provider.rs
