agentforge-graph 0.3.2__py3-none-any.whl

agentforge_graph/serve/tools.py

@@ -0,0 +1,307 @@
+"""The read-only CKG tools (feat-008), thin over feat-006/007/009.
+
+Each is an AgentForge ``Tool`` holding the shared ``_Engine``; ``run`` clamps
+params to ``ServeConfig`` and returns a JSON string (structured) or text (the
+map) — MCP coerces results via ``str()``, so JSON keeps structure intact.
+Every structured envelope carries ``indexed_commit`` + ``dirty`` (staleness)
+and a ``truncated`` flag. Names & schemas are locked at v1.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, ClassVar
+
+from agentforge_core.contracts.tool import Tool
+from pydantic import BaseModel, Field
+
+from agentforge_graph.chunking import estimate_tokens
+from agentforge_graph.core import EdgeKind, GraphQuery
+
+from .engine import TOOL_API_VERSION, _Engine
+
+
+class RepoMapInput(BaseModel):
+    budget_tokens: int = Field(default=2000, description="token budget for the map")
+    focus: list[str] = Field(
+        default_factory=list, description="paths or symbol ids to rank the map around"
+    )
+
+
+class SearchInput(BaseModel):
+    query: str = Field(description="natural-language question about the code")
+    k: int = Field(default=8, description="number of vector hits (clamped to serve.max_k)")
+    mode: str = Field(default="context", description="context | similar | definition | impact")
+
+
+class SymbolInput(BaseModel):
+    symbol_id: str = Field(default="", description="exact symbol id, if known")
+    name: str = Field(default="", description="symbol name (use with path) when id is unknown")
+    path: str = Field(default="", description="file path for a name lookup")
+
+
+class ImpactInput(BaseModel):
+    symbol_id: str = Field(description="symbol whose reverse dependencies to trace")
+    depth: int = Field(default=1, description="hops (clamped to serve.max_depth)")
+
+
+class NeighborsInput(BaseModel):
+    symbol_id: str = Field(description="symbol id")
+    edge_kinds: list[str] = Field(
+        default_factory=list, description="edge kinds to follow (e.g. CALLS, CONTAINS); default all"
+    )
+    depth: int = Field(default=1, description="hops (clamped to serve.max_depth)")
+
+
+class RoutesInput(BaseModel):
+    method: str = Field(default="", description="filter by HTTP method, e.g. GET (optional)")
+    path: str = Field(default="", description="filter by path prefix, e.g. /users (optional)")
+
+
+class DecisionsInput(BaseModel):
+    scope: str = Field(default="", description="restrict to decisions governing a path prefix")
+    status: str = Field(default="", description="filter by status, e.g. accepted (optional)")
+
+
+class ExplainInput(BaseModel):
+    symbol_id: str = Field(description="exact symbol id to explain")
+
+
+class HistoryInput(BaseModel):
+    symbol_id: str = Field(description="exact symbol id whose git history to report")
+
+
+class EmptyInput(BaseModel):
+    pass
+
+
+class _CkgTool(Tool):
+    capabilities: ClassVar[frozenset[str]] = frozenset()
+
+    def __init__(self, engine: _Engine) -> None:
+        self._engine = engine
+
+    async def _pack_json(self, pack: Any) -> str:
+        data = pack.to_dict()
+        data.update(await self._engine.staleness())
+        data["tool_api_version"] = TOOL_API_VERSION
+        cap = self._engine.serve.response_token_cap
+        truncated = False
+        while len(data.get("items", [])) > 1 and estimate_tokens(json.dumps(data)) > cap:
+            data["items"].pop()
+            truncated = True
+        if truncated:
+            data.setdefault("notes", []).append("response truncated to fit response_token_cap")
+        data["truncated"] = truncated
+        return json.dumps(data, indent=2)
+
+
+class CkgRepoMap(_CkgTool):
+    name: ClassVar[str] = "ckg_repo_map"
+    description: ClassVar[str] = (
+        "Orient in a codebase: a budget-aware, centrality-ranked map of the most "
+        "structurally important symbols (signatures), grouped by file. Call this FIRST "
+        "to understand an unfamiliar repo. Pass `focus` (paths/symbol ids) to rank the "
+        "map around a working set."
+    )
+    input_schema: ClassVar[type[BaseModel]] = RepoMapInput
+
+    async def run(self, **kwargs: Any) -> str:
+        repomap = await self._engine.repomap()
+        budget = min(int(kwargs.get("budget_tokens", 2000)), self._engine.serve.response_token_cap)
+        text = await repomap.render(budget_tokens=budget, focus=kwargs.get("focus") or None)
+        return text or "(empty repo map)"
+
+
+class CkgSearch(_CkgTool):
+    name: ClassVar[str] = "ckg_search"
+    description: ClassVar[str] = (
+        "Search the codebase for code relevant to a natural-language question. Returns "
+        "ranked, CONNECTED context: matching chunks plus their symbols and neighbors. "
+        "`mode`: 'context' (default), 'similar' (just similar code), 'definition' (a "
+        "symbol's chunks), 'impact' (reverse deps — prefer ckg_impact). Take a result's "
+        "`id` to chain into ckg_impact / ckg_neighbors / ckg_symbol."
+    )
+    input_schema: ClassVar[type[BaseModel]] = SearchInput
+
+    async def run(self, **kwargs: Any) -> str:
+        retriever = await self._engine.retriever()
+        k = min(int(kwargs.get("k", 8)), self._engine.serve.max_k)
+        pack = await retriever.retrieve(
+            query=kwargs["query"], k=k, mode=kwargs.get("mode", "context")
+        )
+        return await self._pack_json(pack)
+
+
+class CkgSymbol(_CkgTool):
+    name: ClassVar[str] = "ckg_symbol"
+    description: ClassVar[str] = (
+        "Look up a specific symbol's definition: the symbol, its chunks and members. "
+        "Provide `symbol_id` (exact) or `name` + `path`. When the temporal layer is "
+        "enabled, items carry `temporal` (churn_90d, top_authors, introduced, "
+        "last_changed) — recency/ownership signals for triage."
+    )
+    input_schema: ClassVar[type[BaseModel]] = SymbolInput
+
+    async def run(self, **kwargs: Any) -> str:
+        symbol_id = kwargs.get("symbol_id", "")
+        if not symbol_id:
+            cg = await self._engine.code_graph()
+            res = await cg.store.graph.query(
+                GraphQuery(name=kwargs.get("name") or None, path_prefix=kwargs.get("path") or None)
+            )
+            symbol_id = res.nodes[0].id if res.nodes else ""
+        if not symbol_id:
+            return json.dumps({"error": "symbol not found", "tool_api_version": TOOL_API_VERSION})
+        retriever = await self._engine.retriever()
+        depth = min(1, self._engine.serve.max_depth)
+        pack = await retriever.retrieve(symbol=symbol_id, mode="definition", depth=depth)
+        return await self._pack_json(pack)
+
+
+class CkgImpact(_CkgTool):
+    name: ClassVar[str] = "ckg_impact"
+    description: ClassVar[str] = (
+        "Find what DEPENDS ON a symbol (reverse CALLS/IMPORTS/IMPLEMENTS): 'who calls "
+        "this', 'what breaks if I change this'. The impact question grep cannot answer."
+    )
+    input_schema: ClassVar[type[BaseModel]] = ImpactInput
+
+    async def run(self, **kwargs: Any) -> str:
+        retriever = await self._engine.retriever()
+        depth = min(int(kwargs.get("depth", 1)), self._engine.serve.max_depth)
+        pack = await retriever.retrieve(symbol=kwargs["symbol_id"], mode="impact", depth=depth)
+        return await self._pack_json(pack)
+
+
+class CkgNeighbors(_CkgTool):
+    name: ClassVar[str] = "ckg_neighbors"
+    description: ClassVar[str] = (
+        "List a symbol's typed graph neighbors (edges in both directions), optionally "
+        "filtered by edge kind (CALLS, CONTAINS, IMPORTS, INHERITS, REFERENCES, CHUNK_OF)."
+    )
+    input_schema: ClassVar[type[BaseModel]] = NeighborsInput
+
+    async def run(self, **kwargs: Any) -> str:
+        cg = await self._engine.code_graph()
+        kinds = [EdgeKind(k) for k in kwargs.get("edge_kinds", [])] or None
+        depth = min(int(kwargs.get("depth", 1)), self._engine.serve.max_depth)
+        start = kwargs["symbol_id"]
+        seen: set[tuple[str, str, str]] = set()
+        edges: list[dict[str, str]] = []
+        frontier = {start}
+        visited = {start}
+        for _ in range(depth):
+            nxt: set[str] = set()
+            for nid in frontier:
+                for e in await cg.store.graph.adjacent(nid, kinds, "both"):
+                    key = (e.src, e.dst, e.kind.value)
+                    if key not in seen:
+                        seen.add(key)
+                        edges.append(
+                            {
+                                "src": e.src,
+                                "dst": e.dst,
+                                "kind": e.kind.value,
+                                "provenance": e.provenance.source.value,
+                            }
+                        )
+                    other = e.dst if e.src == nid else e.src
+                    if other not in visited:
+                        visited.add(other)
+                        nxt.add(other)
+            frontier = nxt
+        envelope = await self._engine.staleness()
+        return json.dumps(
+            {"symbol_id": start, "edges": edges, **envelope, "tool_api_version": TOOL_API_VERSION},
+            indent=2,
+        )
+
+
+class CkgStatus(_CkgTool):
+    name: ClassVar[str] = "ckg_status"
+    description: ClassVar[str] = (
+        "Report index status: the indexed git commit, whether it is stale vs the working "
+        "tree (`dirty`), node counts by kind, and the tool-API version. If results seem "
+        "out of date and `dirty` is true, tell the user to re-run `ckg index`."
+    )
+    input_schema: ClassVar[type[BaseModel]] = EmptyInput
+
+    async def run(self, **kwargs: Any) -> str:
+        return json.dumps(await self._engine.status(), indent=2)
+
+
+class CkgRoutes(_CkgTool):
+    name: ClassVar[str] = "ckg_routes"
+    description: ClassVar[str] = (
+        "List the app's HTTP API surface: every framework route (e.g. FastAPI) with its "
+        "method, path pattern and handler symbol id. Filter by `method` and/or `path` prefix. "
+        "Take a route's `handler` into ckg_symbol / ckg_impact to see the handler and what it "
+        "touches. Returns empty if no web framework is detected."
+    )
+    input_schema: ClassVar[type[BaseModel]] = RoutesInput
+
+    async def run(self, **kwargs: Any) -> str:
+        data = await self._engine.routes(
+            method=kwargs.get("method", ""), path=kwargs.get("path", "")
+        )
+        return json.dumps(data, indent=2)
+
+
+class CkgDecisions(_CkgTool):
+    name: ClassVar[str] = "ckg_decisions"
+    description: ClassVar[str] = (
+        "List the architecture decisions (ADRs) governing the codebase: each decision's "
+        "status, date, title and the symbols/files it governs. Filter by `scope` (a path "
+        "prefix the decision governs) and/or `status` (e.g. accepted). Call this before a "
+        "refactor to check no documented decision forbids the change. Empty if no ADRs."
+    )
+    input_schema: ClassVar[type[BaseModel]] = DecisionsInput
+
+    async def run(self, **kwargs: Any) -> str:
+        data = await self._engine.decisions(
+            scope=kwargs.get("scope", ""), status=kwargs.get("status", "")
+        )
+        return json.dumps(data, indent=2)
+
+
+class CkgExplain(_CkgTool):
+    name: ClassVar[str] = "ckg_explain"
+    description: ClassVar[str] = (
+        "Explain a symbol: its LLM-derived design-pattern tags (e.g. 'Repository', with "
+        "confidence + rationale) and its 1-hop typed graph facts. Use to learn what role a "
+        "class/function plays before changing it. Tags are [llm]-provenance; empty until "
+        "`ckg enrich` has run."
+    )
+    input_schema: ClassVar[type[BaseModel]] = ExplainInput
+
+    async def run(self, **kwargs: Any) -> str:
+        return json.dumps(await self._engine.explain(kwargs["symbol_id"]), indent=2)
+
+
+class CkgHistory(_CkgTool):
+    name: ClassVar[str] = "ckg_history"
+    description: ClassVar[str] = (
+        "Report a symbol's git evolution (feat-009 temporal): when it was introduced and "
+        "last changed, churn over the last 30/90 days, its top authors, and its lifecycle "
+        "events. Use to gauge recency/ownership before changing a symbol or to triage a "
+        "regression. Returns `available: false` if the temporal layer is off or unindexed."
+    )
+    input_schema: ClassVar[type[BaseModel]] = HistoryInput
+
+    async def run(self, **kwargs: Any) -> str:
+        return json.dumps(await self._engine.history(kwargs["symbol_id"]), indent=2)
+
+
+ALL_TOOLS = [
+    CkgRepoMap,
+    CkgSearch,
+    CkgSymbol,
+    CkgImpact,
+    CkgNeighbors,
+    CkgStatus,
+    CkgRoutes,
+    CkgDecisions,
+    CkgExplain,
+    CkgHistory,
+]

agentforge_graph/store/__init__.py

@@ -0,0 +1,32 @@
+"""agentforge_graph.store — persistent storage adapters (feat-003).
+
+Embedded-first (ADR-0006): the default ``Store`` writes a Kuzu graph DB
+and a LanceDB vector index under ``.ckg/`` with no server. Server adapters
+(Neo4j, FalkorDB) register out-of-tree via entry points. Imports nothing
+from ``agentforge`` (ADR-0001).
+"""
+
+from __future__ import annotations
+
+from .errors import (
+    DriverNotFound,
+    SchemaVersionError,
+    StoreConfigError,
+    StoreError,
+)
+from .facade import Store
+from .kuzu_store import KuzuGraphStore
+from .lance_store import LanceVectorStore
+from .registry import graph_driver, vector_driver
+
+__all__ = [
+    "Store",
+    "KuzuGraphStore",
+    "LanceVectorStore",
+    "graph_driver",
+    "vector_driver",
+    "StoreError",
+    "StoreConfigError",
+    "DriverNotFound",
+    "SchemaVersionError",
+]

agentforge_graph/store/_rowmap.py

@@ -0,0 +1,102 @@
+"""Pure ``Node``/``Edge`` ↔ property-row mapping, shared by every property-graph
+adapter (Kuzu embedded, Neo4j server — ENH-004). No DB driver imports: the open
+schema (arbitrary kinds, free-form ``attrs``) is flattened to a fixed set of
+scalar properties (``kind``, ``name``, span, provenance, ``origin_path``) with
+``attrs`` as a JSON string, so an unrecognized kind round-trips with no DDL
+change (ADR-0005). Each backend supplies the storage; this supplies the shape.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from agentforge_graph.core import Edge, EdgeKind, Node, NodeKind, Provenance, Source, SymbolID
+
+# Trust ordering (PARSED is highest-trust): a node passes a ``min_source`` floor
+# iff its source rank is >= the floor's.
+SOURCE_RANK = {Source.LLM: 0, Source.RESOLVED: 1, Source.MANUAL: 1, Source.PARSED: 2}
+
+
+def dump_attrs(attrs: dict[str, Any]) -> str:
+    return json.dumps(attrs, sort_keys=True)
+
+
+def load_attrs(s: str | None) -> dict[str, Any]:
+    return json.loads(s) if s else {}
+
+
+def acceptable_sources(floor: Source) -> list[str]:
+    threshold = SOURCE_RANK[floor]
+    return [s.value for s, rank in SOURCE_RANK.items() if rank >= threshold]
+
+
+def node_params(node: Node, origin_path: str) -> dict[str, Any]:
+    span_start, span_end = node.span if node.span is not None else (None, None)
+    p = node.provenance
+    return {
+        "id": node.id,
+        "kind": node.kind.value,
+        "name": node.name,
+        "span_start": span_start,
+        "span_end": span_end,
+        "attrs": dump_attrs(node.attrs),
+        "sym_path": SymbolID.parse(node.id).path,
+        "prov_source": p.source.value,
+        "prov_extractor": p.extractor,
+        "prov_commit": p.commit,
+        "prov_confidence": p.confidence,
+        "origin_path": origin_path,
+    }
+
+
+def edge_params(edge: Edge, origin_path: str) -> dict[str, Any]:
+    p = edge.provenance
+    return {
+        "src": edge.src,
+        "dst": edge.dst,
+        "kind": edge.kind.value,
+        "attrs": dump_attrs(edge.attrs),
+        "prov_source": p.source.value,
+        "prov_extractor": p.extractor,
+        "prov_commit": p.commit,
+        "prov_confidence": p.confidence,
+        # An edge that carries its own owner file (resolver edges, feat-004) wins;
+        # otherwise the caller's stamp (the upserted file's path).
+        "origin_path": edge.origin_path or origin_path,
+        "resolved_from": "",
+    }
+
+
+def prov_from_row(d: dict[str, Any]) -> Provenance:
+    # The validating constructor — a corrupt row fails loudly, not silently.
+    return Provenance(
+        source=Source(d["prov_source"]),
+        extractor=d["prov_extractor"],
+        commit=d["prov_commit"],
+        confidence=d["prov_confidence"],
+    )
+
+
+def node_from_row(d: dict[str, Any]) -> Node:
+    # span may be absent (Neo4j drops null properties) or present-but-None (Kuzu).
+    span_start = d.get("span_start")
+    span = (span_start, d.get("span_end")) if span_start is not None else None
+    return Node(
+        id=d["id"],
+        kind=NodeKind(d["kind"]),
+        name=d["name"],
+        span=span,
+        attrs=load_attrs(d["attrs"]),
+        provenance=prov_from_row(d),
+    )
+
+
+def edge_from_row(d: dict[str, Any], src: str, dst: str) -> Edge:
+    return Edge(
+        src=src,
+        dst=dst,
+        kind=EdgeKind(d["kind"]),
+        attrs=load_attrs(d["attrs"]),
+        provenance=prov_from_row(d),
+    )

agentforge_graph/store/errors.py

@@ -0,0 +1,22 @@
+"""Storage-layer errors. All raised at ``Store.open`` / adapter open time
+(fail-at-startup), never mid-index — a half-written graph is worse than a
+refused one."""
+
+from __future__ import annotations
+
+
+class StoreError(Exception):
+    """Base for all storage-adapter errors."""
+
+
+class StoreConfigError(StoreError):
+    """Malformed or unreadable ``ckg.yaml`` ``store:`` block."""
+
+
+class DriverNotFound(StoreConfigError):
+    """Config names a graph/vector driver that isn't registered."""
+
+
+class SchemaVersionError(StoreError):
+    """On-disk index schema version differs from this build's. 0.x policy
+    is to rebuild the index (the data is derivable)."""

agentforge_graph/store/facade.py

@@ -0,0 +1,89 @@
+"""The ``Store`` facade — one ``GraphStore`` + one ``VectorStore`` resolved
+from ``ckg.yaml``, plus the vector→graph join (``expand``) that retrieval
+(feat-006) builds on. Embedded-first: the default writes ``.ckg/graph.kuzu``
+and ``.ckg/vectors.lance`` under the repo root (ADR-0006).
+
+All failure modes (unknown driver, schema mismatch, malformed config) raise
+at ``open`` — never mid-index.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from agentforge_graph.config import StoreConfig
+from agentforge_graph.core import EdgeKind, GraphStore, Node, QueryResult, ScoredRef, VectorStore
+
+from .errors import SchemaVersionError
+from .registry import graph_driver, vector_driver
+
+# Store-level on-disk layout version. Bumped when the .ckg/ layout changes;
+# 0.x policy on mismatch is rebuild (the index is derivable — ADR-0006).
+STORE_SCHEMA_VERSION = 1
+
+
+class Store:
+    """Owns a graph store and a vector store, resolved from config."""
+
+    def __init__(self, graph: GraphStore, vectors: VectorStore, config: StoreConfig) -> None:
+        self.graph = graph
+        self.vectors = vectors
+        self.config = config
+
+    @classmethod
+    async def open(cls, repo_path: str | Path = ".", config: str | Path | None = None) -> Store:
+        """Resolve drivers from ``config`` (a ckg.yaml path) and open the
+        embedded index under ``repo_path``/<store.path>. Raises before any
+        store is opened if the config or on-disk schema is bad."""
+        cfg = StoreConfig.load(config)
+        root = Path(repo_path) / cfg.path
+        _check_or_init_meta(root)  # fail-at-startup on schema mismatch
+        graph_cls = graph_driver(cfg.graph.driver)
+        vector_cls = vector_driver(cfg.vectors.driver)
+        # Embedded drivers use the path under .ckg/; server drivers (ENH-004)
+        # ignore the path and read connection details from their config block.
+        graph: GraphStore = await graph_cls.open(root / "graph.kuzu", config=cfg.graph.config)
+        vectors: VectorStore = await vector_cls.open(
+            root / "vectors.lance", config=cfg.vectors.config
+        )
+        return cls(graph, vectors, cfg)
+
+    async def expand(
+        self,
+        refs: list[ScoredRef],
+        kinds: list[EdgeKind] | None = None,
+        depth: int = 1,
+    ) -> QueryResult:
+        """Join vector hits back into the graph: for each ref, collect the
+        node and its ``kinds``-edge neighborhood within ``depth`` hops. The
+        single place the graph+vector join lives (feat-006)."""
+        nodes: dict[str, Node] = {}
+        for r in refs:
+            hit = await self.graph.get(r.ref)
+            if hit is not None:
+                nodes[hit.id] = hit
+            for nb in await self.graph.neighbors(r.ref, kinds, depth):
+                nodes[nb.id] = nb
+        return QueryResult(nodes=list(nodes.values()))
+
+    async def close(self) -> None:
+        await self.graph.close()
+        await self.vectors.close()
+
+
+def _check_or_init_meta(root: Path) -> None:
+    root.mkdir(parents=True, exist_ok=True)
+    meta = root / "meta.json"
+    if meta.exists():
+        data = json.loads(meta.read_text())
+        on_disk = data.get("schema_version")
+        if on_disk != STORE_SCHEMA_VERSION:
+            raise SchemaVersionError(
+                f"index at {root} is schema v{on_disk}, this build expects "
+                f"v{STORE_SCHEMA_VERSION}; rebuild the index (0.x policy)"
+            )
+    else:
+        meta.write_text(
+            json.dumps({"schema_version": STORE_SCHEMA_VERSION, "indexed_commit": ""}, indent=2)
+        )