litellm_rs/core/embedding/
router.rs

1//! Embedding router - Routes embedding requests to appropriate providers
2//!
3//! This module handles parsing model strings to identify providers
4//! and routes requests to the appropriate embedding implementation.
5
6use crate::core::providers::{Provider, ProviderRegistry};
7use crate::core::types::responses::EmbeddingResponse;
8use crate::core::types::{
9    embedding::EmbeddingInput as TypesEmbeddingInput, embedding::EmbeddingRequest,
10};
11use crate::utils::error::gateway_error::{GatewayError, Result};
12use std::sync::Arc;
13use tokio::sync::OnceCell;
14use tracing::debug;
15
16use super::options::EmbeddingOptions;
17use super::types::EmbeddingInput;
18
19/// Default embedding router using the provider registry
20pub struct EmbeddingRouter {
21    provider_registry: Arc<ProviderRegistry>,
22}
23
24impl EmbeddingRouter {
25    /// Create a new embedding router
26    pub async fn new() -> Result<Self> {
27        let mut provider_registry = ProviderRegistry::new();
28
29        // Add OpenAI provider if API key is available
30        if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
31            use crate::core::providers::base::BaseConfig;
32            use crate::core::providers::openai::OpenAIProvider;
33            use crate::core::providers::openai::config::OpenAIConfig;
34
35            let config = OpenAIConfig {
36                base: BaseConfig {
37                    api_key: Some(api_key),
38                    api_base: Some("https://api.openai.com/v1".to_string()),
39                    organization: std::env::var("OPENAI_ORGANIZATION").ok(),
40                    ..Default::default()
41                },
42                organization: std::env::var("OPENAI_ORGANIZATION").ok(),
43                ..Default::default()
44            };
45
46            if let Ok(openai_provider) = OpenAIProvider::new(config).await {
47                provider_registry.register(Provider::OpenAI(openai_provider));
48            }
49        }
50
51        Ok(Self {
52            provider_registry: Arc::new(provider_registry),
53        })
54    }
55
56    /// Parse model string to extract provider and actual model name
57    ///
58    /// Supports formats like:
59    /// - "openai/text-embedding-ada-002" -> ("openai", "text-embedding-ada-002")
60    /// - "text-embedding-ada-002" -> ("openai", "text-embedding-ada-002") (default)
61    /// - "anthropic/voyage-3" -> ("anthropic", "voyage-3")
62    pub fn parse_model(model: &str) -> (&str, &str) {
63        if let Some(idx) = model.find('/') {
64            let (provider, rest) = model.split_at(idx);
65            (provider, &rest[1..]) // Skip the '/'
66        } else {
67            // Default to OpenAI for models without prefix
68            ("openai", model)
69        }
70    }
71
72    /// Route an embedding request to the appropriate provider
73    pub async fn embed(
74        &self,
75        model: &str,
76        input: EmbeddingInput,
77        options: EmbeddingOptions,
78    ) -> Result<EmbeddingResponse> {
79        let (provider_name, actual_model) = Self::parse_model(model);
80
81        debug!(
82            provider = %provider_name,
83            model = %actual_model,
84            "Routing embedding request"
85        );
86
87        // Check if we can use dynamic provider creation
88        if let Some(response) = self
89            .try_dynamic_provider_embed(provider_name, actual_model, &input, &options)
90            .await?
91        {
92            return Ok(response);
93        }
94
95        // Try to find a registered provider
96        let providers = self.provider_registry.all();
97
98        // Find matching provider from registry
99        for provider in providers.iter() {
100            if provider.name() == provider_name {
101                return self
102                    .execute_embedding(provider, actual_model, &input, &options)
103                    .await;
104            }
105        }
106
107        // No matching provider found
108        Err(GatewayError::not_found(format!(
109            "No embedding provider found for '{}'. Make sure the API key is set.",
110            provider_name
111        )))
112    }
113
114    /// Execute embedding using a specific provider
115    async fn execute_embedding(
116        &self,
117        provider: &Provider,
118        model: &str,
119        input: &EmbeddingInput,
120        options: &EmbeddingOptions,
121    ) -> Result<EmbeddingResponse> {
122        let request = self.build_request(model, input, options);
123
124        match provider {
125            Provider::OpenAI(p) => p
126                .embeddings(request)
127                .await
128                .map_err(|e| GatewayError::internal(format!("OpenAI embedding error: {}", e))),
129            // Add other providers as they support embeddings
130            _ => Err(GatewayError::not_implemented(format!(
131                "Provider '{}' does not support embeddings",
132                provider.name()
133            ))),
134        }
135    }
136
137    /// Build an EmbeddingRequest from input and options
138    fn build_request(
139        &self,
140        model: &str,
141        input: &EmbeddingInput,
142        options: &EmbeddingOptions,
143    ) -> EmbeddingRequest {
144        let types_input = match input {
145            EmbeddingInput::Text(text) => TypesEmbeddingInput::Text(text.clone()),
146            EmbeddingInput::TextArray(texts) => TypesEmbeddingInput::Array(texts.clone()),
147        };
148
149        EmbeddingRequest {
150            model: model.to_string(),
151            input: types_input,
152            user: options.user.clone(),
153            encoding_format: options.encoding_format.clone(),
154            dimensions: options.dimensions,
155            task_type: options.task_type.clone(),
156            truncation: None,
157        }
158    }
159
160    /// Try to create a dynamic provider for the embedding request
161    async fn try_dynamic_provider_embed(
162        &self,
163        provider_name: &str,
164        model: &str,
165        input: &EmbeddingInput,
166        options: &EmbeddingOptions,
167    ) -> Result<Option<EmbeddingResponse>> {
168        // Only proceed if user provided an API key
169        let api_key = match &options.api_key {
170            Some(key) => key.clone(),
171            None => return Ok(None),
172        };
173
174        let api_base = match provider_name {
175            "openai" => options
176                .api_base
177                .clone()
178                .unwrap_or_else(|| "https://api.openai.com/v1".to_string()),
179            "azure" | "azure_ai" | "azure-ai" => options
180                .api_base
181                .clone()
182                .or_else(|| std::env::var("AZURE_AI_API_BASE").ok())
183                .unwrap_or_else(|| "https://api.azure.com".to_string()),
184            _ => match &options.api_base {
185                Some(base) => base.clone(),
186                None => return Ok(None),
187            },
188        };
189
190        debug!(
191            provider = %provider_name,
192            model = %model,
193            "Creating dynamic embedding provider"
194        );
195
196        // Create dynamic OpenAI-compatible provider
197        let response = self
198            .create_dynamic_openai_embedding(&api_key, &api_base, model, input, options)
199            .await?;
200
201        Ok(Some(response))
202    }
203
204    /// Create a dynamic OpenAI-compatible provider for embeddings
205    async fn create_dynamic_openai_embedding(
206        &self,
207        api_key: &str,
208        api_base: &str,
209        model: &str,
210        input: &EmbeddingInput,
211        options: &EmbeddingOptions,
212    ) -> Result<EmbeddingResponse> {
213        use crate::core::providers::base::BaseConfig;
214        use crate::core::providers::openai::OpenAIProvider;
215        use crate::core::providers::openai::config::OpenAIConfig;
216
217        let timeout = options.timeout.unwrap_or(60);
218
219        let config = OpenAIConfig {
220            base: BaseConfig {
221                api_key: Some(api_key.to_string()),
222                api_base: Some(api_base.to_string()),
223                timeout,
224                headers: options.headers.clone().unwrap_or_default(),
225                ..Default::default()
226            },
227            ..Default::default()
228        };
229
230        let provider = OpenAIProvider::new(config).await.map_err(|e| {
231            GatewayError::internal(format!(
232                "Failed to create dynamic embedding provider: {}",
233                e
234            ))
235        })?;
236
237        let request = self.build_request(model, input, options);
238
239        provider
240            .embeddings(request)
241            .await
242            .map_err(|e| GatewayError::internal(format!("Dynamic embedding error: {}", e)))
243    }
244}
245
246/// Global router instance
247static GLOBAL_EMBEDDING_ROUTER: OnceCell<EmbeddingRouter> = OnceCell::const_new();
248
249/// Get or initialize the global embedding router
250pub async fn get_global_embedding_router() -> Result<&'static EmbeddingRouter> {
251    GLOBAL_EMBEDDING_ROUTER
252        .get_or_try_init(|| async {
253            EmbeddingRouter::new().await.map_err(|e| {
254                GatewayError::internal(format!("Failed to initialize embedding router: {}", e))
255            })
256        })
257        .await
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263
264    #[test]
265    fn test_parse_model_with_provider() {
266        let (provider, model) = EmbeddingRouter::parse_model("openai/text-embedding-ada-002");
267        assert_eq!(provider, "openai");
268        assert_eq!(model, "text-embedding-ada-002");
269    }
270
271    #[test]
272    fn test_parse_model_without_provider() {
273        let (provider, model) = EmbeddingRouter::parse_model("text-embedding-ada-002");
274        assert_eq!(provider, "openai");
275        assert_eq!(model, "text-embedding-ada-002");
276    }
277
278    #[test]
279    fn test_parse_model_anthropic() {
280        let (provider, model) = EmbeddingRouter::parse_model("anthropic/voyage-3");
281        assert_eq!(provider, "anthropic");
282        assert_eq!(model, "voyage-3");
283    }
284
285    #[test]
286    fn test_parse_model_azure() {
287        let (provider, model) = EmbeddingRouter::parse_model("azure/text-embedding-3-small");
288        assert_eq!(provider, "azure");
289        assert_eq!(model, "text-embedding-3-small");
290    }
291}
litellm_rs/core/embedding/router.rs

litellm_rs/core/embedding/
router.rs
