403 lines
13 KiB
Python
403 lines
13 KiB
Python
"""Tests for ``pallas.assistant_stream``.
|
|
|
|
Drives the ``after_llm_call`` / ``before_tool_call`` / ``after_tool_call``
|
|
hooks with handcrafted ``PromptMessageExtended`` objects and asserts the
|
|
resulting MCP ``send_log_message`` payload shape. No fast-agent runtime is
|
|
involved — the hooks are pure async functions and the MCP context is faked.
|
|
|
|
Tests use ``asyncio.run`` directly to match the convention in
|
|
``tests/test_health.py`` and ``tests/test_mantle_shims.py`` (pallas has no
|
|
pytest-asyncio dependency).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any
|
|
|
|
from fast_agent.agents.tool_runner import ToolRunnerHooks
|
|
from fast_agent.types import PromptMessageExtended
|
|
from fast_agent.types.llm_stop_reason import LlmStopReason
|
|
from mcp.types import (
|
|
CallToolRequest,
|
|
CallToolRequestParams,
|
|
CallToolResult,
|
|
ImageContent,
|
|
TextContent,
|
|
)
|
|
|
|
from pallas.assistant_stream import (
|
|
KIND,
|
|
KIND_RESULTS,
|
|
LOGGER_NAME,
|
|
SCHEMA_VERSION,
|
|
AssistantChunkEmitter,
|
|
install_for_request,
|
|
)
|
|
|
|
|
|
def _run(coro):
|
|
return asyncio.run(coro)
|
|
|
|
|
|
# ── Fakes ────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class _FakeSession:
|
|
"""Records every call to ``send_log_message`` for later assertion."""
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[dict[str, Any]] = []
|
|
self.fail_with: Exception | None = None
|
|
|
|
async def send_log_message(
|
|
self,
|
|
*,
|
|
level: str,
|
|
data: Any,
|
|
logger: str | None = None,
|
|
related_request_id: Any = None,
|
|
) -> None:
|
|
if self.fail_with is not None:
|
|
raise self.fail_with
|
|
self.calls.append(
|
|
{
|
|
"level": level,
|
|
"data": data,
|
|
"logger": logger,
|
|
"related_request_id": related_request_id,
|
|
}
|
|
)
|
|
|
|
|
|
class _FakeContext:
|
|
def __init__(self, request_id: str = "req-1") -> None:
|
|
self.session = _FakeSession()
|
|
self.request_id = request_id
|
|
|
|
|
|
class _FakeAgent:
|
|
"""Minimal stand-in for a fast-agent agent — only carries hooks."""
|
|
|
|
def __init__(self) -> None:
|
|
self.tool_runner_hooks: ToolRunnerHooks | None = None
|
|
|
|
|
|
def _tool_call(name: str, arguments: dict | None = None) -> CallToolRequest:
|
|
return CallToolRequest(
|
|
method="tools/call",
|
|
params=CallToolRequestParams(name=name, arguments=arguments or {}),
|
|
)
|
|
|
|
|
|
def _tool_result(
|
|
text: str | None = "ok", *, is_error: bool = False
|
|
) -> CallToolResult:
|
|
content: list = []
|
|
if text is not None:
|
|
content.append(TextContent(type="text", text=text))
|
|
return CallToolResult(content=content, isError=is_error)
|
|
|
|
|
|
# ── Tests ────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_emit_text_only_iteration() -> None:
|
|
"""A pure-text assistant turn produces one log message with one text block."""
|
|
ctx = _FakeContext(request_id="req-text")
|
|
emitter = AssistantChunkEmitter(ctx, agent_name="alan", conversation_id="conv-1")
|
|
|
|
msg = PromptMessageExtended(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Fair. You can't size value...")],
|
|
stop_reason=LlmStopReason.END_TURN,
|
|
)
|
|
|
|
_run(emitter._emit_assistant_chunk(msg))
|
|
|
|
assert len(ctx.session.calls) == 1
|
|
call = ctx.session.calls[0]
|
|
assert call["level"] == "info"
|
|
assert call["logger"] == LOGGER_NAME
|
|
assert call["related_request_id"] == "req-text"
|
|
|
|
data = call["data"]
|
|
assert data["kind"] == KIND
|
|
assert data["schema_version"] == SCHEMA_VERSION
|
|
assert data["agent"] == "alan"
|
|
assert data["conversation_id"] == "conv-1"
|
|
assert data["iteration"] == 1
|
|
assert data["stop_reason"] == "endTurn"
|
|
assert data["content"] == [{"type": "text", "text": "Fair. You can't size value..."}]
|
|
assert data["tool_calls"] == []
|
|
|
|
|
|
def test_emit_text_with_tool_call_iteration_carries_args_preview() -> None:
|
|
"""An iteration that requests a tool ships name, server prefix, and args preview."""
|
|
ctx = _FakeContext()
|
|
emitter = AssistantChunkEmitter(ctx, agent_name="alan", conversation_id="conv-1")
|
|
|
|
msg = PromptMessageExtended(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Logging this…")],
|
|
tool_calls={
|
|
"toolu_1": _tool_call(
|
|
"argos-search_web",
|
|
{"query": "ducati v2 vs v4 reliability"},
|
|
),
|
|
},
|
|
stop_reason=LlmStopReason.TOOL_USE,
|
|
)
|
|
|
|
_run(emitter._emit_assistant_chunk(msg))
|
|
|
|
assert len(ctx.session.calls) == 1
|
|
data = ctx.session.calls[0]["data"]
|
|
assert data["stop_reason"] == "toolUse"
|
|
assert data["content"] == [{"type": "text", "text": "Logging this…"}]
|
|
assert data["tool_calls"] == [
|
|
{
|
|
"id": "toolu_1",
|
|
"name": "argos-search_web",
|
|
"server": "argos",
|
|
"arguments_preview": "ducati v2 vs v4 reliability",
|
|
}
|
|
]
|
|
|
|
|
|
def test_emit_skips_completely_empty_iteration() -> None:
|
|
"""A turn with no text blocks and no tool calls emits nothing."""
|
|
ctx = _FakeContext()
|
|
emitter = AssistantChunkEmitter(ctx, agent_name="alan", conversation_id="conv-1")
|
|
|
|
msg = PromptMessageExtended(role="assistant", content=[])
|
|
|
|
_run(emitter._emit_assistant_chunk(msg))
|
|
|
|
assert ctx.session.calls == []
|
|
# Iteration counter still bumps so subsequent chunks aren't mis-numbered.
|
|
assert emitter._iteration == 1
|
|
|
|
|
|
def test_emit_image_block_passes_through_with_mime_type_renamed() -> None:
|
|
"""ImageContent blocks are serialized with ``mime_type`` (snake-case)."""
|
|
ctx = _FakeContext()
|
|
emitter = AssistantChunkEmitter(ctx, agent_name="alan", conversation_id=None)
|
|
|
|
msg = PromptMessageExtended(
|
|
role="assistant",
|
|
content=[ImageContent(type="image", data="ZmFrZQ==", mimeType="image/png")],
|
|
stop_reason=LlmStopReason.END_TURN,
|
|
)
|
|
|
|
_run(emitter._emit_assistant_chunk(msg))
|
|
|
|
data = ctx.session.calls[0]["data"]
|
|
assert data["content"] == [
|
|
{"type": "image", "data": "ZmFrZQ==", "mime_type": "image/png"}
|
|
]
|
|
assert data["conversation_id"] is None
|
|
|
|
|
|
def test_emit_iterations_are_numbered_in_order() -> None:
|
|
"""A multi-iteration loop produces sequentially numbered chunks."""
|
|
ctx = _FakeContext()
|
|
emitter = AssistantChunkEmitter(ctx, agent_name="alan", conversation_id="conv-1")
|
|
|
|
async def drive() -> None:
|
|
await emitter._emit_assistant_chunk(
|
|
PromptMessageExtended(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="first")],
|
|
stop_reason=LlmStopReason.TOOL_USE,
|
|
)
|
|
)
|
|
await emitter._emit_assistant_chunk(
|
|
PromptMessageExtended(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="second")],
|
|
stop_reason=LlmStopReason.TOOL_USE,
|
|
)
|
|
)
|
|
await emitter._emit_assistant_chunk(
|
|
PromptMessageExtended(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="third — done")],
|
|
stop_reason=LlmStopReason.END_TURN,
|
|
)
|
|
)
|
|
|
|
_run(drive())
|
|
|
|
assert [c["data"]["iteration"] for c in ctx.session.calls] == [1, 2, 3]
|
|
assert [c["data"]["stop_reason"] for c in ctx.session.calls] == [
|
|
"toolUse",
|
|
"toolUse",
|
|
"endTurn",
|
|
]
|
|
|
|
|
|
def test_emit_swallows_session_failure() -> None:
|
|
"""If ``send_log_message`` raises, the hook does not propagate."""
|
|
ctx = _FakeContext()
|
|
ctx.session.fail_with = RuntimeError("transport closed")
|
|
emitter = AssistantChunkEmitter(ctx, agent_name="alan", conversation_id="conv-1")
|
|
|
|
msg = PromptMessageExtended(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="hi")],
|
|
stop_reason=LlmStopReason.END_TURN,
|
|
)
|
|
|
|
_run(emitter._emit_assistant_chunk(msg))
|
|
assert ctx.session.calls == []
|
|
|
|
|
|
def test_emit_tool_results_pairs_call_id_with_iteration() -> None:
|
|
"""``after_tool_call`` ships a results payload keyed to the iteration that called."""
|
|
ctx = _FakeContext()
|
|
emitter = AssistantChunkEmitter(ctx, agent_name="alan", conversation_id="conv-1")
|
|
|
|
# Iteration 1: text + tool request.
|
|
iter1 = PromptMessageExtended(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="Searching…")],
|
|
tool_calls={"toolu_1": _tool_call("argos-search_web", {"query": "foo"})},
|
|
stop_reason=LlmStopReason.TOOL_USE,
|
|
)
|
|
|
|
async def drive() -> None:
|
|
# Simulate the runner: after_llm_call → before_tool_call → after_tool_call.
|
|
await emitter._emit_assistant_chunk(iter1)
|
|
await emitter.as_before_tool_call_hook()(None, iter1)
|
|
# The "user" message produced by the tool runner carries tool_results.
|
|
tool_msg = PromptMessageExtended(
|
|
role="user",
|
|
content=[],
|
|
tool_results={"toolu_1": _tool_result("12 results found")},
|
|
)
|
|
await emitter.as_after_tool_call_hook()(None, tool_msg)
|
|
|
|
_run(drive())
|
|
|
|
assert len(ctx.session.calls) == 2
|
|
chunk = ctx.session.calls[0]["data"]
|
|
results = ctx.session.calls[1]["data"]
|
|
assert chunk["kind"] == KIND
|
|
assert chunk["iteration"] == 1
|
|
assert results["kind"] == KIND_RESULTS
|
|
assert results["iteration"] == 1
|
|
assert results["agent"] == "alan"
|
|
assert results["conversation_id"] == "conv-1"
|
|
assert len(results["results"]) == 1
|
|
entry = results["results"][0]
|
|
assert entry["id"] == "toolu_1"
|
|
assert entry["ok"] is True
|
|
assert entry["result_preview"] == "12 results found"
|
|
assert isinstance(entry["duration_ms"], int)
|
|
assert entry["duration_ms"] >= 0
|
|
|
|
|
|
def test_emit_tool_results_marks_error() -> None:
|
|
"""A failing tool call surfaces ``ok: False`` in the results entry."""
|
|
ctx = _FakeContext()
|
|
emitter = AssistantChunkEmitter(ctx, agent_name="alan", conversation_id="c")
|
|
|
|
iter1 = PromptMessageExtended(
|
|
role="assistant",
|
|
content=[],
|
|
tool_calls={"toolu_1": _tool_call("argos-search_web", {"query": "foo"})},
|
|
stop_reason=LlmStopReason.TOOL_USE,
|
|
)
|
|
|
|
async def drive() -> None:
|
|
await emitter._emit_assistant_chunk(iter1)
|
|
await emitter.as_before_tool_call_hook()(None, iter1)
|
|
tool_msg = PromptMessageExtended(
|
|
role="user",
|
|
content=[],
|
|
tool_results={
|
|
"toolu_1": _tool_result("connection refused", is_error=True)
|
|
},
|
|
)
|
|
await emitter.as_after_tool_call_hook()(None, tool_msg)
|
|
|
|
_run(drive())
|
|
|
|
# iter1 is a pure-tool turn (no text + has tool_calls) → still emits an
|
|
# assistant_chunk because tool_calls is non-empty.
|
|
assert len(ctx.session.calls) == 2
|
|
results = ctx.session.calls[1]["data"]
|
|
assert results["results"][0]["ok"] is False
|
|
assert results["results"][0]["result_preview"] == "connection refused"
|
|
|
|
|
|
def test_emit_tool_results_empty_when_no_results() -> None:
|
|
"""An ``after_tool_call`` carrying no tool_results emits nothing."""
|
|
ctx = _FakeContext()
|
|
emitter = AssistantChunkEmitter(ctx, agent_name="alan", conversation_id="c")
|
|
|
|
msg = PromptMessageExtended(role="user", content=[], tool_results=None)
|
|
_run(emitter.as_after_tool_call_hook()(None, msg))
|
|
assert ctx.session.calls == []
|
|
|
|
|
|
def test_install_for_request_merges_with_existing_after_llm_call() -> None:
|
|
"""A pre-existing ``after_llm_call`` hook is composed, not replaced."""
|
|
ctx = _FakeContext()
|
|
agent = _FakeAgent()
|
|
|
|
seen: list[str] = []
|
|
|
|
async def base_after(_runner: Any, message: PromptMessageExtended) -> None:
|
|
seen.append("base")
|
|
|
|
agent.tool_runner_hooks = ToolRunnerHooks(after_llm_call=base_after)
|
|
|
|
restore = install_for_request(
|
|
agent, ctx=ctx, agent_name="alan", conversation_id="conv-1"
|
|
)
|
|
|
|
msg = PromptMessageExtended(
|
|
role="assistant",
|
|
content=[TextContent(type="text", text="hi")],
|
|
stop_reason=LlmStopReason.END_TURN,
|
|
)
|
|
|
|
_run(agent.tool_runner_hooks.after_llm_call(None, msg))
|
|
|
|
assert seen == ["base"]
|
|
assert len(ctx.session.calls) == 1
|
|
assert ctx.session.calls[0]["data"]["content"] == [
|
|
{"type": "text", "text": "hi"}
|
|
]
|
|
|
|
# All four hook slots used by the emitter are now bound; the restore
|
|
# call puts the original hooks back exactly.
|
|
assert agent.tool_runner_hooks.before_tool_call is not None
|
|
assert agent.tool_runner_hooks.after_tool_call is not None
|
|
|
|
restore()
|
|
assert agent.tool_runner_hooks is not None
|
|
assert agent.tool_runner_hooks.after_llm_call is base_after
|
|
assert agent.tool_runner_hooks.before_tool_call is None
|
|
assert agent.tool_runner_hooks.after_tool_call is None
|
|
|
|
|
|
def test_install_for_request_with_no_existing_hooks() -> None:
|
|
"""When the agent has no prior hooks, ours is installed cleanly."""
|
|
ctx = _FakeContext()
|
|
agent = _FakeAgent()
|
|
assert agent.tool_runner_hooks is None
|
|
|
|
restore = install_for_request(
|
|
agent, ctx=ctx, agent_name="alan", conversation_id=None
|
|
)
|
|
|
|
assert agent.tool_runner_hooks is not None
|
|
assert agent.tool_runner_hooks.after_llm_call is not None
|
|
assert agent.tool_runner_hooks.before_tool_call is not None
|
|
assert agent.tool_runner_hooks.after_tool_call is not None
|
|
|
|
restore()
|
|
assert agent.tool_runner_hooks is None
|