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
13pub type LlmResponseStream = Pin<Box<dyn Stream<Item = LlmResult<LlmResponse>> + Send>>;
20
21#[doc = include_str!("docs/provider_factory.md")]
22pub trait ProviderFactory: Sized {
23 fn from_env() -> impl Future<Output = LlmResult<Self>> + Send;
25
26 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 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 fn context_window(&self) -> Option<u32>;
46
47 fn model(&self) -> Option<LlmModel> {
51 None
52 }
53}
54
55pub 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
85pub(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
111pub(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 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}