agentforge-py 0.2.1__py3-none-any.whl

agentforge/memory/in_memory_graph.py

@@ -0,0 +1,262 @@
+"""`InMemoryGraphStore` — process-local `GraphStore` reference impl.
+
+Plain dict of nodes plus an adjacency list. Suitable for tests, demos,
+and small knowledge graphs (~thousands of nodes). Production
+deployments swap to `agentforge-memory-neo4j` or
+`agentforge-memory-surrealdb` — both pass the same
+`run_graph_conformance` suite.
+
+Design notes:
+  - All operations are O(degree) or smaller; pattern matching iterates
+    edges. Fine for small graphs, sluggish past tens of thousands of
+    edges. Native graph DBs declare richer capabilities (`"cypher"`,
+    `"transactions"`); this one does not.
+  - Idempotent upserts: re-adding a node or edge replaces its
+    properties. The contract requires this for both shapes.
+  - Edges that would orphan are blocked by `delete_node(cascade=False)`.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict, deque
+from typing import Any, Literal
+
+from agentforge_core.contracts.graph_store import GraphStore
+from agentforge_core.values.graph import (
+    GraphEdge,
+    GraphNode,
+    GraphPattern,
+    GraphSegment,
+    Path,
+)
+
+
+class InMemoryGraphStore(GraphStore):
+    """In-process `GraphStore` backed by dicts + adjacency lists."""
+
+    def __init__(self) -> None:
+        self._nodes: OrderedDict[str, GraphNode] = OrderedDict()
+        # Outgoing edges keyed by src; ordered for determinism.
+        self._out: dict[str, list[GraphEdge]] = {}
+        # Incoming edges keyed by dst; same ordering.
+        self._in: dict[str, list[GraphEdge]] = {}
+
+    async def add_node(self, node: GraphNode) -> None:
+        self._nodes[node.id] = node
+        self._out.setdefault(node.id, [])
+        self._in.setdefault(node.id, [])
+
+    async def add_edge(self, edge: GraphEdge) -> None:
+        if edge.src not in self._nodes:
+            msg = f"add_edge: source node {edge.src!r} does not exist"
+            raise ValueError(msg)
+        if edge.dst not in self._nodes:
+            msg = f"add_edge: destination node {edge.dst!r} does not exist"
+            raise ValueError(msg)
+        # Idempotent on (src, dst, edge_type): replace existing entry.
+        self._out[edge.src] = [
+            e
+            for e in self._out[edge.src]
+            if not (e.dst == edge.dst and e.edge_type == edge.edge_type)
+        ]
+        self._in[edge.dst] = [
+            e
+            for e in self._in[edge.dst]
+            if not (e.src == edge.src and e.edge_type == edge.edge_type)
+        ]
+        self._out[edge.src].append(edge)
+        self._in[edge.dst].append(edge)
+
+    async def get_node(self, node_id: str) -> GraphNode | None:
+        return self._nodes.get(node_id)
+
+    async def get_edges(
+        self,
+        node_id: str,
+        *,
+        edge_type: str | None = None,
+        direction: Literal["out", "in", "any"] = "out",
+    ) -> list[GraphEdge]:
+        if node_id not in self._nodes:
+            return []
+        candidates: list[GraphEdge]
+        if direction == "out":
+            candidates = list(self._out.get(node_id, ()))
+        elif direction == "in":
+            candidates = list(self._in.get(node_id, ()))
+        else:
+            candidates = [*self._out.get(node_id, ()), *self._in.get(node_id, ())]
+        if edge_type is not None:
+            candidates = [e for e in candidates if e.edge_type == edge_type]
+        return candidates
+
+    async def match(self, pattern: GraphPattern, *, limit: int = 50) -> list[Path]:
+        if limit < 1:
+            msg = f"limit must be >= 1, got {limit}"
+            raise ValueError(msg)
+
+        # Seed with every node — drivers like Neo4j let the optimiser
+        # pick a driving node; we walk all candidates that satisfy the
+        # first segment's source label + node filter.
+        seg0 = pattern.segments[0]
+        node_filter0 = pattern.node_filters[0] if pattern.node_filters else {}
+        starts = [n for n in self._nodes.values() if _node_matches(n, seg0.src_label, node_filter0)]
+
+        results: list[Path] = []
+        for start in starts:
+            self._walk_pattern(
+                start=start,
+                pattern=pattern,
+                seg_index=0,
+                visited_nodes=[start],
+                visited_edges=[],
+                results=results,
+                limit=limit,
+            )
+            if len(results) >= limit:
+                break
+        return results[:limit]
+
+    def _walk_pattern(
+        self,
+        *,
+        start: GraphNode,
+        pattern: GraphPattern,
+        seg_index: int,
+        visited_nodes: list[GraphNode],
+        visited_edges: list[GraphEdge],
+        results: list[Path],
+        limit: int,
+    ) -> None:
+        if seg_index == len(pattern.segments):
+            results.append(Path(nodes=tuple(visited_nodes), edges=tuple(visited_edges)))
+            return
+        if len(results) >= limit:
+            return
+
+        seg = pattern.segments[seg_index]
+        next_filter = pattern.node_filters[seg_index + 1] if pattern.node_filters else {}
+        for edge in self._edges_for_segment(start.id, seg):
+            other_id = edge.dst if edge.src == start.id else edge.src
+            other = self._nodes.get(other_id)
+            if other is None:
+                continue
+            if not _node_matches(other, seg.dst_label, next_filter):
+                continue
+            self._walk_pattern(
+                start=other,
+                pattern=pattern,
+                seg_index=seg_index + 1,
+                visited_nodes=[*visited_nodes, other],
+                visited_edges=[*visited_edges, edge],
+                results=results,
+                limit=limit,
+            )
+            if len(results) >= limit:
+                return
+
+    def _edges_for_segment(self, node_id: str, seg: GraphSegment) -> list[GraphEdge]:
+        edges: list[GraphEdge]
+        if seg.direction == "out":
+            edges = list(self._out.get(node_id, ()))
+        elif seg.direction == "in":
+            edges = list(self._in.get(node_id, ()))
+        else:
+            edges = [*self._out.get(node_id, ()), *self._in.get(node_id, ())]
+        if seg.edge_type is not None:
+            edges = [e for e in edges if e.edge_type == seg.edge_type]
+        return edges
+
+    async def traverse(
+        self,
+        start_id: str,
+        *,
+        edge_types: tuple[str, ...] | None = None,
+        max_depth: int = 3,
+        limit: int = 50,
+    ) -> list[Path]:
+        if max_depth < 1:
+            msg = f"max_depth must be >= 1, got {max_depth}"
+            raise ValueError(msg)
+        if limit < 1:
+            msg = f"limit must be >= 1, got {limit}"
+            raise ValueError(msg)
+        if start_id not in self._nodes:
+            return []
+
+        start_node = self._nodes[start_id]
+        results: list[Path] = []
+        # BFS frontier of (current_node, path-so-far-nodes, path-so-far-edges)
+        frontier: deque[tuple[str, list[GraphNode], list[GraphEdge]]] = deque(
+            [(start_id, [start_node], [])]
+        )
+        while frontier and len(results) < limit:
+            current, path_nodes, path_edges = frontier.popleft()
+            if len(path_edges) >= max_depth:
+                continue
+            for edge in self._out.get(current, ()):
+                if edge_types is not None and edge.edge_type not in edge_types:
+                    continue
+                target = self._nodes.get(edge.dst)
+                if target is None:
+                    continue
+                # Avoid cycles: don't revisit a node already in this path.
+                if any(n.id == target.id for n in path_nodes):
+                    continue
+                new_nodes = [*path_nodes, target]
+                new_edges = [*path_edges, edge]
+                results.append(Path(nodes=tuple(new_nodes), edges=tuple(new_edges)))
+                if len(results) >= limit:
+                    break
+                frontier.append((target.id, new_nodes, new_edges))
+        return results[:limit]
+
+    async def delete_node(self, node_id: str, *, cascade: bool = False) -> bool:
+        if node_id not in self._nodes:
+            return False
+        outgoing = self._out.get(node_id, [])
+        incoming = self._in.get(node_id, [])
+        if not cascade and (outgoing or incoming):
+            msg = (
+                f"delete_node: {node_id!r} has {len(outgoing) + len(incoming)} "
+                f"incident edge(s); pass cascade=True to remove them"
+            )
+            raise ValueError(msg)
+        # Cascade removes incident edges from the *other* side too.
+        for e in list(outgoing):
+            self._in[e.dst] = [x for x in self._in[e.dst] if x is not e]
+        for e in list(incoming):
+            self._out[e.src] = [x for x in self._out[e.src] if x is not e]
+        self._out.pop(node_id, None)
+        self._in.pop(node_id, None)
+        self._nodes.pop(node_id, None)
+        return True
+
+    async def delete_edge(self, src: str, dst: str, *, edge_type: str) -> bool:
+        out = self._out.get(src, [])
+        match = next((e for e in out if e.dst == dst and e.edge_type == edge_type), None)
+        if match is None:
+            return False
+        self._out[src] = [e for e in out if e is not match]
+        self._in[dst] = [e for e in self._in.get(dst, ()) if e is not match]
+        return True
+
+    async def close(self) -> None:
+        self._nodes.clear()
+        self._out.clear()
+        self._in.clear()
+
+
+# ----------------------------------------------------------------------
+# Helpers
+# ----------------------------------------------------------------------
+
+
+def _node_matches(
+    node: GraphNode,
+    label: str | None,
+    properties_filter: dict[str, Any],
+) -> bool:
+    if label is not None and label not in node.labels:
+        return False
+    return all(node.properties.get(k) == v for k, v in properties_filter.items())

agentforge/memory/in_memory_vector.py

@@ -0,0 +1,167 @@
+"""`InMemoryVectorStore` — process-local `VectorStore` reference impl.
+
+Brute-force cosine similarity over a Python dict. Suitable for tests,
+demos, and small RAG corpora (~hundreds of items). Production
+deployments swap to `agentforge-memory-sqlite` (zero-deps persistent)
+or `agentforge-memory-postgres` (scaled, native ANN) — both pass the
+same `run_vector_conformance` suite.
+
+Design notes:
+
+  - Dimensions are fixed at construction. Mismatched vectors raise
+    `ValueError` immediately, before they can corrupt the index.
+  - Search is O(N) per call — fine for small corpora, sluggish past
+    a few thousand items. The capability vocabulary lets future
+    drivers declare `"native_ann"`; this one does not.
+  - Vectors are L2-normalised before storage so similarity is a
+    plain dot product at search time. Callers don't have to
+    pre-normalise.
+"""
+
+from __future__ import annotations
+
+import math
+from collections import OrderedDict
+from typing import Any
+
+from agentforge_core._bm25 import _BM25Index
+from agentforge_core.contracts.vector_store import VectorStore
+from agentforge_core.values.vector import VectorItem, VectorMatch
+
+
+class InMemoryVectorStore(VectorStore):
+    """In-process `VectorStore` backed by a dict + brute-force search."""
+
+    def __init__(self, *, dimensions: int) -> None:
+        if dimensions < 1:
+            raise ValueError(f"dimensions must be >= 1, got {dimensions}")
+        self._dim = dimensions
+        # id -> (normalised vector, text, metadata)
+        self._items: OrderedDict[str, tuple[tuple[float, ...], str, dict[str, Any]]] = OrderedDict()
+        # Lazy BM25 index for hybrid search — rebuilt on demand after
+        # any upsert/delete. `None` means "needs rebuild on next
+        # lexical_search call".
+        self._bm25: _BM25Index | None = None
+
+    def dimensions(self) -> int:
+        return self._dim
+
+    def capabilities(self) -> set[str]:
+        return {"hybrid_search"}
+
+    async def upsert(self, items: list[VectorItem]) -> None:
+        for item in items:
+            if len(item.vector) != self._dim:
+                raise ValueError(
+                    f"vector for id={item.id!r} has length {len(item.vector)} "
+                    f"but store dimensions={self._dim}"
+                )
+            self._items[item.id] = (
+                _l2_normalise(item.vector),
+                item.text,
+                dict(item.metadata),
+            )
+        self._bm25 = None  # invalidate
+
+    async def search(
+        self,
+        query_vector: tuple[float, ...],
+        *,
+        limit: int = 5,
+        filter_metadata: dict[str, Any] | None = None,
+    ) -> list[VectorMatch]:
+        if limit < 1:
+            raise ValueError(f"limit must be >= 1, got {limit}")
+        if len(query_vector) != self._dim:
+            raise ValueError(
+                f"query vector has length {len(query_vector)} but store dimensions={self._dim}"
+            )
+
+        query = _l2_normalise(query_vector)
+        scored: list[tuple[float, str, str, dict[str, Any]]] = []
+        for item_id, (vec, text, meta) in self._items.items():
+            if filter_metadata is not None and not _matches_filter(meta, filter_metadata):
+                continue
+            # Cosine similarity in [-1, 1]. Clamped to [0, 1]: 1.0
+            # means identical direction, 0.0 means orthogonal-or-
+            # anti-correlated. For text embeddings this drops
+            # essentially no information (anti-correlation is rare).
+            similarity = sum(a * b for a, b in zip(query, vec, strict=True))
+            clamped = max(0.0, min(1.0, similarity))
+            scored.append((clamped, item_id, text, meta))
+
+        scored.sort(key=lambda row: row[0], reverse=True)
+        top = scored[:limit]
+        return [
+            VectorMatch(id=item_id, text=text, metadata=dict(meta), score=score)
+            for score, item_id, text, meta in top
+        ]
+
+    async def lexical_search(
+        self,
+        query: str,
+        *,
+        limit: int = 5,
+        filter_metadata: dict[str, Any] | None = None,
+    ) -> list[VectorMatch]:
+        if limit < 1:
+            raise ValueError(f"limit must be >= 1, got {limit}")
+        if not self._items:
+            return []
+
+        if self._bm25 is None:
+            idx = _BM25Index()
+            for item_id, (_vec, text, _meta) in self._items.items():
+                idx.add(item_id, text)
+            self._bm25 = idx
+
+        # Over-fetch by a small factor so the post-filter has room to
+        # work even when many top BM25 hits get filtered out by
+        # metadata. The fuser eventually narrows to `limit` anyway.
+        over_fetch = limit if filter_metadata is None else limit * 4
+        scored = self._bm25.score(query, limit=max(over_fetch, 1))
+        if not scored:
+            return []
+
+        max_score = scored[0][1]
+        matches: list[VectorMatch] = []
+        for doc_id, raw in scored:
+            _vec, text, meta = self._items[doc_id]
+            if filter_metadata is not None and not _matches_filter(meta, filter_metadata):
+                continue
+            normalised = raw / max_score if max_score > 0.0 else 0.0
+            matches.append(VectorMatch(id=doc_id, text=text, metadata=dict(meta), score=normalised))
+            if len(matches) >= limit:
+                break
+        return matches
+
+    async def delete(self, ids: list[str]) -> int:
+        removed = 0
+        for item_id in ids:
+            if self._items.pop(item_id, None) is not None:
+                removed += 1
+        if removed:
+            self._bm25 = None  # invalidate
+        return removed
+
+    async def close(self) -> None:
+        self._items.clear()
+        self._bm25 = None
+
+
+# ----------------------------------------------------------------------
+# Helpers
+# ----------------------------------------------------------------------
+
+
+def _l2_normalise(vector: tuple[float, ...]) -> tuple[float, ...]:
+    """Return `vector / ||vector||`; falls back to the input if zero."""
+    norm = math.sqrt(sum(x * x for x in vector))
+    if norm == 0.0:
+        return vector
+    return tuple(x / norm for x in vector)
+
+
+def _matches_filter(metadata: dict[str, Any], filter_md: dict[str, Any]) -> bool:
+    """True if every (k, v) in `filter_md` matches `metadata`."""
+    return all(metadata.get(k) == v for k, v in filter_md.items())

agentforge/pipeline/__init__.py

@@ -0,0 +1,26 @@
+"""Pipeline engine + Task ABC (feat-015).
+
+`Pipeline` runs a DAG of deterministic `Task` instances before the
+agent's reasoning loop starts; their consolidated findings are
+exposed to the LLM as a system-prompt addendum and via a built-in
+``pipeline_findings`` tool.
+"""
+
+from __future__ import annotations
+
+from agentforge_core.contracts.task import Task
+from agentforge_core.values.pipeline import PipelineResult
+
+from agentforge.pipeline.engine import OnTaskError, Pipeline
+from agentforge.pipeline.errors import PipelineFailure
+from agentforge.pipeline.tool import PipelineFindingsInput, PipelineFindingsTool
+
+__all__ = [
+    "OnTaskError",
+    "Pipeline",
+    "PipelineFailure",
+    "PipelineFindingsInput",
+    "PipelineFindingsTool",
+    "PipelineResult",
+    "Task",
+]

agentforge/pipeline/engine.py

@@ -0,0 +1,189 @@
+"""`Pipeline` — deterministic-task DAG engine (feat-015).
+
+Runs a set of `Task` instances as a DAG (resolved via
+``Task.depends_on``), respecting a ``max_concurrent`` semaphore. Each
+task's execution is bounded by ``Task.timeout_s``. Failures are
+handled per ``on_task_error``:
+
+  - ``"continue"`` (default): the failed task emits a
+    ``SimpleFinding`` with category ``"pipeline.task_failure"``;
+    dependents still run (they see the failure-finding in their
+    merged context under ``"pipeline_findings_so_far"``).
+  - ``"fail"``: the engine cancels outstanding tasks and raises
+    `PipelineFailure`. Findings collected before the failure are
+    lost — the caller treats the run as aborted.
+
+The output is a frozen `PipelineResult` with consolidated findings,
+per-task durations, the failures map, and total cost.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import Mapping
+from typing import Any, Literal
+
+from agentforge_core.contracts.task import Task
+from agentforge_core.values.pipeline import PipelineResult
+
+from agentforge.findings import SimpleFinding
+from agentforge.pipeline.errors import PipelineFailure
+
+OnTaskError = Literal["continue", "fail"]
+
+
+class Pipeline:
+    """A DAG of deterministic (or LLM-using) tasks.
+
+    Construct once, ``run(context)`` once per agent invocation.
+    Validation (duplicate names, missing deps, cycles) runs at
+    construction so misconfigurations fail before the first run.
+    """
+
+    def __init__(
+        self,
+        tasks: list[Task],
+        *,
+        max_concurrent: int = 4,
+        on_task_error: OnTaskError = "continue",
+    ) -> None:
+        if max_concurrent < 1:
+            raise ValueError(f"max_concurrent must be >= 1, got {max_concurrent}")
+        if on_task_error not in ("continue", "fail"):
+            raise ValueError(f"on_task_error must be 'continue' or 'fail', got {on_task_error!r}")
+        self._tasks = list(tasks)
+        self._max_concurrent = max_concurrent
+        self._on_task_error = on_task_error
+        self._by_name: dict[str, Task] = {}
+        for t in self._tasks:
+            tname = type(t).name
+            if tname in self._by_name:
+                raise ValueError(f"duplicate task name in pipeline: {tname!r}")
+            self._by_name[tname] = t
+        self._validate_dag()
+
+    @property
+    def tasks(self) -> tuple[Task, ...]:
+        return tuple(self._tasks)
+
+    @property
+    def max_concurrent(self) -> int:
+        return self._max_concurrent
+
+    @property
+    def on_task_error(self) -> OnTaskError:
+        return self._on_task_error
+
+    def _validate_dag(self) -> None:
+        for t in self._tasks:
+            for dep in type(t).depends_on:
+                if dep not in self._by_name:
+                    raise ValueError(f"task {type(t).name!r} depends on unknown task {dep!r}")
+        # cycle detection via DFS
+        WHITE, GREY, BLACK = 0, 1, 2  # noqa: N806
+        color = dict.fromkeys(self._by_name, WHITE)
+
+        def visit(name: str, stack: list[str]) -> None:
+            if color[name] == GREY:
+                cycle = " -> ".join([*stack, name])
+                raise ValueError(f"cycle in pipeline DAG: {cycle}")
+            if color[name] == BLACK:
+                return
+            color[name] = GREY
+            for dep in type(self._by_name[name]).depends_on:
+                visit(dep, [*stack, name])
+            color[name] = BLACK
+
+        for name in self._by_name:
+            visit(name, [])
+
+    async def run(self, context: Mapping[str, Any]) -> PipelineResult:
+        """Execute the pipeline once and return consolidated output."""
+        loop = asyncio.get_running_loop()
+        state = _RunState(
+            semaphore=asyncio.Semaphore(self._max_concurrent),
+            # Each task gets a future of its emitted findings (or [] on
+            # failure when on_task_error="continue"); dependents await.
+            futures={name: loop.create_future() for name in self._by_name},
+        )
+        runners = [asyncio.create_task(self._run_one(t, context, state)) for t in self._tasks]
+        try:
+            await asyncio.gather(*runners)
+        except PipelineFailure:
+            for r in runners:
+                if not r.done():
+                    r.cancel()
+            await asyncio.gather(*runners, return_exceptions=True)
+            raise
+
+        return PipelineResult(
+            findings=tuple(state.accumulated),
+            task_durations_ms=state.durations,
+            task_failures=state.failures,
+            total_cost_usd=state.total_cost,
+        )
+
+    async def _run_one(self, task: Task, context: Mapping[str, Any], state: _RunState) -> None:
+        tname = type(task).name
+        dep_results = [await state.futures[dep] for dep in type(task).depends_on]
+        prior: list[Any] = []
+        for batch in dep_results:
+            prior.extend(batch)
+        async with state.lock:
+            snapshot = list(state.accumulated)
+        merged: dict[str, Any] = dict(context)
+        existing = merged.get("pipeline_findings_so_far", [])
+        merged["pipeline_findings_so_far"] = [*existing, *prior, *snapshot]
+
+        start = time.monotonic()
+        async with state.semaphore:
+            try:
+                timeout = float(type(task).timeout_s)
+                findings = await asyncio.wait_for(task.run(merged), timeout=timeout)
+            except Exception as exc:
+                await self._record_failure(tname, exc, start, state)
+                return
+            elapsed = int((time.monotonic() - start) * 1000)
+            state.durations[tname] = elapsed
+            state.total_cost += float(type(task).cost_estimate_usd)
+            async with state.lock:
+                state.accumulated.extend(findings)
+            state.futures[tname].set_result(list(findings))
+
+    async def _record_failure(
+        self, tname: str, exc: BaseException, start: float, state: _RunState
+    ) -> None:
+        elapsed = int((time.monotonic() - start) * 1000)
+        state.durations[tname] = elapsed
+        state.failures[tname] = str(exc) if str(exc) else repr(exc)
+        if self._on_task_error == "fail":
+            state.futures[tname].set_exception(PipelineFailure(tname, exc))
+            raise PipelineFailure(tname, exc) from exc
+        failure_finding = SimpleFinding(
+            severity="error",
+            category="pipeline.task_failure",
+            message=f"task {tname!r} failed: {state.failures[tname]}",
+            rule_id=tname,
+        )
+        async with state.lock:
+            state.accumulated.append(failure_finding)
+        state.futures[tname].set_result([failure_finding])
+
+
+class _RunState:
+    """Mutable per-run scratch shared across `_run_one` invocations."""
+
+    def __init__(
+        self,
+        *,
+        semaphore: asyncio.Semaphore,
+        futures: dict[str, asyncio.Future[list[Any]]],
+    ) -> None:
+        self.semaphore = semaphore
+        self.futures = futures
+        self.accumulated: list[Any] = []
+        self.durations: dict[str, int] = {}
+        self.failures: dict[str, str] = {}
+        self.total_cost = 0.0
+        self.lock = asyncio.Lock()

agentforge/pipeline/errors.py

@@ -0,0 +1,19 @@
+"""Pipeline-specific exceptions (feat-015)."""
+
+from __future__ import annotations
+
+from agentforge_core.production.exceptions import AgentForgeError
+
+
+class PipelineFailure(AgentForgeError):  # noqa: N818  # locked name; parity with BudgetExceeded / GuardrailViolation
+    """Raised when a `Pipeline` configured with ``on_task_error="fail"``
+    encounters a task that raised.
+
+    The first failed task's name is on ``task_name``; the original
+    exception is chained via ``__cause__``.
+    """
+
+    def __init__(self, task_name: str, cause: BaseException) -> None:
+        super().__init__(f"pipeline task {task_name!r} failed: {cause}")
+        self.task_name = task_name
+        self.__cause__ = cause