cortiq-gateway 0.2.46

//! Incoming adapter for OpenAI Chat Completions: `POST /v1/chat/completions`.
//! Translates the OpenAI request body into a canonical [`ChatRequest`], passes it
//! to the pipeline, then converts the canonical [`ChatResponse`] back to OpenAI format.
//!
//! Wire-compatibility rules this adapter enforces (see docs/PROTOCOLS.md §2):
//!   * the request body is parsed leniently — every shape a real OpenAI client
//!     sends (`content: null`, content-part arrays, `tools: null`, a string
//!     `reasoning_effort`, UI-private extras) is accepted, and a body we truly
//!     cannot read comes back as an OpenAI error envelope, never a bare 422;
//!   * the response is strict OpenAI JSON — no optional `null` fields, a
//!     spec-legal `finish_reason`, and the *requested* model id echoed back;
//!   * the SSE stream is re-normalised chunk by chunk and always terminates
//!     with `data: [DONE]`, so a client can never sit waiting on a half-formed
//!     stream from an upstream that improvises.

use crate::error::{GatewayError, Result};
use crate::model::{ChatRequest, GenParams, Message, RequestMeta, RouteInfo, RoutingDirective};
use crate::state::SharedState;
use axum::response::{IntoResponse, Response};
use axum::{extract::State, routing::post, Json, Router};
use serde::Deserialize;

/// Build the `X-Cortiq-*` response headers from routing metadata.
pub(crate) fn cortiq_headers(c: &RouteInfo) -> axum::http::HeaderMap {
    use axum::http::{HeaderMap, HeaderValue};
    fn put(h: &mut HeaderMap, k: &'static str, v: &str) {
        if let Ok(val) = HeaderValue::from_str(v) {
            h.insert(k, val);
        }
    }
    let mut h = HeaderMap::new();
    put(&mut h, "X-Cortiq-Task-Label", &c.task_label);
    put(
        &mut h,
        "X-Cortiq-Complexity-Score",
        &c.complexity_score.to_string(),
    );
    put(&mut h, "X-Cortiq-Complexity-Tier", &c.complexity_tier);
    put(&mut h, "X-Cortiq-Selected-Model", &c.selected_model);
    put(&mut h, "X-Cortiq-Route-Source", &c.route_source);
    put(&mut h, "X-Cortiq-Cost-Usd", &c.cost_usd.to_string());
    if let Some(id) = &c.router_request_id {
        put(&mut h, "X-Cortiq-Request-Id", id);
    }
    if c.tools_dropped {
        put(&mut h, "X-Cortiq-Tools-Dropped", "true");
    }
    h
}

/// Headers every SSE response carries. `X-Accel-Buffering: no` matters when the
/// gateway sits behind nginx/Traefik — without it the proxy buffers the whole
/// stream and the client's chat window just spins.
pub(crate) fn sse_headers(info: &RouteInfo) -> axum::http::HeaderMap {
    use axum::http::HeaderValue;
    let mut headers = cortiq_headers(info);
    headers.insert(
        axum::http::header::CONTENT_TYPE,
        HeaderValue::from_static("text/event-stream; charset=utf-8"),
    );
    headers.insert(
        axum::http::header::CACHE_CONTROL,
        HeaderValue::from_static("no-cache"),
    );
    headers.insert("X-Accel-Buffering", HeaderValue::from_static("no"));
    headers
}

pub fn routes() -> Router<SharedState> {
    Router::new().route("/v1/chat/completions", post(handler))
}

/// `reasoning_effort` is a string in the OpenAI spec (`low` | `medium` | `high`)
/// but a token budget in several local runtimes. Accept both.
#[derive(Deserialize)]
#[serde(untagged)]
enum Effort {
    Budget(u32),
    Level(String),
}

impl Effort {
    fn to_budget(&self) -> Option<u32> {
        match self {
            Self::Budget(n) => Some(*n),
            Self::Level(s) => match s.trim().to_ascii_lowercase().as_str() {
                "none" | "off" | "minimal" => Some(0),
                "low" => Some(512),
                "medium" | "default" | "auto" => Some(2048),
                "high" | "max" => Some(8192),
                _ => None,
            },
        }
    }
}

#[derive(Deserialize)]
struct OpenAiChatRequest {
    model: Option<String>,
    messages: Vec<Message>,
    #[serde(default)]
    temperature: Option<f32>,
    #[serde(default)]
    max_tokens: Option<u32>,
    /// The modern spelling of `max_tokens`; clients send one or the other.
    #[serde(default)]
    max_completion_tokens: Option<u32>,
    #[serde(default)]
    top_p: Option<f32>,
    #[serde(default)]
    think_budget: Option<u32>,
    #[serde(default)]
    reasoning_effort: Option<Effort>,
    /// `null` is a legal value here for several clients, hence `Option`.
    #[serde(default)]
    stream: Option<bool>,
    #[serde(default)]
    tools: Option<Vec<serde_json::Value>>,
    #[serde(flatten)]
    rest: serde_json::Map<String, serde_json::Value>,
}

/// Keys chat UIs attach to the payload for their own bookkeeping. They are not
/// part of the OpenAI request schema, and forwarding them to a strict upstream
/// (llama.cpp, vLLM, `cortiq serve`) is a needless way to earn a 400.
const CLIENT_PRIVATE_KEYS: &[&str] = &[
    "chat_id",
    "session_id",
    "id",
    "metadata",
    "background_tasks",
    "features",
    "variables",
    "model_item",
    "tool_ids",
    "filter_ids",
    "files",
    "citations",
    "params",
    "direct",
];

/// Strip client-private bookkeeping from the passthrough map, and drop a
/// non-string `user` (Open WebUI sends an object there for pipelines, which the
/// OpenAI schema does not allow).
fn sanitize_passthrough(
    mut rest: serde_json::Map<String, serde_json::Value>,
) -> serde_json::Map<String, serde_json::Value> {
    for k in CLIENT_PRIVATE_KEYS {
        rest.remove(*k);
    }
    if rest.get("user").map(|u| !u.is_string()).unwrap_or(false) {
        rest.remove("user");
    }
    rest
}

/// Parse the `model` field value into a routing directive.
/// `cortiq-auto[:profile]` → Auto; otherwise → Pinned(real id).
fn parse_routing(model: &str) -> RoutingDirective {
    if let Some(rest) = model.strip_prefix("cortiq-auto") {
        let profile = rest.strip_prefix(':').map(|p| p.to_string());
        RoutingDirective::Auto { profile }
    } else {
        RoutingDirective::Pinned {
            model_id: model.to_string(),
        }
    }
}

fn now_secs() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

/// Map whatever an upstream calls its stop reason onto the closed set the
/// OpenAI schema defines. Clients that validate the field (LiteLLM, the
/// official SDKs, several agent frameworks) reject anything else.
pub(crate) fn normalize_finish_reason(raw: &str) -> &'static str {
    match raw.trim().to_ascii_lowercase().as_str() {
        "length" | "max_tokens" | "max_new_tokens" | "token_limit" => "length",
        "tool_calls" | "tool_use" | "function_call" => "tool_calls",
        "content_filter" | "safety" => "content_filter",
        _ => "stop",
    }
}

/// A response id every client can key on. Upstreams that omit `id` (or send an
/// empty one) would otherwise leave the field blank, which some clients treat
/// as a malformed completion.
fn response_id(upstream: &str) -> String {
    let trimmed = upstream.trim();
    if !trimmed.is_empty() {
        return trimmed.to_string();
    }
    format!("chatcmpl-{}", crate::admin::random_token(12))
}

async fn handler(
    State(state): State<SharedState>,
    acct: Option<axum::Extension<super::AccountTag>>,
    // Taken as raw bytes rather than `Json<T>`: axum's extractor rejection is a
    // plain-text 422 that OpenAI clients cannot parse, and a client that cannot
    // parse the failure has no answer to show — it just hangs.
    body: axum::body::Bytes,
) -> Result<Response> {
    // hot protocol toggle: if the adapter is disabled in config — return 404
    if !state.live().cfg.protocols.openai_chat {
        return Err(GatewayError::InvalidRequest(
            "openai_chat protocol is disabled".into(),
        ));
    }
    let req: OpenAiChatRequest = serde_json::from_slice(&body)
        .map_err(|e| GatewayError::InvalidRequest(format!("invalid request body: {e}")))?;

    if req.messages.is_empty() {
        return Err(GatewayError::InvalidRequest(
            "messages must not be empty".into(),
        ));
    }

    // Echoed back verbatim: OpenAI returns the model the caller asked for, and
    // clients (Open WebUI among them) match the reply against their own id. The
    // model that actually answered travels in `cortiq.selected_model` and the
    // `X-Cortiq-Selected-Model` header.
    let requested_model = req
        .model
        .as_deref()
        .map(str::trim)
        .filter(|m| !m.is_empty())
        .unwrap_or("cortiq-auto")
        .to_string();
    let stream = req.stream.unwrap_or(false);

    let canonical = ChatRequest {
        routing: parse_routing(&requested_model),
        messages: req.messages,
        tools: req.tools.unwrap_or_default(),
        params: GenParams {
            temperature: req.temperature,
            max_tokens: req.max_tokens.or(req.max_completion_tokens),
            top_p: req.top_p,
            think_budget: req
                .think_budget
                .or_else(|| req.reasoning_effort.as_ref().and_then(Effort::to_budget)),
            stop: Vec::new(),
            passthrough: sanitize_passthrough(req.rest),
        },
        stream,
        meta: RequestMeta {
            protocol: "openai_chat".into(),
            account: acct.map(|e| e.0 .0.clone()).unwrap_or_default(),
            ..Default::default()
        },
    };

    // streaming: provider SSE is re-normalised into strict OpenAI chunks
    if stream {
        let (info, raw) = state.pipeline.run_stream(canonical, &state).await?;
        let headers = sse_headers(&info);
        let normalized = normalize_sse(raw, requested_model);
        return Ok((headers, axum::body::Body::from_stream(normalized)).into_response());
    }

    let resp = state.pipeline.run(canonical, &state).await?;
    let headers = cortiq_headers(&resp.cortiq);

    let choices: Vec<serde_json::Value> = resp
        .choices
        .iter()
        .map(|c| {
            let mut message = serde_json::Map::new();
            message.insert("role".into(), serde_json::json!(c.message.role));
            message.insert("content".into(), serde_json::json!(c.message.content));
            // Absent tool calls are omitted entirely — `"tool_calls": null` is
            // not in the schema and trips strict clients.
            if !c.message.tool_calls.is_empty() {
                message.insert(
                    "tool_calls".into(),
                    serde_json::Value::Array(c.message.tool_calls.clone()),
                );
            }
            let finish = if c.message.tool_calls.is_empty() {
                normalize_finish_reason(&c.finish_reason)
            } else {
                "tool_calls"
            };
            serde_json::json!({
                "index": c.index,
                "message": serde_json::Value::Object(message),
                "finish_reason": finish,
            })
        })
        .collect();

    let mut body = serde_json::json!({
        "id": response_id(&resp.id),
        "object": "chat.completion",
        "created": now_secs(),
        "model": requested_model,
        "choices": choices,
        "usage": {
            "prompt_tokens": resp.usage.prompt_tokens,
            "completion_tokens": resp.usage.completion_tokens,
            "total_tokens": resp.usage.total_tokens,
        }
    });

    if state.live().cfg.cortiq.echo {
        let mut echo = serde_json::json!({
            "task_label": resp.cortiq.task_label,
            "complexity": {
                "score": resp.cortiq.complexity_score,
                "tier": resp.cortiq.complexity_tier,
            },
            "selected_model": resp.cortiq.selected_model,
            "route_source": resp.cortiq.route_source,
            "cost_usd": resp.cortiq.cost_usd,
        });
        if let Some(rid) = &resp.cortiq.router_request_id {
            echo["router_request_id"] = serde_json::json!(rid);
        }
        if resp.cortiq.tools_dropped {
            echo["tools_dropped"] = serde_json::json!(true);
        }
        body["cortiq"] = echo;
    }

    Ok((headers, Json(body)).into_response())
}

/// Rewrite an upstream SSE stream into strict OpenAI `chat.completion.chunk`
/// events.
///
/// Every chunk is reissued with an id, `created`, and the *requested* model id;
/// `message`-shaped chunks are converted to `delta`; stop reasons are mapped to
/// the spec's closed set; empty deltas are dropped. The stream is guaranteed to
/// open with a `role` delta, to carry exactly one terminal `finish_reason`, and
/// to end with `data: [DONE]` — even when the upstream dies mid-flight, in
/// which case the failure is delivered as an SSE `error` event so the client
/// shows a message instead of waiting forever.
fn normalize_sse(
    stream: crate::providers::ChatStream,
    model: String,
) -> impl futures::Stream<Item = Result<bytes::Bytes>> + Send + 'static {
    use futures::StreamExt;
    async_stream::stream! {
        let fallback_id = response_id("");
        let created = now_secs();
        let mut buf = String::new();
        let mut id: Option<String> = None;
        // Which choice indices already carried a `role` delta (n > 1 gets one each).
        let mut roled: std::collections::HashSet<u64> = std::collections::HashSet::new();
        let mut sent_finish = false;
        let mut failure: Option<String> = None;

        // Emit one normalised event; `None` means "nothing worth sending".
        macro_rules! chunk {
            ($choices:expr, $extra:expr) => {{
                let mut evt = serde_json::json!({
                    "id": id.clone().unwrap_or_else(|| fallback_id.clone()),
                    "object": "chat.completion.chunk",
                    "created": created,
                    "model": model,
                    "choices": $choices,
                });
                for (k, v) in $extra {
                    evt[k] = v;
                }
                bytes::Bytes::from(format!("data: {evt}\n\n"))
            }};
        }

        futures::pin_mut!(stream);
        'outer: while let Some(item) = stream.next().await {
            let bytes = match item {
                Ok(b) => b,
                Err(e) => {
                    failure = Some(e.to_string());
                    break 'outer;
                }
            };
            buf.push_str(&String::from_utf8_lossy(&bytes));

            while let Some(idx) = buf.find("\n\n") {
                let event: String = buf.drain(..idx + 2).collect();
                for line in event.lines() {
                    let Some(data) = line.trim_start().strip_prefix("data:") else { continue };
                    let data = data.trim();
                    if data.is_empty() || data == "[DONE]" {
                        continue;
                    }
                    let Ok(v) = serde_json::from_str::<serde_json::Value>(data) else { continue };

                    // An upstream error delivered inside the stream: pass it on
                    // and stop — the client needs to see it, not time out.
                    if let Some(err) = v.get("error").filter(|e| !e.is_null()) {
                        failure = Some(err.to_string());
                        break 'outer;
                    }

                    if id.is_none() {
                        id = v["id"].as_str().filter(|s| !s.is_empty()).map(str::to_string);
                    }

                    let mut choices = Vec::new();
                    if let Some(arr) = v["choices"].as_array() {
                        for (i, c) in arr.iter().enumerate() {
                            // Some upstreams send a full `message` per chunk
                            // instead of a `delta`; both mean the same thing here.
                            let src = if c["delta"].is_object() { &c["delta"] } else { &c["message"] };
                            let index = c["index"].as_u64().unwrap_or(i as u64);
                            let mut delta = serde_json::Map::new();
                            if roled.insert(index) {
                                delta.insert("role".into(), serde_json::json!("assistant"));
                            }
                            if let Some(text) = src["content"].as_str() {
                                if !text.is_empty() {
                                    delta.insert("content".into(), serde_json::json!(text));
                                }
                            }
                            // Reasoning traces and citations are not in the base
                            // schema but every client that understands them
                            // ignores them safely, so they are worth keeping.
                            for key in ["reasoning_content", "reasoning", "thinking", "refusal", "annotations"] {
                                if let Some(val) = src.get(key).filter(|v| !v.is_null()) {
                                    delta.insert(key.into(), val.clone());
                                }
                            }
                            if let Some(tc) = src["tool_calls"].as_array().filter(|a| !a.is_empty()) {
                                delta.insert("tool_calls".into(), serde_json::Value::Array(tc.clone()));
                            }

                            let finish = c["finish_reason"].as_str().map(normalize_finish_reason);
                            if finish.is_none() && delta.is_empty() {
                                continue;
                            }
                            let mut choice = serde_json::json!({
                                "index": index,
                                "delta": serde_json::Value::Object(delta),
                            });
                            if let Some(f) = finish {
                                choice["finish_reason"] = serde_json::json!(f);
                                sent_finish = true;
                            }
                            choices.push(choice);
                        }
                    }

                    // Terminal metadata rides on a choice-less chunk, exactly
                    // as OpenAI ships `usage` when `stream_options` asks for it.
                    let mut extra: Vec<(&str, serde_json::Value)> = Vec::new();
                    if let Some(u) = v.get("usage").filter(|u| !u.is_null()) {
                        extra.push(("usage", u.clone()));
                    }
                    for key in ["cortiq_cost_usd", "cortiq_estimated"] {
                        if let Some(val) = v.get(key).filter(|v| !v.is_null()) {
                            extra.push((key, val.clone()));
                        }
                    }

                    if !choices.is_empty() || !extra.is_empty() {
                        yield Ok(chunk!(choices, extra));
                    }
                }
            }
        }

        if let Some(message) = failure {
            let evt = serde_json::json!({
                "error": { "message": message, "type": "upstream_unavailable", "code": "stream_failed" }
            });
            yield Ok(bytes::Bytes::from(format!("data: {evt}\n\n")));
        }

        // Close the stream properly whatever the upstream did: a client that
        // never sees a stop reason and a [DONE] keeps its spinner turning.
        if !sent_finish {
            let mut delta = serde_json::Map::new();
            if roled.insert(0) {
                delta.insert("role".into(), serde_json::json!("assistant"));
            }
            let choices = vec![serde_json::json!({
                "index": 0,
                "delta": serde_json::Value::Object(delta),
                "finish_reason": "stop",
            })];
            let extra: Vec<(&str, serde_json::Value)> = Vec::new();
            yield Ok(chunk!(choices, extra));
        }
        yield Ok(bytes::Bytes::from("data: [DONE]\n\n"));
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn content_parts_and_null_are_accepted() {
        let body = serde_json::json!({
            "model": "cortiq-auto",
            "messages": [
                {"role": "user", "content": [{"type": "text", "text": "hello "}, {"type": "image_url", "image_url": {"url": "x"}}, {"type": "text", "text": "world"}]},
                {"role": "assistant", "content": null},
                {"role": "tool", "content": "42", "tool_call_id": "call_1"}
            ],
            "tools": null,
            "reasoning_effort": "high",
            "chat_id": "owui-123"
        });
        let req: OpenAiChatRequest = serde_json::from_value(body).expect("must parse");
        assert_eq!(req.messages[0].content, "hello world");
        assert_eq!(req.messages[1].content, "");
        assert_eq!(req.messages[2].tool_call_id.as_deref(), Some("call_1"));
        assert!(req.tools.unwrap_or_default().is_empty());
        assert_eq!(
            req.reasoning_effort.as_ref().and_then(Effort::to_budget),
            Some(8192)
        );
        assert!(!sanitize_passthrough(req.rest).contains_key("chat_id"));
    }

    #[test]
    fn finish_reasons_map_to_the_openai_set() {
        assert_eq!(normalize_finish_reason("max_tokens"), "length");
        assert_eq!(normalize_finish_reason("eos"), "stop");
        assert_eq!(normalize_finish_reason("tool_use"), "tool_calls");
    }

    #[tokio::test]
    async fn stream_is_closed_even_when_the_upstream_stops_early() {
        use futures::StreamExt;
        let upstream = futures::stream::iter(vec![Ok(bytes::Bytes::from(
            "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\n",
        ))]);
        let out: Vec<String> = normalize_sse(Box::pin(upstream), "cortiq-auto".into())
            .map(|b| String::from_utf8_lossy(&b.unwrap()).to_string())
            .collect()
            .await;
        let joined = out.join("");
        assert!(joined.contains("\"role\":\"assistant\""));
        assert!(joined.contains("\"model\":\"cortiq-auto\""));
        assert!(joined.contains("\"finish_reason\":\"stop\""));
        assert!(joined.trim_end().ends_with("data: [DONE]"));
    }
}