knowledge-base-framework/kb_framework/query_tools.py

"""The read-only MCP tool contract: definitions and dispatch.

Both front ends (the stdio server and the hosted connector) advertise exactly
these five tools with exactly these schemas, and `call` produces the result
dict (`content` + `structuredContent`) each returns. Keeping the contract here
means an agent skill written against one front end stays valid on the other.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, List

from kb_framework import instructions, query, vault

READ_ONLY_TOOL_NAMES = (
    "read_instructions",
    "graph_summary",
    "graph_search_nodes",
    "graph_neighbors",
    "read_note",
)


def definitions() -> List[Dict[str, Any]]:
    return [
        {
            "name": "read_instructions",
            "description": (
                "Read framework AGENTS.md and AGENT_QUERY_GUIDE.md before working "
                "with the vault. These operating documents live outside an "
                "external vault and do not require a generated graph."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {},
                "additionalProperties": False,
            },
        },
        {
            "name": "graph_summary",
            "description": (
                "Orientation for a cold start: node and edge counts, the repo "
                "scopes, the most-linked notes, and the tag vocabulary as "
                "`by_tag` (most frequent first) plus `tag_count`, the number of "
                "distinct tags. Read `by_tag` before filtering "
                "`graph_search_nodes` by `tag` — that filter matches a tag "
                "verbatim, so a guessed spelling returns nothing."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "topLinked": {
                        "type": "integer",
                        "minimum": 1,
                        "maximum": 50,
                        "description": "Most-linked notes to return (default 10).",
                    },
                    "topTags": {
                        "type": "integer",
                        "minimum": 1,
                        "maximum": 500,
                        "description": "Tags to return, most frequent first (default 50).",
                    },
                },
                "additionalProperties": False,
            },
        },
        {
            "name": "graph_search_nodes",
            "description": "Search graph nodes by metadata and free text.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "repo": {"type": "string"},
                    "type": {"type": "string"},
                    "status": {"type": "string"},
                    "tag": {"type": "string"},
                    "query": {"type": "string"},
                    "limit": {"type": "integer", "minimum": 1, "maximum": 200},
                },
            },
        },
        {
            "name": "graph_neighbors",
            "description": (
                "Show inbound and outbound links for a node ID. Each neighbor "
                "includes the edge `type` (relationship, e.g. supersedes, "
                "depends-on), `anchor_text` (the link's display text), and "
                "`annotation` (what to expect in that note)."
            ),
            "inputSchema": {
                "type": "object",
                "required": ["id"],
                "properties": {
                    "id": {"type": "string"},
                    "limit": {"type": "integer", "minimum": 1, "maximum": 200},
                },
            },
        },
        {
            "name": "read_note",
            "description": (
                "Read note content by graph node ID or vault-relative path. "
                "Provide exactly one of `id` or `path`. "
                "If both or neither are provided, the call returns an error."
            ),
            "inputSchema": {
                "type": "object",
                "properties": {
                    "id": {"type": "string", "description": "Graph node ID (e.g. 'AGENTS')."},
                    "path": {"type": "string", "description": "Vault-relative file path (e.g. 'README.md')."},
                },
                "additionalProperties": False,
            },
        },
    ]


def text_result(payload: Dict[str, Any]) -> Dict[str, Any]:
    """Wrap a payload as an MCP tool result: pretty JSON text plus structuredContent."""
    return {
        "content": [{"type": "text", "text": json.dumps(payload, indent=2)}],
        "structuredContent": payload,
    }


def read_note_path(graph: Dict[str, Any], args: Dict[str, Any]) -> str:
    """Resolve `read_note` arguments to the vault-relative path to read."""
    has_id = bool(args.get("id"))
    has_path = bool(args.get("path"))
    if has_id and has_path:
        raise ValueError("provide exactly one of 'id' or 'path', not both")
    if not has_id and not has_path:
        raise ValueError("exactly one of 'id' or 'path' is required")
    if has_id:
        return query.node_path(graph, str(args["id"]))
    return str(args["path"])


def call(
    name: str, args: Dict[str, Any], *, graph: Dict[str, Any], root: Path
) -> Dict[str, Any]:
    """Run one read-only tool. Raises ValueError for unknown tools or bad input."""
    if name == "read_instructions":
        return text_result(instructions.read_instructions())
    if name == "graph_summary":
        return text_result(
            query.summary(
                graph,
                topLinked=args.get("topLinked", query.DEFAULT_TOP_LINKED),
                topTags=args.get("topTags", query.DEFAULT_TOP_TAGS),
            )
        )
    if name == "graph_search_nodes":
        return text_result(query.search_nodes(graph, args))
    if name == "graph_neighbors":
        return text_result(
            query.neighbors(graph, str(args["id"]), args.get("limit", query.DEFAULT_NEIGHBOR_LIMIT))
        )
    if name == "read_note":
        return text_result(vault.read_note_payload(root, read_note_path(graph, args)))
    raise ValueError(f"unknown tool '{name}'")