Topic 5: Module 4 lab

"""Module 4 lab: resources, a slash-command prompt, two approval gates, and live updates in one run."""
import json
import logging
import shutil
import sys
import tempfile
import time
from pathlib import Path
from typing import Any

import anyio

from mcp import Client
from mcp.client import ClientRequestContext
from mcp.client.subscriptions import ResourcesListChanged, ResourceUpdated, Subscription
from mcp.types import ElicitRequestParams, ElicitResult

from examples.m04_interaction_server import build_interactive_server
from examples.m04_slash import parse
from notes_assistant.host import Host
from notes_assistant.llm import ChatReply, ToolCall
from notes_assistant.store import NoteStore

started = time.perf_counter()
stats = {"elicitations": 0, "host_approvals": 0, "events": 0}


def log(message: str) -> None:
    print(f"[{(time.perf_counter() - started) * 1000:5.0f} ms] {message}")


def scripted_model(messages: list[dict[str, Any]], tools: list[dict[str, Any]]) -> ChatReply:
    """Scripted stand-in for an LLM: search, read the top hit, then save a summary note."""
    last = messages[-1]
    if last["role"] == "user":
        topic = last["content"].split("'")[1]
        return ChatReply(None, [ToolCall("s", "notes__search_notes", {"query": topic, "limit": 1})])
    if last.get("tool_call_id") == "s":
        uri = json.loads(last["content"])["hits"][0]["uri"]
        return ChatReply(None, [ToolCall("r", "read_resource", {"uri": uri})])
    if last.get("tool_call_id") == "r":
        fact = last["content"].split("\n\n", 2)[2].splitlines()[0]
        return ChatReply(None, [ToolCall("w", "notes__create_note", {
            "title": "Sleep summary", "body": f"- {fact} (sleep-and-memory)", "tags": ["sleep", "summary"]})])
    return ChatReply("Saved your summary as notes://sleep-summary.")


def approve(name: str, arguments: dict[str, Any]) -> bool:
    stats["host_approvals"] += 1
    log(f"host gate: allow {name} title={arguments['title']!r}")
    return True


async def on_elicit(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult:
    stats["elicitations"] += 1
    log(f"server gate ({params.mode}): {params.message}")
    return ElicitResult(action="accept", content={"ok": True})


async def watch(client: Client, sub: Subscription) -> None:
    async for event in sub:
        stats["events"] += 1
        match event:
            case ResourcesListChanged():
                log("event: resources list changed")
            case ResourceUpdated(uri=uri):
                index = (await client.read_resource(uri)).contents[0].text
                log(f"event: {uri} updated, now {index.count(chr(10) + '- ')} notes")


async def main() -> None:
    logging.basicConfig(level="WARNING", stream=sys.stderr)
    folder = Path(tempfile.mkdtemp()) / "notes"
    shutil.copytree("notes", folder)
    server = build_interactive_server(NoteStore(folder))

    async with Client(server, elicitation_callback=on_elicit) as client:
        resources = (await client.list_resources()).resources
        templates = (await client.list_resource_templates()).resource_templates
        prompts = {p.name: p for p in (await client.list_prompts()).prompts}
        tools = (await client.list_tools()).tools
        log(f"catalog: tools={[t.name for t in tools]} resources={[str(r.uri) for r in resources]} "
            f"templates={[t.uri_template for t in templates]} prompts={list(prompts)}")

        async with client.listen(resources_list_changed=True, resource_subscriptions=["notes://index"]) as sub, \
                anyio.create_task_group() as tg:
            tg.start_soon(watch, client, sub)

            line = "/summarise_topic sleep"
            log(f"user types: {line}")
            prompt, arguments = parse(line, prompts)
            text = (await client.get_prompt(prompt.name, arguments)).messages[0].content.text
            host = Host({"notes": client}, chat_fn=scripted_model, approve=approve)
            answer = await host.ask(text)
            for call in answer.calls:
                log(f"tool call {call.name} ok={call.ok}")
            log(f"answer: {answer.text}")
            await anyio.sleep(0.1)  # let the watcher print its last event
            tg.cancel_scope.cancel()

        note = (await client.read_resource("notes://sleep-summary")).contents[0].text
        log("notes://sleep-summary reads: " + " | ".join(line for line in note.splitlines() if line))
    log(f"stats: {stats}")


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