agentforge-py 0.2.1__py3-none-any.whl

agentforge/guardrails/engine.py

@@ -0,0 +1,289 @@
+"""Guardrail engine — runs validators + emits audit events (feat-018).
+
+Glues `InputValidator` / `OutputValidator` / `ToolCallGate` lists to
+the agent's lifecycle. `GuardrailEngine` is constructed once per
+`Agent` from the config + the framework-wide policy, then:
+
+- `await engine.check_input(content, ctx)` is invoked once at the
+  start of `agent.run(task)`.
+- The engine wraps the LLM client (`engine.wrap_llm(llm)`) so every
+  `.call(...)` output flows through `OutputValidator`s.
+- The engine wraps each tool (`engine.wrap_tool(tool)`) so every
+  `.run(...)` is gated through `ToolCallGate`s.
+
+All decisions append to `engine.events` (a list of dicts conforming
+to `RunResult.guardrail_events` shape) and emit a structured log
+record on the `agentforge.audit` channel.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from agentforge_core.config.schema import GuardrailPolicy
+from agentforge_core.contracts.guardrails import (
+    InputValidator,
+    OutputValidator,
+    ToolCallGate,
+)
+from agentforge_core.contracts.llm import LLMClient
+from agentforge_core.contracts.tool import Tool
+from agentforge_core.production.exceptions import GuardrailViolation, ModuleError
+from agentforge_core.values.guardrails import ValidationResult
+from agentforge_core.values.messages import LLMResponse, Message, ToolSpec
+
+_audit_log = logging.getLogger("agentforge.audit")
+
+
+class GuardrailEngine:
+    """Per-Agent engine that runs validators + records decisions."""
+
+    def __init__(
+        self,
+        *,
+        input_validators: list[InputValidator],
+        output_validators: list[OutputValidator],
+        tool_gates: list[ToolCallGate],
+        policy: GuardrailPolicy,
+    ) -> None:
+        self._inputs = list(input_validators)
+        self._outputs = list(output_validators)
+        self._gates = list(tool_gates)
+        self._policy = policy
+        self.events: list[dict[str, Any]] = []
+
+    # ------------------------------------------------------------------
+    # Input
+    # ------------------------------------------------------------------
+
+    async def check_input(self, content: str, context: dict[str, Any]) -> str:
+        """Run every input validator. Returns the (possibly redacted)
+        content for downstream use; raises `GuardrailViolation` on
+        block."""
+        return await self._run_stage(
+            stage="input",
+            validators=self._inputs,
+            content=content,
+            context=context,
+            action=self._policy.on_input_violation,
+            call=lambda v, c, ctx: v.validate(c, ctx),
+        )
+
+    async def check_output(self, content: str, context: dict[str, Any]) -> str:
+        """Run every output validator post-LLM call."""
+        return await self._run_stage(
+            stage="output",
+            validators=self._outputs,
+            content=content,
+            context=context,
+            action=self._policy.on_output_violation,
+            call=lambda v, c, ctx: v.validate(c, ctx),
+        )
+
+    # ------------------------------------------------------------------
+    # Wrappers
+    # ------------------------------------------------------------------
+
+    def wrap_llm(self, real: LLMClient, context_factory: Callable[[], dict[str, Any]]) -> LLMClient:
+        """Wrap a real `LLMClient` so every `.call()` result flows
+        through the configured `OutputValidator`s.
+
+        Validators see the model's text content; tool_calls are
+        passed through untouched (those go through `ToolCallGate`s
+        on dispatch).
+        """
+        return _GuardedLLMClient(real=real, engine=self, context_factory=context_factory)
+
+    def wrap_tool(self, real: Tool, context_factory: Callable[[], dict[str, Any]]) -> Tool:
+        """Wrap a real `Tool` so every `.run(...)` is gated."""
+        return _wrap_tool(real, self, context_factory)
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    async def _run_stage(
+        self,
+        *,
+        stage: str,
+        validators: list[Any],
+        content: str,
+        context: dict[str, Any],
+        action: str,
+        call: Callable[[Any, str, dict[str, Any]], Awaitable[ValidationResult]],
+    ) -> str:
+        current = content
+        for v in validators:
+
+            async def _call(
+                validator: Any = v,
+                content_value: str = current,
+                ctx: dict[str, Any] = context,
+            ) -> ValidationResult:
+                return await call(validator, content_value, ctx)
+
+            result = await self._safe_invoke(v, _call)
+            self._emit(stage=stage, validator=v, result=result, content=current, action=action)
+            if result.passed:
+                continue
+            if action == "block":
+                msg = (
+                    f"guardrail block: validator={v.name!r} stage={stage} "
+                    f"violations={list(result.violations)!r}"
+                )
+                raise GuardrailViolation(msg)
+            if action == "redact" and result.redacted_content is not None:
+                current = result.redacted_content
+            # warn / allow: continue without modification
+        return current
+
+    async def authorize_tool(
+        self,
+        tool_name: str,
+        tool: Tool,
+        args: dict[str, Any],
+        context: dict[str, Any],
+    ) -> None:
+        action = self._policy.on_tool_violation
+        for gate in self._gates:
+
+            async def _call(g: ToolCallGate = gate) -> ValidationResult:
+                return await g.authorize(tool_name, tool, args, context)
+
+            result = await self._safe_invoke(gate, _call)
+            self._emit(
+                stage="tool",
+                validator=gate,
+                result=result,
+                content=f"{tool_name}:{sorted(args)}",
+                action=action,
+            )
+            if result.passed:
+                continue
+            if action == "block":
+                msg = (
+                    f"guardrail block: tool gate={gate.name!r} tool={tool_name!r} "
+                    f"violations={list(result.violations)!r}"
+                )
+                raise GuardrailViolation(msg)
+
+    async def _safe_invoke(
+        self,
+        validator: Any,
+        call: Callable[[], Awaitable[ValidationResult]],
+    ) -> ValidationResult:
+        try:
+            return await call()
+        except (RuntimeError, ModuleError, ValueError) as exc:
+            if self._policy.fail_open:
+                _audit_log.warning(
+                    "validator %s raised; fail_open=True so the call proceeds: %s",
+                    validator.name,
+                    exc,
+                )
+                return ValidationResult.ok()
+            return ValidationResult(
+                passed=False,
+                score=0.0,
+                violations=("validator_error",),
+                metadata={"error": str(exc)},
+            )
+
+    def _emit(
+        self,
+        *,
+        stage: str,
+        validator: Any,
+        result: ValidationResult,
+        content: str,
+        action: str,
+    ) -> None:
+        event = {
+            "stage": stage,
+            "validator": validator.name,
+            "passed": result.passed,
+            "violations": list(result.violations),
+            "score": result.score,
+            "action": action,
+            "content_hash": _hash(content),
+        }
+        self.events.append(event)
+        log = _audit_log.info if result.passed else _audit_log.warning
+        log(
+            "guardrail %s: %s passed=%s violations=%s action=%s",
+            stage,
+            validator.name,
+            result.passed,
+            list(result.violations),
+            action,
+        )
+
+
+def _hash(content: str) -> str:
+    return hashlib.sha256(content.encode("utf-8")).hexdigest()[:16]
+
+
+class _GuardedLLMClient(LLMClient):
+    """LLM wrapper that runs output validators after each `.call(...)`."""
+
+    def __init__(
+        self,
+        *,
+        real: LLMClient,
+        engine: GuardrailEngine,
+        context_factory: Callable[[], dict[str, Any]],
+    ) -> None:
+        self._real = real
+        self._engine = engine
+        self._ctx = context_factory
+
+    async def call(
+        self,
+        system: str,
+        messages: list[Message],
+        tools: list[ToolSpec] | None = None,
+    ) -> LLMResponse:
+        response = await self._real.call(system, messages, tools)
+        if not response.content:
+            return response
+        validated = await self._engine.check_output(response.content, self._ctx())
+        if validated == response.content:
+            return response
+        return response.model_copy(update={"content": validated})
+
+    async def close(self) -> None:
+        await self._real.close()
+
+
+def _wrap_tool(
+    real: Tool,
+    engine: GuardrailEngine,
+    context_factory: Callable[[], dict[str, Any]],
+) -> Tool:
+    """Build a per-instance Tool wrapper that consults `engine` before
+    forwarding to `real.run(...)`."""
+
+    real_name = type(real).name
+    real_description = type(real).description
+    real_input_schema = type(real).input_schema
+    real_caps = type(real).capabilities
+
+    async def _run(_self: Tool, **kwargs: Any) -> Any:
+        await engine.authorize_tool(real_name, real, kwargs, context_factory())
+        return await real.run(**kwargs)
+
+    cls_namespace: dict[str, Any] = {
+        "name": real_name,
+        "description": real_description,
+        "input_schema": real_input_schema,
+        "capabilities": real_caps,
+        "run": _run,
+    }
+    synthesized = type(f"Guarded_{type(real).__name__}", (Tool,), cls_namespace)
+    return synthesized()  # type: ignore[no-any-return]
+
+
+__all__ = ["GuardrailEngine"]

agentforge/guardrails/pii_redact_basic.py

@@ -0,0 +1,61 @@
+"""`pii_redact_basic` — regex-based PII detection + redaction
+(feat-018).
+
+Output validator. Replaces detected PII with `<redacted:KIND>`
+placeholders so downstream consumers can still parse around the
+redactions.
+
+Patterns are conservative — they catch obvious cases (RFC-822
+email, US-shaped SSN / phone, common credit-card and IPv4
+formats) without trying to be exhaustive. For richer coverage,
+install `agentforge-guard-presidio`.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any, ClassVar
+
+from agentforge_core.contracts.guardrails import OutputValidator
+from agentforge_core.resolver import register
+from agentforge_core.values.guardrails import ValidationResult
+
+_PATTERNS: tuple[tuple[str, re.Pattern[str]], ...] = (
+    ("email", re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}")),
+    ("phone_us", re.compile(r"\b(?:\+?1[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b")),
+    ("ssn", re.compile(r"\b\d{3}-\d{2}-\d{4}\b")),
+    ("credit_card", re.compile(r"\b(?:\d[ -]*?){13,16}\b")),
+    ("ipv4", re.compile(r"\b\d{1,3}(?:\.\d{1,3}){3}\b")),
+)
+
+
+@register("guardrails.output", "pii_redact_basic")
+class PIIRedactBasic(OutputValidator):
+    """Regex-based PII detector + redactor (basic tier)."""
+
+    name: ClassVar[str] = "pii_redact_basic"
+    description: ClassVar[str] = (
+        "Regex-based PII detector (email / phone / SSN / credit-card / IPv4) "
+        "that emits `<redacted:KIND>` placeholders in the output."
+    )
+    cost_estimate_ms: ClassVar[int] = 2
+
+    async def validate(self, content: str, context: dict[str, Any]) -> ValidationResult:
+        del context
+        violations: list[str] = []
+        redacted = content
+        for kind, pattern in _PATTERNS:
+            if pattern.search(redacted):
+                violations.append(kind)
+                redacted = pattern.sub(f"<redacted:{kind}>", redacted)
+        if not violations:
+            return ValidationResult.ok()
+        return ValidationResult(
+            passed=False,
+            score=0.0,
+            violations=tuple(violations),
+            redacted_content=redacted,
+        )
+
+
+__all__ = ["PIIRedactBasic"]

agentforge/guardrails/prompt_injection_basic.py

@@ -0,0 +1,90 @@
+"""`prompt_injection_basic` — regex pattern matching for the most
+common prompt-injection phrases (feat-018).
+
+This is the *basic* tier — it catches the obvious cases ("ignore
+previous instructions") and almost nothing else. Production
+deployments install `agentforge-guard-llmguard` or
+`agentforge-guard-llamaguard` for richer coverage.
+
+The pattern set is intentionally conservative: false positives are
+expensive (block user) so we only flag phrases with very high
+prior probability of malicious intent.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any, ClassVar
+
+from agentforge_core.contracts.guardrails import InputValidator
+from agentforge_core.resolver import register
+from agentforge_core.values.guardrails import ValidationResult
+
+_PATTERNS: tuple[tuple[str, re.Pattern[str]], ...] = (
+    (
+        "ignore_previous",
+        re.compile(
+            r"\bignore\s+(?:all\s+)?(?:previous|prior)\s+(?:instructions?|prompts?|messages?)\b",
+            re.IGNORECASE,
+        ),
+    ),
+    (
+        "disregard_instructions",
+        re.compile(r"\bdisregard\s+(?:all\s+)?(?:your\s+)?instructions?\b", re.IGNORECASE),
+    ),
+    (
+        "system_prompt_leak",
+        re.compile(
+            r"\b(?:reveal|print|show|repeat|leak|dump)\s+(?:your|the)\s+(?:system\s+)?(?:prompt|instructions)\b",
+            re.IGNORECASE,
+        ),
+    ),
+    (
+        "act_as_jailbreak",
+        re.compile(
+            r"\bact\s+as\s+(?:dan|jailbroken|developer\s+mode|do\s+anything\s+now)\b", re.IGNORECASE
+        ),
+    ),
+    (
+        "new_persona",
+        re.compile(
+            r"\byou\s+are\s+now\s+(?:dan|jailbroken|in\s+developer\s+mode)\b", re.IGNORECASE
+        ),
+    ),
+    (
+        "bypass_safety",
+        re.compile(
+            r"\b(?:bypass|disable|turn\s+off)\s+(?:safety|guardrails?|filter|moderation)\b",
+            re.IGNORECASE,
+        ),
+    ),
+)
+
+
+@register("guardrails.input", "prompt_injection_basic")
+class PromptInjectionBasic(InputValidator):
+    """Regex-based prompt-injection detector (basic tier)."""
+
+    name: ClassVar[str] = "prompt_injection_basic"
+    description: ClassVar[str] = (
+        "Regex-based detector for the most common prompt-injection phrases. "
+        "Conservative pattern set — false positives are expensive."
+    )
+    cost_estimate_ms: ClassVar[int] = 1
+
+    async def validate(self, content: str, context: dict[str, Any]) -> ValidationResult:
+        del context
+        violations: list[str] = []
+        for rule, pattern in _PATTERNS:
+            if pattern.search(content):
+                violations.append(rule)
+        if violations:
+            return ValidationResult(
+                passed=False,
+                score=0.0,
+                violations=tuple(violations),
+            )
+        return ValidationResult.ok()
+
+
+__all__ = ["PromptInjectionBasic"]

agentforge/memory/__init__.py

@@ -0,0 +1,16 @@
+"""Default in-process memory implementations.
+
+`InMemoryStore` (claim audit log) and `InMemoryVectorStore` (semantic
+search) ship with `agentforge` so a fresh `Agent(...)` has
+durable-shaped state out of the box without external infra. Both are
+safe defaults; production deployments swap to real drivers
+(`agentforge-memory-sqlite`, `-postgres`) via `agentforge.yaml`.
+"""
+
+from __future__ import annotations
+
+from agentforge.memory.in_memory import InMemoryStore
+from agentforge.memory.in_memory_graph import InMemoryGraphStore
+from agentforge.memory.in_memory_vector import InMemoryVectorStore
+
+__all__ = ["InMemoryGraphStore", "InMemoryStore", "InMemoryVectorStore"]

agentforge/memory/in_memory.py

@@ -0,0 +1,130 @@
+"""`InMemoryStore` — process-local `MemoryStore` implementation.
+
+Used as the default when no persistence module is configured. Loses
+data on process exit (by design — for tests, demos, ephemeral runs).
+For durable storage swap to a feat-005 driver via `agentforge.yaml`.
+
+Implementation notes:
+
+  - Storage is an `OrderedDict[str, Claim]` keyed by claim id, so
+    insertion order is preserved (sortable ULID ids already provide
+    monotonic ordering, but the dict guarantees it after deletes /
+    supersedes).
+  - All operations are ``async def`` to match the `MemoryStore`
+    contract; they don't actually do I/O so they complete on the
+    current event-loop tick.
+  - Thread safety is NOT a goal — this is for in-process single-loop
+    use. Multi-process / multi-worker deployments use feat-005 drivers.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from collections.abc import AsyncIterator
+from datetime import datetime
+
+from agentforge_core.contracts.memory import MemoryStore
+from agentforge_core.production.exceptions import ModuleError
+from agentforge_core.values.claim import Claim
+
+
+class InMemoryStore(MemoryStore):
+    """Process-local `MemoryStore` backed by an in-memory ``OrderedDict``."""
+
+    def __init__(self) -> None:
+        self._items: OrderedDict[str, Claim] = OrderedDict()
+
+    async def put(self, claim: Claim) -> str:
+        self._items[claim.id] = claim
+        return claim.id
+
+    async def get(self, claim_id: str) -> Claim | None:
+        return self._items.get(claim_id)
+
+    async def query(
+        self,
+        *,
+        project: str | None = None,
+        agent: str | None = None,
+        category: str | None = None,
+        run_id: str | None = None,
+        limit: int = 100,
+    ) -> list[Claim]:
+        results: list[Claim] = []
+        for claim in self._items.values():
+            if project is not None and claim.project != project:
+                continue
+            if agent is not None and claim.agent != agent:
+                continue
+            if category is not None and claim.category != category:
+                continue
+            if run_id is not None and claim.run_id != run_id:
+                continue
+            results.append(claim)
+            if len(results) >= limit:
+                break
+        return results
+
+    async def supersede(self, old_id: str, new_claim: Claim) -> str:
+        if old_id not in self._items:
+            raise ModuleError(f"Cannot supersede unknown claim id: {old_id!r}")
+        if new_claim.supersedes is None:
+            new_claim = new_claim.model_copy(update={"supersedes": old_id})
+        elif new_claim.supersedes != old_id:
+            raise ModuleError(
+                f"new_claim.supersedes={new_claim.supersedes!r} does not match old_id={old_id!r}"
+            )
+        self._items[new_claim.id] = new_claim
+        return new_claim.id
+
+    def stream(
+        self,
+        *,
+        project: str | None = None,
+        agent: str | None = None,
+        category: str | None = None,
+        run_id: str | None = None,
+    ) -> AsyncIterator[Claim]:
+        async def _agen() -> AsyncIterator[Claim]:
+            for claim in list(self._items.values()):
+                if project is not None and claim.project != project:
+                    continue
+                if agent is not None and claim.agent != agent:
+                    continue
+                if category is not None and claim.category != category:
+                    continue
+                if run_id is not None and claim.run_id != run_id:
+                    continue
+                yield claim
+
+        return _agen()
+
+    async def delete(
+        self,
+        *,
+        run_id: str | None = None,
+        older_than: datetime | None = None,
+        category: str | None = None,
+    ) -> int:
+        if run_id is None and older_than is None and category is None:
+            raise ModuleError(
+                "delete() requires at least one filter; refusing to wipe every claim."
+            )
+        keep: OrderedDict[str, Claim] = OrderedDict()
+        removed = 0
+        for cid, claim in self._items.items():
+            if run_id is not None and claim.run_id != run_id:
+                keep[cid] = claim
+                continue
+            if category is not None and claim.category != category:
+                keep[cid] = claim
+                continue
+            if older_than is not None and claim.created_at >= older_than:
+                keep[cid] = claim
+                continue
+            removed += 1
+        self._items = keep
+        return removed
+
+    async def close(self) -> None:
+        self._items.clear()