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

feat: more

Browse Source
This commit is contained in:
Jacob
2026-09-28 00:03:19 -04:00
parent 069bd051f2
commit c233fcda38
5 changed files with 540 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
src/main.rs
View File

@@ -1,7 +1,10 @@
use std::net::{IpAddr, SocketAddr};
use clap::Parser;
use futures::StreamExt as _;
use reqwest_sse::EventSource;
use secrecy::SecretString;
use tokio::{join, pin, task::LocalSet};
use tracing_subscriber::{EnvFilter, fmt::format::FmtSpan};
mod openai;
@@ -60,13 +63,49 @@ async fn main() {
let addr = SocketAddr::new(app_args.ip, app_args.port);
let (tx, rx) = async_channel::unbounded();
let app = server::create_app(
server::AppState::builder()
.app_args(app_args.into())
.stream_requester(tx)
.build(),
);
let listener = tokio::net::TcpListener::bind(addr).await.expect("TODO");
axum::serve(listener, app).await.expect("TODO");
let stream_responses = async move {
let local_set = LocalSet::new();
local_set
.run_until(async {
pin!(rx);
while let Some((response, sse_sink)) = rx.next().await {
local_set.spawn_local(async move {
match response.events().await {
Ok(mut events) => {
while let Some(event_result) = events.next().await {
let message = event_result
.map(|event| {
axum::response::sse::Event::default().data(event.data)
})
.map_err(Into::into);
let _ = sse_sink.send(message).await;
}
}
Err(event_source_error) => {
let _ = sse_sink.send(Err(event_source_error.into())).await;
}
};
});
}
})
.await;
};
join!(stream_responses, async move {
axum::serve(listener, app).await.expect("TODO");
});
}

2
src/openai/chat/mod.rs
View File

@@ -1 +1 @@
pub mod completion;
pub mod completion;

348
src/server.rs
View File

@@ -1,25 +1,54 @@
use std::sync::Arc;
use crate::AppArgs;
use crate::openai;
use axum::Json;
use axum::Router;
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::response::Response;
use axum::response::Sse;
use axum_extra::routing::RouterExt;
use axum_extra::routing::TypedPath;
use ecow::EcoString;
use ecow::string::ToEcoString;
use futures::Stream;
use futures::StreamExt as _;
use futures::TryStreamExt as _;
use headers::Header;
use headers::authorization::Bearer;
use indexmap::IndexMap;
use mitsein::vec1::Vec1;
use rand::Rng;
use rand::RngExt;
use reqwest_sse::EventSource;
use secrecy::ExposeSecret;
use secrecy::SecretString;
use serde::Deserialize;
use serde::Serialize;
use std::collections::VecDeque;
use std::hash::Hash;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::oneshot;
use typed_builder::TypedBuilder;
use url::Url;
type SSEEventResult = Result<axum::response::sse::Event, Box<dyn std::error::Error + Sync + Send>>;
type StreamRequest = (reqwest::Response, async_channel::Sender<SSEEventResult>);
type StreamRequester = async_channel::Sender<StreamRequest>;
#[derive(Debug, Clone, TypedBuilder)]
pub struct AppState {
app_args: Arc<AppArgs>,
stream_requester: StreamRequester,
}
#[derive(Deserialize, TypedPath)]
#[typed_path("/chat/completions")]
struct ChatCompletionsPath;
#[derive(Debug, Clone)]
enum Provider {
CloudflareWorkersAi,
GoogleAiStudio,
@@ -29,7 +58,7 @@ enum Provider {
OpenRouter,
}
#[derive(Debug, Serialize, Deserialize)]
#[derive(Debug, Serialize, Deserialize, Clone)]
enum Shortcut {
XS,
S,
@@ -39,7 +68,7 @@ enum Shortcut {
}
impl Shortcut {
fn suggestions(&self) -> Vec<(Provider, String)> {
fn suggestions(&self) -> Vec<(Provider, EcoString)> {
use Provider as P;
use Shortcut as S;
@@ -52,7 +81,10 @@ impl Shortcut {
(P::OpenRouter, "nvidia/nemotron-3.5-lightning:free".into()),
(P::OpenRouter, "google/gemma-4-26b-a4b-it:free".into()),
(P::OpenRouter, "poolside/laguna-xs.2:free".into()),
(P::OpenRouter, "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:free".into()),
(
P::OpenRouter,
"nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:free".into(),
),
(P::CloudflareWorkersAi, "@cf/openai/gpt-oss-20b".into()),
(P::OpenRouter, "openai/gpt-oss-20b:free".into()),
(P::Groq, "openai/gpt-oss-20b".into()),
@@ -77,7 +109,10 @@ impl Shortcut {
(P::Mistral, "mistral-medium-3.5".into()),
(P::OpenRouter, "inclusionai/ling-3.0-flash-vl:free".into()),
(P::OpenRouter, "poolside/laguna-s-2.1:free".into()),
(P::OpenRouter, "nvidia/nemotron-3-super-120b-a12b:free".into()),
(
P::OpenRouter,
"nvidia/nemotron-3-super-120b-a12b:free".into(),
),
(P::CloudflareWorkersAi, "@cf/openai/gpt-oss-120b".into()),
(P::OpenRouter, "openai/gpt-oss-120b:free".into()),
(P::Groq, "openai/gpt-oss-120b".into()),
@@ -87,8 +122,9 @@ impl Shortcut {
(P::HuggingFace, "zai-org/GLM-5.3-Flash".into()),
(P::HuggingFace, "deepseek-ai/DeepSeek-V4.1-Flash".into()),
(P::OpenRouter, "nex-agi/nex-n2.5-pro:free".into()),
(P::HuggingFace, "MiniMaxAI/MiniMax-M3".into()),
(P::OpenRouter, "thinkingmachines/inkling-small:free".into()),
(P::HuggingFace, "thinkingmachines/Inkling-Small".into())
(P::HuggingFace, "thinkingmachines/Inkling-Small".into()),
],
S::XL => vec![
(P::CloudflareWorkersAi, "@cf/zai-org/glm-5.3".into()),
@@ -96,27 +132,311 @@ impl Shortcut {
(P::Mistral, "zai-glm-5-3".into()),
(P::HuggingFace, "moonshotai/Kimi-K3".into()),
(P::HuggingFace, "Qwen/Qwen3.8-2.4T-A95B".into()),
(P::HuggingFace, "MiniMaxAI/MiniMax-M3".into()),
(P::OpenRouter, "thinkingmachines/inkling:free".into()),
(P::HuggingFace, "thinkingmachines/Inkling".into()),
(P::HuggingFace, "tencent/Hy4-preview".into()),
(P::HuggingFace, "deepseek-ai/DeepSeek-V4-Pro".into()),
],
}
}
}
enum ChatCompletionResponse {
Immediate(openai::chat::completion::ChatCompletion),
Stream(()),
fn take_random_by_swapping<T>(mut vec1: Vec1<T>, mut rng: impl Rng) -> (T, Option<Vec1<T>>) {
let end = vec1.len().get();
let index = rng.random_range(0..end);
vec1.swap(index, end);
let (body, element) = vec1.into_rtail_and_head();
let body = body.try_into().ok();
(element, body)
}
#[derive(Debug, TypedBuilder)]
struct Upstream {
base_url: Url,
#[builder(setter(strip_option))]
api_key: Option<SecretString>,
}
enum ChatCompletionResponse {
Immediate(openai::chat::completion::ChatCompletion),
Stream {
response: reqwest::Response,
stream_requester: StreamRequester,
},
}
impl IntoResponse for ChatCompletionResponse {
fn into_response(self) -> Response {
match self {
ChatCompletionResponse::Immediate(chat_completion) => {
Json(chat_completion).into_response()
}
ChatCompletionResponse::Stream {
response,
stream_requester,
} => {
let (sse_sink, sse_stream) = async_channel::bounded(256);
stream_requester
.try_send((response, sse_sink))
.expect("this is an unbounded channel");
Sse::new(sse_stream).into_response()
}
}
}
}
enum ChatCompletionError {
NoneSucceeded,
}
impl IntoResponse for ChatCompletionError {
fn into_response(self) -> Response {
match self {
ChatCompletionError::NoneSucceeded => StatusCode::SERVICE_UNAVAILABLE.into_response(),
}
}
}
#[axum::debug_handler]
#[tracing::instrument(skip(state, request))]
async fn create_chat_completion(
_: ChatCompletionsPath,
State(state): State<AppState>,
Json(request): Json<openai::chat::completion::CreateChatCompletionRequest<Shortcut>>,
) -> Result<ChatCompletionResponse, ()> {
if request.stream.unwrap_or(false) {
} else {
) -> Result<ChatCompletionResponse, ChatCompletionError> {
let suggestions = request.model.suggestions();
let mut model_upstreams = VecDeque::new();
for (provider, model) in suggestions {
let upstreams = match provider {
Provider::CloudflareWorkersAi => vec![], // TODO
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));
}
let openai::chat::completion::CreateChatCompletionRequest {
model: _shortcut,
messages,
audio,
frequency_penalty,
function_call,
functions,
logit_bias,
logprobs,
max_completion_tokens,
max_tokens,
metadata,
modalities,
moderation,
n,
parallel_tool_calls,
prediction,
presence_penalty,
prompt_cache_key,
prompt_cache_options,
prompt_cache_retention,
reasoning_effort,
response_format,
safety_identifier,
seed,
service_tier,
stop,
store,
stream,
stream_options,
temperature,
tool_choice,
tools,
top_logprobs,
top_p,
user,
verbosity,
web_search_options,
} = request.clone();
let request = openai::chat::completion::CreateChatCompletionRequest {
model,
messages,
audio,
frequency_penalty,
function_call,
functions,
logit_bias,
logprobs,
max_completion_tokens,
max_tokens,
metadata,
modalities,
moderation,
n,
parallel_tool_calls,
prediction,
presence_penalty,
prompt_cache_key,
prompt_cache_options,
prompt_cache_retention,
reasoning_effort,
response_format,
safety_identifier,
seed,
service_tier,
stop,
store,
stream,
stream_options,
temperature,
tool_choice,
tools,
top_logprobs,
top_p,
user,
verbosity,
web_search_options,
};
let client = reqwest::Client::new();
let mut upstream_request = client.post({
let mut url = upstream.base_url.clone();
{
let mut segments = url.path_segments_mut().expect("TODO");
segments.push("chat");
segments.push("completions");
}
url
});
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;
}
};
let upstream_response = match upstream_response.error_for_status() {
Ok(upstream_response) => upstream_response,
Err(upstream_error) => {
tracing::error!(?upstream_error);
continue;
}
};
if request.stream.unwrap_or(false) {
return Ok(ChatCompletionResponse::Stream {
response: upstream_response,
stream_requester: state.stream_requester,
});
} else {
let response = match upstream_response.json().await {
Ok(response) => response,
Err(json_error) => {
tracing::error!(?json_error);
continue;
}
};
return Ok(ChatCompletionResponse::Immediate(response));
}
}
Err(ChatCompletionError::NoneSucceeded)
}
pub fn create_router() -> Router<AppState> {