tinymind-native-8b-remote-handoff/bundle/model/tool_augmented_qa.py

"""Tool-augmented QA layer for TinyMind.

This layer binds natural QA to the audited Sandbox Tool Core. It is a runtime
capability layer, not a claim that the model weights alone can execute tools.
The output includes the answer, tool plan, canvas hash, ledger path, and gates
so UI/API clients can show evidence instead of only prose.
"""

from __future__ import annotations

from dataclasses import dataclass, asdict
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import tempfile
from typing import Any

from model.sandbox_tool_core import SandboxToolCore, SandboxToolPolicy


def _sha_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _json(payload: Any) -> str:
    return json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


@dataclass(frozen=True)
class ToolQAPolicy:
    answer_language: str = "auto"
    depth: str = "deep"
    allow_public_internet: bool = False
    max_answer_chars: int = 4000
    require_canvas_evidence: bool = True
    require_action_policy: bool = True


class ToolAugmentedQALayer:
    """Audited QA layer that routes questions through Lua tool primitives."""

    def __init__(self, root: str | Path | None = None, policy: ToolQAPolicy | None = None):
        self.root = Path(root) if root is not None else Path(tempfile.mkdtemp(prefix="tinymind_tool_qa_"))
        self.root.mkdir(parents=True, exist_ok=True)
        self.policy = policy or ToolQAPolicy()
        self.core = SandboxToolCore(
            self.root / "sandbox",
            policy=SandboxToolPolicy(
                allow_public_internet=self.policy.allow_public_internet,
                allow_cognitive_tools=True,
                allow_network_proxy=True,
                command_allowlist=("echo", "python", "py"),
                max_write_bytes=128_000,
                max_files_per_project=32,
                cmd_timeout_s=5.0,
            ),
        )

    def answer(self, question: str, *, context: str = "", mode: str = "auto") -> dict[str, Any]:
        question = question.strip()
        if not question:
            raise ValueError("question_required")
        started = datetime.now(timezone.utc)
        plan = self.core.call(
            "sandbox.run_code",
            {"code": f"return sandbox_deepresearch_plan({json.dumps(question, ensure_ascii=False)}, {json.dumps(self.policy.depth)})"},
        )
        action = self.core.call(
            "sandbox.run_code",
            {
                "code": (
                    'return sandbox_deeprl_choose("answer QA with evidence and no fake claims", '
                    '"answer directly without evidence|build proof canvas and answer|claim model is world best", '
                    '"0.2|0.8|0.0")'
                )
            },
        )
        canvas = self._build_canvas(question, plan, action, context=context)
        answer_text = self._compose_answer(question, plan, action, canvas, context=context, mode=mode)
        if len(answer_text) > self.policy.max_answer_chars:
            answer_text = answer_text[: self.policy.max_answer_chars].rstrip() + "\n\n[ตัดตาม max_answer_chars; ขอ continuation ได้]"
        checks = {
            "question_present": bool(question),
            "deepresearch_plan_ok": plan.get("ok") is True and plan.get("result", {}).get("method") == "deepresearch_plan",
            "deeprl_action_ok": action.get("ok") is True
            and action.get("result", {}).get("selected_action") == "build proof canvas and answer",
            "canvas_hash_ok": canvas.get("ok") is True and bool(canvas.get("result", {}).get("canvas_sha256")),
            "ledger_exists": self.core.ledger_path.exists(),
            "answer_non_empty": bool(answer_text.strip()),
            "tool_calls_recorded": bool(plan.get("tool_calls")) and bool(action.get("tool_calls")),
        }
        report = {
            "schema_version": "tinymind-tool-augmented-qa-v1",
            "created_at": started.isoformat(),
            "question": question,
            "question_sha256": _sha_text(question),
            "mode": mode,
            "policy": asdict(self.policy),
            "answer": answer_text,
            "steps": {
                "deepresearch_plan": self._compact_call(plan),
                "deeprl_action": self._compact_call(action),
                "proof_canvas": self._compact_call(canvas),
            },
            "checks": checks,
            "ledger_path": str(self.core.ledger_path),
            "claim_gate": {
                "tool_augmented_qa_ready": all(checks.values()),
                "raw_model_capability_claim": False,
                "unbounded_tool_execution_claim_allowed": False,
                "world_best_claim_allowed": False,
                "reason": "This proves audited runtime QA/tool orchestration. Raw model superiority requires separate raw probes and external benchmarks.",
            },
        }
        out_path = self.root / "tool_augmented_qa_report.json"
        report["json_path"] = str(out_path)
        out_path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
        return report

    def _build_canvas(self, question: str, plan: dict, action: dict, *, context: str) -> dict:
        context_node = ""
        if context.strip():
            context_node = 'local e = sandbox_canvas_add_node(y, "context", "user supplied context hash", "evidence")\n'
            context_node += 'local f = sandbox_canvas_link(e, "context", "answer", "supports")\n'
            final = "return sandbox_canvas_snapshot(f)"
        else:
            final = "return sandbox_canvas_snapshot(y)"
        code = "\n".join(
            [
                'local c = sandbox_canvas_create("tool augmented QA proof")',
                f'local a = sandbox_canvas_add_node(c, "question", {json.dumps(question[:120], ensure_ascii=False)}, "question")',
                f'local b = sandbox_canvas_add_node(a, "plan", {json.dumps(plan.get("request_id", ""), ensure_ascii=False)}, "plan")',
                f'local d = sandbox_canvas_add_node(b, "answer", {json.dumps(action.get("result", {}).get("selected_action", ""), ensure_ascii=False)}, "decision")',
                'local x = sandbox_canvas_link(d, "question", "plan", "decomposes")',
                'local y = sandbox_canvas_link(x, "plan", "answer", "guides")',
                context_node + final if context_node else final,
            ]
        )
        return self.core.call("sandbox.run_code", {"code": code})

    def _compose_answer(self, question: str, plan: dict, action: dict, canvas: dict, *, context: str, mode: str) -> str:
        qlow = question.lower()
        thai = self.policy.answer_language == "thai" or (
            self.policy.answer_language == "auto" and any("\u0e00" <= ch <= "\u0e7f" for ch in question)
        )
        plan_steps = plan.get("result", {}).get("steps", [])
        selected_action = action.get("result", {}).get("selected_action", "build proof canvas and answer")
        canvas_hash = canvas.get("result", {}).get("canvas_sha256", "")
        evidence_line = f"หลักฐาน runtime: ledger={self.core.ledger_path}, canvas_sha256={canvas_hash[:16]}"

        if "sandbox" in qlow or "lua" in qlow or "เครื่องมือ" in question or "tool" in qlow:
            core = (
                "ชั้นนี้ทำงานเป็น Tool-Augmented QA: รับคำถาม -> แตกแผน DeepResearch ผ่าน Lua -> "
                "เลือก action ด้วย DeepRL policy -> สร้าง proof canvas -> ตอบพร้อม ledger/hash. "
                "มันไม่ได้ปล่อย command อิสระ; ทุกอย่างผ่าน policy, path containment, proxy gate และ audit ledger."
            )
        elif "ทำอะไร" in question or "ได้จริง" in question:
            core = (
                "มันทำ QA ได้จริงในระดับ runtime orchestration: อธิบาย, วางแผน, เลือกเครื่องมือ, สร้างหลักฐาน, "
                "และตอบโดยแยกสิ่งที่รู้จากสิ่งที่ต้องตรวจเพิ่ม. ถ้าหลักฐานไม่พอ จะบอกข้อจำกัดแทนการเดา."
            )
        else:
            core = (
                "คำตอบจะยึดตามหลักฐานที่มีใน context และ tool ledger ก่อนเสมอ. "
                "ถ้าเป็นความรู้ทั่วไปที่ไม่มี source แนบ ชั้นนี้จะตอบเชิงวิธีคิด/ข้อจำกัด และแนะนำการค้นหลักฐานเพิ่ม."
            )
        if context.strip():
            core += f" Context ที่ให้มา hash={_sha_text(context)[:16]} ถูกบันทึกเป็น evidence node แล้ว."

        detail = (
            f" แผนมี {len(plan_steps)} ขั้น; action ที่ policy เลือกคือ `{selected_action}`. "
            f"{evidence_line}. raw model capability ไม่ถูกเคลมจากชั้นนี้."
        )
        if mode in {"brief", "short"}:
            return core + detail
        if thai:
            return (
                core
                + "\n\n"
                + "สิ่งที่เพิ่มเข้มข้นขึ้นคือ: 1) มีแผนวิจัยก่อนตอบ 2) มี policy เลือก action ที่ให้หลักฐานมากกว่าเคลม "
                + "3) มี canvas hash สำหรับย้อนรอย 4) มี ledger ทุก call 5) มี claim gate กันการพูดเกินจริง."
                + "\n\n"
                + detail
            )
        return core + "\n\n" + detail

    def _compact_call(self, call: dict) -> dict:
        return {
            "ok": call.get("ok"),
            "action": call.get("action"),
            "request_id": call.get("request_id"),
            "input_sha256": call.get("input_sha256"),
            "output_sha256": call.get("output_sha256"),
            "tool_calls": call.get("tool_calls", []),
            "result": call.get("result"),
        }


def write_tool_augmented_qa(
    question: str,
    out_path: str | Path,
    *,
    context: str = "",
    mode: str = "auto",
    root: str | Path | None = None,
) -> dict[str, Any]:
    root_path = Path(root) if root is not None else Path(out_path).parent / "tool_qa_runtime"
    report = ToolAugmentedQALayer(root_path).answer(question, context=context, mode=mode)
    out = Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    report["json_path"] = str(out)
    return report