Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai> Co-authored-by: Guillaume LE GOFF <guillaume.lgf@gmail.com> Co-authored-by: Hdandria <henri.dandria@mistral.ai> Co-authored-by: Ivana Dunisijevic <ivana.dunisijevic@mistral.ai> Co-authored-by: Jean Burellier <sheplu@users.noreply.github.com> Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: Mert Unsal <mert.unsal@mistral.ai> Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com> Co-authored-by: Paul VEZIA <166131032+le-codeur-rapide@users.noreply.github.com> Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai> Co-authored-by: Val <102326092+vdeva@users.noreply.github.com> Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com> Co-authored-by: renovate-mistral[bot] <253709520+renovate-mistral[bot]@users.noreply.github.com> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
111 lines
3.6 KiB
Python
111 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from tests.agent_loop.e2e.conftest import (
|
|
MistralAPI,
|
|
assistant_text,
|
|
build_e2e_agent_loop,
|
|
e2e_config,
|
|
)
|
|
from tests.backend.data.mistral import (
|
|
STREAMED_SIMPLE_CONVERSATION_PARAMS,
|
|
mistral_completion,
|
|
)
|
|
from vibe.core.types import (
|
|
AssistantEvent,
|
|
ToolCallEvent,
|
|
ToolResultEvent,
|
|
UserMessageEvent,
|
|
)
|
|
|
|
TODO_TOOL_CALL = [
|
|
{
|
|
"id": "call_todo_1",
|
|
"function": {"name": "todo", "arguments": '{"action": "read"}'},
|
|
"index": 0,
|
|
}
|
|
]
|
|
|
|
|
|
def _chunked_completion(reasoning: str, answer: str) -> dict[str, Any]:
|
|
completion = mistral_completion(answer)
|
|
completion["choices"][0]["message"]["content"] = [
|
|
{"type": "thinking", "thinking": [{"type": "text", "text": reasoning}]},
|
|
{"type": "text", "text": answer},
|
|
]
|
|
return completion
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_act_completes_against_real_backend(mistral_api: MistralAPI) -> None:
|
|
# A plain prompt yields a user event then the backend's assistant reply, and
|
|
# the loop accumulates prompt+completion tokens from the usage block.
|
|
mistral_api.reply(
|
|
mistral_completion("Hi there!", prompt_tokens=120, completion_tokens=30)
|
|
)
|
|
agent = build_e2e_agent_loop()
|
|
|
|
events = [event async for event in agent.act("Hello")]
|
|
|
|
assert [type(e) for e in events] == [UserMessageEvent, AssistantEvent]
|
|
assert isinstance(events[1], AssistantEvent)
|
|
assert events[1].content == "Hi there!"
|
|
assert agent.stats.context_tokens == 150
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_act_streaming(mistral_api: MistralAPI) -> None:
|
|
# Streamed SSE chunks are reassembled into the final assistant content.
|
|
_, chunks, _ = STREAMED_SIMPLE_CONVERSATION_PARAMS[0]
|
|
mistral_api.reply_stream(chunks)
|
|
agent = build_e2e_agent_loop(enable_streaming=True)
|
|
|
|
events = [event async for event in agent.act("Hi")]
|
|
|
|
assert assistant_text(events).endswith("Some content")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_act_tool_call_round_trip(mistral_api: MistralAPI) -> None:
|
|
# A tool call is parsed, executed, and its result fed back for a final reply.
|
|
mistral_api.reply(
|
|
mistral_completion("", tool_calls=TODO_TOOL_CALL),
|
|
mistral_completion("All done"),
|
|
)
|
|
agent = build_e2e_agent_loop(config=e2e_config(enabled_tools=["todo"]))
|
|
|
|
events = [event async for event in agent.act("Show my todos")]
|
|
|
|
assert any(isinstance(e, ToolCallEvent) for e in events)
|
|
assert any(isinstance(e, ToolResultEvent) for e in events)
|
|
final = [e for e in events if isinstance(e, AssistantEvent)]
|
|
assert final and final[-1].content == "All done"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_act_serializes_tools_in_request_payload(mistral_api: MistralAPI) -> None:
|
|
# Enabled tools are serialized into the outgoing chat-completions request.
|
|
mistral_api.reply(mistral_completion("ok"))
|
|
agent = build_e2e_agent_loop(config=e2e_config(enabled_tools=["todo"]))
|
|
|
|
_ = [event async for event in agent.act("Hello")]
|
|
|
|
names = {t["function"]["name"] for t in mistral_api.request_json.get("tools", [])}
|
|
assert "todo" in names
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_act_extracts_text_from_chunked_content_array(
|
|
mistral_api: MistralAPI,
|
|
) -> None:
|
|
# A thinking+text content array has its text part extracted as the reply.
|
|
mistral_api.reply(_chunked_completion("thinking hard", "the answer"))
|
|
agent = build_e2e_agent_loop()
|
|
|
|
events = [event async for event in agent.act("Why?")]
|
|
|
|
assistant = [e for e in events if isinstance(e, AssistantEvent)]
|
|
assert assistant and assistant[-1].content == "the answer"
|