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()