katniss-ai/src/server/system_one.rs

use crate::CloudflareCredentials;
use crate::openai;
use crate::server::AppState;
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/systemone")]
pub struct V1SystemOnePath;

#[derive(
    Clone, Debug, EnumString, DeserializeFromStr, strum::Display, SerializeDisplay, VariantArray,
)]
#[strum(serialize_all = "lowercase")]
pub enum Shortcut {
    S,
    M,
    L,
}

#[derive(Debug, Clone)]
enum Provider {
    CloudflareWorkersAi,
    LiquidAi,
    OpenCode,
    OpenRouter,
    VercelAiGateway,
}

impl Shortcut {
    fn suggestions(&self) -> Vec<(Provider, EcoString)> {
        use Provider as P;
        use Shortcut as S;

        match self {
            S::S => {
                vec![
                    (P::VercelAiGateway, "perplexity/pplx-decider-flash-v1.2".into()),
                    (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()),
                    (P::VercelAiGateway, "convaiinnovations/laya-free".into()),
                    (P::VercelAiGateway, "convaiinnovations/laya".into()),
                ]
            }

            S::M => {
                vec![
                    (P::OpenCode, "jev-1.13-free".into()),
                    (P::VercelAiGateway, "perplexity/pplx-decider-v1.2".into()),
                    (P::CloudflareWorkersAi, "clef".into()),
                    (P::OpenRouter, "inception/mercury-decide:free".into()),
                ]
            }

            S::L => {
                vec![
                    (P::VercelAiGateway, "liquid/d1-ultra".into())
                ]
            }
        }
    }
}

pub struct SystemOneResponse(pub openai::system_one::SystemOneResponse<String>);

impl IntoResponse for SystemOneResponse {
    fn into_response(self) -> Response {
        let Self(system_one) = self;

        Json(system_one).into_response()
    }
}

pub enum SystemOneError {
    NoneSucceeded,
}

impl IntoResponse for SystemOneError {
    fn into_response(self) -> Response {
        match self {
            SystemOneError::NoneSucceeded => StatusCode::SERVICE_UNAVAILABLE.into_response(),
        }
    }
}

#[tracing::instrument(skip(state, request))]
pub async fn create_system_one(
    _: V1SystemOnePath,
    State(state): State<AppState>,
    Json(request): Json<openai::system_one::SystemOneRequest<Shortcut>>,
) -> Result<SystemOneResponse, SystemOneError> {
    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/run/@cf/cloudflare/{model}")
                                        .parse()
                                        .expect("TODO"),
                                )
                                .api_key(api_key)
                                .build()
                        },
                    )
                    .collect(),
                Provider::LiquidAi => state
                    .app_args
                    .liquid_ai_api_key
                    .iter()
                    .cloned()
                    .map(|api_key| {
                        Upstream::builder()
                            .url("https://api.liquid.ai/decisions/v1/systemone".parse().expect("TODO"))
                            .api_key(api_key)
                            .build()
                    })
                    .collect(),
                Provider::OpenCode => state
                    .app_args
                    .opencode_console_api_key
                    .iter()
                    .cloned()
                    .map(|api_key| {
                        Upstream::builder()
                            .url("https://opencode.ai/inference/systemone/v1/systemone".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/systemone".parse().expect("TODO"))
                            .api_key(api_key)
                            .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(),
            };

        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::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
                    .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: serde_json::Value = match upstream_response.json().await {
            Ok(response) => response,
            Err(json_error) => {
                tracing::error!(?json_error);
                continue;
            }
        };

        match openai::system_one::SystemOneResponse::deserialize(
            response.clone(), // TODO: not need to clone
        ) {
            Ok(response) => return Ok(SystemOneResponse(response)),
            Err(structure_error) => {
                #[derive(Deserialize)]
                struct CloudflareRunResponse {
                    result: openai::system_one::SystemOneResponse<String>,
                    success: bool,
                    errors: Vec<serde_json::Value>,   // TODO
                    messages: Vec<serde_json::Value>, // TODO
                }

                match CloudflareRunResponse::deserialize(response) {
                    Ok(response) => return Ok(SystemOneResponse(response.result)),
                    Err(cloudflare_structure_error) => {
                        tracing::error!(?structure_error);
                        tracing::error!(?cloudflare_structure_error);
                    }
                }
            }
        }
    }
    Err(SystemOneError::NoneSucceeded)
}