agentrust-py 0.0.3__py3-none-any.whl

agentrust_sdk/hooks.py

@@ -0,0 +1,461 @@
+"""
+Auto-instrumentation hooks — monkey-patch AI framework entry points so AgentTrust
+governance fires without any `@harness` decorator changes.
+
+Supported targets
+-----------------
+- **OpenAI SDK** >= 1.0 — ``openai.chat.completions.create`` (sync + async)
+- **LangChain Core (modern LCEL)** — ``langchain_core.runnables.base.Runnable.invoke``
+- **LangChain legacy** — ``langchain.llms.base.BaseLLM.predict`` + ``Chain.__call__``
+- **LangGraph** — ``langgraph.graph.state.CompiledStateGraph.invoke`` (auto-wrap compiled graphs)
+
+Usage
+-----
+Call once at application startup (before frameworks are used)::
+
+    from agentrust_sdk.hooks import auto_instrument
+    auto_instrument()   # patches all installed frameworks; silently skips missing ones
+
+Disable all patches::
+
+    AGENTRUST_AUTO_INSTRUMENT=false   # env var; takes effect before import
+
+Rollback / uninstall
+--------------------
+``remove_patches()`` restores **original callables** (not just clears the set).
+A **process restart** is the safest rollback if the process already served traffic.
+"""
+from __future__ import annotations
+
+import logging
+import os
+import threading
+from typing import Any
+
+from .config import SDK_CONFIG
+
+logger = logging.getLogger(__name__)
+
+# Registry: name → list of (obj, attr, original_callable)
+# Stored so remove_patches() can restore every original.
+# Protected by _REGISTRY_LOCK — auto_instrument() may be called from multiple
+# threads in e.g. gunicorn pre-fork workers.
+_PATCH_REGISTRY: dict[str, list[tuple[Any, str, Any]]] = {}
+_REGISTRY_LOCK = threading.Lock()
+
+
+def _is_patched(name: str) -> bool:
+    with _REGISTRY_LOCK:
+        return name in _PATCH_REGISTRY
+
+
+def auto_instrument(
+    *,
+    langchain: bool = True,
+    openai: bool = True,
+    langgraph: bool = True,
+    agent_id: str = "auto-instrumented",
+    gateway_url: str | None = None,
+    api_key: str | None = None,
+) -> list[str]:
+    """Patch installed AI frameworks to add AgentTrust governance.
+
+    Returns a list of framework names that were successfully patched.
+    Silently skips frameworks that are not installed.
+    """
+    if not SDK_CONFIG.enabled:
+        logger.debug("[AgentTrust] Auto-instrumentation skipped: SDK disabled")
+        return []
+
+    if os.environ.get("AGENTRUST_AUTO_INSTRUMENT", "true").lower() in ("0", "false", "no"):
+        logger.debug("[AgentTrust] Auto-instrumentation skipped: AGENTRUST_AUTO_INSTRUMENT=false")
+        return []
+
+    opts = {"agent_id": agent_id, "gateway_url": gateway_url, "api_key": api_key}
+    results = []
+
+    if openai and patch_openai(**opts):
+        results.append("openai")
+
+    if langchain and patch_langchain(**opts):
+        results.append("langchain")
+
+    if langgraph and patch_langgraph(**opts):
+        results.append("langgraph")
+
+    if results:
+        logger.info("[AgentTrust] Auto-instrumented: %s", ", ".join(results))
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+def _validate_output_sync(
+    output: Any,
+    input_text: str,
+    agent_id: str,
+    framework: str,
+    gateway_url: str | None,
+    api_key: str | None,
+    model: str = "unknown",
+) -> Any:
+    """Fire-and-forget validate call; swallows all errors (fail-open)."""
+    if not SDK_CONFIG.enabled:
+        return output
+    try:
+        from .client import AgentTrustClient
+        payload = output if isinstance(output, dict) else {"result": str(output)[:2000]}
+        with AgentTrustClient(
+            base_url=gateway_url or SDK_CONFIG.gateway_url,
+            api_key=api_key or SDK_CONFIG.api_key,
+        ) as client:
+            client.validate(
+                agent_id=agent_id,
+                user=framework.lower(),
+                input=str(input_text)[:500],
+                output=payload,
+                framework=framework,
+                model=model,
+            )
+    except Exception as exc:
+        logger.debug("[AgentTrust] %s hook skipped: %s", framework, exc)
+    return output
+
+
+async def _validate_output_async(
+    output: Any,
+    input_text: str,
+    agent_id: str,
+    framework: str,
+    gateway_url: str | None,
+    api_key: str | None,
+    model: str = "unknown",
+) -> Any:
+    if not SDK_CONFIG.enabled:
+        return output
+    try:
+        from .client import AsyncAgentTrustClient
+        payload = output if isinstance(output, dict) else {"result": str(output)[:2000]}
+        async with AsyncAgentTrustClient(
+            base_url=gateway_url or SDK_CONFIG.gateway_url,
+            api_key=api_key or SDK_CONFIG.api_key,
+        ) as client:
+            await client.validate(
+                agent_id=agent_id,
+                user=framework.lower(),
+                input=str(input_text)[:500],
+                output=payload,
+                framework=framework,
+                model=model,
+            )
+    except Exception as exc:
+        logger.debug("[AgentTrust] %s async hook skipped: %s", framework, exc)
+    return output
+
+
+# ---------------------------------------------------------------------------
+# OpenAI patch (SDK >= 1.0)
+# ---------------------------------------------------------------------------
+
+def patch_openai(
+    agent_id: str = "openai-auto",
+    gateway_url: str | None = None,
+    api_key: str | None = None,
+) -> bool:
+    """Patch ``openai.chat.completions.create`` (sync + async) to add governance."""
+    if _is_patched("openai"):
+        return True
+
+    try:
+        import openai as _openai
+        _sync_completions = _openai.chat.completions
+        _async_completions = _openai.AsyncOpenAI  # check import exists
+    except (ImportError, AttributeError):
+        return False
+
+    _orig_sync = _sync_completions.create
+
+    def _patched_sync(*args: Any, **kwargs: Any) -> Any:
+        result = _orig_sync(*args, **kwargs)
+        if not SDK_CONFIG.enabled:
+            return result
+        try:
+            output_text = ""
+            if hasattr(result, "choices") and result.choices:
+                msg = result.choices[0].message
+                output_text = getattr(msg, "content", "") or ""
+            input_text = ""
+            for m in kwargs.get("messages", []):
+                if isinstance(m, dict) and m.get("role") == "user":
+                    input_text = str(m.get("content", ""))[:500]
+                    break
+            _validate_output_sync(
+                {"content": output_text[:2000]}, input_text, agent_id,
+                "OpenAI", gateway_url, api_key, model=kwargs.get("model", "unknown"),
+            )
+        except Exception as exc:
+            logger.debug("[AgentTrust] OpenAI sync hook error: %s", exc)
+        return result
+
+    _sync_completions.create = _patched_sync  # type: ignore[method-assign]
+
+    # Also patch AsyncOpenAI client's chat completions create if accessible
+    try:
+        _async_client_cls = _openai.AsyncOpenAI
+        _orig_async_create = _async_client_cls.chat.completions.create
+
+        async def _patched_async(*args: Any, **kwargs: Any) -> Any:
+            result = await _orig_async_create(*args, **kwargs)
+            if not SDK_CONFIG.enabled:
+                return result
+            try:
+                output_text = ""
+                if hasattr(result, "choices") and result.choices:
+                    msg = result.choices[0].message
+                    output_text = getattr(msg, "content", "") or ""
+                input_text = ""
+                for m in kwargs.get("messages", []):
+                    if isinstance(m, dict) and m.get("role") == "user":
+                        input_text = str(m.get("content", ""))[:500]
+                        break
+                await _validate_output_async(
+                    {"content": output_text[:2000]}, input_text, agent_id,
+                    "OpenAI", gateway_url, api_key, model=kwargs.get("model", "unknown"),
+                )
+            except Exception as exc:
+                logger.debug("[AgentTrust] OpenAI async hook error: %s", exc)
+            return result
+
+        _async_client_cls.chat.completions.create = _patched_async  # type: ignore[method-assign]
+        with _REGISTRY_LOCK:
+            _PATCH_REGISTRY["openai"] = [
+                (_sync_completions, "create", _orig_sync),
+                (_async_client_cls.chat.completions, "create", _orig_async_create),
+            ]
+    except Exception:
+        with _REGISTRY_LOCK:
+            _PATCH_REGISTRY["openai"] = [(_sync_completions, "create", _orig_sync)]
+
+    logger.debug("[AgentTrust] OpenAI auto-instrumentation applied (sync + async)")
+    return True
+
+
+# ---------------------------------------------------------------------------
+# LangChain patch — modern LCEL (langchain_core) + legacy fallback
+# ---------------------------------------------------------------------------
+
+def patch_langchain(
+    agent_id: str = "langchain-auto",
+    gateway_url: str | None = None,
+    api_key: str | None = None,
+) -> bool:
+    """Patch LangChain to add governance on invoke calls.
+
+    Targets (in order of preference):
+    1. ``langchain_core.runnables.base.Runnable.invoke``  — modern LCEL (>= 0.1)
+    2. ``langchain.llms.base.BaseLLM.predict``            — legacy LLMs
+    3. ``langchain.chains.base.Chain.__call__``           — legacy chains
+    """
+    if _is_patched("langchain"):
+        return True
+
+    targets: list[tuple[Any, str, Any]] = []
+
+    # ── 1. langchain_core LCEL (modern, preferred) ─────────────────────────
+    try:
+        import langchain_core.runnables.base as _runnables
+        _orig_invoke = _runnables.Runnable.invoke
+
+        def _patched_lcel_invoke(self: Any, input: Any, config: Any = None, **kwargs: Any) -> Any:
+            result = _orig_invoke(self, input, config, **kwargs)
+            if not SDK_CONFIG.enabled:
+                return result
+            input_text = str(input)[:500] if not isinstance(input, dict) else str(input.get("input", input))[:500]
+            _validate_output_sync(result, input_text, agent_id, "LangChain", gateway_url, api_key)
+            return result
+
+        _runnables.Runnable.invoke = _patched_lcel_invoke  # type: ignore[method-assign]
+        targets.append((_runnables.Runnable, "invoke", _orig_invoke))
+        logger.debug("[AgentTrust] LangChain LCEL (langchain_core) patched")
+    except (ImportError, AttributeError):
+        pass
+
+    # ── 2. Legacy BaseLLM.predict ───────────────────────────────────────────
+    try:
+        import langchain.llms.base as _llm_base
+        _orig_predict = getattr(_llm_base.BaseLLM, "predict", None)
+        if _orig_predict:
+            def _patched_predict(self: Any, text: str, **kwargs: Any) -> str:
+                result = _orig_predict(self, text, **kwargs)
+                _validate_output_sync(result, text, agent_id, "LangChain", gateway_url, api_key)
+                return result
+            _llm_base.BaseLLM.predict = _patched_predict  # type: ignore[method-assign]
+            targets.append((_llm_base.BaseLLM, "predict", _orig_predict))
+    except (ImportError, AttributeError):
+        pass
+
+    # ── 3. Legacy Chain.__call__ ────────────────────────────────────────────
+    try:
+        import langchain.chains.base as _chain_base
+        _orig_call = getattr(_chain_base.Chain, "__call__", None)
+        if _orig_call:
+            def _patched_chain_call(self: Any, inputs: Any, **kwargs: Any) -> Any:
+                result = _orig_call(self, inputs, **kwargs)
+                input_text = str(inputs) if not isinstance(inputs, dict) else str(inputs.get("input", inputs))
+                _validate_output_sync(result, input_text, agent_id, "LangChain", gateway_url, api_key)
+                return result
+            _chain_base.Chain.__call__ = _patched_chain_call  # type: ignore[method-assign]
+            targets.append((_chain_base.Chain, "__call__", _orig_call))
+    except (ImportError, AttributeError):
+        pass
+
+    if not targets:
+        return False  # no LangChain variant found
+
+    with _REGISTRY_LOCK:
+        _PATCH_REGISTRY["langchain"] = targets
+    return True
+
+
+# ---------------------------------------------------------------------------
+# LangGraph patch — CompiledStateGraph.invoke / ainvoke (P0-3)
+# ---------------------------------------------------------------------------
+
+def patch_langgraph(
+    agent_id: str = "langgraph-auto",
+    gateway_url: str | None = None,
+    api_key: str | None = None,
+) -> bool:
+    """Patch LangGraph ``CompiledStateGraph.invoke`` and ``ainvoke`` for governance.
+
+    This covers auto-instrumentation without needing ``AgentTrustNode`` wiring.
+    Any ``graph.invoke(state)`` call will be governed.
+    """
+    if _is_patched("langgraph"):
+        return True
+
+    try:
+        from langgraph.graph.state import CompiledStateGraph as _CSG
+    except ImportError:
+        try:
+            # fallback for older langgraph layouts
+            from langgraph.pregel import Pregel as _CSG  # type: ignore[assignment,no-redef]
+        except ImportError:
+            return False
+
+    targets: list[tuple[Any, str, Any]] = []
+
+    _orig_invoke = getattr(_CSG, "invoke", None)
+    _orig_ainvoke = getattr(_CSG, "ainvoke", None)
+
+    if _orig_invoke:
+        def _patched_invoke(self: Any, input: Any, config: Any = None, **kwargs: Any) -> Any:
+            result = _orig_invoke(self, input, config, **kwargs)
+            if not SDK_CONFIG.enabled:
+                return result
+            input_text = str(input)[:500] if not isinstance(input, dict) else str(input.get("input", input))[:500]
+            _validate_output_sync(result, input_text, agent_id, "LangGraph", gateway_url, api_key)
+            return result
+
+        _CSG.invoke = _patched_invoke  # type: ignore[method-assign]
+        targets.append((_CSG, "invoke", _orig_invoke))
+
+    if _orig_ainvoke:
+        async def _patched_ainvoke(self: Any, input: Any, config: Any = None, **kwargs: Any) -> Any:
+            result = await _orig_ainvoke(self, input, config, **kwargs)
+            if not SDK_CONFIG.enabled:
+                return result
+            input_text = str(input)[:500] if not isinstance(input, dict) else str(input.get("input", input))[:500]
+            await _validate_output_async(result, input_text, agent_id, "LangGraph", gateway_url, api_key)
+            return result
+
+        _CSG.ainvoke = _patched_ainvoke  # type: ignore[method-assign]
+        targets.append((_CSG, "ainvoke", _orig_ainvoke))
+
+    if not targets:
+        return False
+
+    with _REGISTRY_LOCK:
+        _PATCH_REGISTRY["langgraph"] = targets
+    logger.debug("[AgentTrust] LangGraph auto-instrumentation applied")
+    return True
+
+
+# ---------------------------------------------------------------------------
+# LangGraph compile-time wrapper (P0-3)
+# ---------------------------------------------------------------------------
+
+def auto_wrap(compiled_graph: Any, agent_id: str = "langgraph-wrapped") -> Any:
+    """Wrap a compiled LangGraph graph so every invoke/ainvoke is governed.
+
+    Usage — zero-touch: pass your compiled graph through this helper once::
+
+        from agentrust_sdk.hooks import auto_wrap
+
+        graph = workflow.compile()
+        graph = auto_wrap(graph, agent_id="my-agent")  # governance on every invoke
+
+        # No other code changes needed
+        result = graph.invoke({"input": "hello", "user": "alice"})
+
+    This is lighter than the global ``patch_langgraph()`` monkey-patch: it wraps
+    only the specific graph instance, not all CompiledStateGraph objects.
+    """
+    _orig_invoke = getattr(compiled_graph, "invoke", None)
+    _orig_ainvoke = getattr(compiled_graph, "ainvoke", None)
+    _orig_stream = getattr(compiled_graph, "stream", None)
+
+    opts = {"agent_id": agent_id, "gateway_url": None, "api_key": SDK_CONFIG.api_key}
+
+    if _orig_invoke:
+        def _wrapped_invoke(input: Any, config: Any = None, **kwargs: Any) -> Any:
+            result = _orig_invoke(input, config, **kwargs)
+            if SDK_CONFIG.enabled:
+                input_text = str(input)[:500] if not isinstance(input, dict) else str(input.get("input", input))[:500]
+                _validate_output_sync(result, input_text, agent_id, "LangGraph", None, SDK_CONFIG.api_key)
+            return result
+        compiled_graph.invoke = _wrapped_invoke
+
+    if _orig_ainvoke:
+        async def _wrapped_ainvoke(input: Any, config: Any = None, **kwargs: Any) -> Any:
+            result = await _orig_ainvoke(input, config, **kwargs)
+            if SDK_CONFIG.enabled:
+                input_text = str(input)[:500] if not isinstance(input, dict) else str(input.get("input", input))[:500]
+                await _validate_output_async(result, input_text, agent_id, "LangGraph", None, SDK_CONFIG.api_key)
+            return result
+        compiled_graph.ainvoke = _wrapped_ainvoke
+
+    logger.debug("[AgentTrust] auto_wrap applied to graph instance (agent_id=%s)", agent_id)
+    return compiled_graph
+
+
+# ---------------------------------------------------------------------------
+# Patch removal — restores original callables (P1-3 fix)
+# ---------------------------------------------------------------------------
+
+def remove_patches() -> list[str]:
+    """Restore all patched callables to their originals.
+
+    Returns the list of framework names that were restored.
+    Note: a **process restart** is still the safest rollback if the process
+    already served traffic through patched paths.
+    """
+    with _REGISTRY_LOCK:
+        snapshot = list(_PATCH_REGISTRY.items())
+
+    restored = []
+    for name, targets in snapshot:
+        for obj, attr, original in targets:
+            try:
+                setattr(obj, attr, original)
+            except Exception as exc:
+                logger.warning("[AgentTrust] Could not restore %s.%s: %s", obj, attr, exc)
+        restored.append(name)
+
+    with _REGISTRY_LOCK:
+        _PATCH_REGISTRY.clear()
+    if restored:
+        logger.info("[AgentTrust] Removed patches: %s", ", ".join(restored))
+    return restored

agentrust_sdk/models.py

@@ -0,0 +1,81 @@
+"""SDK-facing request/response models."""
+from __future__ import annotations
+
+from typing import Any
+from pydantic import BaseModel, Field
+
+
+class ToolCall(BaseModel):
+    name: str
+    arguments: dict[str, Any] = Field(default_factory=dict)
+    result: Any = None
+    latency_ms: float | None = None
+    error: str | None = None
+
+
+class ValidationResult(BaseModel):
+    schema_score: float = 0.0
+    evidence_score: float = 0.0
+    tool_trust_score: float = 0.0
+    consistency_score: float = 0.0
+    policy_score: float = 0.0
+    judge_score: float | None = None
+    final_confidence: float = 0.0
+    failures: list[str] = Field(default_factory=list)
+
+
+class RiskResult(BaseModel):
+    tier: str = "unknown"   # "unknown" = not computed (tier too low)
+    score: float = 0.0
+    reason: str = ""
+
+
+class DecisionResult(BaseModel):
+    outcome: str = "pending"  # "pending" = not computed (tier too low)
+    reason: str = ""
+    policy_version: str = ""
+
+
+class ValidateRequest(BaseModel):
+    agent_id: str
+    framework: str = "REST"
+    version: str = "1"
+    parent_envelope_id: str | None = None
+    user: str
+    input: str
+    output: dict[str, Any] = Field(default_factory=dict)
+    model: str = "unknown"
+    tools_called: list[ToolCall] = Field(default_factory=list)
+    latency_ms: float = 0.0
+    tokens: int = 0
+    session_id: str | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class ValidateResponse(BaseModel):
+    envelope_id: str
+    validation: ValidationResult
+    risk: RiskResult
+    decision: DecisionResult
+    latency_ms: float
+    trust_chain: dict | None = None
+    # Tier metadata — always populated
+    tier_info: str = "unknown"
+    upgrade_hint: str | None = None
+
+    @property
+    def approved(self) -> bool:
+        return self.decision.outcome == "approve"
+
+    @property
+    def blocked(self) -> bool:
+        return self.decision.outcome == "block"
+
+    @property
+    def needs_review(self) -> bool:
+        return self.decision.outcome in ("escalate", "request_evidence")
+
+    @property
+    def schema_valid(self) -> bool:
+        """Available on all tiers including OSS."""
+        return self.validation.schema_score >= 80.0 and not self.validation.failures