codecraft/cli/src/models.rs

//! Model catalog: the server's models ("server:provider/model") and the user's
//! own providers ("provider/model"), and streaming requests to either.

use crate::conf::Config;
use crate::server::Server;
use anyhow::{Result, anyhow};
use codecraft_core::config::ModelConfig;
use codecraft_core::provider::{ChatRequest, StreamEvent, ToolDef};
use codecraft_core::types::{Message, Usage};
use serde_json::Value;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

pub const SERVER_PREFIX: &str = "server:";

#[derive(Debug, Clone)]
pub struct ModelInfo {
    /// Selection key, e.g. "server:gpu/qwen-coder" or "ollama/llama3".
    pub spec: String,
    pub label: String,
    pub config: ModelConfig,
}

impl ModelInfo {
    pub fn cost(&self, u: &Usage) -> f64 {
        self.config.cost(u)
    }
}

#[derive(Default, Clone)]
pub struct Catalog {
    pub models: Vec<ModelInfo>,
    pub currency: String,
}

fn model_from_json(v: &Value) -> ModelConfig {
    serde_json::from_value(v.clone()).unwrap_or_else(|_| ModelConfig { id: v["id"].as_str().unwrap_or("").into(), ..Default::default() })
}

impl Catalog {
    /// Collect models from the server config (if signed in) and local providers.
    pub async fn load(cfg: &Config, server_cfg: Option<&Value>) -> Self {
        let mut models = Vec::new();
        let mut currency = String::new();
        if let Some(sc) = server_cfg {
            currency = sc["currency"].as_str().unwrap_or("").to_string();
            for p in sc["models"].as_array().into_iter().flatten() {
                let provider = p["provider"].as_str().unwrap_or("");
                for m in p["models"].as_array().into_iter().flatten() {
                    let config = model_from_json(m);
                    let name = m["name"].as_str().unwrap_or(&config.id).to_string();
                    models.push(ModelInfo { spec: format!("{SERVER_PREFIX}{provider}/{}", config.id), label: format!("server · {provider} · {name}"), config });
                }
            }
        }
        for (name, pc) in &cfg.providers {
            if !pc.enabled {
                continue;
            }
            let mut list: Vec<ModelConfig> = pc.models.iter().map(|m| m.resolved()).collect();
            if list.is_empty() {
                // Discover (e.g. Ollama, LM Studio) – quick, failures leave the list empty.
                let provider = codecraft_core::provider::build(codecraft_core::provider::http_client(), pc);
                if let Ok(Ok(found)) = tokio::time::timeout(std::time::Duration::from_secs(5), provider.list_models()).await {
                    list = found;
                }
            }
            for m in list {
                let label = format!("{name} · {}", m.name.clone().unwrap_or_else(|| m.id.clone()));
                models.push(ModelInfo { spec: format!("{name}/{}", m.id), label, config: m });
            }
        }
        Self { models, currency }
    }

    pub fn find(&self, spec: &str) -> Option<&ModelInfo> {
        self.models.iter().find(|m| m.spec == spec)
    }
}

/// Stream one model request to the right backend.
#[allow(clippy::too_many_arguments)]
pub async fn stream(
    cfg: &Config,
    server: Option<&Server>,
    model: &ModelInfo,
    system: String,
    messages: Vec<Message>,
    tools: Vec<ToolDef>,
    effort: Option<String>,
    session: &str,
    agent: &str,
    tx: mpsc::Sender<StreamEvent>,
    cancel: CancellationToken,
) -> Result<()> {
    if let Some(rest) = model.spec.strip_prefix(SERVER_PREFIX) {
        let server = server.ok_or_else(|| anyhow!("not signed in to a CodeCraft server (codecraft login)"))?;
        return server.stream(rest, system, messages, tools, effort, session, agent, tx, cancel).await;
    }
    let (pname, _) = model.spec.split_once('/').ok_or_else(|| anyhow!("invalid model '{}'", model.spec))?;
    let pc = cfg.providers.get(pname).ok_or_else(|| anyhow!("unknown provider '{pname}' – run /setup"))?;
    let provider = codecraft_core::provider::build(codecraft_core::provider::http_client(), pc);
    let mut mc = model.config.clone();
    if let Some(e) = effort {
        mc.effort = Some(e);
    }
    let max_tokens = mc.max_output.unwrap_or(16_000);
    let req = ChatRequest { model: mc, system, messages, tools, max_tokens, reasoning: true, session_id: Some(session.to_string()) };
    provider.stream(req, tx, cancel).await
}

/// Known endpoints for the setup wizard.
pub struct Preset {
    pub id: &'static str,
    pub label: &'static str,
    pub kind: codecraft_core::config::ProviderKind,
    pub base_url: &'static str,
    pub needs_key: bool,
    pub key_hint: &'static str,
}

pub fn presets() -> Vec<Preset> {
    use codecraft_core::config::ProviderKind::*;
    vec![
        Preset { id: "anthropic", label: "Anthropic (Claude)", kind: Anthropic, base_url: "", needs_key: true, key_hint: "sk-ant-… from console.anthropic.com" },
        Preset { id: "openai", label: "OpenAI", kind: Openai, base_url: "https://api.openai.com/v1", needs_key: true, key_hint: "sk-… from platform.openai.com" },
        Preset { id: "gemini", label: "Google Gemini", kind: Openai, base_url: "https://generativelanguage.googleapis.com/v1beta/openai", needs_key: true, key_hint: "key from aistudio.google.com" },
        Preset { id: "openrouter", label: "OpenRouter", kind: Openai, base_url: "https://openrouter.ai/api/v1", needs_key: true, key_hint: "sk-or-… from openrouter.ai" },
        Preset { id: "opencode", label: "opencode Zen / Go", kind: Openai, base_url: "https://opencode.ai/zen/v1", needs_key: true, key_hint: "key from opencode.ai" },
        Preset { id: "ollama", label: "Ollama (local)", kind: Openai, base_url: "http://localhost:11434/v1", needs_key: false, key_hint: "" },
        Preset { id: "lmstudio", label: "LM Studio (local)", kind: Openai, base_url: "http://localhost:1234/v1", needs_key: false, key_hint: "" },
        Preset { id: "llamacpp", label: "llama.cpp / vLLM (local)", kind: Openai, base_url: "http://localhost:8080/v1", needs_key: false, key_hint: "" },
        Preset { id: "custom", label: "Other OpenAI-compatible endpoint", kind: Openai, base_url: "", needs_key: false, key_hint: "optional" },
    ]
}

/// Local model servers that answer right now (for the wizard).
pub async fn detect_local() -> Vec<&'static str> {
    let mut found = Vec::new();
    for p in presets().into_iter().filter(|p| p.base_url.starts_with("http://localhost")) {
        let ok = codecraft_core::provider::http_client()
            .get(format!("{}/models", p.base_url))
            .timeout(std::time::Duration::from_millis(800))
            .send()
            .await
            .is_ok_and(|r| r.status().is_success());
        if ok {
            found.push(p.id);
        }
    }
    found
}