From 9dd13f374fdc2a98126ea4987c0f357d9ba9bd04 Mon Sep 17 00:00:00 2001 From: Jacob Date: Fri, 9 Oct 2026 01:52:29 -0400 Subject: [PATCH] feat: implement embeddings --- src/server/embeddings.rs | 275 ++++++++++++++++++++++++++++++++++++--- src/server/mod.rs | 2 + 2 files changed, 259 insertions(+), 18 deletions(-) diff --git a/src/server/embeddings.rs b/src/server/embeddings.rs index 5d5f885..d23b173 100644 --- a/src/server/embeddings.rs +++ b/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")] -pub struct V1ChatCompletionsPath; +#[typed_path("/v1/embeddings")] +pub struct V1EmbeddingsPath; #[derive( Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray, )] #[strum(serialize_all = "lowercase")] pub enum Shortcut { - XS, - S, - M, - L, - XL, + #[strum(serialize = "lfm-2.5-embedding-350m")] + Lfm2_5Embedding350m, } impl Shortcut { @@ -20,17 +42,234 @@ impl Shortcut { use Shortcut as S; match self { - S::XS => vec![], - - S::S => vec![], - - S::M => vec![], - - S::L => vec![], - - S::XL => vec![], - - S::Experimental => vec![], + S::Lfm2_5Embedding350m => { + vec![(P::OpenRouter, "liquid/lfm-2.5-embedding-350m:free".into())] + } } } } + +pub struct EmbeddingsResponse(openai::embeddings::CreateEmbeddingResponse); + +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, + Json(request): Json>, +) -> Result { + let shortcut = request.model.clone(); + let mut model_upstreams = VecDeque::new(); + + for (provider, model) in shortcut.suggestions() { + let upstreams: Vec = + 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::::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) +} diff --git a/src/server/mod.rs b/src/server/mod.rs index c00bead..0f84bc6 100644 --- a/src/server/mod.rs +++ b/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 { Router::new() .typed_post(chat::completions::create_chat_completion) + .typed_post(embeddings::create_embeddings) .typed_get(models::list_models) }