agentforge-graph 0.3.2__py3-none-any.whl

agentforge_graph/retrieve/rerank.py

@@ -0,0 +1,251 @@
+"""Reranking hook (feat-006 seam; ENH-009).
+
+Three modes, all behind ``retrieve.rerank``:
+
+- ``off`` (default) — identity; pure cosine + graph score.
+- ``lexical`` — a deterministic, dependency-free blend of the base score with the
+  **subtoken overlap** between the query and the candidate's name + code, so a
+  chunk whose symbol the query names (``ZodObject``, ``_parse``, ``res.send``)
+  sorts up even when its raw cosine landed *near* the answer. Useful for
+  keyword/symbol-naming queries; measured mixed on prose (hence opt-in).
+- ``cross_encoder`` — a real semantic re-score: a cross-encoder relevance model
+  (``sentence-transformers``, the ``rerank`` extra) scores each (query,
+  candidate) pair, blended with the base score. The highest-leverage lever for
+  natural-language → symbol precision. The model is lazy-loaded so the base
+  install / CI never import torch; the blend logic is injectable (``CrossScorer``)
+  so it is tested without the model. Third-party only — no ``agentforge`` import
+  (ADR-0001).
+
+All rerankers are deterministic given their inputs and order-stable on ties."""
+
+from __future__ import annotations
+
+import asyncio
+import math
+import re
+from typing import Protocol
+
+from .pack import ContextItem
+
+# A capable, small default cross-encoder; overridable via ``retrieve.rerank_model``.
+DEFAULT_CROSS_ENCODER = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+_MAX_CANDIDATE_CHARS = 2000  # cross-encoders truncate anyway; bound the payload
+
+# Function words that carry no retrieval signal — dropped from both sides.
+_STOP = frozenset(
+    [
+        "a",
+        "an",
+        "the",
+        "is",
+        "are",
+        "was",
+        "were",
+        "be",
+        "how",
+        "do",
+        "does",
+        "did",
+        "of",
+        "to",
+        "in",
+        "on",
+        "for",
+        "and",
+        "or",
+        "with",
+        "this",
+        "that",
+        "it",
+        "as",
+        "at",
+        "by",
+        "from",
+        "into",
+        "than",
+        "then",
+        "over",
+        "under",
+        "not",
+        "no",
+        "your",
+        "you",
+        "we",
+        "i",
+        "me",
+        "my",
+        "our",
+        "their",
+        "its",
+        "they",
+        "them",
+        "he",
+        "she",
+        "where",
+        "what",
+        "which",
+        "who",
+        "when",
+        "why",
+        "can",
+        "could",
+        "should",
+        "would",
+        "will",
+        "shall",
+        "may",
+        "might",
+        "must",
+        "have",
+        "has",
+        "had",
+        "get",
+        "set",
+        "up",
+        "out",
+        "off",
+        "via",
+        "per",
+        "use",
+        "used",
+        "using",
+        "return",
+        "returns",
+    ]
+)
+# Split identifiers into subtokens: ALLCAPS, CamelChunk, or a run of lower/digits.
+_CAMEL = re.compile(r"[A-Z]+(?![a-z])|[A-Z][a-z0-9]*|[a-z0-9]+")
+
+
+def _tokens(text: str) -> set[str]:
+    """Lowercased subtokens of ``text``: splits on non-alphanumerics *and*
+    camelCase, so ``ZodObject._parse`` → {zod, object, parse}. Drops stopwords
+    and single chars (noise)."""
+    out: set[str] = set()
+    for raw in re.split(r"[^A-Za-z0-9]+", text):
+        for m in _CAMEL.findall(raw):
+            low = m.lower()
+            if len(low) >= 2 and low not in _STOP:
+                out.add(low)
+    return out
+
+
+class Reranker(Protocol):
+    async def rerank(self, query: str, items: list[ContextItem]) -> list[ContextItem]: ...
+
+
+class NoopReranker:
+    """Identity reranker (rerank: off)."""
+
+    async def rerank(self, query: str, items: list[ContextItem]) -> list[ContextItem]:
+        return items
+
+
+class LexicalReranker:
+    """Blend base score with query↔candidate subtoken overlap, then re-sort.
+
+    ``final = (1 - weight)·base + weight·overlap``, where ``overlap`` is the
+    fraction of (non-stopword) query subtokens present in the candidate's
+    name + code. Deterministic and order-stable on ties."""
+
+    def __init__(self, weight: float = 0.5) -> None:
+        self._w = max(0.0, min(1.0, weight))
+
+    async def rerank(self, query: str, items: list[ContextItem]) -> list[ContextItem]:
+        qtoks = _tokens(query)
+        if not qtoks or not items:
+            return items
+        rescored: list[ContextItem] = []
+        for it in items:
+            itoks = _tokens(f"{it.name} {it.code or ''}")
+            overlap = len(qtoks & itoks) / len(qtoks)
+            final = (1.0 - self._w) * it.score + self._w * overlap
+            rescored.append(
+                it.model_copy(update={"score": final, "why": [*it.why, f"lexical {overlap:.2f}"]})
+            )
+        rescored.sort(key=lambda i: (-i.score, i.id))  # id tiebreak = deterministic
+        return rescored
+
+
+class CrossScorer(Protocol):
+    """Scores (query, candidate-text) pairs — higher = more relevant. The
+    injection seam that keeps the model out of the blend logic (and out of CI)."""
+
+    def score(self, query: str, texts: list[str]) -> list[float]: ...
+
+
+def _candidate_text(it: ContextItem) -> str:
+    body = it.code or it.signature()
+    return f"{it.name}\n{body}"[:_MAX_CANDIDATE_CHARS]
+
+
+class CrossEncoderReranker:
+    """Re-score the top-k candidates with a cross-encoder, then blend and re-sort.
+
+    ``final = (1 - weight)·base + weight·σ(cross_score)`` — the cross-encoder's
+    raw relevance logit is squashed to ``[0, 1]`` (so it is comparable to the
+    cosine-scale base score) and blended. The model call runs off the event loop
+    (``to_thread``); identity on an empty query/candidate set."""
+
+    def __init__(self, scorer: CrossScorer, weight: float = 0.5) -> None:
+        self._scorer = scorer
+        self._w = max(0.0, min(1.0, weight))
+
+    async def rerank(self, query: str, items: list[ContextItem]) -> list[ContextItem]:
+        if not query or not items:
+            return items
+        texts = [_candidate_text(it) for it in items]
+        raw = await asyncio.to_thread(self._scorer.score, query, texts)
+        rescored: list[ContextItem] = []
+        for it, r in zip(items, raw, strict=True):
+            ce = 1.0 / (1.0 + math.exp(-r))  # σ → [0, 1]
+            final = (1.0 - self._w) * it.score + self._w * ce
+            rescored.append(
+                it.model_copy(update={"score": final, "why": [*it.why, f"cross-encoder {ce:.2f}"]})
+            )
+        rescored.sort(key=lambda i: (-i.score, i.id))
+        return rescored
+
+
+class SentenceTransformerScorer:
+    """A ``CrossScorer`` backed by ``sentence_transformers.CrossEncoder``. The
+    model is loaded lazily on first use, so importing this module (and running
+    CI) never pulls torch; the import error names the extra to install."""
+
+    def __init__(self, model_name: str = DEFAULT_CROSS_ENCODER) -> None:
+        self._model_name = model_name
+        self._model: object | None = None
+
+    def _ensure_model(self) -> object:
+        if self._model is None:
+            try:
+                from sentence_transformers import CrossEncoder
+            except ImportError as exc:  # the extra isn't installed
+                raise ImportError(
+                    "cross-encoder rerank needs the 'rerank' extra (uv sync --extra rerank)"
+                ) from exc
+            self._model = CrossEncoder(self._model_name)
+        return self._model
+
+    def score(self, query: str, texts: list[str]) -> list[float]:
+        if not texts:
+            return []
+        model = self._ensure_model()
+        return [float(s) for s in model.predict([(query, t) for t in texts])]  # type: ignore[attr-defined]
+
+
+def reranker_from_config(rerank: str, weight: float = 0.5, model: str = "") -> Reranker:
+    """Resolve the ``retrieve.rerank`` config value to a reranker.
+    ``off``/empty → identity; ``lexical`` → :class:`LexicalReranker`;
+    ``cross_encoder`` → :class:`CrossEncoderReranker` over a lazily-loaded
+    sentence-transformers model (``model`` overrides the default)."""
+    ref = (rerank or "off").strip()
+    if ref in ("", "off"):
+        return NoopReranker()
+    if ref == "lexical":
+        return LexicalReranker(weight)
+    if ref == "cross_encoder":
+        return CrossEncoderReranker(
+            SentenceTransformerScorer(model or DEFAULT_CROSS_ENCODER), weight
+        )
+    raise ValueError(f"unknown reranker {ref!r}; use 'off', 'lexical' or 'cross_encoder'")

agentforge_graph/retrieve/retriever.py

@@ -0,0 +1,286 @@
+"""``Retriever`` — vector entry → typed graph expansion → provenance-weighted
+merge. Deterministic and LLM-free; the single retrieval surface feat-008 and
+the enrichers ride on.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Literal
+
+from agentforge_graph.config import RetrieveConfig
+from agentforge_graph.core import Direction, EdgeKind, Node, NodeKind, Source, SymbolID
+from agentforge_graph.embed import Embedder
+from agentforge_graph.store import Store
+
+from .pack import ContextItem, ContextPack
+from .rerank import NoopReranker, Reranker
+from .scoring import dedupe_max, edge_weight, step_score
+
+Mode = Literal["context", "impact", "definition", "similar"]
+
+# code-symbol kinds an as_of allow-filter constrains (feat-009)
+_SYMBOL_KINDS = frozenset({NodeKind.CLASS, NodeKind.FUNCTION, NodeKind.METHOD})
+
+# A query "smells architectural" when it asks about decisions/rationale/design — the
+# case where ADR/doc prose SHOULD rank with code. Else docs are down-weighted (feat-010).
+_ARCH_TERMS = frozenset(
+    {
+        "why",
+        "decision",
+        "decisions",
+        "rationale",
+        "architecture",
+        "architectural",
+        "design",
+        "convention",
+        "conventions",
+        "adr",
+        "govern",
+        "governs",
+        "principle",
+        "tradeoff",
+        "trade-off",
+        "policy",
+        "constraint",
+        "supposed",
+        "allowed",
+        "forbidden",
+        "deprecated",
+    }
+)
+
+_WORD_RE = re.compile(r"[a-z0-9-]+")
+
+
+def _is_architectural(query: str) -> bool:
+    return any(w in _ARCH_TERMS for w in _WORD_RE.findall(query.lower()))
+
+
+# Trust order for the min_provenance filter: llm < parsed < resolved <= manual
+# (human-asserted facts are trusted; ADR-0004 / spec §2). Distinct from
+# GraphQuery.min_source and from the scoring edge_weights.
+_RANK: dict[Source, int] = {Source.LLM: 0, Source.PARSED: 1, Source.RESOLVED: 2, Source.MANUAL: 3}
+_FLOOR: dict[str, int] = {"parsed": 1, "resolved": 2}
+
+# feat-009 churn/authorship fields denormalised onto a symbol's node.attrs; the
+# Retriever surfaces them on the item without joining the temporal sidecar (it
+# stays in the deterministic core, ADR-0001). Empty → item.temporal stays None.
+_TEMPORAL_KEYS = (
+    "introduced",
+    "introduced_ts",
+    "last_changed",
+    "last_changed_ts",
+    "churn_30d",
+    "churn_90d",
+    "top_authors",
+)
+
+
+def _temporal_attrs(node: Node) -> dict[str, object] | None:
+    out = {k: node.attrs[k] for k in _TEMPORAL_KEYS if k in node.attrs}
+    return out or None
+
+
+_MODE_EDGES: dict[Mode, tuple[list[EdgeKind], Direction]] = {
+    # GOVERNS/DESCRIBES (feat-010) surface the decision/doc governing a retrieved
+    # symbol; TAGGED + SUMMARIZES (feat-012) surface its design-pattern role and
+    # the module summary. The differentiators; llm items obey include_llm_facts.
+    "context": (
+        [
+            EdgeKind.CALLS,
+            EdgeKind.CONTAINS,
+            EdgeKind.INHERITS,
+            EdgeKind.REFERENCES,
+            EdgeKind.GOVERNS,
+            EdgeKind.DESCRIBES,
+            EdgeKind.TAGGED,
+            EdgeKind.SUMMARIZES,
+        ],
+        "both",
+    ),
+    "impact": ([EdgeKind.CALLS, EdgeKind.IMPORTS, EdgeKind.IMPLEMENTS], "in"),
+    "definition": ([EdgeKind.CONTAINS, EdgeKind.CHUNK_OF], "both"),
+    "similar": ([], "both"),
+}
+
+
+class Retriever:
+    def __init__(
+        self,
+        store: Store,
+        embedder: Embedder,
+        config: RetrieveConfig,
+        reranker: Reranker | None = None,
+    ) -> None:
+        self.store = store
+        self.embedder = embedder
+        self.config = config
+        self.reranker: Reranker = reranker or NoopReranker()
+
+    async def retrieve(
+        self,
+        query: str | None = None,
+        symbol: str | None = None,
+        mode: Mode = "context",
+        k: int | None = None,
+        depth: int | None = None,
+        edge_kinds: list[EdgeKind] | None = None,
+        min_provenance: Literal["parsed", "resolved"] | None = None,
+        include_llm_facts: bool = True,
+        allow_ids: set[str] | None = None,
+    ) -> ContextPack:
+        cfg = self.config
+        k = cfg.k if k is None else k
+        depth = cfg.depth if depth is None else depth
+        items: list[ContextItem] = []
+        notes: list[str] = []
+        seeds: dict[str, float] = {}
+
+        # --- entry ---
+        if query is not None:
+            qvec = (await self.embedder.embed([query], "query"))[0]
+            # down-weight ADR/doc prose so code outranks equally-similar docs, unless
+            # the query smells architectural (then docs keep their full score). feat-010.
+            doc_w = 1.0 if _is_architectural(query) else self.config.doc_weight
+            for hit in await self.store.vectors.search(qvec, k):
+                node = await self.store.graph.get(hit.ref)
+                if node is None:
+                    continue
+                score = hit.score * doc_w if node.kind is NodeKind.DOC_CHUNK else hit.score
+                items.append(self._item(node, score, [f"vector hit {score:.2f}"]))
+                if mode != "similar":
+                    # a chunk hit seeds its symbols; a summary hit (feat-012)
+                    # seeds the code it summarizes — concept query → code.
+                    for edge in await self.store.graph.adjacent(
+                        hit.ref, [EdgeKind.CHUNK_OF, EdgeKind.SUMMARIZES], "out"
+                    ):
+                        seeds[edge.dst] = max(seeds.get(edge.dst, 0.0), score)
+                    # a doc-chunk hit (feat-010) seeds the code it attaches to: an
+                    # ADR section seeds its containing Decision (CONTAINS in → then
+                    # GOVERNS to governed code); a docstring seeds the symbol it
+                    # DESCRIBES (out). Either way a prose query reaches the code.
+                    if node.kind is NodeKind.DOC_CHUNK:
+                        for edge in await self.store.graph.adjacent(
+                            hit.ref, [EdgeKind.CONTAINS], "in"
+                        ):
+                            seeds[edge.src] = max(seeds.get(edge.src, 0.0), score)
+                        for edge in await self.store.graph.adjacent(
+                            hit.ref, [EdgeKind.DESCRIBES], "out"
+                        ):
+                            seeds[edge.dst] = max(seeds.get(edge.dst, 0.0), score)
+        if symbol is not None:
+            seeds[symbol] = max(seeds.get(symbol, 0.0), 1.0)
+
+        for sid, score in list(seeds.items()):
+            node = await self.store.graph.get(sid)
+            if node is None:
+                del seeds[sid]
+                continue
+            items.append(self._item(node, score, ["entry"]))
+
+        # --- expand ---
+        kinds, direction = _MODE_EDGES[mode]
+        if edge_kinds is not None:
+            kinds = edge_kinds
+        if mode != "similar" and depth > 0 and seeds:
+            await self._expand(seeds, kinds, direction, depth, items, notes)
+
+        # --- merge ---
+        items = self._filter(items, min_provenance, include_llm_facts)
+        if allow_ids is not None:  # feat-009 as_of: drop symbols not alive at the commit
+            items = [it for it in items if it.kind not in _SYMBOL_KINDS or it.id in allow_ids]
+        items = dedupe_max(items)
+        items = await self.reranker.rerank(query or "", items)
+        return ContextPack(query=query, symbol=symbol, mode=mode, items=items, notes=notes)
+
+    async def _expand(
+        self,
+        seeds: dict[str, float],
+        kinds: list[EdgeKind],
+        direction: Direction,
+        depth: int,
+        items: list[ContextItem],
+        notes: list[str],
+    ) -> None:
+        cfg = self.config
+        frontier = dict(seeds)
+        visited = set(seeds)
+        for hop in range(1, depth + 1):
+            nxt: dict[str, float] = {}
+            for sid, score in frontier.items():
+                parent = await self.store.graph.get(sid)
+                pname = parent.name if parent else sid
+                edges = await self.store.graph.adjacent(sid, kinds, direction)
+                if len(edges) > cfg.fanout_cap:
+                    notes.append(f"fan-out cap {cfg.fanout_cap} at {pname} ({len(edges)} edges)")
+                    edges = edges[: cfg.fanout_cap]
+                for edge in edges:
+                    other = edge.dst if edge.src == sid else edge.src
+                    onode = await self.store.graph.get(other)
+                    if onode is None:
+                        continue
+                    oscore = step_score(
+                        score, cfg.decay, edge_weight(cfg.edge_weights, edge.provenance.source)
+                    )
+                    items.append(
+                        self._item(onode, oscore, [f"{edge.kind.value} of {pname} (hop {hop})"])
+                    )
+                    if other not in visited:
+                        visited.add(other)
+                        nxt[other] = max(nxt.get(other, 0.0), oscore)
+            frontier = nxt
+            if not frontier:
+                break
+
+    def _item(self, node: Node, score: float, why: list[str]) -> ContextItem:
+        return ContextItem(
+            id=node.id,
+            kind=node.kind,
+            name=node.name,
+            score=score,
+            path=SymbolID.parse(node.id).path,
+            span=node.span,
+            code=self._render_code(node),
+            provenance=node.provenance.source,
+            why=list(why),
+            temporal=_temporal_attrs(node),
+        )
+
+    @staticmethod
+    def _render_code(node: Node) -> str | None:
+        """The verbatim block a retrieved item renders. A Decision (feat-010)
+        shows its status/date inline so the agent sees governance at a glance."""
+        if node.kind is NodeKind.DECISION:
+            status = node.attrs.get("status", "")
+            date = node.attrs.get("date", "")
+            adr = node.attrs.get("adr_id", "")
+            stamp = ", ".join(x for x in (status, date) if x)
+            prefix = f"[{stamp}] " if stamp else ""
+            label = f"{adr}: " if adr else ""
+            return f"{prefix}{label}{node.attrs.get('title', node.name)}"
+        if node.kind is NodeKind.PATTERN_TAG:
+            return f"[llm] design pattern: {node.name}"
+        if node.kind is NodeKind.SUMMARY:
+            return f"[summary] {node.attrs.get('text', '')}"
+        if node.kind is NodeKind.DOC_CHUNK:  # feat-010 — ADR/doc prose
+            heading = node.attrs.get("heading", "")
+            text = node.attrs.get("text", "")
+            return f"[doc] {heading}\n{text}".strip() if heading else f"[doc] {text}".strip()
+        return node.attrs.get("code")
+
+    def _filter(
+        self,
+        items: list[ContextItem],
+        min_provenance: Literal["parsed", "resolved"] | None,
+        include_llm_facts: bool,
+    ) -> list[ContextItem]:
+        floor = _FLOOR[min_provenance] if min_provenance else None
+        out: list[ContextItem] = []
+        for it in items:
+            if not include_llm_facts and it.provenance is Source.LLM:
+                continue
+            if floor is not None and _RANK[it.provenance] < floor:
+                continue
+            out.append(it)
+        return out

agentforge_graph/retrieve/scoring.py

@@ -0,0 +1,36 @@
+"""Scoring math for retrieval: provenance edge weights, per-hop decay, and
+dedupe (max score wins, why-traces unioned)."""
+
+from __future__ import annotations
+
+from agentforge_graph.core import Source
+
+from .pack import ContextItem
+
+
+def edge_weight(weights: dict[str, float], source: Source) -> float:
+    """Weight an expansion edge by its provenance — resolved > parsed > llm
+    (ADR-0004). Unknown sources fall back to 0.5."""
+    return weights.get(source.value, 0.5)
+
+
+def step_score(parent_score: float, decay: float, weight: float) -> float:
+    """One hop of decay: ``parent × decay × edge_weight``. Repeated over hops
+    yields the ``decay^hop`` falloff."""
+    return parent_score * decay * weight
+
+
+def dedupe_max(items: list[ContextItem]) -> list[ContextItem]:
+    """Collapse items sharing an id to the highest-scoring one, unioning the
+    why-traces; return sorted by score descending."""
+    best: dict[str, ContextItem] = {}
+    whys: dict[str, list[str]] = {}
+    for it in items:
+        acc = whys.setdefault(it.id, [])
+        for w in it.why:
+            if w not in acc:
+                acc.append(w)
+        if it.id not in best or it.score > best[it.id].score:
+            best[it.id] = it
+    merged = [it.model_copy(update={"why": whys[i]}) for i, it in best.items()]
+    return sorted(merged, key=lambda i: i.score, reverse=True)

agentforge_graph/serve/__init__.py

@@ -0,0 +1,19 @@
+"""agentforge_graph.serve — MCP server & AgentForge tool API (feat-008).
+
+The framework-facing serving layer: the nine read-only tools over feat-006/007,
+bound both as native AgentForge ``Tool`` instances (``code_graph_tools``) and an
+MCP server (``serve_mcp``) over **stdio or streamable-HTTP**, from one definition.
+This package imports ``agentforge`` (the deliberate ADR-0001 exception).
+"""
+
+from __future__ import annotations
+
+from .engine import TOOL_API_VERSION
+from .server import build_mcp_server, code_graph_tools, serve_mcp
+
+__all__ = [
+    "TOOL_API_VERSION",
+    "code_graph_tools",
+    "serve_mcp",
+    "build_mcp_server",
+]