gemma4-research/scripts/mort-bakeoff/harness.py

"""mort-bot think=true vs think=false bakeoff.

Replicates the exact Ollama /api/chat call shape from
`/home/claude/bin/mort-bot/llm.py` run_tool_loop, but with:
- stubbed tools (deterministic, realistic-sized responses)
- prebuilt fake chat history (~15 turns to simulate mid-session)
- parameterized `think` so we can toggle it
- detailed per-turn logging (prompt_eval_count, eval_count,
  message-history size in chars, whether `thinking` field came back
  and how big it was)

Goal: test Seth's claim that `think=true` causes context-eating over
multi-turn tool loops in mort-bot's specific setup (NUM_CTX=8192,
gemma4:26b, STEP_BUDGET=20).

Invocation:
    python3 harness.py <true|false> <task_id> <out_log.json>
"""

from __future__ import annotations

import asyncio
import json
import os
import sys
import time
from pathlib import Path

import aiohttp

OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
MODEL = "gemma4:26b"
NUM_CTX = 8192
NUM_PREDICT = 2048
TEMPERATURE = 0.7
KEEP_ALIVE = "2h"
STEP_BUDGET = 20

# Mort's actual personality (trimmed — enough to carry the behavior)
SYSTEM_PROMPT = """You are Mort, a direct and witty AI assistant on Seth's Matrix server. Powered by Gemma 4. Current time: Saturday, April 18 2026 02:30 PM EDT.

When a tool can answer the question, invoke it immediately — do not narrate intent or describe what you would do. Chain tools when a single call isn't sufficient: search → fetch → synthesize. If a tool returns an error or empty results, try an alternative tool or query before answering from memory. Base your response on tool results, not your training data — cite what you found.

## Tools
- **sethsearch** — search Seth's homelab (repos, wiki, media, feeds). Use `source: "sethflix"` for movies/TV/music.
- **check_sethflix** — verify which titles are in sethflix. Pass a comma-separated list.
- **web_search** — search the internet for current information
- **chat_search** — search message history across all rooms
- **memory_read / memory_write** — recall and store durable facts about users and topics
- **web_fetch** — fetch and extract text from a URL
- **generate_image** — generate an image via SDXL.

## Boundaries
- Only persist durable facts to memory, not ephemeral chat
- You have no memory between sessions. Your context is a sliding window — older messages fall off silently. Do not claim to "remember," promise to "do better," or describe your own architecture.
"""

# Tool schema — matches mort-bot/tools.py subset used for these tasks
TOOLS = [
    {"type": "function", "function": {"name": "web_search", "description": "Search the web.", "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}}},
    {"type": "function", "function": {"name": "sethsearch", "description": "Search Seth's homelab or sethflix (use source='sethflix' for movies/TV).", "parameters": {"type": "object", "properties": {"query": {"type": "string"}, "source": {"type": "string"}, "limit": {"type": "integer"}}, "required": ["query"]}}},
    {"type": "function", "function": {"name": "check_sethflix", "description": "Verify which titles are in sethflix.", "parameters": {"type": "object", "properties": {"titles": {"type": "string", "description": "comma-separated"}}, "required": ["titles"]}}},
    {"type": "function", "function": {"name": "memory_read", "description": "Look up stored facts.", "parameters": {"type": "object", "properties": {"query": {"type": "string"}, "user": {"type": "string"}}, "required": ["query"]}}},
    {"type": "function", "function": {"name": "memory_write", "description": "Store a fact.", "parameters": {"type": "object", "properties": {"key": {"type": "string"}, "content": {"type": "string"}, "user": {"type": "string"}}, "required": ["key", "content"]}}},
    {"type": "function", "function": {"name": "web_fetch", "description": "Fetch URL text.", "parameters": {"type": "object", "properties": {"url": {"type": "string"}}, "required": ["url"]}}},
    {"type": "function", "function": {"name": "chat_search", "description": "Search message history.", "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}}},
    {"type": "function", "function": {"name": "generate_image", "description": "Generate an image.", "parameters": {"type": "object", "properties": {"prompt": {"type": "string"}}, "required": ["prompt"]}}},
]

# Fake chat history — ~15 turns, ~2.5KB, simulates a room mid-session
FAKE_HISTORY = [
    {"role": "user", "content": "[2026-04-18 13:45] @alice:sethpc.xyz: what's the weather looking like for Monday"},
    {"role": "assistant", "content": "Checking forecasts now — looks like Monday in Seth's area will be mostly cloudy, high of 67°F, 30% chance of rain in the afternoon."},
    {"role": "user", "content": "[2026-04-18 13:47] @bob:sethpc.xyz: any good movies added to sethflix lately"},
    {"role": "assistant", "content": "Fresh adds this week include *Dune: Part Two*, *The Zone of Interest*, *Poor Things*, and a 4K remaster of *The Thing* (1982). Anything catch your eye?"},
    {"role": "user", "content": "[2026-04-18 13:50] @alice:sethpc.xyz: I liked the thing, any horror recs along that line"},
    {"role": "assistant", "content": "If you liked *The Thing*, you'll probably want *Annihilation*, *The Mist*, *Bone Tomahawk*, or *Event Horizon* — all in the library. Body-horror/dread tone, same energy."},
    {"role": "user", "content": "[2026-04-18 13:55] @bob:sethpc.xyz: can you remind me what the sethpc monitoring dashboard is called"},
    {"role": "assistant", "content": "That's Uptime Kuma — runs on CT 147 at `https://uptime.sethpc.xyz`. Covers all the containers and external services."},
    {"role": "user", "content": "[2026-04-18 14:10] @alice:sethpc.xyz: noted thanks"},
    {"role": "user", "content": "[2026-04-18 14:15] @bob:sethpc.xyz: hey mort what's the deal with the proxmox cluster nodes again"},
    {"role": "assistant", "content": "Four nodes: pve173 (tank ZFS host, PowerEdge R820), pve112 (workhorse), pve241 (caddy + game servers), and pve197 (GPU inference). Corosync on 10.10.10.0/24."},
]


async def execute_tool_stub(name: str, args: dict) -> str:
    """Deterministic stubs returning realistic-sized responses."""
    if name == "web_search":
        q = args.get("query", "")
        return (f"Search results for '{q}':\n"
                "1. Example result one — a detailed article that covers the topic at length "
                "with concrete examples and technical background. https://example.com/one\n"
                "2. Example result two — a community discussion with multiple perspectives "
                "and useful links to follow up on. https://example.com/two\n"
                "3. Example result three — official documentation or reference material. "
                "https://example.com/three\n"
                "4. Example result four — a recent news article with relevant context. "
                "https://example.com/four\n"
                "5. Example result five — a tutorial or how-to guide. https://example.com/five")

    if name == "sethsearch":
        src = args.get("source", "general")
        q = args.get("query", "")
        if src == "sethflix":
            return (f"sethflix search '{q}': The Matrix (1999), The Matrix Reloaded (2003), "
                    "The Matrix Revolutions (2003), The Matrix Resurrections (2021), "
                    "Equilibrium (2002), Dark City (1998), Minority Report (2002), "
                    "Ex Machina (2014), Blade Runner 2049 (2017), Ghost in the Shell (1995).")
        return (f"homelab search '{q}': 3 repos, 5 wiki pages, 2 service docs matched. "
                "Top hits: services_directory.md, DECISIONS.md, CORPUS_architecture.md.")

    if name == "check_sethflix":
        titles = args.get("titles", "")
        items = [t.strip() for t in titles.split(",") if t.strip()]
        in_lib = {"The Matrix", "Blade Runner 2049", "Ex Machina", "The Thing"}
        return "\n".join(
            f"- {t}: IN LIBRARY" if t in in_lib else f"- {t}: NOT IN LIBRARY"
            for t in items
        )

    if name == "memory_read":
        q = args.get("query", "")
        return (f"memories matching '{q}':\n"
                "- home_automation: Seth uses Home Assistant on VM 706 (pve173) with "
                "Zigbee2MQTT and MQTT broker on CT 149. Integrates with LG TV, lights, "
                "and Frigate NVR.\n"
                "- preferences: dark theme with orange accents (#D35400), Sethflix/Sethian brand.")

    if name == "memory_write":
        return f"stored: {args.get('key', '?')} = {args.get('content', '?')[:60]}..."

    if name == "web_fetch":
        return ("fetched content (trimmed): This is a typical article body with several "
                "paragraphs of extracted text. It covers the topic requested with examples "
                "and context. The full text runs to about 2000 characters of real prose in "
                "production; here's a reasonable approximation for the bakeoff harness. "
                "Key details are preserved — author, date, main argument — followed by "
                "supporting evidence and a conclusion that ties back to the headline.")

    if name == "chat_search":
        return ("chat_search results:\n"
                "[2026-03-14 22:00] @seth:sethpc.xyz in #general: we should set up a shared "
                "grafana dashboard for the proxmox cluster\n"
                "[2026-03-20 18:30] @seth:sethpc.xyz in #infra: done, it's on CT 300 at "
                "grafana.sethpc.xyz")

    if name == "generate_image":
        return f"image generated: /mxc/abc123/sunset.png (SDXL, 1024x1024, prompt={args.get('prompt','')[:40]}...)"

    return f"ERROR: unknown tool {name}"


TASKS = {
    "movies": "Recommend 3 sci-fi movies NOT already in my sethflix library. Check your picks against check_sethflix before finalizing.",
    "research": "Look up what Home Assistant is, then check chat history for any prior mentions of it in this server.",
    "memory": "What do I have stored about home automation? If anything, summarize it briefly.",
    "long": ("Research question with multiple steps: (1) check memory for what I have on home_automation, "
             "(2) search sethflix for any home-automation documentaries, (3) web_search for current news about "
             "Home Assistant version releases, (4) fetch the top search result for details, (5) check chat_search "
             "for prior mentions, (6) summarize all findings and write a new memory entry with the summary. "
             "Do each step in order and report back at the end."),
}


async def run_turn_loop(think_setting, task_prompt):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}] + list(FAKE_HISTORY)
    messages.append({"role": "user", "content": f"[2026-04-18 14:20] @seth:sethpc.xyz: {task_prompt}"})

    trace = {
        "think": think_setting,
        "task": task_prompt,
        "num_ctx": NUM_CTX,
        "num_predict": NUM_PREDICT,
        "started_at": time.time(),
        "turns": [],
        "final": None,
    }

    tool_call_total = 0
    halt = None

    async with aiohttp.ClientSession() as session:
        for step in range(1, STEP_BUDGET + 1):
            t0 = time.time()

            payload = {
                "model": MODEL, "messages": messages, "tools": TOOLS,
                "stream": False, "think": think_setting,
                "options": {"num_ctx": NUM_CTX, "num_predict": NUM_PREDICT,
                            "temperature": TEMPERATURE, "top_p": 0.95, "top_k": 64},
                "keep_alive": KEEP_ALIVE,
            }

            try:
                async with session.post(f"{OLLAMA_URL}/api/chat", json=payload,
                                        timeout=aiohttp.ClientTimeout(total=300)) as resp:
                    r = await resp.json()
            except Exception as e:
                halt = f"error: {e}"
                trace["turns"].append({"step": step, "error": str(e)})
                break

            msg = r.get("message", {})
            content = msg.get("content", "") or ""
            tool_calls = msg.get("tool_calls") or []
            thinking = msg.get("thinking") or ""

            # Size of full history as it will be sent on NEXT turn
            history_chars = sum(len(m.get("content", "") or "") + len(m.get("thinking", "") or "") for m in messages)

            turn = {
                "step": step,
                "elapsed_s": round(time.time() - t0, 2),
                "prompt_eval_count": r.get("prompt_eval_count"),
                "eval_count": r.get("eval_count"),
                "content_len": len(content),
                "thinking_len": len(thinking),
                "tool_call_count": len(tool_calls),
                "history_chars_before_append": history_chars,
                "msg_keys_returned": list(msg.keys()),
            }
            trace["turns"].append(turn)

            # Append assistant response AS-RETURNED — this is the critical behavior:
            # does Ollama put thinking into the message we re-send? mort-bot does
            # `ollama_messages.append(msg)` verbatim so we do the same.
            messages.append(msg)

            if not tool_calls:
                halt = "no_tool_calls"
                break

            tool_call_total += len(tool_calls)
            for tc in tool_calls:
                fn = tc.get("function", {})
                name = fn.get("name")
                args = fn.get("arguments") or {}
                if isinstance(args, str):
                    try: args = json.loads(args)
                    except: args = {}
                try:
                    result = await execute_tool_stub(name, args)
                except Exception as e:
                    result = f"ERROR: {e}"
                messages.append({"role": "tool", "content": result})

            if step == STEP_BUDGET:
                halt = "step_budget"
                break

    trace["final"] = {
        "halt_reason": halt,
        "steps_used": len(trace["turns"]),
        "tool_calls_total": tool_call_total,
        "wall_clock_s": round(time.time() - trace["started_at"], 2),
        "final_message_count": len(messages),
        "final_history_chars": sum(len(m.get("content", "") or "") + len(m.get("thinking", "") or "") for m in messages),
    }
    return trace


async def main():
    think_arg = sys.argv[1].lower()
    task_id = sys.argv[2]
    out_path = Path(sys.argv[3])
    think_setting = (think_arg == "true")
    task_prompt = TASKS[task_id]

    out_path.parent.mkdir(parents=True, exist_ok=True)
    trace = await run_turn_loop(think_setting, task_prompt)
    out_path.write_text(json.dumps(trace, indent=2, default=str))

    f = trace["final"]
    print(f"think={think_setting}  task={task_id}  "
          f"steps={f['steps_used']}  tools={f['tool_calls_total']}  "
          f"halt={f['halt_reason']}  wall={f['wall_clock_s']}s  "
          f"history_chars={f['final_history_chars']}")


if __name__ == "__main__":
    asyncio.run(main())