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
9pub type LlmResponseStream = Pin<Box<dyn Stream<Item = LlmResult<LlmResponse>> + Send>>;
16
17#[doc = include_str!("docs/provider_factory.md")]
18pub trait ProviderFactory: Sized {
19 fn from_env() -> impl Future<Output = LlmResult<Self>> + Send;
21
22 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 fn context_window(&self) -> Option<u32>;
34
35 fn model(&self) -> Option<LlmModel> {
39 None
40 }
41}
42
43pub 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 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}