wehub-resource-sync
153 lines
4.8 KiB
Python
153 lines
4.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
from opensquilla.engine import Agent, AgentConfig, ToolResult
|
|
from opensquilla.engine.types import ToolCall
|
|
from opensquilla.provider import (
|
|
ChatConfig,
|
|
Message,
|
|
ToolDefinition,
|
|
ToolInputSchema,
|
|
)
|
|
from opensquilla.provider import (
|
|
DoneEvent as ProviderDone,
|
|
)
|
|
from opensquilla.provider import (
|
|
TextDeltaEvent as ProviderText,
|
|
)
|
|
from opensquilla.provider import (
|
|
ToolUseEndEvent as ProviderToolUseEnd,
|
|
)
|
|
from opensquilla.provider import (
|
|
ToolUseStartEvent as ProviderToolUseStart,
|
|
)
|
|
|
|
|
|
class _ComplexTaskProvider:
|
|
provider_name = "fake"
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[list[Message]] = []
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[Any] | None = None,
|
|
config: ChatConfig | None = None,
|
|
) -> AsyncIterator[Any]:
|
|
self.calls.append(messages)
|
|
return self._stream(len(self.calls))
|
|
|
|
async def _stream(self, call_number: int) -> AsyncIterator[Any]:
|
|
if call_number == 1:
|
|
yield ProviderToolUseStart(tool_use_id="search-1", tool_name="web_search")
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id="search-1",
|
|
tool_name="web_search",
|
|
arguments={"query": "opensquilla regression"},
|
|
)
|
|
yield ProviderToolUseStart(tool_use_id="read-1", tool_name="read_file")
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id="read-1",
|
|
tool_name="read_file",
|
|
arguments={"path": "README.md"},
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=10, output_tokens=2)
|
|
return
|
|
|
|
if call_number == 2:
|
|
yield ProviderToolUseStart(tool_use_id="cmd-1", tool_name="exec_command")
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id="cmd-1",
|
|
tool_name="exec_command",
|
|
arguments={"cmd": "pytest targeted"},
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=12, output_tokens=2)
|
|
return
|
|
|
|
yield ProviderText(text="complex chain complete")
|
|
yield ProviderDone(stop_reason="stop", input_tokens=14, output_tokens=3)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
def _tool_def(name: str, properties: dict[str, Any]) -> ToolDefinition:
|
|
return ToolDefinition(
|
|
name=name,
|
|
description=f"Mock {name}.",
|
|
input_schema=ToolInputSchema(
|
|
properties=properties,
|
|
required=list(properties),
|
|
),
|
|
)
|
|
|
|
|
|
def test_complex_agent_chain_runs_multi_tool_multi_iteration_task() -> None:
|
|
async def run() -> None:
|
|
provider = _ComplexTaskProvider()
|
|
handled: list[tuple[str, dict[str, Any]]] = []
|
|
|
|
async def tool_handler(call: ToolCall) -> ToolResult:
|
|
handled.append((call.tool_name, dict(call.arguments)))
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content=json.dumps(
|
|
{
|
|
"ok": True,
|
|
"tool": call.tool_name,
|
|
"arguments": call.arguments,
|
|
},
|
|
sort_keys=True,
|
|
),
|
|
)
|
|
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=3),
|
|
tool_definitions=[
|
|
_tool_def("web_search", {"query": {"type": "string"}}),
|
|
_tool_def("read_file", {"path": {"type": "string"}}),
|
|
_tool_def("exec_command", {"cmd": {"type": "string"}}),
|
|
],
|
|
tool_handler=tool_handler,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("diagnose the regression")]
|
|
|
|
assert [name for name, _args in handled] == [
|
|
"web_search",
|
|
"read_file",
|
|
"exec_command",
|
|
]
|
|
assert len(provider.calls) == 3
|
|
assert any(
|
|
event.kind == "done" and event.text == "complex chain complete"
|
|
for event in events
|
|
)
|
|
|
|
second_call = provider.calls[1]
|
|
assert any(
|
|
message.role == "assistant"
|
|
and any(getattr(block, "id", "") == "search-1" for block in message.content)
|
|
for message in second_call
|
|
)
|
|
assert any(
|
|
message.role == "user"
|
|
and any(getattr(block, "tool_use_id", "") == "search-1" for block in message.content)
|
|
for message in second_call
|
|
)
|
|
third_call = provider.calls[2]
|
|
assert any(
|
|
message.role == "user"
|
|
and any(getattr(block, "tool_use_id", "") == "cmd-1" for block in message.content)
|
|
for message in third_call
|
|
)
|
|
|
|
asyncio.run(run())
|