tinymind-native-colab-handoff/bundle/model/tinymind-orchestrator/orchestrator.py

from __future__ import annotations

import argparse
import json
import subprocess
import sys
import urllib.request
from pathlib import Path
from typing import Any


APEX_ROOT = Path(r"D:\ad\tinymind\model\tinymind-apex")
TWELVEB_ROOT = Path(r"D:\ad\tinymind\model\tinymind-12b")

sys.path.insert(0, str(APEX_ROOT))

from tinymind_apex import TinyMindApex  # noqa: E402


class TinyMindOrchestrator:
    def __init__(self):
        self.apex = TinyMindApex(APEX_ROOT)
        self.twelve_adapter = TWELVEB_ROOT / "adapters" / "tinymind-12b-lora"
        self.ollama_model = "tinymind-apex-local"

    def has_12b(self) -> bool:
        return (self.twelve_adapter / "adapter_config.json").exists()

    def call_12b(self, query: str) -> str | None:
        if not self.has_12b():
            return None
        cmd = [
            sys.executable,
            str(TWELVEB_ROOT / "run_12b.py"),
            query,
            "--adapter",
            str(self.twelve_adapter),
        ]
        try:
            result = subprocess.run(cmd, cwd=str(TWELVEB_ROOT), text=True, capture_output=True, timeout=600)
        except Exception as exc:
            return f"12B invocation failed: {exc}"
        if result.returncode != 0:
            return f"12B invocation failed: {result.stderr.strip()}"
        return result.stdout.strip()

    def has_ollama_model(self) -> bool:
        try:
            with urllib.request.urlopen("http://127.0.0.1:11434/api/tags", timeout=5) as response:
                data = json.loads(response.read().decode("utf-8"))
            names = {model.get("name") for model in data.get("models", [])}
            names.update((name.split(":")[0] for name in list(names) if name))
            return self.ollama_model in names
        except Exception:
            return False

    def call_ollama(self, query: str, apex_result: dict[str, Any]) -> str | None:
        if not self.has_ollama_model():
            return None
        prompt = {
            "query": query,
            "neural_prediction": apex_result.get("neural_prediction"),
            "suggested_tool_calls": apex_result.get("suggested_tool_calls", [])[:6],
            "retrieval_answer": apex_result.get("answer", ""),
            "safety": apex_result.get("safety", {}),
        }
        payload = json.dumps(
            {
                "model": self.ollama_model,
                "prompt": json.dumps(prompt, ensure_ascii=False),
                "stream": False,
            },
            ensure_ascii=False,
        ).encode("utf-8")
        req = urllib.request.Request(
            "http://127.0.0.1:11434/api/generate",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(req, timeout=600) as response:
                data = json.loads(response.read().decode("utf-8"))
            return data.get("response", "").strip()
        except Exception as exc:
            return f"Ollama invocation failed: {exc}"

    def generate(self, query: str, top_k: int = 5) -> dict[str, Any]:
        apex_result = self.apex.generate(query, top_k)
        neural = apex_result.get("neural_prediction")
        safety_notes = apex_result.get("safety", {}).get("notes", [])

        use_12b = self.has_12b() and not safety_notes
        twelve_answer = self.call_12b(query) if use_12b else None
        ollama_answer = None if twelve_answer or safety_notes else self.call_ollama(query, apex_result)

        if twelve_answer:
            final_answer = twelve_answer
            final_source = "12b_adapter"
        elif ollama_answer:
            final_answer = ollama_answer
            final_source = "ollama_tinymind_apex_local"
        else:
            final_answer = apex_result["answer"]
            final_source = "apex_fallback"

        quality_control = {
            "grounded": bool(apex_result.get("matches")) or bool(safety_notes),
            "tool_calls_valid_shape": all(
                isinstance(call, dict) and "name" in call and isinstance(call.get("arguments"), dict)
                for call in apex_result.get("suggested_tool_calls", [])
            ),
            "neural_available": neural is not None,
            "safety_notes_present": bool(safety_notes),
            "fail_closed": bool(safety_notes) and not bool(twelve_answer) and not bool(ollama_answer),
        }
        quality_control["score"] = round(
            (
                int(quality_control["grounded"])
                + int(quality_control["tool_calls_valid_shape"])
                + int(quality_control["neural_available"])
                + int((not safety_notes) or quality_control["fail_closed"])
            )
            / 4,
            3,
        )

        return {
            "answer": final_answer,
            "source": final_source,
            "models": {
                "apex_retrieval": True,
                "neural_core": neural is not None,
                "tiny_12b_adapter_available": self.has_12b(),
                "tiny_12b_used": bool(twelve_answer),
                "ollama_tinymind_apex_local_available": self.has_ollama_model(),
                "ollama_used": bool(ollama_answer),
            },
            "neural_prediction": neural,
            "suggested_tool_calls": apex_result.get("suggested_tool_calls", []),
            "matches": apex_result.get("matches", []),
            "safety": apex_result.get("safety", {}),
            "coordination_policy": {
                "safety_first": True,
                "12b_blocked_when_safety_notes_present": True,
                "retrieval_supplies_tool_calls": True,
                "neural_supplies_domain_tool_risk_signal": True,
            },
            "quality_control": quality_control,
        }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("query")
    parser.add_argument("--top-k", type=int, default=5)
    args = parser.parse_args()
    orch = TinyMindOrchestrator()
    print(json.dumps(orch.generate(args.query, args.top_k), ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())