wehub-resource-sync
273 lines
8.5 KiB
Python
273 lines
8.5 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from opensquilla.engine import Agent, AgentConfig, ToolResult
|
|
from opensquilla.engine.agent_injection import ListPendingInputProvider
|
|
from opensquilla.engine.types import ToolCall
|
|
from opensquilla.provider import (
|
|
ChatConfig,
|
|
ContentBlockText,
|
|
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 _ToolBoundaryProvider:
|
|
provider_name = "fake"
|
|
|
|
def __init__(self, *, tool_iterations: int = 1) -> None:
|
|
self.tool_iterations = tool_iterations
|
|
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(list(messages))
|
|
return self._stream(len(self.calls))
|
|
|
|
async def _stream(self, call_number: int) -> AsyncIterator[Any]:
|
|
if call_number <= self.tool_iterations:
|
|
tool_use_id = f"tool-{call_number}"
|
|
yield ProviderToolUseStart(tool_use_id=tool_use_id, tool_name="echo")
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id=tool_use_id,
|
|
tool_name="echo",
|
|
arguments={"value": call_number},
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=1, output_tokens=1)
|
|
return
|
|
|
|
yield ProviderText(text="done")
|
|
yield ProviderDone(stop_reason="stop", input_tokens=1, output_tokens=1)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
class _NoToolProvider:
|
|
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(list(messages))
|
|
return self._stream()
|
|
|
|
async def _stream(self) -> AsyncIterator[Any]:
|
|
yield ProviderText(text="done")
|
|
yield ProviderDone(stop_reason="stop", input_tokens=1, output_tokens=1)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
def _tool_def(name: str) -> ToolDefinition:
|
|
return ToolDefinition(
|
|
name=name,
|
|
description=f"Mock {name}.",
|
|
input_schema=ToolInputSchema(properties={}, required=[]),
|
|
)
|
|
|
|
|
|
async def _tool_handler(call: ToolCall) -> ToolResult:
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content=f"result from {call.tool_name} {call.tool_use_id}",
|
|
)
|
|
|
|
|
|
def _agent(provider: Any, *, max_iterations: int = 3) -> Agent:
|
|
return Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=max_iterations),
|
|
tool_definitions=[_tool_def("echo")],
|
|
tool_handler=_tool_handler,
|
|
)
|
|
|
|
|
|
def _text_block_texts(message: Message) -> list[str]:
|
|
if not isinstance(message.content, list):
|
|
return []
|
|
return [block.text for block in message.content if isinstance(block, ContentBlockText)]
|
|
|
|
|
|
def _is_tool_result_message(message: Message) -> bool:
|
|
return message.role == "user" and isinstance(message.content, list) and any(
|
|
getattr(block, "type", None) == "tool_result" for block in message.content
|
|
)
|
|
|
|
|
|
def _message_text(message: Message) -> str:
|
|
if isinstance(message.content, str):
|
|
return message.content
|
|
if not isinstance(message.content, list):
|
|
return ""
|
|
return "\n".join(
|
|
block.text for block in message.content if isinstance(block, ContentBlockText)
|
|
)
|
|
|
|
|
|
def _tool_result_index(messages: list[Message]) -> int:
|
|
return next(index for index, message in enumerate(messages) if _is_tool_result_message(message))
|
|
|
|
|
|
def _text_message_index(messages: list[Message], texts: list[str]) -> int:
|
|
return next(
|
|
index
|
|
for index, message in enumerate(messages)
|
|
if message.role == "user" and _text_block_texts(message) == texts
|
|
)
|
|
|
|
|
|
def _text_messages(messages: list[Message], texts: list[str]) -> list[Message]:
|
|
return [
|
|
message
|
|
for message in messages
|
|
if message.role == "user" and _text_block_texts(message) == texts
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pending_input_is_injected_after_tool_result_and_seen_by_next_model_request() -> None:
|
|
provider = _ToolBoundaryProvider()
|
|
pending = ListPendingInputProvider()
|
|
pending.append("INJECTED")
|
|
agent = _agent(provider)
|
|
|
|
events = [event async for event in agent.run_turn("run echo", pending_input_provider=pending)]
|
|
|
|
assert len(provider.calls) == 2
|
|
second_request = provider.calls[1]
|
|
assert _text_message_index(second_request, ["INJECTED"]) == (
|
|
_tool_result_index(second_request) + 1
|
|
)
|
|
assert any(event.kind == "done" and event.text == "done" for event in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_pending_inputs_are_merged_into_one_user_message() -> None:
|
|
provider = _ToolBoundaryProvider()
|
|
pending = ListPendingInputProvider()
|
|
pending.append("A")
|
|
pending.append("B")
|
|
agent = _agent(provider)
|
|
|
|
_events = [event async for event in agent.run_turn("run echo", pending_input_provider=pending)]
|
|
|
|
second_request = provider.calls[1]
|
|
injected_messages = _text_messages(second_request, ["A", "B"])
|
|
assert len(injected_messages) == 1
|
|
assert isinstance(injected_messages[0].content, list)
|
|
assert _text_block_texts(injected_messages[0]) == ["A", "B"]
|
|
assert not _text_messages(second_request, ["A"])
|
|
assert not _text_messages(second_request, ["B"])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_pending_provider_anchors_current_request_after_tool_result(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_TURN_OBJECTIVE_REMINDER", "on")
|
|
provider = _ToolBoundaryProvider()
|
|
agent = _agent(provider)
|
|
|
|
_events = [event async for event in agent.run_turn("run echo")]
|
|
|
|
assert len(provider.calls) == 2
|
|
second_request = provider.calls[1]
|
|
tool_result_index = _tool_result_index(second_request)
|
|
reminder_text = _message_text(second_request[tool_result_index + 1])
|
|
assert "Current user request" in reminder_text
|
|
assert "run echo" in reminder_text
|
|
assert not any(
|
|
"Current user request" in _message_text(message) for message in agent._history
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_pending_default_sends_no_reminder_after_tool_result(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.delenv("OPENSQUILLA_TURN_OBJECTIVE_REMINDER", raising=False)
|
|
provider = _ToolBoundaryProvider()
|
|
agent = _agent(provider)
|
|
|
|
_events = [event async for event in agent.run_turn("run echo")]
|
|
|
|
assert len(provider.calls) == 2
|
|
second_request = provider.calls[1]
|
|
assert not any(
|
|
"Current user request" in _message_text(message) for message in second_request
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_drained_pending_input_is_not_injected_again_at_later_tool_boundaries() -> None:
|
|
provider = _ToolBoundaryProvider(tool_iterations=2)
|
|
pending = ListPendingInputProvider()
|
|
pending.append("ONCE")
|
|
agent = _agent(provider, max_iterations=3)
|
|
|
|
_events = [event async for event in agent.run_turn("run echo", pending_input_provider=pending)]
|
|
|
|
assert len(provider.calls) == 3
|
|
assert len(pending) == 0
|
|
assert len(_text_messages(agent._history, ["ONCE"])) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_injected_pending_input_is_persisted_to_successful_turn_history() -> None:
|
|
provider = _ToolBoundaryProvider()
|
|
pending = ListPendingInputProvider()
|
|
pending.append("HISTORY")
|
|
agent = _agent(provider)
|
|
|
|
_events = [event async for event in agent.run_turn("run echo", pending_input_provider=pending)]
|
|
|
|
assert _text_message_index(agent._history, ["HISTORY"]) == (
|
|
_tool_result_index(agent._history) + 1
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pending_input_is_not_drained_without_a_tool_completion_boundary() -> None:
|
|
provider = _NoToolProvider()
|
|
pending = ListPendingInputProvider()
|
|
pending.append("NO_BOUNDARY")
|
|
agent = _agent(provider)
|
|
|
|
_events = [
|
|
event async for event in agent.run_turn("just answer", pending_input_provider=pending)
|
|
]
|
|
|
|
assert len(provider.calls) == 1
|
|
assert len(pending) == 1
|
|
assert not _text_messages(agent._history, ["NO_BOUNDARY"])
|