agentforge-graph 0.3.2__py3-none-any.whl

agentforge_graph/enrich/bedrock_client.py

@@ -0,0 +1,115 @@
+"""Shared AWS Bedrock Claude client for the enrichers (feat-012).
+
+One boto3 ``bedrock-runtime`` client (lazy, optional STS assume-role, sync on a
+worker thread), one ``invoke`` that runs the Anthropic Messages API on Bedrock
+and accumulates cost from token usage. The pattern judge and the summarizer both
+ride this; it is the only model-calling surface (deterministic tests use the
+scripted variants).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Any
+
+# USD per 1M tokens (input, output). Conservative defaults; cheap tier.
+_PRICES: dict[str, tuple[float, float]] = {
+    "haiku-4-5": (1.0, 5.0),
+    "haiku": (0.80, 4.0),
+    "sonnet": (3.0, 15.0),
+}
+_DEFAULT_PRICE = (1.0, 5.0)
+
+
+def price_for(model: str) -> tuple[float, float]:
+    for key, price in _PRICES.items():
+        if key in model:
+            return price
+    return _DEFAULT_PRICE
+
+
+class BedrockClient:
+    def __init__(
+        self,
+        model: str = "us.anthropic.claude-haiku-4-5-20251001-v1:0",
+        region: str = "us-east-1",
+        assume_role_arn: str | None = None,
+        max_tokens: int = 512,
+    ) -> None:
+        self.model = model
+        self.region = region
+        self.assume_role_arn = assume_role_arn
+        self.max_tokens = max_tokens
+        self._client: Any = None
+        self.cost_usd = 0.0
+
+    def _bedrock(self) -> Any:
+        if self._client is None:
+            import boto3
+
+            if self.assume_role_arn:
+                sts = boto3.client("sts", region_name=self.region)
+                creds = sts.assume_role(RoleArn=self.assume_role_arn, RoleSessionName="ckg-enrich")[
+                    "Credentials"
+                ]
+                self._client = boto3.client(
+                    "bedrock-runtime",
+                    region_name=self.region,
+                    aws_access_key_id=creds["AccessKeyId"],
+                    aws_secret_access_key=creds["SecretAccessKey"],
+                    aws_session_token=creds["SessionToken"],
+                )
+            else:
+                self._client = boto3.client("bedrock-runtime", region_name=self.region)
+        return self._client
+
+    async def invoke(
+        self,
+        system: str,
+        user: str,
+        tools: list[dict[str, Any]] | None = None,
+        tool_name: str | None = None,
+    ) -> dict[str, Any]:
+        """One Messages call; accumulates cost from usage. Returns the raw
+        payload (``content`` blocks + ``usage``)."""
+        payload = await asyncio.to_thread(self._invoke, system, user, tools, tool_name)
+        cents_in, cents_out = price_for(self.model)
+        usage = payload.get("usage", {})
+        self.cost_usd += (
+            usage.get("input_tokens", 0) * cents_in + usage.get("output_tokens", 0) * cents_out
+        ) / 1_000_000
+        return payload
+
+    def _invoke(
+        self,
+        system: str,
+        user: str,
+        tools: list[dict[str, Any]] | None,
+        tool_name: str | None,
+    ) -> dict[str, Any]:
+        body: dict[str, Any] = {
+            "anthropic_version": "bedrock-2023-05-31",
+            "max_tokens": self.max_tokens,
+            "system": system,
+            "messages": [{"role": "user", "content": user}],
+        }
+        if tools is not None:
+            body["tools"] = tools
+            if tool_name is not None:
+                body["tool_choice"] = {"type": "tool", "name": tool_name}
+        resp = self._bedrock().invoke_model(
+            modelId=self.model,
+            contentType="application/json",
+            accept="application/json",
+            body=json.dumps(body),
+        )
+        result: dict[str, Any] = json.loads(resp["body"].read())
+        return result
+
+
+def text_of(payload: dict[str, Any]) -> str:
+    """Concatenate the text blocks of a Messages response."""
+    return "".join(
+        b.get("text", "") for b in payload.get("content", []) if b.get("type") == "text"
+    ).strip()

agentforge_graph/enrich/bedrock_summarizer.py

@@ -0,0 +1,23 @@
+"""``BedrockClaudeSummarizer`` — module summaries on AWS Bedrock (feat-012).
+
+Thin endpoint adapter: the summary prompts + plain-text completion live in the
+provider-neutral ``ClaudeSummarizer`` (``claude.py``); this wires it to a Bedrock
+transport (``BedrockClient``). The Anthropic-API sibling is
+``AnthropicClaudeSummarizer`` (``anthropic.py``). Tests use ``ScriptedSummarizer``.
+"""
+
+from __future__ import annotations
+
+from .bedrock_client import BedrockClient
+from .claude import ClaudeSummarizer
+
+
+class BedrockClaudeSummarizer(ClaudeSummarizer):
+    def __init__(
+        self,
+        model: str = "us.anthropic.claude-haiku-4-5-20251001-v1:0",
+        region: str = "us-east-1",
+        assume_role_arn: str | None = None,
+        max_tokens: int = 400,
+    ) -> None:
+        super().__init__(BedrockClient(model, region, assume_role_arn, max_tokens), model)

agentforge_graph/enrich/claude.py

@@ -0,0 +1,172 @@
+"""Provider-neutral Claude pattern judge + summarizer (ENH-003 phase 2).
+
+The judge/summarizer logic is identical whether Claude runs on **AWS Bedrock**
+or the **direct Anthropic API**: both return the Anthropic *Messages* response
+shape (``content`` blocks + ``usage``). Only the transport *client* differs. So
+the prompts + parsing live here once, and the per-endpoint modules
+(``bedrock.py``, ``anthropic.py``) just supply a ``ClaudeClient``.
+
+Tests drive the deterministic ``ScriptedJudge`` / ``ScriptedSummarizer`` instead;
+this base is exercised with a stub client (no network) plus the env-gated live
+tests.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+# ``price_for`` / ``text_of`` are provider-neutral helpers that have lived in
+# ``bedrock_client`` since feat-012; import them here rather than move the file
+# (keeps the public import path stable for existing tests).
+from .bedrock_client import text_of
+from .heuristics import Candidate
+from .judge import Verdict
+from .summarizer import FileContext, Summary
+
+_JUDGE_SYSTEM = (
+    "You classify a code symbol against GoF and architectural design patterns. "
+    "Confirm a pattern ONLY when the symbol's structure clearly supports it; prefer "
+    "rejecting over guessing. For each candidate pattern give is_match, a confidence "
+    "in [0,1], and a one-sentence rationale that cites the structural evidence."
+)
+
+_VERDICT_TOOL = {
+    "name": "submit_verdicts",
+    "description": "Return one verdict per candidate pattern.",
+    "input_schema": {
+        "type": "object",
+        "properties": {
+            "verdicts": {
+                "type": "array",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "pattern": {"type": "string"},
+                        "is_match": {"type": "boolean"},
+                        "confidence": {"type": "number"},
+                        "rationale": {"type": "string"},
+                    },
+                    "required": ["pattern", "is_match", "confidence", "rationale"],
+                },
+            }
+        },
+        "required": ["verdicts"],
+    },
+}
+
+_FILE_SYSTEM = (
+    "You write a one-paragraph summary of a source file for a developer orienting "
+    "in the codebase. State what the file is FOR and the role of its main symbols. "
+    "Summarize only what the signatures and names show — do not invent behaviour. "
+    "No preamble, no bullet lists."
+)
+_REPO_SYSTEM = (
+    "You write a one-paragraph summary of a codebase from its per-file summaries. "
+    "State what the system does and how the major pieces fit. No preamble."
+)
+
+
+@runtime_checkable
+class ClaudeClient(Protocol):
+    """One Messages call + cumulative cost. ``BedrockClient`` and
+    ``AnthropicClient`` both satisfy this."""
+
+    cost_usd: float
+
+    async def invoke(
+        self,
+        system: str,
+        user: str,
+        tools: list[dict[str, Any]] | None = None,
+        tool_name: str | None = None,
+    ) -> dict[str, Any]:
+        """Return the raw Messages payload (``content`` blocks + ``usage``)."""
+        ...
+
+
+class ClaudeJudge:
+    """Pattern judge over any ``ClaudeClient``. A forced ``submit_verdicts`` tool
+    call yields one structured verdict per nominated pattern."""
+
+    def __init__(self, client: ClaudeClient, model: str) -> None:
+        self.model = model
+        self._client = client
+
+    async def judge(self, candidate: Candidate) -> list[Verdict]:
+        if not candidate.patterns:
+            return []
+        payload = await self._client.invoke(
+            _JUDGE_SYSTEM,
+            self._prompt(candidate),
+            tools=[_VERDICT_TOOL],
+            tool_name="submit_verdicts",
+        )
+        nominated = set(candidate.patterns)
+        return [v for v in self._parse(payload) if v.pattern in nominated]
+
+    @property
+    def cost_usd(self) -> float:
+        return self._client.cost_usd
+
+    @staticmethod
+    def _prompt(c: Candidate) -> str:
+        methods = "\n".join(f"  - {n}: {sig}" for n, sig in c.methods[:30]) or "  (none)"
+        return (
+            f"{c.kind} `{c.name}`\nsignature: {c.signature}\nmethods:\n{methods}\n\n"
+            f"candidate patterns (from structural heuristics): {', '.join(c.patterns)}\n"
+            f"evidence: {'; '.join(c.evidence)}\n\n"
+            "Return a verdict for EACH candidate pattern."
+        )
+
+    @staticmethod
+    def _parse(payload: dict[str, Any]) -> list[Verdict]:
+        verdicts: list[Verdict] = []
+        for block in payload.get("content", []):
+            if block.get("type") == "tool_use":
+                for raw in block.get("input", {}).get("verdicts", []):
+                    try:
+                        conf = max(0.0, min(1.0, float(raw.get("confidence", 0.0))))
+                        verdicts.append(
+                            Verdict(
+                                pattern=str(raw.get("pattern", "")),
+                                is_match=bool(raw.get("is_match", False)),
+                                confidence=conf,
+                                rationale=str(raw.get("rationale", "")),
+                            )
+                        )
+                    except (TypeError, ValueError):
+                        continue
+        return verdicts
+
+
+class ClaudeSummarizer:
+    """Module/repo summaries over any ``ClaudeClient`` — plain-text completions."""
+
+    def __init__(self, client: ClaudeClient, model: str) -> None:
+        self.model = model
+        self._client = client
+
+    async def summarize_file(self, ctx: FileContext, max_words: int) -> Summary:
+        symbols = "\n".join(f"  - {n}: {sig}" for n, sig in ctx.symbols[:60]) or "  (no symbols)"
+        imports = ", ".join(ctx.imports[:20]) or "(none)"
+        user = (
+            f"File: {ctx.path}\nImports: {imports}\nSymbols:\n{symbols}\n\n"
+            f"Summarize this file in at most {max_words} words."
+        )
+        payload = await self._client.invoke(_FILE_SYSTEM, user)
+        return Summary(text=text_of(payload), model=self.model)
+
+    async def summarize_repo(
+        self, repo: str, file_summaries: list[tuple[str, str]], max_words: int
+    ) -> Summary:
+        joined = "\n".join(f"- {path}: {text}" for path, text in file_summaries[:200])
+        user = (
+            f"Repository: {repo}\nPer-file summaries:\n{joined}\n\n"
+            f"Summarize the whole codebase in at most {max_words} words."
+        )
+        payload = await self._client.invoke(_REPO_SYSTEM, user)
+        return Summary(text=text_of(payload), model=self.model)
+
+    @property
+    def cost_usd(self) -> float:
+        return self._client.cost_usd

agentforge_graph/enrich/enricher.py

@@ -0,0 +1,108 @@
+"""``PatternTagEnricher`` (feat-012) — orchestrate two-stage pattern tagging.
+
+Stage-1 heuristics nominate; the injected ``PatternJudge`` confirms each under a
+``budget_usd`` cap (the framework ``BudgetPolicy`` breaker — the first feature to
+ride the AgentForge budget rails). Confirmed verdicts above the confidence floor
+become ``PatternTag`` nodes + ``TAGGED`` edges with honest ``llm`` provenance.
+Re-tag is idempotent (clear a judged symbol's old ``TAGGED`` first); a tripped
+budget stops cleanly, leaving unjudged candidates for the next run. This is a
+framework-layer module (ADR-0001: ``enrich`` may import ``agentforge``).
+"""
+
+from __future__ import annotations
+
+import asyncio
+
+from agentforge_core.production.budget import BudgetPolicy
+from agentforge_core.production.exceptions import BudgetExceeded
+
+from agentforge_graph.core import Edge, GraphStore, Node, NodeKind, Provenance
+from agentforge_graph.core.kinds import EdgeKind
+
+from .heuristics import PatternHeuristics
+from .judge import PatternJudge
+from .report import EnrichReport
+from .taxonomy import is_pattern, pattern_tag_id
+
+
+class PatternTagEnricher:
+    version = "pattern-tags@1"  # bump on prompt/taxonomy change → re-tag
+
+    def __init__(
+        self,
+        repo: str,
+        judge: PatternJudge,
+        *,
+        heuristics: PatternHeuristics | None = None,
+        confidence_floor: float = 0.7,
+        budget_usd: float = 2.0,
+        concurrency: int = 6,
+        commit: str = "",
+    ) -> None:
+        self.repo = repo
+        self.judge = judge
+        self.heuristics = heuristics or PatternHeuristics()
+        self.confidence_floor = confidence_floor
+        self.budget_usd = budget_usd
+        self.concurrency = max(1, concurrency)
+        self.commit = commit
+        self.last_judged_ids: list[str] = []
+
+    async def enrich(self, store: GraphStore, symbol_ids: list[str]) -> EnrichReport:
+        report = EnrichReport()
+        candidates = await self.heuristics.nominate(store, symbol_ids)
+        report.candidates = len(candidates)
+        self.last_judged_ids = []
+        if not candidates:
+            return report
+
+        budget = BudgetPolicy(usd=self.budget_usd, max_tokens=10**12, max_iterations=10**12)
+        facts: list[Node | Edge] = []
+
+        # Judge in concurrent batches (ENH-002): cost is accounted per batch —
+        # `budget.check()`/`commit()` sit OUTSIDE the gather, so the shared judge
+        # cost is read atomically (no per-call race). Budget overrun is bounded
+        # to one batch; concurrency=1 reproduces the strict per-call breaker.
+        for start in range(0, len(candidates), self.concurrency):
+            batch = candidates[start : start + self.concurrency]
+            try:
+                budget.check()
+            except BudgetExceeded:
+                report.budget_tripped = True
+                break
+            before = self.judge.cost_usd
+            batch_verdicts = await asyncio.gather(*(self.judge.judge(c) for c in batch))
+            budget.commit(self.judge.cost_usd - before)
+            report.cost_usd = round(self.judge.cost_usd, 6)
+            for cand, verdicts in zip(batch, batch_verdicts, strict=True):
+                report.judged += 1
+                self.last_judged_ids.append(cand.symbol_id)
+                for v in verdicts:
+                    if not (
+                        v.is_match
+                        and v.confidence >= self.confidence_floor
+                        and is_pattern(v.pattern)
+                    ):
+                        continue
+                    prov = Provenance.llm(self.version, round(v.confidence, 4), self.commit)
+                    tag_id = pattern_tag_id(self.repo, v.pattern)
+                    facts.append(
+                        Node(id=tag_id, kind=NodeKind.PATTERN_TAG, name=v.pattern, provenance=prov)
+                    )
+                    facts.append(
+                        Edge(
+                            src=cand.symbol_id,
+                            dst=tag_id,
+                            kind=EdgeKind.TAGGED,
+                            attrs={"confidence": round(v.confidence, 4), "rationale": v.rationale},
+                            provenance=prov,
+                        )
+                    )
+                    report.tagged += 1
+                    report.by_pattern[v.pattern] = report.by_pattern.get(v.pattern, 0) + 1
+
+        # idempotent re-tag: drop judged symbols' old tags, then write the new
+        await store.clear_outgoing(self.last_judged_ids, EdgeKind.TAGGED)
+        if facts:
+            await store.add(facts)
+        return report

agentforge_graph/enrich/governs.py

@@ -0,0 +1,173 @@
+"""The ``infer_governs`` LLM matcher (feat-010 follow-up).
+
+When an ADR's prose does not name the code it governs by path/symbol, the
+deterministic mention parser produces zero ``GOVERNS`` edges. This optional pass
+asks a model to match a decision's text against the repo's candidate symbols and
+proposes ``GOVERNS`` edges with honest ``llm`` provenance + confidence.
+
+The matcher is injectable (the Embedder/PatternJudge pattern): the live
+``ClaudeGovernsMatcher`` runs over any ``ClaudeClient`` (Bedrock or the Anthropic
+API); the ``ScriptedMatcher`` keeps the enricher deterministic and credential-free
+for tests. This is a framework-layer module (ADR-0001: ``enrich`` may import
+``agentforge``); the deterministic ``knowledge`` package stays model-free.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any, Protocol, runtime_checkable
+
+from pydantic import BaseModel, Field
+
+from .claude import ClaudeClient
+
+
+class GovernsCandidate(BaseModel):
+    """A symbol a decision might govern — what the matcher sees per candidate."""
+
+    symbol_id: str
+    name: str
+    kind: str
+    signature: str = ""
+    path: str = ""
+
+
+class GovernsMatch(BaseModel):
+    """A proposed ``GOVERNS`` link from a decision to one candidate symbol."""
+
+    symbol_id: str
+    confidence: float = Field(ge=0.0, le=1.0)
+    rationale: str = ""
+
+
+@runtime_checkable
+class GovernsMatcher(Protocol):
+    async def match(
+        self, title: str, text: str, candidates: list[GovernsCandidate]
+    ) -> list[GovernsMatch]: ...
+
+    @property
+    def cost_usd(self) -> float:
+        """Cumulative USD spent so far (0 for the scripted matcher)."""
+        ...
+
+
+ScriptFn = Callable[[str, str, list[GovernsCandidate]], list[GovernsMatch]]
+
+
+class ScriptedMatcher:
+    """Deterministic matcher for tests. Drive it with a function; the default
+    matches nothing. An optional ``per_call_usd`` exercises the budget breaker."""
+
+    def __init__(self, fn: ScriptFn | None = None, per_call_usd: float = 0.0) -> None:
+        self._fn = fn or (lambda title, text, cands: [])
+        self._per_call_usd = per_call_usd
+        self._cost = 0.0
+
+    async def match(
+        self, title: str, text: str, candidates: list[GovernsCandidate]
+    ) -> list[GovernsMatch]:
+        self._cost += self._per_call_usd
+        return self._fn(title, text, candidates)
+
+    @property
+    def cost_usd(self) -> float:
+        return self._cost
+
+
+_GOVERNS_SYSTEM = (
+    "You match an architecture decision record (ADR) to the code symbols it governs. "
+    "A symbol is governed when the decision's rules plainly constrain how that symbol "
+    "is designed, implemented, or changed. Be conservative: propose a match ONLY when "
+    "the decision clearly applies to the symbol — prefer proposing nothing over "
+    "guessing. For each match give a confidence in [0,1] and a one-sentence rationale "
+    "citing the decision text."
+)
+
+_GOVERNS_TOOL = {
+    "name": "submit_governs",
+    "description": "Return the candidate symbols this decision governs (possibly none).",
+    "input_schema": {
+        "type": "object",
+        "properties": {
+            "matches": {
+                "type": "array",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "symbol_index": {"type": "integer"},
+                        "confidence": {"type": "number"},
+                        "rationale": {"type": "string"},
+                    },
+                    "required": ["symbol_index", "confidence", "rationale"],
+                },
+            }
+        },
+        "required": ["matches"],
+    },
+}
+
+
+class ClaudeGovernsMatcher:
+    """``infer_governs`` matcher over any ``ClaudeClient``. A forced
+    ``submit_governs`` tool call returns candidate indices + confidences."""
+
+    def __init__(self, client: ClaudeClient, model: str) -> None:
+        self.model = model
+        self._client = client
+
+    async def match(
+        self, title: str, text: str, candidates: list[GovernsCandidate]
+    ) -> list[GovernsMatch]:
+        if not candidates:
+            return []
+        payload = await self._client.invoke(
+            _GOVERNS_SYSTEM,
+            self._prompt(title, text, candidates),
+            tools=[_GOVERNS_TOOL],
+            tool_name="submit_governs",
+        )
+        return self._parse(payload, candidates)
+
+    @property
+    def cost_usd(self) -> float:
+        return self._client.cost_usd
+
+    @staticmethod
+    def _prompt(title: str, text: str, candidates: list[GovernsCandidate]) -> str:
+        listing = "\n".join(
+            f"  [{i}] {c.kind} `{c.name}` — {c.path}"
+            + (f" :: {c.signature}" if c.signature else "")
+            for i, c in enumerate(candidates)
+        )
+        return (
+            f"Decision: {title}\n\n{text.strip()}\n\n"
+            f"Candidate symbols (index in brackets):\n{listing}\n\n"
+            "Return the candidates this decision governs by their index. "
+            "If none clearly apply, return an empty list."
+        )
+
+    @staticmethod
+    def _parse(payload: dict[str, Any], candidates: list[GovernsCandidate]) -> list[GovernsMatch]:
+        matches: list[GovernsMatch] = []
+        seen: set[int] = set()
+        for block in payload.get("content", []):
+            if block.get("type") != "tool_use":
+                continue
+            for raw in block.get("input", {}).get("matches", []):
+                try:
+                    idx = int(raw.get("symbol_index", -1))
+                except (TypeError, ValueError):
+                    continue
+                if idx < 0 or idx >= len(candidates) or idx in seen:
+                    continue
+                seen.add(idx)
+                conf = max(0.0, min(1.0, float(raw.get("confidence", 0.0) or 0.0)))
+                matches.append(
+                    GovernsMatch(
+                        symbol_id=candidates[idx].symbol_id,
+                        confidence=conf,
+                        rationale=str(raw.get("rationale", "")),
+                    )
+                )
+        return matches