wehub-resource-sync
624 lines
20 KiB
Python
624 lines
20 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.subagent import SubagentSpec
|
|
from opensquilla.engine.types import ArtifactEvent, ErrorEvent
|
|
from opensquilla.engine.usage import UsageTracker
|
|
from opensquilla.provider import (
|
|
ChatConfig,
|
|
Message,
|
|
ToolDefinition,
|
|
ToolInputSchema,
|
|
)
|
|
from opensquilla.provider import (
|
|
DoneEvent as ProviderDone,
|
|
)
|
|
from opensquilla.provider import (
|
|
ErrorEvent as ProviderError,
|
|
)
|
|
from opensquilla.provider import (
|
|
TextDeltaEvent as ProviderText,
|
|
)
|
|
from opensquilla.provider import (
|
|
ToolUseEndEvent as ProviderToolUseEnd,
|
|
)
|
|
from opensquilla.provider import (
|
|
ToolUseStartEvent as ProviderToolUseStart,
|
|
)
|
|
|
|
|
|
class _LoopingToolProvider:
|
|
provider_name = "fake"
|
|
|
|
def __init__(self, *, final_on_call: int | None = None) -> None:
|
|
self.final_on_call = final_on_call
|
|
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)
|
|
call_number = len(self.calls)
|
|
return self._stream(call_number)
|
|
|
|
async def _stream(self, call_number: int) -> AsyncIterator[Any]:
|
|
if self.final_on_call == call_number:
|
|
yield ProviderText(text="done")
|
|
yield ProviderDone(stop_reason="stop", input_tokens=1, output_tokens=1)
|
|
return
|
|
|
|
tool_id = f"tool-{call_number}"
|
|
yield ProviderToolUseStart(tool_use_id=tool_id, tool_name="echo")
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id=tool_id,
|
|
tool_name="echo",
|
|
arguments={"value": "again"},
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=1, output_tokens=1)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
|
|
|
|
class _ToolThenProviderErrorProvider:
|
|
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)
|
|
call_number = len(self.calls)
|
|
return self._stream(call_number)
|
|
|
|
async def _stream(self, call_number: int) -> AsyncIterator[Any]:
|
|
if call_number == 1:
|
|
yield ProviderToolUseStart(tool_use_id="tool-1", tool_name="echo")
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id="tool-1",
|
|
tool_name="echo",
|
|
arguments={"value": "again"},
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=1, output_tokens=1)
|
|
return
|
|
yield ProviderError(message="request timed out", code="request_error")
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
class _DoneUsageProvider:
|
|
provider_name = "fake"
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
input_tokens: int = 1,
|
|
output_tokens: int = 1,
|
|
billed_cost: float = 0.0,
|
|
) -> None:
|
|
self.input_tokens = input_tokens
|
|
self.output_tokens = output_tokens
|
|
self.billed_cost = billed_cost
|
|
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()
|
|
|
|
async def _stream(self) -> AsyncIterator[Any]:
|
|
yield ProviderText(text="done")
|
|
yield ProviderDone(
|
|
stop_reason="stop",
|
|
input_tokens=self.input_tokens,
|
|
output_tokens=self.output_tokens,
|
|
billed_cost=self.billed_cost,
|
|
)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
class _DoneBreakdownProvider:
|
|
provider_name = "fake"
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[Any] | None = None,
|
|
config: ChatConfig | None = None,
|
|
) -> AsyncIterator[Any]:
|
|
return self._stream()
|
|
|
|
async def _stream(self) -> AsyncIterator[Any]:
|
|
yield ProviderText(text="done")
|
|
yield ProviderDone(
|
|
stop_reason="stop",
|
|
input_tokens=2000,
|
|
output_tokens=1000,
|
|
billed_cost=0.01,
|
|
model="z-ai/glm-5.2",
|
|
cost_source="mixed",
|
|
model_usage_breakdown=[
|
|
{
|
|
"role": "proposer",
|
|
"provider": "openrouter",
|
|
"model": "deepseek/deepseek-v4-pro-20260423",
|
|
"input_tokens": 1000,
|
|
"output_tokens": 1000,
|
|
"billed_cost": 0.0,
|
|
"cost_source": "none",
|
|
},
|
|
{
|
|
"role": "aggregator",
|
|
"provider": "openrouter",
|
|
"model": "z-ai/glm-5.2",
|
|
"input_tokens": 1000,
|
|
"output_tokens": 0,
|
|
"billed_cost": 0.01,
|
|
"cost_source": "provider_billed",
|
|
},
|
|
],
|
|
)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
class _ArtifactThenProviderErrorProvider:
|
|
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)
|
|
call_number = len(self.calls)
|
|
return self._stream(call_number)
|
|
|
|
async def _stream(self, call_number: int) -> AsyncIterator[Any]:
|
|
if call_number == 1:
|
|
tool_id = "publish-1"
|
|
yield ProviderToolUseStart(
|
|
tool_use_id=tool_id,
|
|
tool_name="publish_artifact",
|
|
)
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id=tool_id,
|
|
tool_name="publish_artifact",
|
|
arguments={"path": "report.txt"},
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=1, output_tokens=1)
|
|
return
|
|
|
|
yield ProviderError(message="request timed out", code="request_error")
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
async def _echo_tool(call: Any) -> ToolResult:
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
|
|
async def _error_tool(call: Any) -> ToolResult:
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content="invalid arguments",
|
|
is_error=True,
|
|
)
|
|
|
|
|
|
def _echo_definition() -> ToolDefinition:
|
|
return ToolDefinition(
|
|
name="echo",
|
|
description="Echo.",
|
|
input_schema=ToolInputSchema(
|
|
properties={"value": {"type": "string"}},
|
|
required=["value"],
|
|
),
|
|
)
|
|
|
|
|
|
def _publish_artifact_definition() -> ToolDefinition:
|
|
return ToolDefinition(
|
|
name="publish_artifact",
|
|
description="Publish an artifact.",
|
|
input_schema=ToolInputSchema(
|
|
properties={"path": {"type": "string"}},
|
|
required=["path"],
|
|
),
|
|
)
|
|
|
|
|
|
async def _artifact_tool(call: Any) -> ToolResult:
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content="published",
|
|
artifacts=[
|
|
{
|
|
"id": "art-1",
|
|
"name": "report.txt",
|
|
"mime": "text/plain",
|
|
"size": 12,
|
|
"sha256": "a" * 64,
|
|
"session_id": "session-1",
|
|
"session_key": "agent:main:webchat:session-1",
|
|
"source": "publish_artifact",
|
|
"created_at": "2026-05-06T12:00:00Z",
|
|
"download_url": "/api/v1/artifacts/art-1",
|
|
}
|
|
],
|
|
)
|
|
|
|
|
|
def test_agent_iteration_defaults_are_unbounded() -> None:
|
|
assert AgentConfig().max_iterations == 0
|
|
assert SubagentSpec(task="check").max_iterations == 0
|
|
assert AgentConfig().max_turn_llm_calls == 0
|
|
assert AgentConfig().max_turn_input_tokens == 0
|
|
assert AgentConfig().max_turn_output_tokens == 0
|
|
assert AgentConfig().max_turn_billed_cost_usd == 0.0
|
|
assert AgentConfig().max_turn_tool_errors == 0
|
|
assert AgentConfig().length_capped_continuations == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_default_max_iterations_allows_long_tool_loop_to_finish() -> None:
|
|
provider = _LoopingToolProvider(final_on_call=101)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(),
|
|
tool_definitions=[_echo_definition()],
|
|
tool_handler=_echo_tool,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
|
|
assert len(provider.calls) == 101
|
|
assert not any(event.kind == "error" and event.code == "max_iterations" for event in events)
|
|
assert any(event.kind == "done" and event.text == "done" for event in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_finalizes_when_tool_loop_reaches_max_iterations() -> None:
|
|
provider = _LoopingToolProvider(final_on_call=2)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=1),
|
|
tool_definitions=[_echo_definition()],
|
|
tool_handler=_echo_tool,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
|
|
assert len(provider.calls) == 2
|
|
assert provider.calls[-1][-1].role == "user"
|
|
assert "Do not call tools" in str(provider.calls[-1][-1].content)
|
|
assert "Do not call tools" not in "\n".join(
|
|
str(message.content) for message in agent._history
|
|
)
|
|
assert not any(event.kind == "state" and event.state.value == "error" for event in events)
|
|
assert not any(event.kind == "error" and event.code == "max_iterations" for event in events)
|
|
assert any(event.kind == "done" and event.text == "done" for event in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_reports_partial_max_iterations_after_finalization_attempt_fails() -> None:
|
|
provider = _LoopingToolProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=1),
|
|
tool_definitions=[_echo_definition()],
|
|
tool_handler=_echo_tool,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
|
|
assert len(provider.calls) == 2
|
|
assert not any(event.kind == "state" and event.state.value == "error" for event in events)
|
|
assert any(
|
|
event.kind == "done"
|
|
and "best partial result" in event.text
|
|
for event in events
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_returns_partial_when_max_iteration_finalization_provider_fails() -> None:
|
|
provider = _ToolThenProviderErrorProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=1, max_provider_retries=0),
|
|
tool_definitions=[_echo_definition()],
|
|
tool_handler=_echo_tool,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
|
|
assert len(provider.calls) == 2
|
|
assert not any(event.kind == "state" and event.state.value == "error" for event in events)
|
|
assert not any(event.kind == "error" for event in events)
|
|
done_texts = [event.text for event in events if event.kind == "done"]
|
|
assert len(done_texts) == 1
|
|
assert done_texts[0].count("best partial result") == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_allows_final_response_on_last_iteration() -> None:
|
|
provider = _LoopingToolProvider(final_on_call=2)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_echo_definition()],
|
|
tool_handler=_echo_tool,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
|
|
assert len(provider.calls) == 2
|
|
assert not any(event.kind == "error" and event.code == "max_iterations" for event in events)
|
|
assert any(event.kind == "done" and event.text == "done" for event in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_emits_artifact_event_independent_of_tool_result_text() -> None:
|
|
provider = _LoopingToolProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=1),
|
|
tool_definitions=[_echo_definition()],
|
|
tool_handler=_artifact_tool,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
|
|
artifact_events = [event for event in events if isinstance(event, ArtifactEvent)]
|
|
assert len(artifact_events) == 1
|
|
assert artifact_events[0].id == "art-1"
|
|
assert artifact_events[0].download_url == "/api/v1/artifacts/art-1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_synthesizes_final_artifact_response_without_provider_call() -> None:
|
|
provider = _ArtifactThenProviderErrorProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=3, max_provider_retries=0),
|
|
tool_definitions=[_publish_artifact_definition()],
|
|
tool_handler=_artifact_tool,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("publish a report")]
|
|
|
|
assert len(provider.calls) == 1
|
|
assert any(isinstance(event, ArtifactEvent) for event in events)
|
|
assert not any(isinstance(event, ErrorEvent) for event in events)
|
|
assert not any(event.kind == "warning" for event in events)
|
|
assert any(
|
|
event.kind == "done"
|
|
and event.text == "The generated file is ready: report.txt."
|
|
for event in events
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_synthesizes_final_artifact_response_before_extra_llm_call() -> None:
|
|
provider = _ArtifactThenProviderErrorProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=3, max_turn_llm_calls=1),
|
|
tool_definitions=[_publish_artifact_definition()],
|
|
tool_handler=_artifact_tool,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("publish a report")]
|
|
|
|
assert len(provider.calls) == 1
|
|
assert any(isinstance(event, ArtifactEvent) for event in events)
|
|
assert not any(isinstance(event, ErrorEvent) for event in events)
|
|
assert not any(event.kind == "warning" for event in events)
|
|
assert any(
|
|
event.kind == "done"
|
|
and event.text == "The generated file is ready: report.txt."
|
|
for event in events
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_stops_when_turn_input_token_budget_is_exceeded() -> None:
|
|
provider = _LoopingToolProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(
|
|
max_iterations=100,
|
|
max_turn_input_tokens=1,
|
|
max_turn_output_tokens=0,
|
|
max_turn_billed_cost_usd=0,
|
|
max_turn_tool_errors=0,
|
|
),
|
|
tool_definitions=[_echo_definition()],
|
|
tool_handler=_echo_tool,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
|
|
assert len(provider.calls) == 2
|
|
assert any(
|
|
event.kind == "error" and event.code == "turn_input_token_budget_exceeded"
|
|
for event in events
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_stops_when_turn_llm_call_budget_is_exceeded() -> None:
|
|
provider = _LoopingToolProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(
|
|
max_iterations=100,
|
|
max_turn_llm_calls=1,
|
|
),
|
|
tool_definitions=[_echo_definition()],
|
|
tool_handler=_echo_tool,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
|
|
assert len(provider.calls) == 1
|
|
assert any(
|
|
event.kind == "error" and event.code == "turn_llm_call_budget_exceeded"
|
|
for event in events
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_stops_when_turn_output_token_budget_is_exceeded() -> None:
|
|
provider = _DoneUsageProvider(output_tokens=10)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_turn_output_tokens=5),
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
|
|
assert len(provider.calls) == 1
|
|
assert any(
|
|
event.kind == "error" and event.code == "turn_output_token_budget_exceeded"
|
|
for event in events
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_done_event_uses_current_turn_real_billed_usage_delta() -> None:
|
|
tracker = UsageTracker()
|
|
session_key = "agent:test:webchat:s1"
|
|
tracker.add(
|
|
session_key,
|
|
input_tokens=100,
|
|
output_tokens=10,
|
|
model_id="deepseek/deepseek-v4-pro-20260423",
|
|
billed_cost=0.050,
|
|
)
|
|
provider = _DoneUsageProvider(input_tokens=9, output_tokens=4, billed_cost=0.123)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(model_id="deepseek/deepseek-v4-pro-20260423"),
|
|
usage_tracker=tracker,
|
|
session_key=session_key,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
done = next(event for event in events if event.kind == "done")
|
|
|
|
assert done.input_tokens == 9
|
|
assert done.output_tokens == 4
|
|
assert done.cost_usd == pytest.approx(0.123)
|
|
assert done.billed_cost == pytest.approx(0.123)
|
|
assert done.cost_source == "provider_billed"
|
|
assert done.session_totals is not None
|
|
assert done.session_totals.billed_cost == pytest.approx(0.173)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_enriches_model_usage_breakdown_with_estimated_costs() -> None:
|
|
tracker = UsageTracker()
|
|
session_key = "agent:test:webchat:ensemble-costs"
|
|
agent = Agent(
|
|
provider=_DoneBreakdownProvider(),
|
|
config=AgentConfig(model_id="z-ai/glm-5.2"),
|
|
usage_tracker=tracker,
|
|
session_key=session_key,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
done = next(event for event in events if event.kind == "done")
|
|
|
|
# deepseek/deepseek-v4-pro static price is $0.435/M in, $0.87/M out (see
|
|
# engine/pricing.py): (1000 * 0.435 + 1000 * 0.87) / 1e6 == 0.001305.
|
|
assert done.cost_source == "mixed"
|
|
assert done.cost_usd == pytest.approx(0.011305)
|
|
assert done.billed_cost == pytest.approx(0.01)
|
|
deepseek_row = done.model_usage_breakdown[0]
|
|
assert deepseek_row["model"] == "deepseek/deepseek-v4-pro-20260423"
|
|
assert deepseek_row["cost_usd"] == pytest.approx(0.001305)
|
|
assert deepseek_row["estimated_cost_usd"] == pytest.approx(0.001305)
|
|
assert deepseek_row["billed_cost_usd"] == 0.0
|
|
assert deepseek_row["cost_source"] == "opensquilla_estimate"
|
|
aggregator_row = done.model_usage_breakdown[1]
|
|
assert aggregator_row["cost_usd"] == pytest.approx(0.01)
|
|
assert aggregator_row["billed_cost_usd"] == pytest.approx(0.01)
|
|
assert aggregator_row["cost_source"] == "provider_billed"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_stops_when_turn_billed_cost_budget_is_exceeded() -> None:
|
|
provider = _DoneUsageProvider(billed_cost=0.25)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_turn_billed_cost_usd=0.1),
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
|
|
assert len(provider.calls) == 1
|
|
assert any(
|
|
event.kind == "error" and event.code == "turn_billed_cost_budget_exceeded"
|
|
for event in events
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_stops_when_turn_tool_error_budget_is_exceeded() -> None:
|
|
provider = _LoopingToolProvider()
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(
|
|
max_iterations=100,
|
|
max_turn_input_tokens=0,
|
|
max_turn_output_tokens=0,
|
|
max_turn_billed_cost_usd=0,
|
|
max_turn_tool_errors=2,
|
|
),
|
|
tool_definitions=[_echo_definition()],
|
|
tool_handler=_error_tool,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
|
|
assert len(provider.calls) == 2
|
|
assert any(
|
|
event.kind == "error" and event.code == "turn_tool_error_budget_exceeded"
|
|
for event in events
|
|
)
|