knowledge-base-framework/kb_framework/query.py

"""Pure query functions over an in-memory graph dict.

These back the read-only MCP tools (`graph_summary`, `graph_search_nodes`,
`graph_neighbors`) and the id-to-path lookup `read_note` needs. They take the
graph as data and never read the filesystem, so the stdio server and the
connector return identical payloads for identical graphs.
"""

from __future__ import annotations

from collections import Counter
from typing import Any, Dict, List

# graph_summary defaults: enough most-linked notes to spot the hubs, and enough
# tags to cover the vocabulary an agent actually filters on.
DEFAULT_TOP_LINKED = 10
DEFAULT_TOP_TAGS = 50
DEFAULT_SEARCH_LIMIT = 30
DEFAULT_NEIGHBOR_LIMIT = 50


def nodes_by_id(graph: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
    nodes = graph.get("nodes", [])
    return {str(node.get("id")): node for node in nodes if isinstance(node, dict)}


def match_filters(node: Dict[str, Any], args: Dict[str, Any]) -> bool:
    repo = args.get("repo")
    note_type = args.get("type")
    status = args.get("status")
    tag = args.get("tag")
    query = str(args.get("query", "")).strip().lower()

    if repo and node.get("repo") != repo:
        return False
    if note_type and node.get("type") != note_type:
        return False
    if status and node.get("status") != status:
        return False
    if tag and tag not in (node.get("tags") or []):
        return False
    if query:
        haystacks = [
            str(node.get("id", "")).lower(),
            str(node.get("title", "")).lower(),
            str(node.get("path", "")).lower(),
            " ".join(str(t).lower() for t in (node.get("tags") or [])),
        ]
        if not any(query in h for h in haystacks):
            return False
    return True


def summary(
    graph: Dict[str, Any],
    topLinked: int = DEFAULT_TOP_LINKED,
    topTags: int = DEFAULT_TOP_TAGS,
) -> Dict[str, Any]:
    nodes = graph.get("nodes", [])
    edges = graph.get("edges", [])
    top_n = int(topLinked)
    top_tags = int(topTags)

    by_repo: Counter[str] = Counter(str(node.get("repo", "unknown")) for node in nodes)
    tag_counts: Counter[str] = Counter(
        str(tag) for node in nodes for tag in (node.get("tags") or [])
    )
    # Frequent tags first, alphabetical within a tie, so a truncated list is
    # both the most useful slice and stable across calls.
    ranked_tags = sorted(tag_counts.items(), key=lambda item: (-item[1], item[0]))

    sorted_nodes = sorted(
        nodes,
        key=lambda n: int(n.get("inbound_count", 0)),
        reverse=True,
    )[:top_n]
    return {
        "version": graph.get("version"),
        "generated_at": graph.get("generated_at"),
        "node_count": len(nodes),
        "edge_count": len(edges),
        "by_repo": dict(by_repo),
        "tag_count": len(tag_counts),
        "by_tag": [
            {"tag": tag, "count": count} for tag, count in ranked_tags[:top_tags]
        ],
        "top_linked": [
            {
                "id": node.get("id"),
                "title": node.get("title"),
                "path": node.get("path"),
                "inbound_count": node.get("inbound_count", 0),
            }
            for node in sorted_nodes
        ],
    }


def search_nodes(graph: Dict[str, Any], args: Dict[str, Any]) -> Dict[str, Any]:
    nodes = graph.get("nodes", [])
    limit = int(args.get("limit", DEFAULT_SEARCH_LIMIT))
    filtered = [node for node in nodes if match_filters(node, args)]
    filtered.sort(key=lambda n: str(n.get("id", "")))
    sliced = filtered[:limit]
    return {
        "total_matches": len(filtered),
        "returned": len(sliced),
        "items": sliced,
    }


def neighbors(
    graph: Dict[str, Any], id: str, limit: int = DEFAULT_NEIGHBOR_LIMIT
) -> Dict[str, Any]:
    node_id = str(id)
    limit = int(limit)
    nodes = nodes_by_id(graph)
    if node_id not in nodes:
        raise ValueError(f"node '{node_id}' not found")

    edges = graph.get("edges", [])
    outbound_edges = [e for e in edges if e.get("from") == node_id]
    inbound_edges = [e for e in edges if e.get("to") == node_id]

    def describe(edge_list: List[Dict[str, Any]], direction: str) -> List[Dict[str, Any]]:
        # neighbor id is the "other end" relative to node_id
        other_key = "to" if direction == "outbound" else "from"
        result: List[Dict[str, Any]] = []
        for edge in edge_list:
            item_id = str(edge.get(other_key))
            node = nodes.get(item_id)
            result.append(
                {
                    "id": node.get("id") if node else item_id,
                    "title": node.get("title") if node else "<missing>",
                    "path": node.get("path") if node else "",
                    "direction": direction,
                    "type": edge.get("type", "related"),
                    "anchor_text": edge.get("anchor_text", ""),
                    "annotation": edge.get("annotation", ""),
                }
            )
        return result[:limit]

    return {
        "node": nodes[node_id],
        "inbound_count": len(inbound_edges),
        "outbound_count": len(outbound_edges),
        "inbound": describe(inbound_edges, "inbound"),
        "outbound": describe(outbound_edges, "outbound"),
    }


def node_path(graph: Dict[str, Any], id: str) -> str:
    """Vault-relative path of the note behind a node id."""
    node = nodes_by_id(graph).get(str(id))
    if not node:
        raise ValueError(f"node '{id}' not found")
    return str(node.get("path"))