captchaforge 0.2.40

//! VLM provider wire protocol: per-backend HTTP request/response adapters
//! (Ollama /api/generate, Anthropic /v1/messages, OpenAI-compatible
//! /v1/chat/completions) + the JSON-from-free-text extractor. Split out of
//! vlm.rs (Law 5). The dispatcher `vlm_query`, `openai_request_body`, and
//! `extract_json` are pub(crate) for the solver + tests in the parent module.

use super::*;

/// Raw Ollama `/api/generate` request body.
#[derive(Debug, Serialize)]
struct OllamaRequest<'a> {
    model: &'a str,
    prompt: &'a str,
    images: Vec<String>,
    stream: bool,
}

/// Relevant fields from Ollama `/api/generate` response.
#[derive(Debug, Deserialize)]
struct OllamaResponse {
    response: String,
}

/// Send a screenshot (base64 PNG/JPEG) and a text prompt to a VLM endpoint.
/// Returns the model's raw text reply. Dispatches to the appropriate
/// per-provider adapter; the calling code stays provider-agnostic.
#[allow(clippy::too_many_arguments)]
pub(crate) async fn vlm_query(
    client: &reqwest::Client,
    provider: VlmProvider,
    endpoint: &str,
    model: &str,
    api_key: Option<&str>,
    image_b64: &str,
    prompt: &str,
    timeout_ms: u64,
) -> Result<String> {
    match provider {
        VlmProvider::Ollama => {
            vlm_query_ollama(client, endpoint, model, image_b64, prompt, timeout_ms).await
        }
        VlmProvider::Anthropic => {
            let key = api_key.ok_or_else(|| anyhow!("anthropic provider needs api_key"))?;
            vlm_query_anthropic(client, endpoint, model, key, image_b64, prompt, timeout_ms).await
        }
        VlmProvider::OpenAI => {
            let key = api_key.ok_or_else(|| anyhow!("openai provider needs api_key"))?;
            // Hosted OpenAI rejects unknown body fields, so don't send the
            // local-only thinking toggle.
            vlm_query_openai(
                client, endpoint, model, key, image_b64, prompt, timeout_ms, false,
            )
            .await
        }
        VlmProvider::LlamaCpp => {
            // LlamaCpp uses OpenAI-compatible /v1/chat/completions; no API key needed.
            // Disable thinking: a local reasoning model (e.g. Qwen3.x) otherwise
            // spends the whole token budget on `reasoning_content` and returns an
            // EMPTY `content`, so every solve silently fails. llama-server honours
            // `chat_template_kwargs.enable_thinking=false` (verified live).
            vlm_query_openai(
                client,
                endpoint,
                model,
                api_key.unwrap_or(""),
                image_b64,
                prompt,
                timeout_ms,
                true,
            )
            .await
        }
    }
}

async fn vlm_query_ollama(
    client: &reqwest::Client,
    endpoint: &str,
    model: &str,
    image_b64: &str,
    prompt: &str,
    timeout_ms: u64,
) -> Result<String> {
    let req = OllamaRequest {
        model,
        prompt,
        images: vec![image_b64.to_string()],
        stream: false,
    };
    let resp = client
        .post(format!("{}/api/generate", endpoint.trim_end_matches('/')))
        .json(&req)
        .timeout(Duration::from_millis(timeout_ms))
        .send()
        .await
        .map_err(|e| anyhow!("ollama request failed: {}", e))?;
    if !resp.status().is_success() {
        let status = resp.status();
        let body = resp.text().await.unwrap_or_default();
        return Err(anyhow!("ollama returned {}: {}", status, body));
    }
    let ollama_resp: OllamaResponse = resp
        .json()
        .await
        .map_err(|e| anyhow!("ollama json parse: {}", e))?;
    Ok(ollama_resp.response)
}

/// Anthropic Messages API call with a single image content block.
/// Image must be base64-encoded PNG (the page-screenshot format the
/// solver already produces); we hardcode `image/png` accordingly.
async fn vlm_query_anthropic(
    client: &reqwest::Client,
    endpoint: &str,
    model: &str,
    api_key: &str,
    image_b64: &str,
    prompt: &str,
    timeout_ms: u64,
) -> Result<String> {
    let body = serde_json::json!({
        "model": model,
        "max_tokens": 512,
        "messages": [{
            "role": "user",
            "content": [
                { "type": "image", "source": { "type": "base64", "media_type": "image/png", "data": image_b64 } },
                { "type": "text", "text": prompt },
            ],
        }],
    });
    let resp = client
        .post(format!("{}/v1/messages", endpoint.trim_end_matches('/')))
        .header("x-api-key", api_key)
        .header("anthropic-version", "2023-06-01")
        .header("content-type", "application/json")
        .json(&body)
        .timeout(Duration::from_millis(timeout_ms))
        .send()
        .await
        .map_err(|e| anyhow!("anthropic request failed: {}", e))?;
    if !resp.status().is_success() {
        let status = resp.status();
        let body = resp.text().await.unwrap_or_default();
        return Err(anyhow!("anthropic returned {}: {}", status, body));
    }
    let val: serde_json::Value = resp
        .json()
        .await
        .map_err(|e| anyhow!("anthropic json parse: {}", e))?;
    // Response shape: { content: [{ type: "text", text: "..." }, ...] }
    val["content"]
        .as_array()
        .and_then(|arr| arr.iter().find(|c| c["type"] == "text"))
        .and_then(|c| c["text"].as_str())
        .map(|s| s.to_string())
        .ok_or_else(|| anyhow!("anthropic response missing text content: {val}"))
}

/// Build the OpenAI-compatible `/v1/chat/completions` request body for a single
/// image+text turn. When `disable_thinking`, adds llama.cpp's
/// `chat_template_kwargs.enable_thinking=false` so a local reasoning model
/// returns its answer in `content` instead of an empty string. Pure, unit
/// tested below.
pub(crate) fn openai_request_body(
    model: &str,
    image_b64: &str,
    prompt: &str,
    disable_thinking: bool,
) -> serde_json::Value {
    let data_url = format!("data:image/png;base64,{}", image_b64);
    let mut body = serde_json::json!({
        "model": model,
        "max_tokens": 512,
        "messages": [{
            "role": "user",
            "content": [
                { "type": "text", "text": prompt },
                { "type": "image_url", "image_url": { "url": data_url } },
            ],
        }],
    });
    if disable_thinking {
        body["chat_template_kwargs"] = serde_json::json!({ "enable_thinking": false });
    }
    body
}

#[allow(clippy::too_many_arguments)]
async fn vlm_query_openai(
    client: &reqwest::Client,
    endpoint: &str,
    model: &str,
    api_key: &str,
    image_b64: &str,
    prompt: &str,
    timeout_ms: u64,
    disable_thinking: bool,
) -> Result<String> {
    let body = openai_request_body(model, image_b64, prompt, disable_thinking);
    let resp = client
        .post(format!(
            "{}/v1/chat/completions",
            endpoint.trim_end_matches('/')
        ))
        .header("authorization", format!("Bearer {}", api_key))
        .header("content-type", "application/json")
        .json(&body)
        .timeout(Duration::from_millis(timeout_ms))
        .send()
        .await
        .map_err(|e| anyhow!("openai request failed: {}", e))?;
    if !resp.status().is_success() {
        let status = resp.status();
        let body = resp.text().await.unwrap_or_default();
        return Err(anyhow!("openai returned {}: {}", status, body));
    }
    let val: serde_json::Value = resp
        .json()
        .await
        .map_err(|e| anyhow!("openai json parse: {}", e))?;
    // Response shape: { choices: [{ message: { content: "..." } }] }
    val["choices"][0]["message"]["content"]
        .as_str()
        .map(|s| s.to_string())
        .ok_or_else(|| anyhow!("openai response missing content: {val}"))
}

/// Extract the first JSON object embedded in free-form text (handles markdown fences).
pub(crate) fn extract_json(text: &str) -> Option<&str> {
    if let Some(start) = text.find("```json") {
        let rest = &text[start + 7..];
        if let Some(end) = rest.find("```") {
            return Some(rest[..end].trim());
        }
    }
    if let Some(start) = text.find('{') {
        let rest = &text[start..];
        let mut depth = 0;
        let mut end = None;
        for (i, c) in rest.char_indices() {
            match c {
                '{' => depth += 1,
                '}' => {
                    depth -= 1;
                    if depth == 0 {
                        end = Some(i + 1);
                        break;
                    }
                }
                _ => {}
            }
        }
        if let Some(e) = end {
            return Some(&rest[..e]);
        }
    }
    None
}