server/mcp_tools.py

"""Agent-facing tool wrappers exposed via Gradio's MCP server.

Each function below has a clean signature + docstring on purpose — Gradio's MCP
layer (`mcp_server=True` in app.py) reads the type hints and docstring to build
the JSON-Schema a remote MCP client sees. Keep them stateless and JSON-friendly:
inputs are str / list[dict] / etc., outputs are dict / str / list[dict] (never
pydantic objects, which don't serialise through the MCP boundary).

These wrap the existing pipeline (server/pipeline.run_pipeline) and free/busy
math (calendar_out/freebusy) — no new business logic lives here, just the
shape adaptation an external agent expects.
"""
from __future__ import annotations

import base64
import time
from collections import OrderedDict
from typing import Optional

from calendar_out.freebusy import Busy, check_conflicts as _freebusy_check_conflicts, load_ics_busy
from calendar_out.ics import events_to_ics
from server.pipeline import AgentRequest, run_pipeline
from server.schema import Event

# Short-lived extraction cache. The Agent-tab orchestrator extracts TWICE per
# run — once when the MiniCPM planner calls this tool over MCP, then again when
# the scripted path finalizes — and each call runs the full gemma-cal E4B. With
# identical inputs the second call is pure waste, so memoize on the EXACT inputs
# (thread + images + memory). Different memory/images -> different key -> a fresh
# (correct) extraction; the win is the common no-memory case. TTL is generous so
# the scripted call still hits after a ~2-min planner run; small maxsize bounds
# cross-request staleness (same input -> same output anyway).
_EXTRACT_CACHE: "OrderedDict[tuple, tuple[float, dict]]" = OrderedDict()
_EXTRACT_TTL = 600.0
_EXTRACT_MAX = 8


def extract_events(thread: str, images: Optional[list[str]] = None,
                   memory: Optional[str] = None) -> dict:
    """Extract calendar events from a pasted iMessage thread (and optional screenshots).

    The headline tool. Reads a chat or screenshot, returns an ActionPlan with the
    events found, any conflicts against the user's calendar, and a suggested reply.
    Runs 100% locally inside the Space via llama.cpp — no cloud AI APIs.

    Args:
        thread: Plain-text iMessage conversation, e.g. "Alice: pickup 5pm Thursday".
            Either ``thread`` or ``images`` must be non-empty.
        images: Optional list of base64-encoded screenshots (raw base64 or data URIs).
            Useful when the schedule lives in a screenshot rather than text.
        memory: Optional plain-text recall block about the user (people and their
            roles, preferences like default reminders or days they decline) — used
            to personalize extraction. e.g. "Dana is the soccer coach".

    Returns:
        ActionPlan as a JSON-serialisable dict with keys: ``reasoning``,
        ``events`` (list of {title, start, end, location, attendees, ...}),
        ``conflicts``, ``proposed_times``, ``reply_draft``, ``needs_clarification``.
    """
    key = (thread or "", tuple(images or []), memory or "")
    now = time.monotonic()
    hit = _EXTRACT_CACHE.get(key)
    if hit is not None and now - hit[0] < _EXTRACT_TTL:
        _EXTRACT_CACHE.move_to_end(key)
        return hit[1]
    req = AgentRequest(thread=thread or "", images=images or [], memory=memory,
                       return_ics=False)
    resp = run_pipeline(req)
    plan = resp.plan.model_dump()
    _EXTRACT_CACHE[key] = (now, plan)
    _EXTRACT_CACHE.move_to_end(key)
    while len(_EXTRACT_CACHE) > _EXTRACT_MAX:
        _EXTRACT_CACHE.popitem(last=False)
    return plan


def make_ics(events: list[dict]) -> str:
    """Render a list of event dicts as an .ics file (base64-encoded).

    Args:
        events: List of event dicts in the shape returned by ``extract_events``
            — each needs at least ``title`` and ``start`` (ISO 8601). Optional:
            ``end``, ``location``, ``attendees``, ``reminder_minutes``, ``notes``.

    Returns:
        Base64-encoded VCALENDAR bytes. Decode and write to ``something.ics`` to
        import into any calendar app.
    """
    ev_objs = [Event(**e) for e in events]
    return base64.b64encode(events_to_ics(ev_objs)).decode("ascii")


def check_conflicts(events: list[dict], ics_base64: str) -> list[dict]:
    """Find clashes between proposed events and busy intervals from an .ics calendar.

    Deterministic free/busy math — runs without the LLM, so it's safe for agents
    to call as a fast verification step after ``extract_events``.

    Args:
        events: List of proposed event dicts (same shape as ``extract_events``
            output). Each event needs at least ``title`` and ``start``.
        ics_base64: Base64-encoded .ics calendar to check against. Typically the
            user's current calendar exported from Google/Apple/Outlook.

    Returns:
        List of conflict dicts: ``{event_index, clashes_with, severity}`` where
        severity is one of ``"overlap"``, ``"adjacent"``, ``"tight"``. Empty list
        if nothing clashes.
    """
    if not ics_base64:
        return []
    try:
        busy: list[Busy] = load_ics_busy(base64.b64decode(ics_base64))
    except Exception:  # noqa: BLE001  malformed .ics -> no conflict context
        return []
    ev_objs = [Event(**e) for e in events]
    return [c.model_dump() for c in _freebusy_check_conflicts(ev_objs, busy)]