hypersteeldb 0.3.2

//! 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() };
            let mut body = json!({
                "model": self.model,
                "messages": to_openai_messages(&system_full, msgs),
                "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> {
        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": 2048,
            });
            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}"))?;
            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)
    }
}