From bb67bd86f83d3b4d1c430fcda83e8a8d51cda412 Mon Sep 17 00:00:00 2001 From: Adarsh Prashar Date: Sun, 14 Jun 2026 21:26:45 +0530 Subject: [PATCH] feat(sdk): AutoGen adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bind an AutoGen agent to a governed run with no code change beyond wrapping the model client. GovernedChatCompletionClient wraps any v0.4+ ChatCompletionClient and ticks one governed step per model request before delegating, so a runaway agent is hard-stopped at its loop/time budget; with gate_tools=True each tool call the model requests (a FunctionCall in the result) routes through the approval gate before the agent can run it. Every other protocol method delegates to the wrapped client, so it's a drop-in. A single agent run propagates the halt typed (AutoGen doesn't wrap model-client exceptions there). A team catches the exception and re-raises it as a plain RuntimeError("BudgetExceeded: ..."), so I added governed_run_errors() to restore the typed BudgetExceeded/ApprovalDenied around a team call — verified both paths against autogen-agentchat 0.7.5. Targets autogen-agentchat / autogen-core >= 0.4, < 1 (lazily handled — the wrapper is duck-typed and imports nothing from autogen, so the SDK stays dependency-free). Adds an [autogen] optional extra, a runnable examples/autogen, README + CHANGELOG entries, and stdlib-only unit tests plus a skip-unless integration test that halts a real runaway agent. Closes #84 --- CHANGELOG.md | 13 + examples/autogen/README.md | 112 ++++++ examples/autogen/agent.py | 159 ++++++++ examples/autogen/requirements.txt | 9 + sdks/python/README.md | 10 + sdks/python/pyproject.toml | 1 + sdks/python/riskkernel/adapters/__init__.py | 1 + sdks/python/riskkernel/adapters/autogen.py | 253 ++++++++++++ sdks/python/tests/test_autogen.py | 405 ++++++++++++++++++++ 9 files changed, 963 insertions(+) create mode 100644 examples/autogen/README.md create mode 100644 examples/autogen/agent.py create mode 100644 examples/autogen/requirements.txt create mode 100644 sdks/python/riskkernel/adapters/autogen.py create mode 100644 sdks/python/tests/test_autogen.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c43cde..0ea7686 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,19 @@ surface is governed by [`COMPATIBILITY.md`](COMPATIBILITY.md). ## [Unreleased] ### Added +- **AutoGen adapter (Python SDK).** `from riskkernel.adapters.autogen import + GovernedChatCompletionClient` — wrap your AutoGen model client once and hand it to + your existing `AssistantAgent` (or team) to bind it to a governed run with no other + code change. One model request counts as one governed step, so the deterministic + loop/time budget halts a runaway agent; with `gate_tools=True` each tool call the + model requests (a `FunctionCall` in the result) routes through the human-approval + gate before the agent can run it. Targets the actively maintained v0.4+ line + (`autogen-agentchat` / `autogen-core` >= 0.4), not the legacy `pyautogen` 0.2 API; + `autogen` is lazily handled (the wrapper is duck-typed and imports nothing), so it + stays an optional dependency. A single agent run propagates the halt typed; a team + (`RoundRobinGroupChat`, …) re-raises it as a `RuntimeError`, so `governed_run_errors()` + (from the same module) restores the typed `BudgetExceeded`/`ApprovalDenied` around a + team call. - **OTLP trace ingress.** RiskKernel can now act as an OTLP/HTTP trace endpoint (`POST /v1/traces`), the consume side of the OpenTelemetry surface — point any exporter at it (`OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:7070`) and the GenAI diff --git a/examples/autogen/README.md b/examples/autogen/README.md new file mode 100644 index 0000000..e67d6c7 --- /dev/null +++ b/examples/autogen/README.md @@ -0,0 +1,112 @@ +# autogen — stop a runaway AutoGen agent + +Wrap an AutoGen model client with RiskKernel's `GovernedChatCompletionClient` and +the **deterministic governor caps the agent** — one governed step per model request, +hard-stopped at the loop budget. The halt propagates out of `agent.run()` and ends +the run. The kill comes from RiskKernel, not from the script. + +Targets the actively maintained **AutoGen v0.4+** line (`autogen-agentchat` / +`autogen-core`), not the legacy `pyautogen` 0.2 API. + +**No API key, no model call.** This uses a tiny local model-client stub that always +returns the same tool call (so the agent loops forever), so the loop enforcement +runs with nothing but `riskkernel serve`. (Add a real model — and the dollar/token +ceiling — with a few lines; see [below](#add-the-dollar--token-ceiling-real-model).) + +## Run it in 60 seconds + +```bash +# 1. start the daemon — no key needed for this demo +docker run --rm -p 7070:7070 ghcr.io/prashar32/riskkernel:latest + +# 2. in another terminal, install autogen + the SDK, and run it +cd examples/autogen +pip install -r requirements.txt +python agent.py +``` + +## What you'll see + +A real run (the run-id varies; the structure is the point): + +``` +▶ autogen loop budget = 6 (enforced by the Go governor) + run id: 31ab03d4-e7e4-47ad-bb74-6402999ca4fd + +🛑 RiskKernel halted the AutoGen run — reason: loop_budget_exceeded + ── final ledger (enforced by the governor) ── + model calls (loops) : 7 (budget: 6) + run id : 31ab03d4-e7e4-47ad-bb74-6402999ca4fd + The agent would have looped forever; the governor capped it — and the + halt propagated out of agent.run(). +``` + +The 7th model request is refused: the wrapper's `create()` ticks a governed step, +the daemon returns HTTP `402 loop_budget_exceeded`, and that surfaces as +`rk.BudgetExceeded` — which propagates out of AutoGen and stops the agent. + +## Wrapping your own AutoGen agent + +One object: wrap the model client, then use your agent as-is. + +```python +import riskkernel as rk +from riskkernel.adapters.autogen import GovernedChatCompletionClient +from autogen_agentchat.agents import AssistantAgent + +rt = rk.Runtime() # http://localhost:7070 +with rt.governed_run(budget=rt.budget(loops=50, seconds=600)) as run: + client = GovernedChatCompletionClient(model_client, run) # one step per model call + agent = AssistantAgent("assistant", model_client=client) # otherwise unchanged + await agent.run(task="...") # raises rk.BudgetExceeded at the cap +``` + +Gate side-effecting tools on human approval by constructing it with +`GovernedChatCompletionClient(model_client, run, gate_tools=True)` — each tool the +model requests routes through the approval gate before the agent can run it. + +## Teams re-raise the halt as a RuntimeError + +A single agent run propagates the typed `rk.BudgetExceeded` directly. But a **team** +(`RoundRobinGroupChat`, `SelectorGroupChat`, …) catches an agent's exception and +re-raises it to the caller as a plain `RuntimeError` (its container serializes the +error and the team's `run()` re-raises `RuntimeError(str(error))`). The run still +halts — it is **not** swallowed — but the type is lost. Wrap the team call in +`governed_run_errors()` to get the typed exception back: + +```python +from riskkernel.adapters.autogen import governed_run_errors + +with governed_run_errors(): + await team.run(task="...") # re-raises rk.BudgetExceeded, not RuntimeError +``` + +## Add the dollar / token ceiling (real model) + +The wrapper caps **loops and time**. To also cap **dollars and tokens**, route the +real model client through the run's proxy so every call is priced from real provider +usage and the dollar budget halts the run: + +```python +from autogen_ext.models.openai import OpenAIChatCompletionClient # speaks OpenAI to the proxy + +cfg = run.proxy_config() +real = OpenAIChatCompletionClient( + model="claude-sonnet-4-5", + base_url=cfg["base_url"], # http://localhost:7070/v1 + api_key="sk-unused", # the real provider key lives in the daemon + default_headers=cfg["headers"], # groups calls into this governed run +) +client = GovernedChatCompletionClient(real, run) +# budget=rt.budget(loops=50, dollars=1.00, seconds=600) +# the dollar ceiling trips at the proxy with HTTP 402; the loop/time ceiling +# trips in the wrapper. Start the daemon with your ANTHROPIC_API_KEY. +``` + +## Tuning for a recording + +- `LOOP_BUDGET` (default `6`) — lower it for a faster kill, raise it for more + steps before the halt. + +Nothing about the kill is faked: it's the daemon returning HTTP `402` before the +over-budget model request, surfaced through AutoGen's own model-client path. diff --git a/examples/autogen/agent.py b/examples/autogen/agent.py new file mode 100644 index 0000000..e93383d --- /dev/null +++ b/examples/autogen/agent.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +"""autogen — stop a runaway AutoGen agent at its budget. + +Wraps an AutoGen model client with RiskKernel's ``GovernedChatCompletionClient``: +one governed step per model request. A deliberately runaway agent — one whose model +always asks to call a tool, so it never finishes on its own — is hard-stopped by the +deterministic governor at its loop budget, and the halt propagates out of +``agent.run()``. The kill comes from RiskKernel, not from this script. + +Key-free: the model client here is a tiny local stub that always returns the same +tool call, so the loop enforcement runs with nothing but ``riskkernel serve`` — no +model call, no spend. To add the dollar / token ceiling with a *real* model, build +your real client (e.g. ``OpenAIChatCompletionClient``) pointed at the run's proxy +and wrap that instead; the README shows the few extra lines. + +Run it: + riskkernel serve # in another terminal — no key needed for this demo + pip install -r requirements.txt + python agent.py + +The integration is one object: ``GovernedChatCompletionClient(model_client, run)``. +Hand the wrapped client to your existing ``AssistantAgent`` and the governor caps +the loop the same way it would cap a real agent. +""" + +from __future__ import annotations + +import asyncio +import os + +import riskkernel as rk +from riskkernel.adapters.autogen import GovernedChatCompletionClient + +try: + from autogen_agentchat.agents import AssistantAgent + from autogen_core import CancellationToken, FunctionCall + from autogen_core.models import ( + ChatCompletionClient, + CreateResult, + ModelInfo, + RequestUsage, + ) + from autogen_core.tools import FunctionTool +except ImportError: + raise SystemExit( + "this example needs autogen-agentchat + autogen-core — install them with:\n" + " pip install -r requirements.txt" + ) + +# ───────────────────────────────────────────────────────────────────────────── +# Knobs. The kill is never faked: the governor (in the daemon) enforces the loop +# budget before each model request and the wrapper propagates the halt into AutoGen. +# ───────────────────────────────────────────────────────────────────────────── +DAEMON_URL = os.environ.get("RISKKERNEL_BASE_URL", "http://localhost:7070") +LOOP_BUDGET = 6 # max model requests (loop iterations) the governor allows + + +def _noop(query: str) -> str: + """A do-nothing tool the agent keeps calling.""" + return "keep going" + + +class _LoopingClient(ChatCompletionClient): + """A stand-in model client so the demo needs no key: it always returns the same + tool call, so the agent loops forever — until the governor caps it. Swap this for + a real ``OpenAIChatCompletionClient`` / ``AnthropicChatCompletionClient`` pointed + at ``run.proxy_config()`` to add the dollar/token ceiling (see the README).""" + + def __init__(self) -> None: + self._usage = RequestUsage(prompt_tokens=0, completion_tokens=0) + + async def create(self, messages, *, tools=[], tool_choice="auto", + json_output=None, extra_create_args={}, cancellation_token=None): + return CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", name="noop", arguments='{"query": "again"}')], + usage=RequestUsage(prompt_tokens=1, completion_tokens=1), + cached=False, + ) + + async def create_stream(self, *args, **kwargs): # pragma: no cover + raise NotImplementedError + + async def close(self) -> None: + return None + + def actual_usage(self): + return self._usage + + def total_usage(self): + return self._usage + + def count_tokens(self, messages, *, tools=[]): + return 0 + + def remaining_tokens(self, messages, *, tools=[]): + return 1000 + + @property + def capabilities(self): + return self.model_info + + @property + def model_info(self): + return ModelInfo(vision=False, function_calling=True, json_output=False, + family="unknown", structured_output=False) + + +async def _run() -> int: + rt = rk.Runtime(base_url=DAEMON_URL) + print(f"▶ autogen loop budget = {LOOP_BUDGET} (enforced by the Go governor)") + + tool = FunctionTool(_noop, description="a no-op tool the agent keeps calling") + + try: + with rt.governed_run(name="autogen-demo", + budget=rt.budget(loops=LOOP_BUDGET)) as run: + print(f" run id: {run.id}\n") + # The one integration line: wrap the model client. Everything below is a + # plain AutoGen agent — no other change. + client = GovernedChatCompletionClient(_LoopingClient(), run) + agent = AssistantAgent( + "looper", model_client=client, tools=[tool], + reflect_on_tool_use=False, max_tool_iterations=1000, + ) + try: + await agent.run(task="keep using the tool", + cancellation_token=CancellationToken()) + except rk.BudgetExceeded as halt: + _report_halt(run, halt) + return 0 + print(" (agent finished on its own before the budget — unexpected here)") + return 0 + except rk.APIError as e: + if e.code == "connection_error": + print(f"\n✗ can't reach the daemon at {DAEMON_URL} — start it first:") + print(" riskkernel serve") + print(" # or: docker run --rm -p 7070:7070 ghcr.io/prashar32/riskkernel:latest") + return 1 + raise + + +def _report_halt(run: rk.Run, halt: rk.BudgetExceeded) -> None: + """Print the final, governor-enforced ledger for the halted AutoGen run.""" + usage = run.status().get("usage", {}) + print(f"\n🛑 RiskKernel halted the AutoGen run — reason: {halt.reason}") + print(" ── final ledger (enforced by the governor) ──") + print(f" model calls (loops) : {usage.get('loops', 0):>4} (budget: {LOOP_BUDGET})") + print(f" run id : {run.id}") + print(" The agent would have looped forever; the governor capped it — and the") + print(" halt propagated out of agent.run().") + + +def main() -> int: + return asyncio.run(_run()) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/autogen/requirements.txt b/examples/autogen/requirements.txt new file mode 100644 index 0000000..452d33b --- /dev/null +++ b/examples/autogen/requirements.txt @@ -0,0 +1,9 @@ +# This AutoGen example needs autogen-agentchat + autogen-core (the v0.4+ line — the +# model-client protocol the adapter wraps) and the RiskKernel Python SDK (stdlib-only). +# Tested with autogen-agentchat / autogen-core 0.4+; the adapter supports >=0.4,<1. +autogen-agentchat>=0.4,<1 +autogen-core>=0.4,<1 +riskkernel + +# Working inside a clone and want your local SDK instead? From THIS directory: +# pip install ../../sdks/python diff --git a/sdks/python/README.md b/sdks/python/README.md index ef011d4..4d903d5 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -98,8 +98,18 @@ hooks = RiskKernelRunHooks(run, gate_tools=True) # CrewAI — a step_callback (one governed step per agent step; halts the crew at budget) from riskkernel.adapters.crewai import RiskKernelStepCallback crew = Crew(agents=[...], tasks=[...], step_callback=RiskKernelStepCallback(run)) + +# AutoGen (v0.4+) — wrap the model client; one governed step per model call +from riskkernel.adapters.autogen import GovernedChatCompletionClient +client = GovernedChatCompletionClient(model_client, run) # drop-in for the real client +agent = AssistantAgent("assistant", model_client=client) # halts the agent at budget ``` +> AutoGen halts the run either way, but a *team* (`RoundRobinGroupChat`, …) re-raises +> the halt as a `RuntimeError`; wrap the team call in +> `with governed_run_errors():` (from the same module) to get the typed +> `rk.BudgetExceeded` back. A single agent run propagates it typed already. + ## Configuration `Runtime(base_url=..., token=...)`, or the env vars `RISKKERNEL_BASE_URL` and diff --git a/sdks/python/pyproject.toml b/sdks/python/pyproject.toml index a1a1a36..59b5384 100644 --- a/sdks/python/pyproject.toml +++ b/sdks/python/pyproject.toml @@ -27,6 +27,7 @@ Source = "https://github.com/prashar32/riskkernel" [project.optional-dependencies] langchain = ["langchain-core"] crewai = ["crewai>=0.80,<2"] +autogen = ["autogen-agentchat>=0.4,<1", "autogen-core>=0.4,<1"] dev = ["pytest"] # This project lives in a subdirectory (sdks/python) of the repo. Hatchling's diff --git a/sdks/python/riskkernel/adapters/__init__.py b/sdks/python/riskkernel/adapters/__init__.py index 077cd2b..06f800b 100644 --- a/sdks/python/riskkernel/adapters/__init__.py +++ b/sdks/python/riskkernel/adapters/__init__.py @@ -7,4 +7,5 @@ - ``claude_agent`` — a PreToolUse hook for the Claude Agent SDK (approval gate). - ``openai_agents`` — RunHooks for the OpenAI Agents SDK (steps + approval gate). - ``crewai`` — a step_callback for CrewAI (steps + tool approval gate). +- ``autogen`` — a model-client wrapper for AutoGen v0.4+ (steps + tool gate). """ diff --git a/sdks/python/riskkernel/adapters/autogen.py b/sdks/python/riskkernel/adapters/autogen.py new file mode 100644 index 0000000..0d41d1d --- /dev/null +++ b/sdks/python/riskkernel/adapters/autogen.py @@ -0,0 +1,253 @@ +"""AutoGen adapter: a model-client wrapper that enforces a governed run's loop and +time budgets, ticking one governed step per model request. Wrap your AutoGen model +client once and hand the wrapped client to your existing ``AssistantAgent`` (or +team) — no other code change — and the deterministic governor caps the run; point +the same client at the governing proxy (``run.proxy_config()``) for token/cost +metering, and this wrapper adds the outer-loop enforcement the proxy can't see. + + from autogen_agentchat.agents import AssistantAgent + from autogen_ext.models.openai import OpenAIChatCompletionClient + from riskkernel.adapters.autogen import GovernedChatCompletionClient + + client = GovernedChatCompletionClient(OpenAIChatCompletionClient(model="gpt-4o"), run) + agent = AssistantAgent("assistant", model_client=client) + await agent.run(task="...") # one governed step per model call + +Each ``create()`` / ``create_stream()`` call ticks ``run.step()`` *before* +delegating to the real client — one model request == one governed step — so a +runaway agent is hard-stopped at its loop/time budget. With ``gate_tools=True``, +the wrapper inspects the returned ``CreateResult`` for tool-call requests +(``FunctionCall``) and routes each through the approval gate before the result is +handed back to the agent, so a side-effecting tool the model wants to invoke is +blocked unless approved. + +Propagation (verified against the library — note the asymmetry): + +* **A single agent run directly** (``agent.run()`` / ``agent.on_messages()``) does + NOT wrap model-client exceptions: a ``BudgetExceeded`` raised here propagates + out *unwrapped*, halting the agent. +* **Inside a team** (``RoundRobinGroupChat`` / ``SelectorGroupChat`` / + ``BaseGroupChat.run()`` / ``run_stream()``), the agent's container catches the + exception, serializes it (``SerializableException``) onto the group-chat error + channel, and the team re-raises it to the caller as a plain + ``RuntimeError(str(error))`` — so the run still halts (it is NOT swallowed), but + the original ``BudgetExceeded`` *type* is lost, replaced by a ``RuntimeError`` + whose message reads ``"BudgetExceeded: run halted: loop_budget_exceeded"``. Use + ``governed_run_errors()`` (a context manager) around the team call to restore the + typed ``BudgetExceeded`` / ``ApprovalDenied`` so callers can ``except`` on it:: + + from riskkernel.adapters.autogen import governed_run_errors + with governed_run_errors(): + await team.run(task="...") # re-raises typed BudgetExceeded, not RuntimeError + +Supported API: the autogen-core ``ChatCompletionClient`` protocol +(``autogen_core.models``) used by ``autogen-agentchat`` — i.e. the actively +maintained v0.4+ line (``autogen-agentchat`` / ``autogen-core`` >= 0.4), NOT the +legacy ``pyautogen`` 0.2 ``register_reply`` API. The wrapper delegates every +protocol method to the wrapped client, so it is a drop-in. ``autogen`` is NOT a +dependency of the SDK — nothing here imports it; the wrapper is duck-typed and the +module imports fine without AutoGen present. +""" + +from __future__ import annotations + +from typing import Any, Optional + +from ..approval import ApprovalGate +from ..errors import ApprovalDenied, BudgetExceeded +from ..runtime import Run + + +class GovernedChatCompletionClient: + """Wraps any AutoGen ``ChatCompletionClient`` and binds it to a governed run. + + Pass the wrapped client to your ``AssistantAgent(model_client=...)`` (or any + agent that takes a model client). Each model request ticks one governed step, + enforcing the run's loop/time budget; ``BudgetExceeded`` is raised from the + governor and propagates out of the model call. Every other protocol method is + delegated unchanged to the wrapped client. + + Args: + client: the real AutoGen model client to wrap (an + ``autogen_core.models.ChatCompletionClient``). + run: the governed Run. + gate_tools: if True, every tool call the model requests (a ``FunctionCall`` + in the ``CreateResult``) must pass the approval gate before the result + is returned to the agent; a denial raises ``ApprovalDenied`` and halts. + tool_side_effect: side-effect label reported for gated tools. + timeout: max seconds to await a human decision on a gated tool. + """ + + def __init__(self, client: Any, run: Run, gate_tools: bool = False, + tool_side_effect: str = "tool", timeout: Optional[float] = None): + self._client = client + self.run = run + self.gate_tools = gate_tools + self.tool_side_effect = tool_side_effect + self.timeout = timeout + self._gate = ApprovalGate(run) + + # One model request == one governed step. run.step() raises BudgetExceeded when + # the loop/time budget is spent; AutoGen does not catch it inside the agent, so + # the halt propagates out of the model call (and out of agent.run()). Inside a + # team it is surfaced as a RuntimeError — see governed_run_errors(). + async def create(self, *args: Any, **kwargs: Any) -> Any: + self.run.step() + result = await self._client.create(*args, **kwargs) + if self.gate_tools: + self._gate_result_tools(result) + return result + + async def create_stream(self, *args: Any, **kwargs: Any): + # Tick the step before the stream opens (same as the non-streaming path), so + # a runaway loop is capped before another model request begins. The final + # chunk of an AutoGen model stream is the CreateResult; gate its tool calls. + self.run.step() + async for chunk in self._client.create_stream(*args, **kwargs): + if self.gate_tools and _is_create_result(chunk): + self._gate_result_tools(chunk) + yield chunk + + def _gate_result_tools(self, result: Any) -> None: + """Route each tool call the model requested through the approval gate. + + A ``CreateResult.content`` is either a string (plain text — no tools) or a + list of ``FunctionCall`` objects (each with ``.name`` / ``.arguments``). + Gate each before the agent gets the chance to execute it. Duck-typed so this + works without importing AutoGen.""" + for call in _tool_calls(result): + self._gate.require( + _call_name(call) or "tool", + side_effect=self.tool_side_effect, + arguments={"arguments": _stringify(_call_arguments(call))}, + timeout=self.timeout, + ) + + # ── Delegate the rest of the ChatCompletionClient protocol to the wrapped + # client unchanged, so this stays a drop-in replacement across AutoGen versions. + def actual_usage(self) -> Any: + return self._client.actual_usage() + + def total_usage(self) -> Any: + return self._client.total_usage() + + def count_tokens(self, *args: Any, **kwargs: Any) -> Any: + return self._client.count_tokens(*args, **kwargs) + + def remaining_tokens(self, *args: Any, **kwargs: Any) -> Any: + return self._client.remaining_tokens(*args, **kwargs) + + async def close(self) -> None: + await self._client.close() + + @property + def model_info(self) -> Any: + return self._client.model_info + + @property + def capabilities(self) -> Any: + return self._client.capabilities + + def __getattr__(self, name: str) -> Any: + # Any protocol method or attribute not named above (AutoGen adds/renames + # some across versions) falls through to the wrapped client. __getattr__ is + # only consulted for names not found normally, so it never shadows the + # governed create()/create_stream() above. + return getattr(self._client, name) + + +class governed_run_errors: + """Context manager that restores the typed RiskKernel exception when a governed + halt happens *inside an AutoGen team*. + + A team (``RoundRobinGroupChat`` etc.) catches an agent's exception and re-raises + it to the caller as a plain ``RuntimeError`` whose message is + ``"BudgetExceeded: "`` (the original type name + message). This wraps + the team call and re-raises the original typed ``BudgetExceeded`` / + ``ApprovalDenied`` so callers can ``except`` on it as they do everywhere else:: + + with governed_run_errors(): + await team.run(task="...") + + A single agent run directly doesn't need this (it propagates the typed exception + already), but it is harmless there. Only RuntimeErrors that match a known + RiskKernel error prefix are converted; any other RuntimeError is left untouched. + """ + + def __enter__(self) -> "governed_run_errors": + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> bool: + if exc is None or not isinstance(exc, RuntimeError): + return False + typed = _typed_from_runtime_error(exc) + if typed is None: + return False + raise typed from exc + + +def _typed_from_runtime_error(exc: RuntimeError) -> Optional[Exception]: + """Map an AutoGen-wrapped team RuntimeError back to its RiskKernel type. + + AutoGen's SerializableException stringifies as ``": "`` (and + may append a traceback on following lines), so the first line begins with the + original exception's class name. Match on that prefix to reconstruct the typed + error; return None if it isn't one of ours (leave the RuntimeError alone).""" + text = str(exc) + first = text.split("\n", 1)[0] + if first.startswith("BudgetExceeded:"): + message = first[len("BudgetExceeded:"):].strip() + # The governor's message is "run halted: "; recover the reason. + reason = message.split("run halted:", 1)[-1].strip() if "run halted:" in message else message + return BudgetExceeded(reason or "budget_exceeded", message) + if first.startswith("ApprovalDenied:"): + message = first[len("ApprovalDenied:"):].strip() + # The message is "approval denied for [: ]"; recover the tool. + tool = message + reason = "" + if message.startswith("approval denied for "): + rest = message[len("approval denied for "):] + if ": " in rest: + tool, reason = rest.split(": ", 1) + else: + tool = rest + return ApprovalDenied(tool.strip() or "tool", reason.strip()) + return None + + +def _is_create_result(chunk: Any) -> bool: + """True if a stream chunk is the final CreateResult (has ``.content`` and + ``.finish_reason``); intermediate chunks are plain strings.""" + return hasattr(chunk, "content") and hasattr(chunk, "finish_reason") + + +def _tool_calls(result: Any) -> list: + """The list of FunctionCall objects in a CreateResult, or [] for a text result. + Duck-typed: tool calls are a list under ``.content``; a string is plain text.""" + content = getattr(result, "content", None) + if isinstance(content, list): + return [c for c in content if _is_tool_call(c)] + return [] + + +def _is_tool_call(obj: Any) -> bool: + """True if obj looks like an AutoGen FunctionCall (has a ``.name``).""" + return hasattr(obj, "name") and not isinstance(obj, str) + + +def _call_name(call: Any) -> str: + name = getattr(call, "name", None) + return str(name) if name else "" + + +def _call_arguments(call: Any) -> Any: + return getattr(call, "arguments", None) + + +def _stringify(v: Any) -> Any: + try: + import json + json.dumps(v) + return v + except Exception: + return repr(v) diff --git a/sdks/python/tests/test_autogen.py b/sdks/python/tests/test_autogen.py new file mode 100644 index 0000000..3a9f736 --- /dev/null +++ b/sdks/python/tests/test_autogen.py @@ -0,0 +1,405 @@ +"""AutoGen adapter tests — governance behavior against a fake Run (no daemon, no +third-party deps), plus a real AutoGen integration test gated behind skipUnless. + +The wrapper imports even without autogen installed (it is duck-typed and imports +nothing from autogen), so these exercise the enforcement path on stdlib alone. The +async create()/create_stream() methods are driven with asyncio.run. +""" + +import asyncio +import unittest + +from riskkernel.adapters.autogen import ( + GovernedChatCompletionClient, + governed_run_errors, + _typed_from_runtime_error, +) +from riskkernel.errors import ApprovalDenied, BudgetExceeded + + +def _has_autogen() -> bool: + try: + import autogen_agentchat # noqa: F401 + import autogen_core # noqa: F401 + + return True + except Exception: + return False + + +def _run(coro): + return asyncio.run(coro) + + +async def _drain_stream(client, *args, **kwargs): + """Consume create_stream to completion, returning the list of chunks.""" + out = [] + async for chunk in client.create_stream(*args, **kwargs): + out.append(chunk) + return out + + +class FakeRun: + """A stand-in for runtime.Run: counts steps and halts past a loop budget, + mirroring how the daemon's BeginStep raises BudgetExceeded when the budget is + spent. No HTTP, no daemon.""" + + def __init__(self, loop_budget=None): + self.steps = 0 + self.loop_budget = loop_budget + + def step(self): + self.steps += 1 + if self.loop_budget is not None and self.steps > self.loop_budget: + raise BudgetExceeded("loop_budget_exceeded") + return self.steps + + +class _FakeGate: + """Records gate calls; can be set to deny like a human denying a tool.""" + + def __init__(self, deny=False): + self.deny = deny + self.calls = [] + + def require(self, tool, side_effect="", arguments=None, timeout=None): + self.calls.append({"tool": tool, "side_effect": side_effect, + "arguments": arguments, "timeout": timeout}) + if self.deny: + raise ApprovalDenied(tool, "no") + + +# ── Minimal stand-ins for AutoGen's model client and result objects, so the unit +# tests don't need autogen installed. The wrapper duck-types on these shapes: +# a CreateResult has .content (str OR list[FunctionCall]) and .finish_reason; a +# FunctionCall has .name / .arguments. +class _FunctionCall: + def __init__(self, name, arguments=""): + self.name = name + self.arguments = arguments + + +class _CreateResult: + def __init__(self, content, finish_reason="stop"): + self.content = content + self.finish_reason = finish_reason + + +class _FakeModelClient: + """A model client whose create() returns a canned CreateResult and records the + args it was called with. create_stream() yields a couple of string chunks then + the final CreateResult, like a real AutoGen streaming client.""" + + def __init__(self, result=None): + self.result = result or _CreateResult("hello") + self.create_calls = 0 + self.closed = False + + async def create(self, *args, **kwargs): + self.create_calls += 1 + return self.result + + async def create_stream(self, *args, **kwargs): + self.create_calls += 1 + yield "thinking" + yield "more" + yield self.result + + def actual_usage(self): + return "actual" + + def total_usage(self): + return "total" + + def count_tokens(self, *a, **k): + return 7 + + def remaining_tokens(self, *a, **k): + return 100 + + async def close(self): + self.closed = True + + @property + def model_info(self): + return {"family": "fake"} + + @property + def capabilities(self): + return {"vision": False} + + # An attribute the wrapper doesn't name explicitly, to exercise __getattr__ + # delegation to the wrapped client. + def component_config(self): + return {"provider": "fake"} + + +class AutoGenAdapterTest(unittest.TestCase): + def test_module_imports_without_autogen(self): + # The wrapper is duck-typed and imports nothing from autogen, so the SDK can + # import the adapter with no autogen installed (it is not a dependency). + self.assertTrue(callable(GovernedChatCompletionClient)) + + def test_create_ticks_one_step(self): + run = FakeRun() + client = GovernedChatCompletionClient(_FakeModelClient(), run) + result = _run(client.create([])) + self.assertEqual(run.steps, 1) # one governed step per create() + self.assertEqual(result.content, "hello") # the real result is returned + + def test_create_enforces_loop_budget(self): + # The real proof at the unit level: create() must raise BudgetExceeded when + # the loop budget is spent. AutoGen does not catch model-client exceptions in + # a single agent, so this halt propagates out of the agent. + run = FakeRun(loop_budget=2) + client = GovernedChatCompletionClient(_FakeModelClient(), run) + _run(client.create([])) # step 1 + _run(client.create([])) # step 2 + with self.assertRaises(BudgetExceeded) as cm: + _run(client.create([])) # step 3 -> over budget + self.assertEqual(cm.exception.reason, "loop_budget_exceeded") + + def test_halt_is_not_swallowed(self): + # Guards the propagation contract: create() must let BudgetExceeded escape so + # AutoGen's agent (which does not wrap model-client errors) actually halts. The + # step is ticked BEFORE delegating, so an over-budget call never reaches the + # underlying client. + run = FakeRun(loop_budget=0) + inner = _FakeModelClient() + client = GovernedChatCompletionClient(inner, run) + with self.assertRaises(BudgetExceeded): + _run(client.create([])) + self.assertEqual(inner.create_calls, 0) # halted before the real call + + def test_create_stream_ticks_one_step_and_yields_chunks(self): + run = FakeRun() + client = GovernedChatCompletionClient(_FakeModelClient(), run) + chunks = _run(_drain_stream(client, [])) + self.assertEqual(run.steps, 1) # one governed step per stream + self.assertEqual(chunks[:2], ["thinking", "more"]) + self.assertEqual(chunks[-1].content, "hello") # final chunk is the result + + def test_create_stream_enforces_loop_budget(self): + run = FakeRun(loop_budget=1) + client = GovernedChatCompletionClient(_FakeModelClient(), run) + _run(_drain_stream(client, [])) # step 1 + with self.assertRaises(BudgetExceeded): + _run(_drain_stream(client, [])) # step 2 -> over budget + + def test_tool_gating_off_by_default(self): + # Default: gate_tools is False, so a tool-call result never asks for approval. + run = FakeRun() + result = _CreateResult([_FunctionCall("deploy", '{"x":1}')]) + client = GovernedChatCompletionClient(_FakeModelClient(result), run) + gate = _FakeGate() + client._gate = gate + _run(client.create([])) + self.assertEqual(gate.calls, []) # gate_tools defaults off + self.assertEqual(run.steps, 1) # still ticks the step + + def test_tool_gating_on_gates_each_function_call(self): + # gate_tools=True: every FunctionCall in the result is gated, with its name + # and side-effect label passed through. + run = FakeRun() + result = _CreateResult([_FunctionCall("deploy", '{"x":1}'), + _FunctionCall("notify", "{}")]) + client = GovernedChatCompletionClient( + _FakeModelClient(result), run, gate_tools=True, tool_side_effect="exec") + gate = _FakeGate() + client._gate = gate + _run(client.create([])) + self.assertEqual([c["tool"] for c in gate.calls], ["deploy", "notify"]) + self.assertEqual(gate.calls[0]["side_effect"], "exec") + self.assertEqual(run.steps, 1) + + def test_tool_gating_does_not_gate_text_result(self): + # A plain-text result (content is a str, not a list of FunctionCalls) is not a + # tool call: gate_tools=True must not gate it, though it still ticks a step. + run = FakeRun() + client = GovernedChatCompletionClient( + _FakeModelClient(_CreateResult("final answer")), run, gate_tools=True) + gate = _FakeGate() + client._gate = gate + _run(client.create([])) + self.assertEqual(gate.calls, []) + self.assertEqual(run.steps, 1) + + def test_tool_gating_denied_raises_after_step(self): + # A denied tool must raise ApprovalDenied. The step is ticked first (the model + # request happened); the gate then blocks the side effect from running. + run = FakeRun() + result = _CreateResult([_FunctionCall("deploy")]) + client = GovernedChatCompletionClient( + _FakeModelClient(result), run, gate_tools=True) + client._gate = _FakeGate(deny=True) + with self.assertRaises(ApprovalDenied): + _run(client.create([])) + + def test_tool_gating_in_stream(self): + # The final CreateResult chunk of a stream is gated the same as a create(). + run = FakeRun() + result = _CreateResult([_FunctionCall("deploy")]) + client = GovernedChatCompletionClient( + _FakeModelClient(result), run, gate_tools=True) + gate = _FakeGate() + client._gate = gate + _run(_drain_stream(client, [])) + self.assertEqual([c["tool"] for c in gate.calls], ["deploy"]) + + def test_protocol_methods_delegate_to_wrapped_client(self): + # The wrapper must be a drop-in: every other ChatCompletionClient method/attr + # delegates to the wrapped client unchanged. + run = FakeRun() + inner = _FakeModelClient() + client = GovernedChatCompletionClient(inner, run) + self.assertEqual(client.actual_usage(), "actual") + self.assertEqual(client.total_usage(), "total") + self.assertEqual(client.count_tokens([]), 7) + self.assertEqual(client.remaining_tokens([]), 100) + self.assertEqual(client.model_info, {"family": "fake"}) + self.assertEqual(client.capabilities, {"vision": False}) + # An un-named method falls through __getattr__ to the wrapped client. + self.assertEqual(client.component_config(), {"provider": "fake"}) + _run(client.close()) + self.assertTrue(inner.closed) + + # ── The team-propagation asymmetry: a team re-raises a budget halt as a plain + # RuntimeError("BudgetExceeded: ..."); governed_run_errors() restores the type. + def test_typed_from_runtime_error_budget(self): + e = RuntimeError("BudgetExceeded: run halted: loop_budget_exceeded") + typed = _typed_from_runtime_error(e) + self.assertIsInstance(typed, BudgetExceeded) + self.assertEqual(typed.reason, "loop_budget_exceeded") + + def test_typed_from_runtime_error_budget_with_traceback(self): + # SerializableException may append a traceback on following lines; only the + # first line is matched. + e = RuntimeError( + "BudgetExceeded: run halted: token_budget_exceeded\nTraceback:\n ...") + typed = _typed_from_runtime_error(e) + self.assertIsInstance(typed, BudgetExceeded) + self.assertEqual(typed.reason, "token_budget_exceeded") + + def test_typed_from_runtime_error_approval(self): + e = RuntimeError("ApprovalDenied: approval denied for deploy: nope") + typed = _typed_from_runtime_error(e) + self.assertIsInstance(typed, ApprovalDenied) + self.assertEqual(typed.tool, "deploy") + self.assertEqual(typed.reason, "nope") + + def test_typed_from_runtime_error_ignores_unrelated(self): + # A RuntimeError that isn't one of ours is left alone (returns None). + self.assertIsNone(_typed_from_runtime_error(RuntimeError("some other error"))) + + def test_governed_run_errors_reraises_typed_budget(self): + # The context manager converts a team's wrapped RuntimeError back to the typed + # BudgetExceeded so callers can `except rk.BudgetExceeded`. + with self.assertRaises(BudgetExceeded) as cm: + with governed_run_errors(): + raise RuntimeError("BudgetExceeded: run halted: loop_budget_exceeded") + self.assertEqual(cm.exception.reason, "loop_budget_exceeded") + + def test_governed_run_errors_passes_through_other_errors(self): + # An unrelated RuntimeError is not touched. + with self.assertRaises(RuntimeError) as cm: + with governed_run_errors(): + raise RuntimeError("kaboom") + self.assertEqual(str(cm.exception), "kaboom") + + def test_governed_run_errors_no_error(self): + # No exception -> the block runs and exits cleanly. + with governed_run_errors(): + x = 1 + 1 + self.assertEqual(x, 2) + + @unittest.skipUnless(_has_autogen(), "autogen-agentchat / autogen-core not installed") + def test_autogen_integration_halts_runaway_agent(self): + # The real proof: a real AutoGen AssistantAgent driven by a model client that + # always asks to call a tool (so the agent loops) is hard-stopped at its loop + # budget by the governor. We wrap a real ChatCompletionClient implementation + # with GovernedChatCompletionClient and assert the BudgetExceeded surfaces. + from autogen_agentchat.agents import AssistantAgent + from autogen_core import CancellationToken + from autogen_core.models import ( + ChatCompletionClient, + CreateResult, + FunctionExecutionResultMessage, + ModelInfo, + RequestUsage, + ) + from autogen_core import FunctionCall + from autogen_core.tools import FunctionTool + + def noop(query: str) -> str: + """A do-nothing tool the agent keeps calling.""" + return "keep going" + + tool = FunctionTool(noop, description="a no-op tool") + + # A real ChatCompletionClient that always returns a tool call, so the agent + # never finishes on its own and would loop forever without the governor. No + # network, no key. We implement only what AssistantAgent needs. + class LoopingClient(ChatCompletionClient): + def __init__(self): + self._total = RequestUsage(prompt_tokens=0, completion_tokens=0) + + async def create(self, messages, *, tools=[], tool_choice="auto", + json_output=None, extra_create_args={}, + cancellation_token=None): + # If the last message is a tool result, ask for the tool again — a loop. + return CreateResult( + finish_reason="function_calls", + content=[FunctionCall(id="1", name="noop", + arguments='{"query": "again"}')], + usage=RequestUsage(prompt_tokens=1, completion_tokens=1), + cached=False, + ) + + async def create_stream(self, *args, **kwargs): # pragma: no cover + raise NotImplementedError + + async def close(self): + return None + + def actual_usage(self): + return self._total + + def total_usage(self): + return self._total + + def count_tokens(self, messages, *, tools=[]): + return 0 + + def remaining_tokens(self, messages, *, tools=[]): + return 1000 + + @property + def capabilities(self): + return self.model_info + + @property + def model_info(self): + return ModelInfo(vision=False, function_calling=True, + json_output=False, family="unknown", + structured_output=False) + + run = FakeRun(loop_budget=3) + client = GovernedChatCompletionClient(LoopingClient(), run) + agent = AssistantAgent( + "looper", model_client=client, tools=[tool], + reflect_on_tool_use=False, max_tool_iterations=100, + ) + + async def go(): + await agent.run(task="loop forever", cancellation_token=CancellationToken()) + + # A single agent run propagates the model-client exception unwrapped, so the + # typed BudgetExceeded reaches us directly (no governed_run_errors needed here). + with self.assertRaises(BudgetExceeded): + _run(go()) + # The governor capped the loop: it stopped at the budget + the over-budget tick. + self.assertGreaterEqual(run.steps, 3) + + +if __name__ == "__main__": + unittest.main()