hypersteeldb 0.5.5

//! Paddock provider — a local OpenAI-compatible server (e.g. Qwen3.5-2B). Speaks
//! `POST {base_url}/chat/completions` with function/tool calling. Translates the neutral
//! message/tool types to OpenAI chat messages and back.

use crate::agent::provider::LlmProvider;
use crate::agent::types::{Block, Msg, Role, Stop, ToolSpec, Turn};
use async_trait::async_trait;
use serde_json::{json, Value};

pub struct PaddockProvider {
    client: reqwest::Client,
    base_url: String,
    model: String,
    api_key: Option<String>,
    /// Cached after the server rejects a native `tools` field (e.g. ollama's gemma3n) — subsequent
    /// requests inline the tool schema in the system prompt so the harness can recover `<tool_call>`
    /// blocks from the reply.
    prompt_only_tools: std::sync::atomic::AtomicBool,
}

impl PaddockProvider {
    pub fn new(base_url: String, model: String, api_key: Option<String>) -> PaddockProvider {
        PaddockProvider {
            client: reqwest::Client::new(),
            base_url,
            model,
            api_key,
            prompt_only_tools: std::sync::atomic::AtomicBool::new(false),
        }
    }
}

/// Inline tool schema for models the server won't accept `tools` for. Emit-format matches Hermes,
/// which the harness already recovers.
fn inline_tools_prompt(tools: &[ToolSpec]) -> String {
    let mut out = String::from(
        "\n\nYou have access to these tools. To call one, emit ONLY a single line in this exact form (no prose around it):\n<tool_call>{\"name\":\"<tool_name>\",\"arguments\":{...}}</tool_call>\nAfter the tool result comes back, either call another tool or give the final answer as plain text.\n\nAvailable tools:\n",
    );
    for t in tools {
        out.push_str(&format!("- {} — {}\n  schema: {}\n", t.name, t.description, t.schema));
    }
    out
}

/// Flatten neutral messages into OpenAI chat messages (system first).
fn to_openai_messages(system: &str, msgs: &[Msg]) -> Vec<Value> {
    let mut out = vec![json!({ "role": "system", "content": system })];
    for m in msgs {
        match m.role {
            Role::User => {
                // user text and/or tool results
                let mut text = String::new();
                for b in &m.blocks {
                    match b {
                        Block::Text(t) => {
                            if !text.is_empty() {
                                text.push('\n');
                            }
                            text.push_str(t);
                        }
                        Block::ToolResult { id, content, .. } => {
                            out.push(json!({ "role": "tool", "tool_call_id": id, "content": content }));
                        }
                        _ => {}
                    }
                }
                if !text.is_empty() {
                    out.push(json!({ "role": "user", "content": text }));
                }
            }
            Role::Assistant => {
                let mut text = String::new();
                let mut tool_calls = Vec::new();
                for b in &m.blocks {
                    match b {
                        Block::Text(t) => text.push_str(t),
                        Block::ToolUse { id, name, input } => {
                            tool_calls.push(json!({
                                "id": id,
                                "type": "function",
                                "function": { "name": name, "arguments": input.to_string() }
                            }));
                        }
                        _ => {}
                    }
                }
                let mut msg = json!({ "role": "assistant", "content": if text.is_empty() { Value::Null } else { Value::String(text) } });
                if !tool_calls.is_empty() {
                    msg["tool_calls"] = Value::Array(tool_calls);
                }
                out.push(msg);
            }
        }
    }
    out
}

fn to_openai_tools(tools: &[ToolSpec]) -> Vec<Value> {
    tools
        .iter()
        .map(|t| json!({ "type": "function", "function": { "name": t.name, "description": t.description, "parameters": t.schema } }))
        .collect()
}

#[async_trait]
impl LlmProvider for PaddockProvider {
    fn name(&self) -> &str {
        "paddock"
    }

    async fn chat(&self, system: &str, msgs: &[Msg], tools: &[ToolSpec]) -> Result<Turn, String> {
        use std::sync::atomic::Ordering;
        let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
        let mut prompt_only = self.prompt_only_tools.load(Ordering::Relaxed);
        let resp = loop {
            let system_full = if prompt_only && !tools.is_empty() { format!("{system}{}", inline_tools_prompt(tools)) } else { system.to_string() };
            // `/no_think` is the soft switch Qwen-family reasoning models honour, and it has to ride on the
            // LAST USER turn — in the system prompt it is ignored. It is also the only switch that survives an
            // OpenAI-compatible proxy which drops unknown body fields, which ollama's /v1 endpoint does.
            // Harmless to a model that does not recognise it.
            let mut messages = to_openai_messages(&system_full, msgs);
            if let Some(last_user) = messages
                .iter_mut()
                .rev()
                .find(|m| m.get("role").and_then(|r| r.as_str()) == Some("user"))
            {
                if let Some(c) = last_user.get("content").and_then(|c| c.as_str()) {
                    last_user["content"] = json!(format!("{c}\n\n/no_think"));
                }
            }
            let mut body = json!({
                "model": self.model,
                "messages": messages,
                "max_tokens": 2048
            });
            if !prompt_only && !tools.is_empty() {
                body["tools"] = json!(to_openai_tools(tools));
                body["tool_choice"] = json!("auto");
            }
            let mut req = self.client.post(&url).json(&body);
            if let Some(k) = &self.api_key {
                req = req.bearer_auth(k);
            }
            let resp = req.send().await.map_err(|e| format!("paddock request: {e}"))?;
            if resp.status().is_success() {
                break resp;
            }
            let status = resp.status();
            let text = resp.text().await.unwrap_or_default();
            // fall back to prompt-only when the server rejects a native `tools` field
            if !prompt_only && !tools.is_empty() && (text.contains("does not support tools") || text.contains("tool_choice")) {
                self.prompt_only_tools.store(true, Ordering::Relaxed);
                prompt_only = true;
                continue;
            }
            return Err(format!("paddock {status}: {text}"));
        };
        let v: Value = resp.json().await.map_err(|e| format!("paddock decode: {e}"))?;
        let choice = v.get("choices").and_then(|c| c.get(0)).ok_or("paddock: no choices")?;
        let message = choice.get("message").ok_or("paddock: no message")?;

        let text = message.get("content").and_then(|c| c.as_str()).unwrap_or("").to_string();
        let mut tool_uses = Vec::new();
        if let Some(calls) = message.get("tool_calls").and_then(|c| c.as_array()) {
            for call in calls {
                let id = call.get("id").and_then(|x| x.as_str()).unwrap_or("").to_string();
                let func = call.get("function");
                let name = func.and_then(|f| f.get("name")).and_then(|x| x.as_str()).unwrap_or("").to_string();
                let args_str = func.and_then(|f| f.get("arguments")).and_then(|x| x.as_str()).unwrap_or("{}");
                let input: Value = serde_json::from_str(args_str).unwrap_or(json!({}));
                tool_uses.push((id, name, input));
            }
        }
        let stop = match choice.get("finish_reason").and_then(|f| f.as_str()) {
            Some("tool_calls") => Stop::ToolUse,
            Some("stop") | Some("length") => Stop::EndTurn,
            _ => {
                if tool_uses.is_empty() {
                    Stop::EndTurn
                } else {
                    Stop::ToolUse
                }
            }
        };
        Ok(Turn { text, tool_uses, stop })
    }

    /// Constrain decoding to `schema` via the OpenAI-compatible `response_format`.
    ///
    /// Ollama, llama.cpp's server, vLLM and LM Studio all accept this and compile the schema into a decoding
    /// grammar internally, so the model can only emit tokens the schema permits. That is why this works on
    /// models with no tool-calling ability at all: nothing is being asked of the model except to continue, and
    /// the sampler does the rest.
    ///
    /// Two degradations, because not every server implements the whole thing:
    ///   * a server that rejects `json_schema` is retried with `json_object`, which constrains the output to
    ///     *some* JSON and leaves the shape to the prompt;
    ///   * a server that rejects `response_format` entirely returns `Ok(None)`, so the caller falls back rather
    ///     than seeing an error it cannot act on.
    async fn chat_json(
        &self,
        system: &str,
        msgs: &[Msg],
        schema: &Value,
        name: &str,
    ) -> Result<Option<Value>, String> {
        // Ollama's native endpoint FIRST, when this looks like ollama. Its OpenAI-compatible shim silently
        // drops `think`, `chat_template_kwargs` and the `/no_think` soft switch alike, so a hybrid reasoning
        // model spends the whole completion budget on a `reasoning` field and returns zero characters of
        // content. The native route accepts `think: false` and takes the JSON schema directly as `format`,
        // which is the only way to actually get structured output out of such a model here. Falls through to
        // the portable path when this is not ollama, or when the native call does not produce usable JSON.
        if let Some(base) = self.base_url.trim_end_matches('/').strip_suffix("/v1") {
            if let Some(v) = self.ollama_native_json(base, system, msgs, schema).await? {
                return Ok(Some(v));
            }
        }

        let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));

        // Do NOT put the schema JSON in the prompt. The first version did, on the reasoning that the grammar
        // guarantees shape while only the prompt conveys intent — and `qwen2.5:0.5b` promptly returned the facet
        // names "surface", "token", "head" and "tail", which are the schema's own field names. A small model
        // treats anything in its context as material to copy, so handing it a vocabulary of key names invites
        // exactly that. The grammar already enforces the shape; the caller's own system prompt explains the task.
        let system_full = format!("{system}\n\nReply with JSON only. No prose, no code fence.");

        for mode in ["json_schema", "json_object"] {
            let mut body = json!({
                "model": self.model,
                "messages": to_openai_messages(&system_full, msgs),
                "max_tokens": 4096,
                // Structured output is used for decisions, not prose, so sample greedily: the reference's
                // curator runs at temperature 0. Without this, two curations of the SAME clusters produced
                // materially different ontologies — one run named a facet `sensing-modality`, the next
                // `5g_compatible` — which makes the vocabulary irreproducible for no benefit.
                "temperature": 0,
                // Turn OFF chain-of-thought. A hybrid reasoning model spends its budget thinking before it
                // emits anything, and the grammar constrains only the ANSWER — `qwen3:1.7b` came back
                // `finish_reason: length` having spent all 4096 completion tokens with ZERO characters of
                // content, the reasoning having gone to a separate `reasoning` field. Reasoning buys nothing
                // here: the schema dictates the shape and the judgment wanted is a naming decision.
                //
                // These two flags are sent because some servers honour them — ollama's /v1 endpoint does NOT,
                // which is why the prompt also carries the `/no_think` soft switch below. An unrecognised field
                // is ignored rather than rejected, so sending both costs nothing.
                "think": false,
                "chat_template_kwargs": { "enable_thinking": false },
            });
            body["response_format"] = if mode == "json_schema" {
                json!({ "type": "json_schema",
                        "json_schema": { "name": name, "strict": true, "schema": schema } })
            } else {
                json!({ "type": "json_object" })
            };

            let mut req = self.client.post(&url).json(&body);
            if let Some(k) = &self.api_key {
                req = req.bearer_auth(k);
            }
            let resp = req.send().await.map_err(|e| format!("paddock request: {e}"))?;
            if !resp.status().is_success() {
                let status = resp.status();
                let text = resp.text().await.unwrap_or_default();
                // this server has no structured-output support at all; let the caller fall back
                if text.contains("response_format") || status.as_u16() == 400 {
                    continue;
                }
                return Err(format!("paddock {status}: {text}"));
            }
            let v: Value = resp.json().await.map_err(|e| format!("paddock decode: {e}"))?;
            if std::env::var("STEELDB_DEBUG_JSON").is_ok() {
                eprintln!(
                    "[dbg] mode={mode} finish={:?} completion_tokens={:?} sys_chars={} content_chars={}",
                    v.pointer("/choices/0/finish_reason"),
                    v.pointer("/usage/completion_tokens"),
                    system_full.len(),
                    v.pointer("/choices/0/message/content").and_then(|c| c.as_str()).map(|c| c.len()).unwrap_or(0)
                );
            }
            let content = v
                .get("choices")
                .and_then(|c| c.get(0))
                .and_then(|c| c.get("message"))
                .and_then(|m| m.get("content"))
                .and_then(|c| c.as_str())
                .unwrap_or("");
            if content.trim().is_empty() {
                continue;
            }
            // A constrained reply should be bare JSON, but a server that only honoured `json_object` may still
            // wrap it in prose or a code fence, so reuse the tolerant extractor.
            if let Ok(parsed) = serde_json::from_str::<Value>(content.trim()) {
                return Ok(Some(parsed));
            }
            if let Some(parsed) = crate::vocabulary::extract_json(content) {
                return Ok(Some(parsed));
            }
        }
        Ok(None)
    }
}

impl PaddockProvider {
    /// Ollama's native `/api/chat`: `format` takes a JSON schema directly, and `think: false` actually works.
    ///
    /// Returns `Ok(None)` when this is not ollama or the reply is unusable, so the caller can fall back to the
    /// OpenAI-compatible route rather than failing outright.
    async fn ollama_native_json(
        &self,
        base: &str,
        system: &str,
        msgs: &[Msg],
        schema: &Value,
    ) -> Result<Option<Value>, String> {
        let mut messages = vec![json!({ "role": "system", "content": system })];
        for m in to_openai_messages("", msgs).into_iter().filter(|m| {
            m.get("role").and_then(|r| r.as_str()) != Some("system")
        }) {
            messages.push(m);
        }
        let body = json!({
            "model": self.model,
            "messages": messages,
            "stream": false,
            "think": false,
            "format": schema,
            "options": { "temperature": 0, "num_predict": 4096 },
        });
        let resp = match self.client.post(format!("{base}/api/chat")).json(&body).send().await {
            Ok(r) => r,
            // not ollama, or not listening: let the caller try the portable route
            Err(_) => return Ok(None),
        };
        if !resp.status().is_success() {
            return Ok(None);
        }
        let v: Value = match resp.json().await {
            Ok(v) => v,
            Err(_) => return Ok(None),
        };
        let content = v.pointer("/message/content").and_then(|c| c.as_str()).unwrap_or("");
        if std::env::var("STEELDB_DEBUG_JSON").is_ok() {
            eprintln!(
                "[dbg] ollama-native done={:?} content_chars={}",
                v.get("done_reason"),
                content.len()
            );
        }
        if content.trim().is_empty() {
            return Ok(None);
        }
        Ok(serde_json::from_str::<Value>(content.trim())
            .ok()
            .or_else(|| crate::vocabulary::extract_json(content)))
    }
}