ferrin_core/registry/
custom_provider.rs

1//! Providers assembled from explicit model instances.
2
3use std::collections::HashMap;
4use std::fmt;
5
6use ferrin_spec::EmbeddingModelRef;
7use ferrin_spec::ImageModelRef;
8use ferrin_spec::LanguageModelRef;
9use ferrin_spec::ModelKind;
10use ferrin_spec::NoSuchModelError;
11use ferrin_spec::Provider;
12use ferrin_spec::ProviderId;
13use ferrin_spec::ProviderRef;
14use ferrin_spec::RerankingModelRef;
15use ferrin_spec::SpeechModelRef;
16use ferrin_spec::SpeechTranslationModelRef;
17use ferrin_spec::TranscriptionModelRef;
18use ferrin_spec::VideoModelRef;
19
20/// Starts building a provider named `id`.
21#[must_use]
22pub fn custom_provider(id: impl Into<ProviderId>) -> CustomProviderBuilder {
23    CustomProviderBuilder {
24        provider: CustomProvider {
25            id: id.into(),
26            language_models: HashMap::new(),
27            embedding_models: HashMap::new(),
28            image_models: HashMap::new(),
29            transcription_models: HashMap::new(),
30            speech_models: HashMap::new(),
31            reranking_models: HashMap::new(),
32            video_models: HashMap::new(),
33            speech_translation_models: HashMap::new(),
34            fallback: None,
35        },
36    }
37}
38
39/// A provider that serves pre-built model instances, optionally falling back
40/// to another provider for unknown ids.
41#[derive(Clone)]
42pub struct CustomProvider {
43    id: ProviderId,
44    language_models: HashMap<String, LanguageModelRef>,
45    embedding_models: HashMap<String, EmbeddingModelRef>,
46    image_models: HashMap<String, ImageModelRef>,
47    transcription_models: HashMap<String, TranscriptionModelRef>,
48    speech_models: HashMap<String, SpeechModelRef>,
49    reranking_models: HashMap<String, RerankingModelRef>,
50    video_models: HashMap<String, VideoModelRef>,
51    speech_translation_models: HashMap<String, SpeechTranslationModelRef>,
52    fallback: Option<ProviderRef>,
53}
54
55impl fmt::Debug for CustomProvider {
56    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
57        f.debug_struct("CustomProvider")
58            .field("id", &self.id)
59            .field(
60                "language_models",
61                &self.language_models.keys().collect::<Vec<_>>(),
62            )
63            .field(
64                "embedding_models",
65                &self.embedding_models.keys().collect::<Vec<_>>(),
66            )
67            .field(
68                "image_models",
69                &self.image_models.keys().collect::<Vec<_>>(),
70            )
71            .field("fallback", &self.fallback.is_some())
72            .finish_non_exhaustive()
73    }
74}
75
76macro_rules! resolve {
77    ($self:ident, $map:ident, $fallback:ident, $model_id:expr, $kind:expr) => {{
78        if let Some(model) = $self.$map.get($model_id) {
79            return Ok(model.clone());
80        }
81        match &$self.fallback {
82            Some(fallback) => fallback.$fallback($model_id),
83            None => Err(NoSuchModelError::new($model_id, $kind).with_provider(&$self.id)),
84        }
85    }};
86}
87
88impl Provider for CustomProvider {
89    fn provider_id(&self) -> &ProviderId {
90        &self.id
91    }
92
93    fn language_model(&self, model_id: &str) -> Result<LanguageModelRef, NoSuchModelError> {
94        resolve!(
95            self,
96            language_models,
97            language_model,
98            model_id,
99            ModelKind::Language
100        )
101    }
102
103    fn embedding_model(&self, model_id: &str) -> Result<EmbeddingModelRef, NoSuchModelError> {
104        resolve!(
105            self,
106            embedding_models,
107            embedding_model,
108            model_id,
109            ModelKind::Embedding
110        )
111    }
112
113    fn image_model(&self, model_id: &str) -> Result<ImageModelRef, NoSuchModelError> {
114        resolve!(self, image_models, image_model, model_id, ModelKind::Image)
115    }
116
117    fn transcription_model(
118        &self,
119        model_id: &str,
120    ) -> Result<TranscriptionModelRef, NoSuchModelError> {
121        resolve!(
122            self,
123            transcription_models,
124            transcription_model,
125            model_id,
126            ModelKind::Transcription
127        )
128    }
129
130    fn speech_model(&self, model_id: &str) -> Result<SpeechModelRef, NoSuchModelError> {
131        resolve!(
132            self,
133            speech_models,
134            speech_model,
135            model_id,
136            ModelKind::Speech
137        )
138    }
139
140    fn reranking_model(&self, model_id: &str) -> Result<RerankingModelRef, NoSuchModelError> {
141        resolve!(
142            self,
143            reranking_models,
144            reranking_model,
145            model_id,
146            ModelKind::Reranking
147        )
148    }
149
150    fn video_model(&self, model_id: &str) -> Result<VideoModelRef, NoSuchModelError> {
151        resolve!(self, video_models, video_model, model_id, ModelKind::Video)
152    }
153
154    fn speech_translation_model(
155        &self,
156        model_id: &str,
157    ) -> Result<SpeechTranslationModelRef, NoSuchModelError> {
158        resolve!(
159            self,
160            speech_translation_models,
161            speech_translation_model,
162            model_id,
163            ModelKind::SpeechTranslation
164        )
165    }
166
167    fn realtime(&self) -> Option<ferrin_spec::RealtimeFactoryRef> {
168        self.fallback
169            .as_ref()
170            .and_then(|fallback| fallback.realtime())
171    }
172
173    fn files(&self) -> Option<ferrin_spec::FilesRef> {
174        self.fallback.as_ref().and_then(|fallback| fallback.files())
175    }
176
177    fn skills(&self) -> Option<ferrin_spec::SkillsRef> {
178        self.fallback
179            .as_ref()
180            .and_then(|fallback| fallback.skills())
181    }
182
183    fn batch(&self) -> Option<ferrin_spec::BatchRef> {
184        self.fallback.as_ref().and_then(|fallback| fallback.batch())
185    }
186}
187
188/// Builder of a [`CustomProvider`].
189#[derive(Debug)]
190pub struct CustomProviderBuilder {
191    provider: CustomProvider,
192}
193
194impl CustomProviderBuilder {
195    /// Registers a language model under `id`.
196    #[must_use]
197    pub fn language_model(
198        mut self,
199        id: impl Into<String>,
200        model: impl Into<LanguageModelRef>,
201    ) -> Self {
202        self.provider
203            .language_models
204            .insert(id.into(), model.into());
205        self
206    }
207
208    /// Registers an embedding model under `id`.
209    #[must_use]
210    pub fn embedding_model(
211        mut self,
212        id: impl Into<String>,
213        model: impl Into<EmbeddingModelRef>,
214    ) -> Self {
215        self.provider
216            .embedding_models
217            .insert(id.into(), model.into());
218        self
219    }
220
221    /// Registers an image model under `id`.
222    #[must_use]
223    pub fn image_model(mut self, id: impl Into<String>, model: impl Into<ImageModelRef>) -> Self {
224        self.provider.image_models.insert(id.into(), model.into());
225        self
226    }
227
228    /// Registers a transcription model under `id`.
229    #[must_use]
230    pub fn transcription_model(
231        mut self,
232        id: impl Into<String>,
233        model: impl Into<TranscriptionModelRef>,
234    ) -> Self {
235        self.provider
236            .transcription_models
237            .insert(id.into(), model.into());
238        self
239    }
240
241    /// Registers a speech model under `id`.
242    #[must_use]
243    pub fn speech_model(mut self, id: impl Into<String>, model: impl Into<SpeechModelRef>) -> Self {
244        self.provider.speech_models.insert(id.into(), model.into());
245        self
246    }
247
248    /// Registers a reranking model under `id`.
249    #[must_use]
250    pub fn reranking_model(
251        mut self,
252        id: impl Into<String>,
253        model: impl Into<RerankingModelRef>,
254    ) -> Self {
255        self.provider
256            .reranking_models
257            .insert(id.into(), model.into());
258        self
259    }
260
261    /// Registers a video model under `id`.
262    #[must_use]
263    pub fn video_model(mut self, id: impl Into<String>, model: impl Into<VideoModelRef>) -> Self {
264        self.provider.video_models.insert(id.into(), model.into());
265        self
266    }
267
268    /// Registers a speech translation model under `id`.
269    #[must_use]
270    pub fn speech_translation_model(
271        mut self,
272        id: impl Into<String>,
273        model: impl Into<SpeechTranslationModelRef>,
274    ) -> Self {
275        self.provider
276            .speech_translation_models
277            .insert(id.into(), model.into());
278        self
279    }
280
281    /// Delegates unknown ids (and services) to `provider`.
282    #[must_use]
283    pub fn fallback(mut self, provider: ProviderRef) -> Self {
284        self.provider.fallback = Some(provider);
285        self
286    }
287
288    /// Builds the provider.
289    #[must_use]
290    pub fn build(self) -> CustomProvider {
291        self.provider
292    }
293}
ferrin_core/registry/custom_provider.rs

ferrin_core/registry/
custom_provider.rs
