jacob/katniss-ai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: novita, nvidia nim, vercel ai gateway

Browse Source
This commit is contained in:
Jacob
2026-10-13 23:37:24 -04:00
parent a0d95d52e4
commit f4f657c219
6 changed files with 918 additions and 794 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/lib.rs
View File

@@ -54,9 +54,18 @@ pub struct AppArgs {
#[arg(long)] #[arg(long)]
pub mistral_api_key: Vec<SecretString>, pub mistral_api_key: Vec<SecretString>,
#[arg(long)]
pub novita_ai_api_key: Vec<SecretString>,
#[arg(long)]
pub nvidia_nim_api_key: Vec<SecretString>,
#[arg(long)] #[arg(long)]
pub openrouter_api_key: Vec<SecretString>, pub openrouter_api_key: Vec<SecretString>,
#[arg(long)]
pub vercel_ai_gateway_api_key: Vec<SecretString>,
#[arg(long)] #[arg(long)]
pub z_ai_api_key: Vec<SecretString>, pub z_ai_api_key: Vec<SecretString>,
} }

1110
src/server/chat/completions.rs
View File

File diff suppressed because it is too large Load Diff

9
src/server/decisions.rs
View File

@@ -224,8 +224,13 @@ pub async fn create_decision(
.into_iter() .into_iter()
.map(|(choice_id, probability)| { .map(|(choice_id, probability)| {
openai::decisions::ScoreProbability { openai::decisions::ScoreProbability {
value: choice_id.parse().expect("it should be a number, I think"), value: choice_id
label: legend.get(&choice_id).expect("it should exist, I think").clone(), .parse()
.expect("it should be a number, I think"),
label: legend
.get(&choice_id)
.expect("it should exist, I think")
.clone(),
probability, probability,
} }
}) })

565
src/server/embeddings.rs
View File

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

3
src/server/mod.rs
View File

@@ -28,7 +28,10 @@ enum Provider {
HuggingFace, HuggingFace,
LiquidAi, LiquidAi,
Mistral, Mistral,
NovitaAi,
NvidiaNim,
OpenRouter, OpenRouter,
VercelAiGateway,
ZAi, ZAi,
} }

16
src/server/system_one.rs
View File

@@ -49,6 +49,8 @@ impl Shortcut {
(P::CloudflareWorkersAi, "clef-flash".into()), (P::CloudflareWorkersAi, "clef-flash".into()),
(P::OpenRouter, "respan/span-01-lite:free".into()), (P::OpenRouter, "respan/span-01-lite:free".into()),
(P::OpenRouter, "respan/span-01-lite".into()), (P::OpenRouter, "respan/span-01-lite".into()),
(P::VercelAiGateway, "convaiinnovations/laya-free".into()),
(P::VercelAiGateway, "convaiinnovations/laya".into()),
] ]
} }
S::M => { S::M => {
@@ -132,6 +134,8 @@ pub async fn create_system_one(
}) })
.collect(), .collect(),
Provider::Mistral => unreachable!(), Provider::Mistral => unreachable!(),
Provider::NovitaAi => unreachable!(),
Provider::NvidiaNim => unreachable!(),
Provider::OpenRouter => state Provider::OpenRouter => state
.app_args .app_args
.openrouter_api_key .openrouter_api_key
@@ -144,6 +148,18 @@ pub async fn create_system_one(
.build() .build()
}) })
.collect(), .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!(), Provider::ZAi => unreachable!(),
}; };