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

ferrin_core/registry/
custom_provider.rs
