Skip to main content

llm/
provider.rs

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