feat: novita, nvidia nim, vercel ai gateway
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -224,8 +224,13 @@ pub async fn create_decision(
|
||||
.into_iter()
|
||||
.map(|(choice_id, probability)| {
|
||||
openai::decisions::ScoreProbability {
|
||||
value: choice_id.parse().expect("it should be a number, I think"),
|
||||
label: legend.get(&choice_id).expect("it should exist, I think").clone(),
|
||||
value: choice_id
|
||||
.parse()
|
||||
.expect("it should be a number, I think"),
|
||||
label: legend
|
||||
.get(&choice_id)
|
||||
.expect("it should exist, I think")
|
||||
.clone(),
|
||||
probability,
|
||||
}
|
||||
})
|
||||
|
||||
@@ -1,279 +1,286 @@
|
||||
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/embeddings")]
|
||||
pub struct V1EmbeddingsPath;
|
||||
|
||||
#[derive(
|
||||
Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
|
||||
)]
|
||||
#[strum(serialize_all = "lowercase")]
|
||||
pub enum Shortcut {
|
||||
#[strum(serialize = "codestral-embed-2505")]
|
||||
CodestralEmbed2505,
|
||||
|
||||
#[strum(serialize = "lfm2.5-embedding-350m")]
|
||||
Lfm2_5Embedding350M,
|
||||
|
||||
#[strum(serialize = "mistral-embed-2312")]
|
||||
MistralEmbed2312,
|
||||
|
||||
#[strum(serialize = "nemotron-3-embed-1b")]
|
||||
Nemotron3Embed1B,
|
||||
}
|
||||
|
||||
impl Shortcut {
|
||||
fn suggestions(&self) -> Vec<(Provider, EcoString)> {
|
||||
use Provider as P;
|
||||
use Shortcut as S;
|
||||
|
||||
match self {
|
||||
S::CodestralEmbed2505 => vec![
|
||||
(P::Mistral, "codestral-embed-2505".into())
|
||||
],
|
||||
S::Lfm2_5Embedding350M => {
|
||||
vec![(P::OpenRouter, "liquid/lfm-2.5-embedding-350m:free".into())]
|
||||
},
|
||||
S::MistralEmbed2312 => vec![
|
||||
(P::Mistral, "mistral-embed-2312".into())
|
||||
],
|
||||
S::Nemotron3Embed1B => {
|
||||
vec![(P::OpenRouter, "nvidia/nemotron-3-embed-1b:free".into())]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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()
|
||||
.url(
|
||||
format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1/embeddings")
|
||||
.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()
|
||||
.url(
|
||||
"https://generativelanguage.googleapis.com/v1beta/openai/v1/embeddings"
|
||||
.parse()
|
||||
.expect("TODO"),
|
||||
)
|
||||
.api_key(api_key)
|
||||
.build()
|
||||
})
|
||||
.collect(),
|
||||
Provider::Groq => state
|
||||
.app_args
|
||||
.groq_api_key
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|api_key| {
|
||||
Upstream::builder()
|
||||
.url("https://api.groq.com/openai/v1/embeddings".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()
|
||||
.url("https://router.huggingface.co/v1/embeddings".parse().expect("TODO"))
|
||||
.api_key(api_key)
|
||||
.build()
|
||||
})
|
||||
.collect(),
|
||||
Provider::LiquidAi => unreachable!(),
|
||||
Provider::Mistral => state
|
||||
.app_args
|
||||
.mistral_api_key
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|api_key| {
|
||||
Upstream::builder()
|
||||
.url("https://api.mistral.ai/v1/embeddings".parse().expect("TODO"))
|
||||
.api_key(api_key)
|
||||
.build()
|
||||
})
|
||||
.collect(),
|
||||
Provider::OpenRouter => state
|
||||
.app_args
|
||||
.openrouter_api_key
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|api_key| {
|
||||
Upstream::builder()
|
||||
.url("https://openrouter.ai/api/v1/embeddings".parse().expect("TODO"))
|
||||
.api_key(api_key)
|
||||
.build()
|
||||
})
|
||||
.collect(),
|
||||
Provider::ZAi => unreachable!(),
|
||||
};
|
||||
|
||||
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(upstream.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)
|
||||
}
|
||||
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/embeddings")]
|
||||
pub struct V1EmbeddingsPath;
|
||||
|
||||
#[derive(
|
||||
Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
|
||||
)]
|
||||
#[strum(serialize_all = "lowercase")]
|
||||
pub enum Shortcut {
|
||||
#[strum(serialize = "codestral-embed-2505")]
|
||||
CodestralEmbed2505,
|
||||
|
||||
#[strum(serialize = "lfm2.5-embedding-350m")]
|
||||
Lfm2_5Embedding350M,
|
||||
|
||||
#[strum(serialize = "mistral-embed-2312")]
|
||||
MistralEmbed2312,
|
||||
|
||||
#[strum(serialize = "nemotron-3-embed-1b")]
|
||||
Nemotron3Embed1B,
|
||||
}
|
||||
|
||||
impl Shortcut {
|
||||
fn suggestions(&self) -> Vec<(Provider, EcoString)> {
|
||||
use Provider as P;
|
||||
use Shortcut as S;
|
||||
|
||||
match self {
|
||||
S::CodestralEmbed2505 => vec![(P::Mistral, "codestral-embed-2505".into())],
|
||||
S::Lfm2_5Embedding350M => {
|
||||
vec![(P::OpenRouter, "liquid/lfm-2.5-embedding-350m:free".into())]
|
||||
}
|
||||
S::MistralEmbed2312 => vec![(P::Mistral, "mistral-embed-2312".into())],
|
||||
S::Nemotron3Embed1B => vec![
|
||||
(P::NvidiaNim, "nvidia/nemotron-3-embed-1b".into()),
|
||||
(P::OpenRouter, "nvidia/nemotron-3-embed-1b:free".into()),
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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()
|
||||
.url(
|
||||
format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1/embeddings")
|
||||
.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()
|
||||
.url(
|
||||
"https://generativelanguage.googleapis.com/v1beta/openai/v1/embeddings"
|
||||
.parse()
|
||||
.expect("TODO"),
|
||||
)
|
||||
.api_key(api_key)
|
||||
.build()
|
||||
})
|
||||
.collect(),
|
||||
Provider::Groq => state
|
||||
.app_args
|
||||
.groq_api_key
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|api_key| {
|
||||
Upstream::builder()
|
||||
.url("https://api.groq.com/openai/v1/embeddings".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()
|
||||
.url("https://router.huggingface.co/v1/embeddings".parse().expect("TODO"))
|
||||
.api_key(api_key)
|
||||
.build()
|
||||
})
|
||||
.collect(),
|
||||
Provider::LiquidAi => unreachable!(),
|
||||
Provider::Mistral => state
|
||||
.app_args
|
||||
.mistral_api_key
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|api_key| {
|
||||
Upstream::builder()
|
||||
.url("https://api.mistral.ai/v1/embeddings".parse().expect("TODO"))
|
||||
.api_key(api_key)
|
||||
.build()
|
||||
})
|
||||
.collect(),
|
||||
Provider::NvidiaNim => state
|
||||
.app_args
|
||||
.nvidia_nim_api_key
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|api_key| {
|
||||
Upstream::builder()
|
||||
.url("https://integrate.api.nvidia.com/v1/embeddings".parse().expect("TODO"))
|
||||
.api_key(api_key)
|
||||
.build()
|
||||
})
|
||||
.collect(),
|
||||
Provider::NovitaAi => unreachable!(),
|
||||
Provider::OpenRouter => state
|
||||
.app_args
|
||||
.openrouter_api_key
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|api_key| {
|
||||
Upstream::builder()
|
||||
.url("https://openrouter.ai/api/v1/embeddings".parse().expect("TODO"))
|
||||
.api_key(api_key)
|
||||
.build()
|
||||
})
|
||||
.collect(),
|
||||
Provider::VercelAiGateway => unreachable!(),
|
||||
Provider::ZAi => unreachable!(),
|
||||
};
|
||||
|
||||
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(upstream.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)
|
||||
}
|
||||
|
||||
@@ -28,7 +28,10 @@ enum Provider {
|
||||
HuggingFace,
|
||||
LiquidAi,
|
||||
Mistral,
|
||||
NovitaAi,
|
||||
NvidiaNim,
|
||||
OpenRouter,
|
||||
VercelAiGateway,
|
||||
ZAi,
|
||||
}
|
||||
|
||||
|
||||
@@ -49,6 +49,8 @@ impl Shortcut {
|
||||
(P::CloudflareWorkersAi, "clef-flash".into()),
|
||||
(P::OpenRouter, "respan/span-01-lite:free".into()),
|
||||
(P::OpenRouter, "respan/span-01-lite".into()),
|
||||
(P::VercelAiGateway, "convaiinnovations/laya-free".into()),
|
||||
(P::VercelAiGateway, "convaiinnovations/laya".into()),
|
||||
]
|
||||
}
|
||||
S::M => {
|
||||
@@ -132,6 +134,8 @@ pub async fn create_system_one(
|
||||
})
|
||||
.collect(),
|
||||
Provider::Mistral => unreachable!(),
|
||||
Provider::NovitaAi => unreachable!(),
|
||||
Provider::NvidiaNim => unreachable!(),
|
||||
Provider::OpenRouter => state
|
||||
.app_args
|
||||
.openrouter_api_key
|
||||
@@ -144,6 +148,18 @@ pub async fn create_system_one(
|
||||
.build()
|
||||
})
|
||||
.collect(),
|
||||
Provider::VercelAiGateway => state
|
||||
.app_args
|
||||
.vercel_ai_gateway_api_key
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|api_key| {
|
||||
Upstream::builder()
|
||||
.url("https://ai-gateway.vercel.sh/typesafe/v1/systemone".parse().expect("TODO"))
|
||||
.api_key(api_key)
|
||||
.build()
|
||||
})
|
||||
.collect(),
|
||||
Provider::ZAi => unreachable!(),
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user