serve.py

"""MCP stdio server for the graphanything Skill.

Exposes the same Session core that the CLI does, but via MCP tools so any
MCP-aware agent (Claude Code, Cursor, Gemini CLI, …) can drive the
graph-anything loop.

Tool surface (all under the `graphanything_` prefix to coexist with
external MCP servers):

  graphanything_open_session(inputs, preset?, extractor?)        → session_id
  graphanything_list_presets()                                    → presets
  graphanything_list_extractors()                                 → extractors
  graphanything_propose_schema(session_id, n=3, llm=False)
  graphanything_refine_schema(session_id, instruction, llm=False)
  graphanything_sample(session_id, n=5, extractor?)
  graphanything_review(session_id, actions[])
  graphanything_run(session_id, out_dir?)
  graphanything_status(session_id)
  graphanything_explain(session_id, target)
  graphanything_render(session_id, fmt='mermaid', max_nodes=60)

All tool results are JSON strings. The viz output is returned as a string
the model can paste into chat (mermaid / ASCII) — for HTML/SVG we return
a path the user can open.

Long-running tools (sample / run on large corpora) honour an optional
`max_seconds` arg by polling the work in chunks; the partial state lives
on the Session, so a follow-up call just resumes from where we left off
(`session.inputs[len(session.accepted['nodes_seen_files']):]` etc).
Phase 3 (#11) wires this into versioned snapshots; for now we expose
`status` so the model can poll progress.
"""
from __future__ import annotations

import json
import sys
from dataclasses import asdict
from pathlib import Path
from typing import Any

from . import (
    list_extractors,
    list_presets,
    load_session,
    open_session,
    render,
)
from .session import DEFAULT_SESSIONS_DIR, Session


# ---------------------------------------------------------------------------
# Helpers shared with cli.py — kept private to avoid coupling
# ---------------------------------------------------------------------------


def _make_llm():
    """Same LLM client as cli._make_llm; isolated copy so this module
    has no dependency on cli.py.

    Uses GraphAnything's built-in OpenAI-compatible client; configure
    via env vars (GA_API_BASE / GA_MODEL / GA_API_KEY — see llm_client.py).
    """
    from .llm_client import make_client

    return make_client()


def _maybe_llm(needs_llm: bool):
    if not needs_llm:
        return None
    try:
        return _make_llm()
    except Exception:
        return None


def _extractor_needs_llm(name: str | None) -> bool:
    if not name:
        return False
    try:
        from .extractors import get_extractor
        return get_extractor(name).needs_llm
    except Exception:
        return False


def _resolve(session_id: str, sessions_dir: Path) -> Session:
    return load_session(session_id, sessions_dir=sessions_dir)


def _summarise_session(sess: Session) -> dict:
    return {
        "session_id": sess.id,
        "status": sess.status,
        "inputs": len(sess.inputs),
        "schema": {
            "name": sess.schema.name,
            "version": sess.schema.version,
            "entities": [e.name for e in sess.schema.entities],
            "relations": [r.name for r in sess.schema.relations],
        },
        "accepted": {
            "nodes": len(sess.accepted.get("nodes", [])),
            "edges": len(sess.accepted.get("edges", [])),
        },
        "pending": {
            "nodes": len(sess.pending.get("nodes", [])),
            "edges": len(sess.pending.get("edges", [])),
        },
        "rejected": {"nodes": len(sess.rejected.get("nodes", []))},
        "cost": asdict(sess.cost),
        "extractor": sess.extractor,
    }


# ---------------------------------------------------------------------------
# Tool implementations (pure Python — same callable from tests + MCP)
# ---------------------------------------------------------------------------


def tool_open_session(
    inputs: list[str],
    *,
    preset: str | None = None,
    extractor: str | None = None,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
    budget: dict[str, Any] | None = None,
) -> dict:
    sd = Path(sessions_dir)
    sess = open_session(
        inputs, preset=preset, extractor=extractor,
        sessions_dir=sd, budget=budget,
    )
    return _summarise_session(sess)


def tool_list_presets() -> list[dict]:
    return list_presets()


def tool_list_extractors() -> list[dict]:
    return [
        {"name": m.name, "version": m.version, "handles": m.handles,
         "description": m.description, "needs_llm": m.needs_llm}
        for m in list_extractors()
    ]


def tool_propose_schema(
    session_id: str,
    *,
    n: int = 3,
    llm: bool = False,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    sess = _resolve(session_id, Path(sessions_dir))
    sess.propose(n_samples=n, auto_accept=True, llm=_maybe_llm(llm))
    out = _summarise_session(sess)
    out["schema"] = sess.schema.to_dict()       # full dict, not the summary
    return out


def tool_refine_schema(
    session_id: str,
    instruction: str,
    *,
    llm: bool = False,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    sess = _resolve(session_id, Path(sessions_dir))
    sess.refine(instruction, llm=_maybe_llm(llm))
    out = _summarise_session(sess)
    out["schema"] = sess.schema.to_dict()       # full dict, not the summary
    return out


def tool_sample(
    session_id: str,
    *,
    n: int = 5,
    extractor: str | None = None,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    sess = _resolve(session_id, Path(sessions_dir))
    llm = _maybe_llm(_extractor_needs_llm(extractor or sess.extractor))
    summary = sess.sample(n=n, extractor=extractor, llm=llm)
    out = _summarise_session(sess)
    out["last_sample"] = summary
    out["pending_preview"] = _preview(sess.pending, limit=12)
    return out


def tool_review(
    session_id: str,
    actions: list[dict],
    *,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    sess = _resolve(session_id, Path(sessions_dir))
    sess.review(actions)
    return _summarise_session(sess)


def tool_run(
    session_id: str,
    *,
    out_dir: str | None = None,
    extractor: str | None = None,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    sess = _resolve(session_id, Path(sessions_dir))
    llm = _maybe_llm(_extractor_needs_llm(extractor or sess.extractor))
    result = sess.run(extractor=extractor, llm=llm, out_dir=out_dir)
    return {**result, **_summarise_session(sess)}


def tool_status(
    session_id: str,
    *,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    sess = _resolve(session_id, Path(sessions_dir))
    return _summarise_session(sess)


def tool_explain(
    session_id: str,
    target: str,
    *,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    sess = _resolve(session_id, Path(sessions_dir))
    G = sess.to_viz_graph()
    if target in G.nodes:
        node_data = G.nodes[target]
        return {
            "id": target,
            **node_data,
            "in_edges": [
                {"from": u, **G[u][target]} for u, _ in list(G.in_edges(target))[:10]
            ],
            "out_edges": [
                {"to": v, **G[target][v]} for _, v in list(G.out_edges(target))[:10]
            ],
        }
    cand = [
        n for n in G.nodes
        if str(G.nodes[n].get("label", "")).lower() == target.lower()
    ]
    if cand:
        return tool_explain(session_id, cand[0], sessions_dir=sessions_dir)
    return {"error": f"no node {target!r} in session {session_id}"}


def tool_ask(
    session_id: str,
    question: str,
    *,
    llm: bool = False,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    from .ask import ask
    sess = _resolve(session_id, Path(sessions_dir))
    G = sess.to_viz_graph()
    return ask(
        G, question,
        llm=_maybe_llm(llm),
        schema_entities=[e.name for e in sess.schema.entities],
        schema_relations=[r.name for r in sess.schema.relations],
    )


def tool_update(
    session_id: str,
    *,
    extractor: str | None = None,
    out_dir: str | None = None,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    sess = _resolve(session_id, Path(sessions_dir))
    llm = _maybe_llm(_extractor_needs_llm(extractor or sess.extractor))
    result = sess.update(extractor=extractor, llm=llm, out_dir=out_dir)
    return {**result, **_summarise_session(sess)}


def tool_versions(
    *,
    out_root: str | None = None,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> list[dict]:
    from .temporal import list_versions
    root = Path(out_root) if out_root else Path(sessions_dir).parent
    return [asdict(v) for v in list_versions(root)]


def tool_diff(
    v_old: int,
    v_new: int,
    *,
    out_root: str | None = None,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    from .temporal import diff_graphs, load_version
    root = Path(out_root) if out_root else Path(sessions_dir).parent
    old = load_version(root, v_old)
    new = load_version(root, v_new)
    out = diff_graphs(old, new)
    # Compact node/edge listings so the model isn't drowned.
    return {
        "summary": out["summary"],
        "nodes_added":   [n["id"] for n in out["nodes"]["added"]][:50],
        "nodes_removed": [n["id"] for n in out["nodes"]["removed"]][:50],
        "nodes_modified": [m["id"] for m in out["nodes"]["modified"]][:50],
        "edges_added":   [[e["source"], e["target"], e.get("relation")] for e in out["edges"]["added"]][:50],
        "edges_removed": [[e["source"], e["target"], e.get("relation")] for e in out["edges"]["removed"]][:50],
        "n_nodes_modified": len(out["nodes"]["modified"]),
        "n_edges_modified": len(out["edges"]["modified"]),
    }


def tool_federate(
    graphs: list[str],
    *,
    out: str,
    fuzzy: bool = False,
    fuzzy_threshold: float = 0.7,
    llm: bool = False,
) -> dict:
    from .federate import federate_files
    result = federate_files(
        graphs, fuzzy=fuzzy, fuzzy_threshold=fuzzy_threshold,
        llm_tiebreak=_maybe_llm(llm) if llm else None,
        out_path=out,
    )
    return {"out_path": out, **result["linking_report"]}


def tool_eval(
    session_id: str,
    *,
    out_dir: str = "graphanything-out/quality",
    llm: bool = False,
    judge_n: int = 20,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    from .quality import write_report
    sess = _resolve(session_id, Path(sessions_dir))
    paths = write_report(
        sess.accepted, out_dir=out_dir, schema=sess.schema,
        llm=_maybe_llm(llm) if llm else None, judge_n=judge_n,
    )
    return {k: str(v) for k, v in paths.items()}


def tool_render(
    session_id: str,
    *,
    fmt: str = "mermaid",
    max_nodes: int = 60,
    out_path: str | None = None,
    budget_tokens: int | None = None,
    sessions_dir: str | Path = DEFAULT_SESSIONS_DIR,
) -> dict:
    sess = _resolve(session_id, Path(sessions_dir))
    G = sess.to_viz_graph()
    fmt_opts: dict = {}
    if fmt == "mermaid":
        fmt_opts["max_nodes"] = max_nodes
    out = render(G, fmt=fmt, out_path=out_path, budget_tokens=budget_tokens, **fmt_opts)
    if isinstance(out, Path):
        return {"fmt": fmt, "path": str(out), "n_nodes": G.number_of_nodes()}
    return {"fmt": fmt, "content": out, "n_nodes": G.number_of_nodes()}


def _preview(graph: dict, *, limit: int = 12) -> dict:
    """Compact pending preview for tool results — first N nodes/edges."""
    return {
        "nodes": [
            {"id": n["id"], "label": n.get("label"), "type": n.get("type"),
             "evidence_span": (n.get("evidence_span") or "")[:120]}
            for n in graph.get("nodes", [])[:limit]
        ],
        "edges": [
            {"src": e["source"], "tgt": e["target"], "rel": e.get("relation"),
             "rationale": (e.get("rationale") or "")[:120]}
            for e in graph.get("edges", [])[:limit]
        ],
        "total_nodes": len(graph.get("nodes", [])),
        "total_edges": len(graph.get("edges", [])),
    }


# ---------------------------------------------------------------------------
# MCP server wiring (only imported when serve() is called)
# ---------------------------------------------------------------------------


_TOOL_SCHEMA: dict[str, dict[str, Any]] = {
    "graphanything_open_session": {
        "description": "Start a new graphanything session over one or more inputs (paths or globs). "
                       "Returns session_id and an initial schema (from preset, if any). "
                       "Always call this first when the user asks to graph anything.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "inputs": {"type": "array", "items": {"type": "string"},
                           "description": "Files or directories to graph."},
                "preset": {"type": "string",
                           "description": "Schema preset name (use list_presets to discover)."},
                "extractor": {"type": "string",
                              "description": "Force a specific extractor (else dispatched by suffix)."},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["inputs"],
        },
    },
    "graphanything_list_presets": {
        "description": "List built-in schema presets (contracts, openapi, papers, …). "
                       "Use to pick `preset=` for graphanything_open_session.",
        "inputSchema": {"type": "object", "properties": {}},
    },
    "graphanything_list_extractors": {
        "description": "List available extractors with their handled suffixes and whether they need an LLM.",
        "inputSchema": {"type": "object", "properties": {}},
    },
    "graphanything_propose_schema": {
        "description": "Suggest an initial schema for the session by inspecting the first N inputs. "
                       "Set llm=true to use the OpenAI-compatible client; otherwise rule-based / generic.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "n": {"type": "integer", "default": 3},
                "llm": {"type": "boolean", "default": False},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["session_id"],
        },
    },
    "graphanything_refine_schema": {
        "description": "Apply a free-form schema edit to the session "
                       "(e.g. 'add Amount entity', 'rename Foo to Bar', 'add relation knows from Person to Person').",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "instruction": {"type": "string"},
                "llm": {"type": "boolean", "default": False},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["session_id", "instruction"],
        },
    },
    "graphanything_sample": {
        "description": "Run the current schema over the first N inputs and stage results in `pending`. "
                       "Returns a compact preview of nodes/edges plus suggested merges.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "n": {"type": "integer", "default": 5},
                "extractor": {"type": "string"},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["session_id"],
        },
    },
    "graphanything_review": {
        "description": "Apply review actions to the session's pending nodes/edges. "
                       "Each action is a dict; supported ops: accept_all, accept, reject, merge, rule.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "actions": {"type": "array", "items": {"type": "object"}},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["session_id", "actions"],
        },
    },
    "graphanything_run": {
        "description": "Lock the schema, run all inputs, normalise, and write graph.json. "
                       "Use this when the user is happy with sample + review.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "out_dir": {"type": "string"},
                "extractor": {"type": "string"},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["session_id"],
        },
    },
    "graphanything_status": {
        "description": "Compact status dump for the session (counts, schema, cost). "
                       "Cheap; safe to call between long-running tools to track progress.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["session_id"],
        },
    },
    "graphanything_explain": {
        "description": "Show full provenance for a node id or label, including in/out edges.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "target": {"type": "string"},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["session_id", "target"],
        },
    },
    "graphanything_ask": {
        "description": "Run a natural-language query over the session graph. "
                       "Regex-translates simple queries (type=X, top N, neighbours of 'Y', amount > N) "
                       "and falls back to LLM translation otherwise. Read-only.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "question": {"type": "string"},
                "llm": {"type": "boolean", "default": False},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["session_id", "question"],
        },
    },
    "graphanything_update": {
        "description": "Re-extract only inputs whose source_hash changed since the last run. "
                       "Snapshots the new graph as the next version (use graphanything_diff to compare).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "extractor": {"type": "string"},
                "out_dir": {"type": "string"},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["session_id"],
        },
    },
    "graphanything_versions": {
        "description": "List graph snapshots written by previous run/update calls.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "out_root": {"type": "string"},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
        },
    },
    "graphanything_diff": {
        "description": "Compute added/removed/modified nodes and edges between two graph versions.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "v_old": {"type": "integer"},
                "v_new": {"type": "integer"},
                "out_root": {"type": "string"},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["v_old", "v_new"],
        },
    },
    "graphanything_federate": {
        "description": "Combine multiple graph.json files into one universe graph. "
                       "Same-label entities of the same type are merged; with fuzzy=true "
                       "additionally proposes 'same_as' edges via Jaccard token overlap.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "graphs": {"type": "array", "items": {"type": "string"}},
                "out": {"type": "string"},
                "fuzzy": {"type": "boolean", "default": False},
                "fuzzy_threshold": {"type": "number", "default": 0.7},
                "llm": {"type": "boolean", "default": False},
            },
            "required": ["graphs", "out"],
        },
    },
    "graphanything_eval": {
        "description": "Write a QUALITY_REPORT.md (coverage / dedup / per-extractor stats) for the session graph. "
                       "Pass llm=true to also have the LLM judge a sample of edges against their evidence_span.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "out_dir": {"type": "string", "default": "graphanything-out/quality"},
                "llm": {"type": "boolean", "default": False},
                "judge_n": {"type": "integer", "default": 20},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["session_id"],
        },
    },
    "graphanything_render": {
        "description": "Render the session graph. Default fmt='mermaid' returns inline content "
                       "Claude can paste into chat; fmt='html'/'svg' return a file path the user can open. "
                       "Pass `budget_tokens=N` to PageRank-prune big graphs down to roughly N tokens.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "session_id": {"type": "string"},
                "fmt": {"type": "string", "enum": ["mermaid", "ascii", "json", "html", "svg", "canvas", "timeline", "cypher", "graphml"], "default": "mermaid"},
                "max_nodes": {"type": "integer", "default": 60},
                "budget_tokens": {"type": "integer"},
                "out_path": {"type": "string"},
                "sessions_dir": {"type": "string", "default": str(DEFAULT_SESSIONS_DIR)},
            },
            "required": ["session_id"],
        },
    },
}


_TOOL_FUNCS = {
    "graphanything_open_session": tool_open_session,
    "graphanything_list_presets": tool_list_presets,
    "graphanything_list_extractors": tool_list_extractors,
    "graphanything_propose_schema": tool_propose_schema,
    "graphanything_refine_schema": tool_refine_schema,
    "graphanything_sample": tool_sample,
    "graphanything_review": tool_review,
    "graphanything_run": tool_run,
    "graphanything_status": tool_status,
    "graphanything_explain": tool_explain,
    "graphanything_ask": tool_ask,
    "graphanything_update": tool_update,
    "graphanything_versions": tool_versions,
    "graphanything_diff": tool_diff,
    "graphanything_federate": tool_federate,
    "graphanything_eval": tool_eval,
    "graphanything_render": tool_render,
}


def call_tool(name: str, args: dict | None) -> str:
    """Dispatch a tool by name. Returns a JSON string for MCP `text` content."""
    args = args or {}
    fn = _TOOL_FUNCS.get(name)
    if fn is None:
        return json.dumps({"error": f"unknown tool {name!r}"})
    try:
        result = fn(**args)
    except TypeError as e:
        return json.dumps({"error": f"bad arguments to {name}: {e}"})
    except Exception as e:
        return json.dumps({"error": f"{type(e).__name__}: {e}"})
    return json.dumps(result, default=str, ensure_ascii=False)


def serve() -> None:
    """Start the MCP stdio server. Requires `pip install mcp`."""
    try:
        from mcp.server import Server
        from mcp.server.stdio import stdio_server
        from mcp import types
    except ImportError as e:
        raise ImportError("mcp not installed. Run: pip install mcp") from e

    server = Server("graphanything")

    @server.list_tools()
    async def list_tools() -> list[Any]:
        out = []
        for name, spec in _TOOL_SCHEMA.items():
            out.append(types.Tool(
                name=name,
                description=spec["description"],
                inputSchema=spec["inputSchema"],
            ))
        return out

    @server.call_tool()
    async def call(name: str, arguments: dict | None) -> list[Any]:
        text = call_tool(name, arguments)
        return [types.TextContent(type="text", text=text)]

    import asyncio

    async def _run():
        async with stdio_server() as (read, write):
            from mcp.server.models import InitializationOptions
            from mcp.server import NotificationOptions
            await server.run(
                read, write,
                InitializationOptions(
                    server_name="graphanything",
                    server_version="0.1.0",
                    capabilities=server.get_capabilities(
                        notification_options=NotificationOptions(),
                        experimental_capabilities={},
                    ),
                ),
            )

    asyncio.run(_run())


def main(argv: list[str] | None = None) -> int:
    """Allow `python -m GraphAnything.serve` to start the server."""
    serve()
    return 0


if __name__ == "__main__":  # pragma: no cover
    sys.exit(main())