lakehouse/crates/gateway/src/v1/mod.rs

//! Phase 38 — Universal API skeleton (`/v1/*`).
//!
//! OpenAI-compatible shape on top of the existing aibridge → Ollama
//! path. This is the thin slice: single provider, stateless, no
//! streaming. Phase 39 replaces the direct Ollama call with a
//! `ProviderAdapter` trait dispatch; Phase 40 adds routing + fallback.
//!
//! The shape matches OpenAI's `/v1/chat/completions` closely enough
//! that clients using openai-compatible SDKs can point at us with the
//! URL swap alone. We keep the endpoint path `/v1/chat` (not
//! `/v1/chat/completions`) because our PRD declares the terser form;
//! adding the alias is one line in Phase 39 when it matters.

pub mod ollama;
pub mod ollama_cloud;
pub mod langfuse_trace;

use axum::{
    Router,
    extract::State,
    http::StatusCode,
    response::IntoResponse,
    routing::{get, post},
    Json,
};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Clone)]
pub struct V1State {
    pub ai_client: aibridge::client::AiClient,
    pub usage: Arc<RwLock<Usage>>,
    /// Ollama Cloud bearer token. Loaded at startup via
    /// `ollama_cloud::resolve_cloud_key()`. None = cloud routes 503.
    pub ollama_cloud_key: Option<String>,
    /// Phase 40 early deliverable — Langfuse client. None = tracing
    /// disabled (keys missing or container unreachable). Traces are
    /// fire-and-forget: never block the response path.
    pub langfuse: Option<langfuse_trace::LangfuseClient>,
}

#[derive(Default, Clone, Serialize)]
pub struct Usage {
    pub requests: u64,
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub total_tokens: u64,
}

pub fn router(state: V1State) -> Router {
    Router::new()
        .route("/chat", post(chat))
        .route("/usage", get(usage))
        .route("/sessions", get(sessions))
        .with_state(state)
}

// -- Shared types (OpenAI-compatible) --

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct Message {
    pub role: String,
    pub content: String,
}

#[derive(Deserialize, Debug)]
pub struct ChatRequest {
    pub model: String,
    pub messages: Vec<Message>,
    #[serde(default)]
    pub temperature: Option<f64>,
    #[serde(default)]
    pub max_tokens: Option<u32>,
    /// Accepted for shape-compat but ignored in the thin slice —
    /// Phase 38 returns non-streaming even when the client asked for it.
    /// Phase 39+ wires real streaming.
    #[serde(default)]
    pub stream: Option<bool>,
    /// Non-OpenAI extension. Passes through to the provider's thinking
    /// toggle. Default: **false** — hot-path discipline for thinking
    /// models (qwen3.5, qwen3, gpt-oss) that otherwise burn the token
    /// budget on hidden reasoning before visible output starts,
    /// producing empty responses. Set true explicitly when calling an
    /// overseer / reasoning-heavy path.
    #[serde(default)]
    pub think: Option<bool>,
    /// Non-OpenAI extension. Selects the provider adapter. Accepted:
    ///   - None / "ollama" / "local" → local Ollama via sidecar (default)
    ///   - "ollama_cloud" / "cloud" → Ollama Cloud direct HTTPS
    /// Phase 40 adds a routing engine that picks this automatically
    /// from the model name; Phase 38/39 requires explicit selection.
    #[serde(default)]
    pub provider: Option<String>,
}

#[derive(Serialize)]
pub struct ChatResponse {
    pub id: String,
    pub object: &'static str,
    pub created: i64,
    pub model: String,
    pub choices: Vec<Choice>,
    pub usage: UsageBlock,
}

#[derive(Serialize)]
pub struct Choice {
    pub index: u32,
    pub message: Message,
    pub finish_reason: String,
}

#[derive(Serialize, Deserialize, Clone)]
pub struct UsageBlock {
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
    pub total_tokens: u32,
}

// -- Handlers --

async fn chat(
    State(state): State<V1State>,
    Json(req): Json<ChatRequest>,
) -> Result<Json<ChatResponse>, (StatusCode, String)> {
    if req.messages.is_empty() {
        return Err((StatusCode::BAD_REQUEST, "messages must be non-empty".into()));
    }
    if req.stream.unwrap_or(false) {
        tracing::warn!("/v1/chat: stream=true requested but Phase 38 returns non-streaming");
    }

    let provider = req.provider.as_deref().unwrap_or("ollama").to_ascii_lowercase();
    let start_time = chrono::Utc::now();
    let start_instant = std::time::Instant::now();

    let resp = match provider.as_str() {
        "ollama" | "local" | "" => ollama::chat(&state.ai_client, &req)
            .await
            .map_err(|e| (StatusCode::BAD_GATEWAY, format!("ollama local: {e}")))?,
        "ollama_cloud" | "cloud" => {
            let key = state.ollama_cloud_key.as_deref().ok_or((
                StatusCode::SERVICE_UNAVAILABLE,
                "OLLAMA_CLOUD_KEY not configured — set env or populate /root/llm_team_config.json".to_string(),
            ))?;
            ollama_cloud::chat(key, &req)
                .await
                .map_err(|e| (StatusCode::BAD_GATEWAY, format!("ollama cloud: {e}")))?
        }
        other => {
            return Err((
                StatusCode::BAD_REQUEST,
                format!("unknown provider '{other}' — supported: ollama, ollama_cloud"),
            ));
        }
    };

    let end_time = chrono::Utc::now();
    let latency_ms = start_instant.elapsed().as_millis() as u64;

    // Phase 40 — emit Langfuse trace. Fire-and-forget: the clone is
    // cheap (Arc inside), the tokio::spawn never blocks us, a dead
    // Langfuse just logs a warn. Client-visible response latency is
    // untouched.
    if let Some(lf) = &state.langfuse {
        let output = resp.choices.first()
            .map(|c| c.message.content.clone())
            .unwrap_or_default();
        lf.emit_chat(langfuse_trace::ChatTrace {
            provider: provider.clone(),
            model: resp.model.clone(),
            input: req.messages.clone(),
            output,
            prompt_tokens: resp.usage.prompt_tokens,
            completion_tokens: resp.usage.completion_tokens,
            temperature: req.temperature,
            max_tokens: req.max_tokens,
            think: req.think,
            start_time: start_time.to_rfc3339(),
            end_time: end_time.to_rfc3339(),
            latency_ms,
        });
    }

    {
        let mut u = state.usage.write().await;
        u.requests += 1;
        u.prompt_tokens += resp.usage.prompt_tokens as u64;
        u.completion_tokens += resp.usage.completion_tokens as u64;
        u.total_tokens += resp.usage.total_tokens as u64;
    }

    Ok(Json(resp))
}

async fn usage(State(state): State<V1State>) -> impl IntoResponse {
    let snapshot = state.usage.read().await.clone();
    Json(snapshot)
}

// Phase 38 is stateless — no session persistence yet. Return an empty
// list in OpenAI-ish shape so clients that probe this endpoint don't
// 404. Real session state lands in Phase 41 with the profile-system
// expansion.
async fn sessions() -> impl IntoResponse {
    Json(serde_json::json!({
        "data": [],
        "object": "list",
        "note": "Phase 38: stateless. Session state lands in Phase 41.",
    }))
}

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

    #[test]
    fn chat_request_parses_openai_shape() {
        let raw = r#"{
            "model": "qwen3.5:latest",
            "messages": [
                {"role": "system", "content": "You are helpful."},
                {"role": "user", "content": "Hi"}
            ],
            "temperature": 0.2,
            "max_tokens": 100
        }"#;
        let r: ChatRequest = serde_json::from_str(raw).unwrap();
        assert_eq!(r.model, "qwen3.5:latest");
        assert_eq!(r.messages.len(), 2);
        assert_eq!(r.messages[0].role, "system");
        assert_eq!(r.messages[1].content, "Hi");
        assert_eq!(r.temperature, Some(0.2));
        assert_eq!(r.max_tokens, Some(100));
    }

    #[test]
    fn chat_request_accepts_minimal() {
        let raw = r#"{
            "model": "any",
            "messages": [{"role": "user", "content": "hi"}]
        }"#;
        let r: ChatRequest = serde_json::from_str(raw).unwrap();
        assert_eq!(r.temperature, None);
        assert_eq!(r.max_tokens, None);
        assert_eq!(r.stream, None);
    }

    #[test]
    fn usage_counter_default_is_zero() {
        let u = Usage::default();
        assert_eq!(u.requests, 0);
        assert_eq!(u.total_tokens, 0);
    }
}