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
477
tests/test_agent_tool_call.py
Normal file
477
tests/test_agent_tool_call.py
Normal file
|
|
@ -0,0 +1,477 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.mock.utils import mock_llm_chunk
|
||||
from tests.stubs.fake_backend import FakeBackend
|
||||
from tests.stubs.fake_tool import FakeTool
|
||||
from vibe.core.agent import Agent
|
||||
from vibe.core.config import SessionLoggingConfig, VibeConfig
|
||||
from vibe.core.tools.base import BaseToolConfig, ToolPermission
|
||||
from vibe.core.tools.builtins.todo import TodoItem
|
||||
from vibe.core.types import (
|
||||
AssistantEvent,
|
||||
BaseEvent,
|
||||
FunctionCall,
|
||||
LLMMessage,
|
||||
Role,
|
||||
SyncApprovalCallback,
|
||||
ToolCall,
|
||||
ToolCallEvent,
|
||||
ToolResultEvent,
|
||||
)
|
||||
from vibe.core.utils import ApprovalResponse
|
||||
|
||||
|
||||
async def act_and_collect_events(agent: Agent, prompt: str) -> list[BaseEvent]:
|
||||
return [ev async for ev in agent.act(prompt)]
|
||||
|
||||
|
||||
def make_config(todo_permission: ToolPermission = ToolPermission.ALWAYS) -> VibeConfig:
|
||||
return VibeConfig(
|
||||
session_logging=SessionLoggingConfig(enabled=False),
|
||||
auto_compact_threshold=0,
|
||||
enabled_tools=["todo"],
|
||||
tools={"todo": BaseToolConfig(permission=todo_permission)},
|
||||
system_prompt_id="tests",
|
||||
include_project_context=False,
|
||||
include_prompt_detail=False,
|
||||
)
|
||||
|
||||
|
||||
def make_todo_tool_call(call_id: str, action: str = "read") -> ToolCall:
|
||||
return ToolCall(
|
||||
id=call_id,
|
||||
function=FunctionCall(name="todo", arguments=f'{{"action": "{action}"}}'),
|
||||
)
|
||||
|
||||
|
||||
def make_agent(
|
||||
*,
|
||||
auto_approve: bool = True,
|
||||
todo_permission: ToolPermission = ToolPermission.ALWAYS,
|
||||
backend: FakeBackend,
|
||||
approval_callback: SyncApprovalCallback | None = None,
|
||||
) -> Agent:
|
||||
agent = Agent(
|
||||
make_config(todo_permission=todo_permission),
|
||||
auto_approve=auto_approve,
|
||||
backend=backend,
|
||||
)
|
||||
if approval_callback:
|
||||
agent.set_approval_callback(approval_callback)
|
||||
return agent
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_tool_call_executes_under_auto_approve() -> None:
|
||||
mocked_tool_call_id = "call_1"
|
||||
tool_call = make_todo_tool_call(mocked_tool_call_id)
|
||||
backend = FakeBackend([
|
||||
mock_llm_chunk(content="Let me check your todos.", tool_calls=[tool_call]),
|
||||
mock_llm_chunk(content="I retrieved 0 todos.", finish_reason="stop"),
|
||||
])
|
||||
agent = make_agent(auto_approve=True, backend=backend)
|
||||
|
||||
events = await act_and_collect_events(agent, "What's my todo list?")
|
||||
|
||||
assert [type(e) for e in events] == [
|
||||
AssistantEvent,
|
||||
ToolCallEvent,
|
||||
ToolResultEvent,
|
||||
AssistantEvent,
|
||||
]
|
||||
assert isinstance(events[0], AssistantEvent)
|
||||
assert events[0].content == "Let me check 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 events[2].result is not None
|
||||
assert isinstance(events[3], AssistantEvent)
|
||||
assert events[3].content == "I retrieved 0 todos."
|
||||
# check conversation history
|
||||
tool_msgs = [m for m in agent.messages if m.role == Role.tool]
|
||||
assert len(tool_msgs) == 1
|
||||
assert tool_msgs[-1].tool_call_id == mocked_tool_call_id
|
||||
assert "total_count" in (tool_msgs[-1].content or "")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_requires_approval_if_not_auto_approved() -> None:
|
||||
agent = make_agent(
|
||||
auto_approve=False,
|
||||
todo_permission=ToolPermission.ASK,
|
||||
backend=FakeBackend([
|
||||
mock_llm_chunk(
|
||||
content="Let me check your todos.",
|
||||
tool_calls=[make_todo_tool_call("call_2")],
|
||||
),
|
||||
mock_llm_chunk(
|
||||
content="I cannot execute the tool without approval.",
|
||||
finish_reason="stop",
|
||||
),
|
||||
]),
|
||||
)
|
||||
|
||||
events = await act_and_collect_events(agent, "What's my todo list?")
|
||||
|
||||
assert isinstance(events[1], ToolCallEvent)
|
||||
assert events[1].tool_name == "todo"
|
||||
assert isinstance(events[2], ToolResultEvent)
|
||||
assert events[2].skipped is True
|
||||
assert events[2].error is None
|
||||
assert events[2].result is None
|
||||
assert events[2].skip_reason is not None
|
||||
assert "not permitted" in events[2].skip_reason.lower()
|
||||
assert isinstance(events[3], AssistantEvent)
|
||||
assert events[3].content == "I cannot execute the tool without approval."
|
||||
assert agent.stats.tool_calls_rejected == 1
|
||||
assert agent.stats.tool_calls_agreed == 0
|
||||
assert agent.stats.tool_calls_succeeded == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_approved_by_callback() -> None:
|
||||
def approval_callback(
|
||||
_tool_name: str, _args: dict[str, Any], _tool_call_id: str
|
||||
) -> tuple[str, str | None]:
|
||||
return (ApprovalResponse.YES, None)
|
||||
|
||||
agent = make_agent(
|
||||
auto_approve=False,
|
||||
todo_permission=ToolPermission.ASK,
|
||||
approval_callback=approval_callback,
|
||||
backend=FakeBackend([
|
||||
mock_llm_chunk(
|
||||
content="Let me check your todos.",
|
||||
tool_calls=[make_todo_tool_call("call_3")],
|
||||
),
|
||||
mock_llm_chunk(content="I retrieved 0 todos.", finish_reason="stop"),
|
||||
]),
|
||||
)
|
||||
|
||||
events = await act_and_collect_events(agent, "What's my todo list?")
|
||||
|
||||
assert isinstance(events[2], ToolResultEvent)
|
||||
assert events[2].skipped is False
|
||||
assert events[2].error is None
|
||||
assert events[2].result is not None
|
||||
assert agent.stats.tool_calls_agreed == 1
|
||||
assert agent.stats.tool_calls_rejected == 0
|
||||
assert agent.stats.tool_calls_succeeded == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_rejected_when_auto_approve_disabled_and_rejected_by_callback() -> (
|
||||
None
|
||||
):
|
||||
custom_feedback = "User declined tool execution"
|
||||
|
||||
def approval_callback(
|
||||
_tool_name: str, _args: dict[str, Any], _tool_call_id: str
|
||||
) -> tuple[str, str | None]:
|
||||
return (ApprovalResponse.NO, custom_feedback)
|
||||
|
||||
agent = make_agent(
|
||||
auto_approve=False,
|
||||
todo_permission=ToolPermission.ASK,
|
||||
approval_callback=approval_callback,
|
||||
backend=FakeBackend([
|
||||
mock_llm_chunk(
|
||||
content="Let me check your todos.",
|
||||
tool_calls=[make_todo_tool_call("call_4")],
|
||||
),
|
||||
mock_llm_chunk(
|
||||
content="Understood, I won't check the todos.", finish_reason="stop"
|
||||
),
|
||||
]),
|
||||
)
|
||||
|
||||
events = await act_and_collect_events(agent, "What's my todo list?")
|
||||
|
||||
assert isinstance(events[2], ToolResultEvent)
|
||||
assert events[2].skipped is True
|
||||
assert events[2].error is None
|
||||
assert events[2].result is None
|
||||
assert events[2].skip_reason == custom_feedback
|
||||
assert agent.stats.tool_calls_rejected == 1
|
||||
assert agent.stats.tool_calls_agreed == 0
|
||||
assert agent.stats.tool_calls_succeeded == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_skipped_when_permission_is_never() -> None:
|
||||
agent = make_agent(
|
||||
auto_approve=False,
|
||||
todo_permission=ToolPermission.NEVER,
|
||||
backend=FakeBackend([
|
||||
mock_llm_chunk(
|
||||
content="Let me check your todos.",
|
||||
tool_calls=[make_todo_tool_call("call_never")],
|
||||
),
|
||||
mock_llm_chunk(content="Tool is disabled.", finish_reason="stop"),
|
||||
]),
|
||||
)
|
||||
|
||||
events = await act_and_collect_events(agent, "What's my todo list?")
|
||||
|
||||
assert isinstance(events[2], ToolResultEvent)
|
||||
assert events[2].skipped is True
|
||||
assert events[2].error is None
|
||||
assert events[2].result is None
|
||||
assert events[2].skip_reason is not None
|
||||
assert "permanently disabled" in events[2].skip_reason.lower()
|
||||
tool_msgs = [m for m in agent.messages if m.role == Role.tool and m.name == "todo"]
|
||||
assert len(tool_msgs) == 1
|
||||
assert tool_msgs[0].name == "todo"
|
||||
assert events[2].skip_reason in (tool_msgs[-1].content or "")
|
||||
assert agent.stats.tool_calls_rejected == 1
|
||||
assert agent.stats.tool_calls_agreed == 0
|
||||
assert agent.stats.tool_calls_succeeded == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_approval_always_flips_auto_approve_for_subsequent_calls() -> None:
|
||||
callback_invocations = []
|
||||
|
||||
def approval_callback(
|
||||
tool_name: str, _args: dict[str, Any], _tool_call_id: str
|
||||
) -> tuple[str, str | None]:
|
||||
callback_invocations.append(tool_name)
|
||||
return (ApprovalResponse.ALWAYS, None)
|
||||
|
||||
agent = make_agent(
|
||||
auto_approve=False,
|
||||
todo_permission=ToolPermission.ASK,
|
||||
approval_callback=approval_callback,
|
||||
backend=FakeBackend([
|
||||
mock_llm_chunk(
|
||||
content="First check.", tool_calls=[make_todo_tool_call("call_first")]
|
||||
),
|
||||
mock_llm_chunk(content="First done.", finish_reason="stop"),
|
||||
mock_llm_chunk(
|
||||
content="Second check.", tool_calls=[make_todo_tool_call("call_second")]
|
||||
),
|
||||
mock_llm_chunk(content="Second done.", finish_reason="stop"),
|
||||
]),
|
||||
)
|
||||
|
||||
events1 = await act_and_collect_events(agent, "First request")
|
||||
events2 = await act_and_collect_events(agent, "Second request")
|
||||
|
||||
assert agent.auto_approve is True
|
||||
assert len(callback_invocations) == 1
|
||||
assert callback_invocations[0] == "todo"
|
||||
assert isinstance(events1[2], ToolResultEvent)
|
||||
assert events1[2].skipped is False
|
||||
assert events1[2].result is not None
|
||||
assert isinstance(events2[2], ToolResultEvent)
|
||||
assert events2[2].skipped is False
|
||||
assert events2[2].result is not None
|
||||
assert agent.stats.tool_calls_rejected == 0
|
||||
assert agent.stats.tool_calls_succeeded == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_with_invalid_action() -> None:
|
||||
tool_call = ToolCall(
|
||||
id="call_5",
|
||||
function=FunctionCall(name="todo", arguments='{"action": "invalid_action"}'),
|
||||
)
|
||||
agent = make_agent(
|
||||
auto_approve=True,
|
||||
backend=FakeBackend([
|
||||
mock_llm_chunk(content="Let me check your todos.", tool_calls=[tool_call]),
|
||||
mock_llm_chunk(
|
||||
content="I encountered an error with the action.", finish_reason="stop"
|
||||
),
|
||||
]),
|
||||
)
|
||||
|
||||
events = await act_and_collect_events(agent, "What's my todo list?")
|
||||
|
||||
assert isinstance(events[2], ToolResultEvent)
|
||||
assert events[2].error is not None
|
||||
assert events[2].result is None
|
||||
assert "tool_error" in events[2].error.lower()
|
||||
assert agent.stats.tool_calls_failed == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_with_duplicate_todo_ids() -> None:
|
||||
duplicate_todos = [
|
||||
TodoItem(id="duplicate", content="Task 1"),
|
||||
TodoItem(id="duplicate", content="Task 2"),
|
||||
]
|
||||
tool_call = ToolCall(
|
||||
id="call_6",
|
||||
function=FunctionCall(
|
||||
name="todo",
|
||||
arguments=json.dumps({
|
||||
"action": "write",
|
||||
"todos": [t.model_dump() for t in duplicate_todos],
|
||||
}),
|
||||
),
|
||||
)
|
||||
agent = make_agent(
|
||||
auto_approve=True,
|
||||
backend=FakeBackend([
|
||||
mock_llm_chunk(content="Let me write todos.", tool_calls=[tool_call]),
|
||||
mock_llm_chunk(
|
||||
content="I couldn't write todos with duplicate IDs.",
|
||||
finish_reason="stop",
|
||||
),
|
||||
]),
|
||||
)
|
||||
|
||||
events = await act_and_collect_events(agent, "Add todos")
|
||||
|
||||
assert isinstance(events[2], ToolResultEvent)
|
||||
assert events[2].error is not None
|
||||
assert events[2].result is None
|
||||
assert "unique" in events[2].error.lower()
|
||||
assert agent.stats.tool_calls_failed == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_with_exceeding_max_todos() -> None:
|
||||
many_todos = [TodoItem(id=f"todo_{i}", content=f"Task {i}") for i in range(150)]
|
||||
tool_call = ToolCall(
|
||||
id="call_7",
|
||||
function=FunctionCall(
|
||||
name="todo",
|
||||
arguments=json.dumps({
|
||||
"action": "write",
|
||||
"todos": [t.model_dump() for t in many_todos],
|
||||
}),
|
||||
),
|
||||
)
|
||||
agent = make_agent(
|
||||
auto_approve=True,
|
||||
backend=FakeBackend([
|
||||
mock_llm_chunk(content="Let me write todos.", tool_calls=[tool_call]),
|
||||
mock_llm_chunk(
|
||||
content="I couldn't write that many todos.", finish_reason="stop"
|
||||
),
|
||||
]),
|
||||
)
|
||||
|
||||
events = await act_and_collect_events(agent, "Add todos")
|
||||
|
||||
assert isinstance(events[2], ToolResultEvent)
|
||||
assert events[2].error is not None
|
||||
assert events[2].result is None
|
||||
assert "100" in events[2].error
|
||||
assert agent.stats.tool_calls_failed == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"exception_class",
|
||||
[
|
||||
pytest.param(KeyboardInterrupt, id="keyboard_interrupt"),
|
||||
pytest.param(asyncio.CancelledError, id="asyncio_cancelled"),
|
||||
],
|
||||
)
|
||||
async def test_tool_call_can_be_interrupted(
|
||||
exception_class: type[BaseException],
|
||||
) -> None:
|
||||
tool_call = ToolCall(
|
||||
id="call_8", function=FunctionCall(name="stub_tool", arguments="{}")
|
||||
)
|
||||
config = VibeConfig(
|
||||
session_logging=SessionLoggingConfig(enabled=False),
|
||||
auto_compact_threshold=0,
|
||||
enabled_tools=["stub_tool"],
|
||||
)
|
||||
agent = Agent(
|
||||
config,
|
||||
auto_approve=True,
|
||||
backend=FakeBackend([
|
||||
mock_llm_chunk(content="Let me use the tool.", tool_calls=[tool_call]),
|
||||
mock_llm_chunk(content="Tool execution completed.", finish_reason="stop"),
|
||||
]),
|
||||
)
|
||||
# no dependency injection available => monkey patch
|
||||
agent.tool_manager._available["stub_tool"] = FakeTool
|
||||
stub_tool_instance = agent.tool_manager.get("stub_tool")
|
||||
assert isinstance(stub_tool_instance, FakeTool)
|
||||
stub_tool_instance._exception_to_raise = exception_class()
|
||||
|
||||
events: list[BaseEvent] = []
|
||||
with pytest.raises(exception_class):
|
||||
async for ev in agent.act("Execute tool"):
|
||||
events.append(ev)
|
||||
|
||||
tool_result_event = next(
|
||||
(e for e in events if isinstance(e, ToolResultEvent)), None
|
||||
)
|
||||
assert tool_result_event is not None
|
||||
assert tool_result_event.error is not None
|
||||
assert "execution interrupted by user" in tool_result_event.error.lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fill_missing_tool_responses_inserts_placeholders() -> None:
|
||||
agent = Agent(
|
||||
make_config(),
|
||||
auto_approve=True,
|
||||
backend=FakeBackend([mock_llm_chunk(content="ok", finish_reason="stop")]),
|
||||
)
|
||||
tool_calls_messages = [
|
||||
ToolCall(
|
||||
id="tc1", function=FunctionCall(name="todo", arguments='{"action": "read"}')
|
||||
),
|
||||
ToolCall(
|
||||
id="tc2", function=FunctionCall(name="todo", arguments='{"action": "read"}')
|
||||
),
|
||||
]
|
||||
assistant_msg = LLMMessage(
|
||||
role=Role.assistant, content="Calling tools...", tool_calls=tool_calls_messages
|
||||
)
|
||||
agent.messages = [
|
||||
agent.messages[0],
|
||||
assistant_msg,
|
||||
# only one tool responded: the second is missing
|
||||
LLMMessage(
|
||||
role=Role.tool, tool_call_id="tc1", name="todo", content="Retrieved 0 todos"
|
||||
),
|
||||
]
|
||||
|
||||
await act_and_collect_events(agent, "Proceed")
|
||||
|
||||
tool_msgs = [m for m in agent.messages if m.role == Role.tool]
|
||||
assert any(m.tool_call_id == "tc2" for m in tool_msgs)
|
||||
# find placeholder message for tc2
|
||||
placeholder = next(m for m in tool_msgs if m.tool_call_id == "tc2")
|
||||
assert placeholder.name == "todo"
|
||||
assert (
|
||||
placeholder.content
|
||||
== "<user_cancellation>Tool execution interrupted - no response available</user_cancellation>"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_assistant_after_tool_appends_understood() -> None:
|
||||
agent = Agent(
|
||||
make_config(),
|
||||
auto_approve=True,
|
||||
backend=FakeBackend([mock_llm_chunk(content="ok", finish_reason="stop")]),
|
||||
)
|
||||
tool_msg = LLMMessage(
|
||||
role=Role.tool, tool_call_id="tc_z", name="todo", content="Done"
|
||||
)
|
||||
agent.messages = [agent.messages[0], tool_msg]
|
||||
|
||||
await act_and_collect_events(agent, "Next")
|
||||
|
||||
# find the seeded tool message and ensure the next message is "Understood."
|
||||
idx = next(i for i, m in enumerate(agent.messages) if m.role == Role.tool)
|
||||
assert agent.messages[idx + 1].role == Role.assistant
|
||||
assert agent.messages[idx + 1].content == "Understood."
|
||||
Loading…
Add table
Add a link
Reference in a new issue