v1.3.0
Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai> Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai> Co-Authored-By: Clément Drouin <clement.drouin@mistral.ai> Co-Authored-by: Thiago Padilha <thiago@coplane.com>
This commit is contained in:
Mathias Gesbert
Mathias Gesbert
parent
2e1e15120d
commit
078693fc64
67 changed files with 3959 additions and 819 deletions
360
tests/test_reasoning_content.py
Normal file
360
tests/test_reasoning_content.py
Normal file
|
|
@ -0,0 +1,360 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import mistralai
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
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 (
|
||||
ModelConfig,
|
||||
ProviderConfig,
|
||||
SessionLoggingConfig,
|
||||
VibeConfig,
|
||||
)
|
||||
from vibe.core.llm.backend.generic import GenericBackend
|
||||
from vibe.core.llm.backend.mistral import MistralMapper, ParsedContent
|
||||
from vibe.core.llm.format import APIToolFormatHandler
|
||||
from vibe.core.types import AssistantEvent, LLMMessage, ReasoningEvent, Role
|
||||
|
||||
|
||||
def make_config() -> VibeConfig:
|
||||
return VibeConfig(
|
||||
session_logging=SessionLoggingConfig(enabled=False),
|
||||
auto_compact_threshold=0,
|
||||
system_prompt_id="tests",
|
||||
include_project_context=False,
|
||||
include_prompt_detail=False,
|
||||
include_model_info=False,
|
||||
include_commit_signature=False,
|
||||
enabled_tools=[],
|
||||
tools={},
|
||||
)
|
||||
|
||||
|
||||
class TestMistralMapperParseContent:
|
||||
def test_parse_content_string_returns_content_only(self):
|
||||
mapper = MistralMapper()
|
||||
result = mapper.parse_content("Hello, world!")
|
||||
|
||||
assert result == ParsedContent(content="Hello, world!", reasoning_content=None)
|
||||
|
||||
def test_parse_content_text_chunk_returns_content_only(self):
|
||||
mapper = MistralMapper()
|
||||
content: list[mistralai.ContentChunk] = [
|
||||
mistralai.TextChunk(type="text", text="Hello from text chunk")
|
||||
]
|
||||
|
||||
result = mapper.parse_content(content)
|
||||
|
||||
assert result == ParsedContent(
|
||||
content="Hello from text chunk", reasoning_content=None
|
||||
)
|
||||
|
||||
def test_parse_content_thinking_chunk_extracts_reasoning(self):
|
||||
mapper = MistralMapper()
|
||||
content: list[mistralai.ContentChunk] = [
|
||||
mistralai.ThinkChunk(
|
||||
type="thinking",
|
||||
thinking=[mistralai.TextChunk(type="text", text="Let me think...")],
|
||||
),
|
||||
mistralai.TextChunk(type="text", text="The answer is 42."),
|
||||
]
|
||||
|
||||
result = mapper.parse_content(content)
|
||||
|
||||
assert result == ParsedContent(
|
||||
content="The answer is 42.", reasoning_content="Let me think..."
|
||||
)
|
||||
|
||||
def test_parse_content_multiple_thinking_chunks_concatenates(self):
|
||||
mapper = MistralMapper()
|
||||
content: list[mistralai.ContentChunk] = [
|
||||
mistralai.ThinkChunk(
|
||||
type="thinking",
|
||||
thinking=[mistralai.TextChunk(type="text", text="First thought. ")],
|
||||
),
|
||||
mistralai.ThinkChunk(
|
||||
type="thinking",
|
||||
thinking=[mistralai.TextChunk(type="text", text="Second thought.")],
|
||||
),
|
||||
mistralai.TextChunk(type="text", text="Final answer."),
|
||||
]
|
||||
|
||||
result = mapper.parse_content(content)
|
||||
|
||||
assert result == ParsedContent(
|
||||
content="Final answer.", reasoning_content="First thought. Second thought."
|
||||
)
|
||||
|
||||
def test_parse_content_thinking_only_returns_empty_content(self):
|
||||
mapper = MistralMapper()
|
||||
content: list[mistralai.ContentChunk] = [
|
||||
mistralai.ThinkChunk(
|
||||
type="thinking",
|
||||
thinking=[mistralai.TextChunk(type="text", text="Just thinking...")],
|
||||
)
|
||||
]
|
||||
|
||||
result = mapper.parse_content(content)
|
||||
|
||||
assert result == ParsedContent(content="", reasoning_content="Just thinking...")
|
||||
|
||||
def test_parse_content_empty_list_returns_empty(self):
|
||||
mapper = MistralMapper()
|
||||
content: list[mistralai.ContentChunk] = []
|
||||
|
||||
result = mapper.parse_content(content)
|
||||
|
||||
assert result == ParsedContent(content="", reasoning_content=None)
|
||||
|
||||
|
||||
class TestMistralMapperPrepareMessage:
|
||||
def test_prepare_assistant_message_without_reasoning(self):
|
||||
mapper = MistralMapper()
|
||||
msg = LLMMessage(role=Role.assistant, content="Hello!")
|
||||
|
||||
result = mapper.prepare_message(msg)
|
||||
|
||||
assert isinstance(result, mistralai.AssistantMessage)
|
||||
assert result.content == "Hello!"
|
||||
|
||||
def test_prepare_assistant_message_with_reasoning_creates_chunks(self):
|
||||
mapper = MistralMapper()
|
||||
msg = LLMMessage(
|
||||
role=Role.assistant,
|
||||
content="The answer is 42.",
|
||||
reasoning_content="Let me calculate...",
|
||||
)
|
||||
|
||||
result = mapper.prepare_message(msg)
|
||||
|
||||
assert isinstance(result, mistralai.AssistantMessage)
|
||||
assert isinstance(result.content, list)
|
||||
assert len(result.content) == 2
|
||||
|
||||
think_chunk = result.content[0]
|
||||
assert isinstance(think_chunk, mistralai.ThinkChunk)
|
||||
assert think_chunk.type == "thinking"
|
||||
assert len(think_chunk.thinking) == 1
|
||||
inner_chunk = think_chunk.thinking[0]
|
||||
assert isinstance(inner_chunk, mistralai.TextChunk)
|
||||
assert inner_chunk.text == "Let me calculate..."
|
||||
|
||||
text_chunk = result.content[1]
|
||||
assert isinstance(text_chunk, mistralai.TextChunk)
|
||||
assert text_chunk.type == "text"
|
||||
assert text_chunk.text == "The answer is 42."
|
||||
|
||||
def test_prepare_assistant_message_with_reasoning_and_none_content(self):
|
||||
mapper = MistralMapper()
|
||||
msg = LLMMessage(
|
||||
role=Role.assistant, content=None, reasoning_content="Just thinking..."
|
||||
)
|
||||
|
||||
result = mapper.prepare_message(msg)
|
||||
|
||||
assert isinstance(result, mistralai.AssistantMessage)
|
||||
assert isinstance(result.content, list)
|
||||
assert len(result.content) == 2
|
||||
|
||||
think_chunk = result.content[0]
|
||||
assert isinstance(think_chunk, mistralai.ThinkChunk)
|
||||
assert think_chunk.type == "thinking"
|
||||
assert len(think_chunk.thinking) == 1
|
||||
inner_chunk = think_chunk.thinking[0]
|
||||
assert isinstance(inner_chunk, mistralai.TextChunk)
|
||||
assert inner_chunk.text == "Just thinking..."
|
||||
|
||||
text_chunk = result.content[1]
|
||||
assert isinstance(text_chunk, mistralai.TextChunk)
|
||||
assert text_chunk.text == ""
|
||||
|
||||
|
||||
class TestGenericBackendReasoningContent:
|
||||
@pytest.mark.asyncio
|
||||
async def test_complete_extracts_reasoning_content(self):
|
||||
base_url = "https://api.example.com"
|
||||
json_response = {
|
||||
"id": "fake_id",
|
||||
"created": 1234567890,
|
||||
"model": "test-model",
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
||||
"object": "chat.completion",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "The answer is 42.",
|
||||
"reasoning_content": "Let me think step by step...",
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
with respx.mock(base_url=base_url) as mock_api:
|
||||
mock_api.post("/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(status_code=200, json=json_response)
|
||||
)
|
||||
provider = ProviderConfig(
|
||||
name="test", api_base=f"{base_url}/v1", api_key_env_var="API_KEY"
|
||||
)
|
||||
backend = GenericBackend(provider=provider)
|
||||
model = ModelConfig(name="test-model", provider="test", alias="test")
|
||||
messages = [LLMMessage(role=Role.user, content="What is the answer?")]
|
||||
|
||||
result = await backend.complete(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.2,
|
||||
tools=None,
|
||||
max_tokens=None,
|
||||
tool_choice=None,
|
||||
extra_headers=None,
|
||||
)
|
||||
|
||||
assert result.message.content == "The answer is 42."
|
||||
assert result.message.reasoning_content == "Let me think step by step..."
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_complete_streaming_extracts_reasoning_content(self):
|
||||
base_url = "https://api.example.com"
|
||||
chunks = [
|
||||
b'data: {"id":"id1","object":"chat.completion.chunk","created":123,"model":"test","choices":[{"index":0,"delta":{"role":"assistant","reasoning_content":"Thinking..."},"finish_reason":null}]}',
|
||||
b'data: {"id":"id1","object":"chat.completion.chunk","created":123,"model":"test","choices":[{"index":0,"delta":{"content":"Answer"},"finish_reason":null}]}',
|
||||
b'data: {"id":"id1","object":"chat.completion.chunk","created":123,"model":"test","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}',
|
||||
b"data: [DONE]",
|
||||
]
|
||||
|
||||
with respx.mock(base_url=base_url) as mock_api:
|
||||
mock_api.post("/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(
|
||||
status_code=200,
|
||||
stream=httpx.ByteStream(stream=b"\n\n".join(chunks)),
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
)
|
||||
)
|
||||
provider = ProviderConfig(
|
||||
name="test", api_base=f"{base_url}/v1", api_key_env_var="API_KEY"
|
||||
)
|
||||
backend = GenericBackend(provider=provider)
|
||||
model = ModelConfig(name="test-model", provider="test", alias="test")
|
||||
messages = [LLMMessage(role=Role.user, content="Stream please")]
|
||||
|
||||
results = []
|
||||
async for chunk in backend.complete_streaming(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.2,
|
||||
tools=None,
|
||||
max_tokens=None,
|
||||
tool_choice=None,
|
||||
extra_headers=None,
|
||||
):
|
||||
results.append(chunk)
|
||||
|
||||
assert results[0].message.reasoning_content == "Thinking..."
|
||||
assert results[0].message.content == ""
|
||||
assert results[1].message.content == "Answer"
|
||||
assert results[1].message.reasoning_content is None
|
||||
|
||||
|
||||
class TestAPIToolFormatHandlerReasoningContent:
|
||||
def test_process_api_response_message_extracts_reasoning_content(self):
|
||||
handler = APIToolFormatHandler()
|
||||
|
||||
mock_message = MagicMock()
|
||||
mock_message.role = "assistant"
|
||||
mock_message.content = "The answer is 42."
|
||||
mock_message.reasoning_content = "Let me think..."
|
||||
mock_message.tool_calls = None
|
||||
|
||||
result = handler.process_api_response_message(mock_message)
|
||||
|
||||
assert result.content == "The answer is 42."
|
||||
assert result.reasoning_content == "Let me think..."
|
||||
|
||||
def test_process_api_response_message_handles_missing_reasoning_content(self):
|
||||
handler = APIToolFormatHandler()
|
||||
|
||||
mock_message = MagicMock(spec=["role", "content", "tool_calls"])
|
||||
mock_message.role = "assistant"
|
||||
mock_message.content = "Hello"
|
||||
mock_message.tool_calls = None
|
||||
|
||||
result = handler.process_api_response_message(mock_message)
|
||||
|
||||
assert result.content == "Hello"
|
||||
assert result.reasoning_content is None
|
||||
|
||||
|
||||
class TestAgentStreamingReasoningEvents:
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_accumulates_reasoning_in_message(self):
|
||||
backend = FakeBackend([
|
||||
mock_llm_chunk(content="", reasoning_content="First thought. "),
|
||||
mock_llm_chunk(content="", reasoning_content="Second thought."),
|
||||
mock_llm_chunk(content="Final answer."),
|
||||
])
|
||||
agent = Agent(make_config(), backend=backend, enable_streaming=True)
|
||||
|
||||
[_ async for _ in agent.act("Think and answer")]
|
||||
|
||||
assistant_msg = next(m for m in agent.messages if m.role == Role.assistant)
|
||||
assert assistant_msg.reasoning_content == "First thought. Second thought."
|
||||
assert assistant_msg.content == "Final answer."
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_content_only_no_reasoning(self):
|
||||
backend = FakeBackend([
|
||||
mock_llm_chunk(content="Hello "),
|
||||
mock_llm_chunk(content="world!"),
|
||||
])
|
||||
agent = Agent(make_config(), backend=backend, enable_streaming=True)
|
||||
|
||||
events = [event async for event in agent.act("Say hello")]
|
||||
|
||||
reasoning_events = [e for e in events if isinstance(e, ReasoningEvent)]
|
||||
assert len(reasoning_events) == 0
|
||||
|
||||
assistant_events = [e for e in events if isinstance(e, AssistantEvent)]
|
||||
assert len(assistant_events) == 1
|
||||
|
||||
assistant_msg = next(m for m in agent.messages if m.role == Role.assistant)
|
||||
assert assistant_msg.reasoning_content is None
|
||||
assert assistant_msg.content == "Hello world!"
|
||||
|
||||
|
||||
class TestLLMMessageReasoningContent:
|
||||
def test_llm_message_from_dict_with_reasoning_content(self):
|
||||
data = {
|
||||
"role": "assistant",
|
||||
"content": "Answer",
|
||||
"reasoning_content": "Thinking...",
|
||||
}
|
||||
|
||||
msg = LLMMessage.model_validate(data)
|
||||
|
||||
assert msg.reasoning_content == "Thinking..."
|
||||
|
||||
def test_llm_message_model_dump_includes_reasoning_content(self):
|
||||
msg = LLMMessage(
|
||||
role=Role.assistant, content="Answer", reasoning_content="Thinking..."
|
||||
)
|
||||
|
||||
dumped = msg.model_dump(exclude_none=True)
|
||||
|
||||
assert dumped["reasoning_content"] == "Thinking..."
|
||||
|
||||
def test_llm_message_model_dump_excludes_none_reasoning_content(self):
|
||||
msg = LLMMessage(role=Role.assistant, content="Answer")
|
||||
|
||||
dumped = msg.model_dump(exclude_none=True)
|
||||
|
||||
assert "reasoning_content" not in dumped
|
||||
Loading…
Add table
Add a link
Reference in a new issue