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

feat: more models, make decisions a wrapper for system one

Browse Source
This commit is contained in:
Jacob
2026-10-10 19:06:02 -04:00
parent 9a848e2c78
commit e0e1f0781e
7 changed files with 248 additions and 167 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/openai/chat/completion.rs
View File

@@ -24,7 +24,7 @@ pub struct CreateChatCompletionRequest<Model> {
/// Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing
/// frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim.
#[serde(skip_serializing_if = "Option::is_none")]
pub frequency_penalty: Option<f32>,
pub frequency_penalty: Option<f64>,
/// Deprecated in favor of `tool_choice`. Controls which (if any) function is called by the model.
#[serde(skip_serializing_if = "Option::is_none")]
@@ -37,7 +37,7 @@ pub struct CreateChatCompletionRequest<Model> {
/// Modify the likelihood of specified tokens appearing in the completion. Accepts a JSON object
/// that maps tokens (specified by their token ID in the tokenizer) to an associated bias value from -100 to 100.
#[serde(skip_serializing_if = "Option::is_none")]
pub logit_bias: Option<HashMap<String, f32>>,
pub logit_bias: Option<HashMap<String, f64>>,
/// Whether to return log probabilities of the output tokens or not. If true, returns the log
/// probabilities of each output token returned in the `content` of `message`.
@@ -83,7 +83,7 @@ pub struct CreateChatCompletionRequest<Model> {
/// Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in
/// the text so far, increasing the model's likelihood to talk about new topics.
#[serde(skip_serializing_if = "Option::is_none")]
pub presence_penalty: Option<f32>,
pub presence_penalty: Option<f64>,
/// Used by OpenAI to cache responses for similar requests to optimize your cache hit rates. Replaces the `user` field.
#[serde(skip_serializing_if = "Option::is_none")]
@@ -137,7 +137,7 @@ pub struct CreateChatCompletionRequest<Model> {
/// What sampling temperature to use, between 0 and 2. Higher values like 0.8 will make the output more random.
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
pub temperature: Option<f64>,
/// Controls which (if any) tool is called by the model.
#[serde(skip_serializing_if = "Option::is_none")]
@@ -155,7 +155,7 @@ pub struct CreateChatCompletionRequest<Model> {
/// An alternative to sampling with temperature, called nucleus sampling, where the model considers the
/// results of the tokens with top_p probability mass.
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f32>,
pub top_p: Option<f64>,
/// Deprecated in favor of `safety_identifier` and `prompt_cache_key`. A stable identifier for your end-users.
#[serde(skip_serializing_if = "Option::is_none")]

1
src/openai/decisions.rs
View File

@@ -8,7 +8,6 @@ pub struct DecisionRequest<Model> {
/// The text or images to evaluate for every question.
/// Can be a simple string or an array of user messages containing text and images.
#[serde(flatten)]
pub input: DecisionInput,
/// Array of questions to evaluate against the input

2
src/openai/embeddings.rs
View File

@@ -73,7 +73,7 @@ pub struct Embedding {
#[serde(untagged)]
pub enum EmbeddingOutput {
/// Array of floats when encoding_format is "float"
FloatArray(Vec<f32>),
FloatArray(Vec<f64>),
/// Base64-encoded string when encoding_format is "base64"
Base64String(String),
}

58
src/openai/system_one.rs
View File

@@ -6,14 +6,27 @@ use std::collections::HashMap;
/// This represents the complete request body that can be sent to the API.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct SystemOneRequest<Model> {
/// The model that handles the request.
pub model: Model,
/// The content to evaluate. Can be a string, object, or array.
///
/// For text content, use a plain string. For structured data (like chat logs or application state),
/// use an object or array.
pub state: State,
/// The model that handles the request.
pub model: Model,
/// Optional file inputs (images, audio, etc.).
///
/// Array of data URLs for input files.
///
/// For images: data:image/...;base64,...
/// For audio: data:audio/...;base64,...
#[serde(skip_serializing_if = "Option::is_none")]
pub files: Option<Vec<String>>,
/// Alternative name for file inputs (alias for `files`).
#[serde(skip_serializing_if = "Option::is_none")]
pub images: Option<Vec<String>>,
/// A map of typed questions to evaluate against the state.
///
@@ -171,7 +184,7 @@ pub enum Answer {
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct NoulAnswer {
/// The yes/no answer on a scale from 0 (no) to 1 (yes).
pub noul: f32,
pub noul: f64,
}
/// Answer to a Choice question (multiple choice).
@@ -181,12 +194,12 @@ pub struct ChoiceAnswer {
pub choice: String,
/// The probability of each option, they sum to 1.
pub probabilities: HashMap<String, f32>,
pub probabilities: HashMap<String, f64>,
/// How certain the model is (0 to 1).
///
/// 0 means all options are equally likely, 1 means maximum certainty.
pub confidence: f32,
pub confidence: f64,
}
/// Answer to a Score question.
@@ -195,18 +208,18 @@ pub struct ScoreAnswer {
/// The expected level index, weighted by probability.
///
/// Can be between two levels (e.g., 1.5 would be halfway between level 1 and 2).
pub score: f32,
pub score: f64,
/// The description of each level index.
pub legend: HashMap<String, String>,
/// The probability of each level index, they sum to 1.
pub probabilities: HashMap<String, f32>,
pub probabilities: HashMap<String, f64>,
/// How certain the model is (0 to 1).
///
/// Derived from the probabilities distribution.
pub confidence: f32,
pub confidence: f64,
}
/// Token usage information for the request.
@@ -218,32 +231,3 @@ pub struct Usage {
/// Number of output tokens (always 0 for this API).
pub output_tokens: u32,
}
/// Complete request structure that includes optional file inputs.
///
/// This is a convenience structure that combines the basic request with optional
/// file inputs for multimodal models.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct MultimodalRequest {
/// The content to evaluate (same as in Request).
pub state: State,
/// Optional file inputs (images, audio, etc.).
///
/// Array of data URLs for input files.
///
/// For images: data:image/...;base64,...
/// For audio: data:audio/...;base64,...
#[serde(skip_serializing_if = "Option::is_none")]
pub files: Option<Vec<String>>,
/// Alternative name for file inputs (alias for `files`).
#[serde(skip_serializing_if = "Option::is_none")]
pub images: Option<Vec<String>>,
/// The model that handles the request (same as in Request).
pub model: String,
/// A map of typed questions to evaluate (same as in Request).
pub questions: HashMap<String, Question>,
}

251
src/server/decisions.rs
View File

@@ -1,25 +1,16 @@
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 std::collections::HashMap;
use std::hash::RandomState;
use strum::EnumString;
use strum::VariantArray;
@@ -31,20 +22,12 @@ pub struct V1DecisionsPath;
Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
)]
#[strum(serialize_all = "lowercase")]
pub enum Shortcut {}
impl Shortcut {
fn suggestions(&self) -> Vec<(Provider, EcoString)> {
use Provider as P;
use Shortcut as S;
match self {
_ => unreachable!()
}
}
pub enum Shortcut {
S,
M,
}
pub struct DecisionsResponse(openai::embeddings::CreateEmbeddingResponse<String>);
pub struct DecisionsResponse(openai::decisions::DecisionResponse<String>);
impl IntoResponse for DecisionsResponse {
fn into_response(self) -> Response {
@@ -66,95 +49,175 @@ impl IntoResponse for DecisionsError {
}
}
#[tracing::instrument(skip(state, request))]
#[tracing::instrument(skip(app_state, request))]
pub async fn create_decision(
_: V1DecisionsPath,
State(state): State<AppState>,
State(app_state): State<AppState>,
Json(request): Json<openai::decisions::DecisionRequest<Shortcut>>,
) -> Result<DecisionsResponse, DecisionsError> {
let shortcut = request.model.clone();
let mut model_upstreams = VecDeque::new();
use super::system_one;
for (provider, model) in shortcut.suggestions() {
let upstreams: Vec<Upstream> = match provider {
Provider::CloudflareWorkersAi => unreachable!(),
Provider::GoogleAiStudio => unreachable!(),
Provider::Groq => unreachable!(),
Provider::HuggingFace => unreachable!(),
Provider::LiquidAi => unreachable!(),
Provider::Mistral => unreachable!(),
Provider::OpenRouter => unreachable!(),
Provider::ZAi => unreachable!(),
let openai::decisions::DecisionRequest {
model,
input,
questions,
} = request;
let model = match model {
Shortcut::S => system_one::Shortcut::S,
Shortcut::M => system_one::Shortcut::M,
};
if let Ok(upstreams) = Vec1::try_from(upstreams) {
model_upstreams.push_back((model, upstreams));
let (state, images, files) = match input {
openai::decisions::DecisionInput::Messages(messages) => {
let mut images = vec![];
let mut texts = vec![];
for openai::decisions::DecisionInputMessage { role, content } in messages {
match content {
openai::decisions::DecisionInputContent::Text(text) => texts.push(text),
openai::decisions::DecisionInputContent::Parts(decision_input_parts) => {
for decision_input_part in decision_input_parts {
match decision_input_part {
openai::decisions::DecisionInputPart::Image { image_url, detail } => images.push(image_url),
openai::decisions::DecisionInputPart::Text { text } => texts.push(text),
}
}
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 text = texts.join("\n\n");
let openai::decisions::DecisionRequest { model: _shortcut, input, questions } = request.clone();
let files = images.clone();
let request = openai::decisions::DecisionRequest { model, input, questions };
(openai::system_one::State::Text(text), Some(images), Some(files))
}
openai::decisions::DecisionInput::Text(text) => (openai::system_one::State::Text(text), None, None),
};
let client = reqwest::Client::builder()
.tls_certs_only(
webpki_root_certs::TLS_SERVER_ROOT_CERTS
.iter()
.filter_map(|cert| Certificate::from_der(cert).ok()),
let questions = questions.into_iter().map(|question| match question {
openai::decisions::DecisionQuestion::Predicate(predicate_question) => todo!(),
openai::decisions::DecisionQuestion::Choice(choice_question) => todo!(),
openai::decisions::DecisionQuestion::Score(score_question) => todo!(),
} ).collect();
let request = openai::system_one::SystemOneRequest {
model,
state,
files,
images,
questions,
};
system_one::create_system_one(system_one::V1SystemOnePath, State(app_state), Json(request))
.await
.map(
|system_one::SystemOneResponse(openai::system_one::SystemOneResponse {
model,
answers,
usage,
})| {
let answers = answers
.into_iter()
.map(|(id, answer)| match answer {
openai::system_one::Answer::Noul(openai::system_one::NoulAnswer {
noul,
}) => openai::decisions::DecisionAnswer::Predicate(
openai::decisions::PredicateAnswer {
name: Some(id),
probability: noul,
},
),
openai::system_one::Answer::Choice(openai::system_one::ChoiceAnswer {
choice,
probabilities,
confidence,
}) => openai::decisions::DecisionAnswer::Choice(
openai::decisions::ChoiceAnswer {
name: Some(id),
choice: openai::decisions::Choice::String(choice),
confidence,
probabilities: probabilities
.into_iter()
.map(|(choice_id, probability)| {
openai::decisions::ChoiceProbability {
value: openai::decisions::Choice::String(choice_id),
probability,
}
})
.collect(),
},
),
openai::system_one::Answer::Score(openai::system_one::ScoreAnswer {
score,
legend,
probabilities,
confidence,
}) => openai::decisions::DecisionAnswer::Score(
openai::decisions::ScoreAnswer {
name: Some(id),
score,
confidence,
probabilities: {
let reverse_legend = HashMap::<_, _, RandomState>::from_iter(
legend.into_iter().map(|(index, value)| {
(
value,
index
.parse::<i32>()
.expect("this should be an integer, I believe"),
)
.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
});
probabilities
.into_iter()
.map(|(choice_id, probability)| {
openai::decisions::ScoreProbability {
value: reverse_legend
.get(&choice_id)
.expect("it should exist, I think")
.clone(),
label: choice_id,
probability,
}
})
.collect()
},
},
),
})
.collect();
let upstream_response = match upstream_request.send().await {
Ok(upstream_response) => upstream_response,
Err(upstream_error) => {
tracing::error!(?upstream_error);
continue;
}
let openai::system_one::Usage {
input_tokens,
output_tokens,
} = usage;
let input_tokens = input_tokens as i32;
let output_tokens = output_tokens as i32;
let total_tokens = input_tokens + output_tokens;
let input_tokens_details = openai::decisions::TokenUsageDetails {
cache_write_tokens: 0,
cached_tokens: 0,
};
let output_tokens_details = openai::decisions::OutputTokenDetails {
reasoning_tokens: 0,
};
let usage = openai::decisions::TokenUsage {
input_tokens,
input_tokens_details,
output_tokens,
output_tokens_details,
total_tokens,
};
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(DecisionsResponse(response));
}
Err(DecisionsError::NoneSucceeded)
DecisionsResponse(openai::decisions::DecisionResponse {
model,
answers,
usage,
})
},
)
.map_err(|_| DecisionsError::NoneSucceeded)
}

24
src/server/embeddings.rs
View File

@@ -32,8 +32,17 @@ pub struct V1EmbeddingsPath;
)]
#[strum(serialize_all = "lowercase")]
pub enum Shortcut {
#[strum(serialize = "lfm-2.5-embedding-350m")]
Lfm2_5Embedding350m,
#[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 {
@@ -42,8 +51,17 @@ impl Shortcut {
use Shortcut as S;
match self {
S::Lfm2_5Embedding350m => {
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())]
}
}
}

35
src/server/system_one.rs
View File

@@ -32,6 +32,7 @@ pub struct V1SystemOnePath;
)]
#[strum(serialize_all = "lowercase")]
pub enum Shortcut {
S,
M,
}
@@ -41,14 +42,26 @@ impl Shortcut {
use Shortcut as S;
match self {
S::S => {
vec![
(P::LiquidAi, "d1:free".into()),
(P::CloudflareWorkersAi, "clef-omni".into()),
(P::CloudflareWorkersAi, "clef-flash".into()),
(P::OpenRouter, "respan/span-01-lite:free".into()),
(P::OpenRouter, "respan/span-01-lite".into()),
]
}
S::M => {
vec![(P::LiquidAi, "d1:free".into())]
vec![
(P::CloudflareWorkersAi, "clef".into()),
(P::OpenRouter, "inception/mercury-decide:free".into()),
]
}
}
}
}
pub struct SystemOneResponse(openai::system_one::SystemOneResponse<String>);
pub struct SystemOneResponse(pub openai::system_one::SystemOneResponse<String>);
impl IntoResponse for SystemOneResponse {
fn into_response(self) -> Response {
@@ -94,7 +107,7 @@ pub async fn create_system_one(
}| {
Upstream::builder()
.url(
format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1/systemone")
format!("https://api.cloudflare.com/client/v4/accounts/{account_id}/run/@cf/cloudflare/{model}")
.parse()
.expect("TODO"),
)
@@ -147,15 +160,22 @@ pub async fn create_system_one(
tracing::info!(?model, ?upstream);
let openai::system_one::SystemOneRequest { model: _shortcut, state, questions } = request.clone();
let openai::system_one::SystemOneRequest {
model: _shortcut,
state,
files,
images,
questions,
} = request.clone();
let request = openai::system_one::SystemOneRequest {
model,
state,
files,
images,
questions,
};
let client = reqwest::Client::builder()
.tls_certs_only(
webpki_root_certs::TLS_SERVER_ROOT_CERTS
@@ -165,9 +185,7 @@ pub async fn create_system_one(
.build()
.expect("TODO");
let mut upstream_request = client
.post(upstream.url)
.json(&request);
let mut upstream_request = client.post(upstream.url).json(&request);
if let Some(api_key) = upstream.api_key {
let authorization =
@@ -200,7 +218,6 @@ pub async fn create_system_one(
continue;
};
let response = match upstream_response.json().await {
Ok(response) => response,
Err(json_error) => {