agentrust-py 0.0.6__tar.gz → 0.0.7__tar.gz

{agentrust_py-0.0.6 → agentrust_py-0.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: agentrust-py
-Version: 0.0.6
+Version: 0.0.7
 Summary: AgentTrust Edge — Full Governance Platform. Policy Engine + MAD Classifier + HITL + Kill Switch.
 Project-URL: Homepage, https://agentrust.io
 Project-URL: Documentation, https://docs.agentrust.io
@@ -23,6 +23,8 @@ Requires-Python: >=3.10
 Requires-Dist: httpx<1.0.0,>=0.27.0
 Requires-Dist: pydantic<3.0.0,>=2.0.0
 Requires-Dist: pyyaml<7.0,>=6.0.3
+Provides-Extra: ag2
+Requires-Dist: ag2<1.0,>=0.9; extra == 'ag2'
 Provides-Extra: autoload
 Provides-Extra: crewai
 Requires-Dist: crewai<2.0,>=1.0; extra == 'crewai'

{agentrust_py-0.0.6 → agentrust_py-0.0.7}/agentrust_sdk/__init__.py

@@ -38,7 +38,12 @@ from .client import (
 from .config import SDK_CONFIG
 from .hooks import auto_instrument, auto_wrap
 from .adapters.mcp import mcp_tool_guard, MCPServerGuard
-from .adapters.autogen import AgentTrustReplyHook, AgentTrustGroupChatMonitor
+from .adapters.autogen import (
+    AgentTrustGroupChatMonitor,
+    AgentTrustReplyHook,
+    governed_function,
+    governed_functions,
+)
 from .adapters.claude_agents import ClaudeAgentGuard
 from .adapters.openai_agents import agent_guard as openai_agent_guard, OpenAIAgentsGuard, SwarmGuard
 from .adapters.crewai import AgentTrustCallback
@@ -87,6 +92,8 @@ from .webhooks import (
 __all__ = [
     # Primary API
     "harness",
+    "governed_function",
+    "governed_functions",
     # Embedded gateway (zero-dep plug-and-play)
     "embed_gateway",
     "EmbeddedGateway",
@@ -147,4 +154,4 @@ __all__ = [
     "validate",
 ]
 
-__version__ = "0.0.6"
+__version__ = "0.0.7"

{agentrust_py-0.0.6 → agentrust_py-0.0.7}/agentrust_sdk/adapters/autogen.py

@@ -5,6 +5,11 @@ Architecture doc (two-pager): "Connects to … AutoGen agents."
 
 Supports both AutoGen v0.2 (pyautogen) and AutoGen v0.4 (autogen-agentchat):
 
+Pattern 0 — governed_function (AG2 / pyautogen 0.2 tools):
+  Wrap each tool before register_function so every call is pre-checked,
+  held for approval when needed, run, post-checked and audited. A refusal
+  is returned as the tool result, so the chat keeps going.
+
 Pattern 1 — AgentTrustReplyHook (v0.2 + v0.4):
   Attach to any ConversableAgent / AssistantAgent to validate every reply
   before it is sent.
@@ -25,19 +30,111 @@ Requires Team tier.
 """
 from __future__ import annotations
 
+import functools
+import inspect
 import logging
 import time
-from typing import Any, Callable
+from typing import Any, Callable, Sequence
+
+from ..tiers import Capability
+from ._core import Allowed, Governor, OnBlock, Refused, check_adapter_tier
 
 logger = logging.getLogger(__name__)
 
+_FRAMEWORK = "ag2"
+
+
+def governed_function(
+    fn: Callable[..., Any],
+    *,
+    agent_id: str,
+    name: str | None = None,
+    on_block: OnBlock = "refuse",
+    review_timeout_s: float = 300.0,
+    user: str = "ag2",
+    base_url: str | None = None,
+    api_key: str | None = None,
+    governor: Governor | None = None,
+) -> Callable[..., Any]:
+    """Wrap an AG2 tool function so AgentTrust governs every call.
+
+    The result keeps the name, docstring and signature of *fn*, so AG2 builds
+    the same tool schema from it. Pass it to ``register_function`` or
+    ``register_for_llm``. Async functions stay async.
+    """
+    check_adapter_tier(Capability.AUTOGEN_ADAPTER, api_key)
+    gov = governor or Governor(
+        agent_id,
+        framework=_FRAMEWORK,
+        on_block=on_block,
+        review_timeout_s=review_timeout_s,
+        user=user,
+        base_url=base_url,
+        api_key=api_key,
+    )
+    tool_name = name or fn.__name__
+    signature = inspect.signature(fn)
+
+    def call_args(args: tuple, kwargs: dict) -> dict[str, Any]:
+        bound = signature.bind(*args, **kwargs)
+        bound.apply_defaults()
+        return dict(bound.arguments)
+
+    def unwrap(outcome: Allowed | Refused) -> Any:
+        # The refusal text is the tool result the agent reads next.
+        return outcome.message if isinstance(outcome, Refused) else outcome.result
+
+    if inspect.iscoroutinefunction(fn):
+
+        @functools.wraps(fn)
+        async def governed(*args: Any, **kwargs: Any) -> Any:
+            params = call_args(args, kwargs)
+            outcome = await gov.aguard_tool(
+                tool_name, params, lambda: fn(*args, **kwargs), input=_describe_call(tool_name, params),
+            )
+            return unwrap(outcome)
+
+    else:
 
-def _check_tier(api_key: str | None = None) -> None:
-    from agentrust_sdk.auth import resolve_key
-    from agentrust_sdk.tiers import Capability, is_allowed
-    info = resolve_key(api_key)
-    if not is_allowed(Capability.CREWAI_ADAPTER, info.tier):
-        raise RuntimeError("[AgentTrust] AutoGen adapter requires Team tier or higher.")
+        @functools.wraps(fn)
+        def governed(*args: Any, **kwargs: Any) -> Any:
+            params = call_args(args, kwargs)
+            outcome = gov.guard_tool(
+                tool_name, params, lambda: fn(*args, **kwargs), input=_describe_call(tool_name, params),
+            )
+            return unwrap(outcome)
+
+    return governed
+
+
+def governed_functions(
+    fns: Sequence[Callable[..., Any]],
+    *,
+    agent_id: str,
+    on_block: OnBlock = "refuse",
+    review_timeout_s: float = 300.0,
+    user: str = "ag2",
+    base_url: str | None = None,
+    api_key: str | None = None,
+    governor: Governor | None = None,
+) -> list[Callable[..., Any]]:
+    """Wrap each function with :func:`governed_function`, sharing one Governor."""
+    check_adapter_tier(Capability.AUTOGEN_ADAPTER, api_key)
+    gov = governor or Governor(
+        agent_id,
+        framework=_FRAMEWORK,
+        on_block=on_block,
+        review_timeout_s=review_timeout_s,
+        user=user,
+        base_url=base_url,
+        api_key=api_key,
+    )
+    return [governed_function(fn, agent_id=agent_id, governor=gov) for fn in fns]
+
+
+def _describe_call(name: str, args: dict[str, Any]) -> str:
+    """The proposed call as the pre-check's request input."""
+    return f"tool call {name}({', '.join(f'{k}={v!r}' for k, v in args.items())})"
 
 
 def _validate_sync(
@@ -50,7 +147,8 @@ def _validate_sync(
     framework: str,
     block_on_block: bool,
 ) -> None:
-    from agentrust_sdk.client import AgentTrustClient
+    from agentrust_sdk.client import AgentTrustClient, GatewayUnavailableError
+    from agentrust_sdk.config import SDK_CONFIG
     from agentrust_sdk.decorator import _check_decision, BlockedError
 
     try:
@@ -63,9 +161,11 @@ def _validate_sync(
                 framework=framework,
             )
         _check_decision(resp, block_on_block, False)
-    except BlockedError:
+    except (BlockedError, GatewayUnavailableError):
         raise
     except Exception as exc:
+        if SDK_CONFIG.failure_mode == "closed":
+            raise
         logger.warning("[AgentTrust] AutoGen validate failed (non-fatal): %s", exc)
 
 
@@ -99,7 +199,7 @@ class AgentTrustReplyHook:
         block_on_block: bool = True,
         framework: str = "AutoGen",
     ) -> None:
-        _check_tier(api_key)
+        check_adapter_tier(Capability.AUTOGEN_ADAPTER, api_key)
         self._agent_id = agent_id
         self._base_url = base_url
         self._api_key = api_key
@@ -111,47 +211,35 @@ class AgentTrustReplyHook:
     # AutoGen v0.2
     # ------------------------------------------------------------------
 
-    def _reply_func(
-        self,
-        recipient: Any,
-        messages: list[dict] | None = None,
-        sender: Any = None,
-        config: Any = None,
-    ) -> tuple[bool, str | None]:
-        """AutoGen v0.2 reply function — registered via register_reply()."""
-        if not messages:
-            return False, None
-
-        last = messages[-1]
-        content = last.get("content", "") or ""
-        role = last.get("role", "assistant")
-
-        if role not in ("assistant", "user"):
-            return False, None
+    def _check_outgoing(
+        self, sender: Any = None, message: Any = None, recipient: Any = None, silent: bool = False,
+    ) -> Any:
+        """AG2 ``process_message_before_send`` hook: check the agent's own reply."""
+        content = message.get("content") if isinstance(message, dict) else message
+        if not isinstance(content, str) or not content:
+            return message
 
         t0 = time.perf_counter()
-        output = {"content": content[:4000], "role": role}
         _validate_sync(
-            self._agent_id, self._base_url, self._api_key,
-            self._user, content[:500], output, self._framework, self._block_on_block,
+            self._agent_id, self._base_url, self._api_key, self._user,
+            content[:500], {"content": content[:4000], "role": "assistant"},
+            self._framework, self._block_on_block,
         )
         logger.debug(
             "[AgentTrust] AutoGen reply validated in %.1f ms", (time.perf_counter() - t0) * 1000
         )
-        return False, None  # False = don't override reply; None = no replacement content
+        return message
 
     def attach(self, agent: Any) -> None:
-        """Attach to an AutoGen v0.2 ConversableAgent / AssistantAgent."""
-        if not hasattr(agent, "register_reply"):
+        """Attach to an AG2 / AutoGen v0.2 ConversableAgent.
+
+        Checks every message the agent sends, before it leaves the agent.
+        """
+        if not hasattr(agent, "register_hook"):
             raise TypeError(
-                f"Expected a ConversableAgent with register_reply(), got {type(agent).__name__}"
+                f"Expected a ConversableAgent with register_hook(), got {type(agent).__name__}"
             )
-        agent.register_reply(
-            trigger=[],
-            reply_func=self._reply_func,
-            position=0,
-            name="agentrust_reply_hook",
-        )
+        agent.register_hook("process_message_before_send", self._check_outgoing)
         logger.info(
             "[AgentTrust] AutoGen reply hook attached to agent '%s'", getattr(agent, "name", "?")
         )
@@ -209,7 +297,7 @@ class AgentTrustGroupChatMonitor:
         block_on_block: bool = False,  # Default False — group chats tolerate review
         framework: str = "AutoGen",
     ) -> None:
-        _check_tier(api_key)
+        check_adapter_tier(Capability.AUTOGEN_ADAPTER, api_key)
         self._agent_id = agent_id
         self._base_url = base_url
         self._api_key = api_key

{agentrust_py-0.0.6 → agentrust_py-0.0.7}/agentrust_sdk/config.py

@@ -47,7 +47,7 @@ from pathlib import Path
 logger = logging.getLogger(__name__)
 
 _VALID_FAILURE_MODES = frozenset({"open", "closed", "queue"})
-_SDK_VERSION = "0.0.6"
+_SDK_VERSION = "0.0.7"
 _MIN_GATEWAY_VERSION = "0.0.1a1"
 
 

{agentrust_py-0.0.6 → agentrust_py-0.0.7}/agentrust_sdk/decorator.py

@@ -329,6 +329,17 @@ def _check_decision(resp: Any, block_on_block: bool, block_on_review: bool) -> N
         raise BlockedError(outcome, resp.decision.reason, resp.envelope_id)
 
 
+def _ag2_installed() -> bool:
+    """AG2 installs the ``autogen`` module under the distribution name ``ag2``."""
+    from importlib.metadata import PackageNotFoundError, version
+
+    try:
+        version("ag2")
+    except PackageNotFoundError:
+        return False
+    return True
+
+
 def _detect_framework() -> str:
     """Best-effort framework detection from installed packages."""
     try:
@@ -341,7 +352,7 @@ def _detect_framework() -> str:
             ("llama_index", "LlamaIndex"),
         ]:
             if importlib.util.find_spec(name):
-                return label
+                return "AG2" if name == "autogen" and _ag2_installed() else label
     except Exception:
         pass
     return "REST"

{agentrust_py-0.0.6 → agentrust_py-0.0.7}/agentrust_sdk/manifest.py

@@ -50,7 +50,7 @@ class CertifyManifest(BaseModel):
     display_name: str = ""
     framework: str = Field(
         default="rest",
-        pattern=r"^(langchain|langgraph|crewai|autogen|openai-agents|mcp|rest)$",
+        pattern=r"^(langchain|langgraph|crewai|autogen|ag2|openai-agents|mcp|rest)$",
     )
     language: str = Field(default="python", pattern=r"^(python|typescript|rest)$")
     model: ModelDeclaration = Field(default_factory=ModelDeclaration)

{agentrust_py-0.0.6 → agentrust_py-0.0.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "agentrust-py"
-version = "0.0.6"
+version = "0.0.7"
 description = "AgentTrust Edge — Full Governance Platform. Policy Engine + MAD Classifier + HITL + Kill Switch."
 readme = "README.md"
 license = { text = "Apache-2.0" }
@@ -41,6 +41,8 @@ llm = [
 # Framework adapters (optional, tier-gated)
 langgraph = ["langgraph>=0.2.0,<2.0.0"]
 crewai    = ["crewai>=1.0,<2.0"]
+# AG2 0.x keeps the ConversableAgent API (module `autogen`); 1.x is a new API
+ag2       = ["ag2>=0.9,<1.0"]
 openai    = ["openai>=2.41.1,<3.0.0"]
 langchain = ["langchain>=0.2.0,<1.0.0"]
 # Retry/backoff (recommended for production)

agentrust_py-0.0.7/tests/adapters/_ag2.py

@@ -0,0 +1,57 @@
+"""A scripted chat model for real AG2 agents, shared by the AG2 binding tests."""
+from __future__ import annotations
+
+import json
+from types import SimpleNamespace as NS
+from typing import Any
+
+
+def tool_call(name: str, args: dict[str, Any]) -> tuple[str, dict[str, Any]]:
+    """The reply in which the agent calls a tool."""
+    return name, args
+
+
+def final(text: str = "done DONE") -> str:
+    return text
+
+
+class ScriptedModel:
+    """Stands in for ``OpenAIWrapper.create``: replays *replies* in order."""
+
+    def __init__(self, *replies):
+        self.replies = list(replies)
+        self.prompts: list[str] = []
+
+    def create(self, wrapper, **params):
+        self.prompts.append(json.dumps(params.get("messages", []), default=str))
+        reply = self.replies.pop(0)
+        if isinstance(reply, str):
+            message = {"role": "assistant", "content": reply}
+            msg = NS(content=reply, tool_calls=None, function_call=None, role="assistant")
+        else:
+            name, args = reply
+            call = {"id": "call-1", "type": "function",
+                    "function": {"name": name, "arguments": json.dumps(args)}}
+            message = {"role": "assistant", "content": None, "tool_calls": [call]}
+            msg = NS(content=None, tool_calls=[call], function_call=None, role="assistant")
+        msg.model_dump = lambda: message
+        return NS(
+            choices=[NS(message=msg, finish_reason="stop")], model="scripted", cost=0, id="r",
+            usage=NS(prompt_tokens=0, completion_tokens=0, total_tokens=0),
+        )
+
+    @staticmethod
+    def extract(response):
+        msg = response.choices[0].message
+        return [msg.model_dump()] if msg.tool_calls else [msg.content]
+
+
+def install(monkeypatch, *replies) -> ScriptedModel:
+    from autogen.oai.client import OpenAIWrapper
+
+    model = ScriptedModel(*replies)
+    monkeypatch.setattr(OpenAIWrapper, "create", lambda self, **kw: model.create(self, **kw))
+    monkeypatch.setattr(
+        OpenAIWrapper, "extract_text_or_completion_object", classmethod(lambda cls, r: model.extract(r)),
+    )
+    return model

agentrust_py-0.0.7/tests/adapters/test_ag2_governed_function.py

@@ -0,0 +1,232 @@
+"""governed_function: every AG2 tool call is pre-checked before it runs."""
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+
+import httpx
+import pytest
+
+pytest.importorskip("autogen")
+
+from autogen import ConversableAgent, register_function
+
+from agentrust_sdk.adapters.autogen import governed_function, governed_functions
+from tests.test_governor import FakeAsyncClient, FakeClient, _gov, _resp
+
+AGENT = "ag2-agent"
+
+
+def _refund(calls):
+    def refund(amount: float, reason: str = "") -> str:
+        """Refund an amount to the customer."""
+        calls.append(amount)
+        return f"refunded {amount}"
+
+    return refund
+
+
+# ---------------------------------------------------------------------------
+# Unit: a real Governor on a scripted client
+# ---------------------------------------------------------------------------
+
+def test_wrapper_keeps_name_docstring_and_signature():
+    inner = _refund([])
+    wrapped = governed_function(inner, agent_id="a", governor=_gov(FakeClient()))
+
+    assert wrapped.__name__ == "refund"
+    assert wrapped.__doc__ == inner.__doc__
+    assert inspect.signature(wrapped) == inspect.signature(inner)
+
+
+def test_allowed_call_is_pre_checked_runs_and_is_post_checked():
+    calls, client = [], FakeClient()
+    wrapped = governed_function(_refund(calls), agent_id="a", governor=_gov(client))
+
+    assert wrapped(5, reason="late") == "refunded 5"
+
+    assert calls == [5]
+    [pre] = client.pre_calls
+    assert (pre["action"], pre["params"]) == ("refund", {"amount": 5, "reason": "late"})
+    assert pre["input"] == "tool call refund(amount=5, reason='late')"
+    assert len(client.post_calls) == 1
+
+
+def test_defaults_are_part_of_the_checked_arguments():
+    client = FakeClient()
+    governed_function(_refund([]), agent_id="a", governor=_gov(client))(amount=1.5)
+
+    assert client.pre_calls[0]["params"] == {"amount": 1.5, "reason": ""}
+
+
+def test_blocked_call_does_not_run_and_returns_the_refusal_text():
+    calls = []
+    client = FakeClient(pre=[_resp("block", reason="over the limit")])
+    wrapped = governed_function(_refund(calls), agent_id="a", governor=_gov(client))
+
+    result = wrapped(9999)
+
+    assert calls == []
+    assert "over the limit" in result
+    assert client.post_calls == []
+
+
+def test_held_call_runs_once_a_reviewer_approves():
+    calls = []
+    client = FakeClient(
+        pre=[_resp("escalate", review_id="r1")], resolutions=["pending", "approved"],
+    )
+    wrapped = governed_function(_refund(calls), agent_id="a", governor=_gov(client))
+
+    assert wrapped(5) == "refunded 5"
+    assert calls == [5]
+
+
+def test_held_call_is_refused_when_the_review_times_out():
+    calls = []
+    client = FakeClient(pre=[_resp("escalate", review_id="r1")])
+    wrapped = governed_function(
+        _refund(calls), agent_id="a", governor=_gov(client, review_timeout_s=10),
+    )
+
+    result = wrapped(5)
+
+    assert calls == []
+    assert isinstance(result, str) and result
+
+
+def test_async_function_stays_async_and_is_governed():
+    calls, client = [], FakeAsyncClient()
+
+    async def lookup(order_id: str) -> str:
+        """Look up an order."""
+        calls.append(order_id)
+        return f"order {order_id}"
+
+    wrapped = governed_function(lookup, agent_id="a", governor=_gov(client))
+
+    assert inspect.iscoroutinefunction(wrapped)
+    assert asyncio.run(wrapped("42")) == "order 42"
+    assert calls == ["42"] and len(client.pre_calls) == 1
+
+
+def test_blocked_async_call_does_not_run():
+    calls = []
+
+    async def lookup(order_id: str) -> str:
+        """Look up an order."""
+        calls.append(order_id)
+        return "x"
+
+    client = FakeAsyncClient(pre=[_resp("block", reason="nope")])
+    wrapped = governed_function(lookup, agent_id="a", governor=_gov(client))
+
+    assert "nope" in asyncio.run(wrapped("42"))
+    assert calls == []
+
+
+def test_function_without_arguments():
+    client = FakeClient()
+
+    def ping() -> str:
+        """Ping."""
+        return "pong"
+
+    assert governed_function(ping, agent_id="a", governor=_gov(client))() == "pong"
+    assert client.pre_calls[0]["params"] == {}
+
+
+def test_exception_from_the_tool_propagates_and_skips_the_post_check():
+    client = FakeClient()
+
+    def boom() -> str:
+        """Fail."""
+        raise ValueError("bad")
+
+    with pytest.raises(ValueError):
+        governed_function(boom, agent_id="a", governor=_gov(client))()
+    assert client.post_calls == []
+
+
+def test_kill_switch_off_runs_the_function_plainly(monkeypatch):
+    monkeypatch.setenv("AGENTRUST_ENABLED", "false")
+    calls, client = [], FakeClient()
+    wrapped = governed_function(_refund(calls), agent_id="a", governor=_gov(client))
+
+    assert wrapped(5) == "refunded 5"
+    assert client.pre_calls == []
+
+
+def test_governed_functions_share_one_governor():
+    client = FakeClient()
+    a, b = governed_functions([_refund([]), _refund([])], agent_id="a", governor=_gov(client))
+
+    assert a.__name__ == b.__name__ == "refund"
+
+
+def test_gateway_down_in_closed_mode_refuses(monkeypatch):
+    from agentrust_sdk.client import GatewayUnavailableError
+
+    calls = []
+    client = FakeClient(pre_error=GatewayUnavailableError("down"))
+    wrapped = governed_function(_refund(calls), agent_id="a", governor=_gov(client))
+
+    assert "gateway unavailable" in wrapped(5)
+    assert calls == []
+
+
+# ---------------------------------------------------------------------------
+# Real AG2: the tool schema is unchanged and registration works
+# ---------------------------------------------------------------------------
+
+def _schema(fn):
+    from autogen import LLMConfig
+
+    cfg = LLMConfig({"model": "m", "api_key": "k", "api_type": "openai"})
+    caller = ConversableAgent("caller", llm_config=cfg)
+    executor = ConversableAgent("executor", llm_config=False)
+    register_function(fn, caller=caller, executor=executor, name=fn.__name__, description="d")
+    return caller.llm_config.tools if hasattr(caller.llm_config, "tools") else caller.llm_config["tools"]
+
+
+def test_ag2_builds_the_same_tool_schema_from_the_wrapper():
+    inner = _refund([])
+    wrapped = governed_function(inner, agent_id="a", governor=_gov(FakeClient()))
+
+    assert _schema(wrapped) == _schema(inner)
+
+
+# ---------------------------------------------------------------------------
+# Embedded gateway: what lands in the ledger
+# ---------------------------------------------------------------------------
+
+def _ledger(gw):
+    rows = httpx.get(f"{gw.url}/v1/audit/executions", params={"limit": 100},
+                     headers={"Authorization": f"Bearer {gw._token}"}).json()["items"]
+    return sorted(rows, key=lambda r: r["id"])
+
+
+def test_declared_tool_leaves_pre_and_post_check_rows(embedded_gw):
+    gw = embedded_gw(manifests={AGENT: {"declared_tools": [{"name": "refund"}]}})
+    calls = []
+
+    result = governed_function(_refund(calls), agent_id=AGENT)(5)
+
+    assert result == "refunded 5" and calls == [5]
+    assert [(r["phase"], r["decision"]) for r in _ledger(gw)] == [
+        ("pre_check", "approve"), ("post_check", "approve"),
+    ]
+
+
+def test_undeclared_tool_is_refused_before_it_runs(embedded_gw):
+    gw = embedded_gw(manifests={AGENT: {"declared_tools": [{"name": "calculator"}]}})
+    calls = []
+
+    result = governed_function(_refund(calls), agent_id=AGENT)(5)
+
+    assert calls == []
+    assert "is not in this agent's declared_tools" in result
+    first = _ledger(gw)[0]
+    assert (first["phase"], first["decision"]) == ("pre_check", "block")
+    assert json.loads(first["request_json"])["metadata"]["action"] == "refund"