feat: implement embeddings

src/server/embeddings.rs

@@ -1,17 +1,39 @@
+use crate::CloudflareCredentials;
+use crate::openai;
+use crate::server::AppState;
+use crate::server::Provider;
+use crate::server::Upstream;
+use crate::server::take_random_by_swapping;
+use axum::Json;
+use axum::extract::State;
+use axum::http::StatusCode;
+use axum::response::IntoResponse;
+use axum::response::Response;
+use axum_extra::routing::TypedPath;
+use ecow::EcoString;
+use headers::Header;
+use headers::authorization::Bearer;
+use mitsein::vec1::Vec1;
+use reqwest::Certificate;
+use secrecy::ExposeSecret;
+use serde::Deserialize;
+use serde_with::DeserializeFromStr;
+use serde_with::SerializeDisplay;
+use std::collections::VecDeque;
+use strum::EnumString;
+use strum::VariantArray;
+
 #[derive(Deserialize, TypedPath)]
-#[typed_path("/v1/chat/completions")]
+#[typed_path("/v1/embeddings")]
-pub struct V1ChatCompletionsPath;
+pub struct V1EmbeddingsPath;
 
 #[derive(
     Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
 )]
 #[strum(serialize_all = "lowercase")]
 pub enum Shortcut {
-    XS,
+    #[strum(serialize = "lfm-2.5-embedding-350m")]
-    S,
+    Lfm2_5Embedding350m,
-    M,
-    L,
-    XL,
 }
 
 impl Shortcut {
@@ -20,17 +42,234 @@ impl Shortcut {
         use Shortcut as S;
 
         match self {
-            S::XS => vec![],
+            S::Lfm2_5Embedding350m => {
-
+                vec![(P::OpenRouter, "liquid/lfm-2.5-embedding-350m:free".into())]
-            S::S => vec![],
-
-            S::M => vec![],
-
-            S::L => vec![],
-
-            S::XL => vec![],
-
-            S::Experimental => vec![],
             }
         }
     }
+}
+
+pub struct EmbeddingsResponse(openai::embeddings::CreateEmbeddingResponse<String>);
+
+impl IntoResponse for EmbeddingsResponse {
+    fn into_response(self) -> Response {
+        let Self(embeddings) = self;
+
+        Json(embeddings).into_response()
+    }
+}
+
+pub enum EmbeddingsError {
+    NoneSucceeded,
+}
+
+impl IntoResponse for EmbeddingsError {
+    fn into_response(self) -> Response {
+        match self {
+            EmbeddingsError::NoneSucceeded => StatusCode::SERVICE_UNAVAILABLE.into_response(),
+        }
+    }
+}
+
+#[tracing::instrument(skip(state, request))]
+pub async fn create_embeddings(
+    _: V1EmbeddingsPath,
+    State(state): State<AppState>,
+    Json(request): Json<openai::embeddings::CreateEmbeddingRequest<Shortcut>>,
+) -> Result<EmbeddingsResponse, EmbeddingsError> {
+    let shortcut = request.model.clone();
+    let mut model_upstreams = VecDeque::new();
+
+    for (provider, model) in shortcut.suggestions() {
+        let upstreams: Vec<Upstream> =
+            match provider {
+                Provider::CloudflareWorkersAi => state
+                    .app_args
+                    .cloudflare_workers_ai_credentials
+                    .iter()
+                    .cloned()
+                    .map(
+                        |CloudflareCredentials {
+                             account_id,
+                             api_key,
+                         }| {
+                            Upstream::builder()
+                        .base_url(
+                            format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai")
+                                .parse()
+                                .expect("TODO"),
+                        )
+                        .api_key(api_key)
+                        .build()
+                        },
+                    )
+                    .collect(),
+                Provider::GoogleAiStudio => state
+                    .app_args
+                    .google_ai_studio_api_key
+                    .iter()
+                    .cloned()
+                    .map(|api_key| {
+                        Upstream::builder()
+                            .base_url(
+                                "https://generativelanguage.googleapis.com/v1beta/openai"
+                                    .parse()
+                                    .expect("TODO"),
+                            )
+                            .api_key(api_key)
+                            .build()
+                    })
+                    .collect(),
+                Provider::Groq => state
+                    .app_args
+                    .groq_api_key
+                    .iter()
+                    .cloned()
+                    .map(|api_key| {
+                        Upstream::builder()
+                            .base_url("https://api.groq.com/openai".parse().expect("TODO"))
+                            .api_key(api_key)
+                            .build()
+                    })
+                    .collect(),
+                Provider::HuggingFace => state
+                    .app_args
+                    .hugging_face_api_key
+                    .iter()
+                    .cloned()
+                    .map(|api_key| {
+                        Upstream::builder()
+                            .base_url("https://router.huggingface.co".parse().expect("TODO"))
+                            .api_key(api_key)
+                            .build()
+                    })
+                    .collect(),
+                Provider::Mistral => state
+                    .app_args
+                    .mistral_api_key
+                    .iter()
+                    .cloned()
+                    .map(|api_key| {
+                        Upstream::builder()
+                            .base_url("https://api.mistral.ai".parse().expect("TODO"))
+                            .api_key(api_key)
+                            .build()
+                    })
+                    .collect(),
+                Provider::OpenRouter => state
+                    .app_args
+                    .openrouter_api_key
+                    .iter()
+                    .cloned()
+                    .map(|api_key| {
+                        Upstream::builder()
+                            .base_url("https://openrouter.ai/api".parse().expect("TODO"))
+                            .api_key(api_key)
+                            .build()
+                    })
+                    .collect(),
+            };
+
+        if let Ok(upstreams) = Vec1::try_from(upstreams) {
+            model_upstreams.push_back((model, upstreams));
+        }
+    }
+
+    while let Some((model, upstreams)) = model_upstreams.pop_front() {
+        let (upstream, upstreams) = take_random_by_swapping(upstreams, &mut rand::rng());
+        if let Some(upstreams) = upstreams {
+            model_upstreams.push_back((model.clone(), upstreams));
+        }
+
+        tracing::info!(?model, ?upstream);
+
+        let openai::embeddings::CreateEmbeddingRequest {
+            model: _shortcut,
+            input,
+            dimensions,
+            encoding_format,
+            user,
+        } = request.clone();
+
+        let request = openai::embeddings::CreateEmbeddingRequest {
+            model,
+            input,
+            dimensions,
+            encoding_format,
+            user,
+        };
+
+        
+        let client = reqwest::Client::builder()
+            .tls_certs_only(
+                webpki_root_certs::TLS_SERVER_ROOT_CERTS
+                    .iter()
+                    .filter_map(|cert| Certificate::from_der(cert).ok()),
+            )
+            .build()
+            .expect("TODO");
+
+        let mut upstream_request = client
+            .post({
+                let mut url = upstream.base_url;
+
+                {
+                    let mut path_segments = url.path_segments_mut().expect("TODO");
+
+                    let s = V1EmbeddingsPath.to_string();
+                    let s = s.strip_prefix("/").unwrap_or(&s);
+                    let s = s.strip_suffix("/").unwrap_or(&s);
+
+                    path_segments.extend(s.split("/"));
+                }
+
+                tracing::info!(?url);
+
+                url
+            })
+            .json(&request);
+
+        if let Some(api_key) = upstream.api_key {
+            let authorization =
+                headers::Authorization::bearer(api_key.expose_secret()).expect("TODO");
+            upstream_request = upstream_request.header(headers::Authorization::<Bearer>::name(), {
+                let mut vec = vec![];
+
+                authorization.encode(&mut vec);
+
+                let [one] = vec.try_into().expect("TODO");
+
+                one
+            });
+        }
+
+        let upstream_response = match upstream_request.send().await {
+            Ok(upstream_response) => upstream_response,
+            Err(upstream_error) => {
+                tracing::error!(?upstream_error);
+                continue;
+            }
+        };
+
+        if let Err(upstream_error) = upstream_response.error_for_status_ref() {
+            tracing::error!(?upstream_error);
+            match upstream_response.text().await {
+                Ok(body) => tracing::error!(?body),
+                Err(body_error) => tracing::error!(?body_error),
+            };
+            continue;
+        };
+
+        
+            let response = match upstream_response.json().await {
+                Ok(response) => response,
+                Err(json_error) => {
+                    tracing::error!(?json_error);
+                    continue;
+                }
+            };
+
+            return Ok(EmbeddingsResponse(response));
+    }
+    Err(EmbeddingsError::NoneSucceeded)
+}

src/server/mod.rs

@@ -10,6 +10,7 @@ use typed_builder::TypedBuilder;
 use url::Url;
 
 mod chat;
+mod embeddings;
 mod models;
 
 #[derive(Debug, Clone, TypedBuilder)]
@@ -50,6 +51,7 @@ struct Upstream {
 pub fn create_router() -> Router<AppState> {
     Router::new()
         .typed_post(chat::completions::create_chat_completion)
+        .typed_post(embeddings::create_embeddings)
         .typed_get(models::list_models)
 }