llm/
provider.rs

1use crate::LlmError;
2use crate::LlmModel;
3use crate::ProviderConnectionConfig;
4use crate::Result as LlmResult;
5use std::future::Future;
6use std::pin::Pin;
7use tokio_stream::{Stream, StreamExt};
8
9use super::{Context, LlmResponse};
10
11/// A stream of [`LlmResponse`] events from an LLM provider.
12///
13/// This is a pinned, boxed, `Send` stream used as the return type of
14/// [`StreamingModelProvider::stream_response`]. Boxing is required to support
15/// trait objects (`Vec<Box<dyn StreamingModelProvider>>`) in types like
16/// [`AlloyedModelProvider`](crate::alloyed::AlloyedModelProvider).
17pub type LlmResponseStream = Pin<Box<dyn Stream<Item = LlmResult<LlmResponse>> + Send>>;
18
19#[doc = include_str!("docs/provider_factory.md")]
20pub trait ProviderFactory: Sized {
21    /// Create provider from environment variables and default configuration
22    fn from_env() -> impl Future<Output = LlmResult<Self>> + Send;
23
24    /// Create provider from environment variables with provider connection overrides.
25    fn from_env_with_connection(connection: ProviderConnectionConfig) -> impl Future<Output = LlmResult<Self>> + Send {
26        async move {
27            let _ = connection;
28            Self::from_env().await
29        }
30    }
31
32    /// Set or update the model for this provider (builder pattern)
33    fn with_model(self, model: &str) -> Self;
34}
35
36#[doc = include_str!("docs/streaming_model_provider.md")]
37pub trait StreamingModelProvider: Send + Sync {
38    fn stream_response(&self, context: &Context) -> LlmResponseStream;
39    fn display_name(&self) -> String;
40
41    /// Context window size in tokens for the current model.
42    /// Returns `None` for unknown models (e.g. Ollama, `LlamaCpp`).
43    fn context_window(&self) -> Option<u32>;
44
45    /// The `LlmModel` this provider is currently configured to use.
46    /// Returns `None` for providers where the model is unknown at compile time
47    /// (e.g. test fakes).
48    fn model(&self) -> Option<LlmModel> {
49        None
50    }
51}
52
53/// Look up context window for a known provider + model ID combo via the catalog.
54///
55/// Returns `None` if the model is not in the catalog.
56pub fn get_context_window(provider: &str, model_id: &str) -> Option<u32> {
57    let key = format!("{provider}:{model_id}");
58    key.parse::<LlmModel>().ok().and_then(|m| m.context_window())
59}
60
61/// Bridge a fallible request setup into an [`LlmResponseStream`].
62///
63/// `open` issues the request; `process` turns what it returns into a response
64/// stream. A setup failure becomes the stream's single item, so providers never
65/// hand-roll the yield-then-return dance and cannot drop an error on the way.
66pub(crate) fn stream_from<T, S>(
67    open: impl Future<Output = LlmResult<T>> + Send + 'static,
68    process: impl FnOnce(T) -> S + Send + 'static,
69) -> LlmResponseStream
70where
71    T: Send,
72    S: Stream<Item = LlmResult<LlmResponse>> + Send + 'static,
73{
74    Box::pin(async_stream::stream! {
75        match open.await {
76            Ok(opened) => {
77                let mut stream = Box::pin(process(opened));
78                while let Some(item) = stream.next().await {
79                    yield item;
80                }
81            }
82            Err(error) => yield Err(error),
83        }
84    })
85}
86
87/// A response stream whose only item is `error`.
88pub(crate) fn error_stream(error: LlmError) -> LlmResponseStream {
89    Box::pin(tokio_stream::once(Err(error)))
90}
91
92impl StreamingModelProvider for Box<dyn StreamingModelProvider> {
93    fn stream_response(&self, context: &Context) -> LlmResponseStream {
94        (**self).stream_response(context)
95    }
96
97    fn display_name(&self) -> String {
98        (**self).display_name()
99    }
100
101    fn context_window(&self) -> Option<u32> {
102        (**self).context_window()
103    }
104
105    fn model(&self) -> Option<LlmModel> {
106        (**self).model()
107    }
108}
109
110impl<T: StreamingModelProvider + ?Sized> StreamingModelProvider for std::sync::Arc<T> {
111    fn stream_response(&self, context: &Context) -> LlmResponseStream {
112        (**self).stream_response(context)
113    }
114
115    fn display_name(&self) -> String {
116        (**self).display_name()
117    }
118
119    fn context_window(&self) -> Option<u32> {
120        (**self).context_window()
121    }
122
123    fn model(&self) -> Option<LlmModel> {
124        (**self).model()
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    #[test]
133    fn lookup_context_window_known_model() {
134        assert_eq!(get_context_window("anthropic", "claude-opus-4-6"), Some(1_000_000));
135    }
136
137    #[test]
138    fn lookup_context_window_openrouter_model() {
139        // OpenRouter Qwen models should resolve from catalog
140        let result = get_context_window("openrouter", "anthropic/claude-opus-4");
141        assert_eq!(result, Some(200_000));
142    }
143
144    #[test]
145    fn lookup_context_window_unknown_model() {
146        assert_eq!(get_context_window("anthropic", "unknown-model-xyz"), None);
147    }
148
149    #[test]
150    fn lookup_context_window_unknown_provider() {
151        assert_eq!(get_context_window("unknown-provider", "some-model"), None);
152    }
153}
llm/provider.rs

llm/
provider.rs
