From da367f3994e36c79b18d75b732ab07648043743e Mon Sep 17 00:00:00 2001 From: vsoch Date: Thu, 29 Oct 2026 11:46:06 -0700 Subject: [PATCH] feat: send manifest schema to setup And also validate and add record_tool. Signed-off-by: vsoch --- .github/workflows/test.yaml | 24 ++++- README.md | 20 ++++ behalf/__init__.py | 8 ++ behalf/core.py | 60 +++++++++++- behalf/runner/adk.py | 47 ++++----- behalf/runner/sdk.py | 40 +++++--- behalf/runner/strands.py | 39 ++++---- pyproject.toml | 2 +- tests/test_adk.py | 108 +++++++++++++++++++++ tests/test_core.py | 78 ++++++++++++++- tests/test_sdk.py | 131 +++++++++++++++++++++++++ tests/test_strands.py | 187 ++++++++++++++++++++++++++++++++++++ 12 files changed, 677 insertions(+), 67 deletions(-) create mode 100644 tests/test_sdk.py create mode 100644 tests/test_strands.py diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index dac9573..aed02a7 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -21,7 +21,7 @@ jobs: - run: python -m pytest -v tests/test_core.py gemini: - # The ADK runner tests stub the model, so no GOOGLE_API_KEY is needed. + # The ADK runner tests use a scripted model, so no GOOGLE_API_KEY is needed. runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 @@ -31,6 +31,28 @@ jobs: - run: pip install ".[gemini,dev]" - run: python -m pytest -v tests/test_adk.py + claude: + # The SDK runner tests stub the client, so no ANTHROPIC_API_KEY or CLI is needed. + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - run: pip install ".[claude,dev]" + - run: python -m pytest -v tests/test_sdk.py + + aws: + # The Strands runner tests use a scripted model, so no AWS credentials are needed. + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - run: pip install ".[aws,dev]" + - run: python -m pytest -v tests/test_strands.py + format: runs-on: ubuntu-latest steps: diff --git a/README.md b/README.md index c90aada..d5620e4 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,26 @@ from behalf import run_task, make_runner outcome = await run_task(MyTask(), make_runner("aws")) ``` +### Manifests + +`manifest_schema()` is sent to the model during setup (see `render_setup_prompt`), +and `finalize_setup` only accepts a manifest that has every schema field. `run_task` +raises `ValueError` if a manifest, conversed or passed in, is missing a field. + +### Results + +Runners return `None` from `run_agent`; results come from tool side effects. +`record_tool(sink, fields)` returns a tool that appends each call's arguments to `sink`: + +```python +from behalf import record_tool + +class MyTask(Task): + def build_tools(self, manifest): + self.results = [] + return [record_tool(self.results, {"job": str, "site": str})] +``` + ## Install diff --git a/behalf/__init__.py b/behalf/__init__.py index 0e03eb4..c530bee 100644 --- a/behalf/__init__.py +++ b/behalf/__init__.py @@ -11,7 +11,11 @@ RunOutcome, Task, ToolSpec, + accept_manifest, default_confirm, + missing_manifest_fields, + record_tool, + render_setup_prompt, run_task, ) @@ -21,7 +25,11 @@ "RunOutcome", "Task", "ToolSpec", + "accept_manifest", "default_confirm", + "missing_manifest_fields", + "record_tool", + "render_setup_prompt", "run_task", "make_runner", ] diff --git a/behalf/core.py b/behalf/core.py index 46ddaae..64c63c0 100644 --- a/behalf/core.py +++ b/behalf/core.py @@ -86,7 +86,7 @@ class AgentRunner(Protocol): fake. This seam is what lets the whole flow be tested without a key.""" async def converse(self, task: "Task") -> dict: - """Run the seeded setup conversation; return the manifest.""" + """Run the setup conversation from render_setup_prompt(task); return the manifest.""" ... async def run_agent( @@ -100,6 +100,56 @@ async def run_agent( ... +def render_setup_prompt(task: "Task") -> str: + """Return the task's setup prompt with its manifest schema and finalize instruction appended.""" + fields = "\n".join( + f"- {name}: {getattr(typ, '__name__', str(typ))}" + for name, typ in task.manifest_schema().items() + ) + return ( + f"{task.setup_system_prompt()}\n\n" + f"The manifest is a JSON object with these fields:\n{fields}\n\n" + "When every field is known, call finalize_setup with the manifest." + ) + + +def missing_manifest_fields(task: "Task", manifest: Any) -> list[str]: + """Return schema fields absent from the manifest (all fields if it is not a dict).""" + schema = task.manifest_schema() + if not isinstance(manifest, dict): + return list(schema) + return [name for name in schema if name not in manifest] + + +def accept_manifest(task: "Task", manifest: Any, captured: dict) -> str: + """Store the manifest in captured if complete; return the message for the model.""" + missing = missing_manifest_fields(task, manifest) + if missing: + return ( + f"manifest is missing fields: {', '.join(missing)}. " + "Ask the user for them, then call finalize_setup again." + ) + captured["manifest"] = manifest + return "setup finalized" + + +def record_tool( + sink: list, + fields: dict, + name: str = "record_result", + description: str = "Record one result.", +) -> ToolSpec: + """Return a tool that appends each call's arguments to sink.""" + + async def _record(args: dict) -> dict: + sink.append(dict(args)) + return { + "content": [{"type": "text", "text": json.dumps({"recorded": len(sink)})}] + } + + return ToolSpec(name, description, fields, _record, kind="action", confirm=False) + + @dataclass class RunOutcome: manifest: dict @@ -122,10 +172,14 @@ async def run_task( approve_fn: Optional[Callable[[dict], bool]] = None, ) -> RunOutcome: """The one entrypoint every task flows through. manifest given -> skip the - conversation (reproducible/batch); None -> converse to produce it. Then - optional setup approval, then execute; action tools gate on confirm_fn.""" + conversation (reproducible/batch); None -> converse to produce it. The + manifest must contain every manifest_schema() field. Then optional setup + approval, then execute; action tools gate on confirm_fn.""" if manifest is None: manifest = await runner.converse(task) + missing = missing_manifest_fields(task, manifest) + if missing: + raise ValueError(f"{task.name}: manifest is missing fields: {missing}") approved = True if task.requires_setup_approval: diff --git a/behalf/runner/adk.py b/behalf/runner/adk.py index 84a9dff..43f965f 100644 --- a/behalf/runner/adk.py +++ b/behalf/runner/adk.py @@ -1,16 +1,10 @@ """A Gemini backend via Google's Agent Development Kit (ADK), as an alternative AgentRunner to SDKRunner. -STATUS: scaffolded, NOT yet run. Google ADK isn't a base dependency and isn't -installed in the dev image, so this module is untested end to end. The shape -follows ADK's documented API (Agent + FunctionTool + Runner + session service), -but the spots marked TODO are where the ADK specifics must be confirmed against -the installed version before relying on it. Install with `pip install google-adk`. - -It exists because the framework already abstracts the model behind AgentRunner, -so adding Gemini is a contained job: convert our ToolSpecs to ADK function tools -(carrying the same read/action confirmation gate) and drive ADK's runner. None -of ProfileTask / tools.py / the container code changes. +Converts ToolSpecs to ADK function tools (carrying the same read/action +confirmation gate) and drives ADK's Runner. Tested end to end with a scripted +BaseLlm in place of Gemini (tests/test_adk.py). Install with `pip install +behalf[gemini]`. """ from __future__ import annotations @@ -29,7 +23,14 @@ from google.adk.tools import FunctionTool from google.genai import types -from ..core import AgentRunner, ConfirmFn, Task, ToolSpec +from ..core import ( + AgentRunner, + ConfirmFn, + Task, + ToolSpec, + accept_manifest, + render_setup_prompt, +) _PY_TYPE = {str: str, int: int, float: float, bool: bool, list: list, dict: dict} @@ -77,15 +78,11 @@ async def _invoke(**kwargs) -> Any: ] _invoke.__name__ = ts.name _invoke.__doc__ = ts.description - _invoke.__signature__ = inspect.Signature( - params - ) # TODO: confirm ADK reads __signature__ (vs real params) + _invoke.__signature__ = inspect.Signature(params) _invoke.__annotations__ = { n: _PY_TYPE.get(t, str) for n, t in ts.input_schema.items() } - return FunctionTool( - _invoke - ) # TODO: confirm FunctionTool(func) is the current constructor + return FunctionTool(_invoke) class ADKRunner(AgentRunner): @@ -93,7 +90,7 @@ class ADKRunner(AgentRunner): def __init__( self, - model: str = "gemini-2.5-flash", + model: Any = "gemini-2.5-flash", app_name: str = "behalf", verbose: bool = True, ): @@ -117,7 +114,6 @@ async def _run_once( msg = types.Content(role="user", parts=[types.Part(text=user_prompt)]) texts: list[str] = [] - # TODO: confirm run_async signature + how final text is surfaced on events async for event in runner.run_async( user_id="local", session_id=session.id, new_message=msg ): @@ -139,26 +135,23 @@ async def run_agent( ) -> Any: adk_tools = [_to_adk_tool(t, confirm_fn) for t in tools] await self._run_once(system_prompt, user_prompt, adk_tools) - # results are collected via the tools' side effects (e.g. record_artifact - # writing into ProfileTask's sink), same as SDKRunner. + # Results are collected through tool side effects (e.g. a record_tool + # sink), same as SDKRunner. return None async def converse(self, task: Task) -> dict: - """Interactive setup: give the agent a finalize tool and loop until it's - called. Mirrors SDKRunner.converse. TODO: verify ADK multi-turn session - reuse (re-running the runner with the same session_id to continue).""" + """Interactive setup over one ADK session; loops until finalize_setup accepts a manifest.""" captured: dict = {} async def finalize_setup(manifest: dict) -> str: """Finalize the setup manifest (matching the task's schema) and end setup.""" - captured["manifest"] = manifest - return "setup finalized" + return accept_manifest(task, manifest, captured) finalize = FunctionTool(finalize_setup) agent = Agent( name="setup", model=self.model, - instruction=task.setup_system_prompt(), + instruction=render_setup_prompt(task), tools=[finalize], ) session_service = InMemorySessionService() diff --git a/behalf/runner/sdk.py b/behalf/runner/sdk.py index cfec766..5a2baf2 100644 --- a/behalf/runner/sdk.py +++ b/behalf/runner/sdk.py @@ -22,7 +22,14 @@ tool, ) -from ..core import AgentRunner, ConfirmFn, Task, ToolSpec +from ..core import ( + AgentRunner, + ConfirmFn, + Task, + ToolSpec, + accept_manifest, + render_setup_prompt, +) def _say(text: str) -> dict: @@ -43,6 +50,20 @@ async def _sdk_tool(args): return _sdk_tool +def _finalize_tool(task: Task, captured: dict): + """SDK tool that stores a complete manifest in captured, or reports missing fields.""" + + @tool( + "finalize_setup", + "Finalize the setup manifest (matching the task's schema) and end setup.", + {"manifest": dict}, + ) + async def finalize(args): + return _say(accept_manifest(task, args.get("manifest", args), captured)) + + return finalize + + class SDKRunner(AgentRunner): def __init__(self, model: str | None = None, verbose: bool = True): self.model = model @@ -83,25 +104,16 @@ async def run_agent( async for message in query(prompt=user_prompt, options=options): self._echo(message) - # Results are collected through the tools' side effects (e.g. - # record_artifact writing into the task's sink), not the return value. + # Results are collected through tool side effects (e.g. a record_tool + # sink), not the return value. return None async def converse(self, task: Task) -> dict: captured: dict = {} - - @tool( - "finalize_setup", - "Finalize the setup manifest (matching the task's schema) and end setup.", - {"manifest": dict}, - ) - async def finalize(args): - captured["manifest"] = args.get("manifest", args) - return _say("setup finalized") - + finalize = _finalize_tool(task, captured) server = create_sdk_mcp_server(name="setup", version="0.1.0", tools=[finalize]) options = self._options( - "setup", server, ["mcp__setup__finalize_setup"], task.setup_system_prompt() + "setup", server, ["mcp__setup__finalize_setup"], render_setup_prompt(task) ) # The model asks the user for what it needs and drives the structure; we diff --git a/behalf/runner/strands.py b/behalf/runner/strands.py index 5ed8a63..8860201 100644 --- a/behalf/runner/strands.py +++ b/behalf/runner/strands.py @@ -5,11 +5,9 @@ role) and a region; the model is a Bedrock model id, which can still be Claude (a Bedrock Anthropic inference profile), just reached through Bedrock. -STATUS: scaffolded, NOT yet run here (needs `pip install strands-agents` and AWS -creds). Shape follows the documented Strands API; spots to confirm against the -installed version are marked TODO. Strands tools are plain Python functions and -it exposes tool calls, so our ToolSpec -> function conversion (with the -read/action confirm gate) maps over directly. +Strands tools are plain Python functions, so ToolSpecs convert directly (with +the read/action confirm gate). Tested end to end with a scripted Model in place +of Bedrock (tests/test_strands.py). Install with `pip install behalf[aws]`. """ from __future__ import annotations @@ -20,11 +18,16 @@ from typing import Any from strands import Agent -from strands.models import ( - BedrockModel, -) # TODO: confirm class name in the installed version +from strands.models import BedrockModel -from ..core import AgentRunner, ConfirmFn, Task, ToolSpec +from ..core import ( + AgentRunner, + ConfirmFn, + Task, + ToolSpec, + accept_manifest, + render_setup_prompt, +) _PY_TYPE = {str: str, int: int, float: float, bool: bool, list: list, dict: dict} @@ -79,7 +82,7 @@ def _invoke(**kwargs) -> Any: _invoke.__annotations__ = { n: _PY_TYPE.get(t, str) for n, t in ts.input_schema.items() } - return tool(_invoke) # TODO: confirm strands.tool(func) call form + return tool(_invoke) class StrandsRunner(AgentRunner): @@ -95,9 +98,7 @@ def __init__( def _model(self): if self.model is None: return None - return BedrockModel( - model_id=self.model, region_name=self.region - ) # TODO: confirm kwargs + return BedrockModel(model_id=self.model, region_name=self.region) def _render(self, event) -> None: """Stream assistant text (dimmed). Tool lines print from the wrapper.""" @@ -140,9 +141,8 @@ async def run_agent( stream = getattr(agent, "stream_async", None) if stream is None: - result = await asyncio.to_thread( - agent, user_prompt - ) # TODO: prefer invoke_async + # Older Strands without stream_async. + result = await asyncio.to_thread(agent, user_prompt) if self.verbose: print(result) else: @@ -150,7 +150,7 @@ async def run_agent( self._render(event) if self.verbose and self._mid_line: print() # newline after trailing streamed text - return None # results collected via tool side effects (record_artifact -> sink) + return None # results collected via tool side effects (e.g. a record_tool sink) async def converse(self, task: Task) -> dict: from strands import tool @@ -159,11 +159,10 @@ async def converse(self, task: Task) -> dict: def finalize_setup(manifest: dict) -> str: """Finalize the setup manifest (matching the task's schema) and end setup.""" - captured["manifest"] = manifest - return "setup finalized" + return accept_manifest(task, manifest, captured) kwargs = { - "system_prompt": task.setup_system_prompt(), + "system_prompt": render_setup_prompt(task), "tools": [tool(finalize_setup)], } model = self._model() diff --git a/pyproject.toml b/pyproject.toml index cf00083..d7f75c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "behalf" -version = "0.1.1" +version = "0.2.0" description = "A small framework for scoped agents: a brief, a typed toolset (read/action), a pluggable model backend." readme = "README.md" requires-python = ">=3.10" diff --git a/tests/test_adk.py b/tests/test_adk.py index d7dd68e..8308b35 100644 --- a/tests/test_adk.py +++ b/tests/test_adk.py @@ -260,5 +260,113 @@ def test_make_runner_gemini(): assert make_runner("gemini").model == "gemini-2.5-flash" +# Real ADK Runner with a scripted BaseLlm in place of Gemini: exercises ADK's +# own tool dispatch, session handling, and multi-turn reuse without a key. + +from google.adk.models.base_llm import BaseLlm # noqa: E402 +from google.adk.models.llm_response import LlmResponse # noqa: E402 +from google.genai import types as genai_types # noqa: E402 + +from behalf import Task # noqa: E402 + + +class ScriptedLlm(BaseLlm): + """Replies with the next scripted step: a (tool, args) call or plain text.""" + + script: list = [] + requests: list = [] + + async def generate_content_async(self, llm_request, stream=False): + self.requests.append(llm_request) + step = self.script.pop(0) + if isinstance(step, tuple): + name, args = step + part = genai_types.Part( + function_call=genai_types.FunctionCall(name=name, args=args) + ) + else: + part = genai_types.Part(text=step) + yield LlmResponse(content=genai_types.Content(role="model", parts=[part])) + + +def _function_responses(request): + return [ + p.function_response.response + for c in request.contents + for p in c.parts or [] + if p.function_response + ] + + +def test_run_agent_with_real_adk_runner(): + log = [] + + async def look(args): + log.append(("look", args)) + return _text({"saw": args["q"]}) + + async def act(args): + log.append(("act", args)) + return _text("acted") + + tools = [ + ToolSpec("look", "look", {"q": str}, look), + ToolSpec("act", "act", {"n": int}, act, kind="action", confirm=True), + ] + llm = ScriptedLlm( + model="scripted", script=[("look", {"q": "hi"}), ("act", {"n": 3}), "done"] + ) + denied = [] + + def deny(name, args): + denied.append((name, args)) + return False + + runner = ADKRunner(model=llm, verbose=False) + assert asyncio.run(runner.run_agent("system", "go", tools, deny)) is None + assert log == [("look", {"q": "hi"})] + assert denied == [("act", {"n": 3})] + responses = _function_responses(llm.requests[-1]) + assert {"saw": "hi"} in responses + assert any("cancelled by user" in json.dumps(r) for r in responses) + + +class SetupTask(Task): + name = "setup" + + def manifest_schema(self): + return {"goal": str, "sites": list} + + def setup_system_prompt(self): + return "Plan it." + + def execute_system_prompt(self, manifest): + return "run" + + +def test_converse_with_real_adk_runner_reuses_session(monkeypatch): + llm = ScriptedLlm( + model="scripted", + script=[ + "Which sites?", + ("finalize_setup", {"manifest": {"goal": "g"}}), + ("finalize_setup", {"manifest": {"goal": "g", "sites": ["a"]}}), + "ok", + ], + ) + replies = iter(["site a"]) + monkeypatch.setattr("builtins.input", lambda prompt="": next(replies)) + manifest = asyncio.run(ADKRunner(model=llm, verbose=False).converse(SetupTask())) + assert manifest == {"goal": "g", "sites": ["a"]} + + first, *_, last = llm.requests + assert "- sites: list" in first.config.system_instruction + texts = [p.text for c in last.contents for p in c.parts or [] if p.text] + assert "Which sites?" in texts and "site a" in texts, "second turn must see first" + assert any( + "missing fields: sites" in json.dumps(r) for r in _function_responses(last) + ) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/test_core.py b/tests/test_core.py index 818369c..b3e5a53 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -3,7 +3,20 @@ import sys sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) -from behalf import RunOutcome, Task, ToolSpec, run_task +import json + +import pytest + +from behalf import ( + RunOutcome, + Task, + ToolSpec, + accept_manifest, + missing_manifest_fields, + record_tool, + render_setup_prompt, + run_task, +) class DemoTask(Task): @@ -101,6 +114,69 @@ def test_setup_approval_can_abort(): print("OK setup approval can abort before execution") +class SchemaTask(DemoTask): + name = "schema" + + def manifest_schema(self): + return {"goal": str, "sites": list, "budget": float} + + def setup_system_prompt(self): + return "Plan an experiment." + + +def test_setup_prompt_carries_schema(): + prompt = render_setup_prompt(SchemaTask()) + assert prompt.startswith("Plan an experiment.") + for line in ("- goal: str", "- sites: list", "- budget: float", "finalize_setup"): + assert line in prompt + + +def test_missing_manifest_fields(): + task = SchemaTask() + assert missing_manifest_fields(task, {"goal": "g"}) == ["sites", "budget"] + assert missing_manifest_fields(task, {"goal": 1, "sites": [], "budget": 0}) == [] + assert missing_manifest_fields(task, "not a dict") == ["goal", "sites", "budget"] + + +def test_accept_manifest_only_captures_complete_manifests(): + task, captured = SchemaTask(), {} + message = accept_manifest(task, {"goal": "g"}, captured) + assert "sites, budget" in message and captured == {} + full = {"goal": "g", "sites": ["a"], "budget": 1.0} + assert accept_manifest(task, full, captured) == "setup finalized" + assert captured == {"manifest": full} + + +def test_run_task_rejects_incomplete_manifest(): + task = SchemaTask() + with pytest.raises(ValueError, match="sites"): + asyncio.run( + run_task( + task, + FakeRunner({}), + manifest={"goal": "g"}, + confirm_fn=lambda n, a: True, + ) + ) + assert not hasattr(task, "acted"), "execute must not run" + + +def test_run_task_rejects_incomplete_conversed_manifest(): + with pytest.raises(ValueError, match="goal"): + asyncio.run(run_task(DemoTask(), FakeRunner({}), confirm_fn=lambda n, a: True)) + + +def test_record_tool_appends_to_sink(): + sink = [] + spec = record_tool(sink, {"job": str, "site": str}, name="record_placement") + assert (spec.name, spec.kind, spec.confirm) == ("record_placement", "action", False) + assert spec.input_schema == {"job": str, "site": str} + out = asyncio.run(spec.handler({"job": "j1", "site": "a"})) + asyncio.run(spec.handler({"job": "j2", "site": "b"})) + assert json.loads(out["content"][0]["text"]) == {"recorded": 1} + assert sink == [{"job": "j1", "site": "a"}, {"job": "j2", "site": "b"}] + + if __name__ == "__main__": test_manifest_passthrough_runs_execute() test_action_gate_blocks_on_denial() diff --git a/tests/test_sdk.py b/tests/test_sdk.py new file mode 100644 index 0000000..f5f22e4 --- /dev/null +++ b/tests/test_sdk.py @@ -0,0 +1,131 @@ +"""Tests for the Claude (Agent SDK) runner that need no API key or CLI. + +They are skipped when claude-agent-sdk is not installed (pip install behalf[claude]). +""" + +import asyncio +import json +import os +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) + +pytest.importorskip("claude_agent_sdk") + +import behalf.runner.sdk as sdk # noqa: E402 +from behalf import Task, ToolSpec # noqa: E402 +from behalf.runner.sdk import SDKRunner, _finalize_tool, _to_sdk_tool # noqa: E402 + + +def _text(obj) -> dict: + return {"content": [{"type": "text", "text": json.dumps(obj)}]} + + +class SetupTask(Task): + name = "setup" + + def manifest_schema(self): + return {"goal": str, "sites": list} + + def setup_system_prompt(self): + return "Plan it." + + def execute_system_prompt(self, manifest): + return "run" + + +def test_sdk_tool_carries_spec_and_gates_actions(): + acted = [] + + async def act(args): + acted.append(args) + return _text("acted") + + spec = ToolSpec("act", "do it", {"n": int}, act, kind="action", confirm=True) + denied = _to_sdk_tool(spec, lambda n, a: False) + assert (denied.name, denied.description, denied.input_schema) == ( + "act", + "do it", + {"n": int}, + ) + assert "cancelled by user" in json.dumps(asyncio.run(denied.handler({"n": 1}))) + assert acted == [] + asyncio.run(_to_sdk_tool(spec, lambda n, a: True).handler({"n": 2})) + assert acted == [{"n": 2}] + + +def test_finalize_tool_rejects_incomplete_manifest(): + captured = {} + finalize = _finalize_tool(SetupTask(), captured) + out = asyncio.run(finalize.handler({"manifest": {"goal": "g"}})) + assert "missing fields: sites" in json.dumps(out) and captured == {} + asyncio.run(finalize.handler({"manifest": {"goal": "g", "sites": ["a"]}})) + assert captured == {"manifest": {"goal": "g", "sites": ["a"]}} + + +def test_run_agent_passes_prompt_and_allowed_tools(monkeypatch): + seen = {} + + async def fake_query(prompt, options): + seen["prompt"], seen["options"] = prompt, options + yield "message" + + monkeypatch.setattr(sdk, "query", fake_query) + spec = ToolSpec("look", "look", {"q": str}, lambda a: None) + result = asyncio.run( + SDKRunner(verbose=False).run_agent("system", "go", [spec], lambda n, a: True) + ) + assert result is None and seen["prompt"] == "go" + assert seen["options"].system_prompt == "system" + assert seen["options"].allowed_tools == ["mcp__behalf__look"] + + +def test_converse_uses_schema_prompt_and_retries_incomplete(monkeypatch): + tools = [] + real_finalize = sdk._finalize_tool + + def spy(task, captured): + tools.append(real_finalize(task, captured)) + return tools[-1] + + class StubClient: + """Stands in for ClaudeSDKClient: each reply turn calls finalize_setup once.""" + + attempts = iter([{"goal": "g"}, {"goal": "g", "sites": ["a"]}]) + queries = [] + + def __init__(self, options): + StubClient.options = options + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def query(self, text): + StubClient.queries.append(text) + + async def receive_response(self): + await tools[0].handler({"manifest": next(StubClient.attempts)}) + yield "message" + + monkeypatch.setattr(sdk, "_finalize_tool", spy) + monkeypatch.setattr(sdk, "ClaudeSDKClient", StubClient) + monkeypatch.setattr("builtins.input", lambda prompt="": "site a") + manifest = asyncio.run(SDKRunner(verbose=False).converse(SetupTask())) + assert manifest == {"goal": "g", "sites": ["a"]} + assert "- sites: list" in StubClient.options.system_prompt + assert StubClient.queries[1:] == ["site a"] + + +def test_make_runner_claude(): + from behalf import make_runner + + assert make_runner("claude", "claude-test").model == "claude-test" + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/test_strands.py b/tests/test_strands.py new file mode 100644 index 0000000..ee4bbf9 --- /dev/null +++ b/tests/test_strands.py @@ -0,0 +1,187 @@ +"""Tests for the AWS (Strands) runner that need no AWS credentials. + +They are skipped when strands-agents is not installed (pip install behalf[aws]). +""" + +import asyncio +import json +import os +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) + +pytest.importorskip("strands") + +from strands.models import Model # noqa: E402 + +from behalf import Task, ToolSpec # noqa: E402 +from behalf.runner.strands import StrandsRunner, _to_strands_tool # noqa: E402 + + +def _text(obj) -> dict: + return {"content": [{"type": "text", "text": json.dumps(obj)}]} + + +class ScriptedModel(Model): + """Streams the next scripted step: a (tool, args) call or plain text.""" + + def __init__(self, script): + self.script = list(script) + self.calls = [] + + def get_config(self): + return {} + + def update_config(self, **model_config): + pass + + async def structured_output(self, output_model, prompt, system_prompt=None, **kw): + raise NotImplementedError + yield + + async def stream(self, messages, tool_specs=None, system_prompt=None, **kwargs): + self.calls.append( + { + "messages": json.loads(json.dumps(messages, default=str)), + "tools": [t["name"] for t in tool_specs or []], + "system_prompt": system_prompt, + } + ) + step = self.script.pop(0) + yield {"messageStart": {"role": "assistant"}} + if isinstance(step, tuple): + name, args = step + tool_use = {"toolUseId": f"t{len(self.calls)}", "name": name} + yield {"contentBlockStart": {"start": {"toolUse": tool_use}}} + yield { + "contentBlockDelta": {"delta": {"toolUse": {"input": json.dumps(args)}}} + } + yield {"contentBlockStop": {}} + yield {"messageStop": {"stopReason": "tool_use"}} + else: + yield {"contentBlockDelta": {"delta": {"text": step}}} + yield {"contentBlockStop": {}} + yield {"messageStop": {"stopReason": "end_turn"}} + + +def _runner(model) -> StrandsRunner: + runner = StrandsRunner(verbose=False) + runner._model = lambda: model + return runner + + +def _tool_results(call): + return [ + block["toolResult"] + for message in call["messages"] + for block in message["content"] + if "toolResult" in block + ] + + +def test_tool_spec_carries_every_argument_and_type(): + spec = ToolSpec( + "run", + "run something", + { + "name": str, + "n": int, + "ratio": float, + "flag": bool, + "items": list, + "opts": dict, + }, + lambda args: None, + ) + schema = _to_strands_tool(spec, lambda n, a: True).tool_spec + props = schema["inputSchema"]["json"]["properties"] + assert schema["name"] == "run" and schema["description"] == "run something" + assert {k: v["type"] for k, v in props.items()} == { + "name": "string", + "n": "integer", + "ratio": "number", + "flag": "boolean", + "items": "array", + "opts": "object", + } + + +def test_run_agent_with_real_strands_agent(): + log = [] + + async def look(args): + log.append(("look", args)) + return _text({"saw": args["q"]}) + + async def act(args): + log.append(("act", args)) + return _text("acted") + + tools = [ + ToolSpec("look", "look", {"q": str}, look), + ToolSpec("act", "act", {"n": int}, act, kind="action", confirm=True), + ] + model = ScriptedModel([("look", {"q": "hi"}), ("act", {"n": 3}), "done"]) + denied = [] + + def deny(name, args): + denied.append((name, args)) + return False + + assert asyncio.run(_runner(model).run_agent("system", "go", tools, deny)) is None + assert log == [("look", {"q": "hi"})] + assert denied == [("act", {"n": 3})] + first, *_, last = model.calls + assert first["system_prompt"] == "system" and first["tools"] == ["look", "act"] + results = json.dumps(_tool_results(last)) + assert "hi" in results and "cancelled by user" in results + + +class SetupTask(Task): + name = "setup" + + def manifest_schema(self): + return {"goal": str, "sites": list} + + def setup_system_prompt(self): + return "Plan it." + + def execute_system_prompt(self, manifest): + return "run" + + +def test_converse_with_real_strands_agent(monkeypatch): + model = ScriptedModel( + [ + "Which sites?", + ("finalize_setup", {"manifest": {"goal": "g"}}), + ("finalize_setup", {"manifest": {"goal": "g", "sites": ["a"]}}), + "ok", + ] + ) + replies = iter(["site a"]) + monkeypatch.setattr("builtins.input", lambda prompt="": next(replies)) + manifest = asyncio.run(_runner(model).converse(SetupTask())) + assert manifest == {"goal": "g", "sites": ["a"]} + + first, *_, last = model.calls + assert "- sites: list" in first["system_prompt"] + history = json.dumps(last["messages"]) + assert "Which sites?" in history and "site a" in history + assert "missing fields: sites" in json.dumps(_tool_results(last)) + + +def test_make_runner_aws(): + from behalf import make_runner + + runner = make_runner("aws", "us.anthropic.claude-test") + assert ( + isinstance(runner, StrandsRunner) and runner.model == "us.anthropic.claude-test" + ) + assert runner._model().get_config()["model_id"] == "us.anthropic.claude-test" + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))