agentrust-py 0.0.3__py3-none-any.whl

agentrust_sdk/embedded.py

@@ -0,0 +1,428 @@
+"""
+AgentTrust EmbeddedGateway — zero-dependency local governance.
+
+Ships as part of agentrust[embedded].  Starts a minimal FastAPI app in a
+background thread backed by SQLite — no PostgreSQL, no Redis, no Docker.
+
+Usage::
+
+    from agentrust import embed_gateway, harness
+
+    embed_gateway()          # starts in-process gateway on :8765
+
+    @harness
+    def my_agent(user: str, input: str) -> dict:
+        return {"answer": "42"}
+
+The gateway runs only the fast deterministic checks (schema + basic policy).
+For full risk scoring, LLM judge, and analytics, deploy the full gateway.
+
+Stop it::
+
+    gw = embed_gateway()     # returns EmbeddedGateway instance
+    gw.stop()
+
+Or use as context manager::
+
+    with EmbeddedGateway() as gw:
+        ...
+
+Env vars respected
+------------------
+AGENTRUST_EMBED_PORT     Port (default 8765)
+AGENTRUST_EMBED_DB       SQLite db path (default ~/.agentrust/embedded.db)
+                         Use ":memory:" for ephemeral (lost on stop)
+AGENTRUST_EMBED_TOKEN    Bearer token required on all non-health endpoints.
+                         When unset, a random token is generated at startup and
+                         logged at INFO level. Set explicitly for reproducible auth.
+AGENTRUST_ENABLED        When false, embed_gateway() is a no-op
+"""
+from __future__ import annotations
+
+import logging
+import os
+import secrets
+import sqlite3
+import threading
+import time
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_PORT = int(os.environ.get("AGENTRUST_EMBED_PORT", "8765"))
+_DEFAULT_DB = os.environ.get(
+    "AGENTRUST_EMBED_DB",
+    str(Path.home() / ".agentrust" / "embedded.db"),
+)
+_MAX_LIST_LIMIT = 200  # server-side cap on /v1/audit/executions?limit=
+
+# Module-level singleton so embed_gateway() is idempotent
+_INSTANCE: "EmbeddedGateway | None" = None
+
+
+class EmbeddedGateway:
+    """
+    Minimal in-process governance gateway backed by SQLite.
+
+    Supports:
+      - Schema validation (output must be non-empty JSON object)
+      - Basic policy checks (output keys, type enforcement)
+      - Append-only SQLite audit log
+      - governance_disclosure + confidence_rationale fields
+      - All 5 decision outcomes: approve/retry/request_evidence/escalate/block
+
+    Does NOT support (requires full gateway):
+      - LLM judge slow-path
+      - Risk engine (scores all LOW)
+      - Human review queue
+      - Webhook notifications
+      - Multi-tenant isolation
+    """
+
+    def __init__(
+        self,
+        port: int = _DEFAULT_PORT,
+        db_path: str = _DEFAULT_DB,
+        token: str | None = None,
+    ) -> None:
+        self._port = port
+        self._db_path = db_path
+        self._server: Any = None
+        self._thread: threading.Thread | None = None
+        self._ready = threading.Event()
+        # A-3: bearer token for all non-health endpoints.
+        # Caller may pass an explicit token; otherwise a random one is generated
+        # and emitted to logs so the SDK client can pick it up.
+        env_token = os.environ.get("AGENTRUST_EMBED_TOKEN", "").strip()
+        self._token: str = token or env_token or secrets.token_hex(32)
+        if not (token or env_token):
+            logger.info(
+                "[AgentTrust] EmbeddedGateway auto-generated bearer token: %s  "
+                "(set AGENTRUST_EMBED_TOKEN env var to pin this across restarts)",
+                self._token,
+            )
+        self._app = self._build_app()
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def start(self) -> "EmbeddedGateway":
+        if self._thread and self._thread.is_alive():
+            return self
+        self._init_db()
+        self._thread = threading.Thread(
+            target=self._run_server, daemon=True, name="agentrust-embedded"
+        )
+        self._thread.start()
+        if not self._ready.wait(timeout=10):
+            raise RuntimeError(
+                f"EmbeddedGateway did not start within 10s on port {self._port}"
+            )
+        logger.info(
+            "[AgentTrust] EmbeddedGateway started on http://127.0.0.1:%d "
+            "(SQLite: %s)", self._port, self._db_path
+        )
+        return self
+
+    def stop(self) -> None:
+        if self._server:
+            self._server.should_exit = True
+        if self._thread:
+            self._thread.join(timeout=5)
+        logger.info("[AgentTrust] EmbeddedGateway stopped.")
+
+    def __enter__(self) -> "EmbeddedGateway":
+        return self.start()
+
+    def __exit__(self, *args: Any) -> None:
+        self.stop()
+
+    # ------------------------------------------------------------------
+    # Internal: SQLite
+    # ------------------------------------------------------------------
+
+    def _init_db(self) -> None:
+        if self._db_path == ":memory:":
+            # Will be created per-connection; not shareable across threads,
+            # but fine for unit tests. Use file path for persistence.
+            return
+        Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
+        con = sqlite3.connect(self._db_path)
+        con.execute("PRAGMA journal_mode=WAL")
+        con.execute("PRAGMA synchronous=NORMAL")
+        con.execute(
+            """
+            CREATE TABLE IF NOT EXISTS executions (
+                id          INTEGER PRIMARY KEY AUTOINCREMENT,
+                envelope_id TEXT    NOT NULL,
+                agent_id    TEXT    NOT NULL,
+                decision    TEXT    NOT NULL,
+                risk_tier   TEXT    NOT NULL DEFAULT 'low',
+                final_confidence REAL NOT NULL DEFAULT 100.0,
+                request_json     TEXT,
+                output_json      TEXT,
+                timestamp   TEXT    NOT NULL DEFAULT (datetime('now'))
+            )
+            """
+        )
+        con.commit()
+        con.close()
+
+    def _save_execution(self, envelope_id: str, agent_id: str,
+                        decision: str, payload: dict[str, Any]) -> None:
+        import json
+        if self._db_path == ":memory:":
+            return
+        try:
+            con = sqlite3.connect(self._db_path)
+            con.execute(
+                "INSERT INTO executions "
+                "(envelope_id, agent_id, decision, request_json, output_json) "
+                "VALUES (?, ?, ?, ?, ?)",
+                (
+                    envelope_id, agent_id, decision,
+                    json.dumps(payload.get("request", {})),
+                    json.dumps(payload.get("output", {})),
+                ),
+            )
+            con.commit()
+            con.close()
+        except Exception as exc:
+            logger.warning("[AgentTrust] Embedded audit write failed: %s", exc)
+
+    # ------------------------------------------------------------------
+    # Internal: FastAPI app
+    # ------------------------------------------------------------------
+
+    def _build_app(self) -> Any:
+        try:
+            from fastapi import FastAPI, Request
+            from fastapi.responses import JSONResponse
+            from starlette.middleware.base import BaseHTTPMiddleware
+        except ImportError:
+            raise ImportError(
+                "FastAPI is required for EmbeddedGateway. "
+                "Install it: pip install 'agentrust-sdk[embedded]'"
+            )
+
+        app = FastAPI(
+            title="AgentTrust EmbeddedGateway",
+            version="embedded",
+            docs_url=None,
+            redoc_url=None,
+            openapi_url=None,
+        )
+        _gw = self  # closure reference
+
+        # Reject oversized request bodies (default 1 MB) to prevent memory exhaustion
+        _MAX_BODY_BYTES = 1 * 1024 * 1024
+
+        class _BodySizeLimiter(BaseHTTPMiddleware):
+            async def dispatch(self, request: Request, call_next):
+                content_length = request.headers.get("content-length")
+                if content_length and int(content_length) > _MAX_BODY_BYTES:
+                    return JSONResponse(
+                        {"detail": f"Request body too large (max {_MAX_BODY_BYTES} bytes)"},
+                        status_code=413,
+                    )
+                return await call_next(request)
+
+        # A-3: bearer token enforcement on all non-health endpoints
+        _expected_token = _gw._token
+
+        class _TokenAuth(BaseHTTPMiddleware):
+            async def dispatch(self, request: Request, call_next):
+                if request.url.path == "/v1/health":
+                    return await call_next(request)
+                auth = request.headers.get("authorization", "")
+                provided = auth.removeprefix("Bearer ").strip()
+                if not secrets.compare_digest(provided, _expected_token):
+                    return JSONResponse(
+                        {"detail": "Invalid or missing Bearer token"},
+                        status_code=401,
+                    )
+                return await call_next(request)
+
+        app.add_middleware(_BodySizeLimiter)
+        app.add_middleware(_TokenAuth)
+
+        @app.get("/v1/health")
+        async def health() -> dict:
+            return {"status": "ok", "mode": "embedded", "db": _gw._db_path}
+
+        @app.post("/v1/runtime/validate")
+        async def validate(body: dict) -> dict:
+            return await _gw._handle_validate(body)
+
+        @app.get("/v1/audit/executions")
+        async def list_executions(limit: int = 20) -> dict:
+            # Server-side cap to prevent runaway memory reads
+            return _gw._list_executions(min(limit, _MAX_LIST_LIMIT))
+
+        return app
+
+    async def _handle_validate(self, body: dict[str, Any]) -> dict[str, Any]:
+        envelope_id = str(uuid4())
+        agent_id = body.get("agent_id", "unknown")
+        output = body.get("output", {})
+        request = body.get("request", {})
+
+        failures: list[str] = []
+        schema_score = 100.0
+
+        # Schema check
+        if not isinstance(output, dict):
+            failures.append("schema: output must be a JSON object")
+            schema_score = 0.0
+        elif not output:
+            failures.append("schema: output is empty — agent produced no result")
+            schema_score = 40.0
+
+        # Basic policy check: no error keys in output
+        if isinstance(output, dict) and output.get("error"):
+            failures.append("policy: output contains error field")
+            schema_score = min(schema_score, 50.0)
+
+        final_confidence = schema_score
+        decision = "approve" if not failures else ("block" if schema_score == 0 else "retry")
+
+        governance_disclosure = (
+            f"Evaluated by AgentTrust EmbeddedGateway "
+            f"(envelope_id={envelope_id}, agent_id={agent_id}). "
+            f"Decision: {decision}. Confidence: {final_confidence:.1f}%. "
+            f"Mode: embedded (schema + basic policy checks only)."
+        )
+
+        self._save_execution(envelope_id, agent_id, decision, body)
+
+        return {
+            "envelope_id": envelope_id,
+            "validation": {
+                "schema_score": schema_score,
+                "evidence_score": 0.0,
+                "tool_trust_score": 0.0,
+                "consistency_score": 0.0,
+                "policy_score": schema_score,
+                "final_confidence": final_confidence,
+                "failures": failures,
+            },
+            "risk": {
+                "tier": "low",
+                "score": 0.0,
+                "reason": "Embedded mode: risk engine not available",
+            },
+            "decision": {
+                "outcome": decision,
+                "reason": failures[0] if failures else "All embedded checks passed",
+                "policy_version": "embedded-v1",
+            },
+            "latency_ms": 0.0,
+            "governance_disclosure": governance_disclosure,
+            "confidence_rationale": (
+                f"Schema check: {schema_score:.0f}/100. "
+                "Full confidence engine requires full gateway deployment."
+            ),
+            "trust_chain": None,
+        }
+
+    def _list_executions(self, limit: int) -> dict[str, Any]:
+        if self._db_path == ":memory:":
+            return {"items": [], "note": "In-memory mode: no persisted records"}
+        try:
+            con = sqlite3.connect(self._db_path)
+            con.row_factory = sqlite3.Row
+            rows = con.execute(
+                "SELECT * FROM executions ORDER BY timestamp DESC LIMIT ?", (limit,)
+            ).fetchall()
+            con.close()
+            return {"items": [dict(r) for r in rows], "total": len(rows)}
+        except Exception as exc:
+            return {"items": [], "error": str(exc)}
+
+    # ------------------------------------------------------------------
+    # Internal: server runner
+    # ------------------------------------------------------------------
+
+    def _run_server(self) -> None:
+        try:
+            import uvicorn
+        except ImportError:
+            raise ImportError(
+                "uvicorn is required for EmbeddedGateway. "
+                "Install it: pip install 'agentrust-sdk[embedded]'"
+            )
+
+        config = uvicorn.Config(
+            app=self._app,
+            host="127.0.0.1",
+            port=self._port,
+            log_level="error",
+            access_log=False,
+        )
+        self._server = uvicorn.Server(config)
+
+        # Signal ready once the server loop is about to start
+        original_startup = self._server.startup
+
+        async def _startup_with_signal(*args: Any, **kw: Any) -> None:
+            await original_startup(*args, **kw)
+            self._ready.set()
+
+        self._server.startup = _startup_with_signal
+        self._server.run()
+
+
+# ---------------------------------------------------------------------------
+# Public helper
+# ---------------------------------------------------------------------------
+
+def embed_gateway(
+    port: int = _DEFAULT_PORT,
+    db_path: str = _DEFAULT_DB,
+    set_env: bool = True,
+) -> "EmbeddedGateway":
+    """
+    Start an in-process EmbeddedGateway and (optionally) set AGENTRUST_GATEWAY_URL
+    so all SDK clients automatically route to it.
+
+    Idempotent — calling twice returns the same instance.
+
+    Parameters
+    ----------
+    port:     Port to listen on (default 8765, configurable via AGENTRUST_EMBED_PORT)
+    db_path:  SQLite database file path (default ~/.agentrust/embedded.db)
+    set_env:  When True, sets AGENTRUST_GATEWAY_URL=http://127.0.0.1:{port}
+              so @harness picks it up automatically.
+
+    Returns the running EmbeddedGateway instance.
+    """
+    from .config import SDK_CONFIG
+
+    if not SDK_CONFIG.enabled:
+        logger.info("[AgentTrust] AGENTRUST_ENABLED=false — embed_gateway() is a no-op")
+
+        class _NoOp:
+            def stop(self) -> None: ...
+            def __enter__(self): return self
+            def __exit__(self, *a: Any): ...
+
+        return _NoOp()  # type: ignore[return-value]
+
+    global _INSTANCE
+    if _INSTANCE is not None:
+        return _INSTANCE
+
+    gw = EmbeddedGateway(port=port, db_path=db_path)
+    gw.start()
+
+    if set_env:
+        # SDK_CONFIG.gateway_url and api_key are properties that read from os.environ,
+        # so setting env vars here is sufficient for all clients created afterwards.
+        os.environ["AGENTRUST_GATEWAY_URL"] = f"http://127.0.0.1:{port}"
+        os.environ["AGENTRUST_KEY"] = gw._token
+
+    _INSTANCE = gw
+    return gw