scripts/check_llama_cpp_smoke.py

"""Smoke-test the optional llama.cpp text runtime with an external GGUF model."""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from src.models.llama_cpp_runner import get_text_runtime_fallbacks, reply_as_object
from src.pipeline import generate_object_diary


DEFAULT_GGUF_REPO = "Qwen/Qwen2.5-1.5B-Instruct-GGUF"
DEFAULT_GGUF_FILE = "qwen2.5-1.5b-instruct-q4_k_m.gguf"
TEXT_FALLBACK_MARKER = "text-fallback-to-mock"


def run_llama_cpp_smoke(
    *,
    model_path: Path,
    description: str,
    mode: str,
    save_trace: bool,
) -> dict[str, Any]:
    if not model_path.exists():
        raise FileNotFoundError(f"GGUF model path does not exist: {model_path}")

    previous_text_backend = os.environ.get("OBJECTVERSE_TEXT_BACKEND")
    previous_text_model_path = os.environ.get("TEXT_MODEL_PATH")
    try:
        os.environ["OBJECTVERSE_TEXT_BACKEND"] = "llama-cpp"
        os.environ["TEXT_MODEL_PATH"] = str(model_path)
        result = generate_object_diary(
            None,
            description,
            mode,
            save=save_trace,
        )
        chat_reply = reply_as_object(
            result.persona.model_dump(mode="json"),
            "What did you see today?",
        )
        chat_fallbacks = get_text_runtime_fallbacks()
    finally:
        _restore_env("OBJECTVERSE_TEXT_BACKEND", previous_text_backend)
        _restore_env("TEXT_MODEL_PATH", previous_text_model_path)

    payload = {
        "status": "pass",
        "model_path": _display_model_path(model_path),
        "description": description,
        "mode": mode,
        "trace_id": result.trace.trace_id,
        "trace_path": result.trace_path,
        "model_runtime": result.trace.model_runtime,
        "fallbacks": result.trace.fallbacks,
        "object_name": result.object_understanding.object.name,
        "character_name": result.persona.persona.character_name,
        "diary_title": result.diary.title,
        "chat_reply_preview": chat_reply[:160],
        "chat_fallbacks": chat_fallbacks,
    }
    if result.trace.model_runtime.get("text") != "llama-cpp text generation":
        payload["status"] = "fail"
        payload["error"] = "trace did not record llama-cpp text generation"
    if TEXT_FALLBACK_MARKER in result.trace.fallbacks:
        payload["status"] = "fail"
        payload["error"] = "trace included text-fallback-to-mock"
    if TEXT_FALLBACK_MARKER in chat_fallbacks:
        payload["status"] = "fail"
        payload["error"] = "chat included text-fallback-to-mock"
    return payload


def _restore_env(key: str, previous_value: str | None) -> None:
    if previous_value is None:
        os.environ.pop(key, None)
    else:
        os.environ[key] = previous_value


def _print_json(payload: dict[str, Any]) -> None:
    print(json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True), flush=True)


def _display_model_path(model_path: Path) -> str:
    try:
        return str(model_path.resolve().relative_to(PROJECT_ROOT))
    except ValueError:
        return model_path.name


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--model-path",
        type=Path,
        default=Path("models") / DEFAULT_GGUF_FILE,
        help=f"Path to {DEFAULT_GGUF_FILE} or another external GGUF file.",
    )
    parser.add_argument(
        "--description",
        default="old white coffee mug on a developer desk",
    )
    parser.add_argument("--mode", default="Cynical")
    parser.add_argument("--save-trace", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = _parse_args()
    try:
        payload = run_llama_cpp_smoke(
            model_path=args.model_path,
            description=args.description,
            mode=args.mode,
            save_trace=args.save_trace,
        )
    except Exception as exc:
        _print_json(
            {
                "status": "fail",
                "model_path": _display_model_path(args.model_path),
                "recommended_repo": DEFAULT_GGUF_REPO,
                "recommended_file": DEFAULT_GGUF_FILE,
                "error_type": type(exc).__name__,
                "error": str(exc),
            }
        )
        raise SystemExit(1) from exc

    _print_json(payload)
    if payload["status"] != "pass":
        raise SystemExit(1)


if __name__ == "__main__":
    main()