agentforge-graph 0.3.2__py3-none-any.whl

agentforge_graph/embed/fake.py

@@ -0,0 +1,34 @@
+"""Deterministic, dependency-free embedder for tests and CI — no creds, no
+network. Same text always yields the same L2-normalized vector, so retrieval
+tests are reproducible."""
+
+from __future__ import annotations
+
+import hashlib
+import math
+import struct
+
+from .base import Embedder, InputType
+
+
+class FakeEmbedder(Embedder):
+    def __init__(self, dim: int = 256) -> None:
+        self.name = "fake"
+        self.dim = dim
+
+    async def embed(
+        self, texts: list[str], input_type: InputType = "document"
+    ) -> list[list[float]]:
+        return [self._vector(t) for t in texts]
+
+    def _vector(self, text: str) -> list[float]:
+        buf = b""
+        counter = 0
+        need = self.dim * 4
+        while len(buf) < need:
+            buf += hashlib.sha256(text.encode("utf-8") + counter.to_bytes(4, "big")).digest()
+            counter += 1
+        words = struct.unpack(f">{self.dim}I", buf[:need])
+        vals = [(w / 2**32) * 2.0 - 1.0 for w in words]  # finite, in [-1, 1)
+        norm = math.sqrt(sum(v * v for v in vals)) or 1.0
+        return [v / norm for v in vals]

agentforge_graph/embed/openai.py

@@ -0,0 +1,67 @@
+"""``OpenAIEmbedder`` — OpenAI (and OpenAI-compatible) embeddings (ENH-003
+phase 2; the most-requested non-AWS path, and the **local model** path).
+
+Lazy-imports the ``openai`` SDK (the ``openai`` extra); synchronous calls run on
+a worker thread, mirroring ``BedrockEmbedder``. Setting ``embed.base_url`` points
+the same adapter at any OpenAI-compatible server — a local Ollama
+(``http://localhost:11434/v1``), vLLM, LM Studio, or a gateway — so "bring your
+own / run it locally" is a config line, not a new adapter.
+
+``text-embedding-3-*`` models support arbitrary output ``dimensions``; ``dim``
+is passed through. Credentials come from ``OPENAI_API_KEY`` (the SDK default)
+unless ``api_key_env`` overrides it. Imports nothing from ``agentforge``
+(ADR-0001).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from typing import Any
+
+from .base import Embedder, InputType
+
+
+class OpenAIEmbedder(Embedder):
+    def __init__(
+        self,
+        model: str = "text-embedding-3-small",
+        dim: int = 1536,
+        batch_size: int = 96,
+        base_url: str = "",
+        api_key_env: str = "OPENAI_API_KEY",
+    ) -> None:
+        self.name = f"openai:{model}"
+        self.model = model
+        self.dim = dim
+        self.batch_size = batch_size
+        self.base_url = base_url
+        self.api_key_env = api_key_env
+        self._client: Any = None
+
+    def _openai(self) -> Any:
+        if self._client is None:
+            import openai
+
+            kwargs: dict[str, Any] = {}
+            key = os.environ.get(self.api_key_env)
+            if key:
+                kwargs["api_key"] = key
+            if self.base_url:
+                kwargs["base_url"] = self.base_url
+            self._client = openai.OpenAI(**kwargs)
+        return self._client
+
+    async def embed(
+        self, texts: list[str], input_type: InputType = "document"
+    ) -> list[list[float]]:
+        # OpenAI embeddings are symmetric — ``input_type`` is ignored.
+        out: list[list[float]] = []
+        for i in range(0, len(texts), self.batch_size):
+            batch = texts[i : i + self.batch_size]
+            out.extend(await asyncio.to_thread(self._invoke, batch))
+        return out
+
+    def _invoke(self, batch: list[str]) -> list[list[float]]:
+        resp = self._openai().embeddings.create(model=self.model, input=batch, dimensions=self.dim)
+        return [[float(x) for x in item.embedding] for item in resp.data]

agentforge_graph/embed/pipeline.py

@@ -0,0 +1,184 @@
+"""``EmbedPipeline`` — chunk the indexed code and embed the chunks.
+
+Per file: pull its symbol nodes from the graph, chunk them, write `CHUNK`
+nodes + `CHUNK_OF` edges, embed the chunk texts, and upsert vectors. Coarse
+incrementality at 0.1: if a file's chunk-hash set is unchanged, skip
+re-embedding (saves cost); otherwise clean-replace the file's chunk vectors.
+feat-004 will scope this to a DirtySet.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+
+from agentforge_graph.chunking import Chunker
+from agentforge_graph.core import (
+    Edge,
+    EdgeKind,
+    Embedded,
+    GraphQuery,
+    Node,
+    NodeKind,
+    Provenance,
+    SymbolID,
+)
+from agentforge_graph.ingest import PackRegistry, RepoSource
+from agentforge_graph.store import Store
+
+from .base import Embedder
+from .report import EmbedReport
+
+_ALL = 10_000_000
+
+
+class EmbedPipeline:
+    def __init__(self, chunker: Chunker, embedder: Embedder, commit: str = "") -> None:
+        self.chunker = chunker
+        self.embedder = embedder
+        self.commit = commit
+        self.name = "cast-chunker"
+
+    async def run(
+        self,
+        store: Store,
+        source: RepoSource,
+        registry: PackRegistry,
+        only_paths: set[str] | None = None,
+        doc_root: Path | None = None,
+    ) -> EmbedReport:
+        """Embed the indexed code. When ``only_paths`` is given (feat-004: the
+        files a refresh dirtied), only those files are re-chunked/embedded;
+        otherwise every file is visited (the chunk-hash skip still avoids
+        re-embedding unchanged files)."""
+        report = EmbedReport(model=self.embedder.name, dim=self.embedder.dim)
+        prov = Provenance.parsed(self.name, self.commit)
+
+        for sf in source.iter_files(registry):
+            if only_paths is not None and sf.path not in only_paths:
+                continue
+            nodes_for_path = [
+                n
+                for n in (
+                    await store.graph.query(GraphQuery(path_prefix=sf.path, limit=_ALL))
+                ).nodes
+                if SymbolID.parse(n.id).path == sf.path
+            ]
+            symbols = [n for n in nodes_for_path if n.kind is not NodeKind.CHUNK]
+            if not symbols:
+                continue
+            chunks = self.chunker.chunk(sf, symbols)
+            if not chunks:
+                continue
+            report.files += 1
+            report.chunks += len(chunks)
+
+            prior = {
+                n.attrs.get("content_hash") for n in nodes_for_path if n.kind is NodeKind.CHUNK
+            }
+            if prior and prior == {c.content_hash for c in chunks}:
+                report.skipped_unchanged += 1
+                continue
+
+            repo = SymbolID.parse(symbols[0].id).repo
+            file_id = SymbolID.for_symbol(sf.language, repo, sf.path, "")
+            graph_items: list[Node | Edge] = []
+            for ch in chunks:
+                graph_items.append(
+                    Node(
+                        id=ch.id,
+                        kind=NodeKind.CHUNK,
+                        name=f"chunk{ch.seq}",
+                        span=ch.span,
+                        attrs={
+                            "path": ch.path,
+                            "token_count": ch.token_count,
+                            "content_hash": ch.content_hash,
+                            "seq": ch.seq,
+                            "code": ch.code,  # carried for retrieval rendering (feat-006)
+                        },
+                        provenance=prov,
+                    )
+                )
+                for target in ch.symbol_ids or [file_id]:
+                    graph_items.append(
+                        Edge(src=ch.id, dst=target, kind=EdgeKind.CHUNK_OF, provenance=prov)
+                    )
+            await store.graph.add(graph_items)
+
+            await store.vectors.delete_where({"path": sf.path})  # clean-replace this file
+            vectors = await self.embedder.embed([c.text for c in chunks], input_type="document")
+            await store.vectors.upsert(
+                [
+                    Embedded(
+                        ref=ch.id,
+                        vector=vec,
+                        kind=NodeKind.CHUNK,
+                        attrs={
+                            "path": ch.path,
+                            "span": list(ch.span),
+                            "symbol_ids": ch.symbol_ids,
+                            "source_type": "code",  # vs "doc" (feat-010) — lets
+                            "model": self.embedder.name,  # retrieval tell them apart
+                        },
+                    )
+                    for ch, vec in zip(chunks, vectors, strict=True)
+                ]
+            )
+            report.embedded += len(chunks)
+
+        report.doc_chunks = await self._embed_docs(store, doc_root)
+        return report
+
+    async def _embed_docs(self, store: Store, doc_root: Path | None = None) -> int:
+        """Embed ADR/doc ``DocChunk`` prose so an architectural query surfaces the
+        governing decision / documented symbol (feat-010). A ``source_type: doc``
+        tag keeps these distinct from code chunks. Incremental: a fingerprint of all
+        doc chunks (ids + content hashes + embedder) is recorded under ``doc_root``;
+        when it is unchanged the whole pass is skipped (no API calls). On any change
+        it clean-replaces every doc vector (the simple, orphan-safe path for the
+        small doc set)."""
+        docs = (await store.graph.query(GraphQuery(kinds=[NodeKind.DOC_CHUNK], limit=_ALL))).nodes
+        manifest = (doc_root / "doc_embed.hash") if doc_root is not None else None
+        if not docs:
+            await store.vectors.delete_where({"kind": NodeKind.DOC_CHUNK.value})
+            if manifest is not None and manifest.exists():
+                manifest.unlink()
+            return 0
+        fp_body = "".join(
+            f"{n.id}|{n.attrs.get('content_hash', '')};" for n in sorted(docs, key=lambda z: z.id)
+        )
+        fingerprint = hashlib.sha256(
+            f"{self.embedder.name}:{self.embedder.dim}:{fp_body}".encode()
+        ).hexdigest()
+        if (
+            manifest is not None
+            and manifest.exists()
+            and manifest.read_text().strip() == fingerprint
+        ):
+            return 0  # docs unchanged since the last embed → skip the re-embed
+        # clean-replace via the DocChunk kind (a filterable vector column) — this
+        # also GCs vectors for docs/ADRs that were removed since the last embed.
+        await store.vectors.delete_where({"kind": NodeKind.DOC_CHUNK.value})
+        texts = [f"{n.attrs.get('heading', '')}\n{n.attrs.get('text', '')}".strip() for n in docs]
+        vectors = await self.embedder.embed(texts, input_type="document")
+        await store.vectors.upsert(
+            [
+                Embedded(
+                    ref=n.id,
+                    vector=vec,
+                    kind=NodeKind.DOC_CHUNK,
+                    attrs={
+                        "path": n.attrs.get("path", ""),
+                        "source_type": "doc",
+                        "heading": n.attrs.get("heading", ""),
+                        "model": self.embedder.name,
+                    },
+                )
+                for n, vec in zip(docs, vectors, strict=True)
+            ]
+        )
+        if manifest is not None:
+            manifest.parent.mkdir(parents=True, exist_ok=True)
+            manifest.write_text(fingerprint)
+        return len(docs)

agentforge_graph/embed/registry.py

@@ -0,0 +1,66 @@
+"""Resolve an ``Embedder`` from ``EmbedConfig`` via the provider registry.
+
+Built-ins (``fake``, ``bedrock``, ``openai``) are registered below; third-party
+embedders register out-of-tree under the ``agentforge_graph.embedder_providers``
+entry-point group (``pip install`` + one ``embed.driver`` line, no core change).
+Each live driver lazy-imports its SDK so the base/fake path needs neither boto3
+nor openai. ``openai`` also covers OpenAI-compatible local servers via
+``embed.base_url`` (ENH-003 phase 2).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from agentforge_graph.config import EmbedConfig
+from agentforge_graph.providers import resolve_provider
+
+from .base import Embedder
+
+EMBEDDER_GROUP = "agentforge_graph.embedder_providers"
+
+# A builder takes the parsed ``embed:`` block and returns a ready Embedder.
+EmbedderBuilder = Callable[[EmbedConfig], Embedder]
+
+
+def _build_fake(cfg: EmbedConfig) -> Embedder:
+    from .fake import FakeEmbedder
+
+    return FakeEmbedder(dim=cfg.dim)
+
+
+def _build_bedrock(cfg: EmbedConfig) -> Embedder:
+    from .bedrock import BedrockEmbedder  # lazy: only needs boto3 on this path
+
+    return BedrockEmbedder(
+        model=cfg.model,
+        region=cfg.region,
+        dim=cfg.dim,
+        batch_size=cfg.batch_size,
+        assume_role_arn=cfg.assume_role_arn or None,
+    )
+
+
+def _build_openai(cfg: EmbedConfig) -> Embedder:
+    from .openai import OpenAIEmbedder  # lazy: only needs the openai SDK on this path
+
+    return OpenAIEmbedder(
+        model=cfg.model,
+        dim=cfg.dim,
+        batch_size=cfg.batch_size,
+        base_url=cfg.base_url,
+        api_key_env=cfg.api_key_env or "OPENAI_API_KEY",
+    )
+
+
+_EMBEDDER_BUILTINS: dict[str, EmbedderBuilder] = {
+    "fake": _build_fake,
+    "bedrock": _build_bedrock,
+    "openai": _build_openai,
+}
+
+
+def embedder_from_config(cfg: EmbedConfig) -> Embedder:
+    """Construct the ``Embedder`` selected by ``cfg.driver`` via the registry."""
+    builder = resolve_provider(cfg.driver, _EMBEDDER_BUILTINS, EMBEDDER_GROUP, role="embedder")
+    return builder(cfg)

agentforge_graph/embed/report.py

@@ -0,0 +1,15 @@
+"""Result type for an embedding run."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+
+class EmbedReport(BaseModel):
+    files: int = 0
+    chunks: int = 0
+    embedded: int = 0
+    skipped_unchanged: int = 0  # files whose chunk set was unchanged (hash-skip)
+    doc_chunks: int = 0  # ADR/doc DocChunks embedded for semantic search (feat-010)
+    model: str = ""
+    dim: int = 0

agentforge_graph/enrich/__init__.py

@@ -0,0 +1,70 @@
+"""LLM enrichment (feat-012): turn the code graph into a knowledge graph.
+
+MVP: **design-pattern tagging** — deterministic structural heuristics nominate
+candidates, a budgeted LLM judge confirms them, and confirmed verdicts become
+``TAGGED`` edges to a fixed v1 ``PatternTag`` taxonomy with honest ``llm``
+provenance + confidence + rationale. The judge is injectable
+(``ScriptedJudge`` for tests, ``BedrockClaudeJudge`` live), so all orchestration
+is deterministic. This is the framework layer (ADR-0001: ``enrich`` may import
+``agentforge``); never runs implicitly (``ckg enrich`` only).
+"""
+
+from __future__ import annotations
+
+from .enricher import PatternTagEnricher
+from .governs import (
+    ClaudeGovernsMatcher,
+    GovernsCandidate,
+    GovernsMatch,
+    GovernsMatcher,
+    ScriptedMatcher,
+)
+from .governs_enricher import DecisionGovernsInferencer
+from .heuristics import Candidate, PatternHeuristics
+from .judge import PatternJudge, ScriptedJudge, Verdict
+from .registry import (
+    JUDGE_GROUP,
+    SUMMARIZER_GROUP,
+    governs_matcher_from_config,
+    judge_from_config,
+    summarizer_from_config,
+)
+from .report import EnrichReport, GovernsReport, SummaryInfo, SummaryReport, TaggedInfo
+from .summarizer import FileContext, ScriptedSummarizer, Summarizer, Summary
+from .summary_enricher import SummaryEnricher, repo_node_id, summary_id
+from .taxonomy import TAXONOMY_V1, is_pattern, pattern_tag_id
+
+__all__ = [
+    "PatternTagEnricher",
+    "PatternHeuristics",
+    "Candidate",
+    "PatternJudge",
+    "ScriptedJudge",
+    "Verdict",
+    "judge_from_config",
+    "summarizer_from_config",
+    "governs_matcher_from_config",
+    "DecisionGovernsInferencer",
+    "GovernsMatcher",
+    "ScriptedMatcher",
+    "ClaudeGovernsMatcher",
+    "GovernsCandidate",
+    "GovernsMatch",
+    "GovernsReport",
+    "JUDGE_GROUP",
+    "SUMMARIZER_GROUP",
+    "EnrichReport",
+    "TaggedInfo",
+    "SummaryReport",
+    "SummaryInfo",
+    "SummaryEnricher",
+    "Summarizer",
+    "ScriptedSummarizer",
+    "Summary",
+    "FileContext",
+    "summary_id",
+    "repo_node_id",
+    "TAXONOMY_V1",
+    "is_pattern",
+    "pattern_tag_id",
+]

agentforge_graph/enrich/anthropic.py

@@ -0,0 +1,38 @@
+"""``AnthropicClaudeJudge`` / ``AnthropicClaudeSummarizer`` — the **direct
+Anthropic API** enrichers (ENH-003 phase 2; the non-AWS Claude path).
+
+Thin endpoint adapters: the judging/summary logic is the provider-neutral
+``ClaudeJudge`` / ``ClaudeSummarizer`` (``claude.py``); these wire it to an
+``AnthropicClient`` transport. Pick them with ``enrich.provider: anthropic`` and
+set ``ANTHROPIC_API_KEY`` (``enrich.model`` may stay the Bedrock default — the
+id is normalised to its API form). Tests use the ``Scripted*`` variants.
+"""
+
+from __future__ import annotations
+
+from .anthropic_client import AnthropicClient
+from .claude import ClaudeJudge, ClaudeSummarizer
+
+
+class AnthropicClaudeJudge(ClaudeJudge):
+    def __init__(
+        self,
+        model: str = "claude-haiku-4-5-20251001",
+        api_key_env: str = "ANTHROPIC_API_KEY",
+        base_url: str = "",
+        max_tokens: int = 512,
+    ) -> None:
+        client = AnthropicClient(model, api_key_env, base_url, max_tokens)
+        super().__init__(client, client.model)
+
+
+class AnthropicClaudeSummarizer(ClaudeSummarizer):
+    def __init__(
+        self,
+        model: str = "claude-haiku-4-5-20251001",
+        api_key_env: str = "ANTHROPIC_API_KEY",
+        base_url: str = "",
+        max_tokens: int = 400,
+    ) -> None:
+        client = AnthropicClient(model, api_key_env, base_url, max_tokens)
+        super().__init__(client, client.model)

agentforge_graph/enrich/anthropic_client.py

@@ -0,0 +1,109 @@
+"""Shared **direct Anthropic API** Claude client for the enrichers (ENH-003
+phase 2 — the non-AWS path).
+
+Mirrors ``BedrockClient``: one lazily-built ``anthropic.Anthropic`` client, a
+synchronous Messages call run on a worker thread, cost accumulated from token
+usage. Returns ``resp.model_dump()`` — the same ``content``/``usage`` dict shape
+Bedrock returns — so ``ClaudeJudge`` / ``ClaudeSummarizer`` parse it unchanged.
+
+The ``anthropic`` SDK ships with the base install (pulled by
+``agentforge-anthropic[anthropic]``); it is imported lazily so the
+scripted/offline path never needs it. Credentials come from ``ANTHROPIC_API_KEY``
+(the SDK's default env var) unless overridden.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from typing import Any
+
+from .bedrock_client import price_for  # provider-neutral USD-per-token table
+
+# Bedrock ids carry an inference-profile prefix (``us.``/``eu.``/…) and a
+# ``-v1:0`` suffix; the direct API wants the bare model id. Normalising lets the
+# same ``enrich.model`` default work on either provider.
+_PROFILE_PREFIXES = ("us.", "eu.", "apac.", "us-gov.")
+
+
+def api_model_id(model: str) -> str:
+    """Map a Bedrock model id to its Anthropic-API equivalent (idempotent on an
+    id that is already an API id). ``us.anthropic.claude-haiku-4-5-20251001-v1:0``
+    → ``claude-haiku-4-5-20251001``."""
+    m = model
+    for prefix in _PROFILE_PREFIXES:
+        if m.startswith(prefix):
+            m = m[len(prefix) :]
+            break
+    if m.startswith("anthropic."):
+        m = m[len("anthropic.") :]
+    if m.endswith("-v1:0"):
+        m = m[: -len("-v1:0")]
+    return m
+
+
+class AnthropicClient:
+    def __init__(
+        self,
+        model: str = "claude-haiku-4-5-20251001",
+        api_key_env: str = "ANTHROPIC_API_KEY",
+        base_url: str = "",
+        max_tokens: int = 512,
+    ) -> None:
+        self.model = api_model_id(model)
+        self.api_key_env = api_key_env
+        self.base_url = base_url
+        self.max_tokens = max_tokens
+        self._client: Any = None
+        self.cost_usd = 0.0
+
+    def _anthropic(self) -> Any:
+        if self._client is None:
+            import anthropic
+
+            kwargs: dict[str, Any] = {}
+            key = os.environ.get(self.api_key_env)
+            if key:
+                kwargs["api_key"] = key
+            if self.base_url:
+                kwargs["base_url"] = self.base_url
+            self._client = anthropic.Anthropic(**kwargs)
+        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]:
+        kwargs: dict[str, Any] = {
+            "model": self.model,
+            "max_tokens": self.max_tokens,
+            "system": system,
+            "messages": [{"role": "user", "content": user}],
+        }
+        if tools is not None:
+            kwargs["tools"] = tools
+            if tool_name is not None:
+                kwargs["tool_choice"] = {"type": "tool", "name": tool_name}
+        resp = self._anthropic().messages.create(**kwargs)
+        result: dict[str, Any] = resp.model_dump()
+        return result

agentforge_graph/enrich/bedrock.py

@@ -0,0 +1,24 @@
+"""``BedrockClaudeJudge`` — the pattern judge on AWS Bedrock (feat-012).
+
+Thin endpoint adapter: the judging logic (prompts, forced ``submit_verdicts``
+tool call, verdict parsing, cost) lives in the provider-neutral ``ClaudeJudge``
+(``claude.py``); this just wires it to a Bedrock transport (``BedrockClient``).
+The Anthropic-API sibling is ``AnthropicClaudeJudge`` (``anthropic.py``). Tests
+use the ``ScriptedJudge`` instead.
+"""
+
+from __future__ import annotations
+
+from .bedrock_client import BedrockClient
+from .claude import ClaudeJudge
+
+
+class BedrockClaudeJudge(ClaudeJudge):
+    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:
+        super().__init__(BedrockClient(model, region, assume_role_arn, max_tokens), model)