agentforge-graph 0.3.2__py3-none-any.whl

agentforge_graph/core/contracts.py

@@ -0,0 +1,163 @@
+"""The locked ABCs every later feature plugs into.
+
+- ``Extractor`` — turns a file into a ``FileSubgraph`` (feat-002, feat-011).
+- ``GraphStore`` — persists subgraphs and enrichment facts, answers
+  queries and neighborhood walks (feat-003 adapters).
+- ``Enricher`` — derives new nodes/edges from the existing graph
+  (feat-010/012).
+
+Signatures only; implementations ship with their owning features. The
+constructor/method surface here is the stable contract — additions are
+minor bumps, removals/renames are major. See ADR-0001 (layering: this
+module imports nothing from ``agentforge``).
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Literal
+
+from .kinds import EdgeKind
+from .models import (
+    Edge,
+    Embedded,
+    FileSubgraph,
+    GraphQuery,
+    Node,
+    QueryResult,
+    ScoredRef,
+    SourceFile,
+)
+
+# Direction of a 1-hop edge walk: out = node is src, in = node is dst.
+Direction = Literal["out", "in", "both"]
+
+
+class Extractor(ABC):
+    """Produces a ``FileSubgraph`` from a single file, in isolation.
+
+    Extraction must not read other files (per-file isolation is what
+    makes feat-004 incremental); cross-file edges are emitted as
+    candidate references and resolved in a later pass.
+    """
+
+    name: str
+
+    @abstractmethod
+    def extract(self, file: SourceFile) -> FileSubgraph: ...
+
+
+class GraphStore(ABC):
+    """Persistence + query contract. feat-003 ships the adapters."""
+
+    @abstractmethod
+    async def upsert(self, subgraph: FileSubgraph) -> None:
+        """Insert/replace all nodes & edges for ``subgraph.path``
+        transactionally (delete prior content for that path, add new)."""
+
+    @abstractmethod
+    async def add(self, items: list[Node | Edge]) -> None:
+        """Persist facts not tied to a single file (enrichment, resolved
+        cross-file edges). These survive ``delete_file`` of code files."""
+
+    @abstractmethod
+    async def delete_file(self, path: str) -> None:
+        """Remove everything previously upserted for ``path``."""
+
+    @abstractmethod
+    async def clear_resolved(self, paths: list[str]) -> None:
+        """Delete resolved-provenance edges whose ``origin_path`` is in
+        ``paths`` — the inverse of a scoped re-resolve (feat-004). Parsed
+        nodes/edges are untouched. Also garbage-collects external ``PACKAGE``
+        nodes left with no inbound edge, so an incremental re-resolve converges
+        to the same graph a full re-index would produce."""
+
+    @abstractmethod
+    async def clear_outgoing(self, src_ids: list[str], kind: EdgeKind) -> None:
+        """Delete edges of ``kind`` whose ``src`` is in ``src_ids`` — lets an
+        enricher (feat-012) re-derive a symbol's facts idempotently (re-tag
+        without duplicating ``TAGGED``/``SUMMARIZES`` edges)."""
+
+    @abstractmethod
+    async def query(self, q: GraphQuery) -> QueryResult:
+        """Exact-match node lookup with the flat ``GraphQuery`` filter."""
+
+    @abstractmethod
+    async def neighbors(
+        self,
+        node_id: str,
+        kinds: list[EdgeKind] | None = None,
+        depth: int = 1,
+    ) -> list[Node]:
+        """Nodes reachable from ``node_id`` over edges of ``kinds`` within
+        ``depth`` hops (either direction)."""
+
+    @abstractmethod
+    async def get(self, node_id: str) -> Node | None:
+        """Fetch a node by id, or ``None``."""
+
+    @abstractmethod
+    async def set_attrs(self, node_id: str, attrs: dict[str, Any]) -> None:
+        """Merge ``attrs`` into an existing node's ``attrs`` (a partial update —
+        other fields, including the file-ownership ``origin_path`` that drives
+        ``delete_file``, are untouched). No-op if the node is absent. The
+        denormalisation channel for derived facts (feat-009 churn/authorship)
+        that must not detach a file-owned node from its file."""
+
+    @abstractmethod
+    async def adjacent(
+        self,
+        node_id: str,
+        kinds: list[EdgeKind] | None = None,
+        direction: Direction = "both",
+    ) -> list[Edge]:
+        """The 1-hop edges touching ``node_id`` (``out``: it is the src;
+        ``in``: it is the dst; ``both``), optionally filtered by edge kind.
+        Returns full ``Edge`` objects, so the caller sees each edge's kind,
+        direction and provenance (feat-006 retrieval scoring)."""
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Release resources. Safe to call more than once."""
+
+
+class VectorStore(ABC):
+    """Vector persistence + similarity search. feat-003 ships the LanceDB
+    adapter; feat-005 produces the ``Embedded`` items it stores. A peer of
+    ``GraphStore`` — the ``Store`` facade (feat-003) owns one of each and
+    joins them (vector hit -> graph expansion) for retrieval (feat-006)."""
+
+    @abstractmethod
+    async def upsert(self, items: list[Embedded]) -> None:
+        """Insert/replace vectors keyed by ``Embedded.ref``."""
+
+    @abstractmethod
+    async def search(
+        self,
+        vector: list[float],
+        k: int,
+        filter: dict[str, Any] | None = None,
+    ) -> list[ScoredRef]:
+        """Top-``k`` nearest refs, optionally constrained by an attribute
+        ``filter`` (e.g. ``{"kind": "Chunk"}``)."""
+
+    @abstractmethod
+    async def delete_where(self, filter: dict[str, Any]) -> None:
+        """Drop vectors matching ``filter`` (feat-004 invalidation)."""
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Release resources. Safe to call more than once."""
+
+
+class Enricher(ABC):
+    """Derives new nodes/edges from the existing graph (feat-010/012).
+
+    Returns the facts it derived; the caller persists them via
+    ``GraphStore.add``. Derived facts must carry ``source=llm``
+    provenance with a confidence (ADR-0004)."""
+
+    name: str
+
+    @abstractmethod
+    async def enrich(self, store: GraphStore) -> list[Node | Edge]: ...

agentforge_graph/core/kinds.py

@@ -0,0 +1,68 @@
+"""Node and edge kind vocabularies for the code knowledge graph.
+
+The full vocabulary is locked at 0.1 — including the higher-level kinds
+whose producers ship later (feat-010/011/012) — so stores and queries
+handle every kind from day one and no schema migration is needed when a
+later producer lands. See ADR-0005.
+"""
+
+from __future__ import annotations
+
+from enum import StrEnum
+
+
+class NodeKind(StrEnum):
+    """Every node kind the graph may contain. Locked at 0.1 (ADR-0005)."""
+
+    # --- structural (produced by feat-002) ---
+    REPOSITORY = "Repository"
+    PACKAGE = "Package"
+    FILE = "File"
+    CLASS = "Class"
+    INTERFACE = "Interface"
+    FUNCTION = "Function"
+    METHOD = "Method"
+    VARIABLE = "Variable"
+    TYPE_ALIAS = "TypeAlias"
+
+    # --- retrieval (feat-005 / feat-010) ---
+    CHUNK = "Chunk"
+    DOC_CHUNK = "DocChunk"
+
+    # --- higher-level: reserved now, produced later (ADR-0005) ---
+    DECISION = "Decision"  # feat-010
+    ROUTE = "Route"  # feat-011
+    DATA_MODEL = "DataModel"  # feat-011
+    SERVICE = "Service"  # feat-011
+    SUMMARY = "Summary"  # feat-012
+    PATTERN_TAG = "PatternTag"  # feat-012
+
+
+class EdgeKind(StrEnum):
+    """Every edge kind the graph may contain. Locked at 0.1 (ADR-0005)."""
+
+    # --- structural (feat-002) ---
+    CONTAINS = "CONTAINS"
+    IMPORTS = "IMPORTS"
+    CALLS = "CALLS"
+    INHERITS = "INHERITS"
+    IMPLEMENTS = "IMPLEMENTS"
+    REFERENCES = "REFERENCES"
+
+    # --- retrieval / docs (feat-005 / feat-010) ---
+    CHUNK_OF = "CHUNK_OF"
+    DESCRIBES = "DESCRIBES"
+
+    # --- decisions (feat-010) ---
+    GOVERNS = "GOVERNS"
+    SUPERSEDES = "SUPERSEDES"
+
+    # --- framework (feat-011) ---
+    HANDLED_BY = "HANDLED_BY"
+    INJECTED_INTO = "INJECTED_INTO"
+    HAS_FIELD = "HAS_FIELD"
+    RELATES_TO = "RELATES_TO"
+
+    # --- enrichment (feat-012) ---
+    SUMMARIZES = "SUMMARIZES"
+    TAGGED = "TAGGED"

agentforge_graph/core/models.py

@@ -0,0 +1,134 @@
+"""Value types for the graph: nodes, edges, the per-file subgraph, and
+the minimal query/result shapes.
+
+``Node``/``Edge`` validate their IDs and (via ``Provenance``) their
+attribution at construction, so the graph cannot hold a malformed or
+unattributed fact. ``FileSubgraph`` is the unit of ingestion *and*
+deletion — keyed by ``(path, content_hash)`` — which is what makes
+incremental indexing (feat-004) a thin layer. See ADR-0003/0004.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+from .kinds import EdgeKind, NodeKind
+from .provenance import Provenance, Source
+from .symbols import SymbolID
+
+
+def _require_symbol_id(value: str) -> str:
+    SymbolID.parse(value)  # raises ValueError if malformed
+    return value
+
+
+class SourceFile(BaseModel):
+    """A single file handed to an ``Extractor``."""
+
+    model_config = ConfigDict(frozen=True)
+
+    path: str  # repo-relative, posix
+    text: str
+    language: str
+    content_hash: str  # sha256 of the file bytes
+
+
+class Node(BaseModel):
+    """A typed entity in the graph."""
+
+    id: str  # a SymbolID string
+    kind: NodeKind
+    name: str
+    span: tuple[int, int] | None = None  # (start_line, end_line), 1-based
+    attrs: dict[str, Any] = Field(default_factory=dict)
+    provenance: Provenance
+
+    @field_validator("id")
+    @classmethod
+    def _check_id(cls, v: str) -> str:
+        return _require_symbol_id(v)
+
+
+class Edge(BaseModel):
+    """A typed relationship between two symbols."""
+
+    src: str
+    dst: str
+    kind: EdgeKind
+    attrs: dict[str, Any] = Field(default_factory=dict)
+    provenance: Provenance
+    # The file whose content produced this edge (the import/call site's file).
+    # Lets resolver-produced edges be invalidated per file on an incremental
+    # re-resolve (feat-004). Empty for file-parsed edges, where the store
+    # stamps the owning file's path at upsert time.
+    origin_path: str = ""
+
+    @field_validator("src", "dst")
+    @classmethod
+    def _check_endpoint(cls, v: str) -> str:
+        return _require_symbol_id(v)
+
+
+class FileSubgraph(BaseModel):
+    """Everything extracted from one file — the ingestion/deletion unit."""
+
+    path: str  # repo-relative, posix
+    content_hash: str
+    nodes: list[Node] = Field(default_factory=list)
+    edges: list[Edge] = Field(default_factory=list)
+
+
+class GraphQuery(BaseModel):
+    """A minimal, flat node filter (0.1). Graph traversal lives in
+    ``GraphStore.neighbors``, not here. Extends by minor bump."""
+
+    kinds: list[NodeKind] | None = None
+    name: str | None = None  # exact match
+    path_prefix: str | None = None
+    edge_kind: EdgeKind | None = None
+    min_source: Source | None = None  # provenance floor
+    limit: int = 100
+
+    @field_validator("limit")
+    @classmethod
+    def _positive_limit(cls, v: int) -> int:
+        if v <= 0:
+            raise ValueError("limit must be > 0")
+        return v
+
+
+class QueryResult(BaseModel):
+    """The result of a ``GraphStore.query``."""
+
+    nodes: list[Node] = Field(default_factory=list)
+    edges: list[Edge] = Field(default_factory=list)
+    truncated: bool = False  # True if `limit` clipped the result
+
+
+class Embedded(BaseModel):
+    """A vector plus the symbol/chunk it represents — the ``VectorStore``
+    write unit. ``ref`` is the id of the node the vector stands in for
+    (a Chunk, DocChunk, Summary…); the producer is feat-005."""
+
+    ref: str  # symbol/chunk id this vector represents
+    vector: list[float]
+    kind: NodeKind
+    attrs: dict[str, Any] = Field(default_factory=dict)
+
+    @field_validator("vector")
+    @classmethod
+    def _non_empty(cls, v: list[float]) -> list[float]:
+        if not v:
+            raise ValueError("vector must be non-empty")
+        return v
+
+
+class ScoredRef(BaseModel):
+    """A vector-search hit: a ref and its similarity score (higher =
+    closer). feat-006 expands these into a graph neighborhood."""
+
+    ref: str
+    score: float
+    attrs: dict[str, Any] = Field(default_factory=dict)

agentforge_graph/core/provenance.py

@@ -0,0 +1,62 @@
+"""Provenance — attribution carried by every node and edge.
+
+Every fact in the graph records where it came from, so an agent can
+tell parsed ground truth from a heuristic resolution or an LLM guess,
+and rank/filter accordingly. Enforced at construction so no
+unattributed fact can exist. See ADR-0004.
+"""
+
+from __future__ import annotations
+
+from enum import StrEnum
+
+from pydantic import BaseModel, ConfigDict, model_validator
+
+
+class Source(StrEnum):
+    """How a fact was derived, in decreasing reliability."""
+
+    PARSED = "parsed"  # straight from the syntax tree
+    RESOLVED = "resolved"  # upgraded via the import graph / resolver
+    LLM = "llm"  # generated by a model (summaries, pattern tags, inferred links)
+    MANUAL = "manual"  # human-asserted
+
+
+class Provenance(BaseModel):
+    """Attribution for a single node or edge."""
+
+    model_config = ConfigDict(frozen=True)
+
+    source: Source
+    extractor: str  # producer name + version, e.g. "tree-sitter-python@0.23"
+    commit: str = ""  # git sha the fact was derived at; "" if non-git / unstaged
+    confidence: float = 1.0  # < 1.0 only meaningful for source=llm
+
+    @model_validator(mode="after")
+    def _check_confidence(self) -> Provenance:
+        if not 0.0 <= self.confidence <= 1.0:
+            raise ValueError("confidence must be in [0.0, 1.0]")
+        if self.source is not Source.LLM and self.confidence != 1.0:
+            raise ValueError("confidence < 1.0 is only valid for source=llm")
+        return self
+
+    @classmethod
+    def parsed(cls, extractor: str, commit: str = "") -> Provenance:
+        return cls(source=Source.PARSED, extractor=extractor, commit=commit)
+
+    @classmethod
+    def resolved(cls, extractor: str, commit: str = "") -> Provenance:
+        return cls(source=Source.RESOLVED, extractor=extractor, commit=commit)
+
+    @classmethod
+    def manual(cls, extractor: str, commit: str = "") -> Provenance:
+        return cls(source=Source.MANUAL, extractor=extractor, commit=commit)
+
+    @classmethod
+    def llm(cls, extractor: str, confidence: float, commit: str = "") -> Provenance:
+        return cls(
+            source=Source.LLM,
+            extractor=extractor,
+            confidence=confidence,
+            commit=commit,
+        )

agentforge_graph/core/symbols.py

@@ -0,0 +1,116 @@
+"""Symbol identity — stable, human-readable, deterministic node IDs.
+
+A symbol ID is a single string derived from
+``(scheme, lang, repo, path, descriptor)``. It has no global counters
+and no ordering constraints, so per-file extraction can run in any order
+and merge, and the same symbol keeps its ID across commits — the
+property incremental indexing (feat-004) and history (feat-009) depend
+on. The descriptor grammar is SCIP-derived. See ADR-0003 and
+``docs/design/design-001-core-contracts-module.md`` §4.4.
+"""
+
+from __future__ import annotations
+
+import hashlib
+
+from pydantic import BaseModel, ConfigDict
+
+SCHEME = "ckg"
+_FIELD_SEP = " "
+
+
+def _encode(field: str) -> str:
+    """Escape the field separator and the escape char so IDs round-trip."""
+    return field.replace("%", "%25").replace(_FIELD_SEP, "%20")
+
+
+def _decode(field: str) -> str:
+    # %20 before %25 so an escaped escape (%2520) decodes back to "%20".
+    return field.replace("%20", _FIELD_SEP).replace("%25", "%")
+
+
+def normalize_path(path: str) -> str:
+    """Repo-relative, posix-separated, no leading ``./`` or ``/``.
+
+    Ensures the same file yields the same ID on any OS.
+    """
+    p = path.replace("\\", "/")
+    while p.startswith("./"):
+        p = p[2:]
+    return p.lstrip("/")
+
+
+class ParsedSymbol(BaseModel):
+    """The structured form of a symbol ID."""
+
+    model_config = ConfigDict(frozen=True)
+
+    scheme: str
+    lang: str
+    repo: str
+    path: str
+    descriptor: str
+
+
+class SymbolID:
+    """Format and parse symbol-ID strings. Stateless."""
+
+    SCHEME = SCHEME
+
+    @staticmethod
+    def for_symbol(lang: str, repo: str, path: str, descriptor: str) -> str:
+        parts = [SCHEME, lang, repo, normalize_path(path), descriptor]
+        return _FIELD_SEP.join(_encode(p) for p in parts)
+
+    @staticmethod
+    def parse(symbol_id: str) -> ParsedSymbol:
+        raw = symbol_id.split(_FIELD_SEP)
+        if len(raw) != 5:
+            raise ValueError(
+                f"malformed symbol id (expected 5 space-separated fields): {symbol_id!r}"
+            )
+        scheme, lang, repo, path, descriptor = (_decode(p) for p in raw)
+        if scheme != SCHEME:
+            raise ValueError(f"unknown symbol-id scheme {scheme!r} (expected {SCHEME!r})")
+        return ParsedSymbol(scheme=scheme, lang=lang, repo=repo, path=path, descriptor=descriptor)
+
+
+class Descriptor:
+    """Builders for the SCIP-derived descriptor segments.
+
+    Segments compose by concatenation — a method on a class is
+    ``Descriptor.type("Auth") + Descriptor.method("login")`` →
+    ``"Auth#login()."``. Language packs (feat-002) map AST nodes to
+    these; core only owns the string format.
+    """
+
+    @staticmethod
+    def namespace(name: str) -> str:
+        return f"{name}/"
+
+    @staticmethod
+    def type(name: str) -> str:
+        return f"{name}#"
+
+    @staticmethod
+    def term(name: str) -> str:
+        return f"{name}."
+
+    @staticmethod
+    def method(name: str, disambiguator: int = 0) -> str:
+        """A method/function. ``disambiguator`` n>=1 marks the nth overload."""
+        if disambiguator < 0:
+            raise ValueError("disambiguator must be >= 0")
+        suffix = f"(+{disambiguator})" if disambiguator else ""
+        return f"{name}{suffix}()."
+
+    @staticmethod
+    def local(seed: str) -> str:
+        """Descriptor for an anonymous/local symbol with no stable name.
+
+        ``seed`` should be derived from the symbol's position within its
+        nearest *named* ancestor so edits above the ancestor don't shift
+        it. Inherently less stable than named symbols (ADR-0003 §risks).
+        """
+        digest = hashlib.sha1(seed.encode("utf-8")).hexdigest()[:8]
+        return f"local({digest})"

agentforge_graph/embed/__init__.py

@@ -0,0 +1,28 @@
+"""agentforge_graph.embed — chunk embedding (feat-005).
+
+Default real backend is AWS Bedrock Cohere embed-v4 (`BedrockEmbedder`);
+`OpenAIEmbedder` is the non-AWS / local-server path (ENH-003 phase 2); tests/CI
+use the deterministic `FakeEmbedder`. Imports nothing from ``agentforge``
+(ADR-0001); each driver's SDK (boto3 / openai) is lazy-imported in its module.
+"""
+
+from __future__ import annotations
+
+from .base import Embedder, InputType
+from .bedrock import BedrockEmbedder
+from .fake import FakeEmbedder
+from .openai import OpenAIEmbedder
+from .pipeline import EmbedPipeline
+from .registry import embedder_from_config
+from .report import EmbedReport
+
+__all__ = [
+    "Embedder",
+    "InputType",
+    "FakeEmbedder",
+    "BedrockEmbedder",
+    "OpenAIEmbedder",
+    "EmbedPipeline",
+    "EmbedReport",
+    "embedder_from_config",
+]

agentforge_graph/embed/base.py

@@ -0,0 +1,22 @@
+"""The ``Embedder`` contract. Implementations: ``FakeEmbedder`` (CI default,
+deterministic) and ``BedrockEmbedder`` (Cohere embed-v4). Imports nothing
+from ``agentforge`` (ADR-0001)."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Literal
+
+InputType = Literal["document", "query"]
+
+
+class Embedder(ABC):
+    name: str
+    dim: int
+
+    @abstractmethod
+    async def embed(
+        self, texts: list[str], input_type: InputType = "document"
+    ) -> list[list[float]]:
+        """Embed ``texts``. ``input_type`` distinguishes stored documents from
+        search queries (asymmetric models use it; symmetric ones ignore it)."""

agentforge_graph/embed/bedrock.py

@@ -0,0 +1,85 @@
+"""``BedrockEmbedder`` — AWS Bedrock Cohere embed-v4 via boto3.
+
+boto3 is imported lazily (only this driver needs it; it lives in the
+``bedrock`` extra). Synchronous Bedrock calls run on a worker thread.
+Supports an optional STS assume-role (the CI path); otherwise the default
+AWS credential chain (a developer's configured CLI). Voyage is *not* on
+Bedrock — see memory `embeddings-bedrock`.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Any
+
+from .base import Embedder, InputType
+
+_INPUT_MAP: dict[str, str] = {"document": "search_document", "query": "search_query"}
+
+
+class BedrockEmbedder(Embedder):
+    def __init__(
+        self,
+        model: str = "cohere.embed-v4:0",
+        region: str = "us-east-1",
+        dim: int = 1024,
+        batch_size: int = 96,
+        assume_role_arn: str | None = None,
+    ) -> None:
+        self.name = f"bedrock:{model}"
+        self.model = model
+        self.region = region
+        self.dim = dim
+        self.batch_size = batch_size
+        self.assume_role_arn = assume_role_arn
+        self._client: Any = None
+
+    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-embed")[
+                    "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 embed(
+        self, texts: list[str], input_type: InputType = "document"
+    ) -> list[list[float]]:
+        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, input_type))
+        return out
+
+    def _invoke(self, batch: list[str], input_type: InputType) -> list[list[float]]:
+        body = json.dumps(
+            {
+                "texts": batch,
+                "input_type": _INPUT_MAP[input_type],
+                "embedding_types": ["float"],
+                "output_dimension": self.dim,
+            }
+        )
+        resp = self._bedrock().invoke_model(
+            modelId=self.model,
+            contentType="application/json",
+            accept="application/json",
+            body=body,
+        )
+        payload = json.loads(resp["body"].read())
+        embeddings = payload["embeddings"]
+        floats = embeddings["float"] if isinstance(embeddings, dict) else embeddings
+        return [[float(x) for x in vec] for vec in floats]