Topic 7: Module 1 lab

"""Module 1 lab: one script that exercises everything from the module.

1. Loads the notes with NoteStore.
2. Prints the N x M integration arithmetic.
3. Compares the hand-written native tool schema with the one MCP generates.
4. Reaches the same server in memory, over stdio, and over Streamable HTTP,
   checks all three return the same hits, and times them.
5. Runs one scripted-model turn through MCP instead of a local dispatcher.

Run from the repository root:  PYTHONPATH=. python examples/m01_lab.py
"""
import json
import logging
import statistics
import subprocess
import sys
import time

import anyio
import httpx2

from mcp import Client, StdioServerParameters

sys.path.insert(0, "examples")
from m01_first_server import mcp  # noqa: E402
from m01_native_tools import SEARCH_TOOL, scripted_model  # noqa: E402
from notes_assistant.llm import ChatReply  # noqa: E402
from notes_assistant.store import NoteStore  # noqa: E402

logging.getLogger().setLevel(logging.WARNING)  # keep per-request HTTP logs out of the lab output

HTTP_PORT = 8011
HTTP_URL = f"http://127.0.0.1:{HTTP_PORT}/mcp"
CALLS = 50
QUERY = {"query": "nap study", "limit": 2}


def step1_store() -> None:
    store = NoteStore("notes")
    print(f"1. NoteStore: {len(store.list_notes())} notes, top hit for 'nap study' is {store.search('nap study')[0].note_id}")


def step2_arithmetic() -> None:
    print("2. Integrations needed (hosts x tools versus hosts + tools):")
    for hosts, tools in [(4, 6), (10, 50), (30, 400)]:
        print(f"   {hosts:>3} hosts, {tools:>3} tools: {hosts * tools:>6} bespoke adapters versus {hosts + tools:>4} MCP implementations")


def to_chat_tool(tool) -> dict:
    """Minimal MCP-to-Chat-Completions conversion (Module 6 builds the real one)."""
    return {"type": "function", "function": {"name": tool.name, "description": tool.description, "parameters": tool.input_schema}}


async def step3_schema() -> None:
    async with Client(mcp) as client:
        generated = to_chat_tool((await client.list_tools()).tools[0])
    for label, schema in [("hand-written native", SEARCH_TOOL), ("generated by MCP", generated)]:
        chars = len(json.dumps(schema))
        print(f"3. {label:<20} {chars} characters, about {chars // 4} tokens (estimate: characters / 4)")


async def time_target(label: str, target: object) -> list[str]:
    start = time.perf_counter()
    async with Client(target) as client:
        connect_ms = 1000 * (time.perf_counter() - start)
        first = await client.call_tool("search_notes", QUERY)  # warm-up call, not timed
        samples = []
        for _ in range(CALLS):
            t0 = time.perf_counter()
            await client.call_tool("search_notes", QUERY)
            samples.append(1000 * (time.perf_counter() - t0))
    ids = [hit["note_id"] for hit in first.structured_content["result"]]
    print(f"   {label:<7} connect {connect_ms:7.1f} ms   median call {statistics.median(samples):5.2f} ms   hits {ids}")
    return ids


async def wait_for_http() -> None:
    for _ in range(50):
        try:
            async with httpx2.AsyncClient() as http:
                await http.get(HTTP_URL)
            return
        except httpx2.TransportError:
            await anyio.sleep(0.1)
    raise RuntimeError("HTTP server did not start")


async def step4_transports() -> None:
    print(f"4. Same server, three transports ({CALLS} timed calls each; timings vary run to run):")
    stdio = StdioServerParameters(command=sys.executable, args=["examples/m01_first_server.py"], env={"PYTHONPATH": "."})
    server = subprocess.Popen(
        [sys.executable, "examples/m01_first_server.py", "--http", str(HTTP_PORT)],
        env={"PYTHONPATH": "."}, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )
    try:
        await wait_for_http()
        results = [await time_target("memory", mcp), await time_target("stdio", stdio), await time_target("http", HTTP_URL)]
    finally:
        server.terminate()
        server.wait()
    print(f"   all three transports agree: {results[0] == results[1] == results[2]}")


async def step5_scripted_turn() -> None:
    async with Client(mcp) as client:
        tools = [to_chat_tool(t) for t in (await client.list_tools()).tools]
        messages = [{"role": "user", "content": "How many participants will the nap study have?"}]
        reply: ChatReply = scripted_model(messages, tools)
        call = reply.tool_calls[0]
        result = await client.call_tool(call.name, call.arguments)
        messages += [reply.as_message(), {"role": "tool", "tool_call_id": call.id, "content": json.dumps(result.structured_content["result"])}]
        final = scripted_model(messages, tools)
    print(f"5. Scripted stand-in model through MCP: {final.content}")


async def main() -> None:
    step1_store()
    step2_arithmetic()
    await step3_schema()
    await step4_transports()
    await step5_scripted_turn()


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