Initial commit
Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai> Co-Authored-By: Laure Hugo <laure.hugo@mistral.ai> Co-Authored-By: Benjamin Trom <benjamin.trom@mistral.ai> Co-Authored-By: Mathias Gesbert <mathias.gesbert@ext.mistral.ai> Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai> Co-Authored-By: Clément Drouin <clement.drouin@mistral.ai> Co-Authored-By: Vincent Guilloux <vincent.guilloux@mistral.ai> Co-Authored-By: Valentin Berard <val@mistral.ai> Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Quentin Torroba
Quentin Torroba
commit
fa15fc977b
200 changed files with 30484 additions and 0 deletions
430
tests/test_agent_observer_streaming.py
Normal file
430
tests/test_agent_observer_streaming.py
Normal file
|
|
@ -0,0 +1,430 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import cast
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.mock.utils import mock_llm_chunk
|
||||
from tests.stubs.fake_backend import FakeBackend
|
||||
from vibe.core.agent import Agent
|
||||
from vibe.core.config import SessionLoggingConfig, VibeConfig
|
||||
from vibe.core.llm.types import BackendLike
|
||||
from vibe.core.middleware import (
|
||||
ConversationContext,
|
||||
MiddlewareAction,
|
||||
MiddlewarePipeline,
|
||||
MiddlewareResult,
|
||||
ResetReason,
|
||||
)
|
||||
from vibe.core.tools.base import BaseToolConfig, ToolPermission
|
||||
from vibe.core.tools.builtins.todo import TodoArgs
|
||||
from vibe.core.types import (
|
||||
AssistantEvent,
|
||||
FunctionCall,
|
||||
LLMChunk,
|
||||
LLMMessage,
|
||||
Role,
|
||||
ToolCall,
|
||||
ToolCallEvent,
|
||||
ToolResultEvent,
|
||||
)
|
||||
from vibe.core.utils import (
|
||||
ApprovalResponse,
|
||||
CancellationReason,
|
||||
get_user_cancellation_message,
|
||||
)
|
||||
|
||||
|
||||
class InjectBeforeMiddleware:
|
||||
injectedMessage = "<injected>"
|
||||
|
||||
async def before_turn(self, context: ConversationContext) -> MiddlewareResult:
|
||||
"Inject a message just before the current step executes."
|
||||
return MiddlewareResult(
|
||||
action=MiddlewareAction.INJECT_MESSAGE, message=self.injectedMessage
|
||||
)
|
||||
|
||||
async def after_turn(self, context: ConversationContext) -> MiddlewareResult:
|
||||
return MiddlewareResult()
|
||||
|
||||
def reset(self, reset_reason: ResetReason = ResetReason.STOP) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def make_config(
|
||||
*,
|
||||
disable_logging: bool = True,
|
||||
enabled_tools: list[str] | None = None,
|
||||
tools: dict[str, BaseToolConfig] | None = None,
|
||||
) -> VibeConfig:
|
||||
cfg = VibeConfig(
|
||||
session_logging=SessionLoggingConfig(enabled=not disable_logging),
|
||||
auto_compact_threshold=0,
|
||||
system_prompt_id="tests",
|
||||
include_project_context=False,
|
||||
include_prompt_detail=False,
|
||||
include_model_info=False,
|
||||
enabled_tools=enabled_tools or [],
|
||||
tools=tools or {},
|
||||
)
|
||||
return cfg
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def observer_capture() -> tuple[
|
||||
list[tuple[Role, str | None]], Callable[[LLMMessage], None]
|
||||
]:
|
||||
observed: list[tuple[Role, str | None]] = []
|
||||
|
||||
def observer(msg: LLMMessage) -> None:
|
||||
observed.append((msg.role, msg.content))
|
||||
|
||||
return observed, observer
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_act_flushes_batched_messages_with_injection_middleware(
|
||||
observer_capture,
|
||||
) -> None:
|
||||
observed, observer = observer_capture
|
||||
|
||||
backend = FakeBackend([mock_llm_chunk(content="I can write very efficient code.")])
|
||||
agent = Agent(make_config(), message_observer=observer, backend=backend)
|
||||
agent.middleware_pipeline.add(InjectBeforeMiddleware())
|
||||
|
||||
async for _ in agent.act("How can you help?"):
|
||||
pass
|
||||
|
||||
assert len(observed) == 3
|
||||
assert [r for r, _ in observed] == [Role.system, Role.user, Role.assistant]
|
||||
assert observed[0][1] == "You are Vibe, a super useful programming assistant."
|
||||
# injected content should be appended to the user's message before emission
|
||||
assert (
|
||||
observed[1][1]
|
||||
== f"How can you help?\n\n{InjectBeforeMiddleware.injectedMessage}"
|
||||
)
|
||||
assert observed[2][1] == "I can write very efficient code."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_action_flushes_user_msg_before_returning(observer_capture) -> None:
|
||||
observed, observer = observer_capture
|
||||
|
||||
# max_turns=0 forces an immediate STOP on the first before_turn
|
||||
backend = FakeBackend([
|
||||
mock_llm_chunk(content="My response will never reach you...")
|
||||
])
|
||||
agent = Agent(
|
||||
make_config(), message_observer=observer, max_turns=0, backend=backend
|
||||
)
|
||||
|
||||
async for _ in agent.act("Greet."):
|
||||
pass
|
||||
|
||||
assert len(observed) == 2
|
||||
# user's message should have been flushed before returning
|
||||
assert [r for r, _ in observed] == [Role.system, Role.user]
|
||||
assert observed[0][1] == "You are Vibe, a super useful programming assistant."
|
||||
assert observed[1][1] == "Greet."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_act_emits_user_and_assistant_msgs(observer_capture) -> None:
|
||||
observed, observer = observer_capture
|
||||
|
||||
backend = FakeBackend([mock_llm_chunk(content="Pong!")])
|
||||
agent = Agent(make_config(), message_observer=observer, backend=backend)
|
||||
|
||||
async for _ in agent.act("Ping?"):
|
||||
pass
|
||||
|
||||
assert len(observed) == 3
|
||||
assert [r for r, _ in observed] == [Role.system, Role.user, Role.assistant]
|
||||
assert observed[1][1] == "Ping?"
|
||||
assert observed[2][1] == "Pong!"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_act_yields_assistant_event_with_usage_stats() -> None:
|
||||
backend = FakeBackend([mock_llm_chunk(content="Pong!")])
|
||||
agent = Agent(make_config(), backend=backend)
|
||||
|
||||
events = [ev async for ev in agent.act("Ping?")]
|
||||
|
||||
assert len(events) == 1
|
||||
ev = events[-1]
|
||||
assert isinstance(ev, AssistantEvent)
|
||||
assert ev.content == "Pong!"
|
||||
# stats come from tests.mock.utils.mock_llm_result (prompt=10, completion=5)
|
||||
assert ev.prompt_tokens == 10
|
||||
assert ev.completion_tokens == 5
|
||||
assert ev.session_total_tokens == 15
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_act_streams_batched_chunks_in_order() -> None:
|
||||
backend = FakeBackend([
|
||||
mock_llm_chunk(content="Hello"),
|
||||
mock_llm_chunk(content=" from"),
|
||||
mock_llm_chunk(content=" Vibe"),
|
||||
mock_llm_chunk(content="! "),
|
||||
mock_llm_chunk(content="More"),
|
||||
mock_llm_chunk(content=" and"),
|
||||
mock_llm_chunk(content=" end"),
|
||||
])
|
||||
agent = Agent(make_config(), backend=backend, enable_streaming=True)
|
||||
|
||||
events = [event async for event in agent.act("Stream, please.")]
|
||||
|
||||
assert len(events) == 2
|
||||
assert [event.content for event in events if isinstance(event, AssistantEvent)] == [
|
||||
"Hello from Vibe! More",
|
||||
" and end",
|
||||
]
|
||||
assert agent.messages[-1].role == Role.assistant
|
||||
assert agent.messages[-1].content == "Hello from Vibe! More and end"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_act_handles_streaming_with_tool_call_events_in_sequence() -> None:
|
||||
todo_tool_call = ToolCall(
|
||||
id="tc_stream",
|
||||
index=0,
|
||||
function=FunctionCall(name="todo", arguments='{"action": "read"}'),
|
||||
)
|
||||
backend = FakeBackend([
|
||||
mock_llm_chunk(content="Checking your todos."),
|
||||
mock_llm_chunk(content="", tool_calls=[todo_tool_call]),
|
||||
mock_llm_chunk(content="", finish_reason="stop"),
|
||||
mock_llm_chunk(content="Done reviewing todos."),
|
||||
])
|
||||
agent = Agent(
|
||||
make_config(
|
||||
enabled_tools=["todo"],
|
||||
tools={"todo": BaseToolConfig(permission=ToolPermission.ALWAYS)},
|
||||
),
|
||||
backend=backend,
|
||||
auto_approve=True,
|
||||
enable_streaming=True,
|
||||
)
|
||||
|
||||
events = [event async for event in agent.act("What about my todos?")]
|
||||
|
||||
assert [type(event) for event in events] == [
|
||||
AssistantEvent,
|
||||
ToolCallEvent,
|
||||
ToolResultEvent,
|
||||
AssistantEvent,
|
||||
]
|
||||
assert isinstance(events[0], AssistantEvent)
|
||||
assert events[0].content == "Checking your todos."
|
||||
assert isinstance(events[1], ToolCallEvent)
|
||||
assert events[1].tool_name == "todo"
|
||||
assert isinstance(events[2], ToolResultEvent)
|
||||
assert events[2].error is None
|
||||
assert events[2].skipped is False
|
||||
assert isinstance(events[3], AssistantEvent)
|
||||
assert events[3].content == "Done reviewing todos."
|
||||
assert agent.messages[-1].content == "Done reviewing todos."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_act_handles_tool_call_chunk_with_content() -> None:
|
||||
todo_tool_call = ToolCall(
|
||||
id="tc_content",
|
||||
index=0,
|
||||
function=FunctionCall(name="todo", arguments='{"action": "read"}'),
|
||||
)
|
||||
backend = FakeBackend([
|
||||
mock_llm_chunk(content="Preparing "),
|
||||
mock_llm_chunk(content="todo request", tool_calls=[todo_tool_call]),
|
||||
mock_llm_chunk(content=" complete", finish_reason="stop"),
|
||||
])
|
||||
agent = Agent(
|
||||
make_config(
|
||||
enabled_tools=["todo"],
|
||||
tools={"todo": BaseToolConfig(permission=ToolPermission.ALWAYS)},
|
||||
),
|
||||
backend=backend,
|
||||
auto_approve=True,
|
||||
enable_streaming=True,
|
||||
)
|
||||
|
||||
events = [event async for event in agent.act("Check todos with content.")]
|
||||
|
||||
assert [type(event) for event in events] == [
|
||||
AssistantEvent,
|
||||
AssistantEvent,
|
||||
ToolCallEvent,
|
||||
ToolResultEvent,
|
||||
]
|
||||
assert isinstance(events[0], AssistantEvent)
|
||||
assert events[0].content == "Preparing todo request"
|
||||
assert isinstance(events[1], AssistantEvent)
|
||||
assert events[1].content == " complete"
|
||||
assert any(
|
||||
m.role == Role.assistant and m.content == "Preparing todo request complete"
|
||||
for m in agent.messages
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_act_merges_streamed_tool_call_arguments() -> None:
|
||||
tool_call_part_one = ToolCall(
|
||||
id="tc_merge",
|
||||
index=0,
|
||||
function=FunctionCall(
|
||||
name="todo", arguments='{"action": "read", "note": "First '
|
||||
),
|
||||
)
|
||||
tool_call_part_two = ToolCall(
|
||||
id="tc_merge", index=0, function=FunctionCall(name="todo", arguments='part"}')
|
||||
)
|
||||
backend = FakeBackend([
|
||||
mock_llm_chunk(content="Planning: "),
|
||||
mock_llm_chunk(content="", tool_calls=[tool_call_part_one]),
|
||||
mock_llm_chunk(content="", tool_calls=[tool_call_part_two]),
|
||||
])
|
||||
agent = Agent(
|
||||
make_config(
|
||||
enabled_tools=["todo"],
|
||||
tools={"todo": BaseToolConfig(permission=ToolPermission.ALWAYS)},
|
||||
),
|
||||
backend=backend,
|
||||
auto_approve=True,
|
||||
enable_streaming=True,
|
||||
)
|
||||
|
||||
events = [event async for event in agent.act("Merge streamed tool call args.")]
|
||||
|
||||
assert [type(event) for event in events] == [
|
||||
AssistantEvent,
|
||||
ToolCallEvent,
|
||||
ToolResultEvent,
|
||||
]
|
||||
call_event = events[1]
|
||||
assert isinstance(call_event, ToolCallEvent)
|
||||
assert call_event.tool_call_id == "tc_merge"
|
||||
call_args = cast(TodoArgs, call_event.args)
|
||||
assert call_args.action == "read"
|
||||
assert isinstance(events[2], ToolResultEvent)
|
||||
assert events[2].error is None
|
||||
assert events[2].skipped is False
|
||||
assistant_with_calls = next(
|
||||
m for m in agent.messages if m.role == Role.assistant and m.tool_calls
|
||||
)
|
||||
reconstructed_calls = assistant_with_calls.tool_calls or []
|
||||
assert len(reconstructed_calls) == 1
|
||||
assert reconstructed_calls[0].function.arguments == (
|
||||
'{"action": "read", "note": "First part"}'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_act_raises_when_stream_never_signals_finish() -> None:
|
||||
class IncompleteStreamingBackend(BackendLike):
|
||||
def __init__(self, chunks: list[LLMChunk]) -> None:
|
||||
self._chunks = list(chunks)
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return None
|
||||
|
||||
async def complete_streaming(self, **_: object):
|
||||
while self._chunks:
|
||||
yield self._chunks.pop(0)
|
||||
|
||||
async def complete(self, **_: object):
|
||||
return mock_llm_chunk(content="", finish_reason="stop")
|
||||
|
||||
async def count_tokens(self, **_: object) -> int:
|
||||
return 0
|
||||
|
||||
backend = IncompleteStreamingBackend([mock_llm_chunk(content="partial")])
|
||||
agent = Agent(make_config(), backend=backend, enable_streaming=True)
|
||||
|
||||
with pytest.raises(RuntimeError, match="Streamed completion returned no chunks"):
|
||||
[event async for event in agent.act("Will this finish?")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_act_handles_user_cancellation_during_streaming() -> None:
|
||||
class CountingMiddleware(MiddlewarePipeline):
|
||||
def __init__(self) -> None:
|
||||
self.before_calls = 0
|
||||
self.after_calls = 0
|
||||
|
||||
async def before_turn(self, context: ConversationContext) -> MiddlewareResult:
|
||||
self.before_calls += 1
|
||||
return MiddlewareResult()
|
||||
|
||||
async def after_turn(self, context: ConversationContext) -> MiddlewareResult:
|
||||
self.after_calls += 1
|
||||
return MiddlewareResult()
|
||||
|
||||
def reset(self, reset_reason: ResetReason = ResetReason.STOP) -> None:
|
||||
return None
|
||||
|
||||
todo_tool_call = ToolCall(
|
||||
id="tc_cancel",
|
||||
index=0,
|
||||
function=FunctionCall(name="todo", arguments='{"action": "read"}'),
|
||||
)
|
||||
backend = FakeBackend([
|
||||
mock_llm_chunk(content="Preparing "),
|
||||
mock_llm_chunk(content="todo request", tool_calls=[todo_tool_call]),
|
||||
mock_llm_chunk(content="", finish_reason="stop"),
|
||||
])
|
||||
agent = Agent(
|
||||
make_config(
|
||||
enabled_tools=["todo"],
|
||||
tools={"todo": BaseToolConfig(permission=ToolPermission.ASK)},
|
||||
),
|
||||
backend=backend,
|
||||
auto_approve=False,
|
||||
enable_streaming=True,
|
||||
)
|
||||
middleware = CountingMiddleware()
|
||||
agent.middleware_pipeline.add(middleware)
|
||||
agent.set_approval_callback(
|
||||
lambda _name, _args, _id: (
|
||||
ApprovalResponse.NO,
|
||||
str(get_user_cancellation_message(CancellationReason.OPERATION_CANCELLED)),
|
||||
)
|
||||
)
|
||||
agent.interaction_logger.save_interaction = AsyncMock(return_value=None)
|
||||
|
||||
events = [event async for event in agent.act("Cancel mid stream?")]
|
||||
|
||||
assert [type(event) for event in events] == [
|
||||
AssistantEvent,
|
||||
ToolCallEvent,
|
||||
ToolResultEvent,
|
||||
]
|
||||
assert middleware.before_calls == 1
|
||||
assert middleware.after_calls == 0
|
||||
assert isinstance(events[-1], ToolResultEvent)
|
||||
assert events[-1].skipped is True
|
||||
assert events[-1].skip_reason is not None
|
||||
assert "<user_cancellation>" in events[-1].skip_reason
|
||||
assert agent.interaction_logger.save_interaction.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_act_flushes_and_logs_when_streaming_errors(observer_capture) -> None:
|
||||
observed, observer = observer_capture
|
||||
backend = FakeBackend(exception_to_raise=RuntimeError("boom in streaming"))
|
||||
agent = Agent(
|
||||
make_config(), backend=backend, message_observer=observer, enable_streaming=True
|
||||
)
|
||||
agent.interaction_logger.save_interaction = AsyncMock(return_value=None)
|
||||
|
||||
with pytest.raises(RuntimeError, match="boom in streaming"):
|
||||
[_ async for _ in agent.act("Trigger stream failure")]
|
||||
|
||||
assert [role for role, _ in observed] == [Role.system, Role.user]
|
||||
assert agent.interaction_logger.save_interaction.await_count == 1
|
||||
Loading…
Add table
Add a link
Reference in a new issue