wehub-resource-sync
264 lines
9.5 KiB
Python
264 lines
9.5 KiB
Python
"""Agent re-emit hop: provider ReasoningDeltaEvent -> engine ThinkingEvent.
|
|
|
|
The agent layer translates provider-level stream events into engine-level
|
|
events. Reasoning must cross this hop as a first-class ThinkingEvent (not folded
|
|
into assistant text), so the TUI can render it as a distinct thinking stream.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
from opensquilla.engine import Agent, AgentConfig
|
|
from opensquilla.engine.types import ThinkingEvent
|
|
from opensquilla.provider import ChatConfig, Message
|
|
from opensquilla.provider import DoneEvent as ProviderDone
|
|
from opensquilla.provider import ReasoningDeltaEvent as ProviderReasoning
|
|
from opensquilla.provider import TextDeltaEvent as ProviderText
|
|
|
|
|
|
class _ReasoningProvider:
|
|
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 ProviderReasoning(text="Let me ")
|
|
yield ProviderReasoning(text="think.")
|
|
yield ProviderText(text="The answer.")
|
|
yield ProviderDone(
|
|
stop_reason="stop",
|
|
input_tokens=5,
|
|
output_tokens=3,
|
|
reasoning_content="Let me think.",
|
|
)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
def test_agent_reemits_reasoning_as_thinking_event() -> None:
|
|
import asyncio
|
|
|
|
async def run() -> list[Any]:
|
|
agent = Agent(provider=_ReasoningProvider(), config=AgentConfig(max_iterations=1))
|
|
return [event async for event in agent.run_turn("hi")]
|
|
|
|
events = asyncio.run(run())
|
|
|
|
thinking = [e for e in events if isinstance(e, ThinkingEvent)]
|
|
assert [t.text for t in thinking] == ["Let me ", "think."]
|
|
|
|
# reasoning must NOT be folded into the assistant answer text
|
|
answer = "".join(
|
|
e.text for e in events if type(e).__name__ == "TextDeltaEvent"
|
|
)
|
|
assert answer == "The answer."
|
|
assert "think" not in answer
|
|
|
|
|
|
from opensquilla.engine import ToolResult # noqa: E402
|
|
from opensquilla.engine.types import TextDeltaEvent, ToolCall, ToolUseStartEvent # noqa: E402
|
|
from opensquilla.provider import ( # noqa: E402
|
|
ToolDefinition,
|
|
ToolInputSchema,
|
|
)
|
|
from opensquilla.provider import ToolUseEndEvent as ProviderToolEnd # noqa: E402
|
|
from opensquilla.provider import ToolUseStartEvent as ProviderToolStart # noqa: E402
|
|
|
|
|
|
class _TextThenToolThenAnswerProvider:
|
|
"""Round 1: narration text + a tool call. Round 2: final answer text."""
|
|
|
|
provider_name = "fake"
|
|
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[Any] | None = None,
|
|
config: ChatConfig | None = None,
|
|
) -> AsyncIterator[Any]:
|
|
self.calls += 1
|
|
return self._stream(self.calls)
|
|
|
|
async def _stream(self, call: int) -> AsyncIterator[Any]:
|
|
if call == 1:
|
|
# text the model speaks before deciding to call a tool
|
|
yield ProviderText(text="Let me look ")
|
|
yield ProviderText(text="at the file.")
|
|
yield ProviderToolStart(tool_use_id="t1", tool_name="read")
|
|
yield ProviderToolEnd(tool_use_id="t1", tool_name="read", arguments={})
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=5, output_tokens=2)
|
|
return
|
|
# final answer round: no tool
|
|
yield ProviderText(text="Here is the answer.")
|
|
yield ProviderDone(stop_reason="end_turn", input_tokens=6, output_tokens=3)
|
|
|
|
|
|
def test_text_streams_live_as_answer_until_a_tool_appears() -> None:
|
|
"""Issue #358 regression guard: assistant text streams live, token by token, the
|
|
moment it arrives instead of being buffered until the provider call ends. Until a
|
|
tool appears the running text is the answer (the common plain-Q&A case), so nothing
|
|
is held back and the UI updates incrementally rather than freezing then dumping the
|
|
whole reply at once."""
|
|
import asyncio
|
|
|
|
async def tool_handler(call: ToolCall) -> ToolResult:
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content="{}",
|
|
)
|
|
|
|
async def run() -> list[Any]:
|
|
agent = Agent(
|
|
provider=_TextThenToolThenAnswerProvider(),
|
|
config=AgentConfig(max_iterations=3),
|
|
tool_definitions=[
|
|
ToolDefinition(
|
|
name="read",
|
|
description="read a file",
|
|
input_schema=ToolInputSchema(properties={}, required=[]),
|
|
)
|
|
],
|
|
tool_handler=tool_handler,
|
|
)
|
|
return [event async for event in agent.run_turn("hi")]
|
|
|
|
events = asyncio.run(run())
|
|
text_events = [e for e in events if isinstance(e, TextDeltaEvent)]
|
|
|
|
# Every provider text delta is surfaced as its own live event — no buffering and
|
|
# no end-of-call flush that lumps the reply into one blob.
|
|
assert [e.text for e in text_events] == [
|
|
"Let me look ",
|
|
"at the file.",
|
|
"Here is the answer.",
|
|
]
|
|
# No tool has appeared when any of this text arrives, so all of it streams as the
|
|
# answer; nothing is withheld as intermediate.
|
|
intermediate = "".join(e.text for e in text_events if e.presentation == "intermediate")
|
|
answer = "".join(e.text for e in text_events if e.presentation == "answer")
|
|
assert intermediate == ""
|
|
assert answer == "Let me look at the file.Here is the answer."
|
|
|
|
|
|
class _ToolThenTextThenAnswerProvider:
|
|
"""Round 1: a tool call, then narration text *after* it. Round 2: final answer."""
|
|
|
|
provider_name = "fake"
|
|
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[Any] | None = None,
|
|
config: ChatConfig | None = None,
|
|
) -> AsyncIterator[Any]:
|
|
self.calls += 1
|
|
return self._stream(self.calls)
|
|
|
|
async def _stream(self, call: int) -> AsyncIterator[Any]:
|
|
if call == 1:
|
|
yield ProviderToolStart(tool_use_id="t1", tool_name="read")
|
|
yield ProviderToolEnd(tool_use_id="t1", tool_name="read", arguments={})
|
|
# narration the model speaks between tool calls
|
|
yield ProviderText(text="Now let me ")
|
|
yield ProviderText(text="explain.")
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=5, output_tokens=2)
|
|
return
|
|
yield ProviderText(text="Final answer.")
|
|
yield ProviderDone(stop_reason="end_turn", input_tokens=6, output_tokens=3)
|
|
|
|
|
|
def test_text_after_a_tool_is_intermediate_narration() -> None:
|
|
"""Once a tool has appeared in a round, later text in that round is intermediate
|
|
narration between tools (purple), while the final no-tool round is the answer."""
|
|
import asyncio
|
|
|
|
async def tool_handler(call: ToolCall) -> ToolResult:
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content="{}",
|
|
)
|
|
|
|
async def run() -> list[Any]:
|
|
agent = Agent(
|
|
provider=_ToolThenTextThenAnswerProvider(),
|
|
config=AgentConfig(max_iterations=3),
|
|
tool_definitions=[
|
|
ToolDefinition(
|
|
name="read",
|
|
description="read a file",
|
|
input_schema=ToolInputSchema(properties={}, required=[]),
|
|
)
|
|
],
|
|
tool_handler=tool_handler,
|
|
)
|
|
return [event async for event in agent.run_turn("hi")]
|
|
|
|
events = asyncio.run(run())
|
|
text_events = [e for e in events if isinstance(e, TextDeltaEvent)]
|
|
|
|
intermediate = "".join(e.text for e in text_events if e.presentation == "intermediate")
|
|
answer = "".join(e.text for e in text_events if e.presentation == "answer")
|
|
|
|
assert intermediate == "Now let me explain."
|
|
assert answer == "Final answer."
|
|
|
|
|
|
def test_tool_use_start_carries_server_start_timestamp() -> None:
|
|
"""Issue #329: the engine stamps ToolUseStartEvent with a server wall-clock start
|
|
time (epoch ms). A client seeds the running tool's elapsed timer from this instead
|
|
of its own clock, so the timer survives page switches / stream replay (where the
|
|
component remounts) rather than restarting from zero on every remount."""
|
|
import asyncio
|
|
import time
|
|
|
|
async def tool_handler(call: ToolCall) -> ToolResult:
|
|
return ToolResult(
|
|
tool_use_id=call.tool_use_id,
|
|
tool_name=call.tool_name,
|
|
content="{}",
|
|
)
|
|
|
|
async def run() -> list[Any]:
|
|
agent = Agent(
|
|
provider=_ToolThenTextThenAnswerProvider(),
|
|
config=AgentConfig(max_iterations=3),
|
|
tool_definitions=[
|
|
ToolDefinition(
|
|
name="read",
|
|
description="read a file",
|
|
input_schema=ToolInputSchema(properties={}, required=[]),
|
|
)
|
|
],
|
|
tool_handler=tool_handler,
|
|
)
|
|
return [event async for event in agent.run_turn("hi")]
|
|
|
|
before_ms = int(time.time() * 1000)
|
|
events = asyncio.run(run())
|
|
after_ms = int(time.time() * 1000)
|
|
|
|
starts = [e for e in events if isinstance(e, ToolUseStartEvent)]
|
|
assert starts, "expected at least one tool_use_start event"
|
|
for ev in starts:
|
|
# Stamped with a real server time within the turn's wall-clock window, not the
|
|
# 0 "unstamped" sentinel.
|
|
assert ev.started_at > 0
|
|
assert before_ms <= ev.started_at <= after_ms
|