agentrust-py 0.0.3__py3-none-any.whl

agentrust_sdk/auto.py

@@ -0,0 +1,397 @@
+"""
+AgentTrust transparent auto-instrumentation.
+
+Drop-in patching for OpenAI, LangChain, and Anthropic SDKs — no changes
+required in user code.
+
+Quick start::
+
+    import agentrust_sdk.auto as at_auto
+    at_auto.install(agent_id="my-agent", api_key="at-...")
+
+Or use the convenience alias::
+
+    from agentrust_sdk.auto import auto_install
+    auto_install(agent_id="my-agent")
+
+After calling ``install()`` every call to:
+
+* ``openai.OpenAI().chat.completions.create`` / ``AsyncOpenAI`` variant
+* ``anthropic.Anthropic().messages.create`` / ``AsyncAnthropic`` variant
+* Any LangChain LLM / chain / tool invocation (via ``BaseCallbackHandler``)
+
+…will be transparently wrapped with ``AgentTrustClient.validate()``.
+
+The original return value is **always** preserved.  Validation errors are
+logged and swallowed (or raise, depending on ``failure_mode``).
+"""
+from __future__ import annotations
+
+import functools
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+# Module-level set that tracks which patches have been applied so that
+# calling install() multiple times is a no-op.
+_INSTALLED: set[str] = set()
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_client(agent_id: str, api_key: str | None, gateway_url: str | None, failure_mode: str):
+    """Lazily import and construct an AgentTrustClient."""
+    from .client import AgentTrustClient  # local import to avoid circular deps
+
+    kwargs: dict[str, Any] = {"failure_mode": failure_mode}
+    if api_key:
+        kwargs["api_key"] = api_key
+    if gateway_url:
+        kwargs["gateway_url"] = gateway_url
+    return AgentTrustClient(**kwargs), agent_id
+
+
+def _safe_validate(client, agent_id: str, user: str, input_text: str, output_text: str) -> None:
+    """Call validate() and swallow any exception so the caller is unaffected."""
+    try:
+        client.validate(
+            agent_id=agent_id,
+            user=user,
+            input=input_text,
+            output=output_text,
+        )
+    except Exception as exc:  # noqa: BLE001
+        logger.debug("AgentTrust validation error (non-fatal): %s", exc)
+
+
+async def _safe_validate_async(client, agent_id: str, user: str, input_text: str, output_text: str) -> None:
+    """Async version of _safe_validate."""
+    try:
+        await client.validate_async(
+            agent_id=agent_id,
+            user=user,
+            input=input_text,
+            output=output_text,
+        )
+    except Exception as exc:  # noqa: BLE001
+        logger.debug("AgentTrust async validation error (non-fatal): %s", exc)
+
+
+def _extract_openai_prompt(kwargs: dict[str, Any]) -> str:
+    """Best-effort extraction of a human-readable prompt from OpenAI messages."""
+    messages = kwargs.get("messages", [])
+    if not messages:
+        return ""
+    parts = []
+    for m in messages:
+        role = m.get("role", "")
+        content = m.get("content", "")
+        if isinstance(content, list):
+            # vision / multi-part content
+            content = " ".join(
+                p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text"
+            )
+        parts.append(f"{role}: {content}")
+    return "\n".join(parts)
+
+
+def _extract_openai_response(response: Any) -> str:
+    """Extract text from an OpenAI ChatCompletion response object."""
+    try:
+        return response.choices[0].message.content or ""
+    except Exception:
+        return str(response)
+
+
+# ---------------------------------------------------------------------------
+# Patch 1 — OpenAI SDK (openai >= 1.0)
+# ---------------------------------------------------------------------------
+
+def _patch_openai(client_at, agent_id: str) -> None:
+    """Wrap openai.OpenAI and openai.AsyncOpenAI completions.create."""
+    if "openai" in _INSTALLED:
+        return
+    try:
+        import openai  # noqa: PLC0415
+    except ImportError:
+        logger.debug("openai not installed — skipping AgentTrust OpenAI patch")
+        return
+
+    # --- sync ---
+    _orig_sync_create = openai.resources.chat.completions.Completions.create
+
+    @functools.wraps(_orig_sync_create)
+    def _sync_create(self, *args: Any, **kwargs: Any) -> Any:
+        result = _orig_sync_create(self, *args, **kwargs)
+        prompt = _extract_openai_prompt(kwargs)
+        output = _extract_openai_response(result)
+        _safe_validate(client_at, agent_id, user="openai-sync", input_text=prompt, output_text=output)
+        return result
+
+    openai.resources.chat.completions.Completions.create = _sync_create  # type: ignore[method-assign]
+
+    # --- async ---
+    try:
+        _orig_async_create = openai.resources.chat.completions.AsyncCompletions.create
+
+        @functools.wraps(_orig_async_create)
+        async def _async_create(self, *args: Any, **kwargs: Any) -> Any:
+            result = await _orig_async_create(self, *args, **kwargs)
+            prompt = _extract_openai_prompt(kwargs)
+            output = _extract_openai_response(result)
+            await _safe_validate_async(client_at, agent_id, user="openai-async", input_text=prompt, output_text=output)
+            return result
+
+        openai.resources.chat.completions.AsyncCompletions.create = _async_create  # type: ignore[method-assign]
+    except AttributeError:
+        logger.debug("AsyncCompletions not found in openai — async patch skipped")
+
+    _INSTALLED.add("openai")
+    logger.info("AgentTrust: OpenAI patch installed")
+
+
+# ---------------------------------------------------------------------------
+# Patch 2 — LangChain (langchain >= 0.1)
+# ---------------------------------------------------------------------------
+
+class AgentTrustCallbackHandler:
+    """
+    LangChain BaseCallbackHandler that validates LLM / chain / tool outputs
+    via AgentTrust.
+
+    This class inherits from ``langchain_core.callbacks.BaseCallbackHandler``
+    only when LangChain is available.  When LangChain is absent the class is
+    defined as a plain Python class so that the module can still be imported.
+    """
+
+    # Filled in by _patch_langchain() after the real base class is resolved.
+    _client_at: Any = None
+    _agent_id: str = "auto"
+
+    def on_llm_end(self, response: Any, **kwargs: Any) -> None:  # type: ignore[override]
+        try:
+            # LLMResult stores generations as a list-of-list of Generation objects
+            for gen_list in getattr(response, "generations", []):
+                for gen in gen_list:
+                    text = getattr(gen, "text", "") or ""
+                    _safe_validate(
+                        self._client_at,
+                        self._agent_id,
+                        user="langchain-llm",
+                        input_text="",
+                        output_text=text,
+                    )
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("AgentTrust LangChain on_llm_end error: %s", exc)
+
+    def on_chain_end(self, outputs: Any, **kwargs: Any) -> None:  # type: ignore[override]
+        try:
+            output_text = str(outputs) if outputs is not None else ""
+            _safe_validate(
+                self._client_at,
+                self._agent_id,
+                user="langchain-chain",
+                input_text="",
+                output_text=output_text,
+            )
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("AgentTrust LangChain on_chain_end error: %s", exc)
+
+    def on_tool_end(self, output: Any, **kwargs: Any) -> None:  # type: ignore[override]
+        try:
+            output_text = str(output) if output is not None else ""
+            _safe_validate(
+                self._client_at,
+                self._agent_id,
+                user="langchain-tool",
+                input_text="",
+                output_text=output_text,
+            )
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("AgentTrust LangChain on_tool_end error: %s", exc)
+
+
+def _patch_langchain(client_at, agent_id: str) -> None:
+    """Register AgentTrustCallbackHandler as a global LangChain callback."""
+    if "langchain" in _INSTALLED:
+        return
+    try:
+        from langchain_core.callbacks import BaseCallbackHandler  # noqa: PLC0415
+    except ImportError:
+        try:
+            from langchain.callbacks import BaseCallbackHandler  # type: ignore[no-redef]  # noqa: PLC0415
+        except ImportError:
+            logger.debug("langchain not installed — skipping AgentTrust LangChain patch")
+            return
+
+    # Re-define the handler as a proper subclass now that we have the base.
+    class _Handler(BaseCallbackHandler):  # type: ignore[misc]
+        def on_llm_end(self, response: Any, **kwargs: Any) -> None:
+            AgentTrustCallbackHandler.on_llm_end(self, response, **kwargs)
+
+        def on_chain_end(self, outputs: Any, **kwargs: Any) -> None:
+            AgentTrustCallbackHandler.on_chain_end(self, outputs, **kwargs)
+
+        def on_tool_end(self, output: Any, **kwargs: Any) -> None:
+            AgentTrustCallbackHandler.on_tool_end(self, output, **kwargs)
+
+    handler = _Handler()
+    handler._client_at = client_at  # type: ignore[attr-defined]
+    handler._agent_id = agent_id  # type: ignore[attr-defined]
+
+    # Expose the concrete subclass so users can import it if needed.
+    AgentTrustCallbackHandler.__bases__ = (BaseCallbackHandler,)  # type: ignore[misc]
+
+    # Try to register globally via the callback manager.
+    _registered = False
+    try:
+        from langchain.callbacks import set_handler  # noqa: PLC0415
+        set_handler(handler)
+        _registered = True
+    except (ImportError, AttributeError):
+        pass
+
+    if not _registered:
+        try:
+            from langchain_core.callbacks import get_callback_manager  # noqa: PLC0415
+            get_callback_manager().add_handler(handler)
+            _registered = True
+        except (ImportError, AttributeError):
+            pass
+
+    if not _registered:
+        logger.debug(
+            "AgentTrust: could not find a global LangChain callback registry; "
+            "add AgentTrustCallbackHandler() manually to your chain's callbacks."
+        )
+
+    _INSTALLED.add("langchain")
+    logger.info("AgentTrust: LangChain patch installed (global registration=%s)", _registered)
+
+
+# ---------------------------------------------------------------------------
+# Patch 3 — Anthropic SDK (anthropic >= 0.20)
+# ---------------------------------------------------------------------------
+
+def _extract_anthropic_prompt(kwargs: dict[str, Any]) -> str:
+    """Best-effort extraction of prompt text from anthropic.messages.create kwargs."""
+    messages = kwargs.get("messages", [])
+    parts = []
+    for m in messages:
+        role = m.get("role", "")
+        content = m.get("content", "")
+        if isinstance(content, list):
+            content = " ".join(
+                block.get("text", "") for block in content
+                if isinstance(block, dict) and block.get("type") == "text"
+            )
+        parts.append(f"{role}: {content}")
+    system = kwargs.get("system", "")
+    if system:
+        parts.insert(0, f"system: {system}")
+    return "\n".join(parts)
+
+
+def _extract_anthropic_response(response: Any) -> str:
+    """Extract text from an anthropic Message response."""
+    try:
+        parts = []
+        for block in response.content:
+            if hasattr(block, "text"):
+                parts.append(block.text)
+        return "\n".join(parts)
+    except Exception:
+        return str(response)
+
+
+def _patch_anthropic(client_at, agent_id: str) -> None:
+    """Wrap anthropic.Anthropic().messages.create and AsyncAnthropic variant."""
+    if "anthropic" in _INSTALLED:
+        return
+    try:
+        import anthropic  # noqa: PLC0415
+    except ImportError:
+        logger.debug("anthropic not installed — skipping AgentTrust Anthropic patch")
+        return
+
+    # --- sync ---
+    try:
+        _orig_sync = anthropic.resources.messages.Messages.create
+
+        @functools.wraps(_orig_sync)
+        def _sync_create(self, *args: Any, **kwargs: Any) -> Any:
+            result = _orig_sync(self, *args, **kwargs)
+            prompt = _extract_anthropic_prompt(kwargs)
+            output = _extract_anthropic_response(result)
+            _safe_validate(client_at, agent_id, user="anthropic-sync", input_text=prompt, output_text=output)
+            return result
+
+        anthropic.resources.messages.Messages.create = _sync_create  # type: ignore[method-assign]
+    except AttributeError:
+        logger.debug("anthropic.resources.messages.Messages not found — sync patch skipped")
+
+    # --- async ---
+    try:
+        _orig_async = anthropic.resources.messages.AsyncMessages.create
+
+        @functools.wraps(_orig_async)
+        async def _async_create(self, *args: Any, **kwargs: Any) -> Any:
+            result = await _orig_async(self, *args, **kwargs)
+            prompt = _extract_anthropic_prompt(kwargs)
+            output = _extract_anthropic_response(result)
+            await _safe_validate_async(client_at, agent_id, user="anthropic-async", input_text=prompt, output_text=output)
+            return result
+
+        anthropic.resources.messages.AsyncMessages.create = _async_create  # type: ignore[method-assign]
+    except AttributeError:
+        logger.debug("anthropic.resources.messages.AsyncMessages not found — async patch skipped")
+
+    _INSTALLED.add("anthropic")
+    logger.info("AgentTrust: Anthropic patch installed")
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+def install(
+    agent_id: str = "auto",
+    api_key: str | None = None,
+    gateway_url: str | None = None,
+    failure_mode: str = "open",
+) -> None:
+    """Install all available AgentTrust auto-instrumentation patches.
+
+    This function is **idempotent** — calling it more than once has no effect.
+
+    Parameters
+    ----------
+    agent_id:
+        The agent identifier reported to AgentTrust for every validation call.
+        Defaults to ``"auto"``.
+    api_key:
+        AgentTrust API key.  When *None* the key is resolved from the
+        environment variable ``AGENTRUST_API_KEY`` or ``~/.agentrust/config.yaml``.
+    gateway_url:
+        Override the AgentTrust gateway URL (useful for on-premise deployments).
+    failure_mode:
+        One of ``"open"`` (default — allow on error), ``"closed"`` (deny on
+        error), or ``"queue"`` (buffer calls for later retry).
+    """
+    if _INSTALLED == {"openai", "langchain", "anthropic"}:
+        # All known patches already applied — fast-path exit.
+        return
+
+    client_at, resolved_agent_id = _make_client(agent_id, api_key, gateway_url, failure_mode)
+
+    _patch_openai(client_at, resolved_agent_id)
+    _patch_langchain(client_at, resolved_agent_id)
+    _patch_anthropic(client_at, resolved_agent_id)
+
+    logger.info("AgentTrust auto-instrumentation complete. Installed: %s", _INSTALLED)
+
+
+auto_install = install  # alias for convenience

agentrust_sdk/autoload.py

@@ -0,0 +1,95 @@
+"""
+AgentTrust autoload bootstrap — zero-code instrumentation entry point.
+
+Three ways to activate without editing application code
+-------------------------------------------------------
+
+**1. PYTHONSTARTUP (interactive / script startup):**
+
+    export PYTHONSTARTUP=/path/to/agentrust_autoload.py
+    # or point to this module:
+    export PYTHONSTARTUP=$(python -c "import agentrust_sdk.autoload; print(agentrust_sdk.autoload.__file__)")
+
+**2. sitecustomize / usercustomize:**
+
+    # /path/to/site-packages/sitecustomize.py (system) or usercustomize.py (user)
+    import agentrust_sdk.autoload  # noqa: F401
+
+**3. PYTHONPATH + sitecustomize trick:**
+
+    # Create a sitecustomize.py in a directory and prepend to PYTHONPATH:
+    echo "import agentrust_sdk.autoload" > /tmp/agentrust_site/sitecustomize.py
+    export PYTHONPATH=/tmp/agentrust_site:$PYTHONPATH
+
+**4. Framework entry point (gunicorn/uvicorn pre-exec):**
+
+    gunicorn myapp:app --preload --worker-class uvicorn.workers.UvicornWorker \\
+        --config agentrust_sdk.autoload.gunicorn_conf
+
+**5. `[autoload]` setuptools extra:**
+
+    pip install agentrust-sdk[autoload]
+
+    Then in your app's own `sitecustomize.py` or startup:
+        import agentrust_sdk.autoload
+
+Environment controls
+--------------------
+``AGENTRUST_ENABLED=false``                → skip all instrumentation
+``AGENTRUST_AUTO_INSTRUMENT=false``        → skip hook patching
+``AGENTRUST_AUTOLOAD_LOG_LEVEL=DEBUG``     → verbose bootstrap logging
+``AGENTRUST_AUTOLOAD_EMBED=true``          → also start embedded gateway
+``AGENTRUST_AUTOLOAD_AGENT_ID=<string>``   → agent_id label for auto patches
+"""
+from __future__ import annotations
+
+import logging
+import os
+
+_log_level = os.environ.get("AGENTRUST_AUTOLOAD_LOG_LEVEL", "INFO").upper()
+logging.basicConfig(level=getattr(logging, _log_level, logging.INFO))
+logger = logging.getLogger("agentrust.autoload")
+
+# Respect kill-switch — do nothing if disabled
+if os.environ.get("AGENTRUST_ENABLED", "true").lower() in ("0", "false", "no"):
+    logger.debug("[AgentTrust autoload] AGENTRUST_ENABLED=false — skipping bootstrap")
+else:
+    try:
+        from agentrust_sdk.hooks import auto_instrument
+        from agentrust_sdk.config import SDK_CONFIG
+
+        _agent_id = os.environ.get("AGENTRUST_AUTOLOAD_AGENT_ID", "autoload")
+
+        patched = auto_instrument(agent_id=_agent_id)
+        if patched:
+            logger.info("[AgentTrust autoload] Patched frameworks: %s", ", ".join(patched))
+        else:
+            logger.debug("[AgentTrust autoload] No frameworks to patch (none installed or AGENTRUST_AUTO_INSTRUMENT=false)")
+
+        # Optionally start embedded gateway (SQLite, no external deps)
+        if os.environ.get("AGENTRUST_AUTOLOAD_EMBED", "false").lower() in ("1", "true", "yes"):
+            from agentrust_sdk.embedded import embed_gateway
+            gw = embed_gateway()
+            logger.info("[AgentTrust autoload] Embedded gateway started on %s", SDK_CONFIG.gateway_url)
+
+    except Exception as exc:
+        # Bootstrap must never crash the host application
+        logger.warning("[AgentTrust autoload] Bootstrap error (non-fatal): %s", exc)
+
+
+# ---------------------------------------------------------------------------
+# gunicorn config helper
+# ---------------------------------------------------------------------------
+
+def _post_fork(server, worker):  # noqa: ANN001
+    """gunicorn post_fork hook — re-apply patches after fork."""
+    try:
+        from agentrust_sdk.hooks import auto_instrument
+        auto_instrument(agent_id=os.environ.get("AGENTRUST_AUTOLOAD_AGENT_ID", "autoload"))
+    except Exception as exc:
+        import logging as _log
+        _log.getLogger("agentrust.autoload").warning("post_fork patch error: %s", exc)
+
+
+# gunicorn config module surface (used via --config agentrust_sdk.autoload)
+post_fork = _post_fork