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

feat: ollama cloud support

Browse Source
This commit is contained in:
J / Jacob Babich
2026-10-16 12:57:27 -04:00
parent e74b5f6a4b
commit 6e3a283d81
5 changed files with 102 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/lib.rs
View File

@@ -63,6 +63,9 @@ pub struct AppArgs {
#[arg(long)]
pub nvidia_nim_api_key: Vec<SecretString>,
#[arg(long)]
pub ollama_cloud_api_key: Vec<SecretString>,
#[arg(long)]
pub openrouter_api_key: Vec<SecretString>,

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

@@ -1,7 +1,6 @@
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;
@@ -86,6 +85,22 @@ impl Display for ShortcutAndMaybeReasoning {
}
}
#[derive(Debug, Clone)]
enum Provider {
CloudflareWorkersAi,
GoogleAiStudio,
Groq,
HuggingFace,
Kilo,
Mistral,
NovitaAi,
NvidiaNim,
Ollama,
OpenRouter,
VercelAiGateway,
ZAi,
}
impl Shortcut {
fn suggestions(&self) -> Vec<(Provider, EcoString)> {
use Provider as P;
@@ -125,7 +140,9 @@ impl Shortcut {
P::OpenRouter,
"nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:free".into(),
),
(P::Ollama, "nemotron-3-nano:30b".into()),
(P::Groq, "openai/gpt-oss-20b".into()),
(P::Ollama, "gpt-oss:20b".into()),
(P::NvidiaNim, "openai/gpt-oss-20b".into()),
(P::CloudflareWorkersAi, "@cf/openai/gpt-oss-20b".into()),
(P::OpenRouter, "openai/gpt-oss-20b:free".into()),
@@ -145,6 +162,7 @@ impl Shortcut {
(P::NvidiaNim, "meta/muse-glimmer-30b".into()),
(P::HuggingFace, "meta-models/Muse-Glimmer-30B".into()),
(P::GoogleAiStudio, "gemma-4-31b-it".into()),
(P::Ollama, "gemma4:31b".into()),
(P::NvidiaNim, "google/gemma-4-31b-it".into()),
(P::OpenRouter, "google/gemma-4-31b-it:free".into()),
(P::HuggingFace, "google/gemma-4-31B-it".into()),
@@ -185,6 +203,7 @@ impl Shortcut {
"inclusionai/ling-3.0-flash-sante".into(),
),
(P::VercelAiGateway, "mistral/mistral-large-3".into()),
(P::Ollama, "nemotron-3-super".into()),
(P::NvidiaNim, "nvidia/nemotron-3-super-120b-a12b".into()),
(P::Kilo, "nvidia-nemotron-3-super-120b-a12b-free".into()),
(
@@ -192,6 +211,7 @@ impl Shortcut {
"nvidia/nemotron-3-super-120b-a12b:free".into(),
),
(P::Groq, "openai/gpt-oss-120b".into()),
(P::Ollama, "gpt-oss:120b".into()),
(P::CloudflareWorkersAi, "@cf/openai/gpt-oss-120b".into()),
(P::OpenRouter, "openai/gpt-oss-120b:free".into()),
(P::HuggingFace, "openai/gpt-oss-120b".into()),
@@ -208,10 +228,11 @@ impl Shortcut {
(P::VercelAiGateway, "openai/gpt-5.2".into()),
(P::NvidiaNim, "deepseek-ai/deepseek-v4.1-flash".into()),
(P::HuggingFace, "deepseek-ai/DeepSeek-V4.1-Flash".into()),
(P::VercelAiGateway, "xiaomi/mimo-v2.6-flash".into()),
(P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Flash-RL".into()),
(P::VercelAiGateway, "inclusionai/ling-3.1-flash".into()),
(P::HuggingFace, "inclusionAI/Ling-3.1-flash".into()),
(P::Ollama, "beam".into()),
(P::VercelAiGateway, "xiaomi/mimo-v2.6-flash".into()),
(P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Flash-RL".into()),
(P::Kilo, "nex-agi-nex-n2-5-pro-free".into()),
(P::VercelAiGateway, "xiaomi/mimo-v2.5".into()),
(P::Kilo, "dots-studio-dots-3-note-preview-free".into()),
@@ -223,6 +244,7 @@ impl Shortcut {
(P::Kilo, "tencent-hy3-preview-free".into()),
(P::OpenRouter, "thinkingmachines/inkling-small:free".into()),
(P::HuggingFace, "thinkingmachines/Inkling-Small".into()),
(P::Ollama, "nemotron-3-ultra".into()),
(P::NvidiaNim, "nvidia/nemotron-3-ultra-550b-a55b".into()),
(P::Kilo, "nvidia-nemotron-3-ultra-550b-a55b-free".into()),
],
@@ -232,11 +254,11 @@ impl Shortcut {
(P::HuggingFace, "zai-org/GLM-5.3".into()),
(P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3".into()),
(P::Mistral, "zai-glm-5-3".into()),
(P::NvidiaNim, "moonshotai/kimi-k3".into()),
(P::HuggingFace, "moonshotai/Kimi-K3".into()),
(P::HuggingFace, "tencent/Hy4-preview".into()),
(P::HuggingFace, "XiaomiMiMo/MiMo-V2.6-Pro-RL".into()),
(P::Mistral, "mistral-large-4".into()),
(P::NvidiaNim, "moonshotai/kimi-k3".into()),
(P::HuggingFace, "moonshotai/Kimi-K3".into()),
(P::HuggingFace, "Qwen/Qwen3.8-2.4T-A95B".into()),
(P::VercelAiGateway, "xiaomi/mimo-v2.5-pro".into()),
(P::OpenRouter, "thinkingmachines/inkling:free".into()),
@@ -397,7 +419,6 @@ pub async fn create_chat_completion(
.build()
})
.collect(),
Provider::LiquidAi => unreachable!(),
Provider::Mistral => state
.app_args
.mistral_api_key
@@ -434,6 +455,18 @@ pub async fn create_chat_completion(
.build()
})
.collect(),
Provider::Ollama => state
.app_args
.ollama_cloud_api_key
.iter()
.cloned()
.map(|api_key| {
Upstream::builder()
.url("https://ollama.com/v1/chat/completions".parse().expect("TODO"))
.api_key(api_key)
.build()
})
.collect(),
Provider::OpenRouter => state
.app_args
.openrouter_api_key

151
src/server/embeddings.rs
View File

@@ -1,7 +1,6 @@
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;
@@ -45,6 +44,13 @@ pub enum Shortcut {
Nemotron3Embed1B,
}
#[derive(Debug, Clone)]
enum Provider {
Mistral,
NvidiaNim,
OpenRouter,
}
impl Shortcut {
fn suggestions(&self) -> Vec<(Provider, EcoString)> {
use Provider as P;
@@ -96,111 +102,56 @@ pub async fn create_embeddings(
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()
let upstreams: Vec<Upstream> = match provider {
Provider::Mistral => state
.app_args
.mistral_api_key
.iter()
.cloned()
.map(|api_key| {
Upstream::builder()
.url(
format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1/embeddings")
"https://api.mistral.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::Kilo => unreachable!(),
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!(),
};
})
.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::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(),
};
if let Ok(upstreams) = Vec1::try_from(upstreams) {
model_upstreams.push_back((model, upstreams));

17
src/server/mod.rs
View File

@@ -19,23 +19,6 @@ mod system_one;
pub struct AppState {
app_args: Arc<AppArgs>,
}
#[derive(Debug, Clone)]
enum Provider {
CloudflareWorkersAi,
GoogleAiStudio,
Groq,
HuggingFace,
Kilo,
LiquidAi,
Mistral,
NovitaAi,
NvidiaNim,
OpenRouter,
VercelAiGateway,
ZAi,
}
fn take_random_by_swapping<T>(mut vec1: Vec1<T>, mut rng: impl Rng) -> (T, Option<Vec1<T>>) {
let end = vec1.len().get() - 1;
let index = rng.random_range(0..=end);

18
src/server/system_one.rs
View File

@@ -1,7 +1,6 @@
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;
@@ -36,6 +35,14 @@ pub enum Shortcut {
M,
}
#[derive(Debug, Clone)]
enum Provider {
CloudflareWorkersAi,
LiquidAi,
OpenRouter,
VercelAiGateway,
}
impl Shortcut {
fn suggestions(&self) -> Vec<(Provider, EcoString)> {
use Provider as P;
@@ -53,6 +60,7 @@ impl Shortcut {
(P::VercelAiGateway, "convaiinnovations/laya".into()),
]
}
S::M => {
vec![
(P::CloudflareWorkersAi, "clef".into()),
@@ -118,10 +126,6 @@ pub async fn create_system_one(
},
)
.collect(),
Provider::GoogleAiStudio => unreachable!(),
Provider::Groq => unreachable!(),
Provider::HuggingFace => unreachable!(),
Provider::Kilo => unreachable!(),
Provider::LiquidAi => state
.app_args
.liquid_ai_api_key
@@ -134,9 +138,6 @@ pub async fn create_system_one(
.build()
})
.collect(),
Provider::Mistral => unreachable!(),
Provider::NovitaAi => unreachable!(),
Provider::NvidiaNim => unreachable!(),
Provider::OpenRouter => state
.app_args
.openrouter_api_key
@@ -161,7 +162,6 @@ pub async fn create_system_one(
.build()
})
.collect(),
Provider::ZAi => unreachable!(),
};
if let Ok(upstreams) = Vec1::try_from(upstreams) {