Topic 5: Module 10 lab

"""Module 10 lab: the patterns of this module, end to end, in one run.

1. An aggregator (gateway) fronts the notes and calendar servers.
2. The eval harness scores a single agent that reaches both through the gateway.
3. An orchestrator splits a compound question across concurrent workers that share the gateway.
4. The notes agent is exposed as an MCP server, and an outer agent calls it as a tool.
Everything runs in memory with the KeywordAgent stand-in; the JSON report lands in m10_lab_report.json.
"""
from __future__ import annotations

import json
import shutil
import tempfile
import time
from pathlib import Path

import anyio

from mcp import Client

from examples.m10_aggregator import Aggregator
from examples.m10_agent_server import build_agent_server
from examples.m10_calendar_server import build_calendar_server
from examples.m10_common import NOTES_DIR, KeywordAgent, estimate_tokens
from examples.m10_eval import CASES, run_eval, stand_in_factory
from examples.m10_orchestrator import QUESTION, split_stand_in
from notes_assistant.host import Host
from notes_assistant.server import build_server
from notes_assistant.store import NoteStore


def through_gateway(case: dict) -> dict:
    """The gateway renames notes__search_notes to gateway__notes_search_notes, and so on."""
    tool = case["expect_tool"]
    renamed = tool if tool == "none" else "gateway__" + tool.replace("__", "_")
    return {**case, "expect_tool": renamed}


async def main() -> None:
    work = Path(tempfile.mkdtemp(prefix="m10-lab-"))
    shutil.copytree(NOTES_DIR, work / "notes")  # the lab never touches the real notes folder
    store = NoteStore(work / "notes")
    report: dict = {}
    gateway = Aggregator({"notes": build_server(store), "calendar": build_calendar_server()})

    async with Client(gateway) as gw:
        specs = await Host({"gateway": gw}).load_tools()
        report["gateway_tools"] = [s["function"]["name"] for s in specs]
        report["definition_tokens_estimate"] = estimate_tokens(specs)[1]
        print(f"1. gateway exposes {len(specs)} tools, about {report['definition_tokens_estimate']} tokens")

        ev = await run_eval({"gateway": gw}, stand_in_factory, "keyword stand-in", [through_gateway(c) for c in CASES])
        report["eval"] = {k: ev[k] for k in ("cases", "selection_accuracy", "task_success", "selection_ci95")}
        print(f"2. eval through the gateway: selection {ev['selection_accuracy']:.3f}, success {ev['task_success']:.3f} "
              f"over {ev['cases']} cases")

        subs = split_stand_in(QUESTION)
        answers: dict[int, str] = {}

        async def work_on(i: int) -> None:
            answers[i] = (await Host({"gateway": gw}, chat_fn=KeywordAgent()).ask(subs[i])).text

        started = time.perf_counter()
        async with anyio.create_task_group() as tg:
            for i in range(len(subs)):
                tg.start_soon(work_on, i)
        report["orchestrator"] = [{"sub_question": subs[i], "answer": " ".join(answers[i].split())} for i in sorted(answers)]
        print(f"3. orchestrator: {len(subs)} workers finished in {(time.perf_counter() - started) * 1000:.0f} ms")
        for item in report["orchestrator"]:
            print(f"   - {item['sub_question']} => {item['answer'][:70]}")

    async with Client(build_agent_server()) as agent:
        outer = Host({"agent": agent}, chat_fn=KeywordAgent())
        result = await outer.ask("What do my notes say about spaced repetition before bed?")
        inner = json.loads(result.text)
        report["agent_as_server"] = {"outer_calls": [c.name for c in result.calls], "inner": inner}
        print(f"4. outer agent called {report['agent_as_server']['outer_calls']}; "
              f"inner agent used {inner['tools_used']}")
        print(f"   answer: {inner['answer']}")

    Path("m10_lab_report.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    shutil.rmtree(work)
    print("report written to m10_lab_report.json")


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