Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai>
Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai>
This commit is contained in:
Mathias Gesbert 2025-12-22 13:33:20 +01:00 committed by Mathias Gesbert
parent 402e898f39
commit 2e1e15120d
32 changed files with 391 additions and 549 deletions

View file

@ -108,12 +108,17 @@ jobs:
with:
path: artifacts
- name: Zip artifacts like GitHub UI
- name: Zip artifacts
run: |
mkdir release-assets
for dir in artifacts/*; do
name=$(basename "$dir")
(cd artifacts && zip -r "../release-assets/${name}.zip" "$name")
if [ -f "$dir/vibe-acp" ]; then
chmod +x "$dir/vibe-acp"
zip -j "release-assets/${name}.zip" "$dir/vibe-acp"
elif [ -f "$dir/vibe-acp.exe" ]; then
zip -j "release-assets/${name}.zip" "$dir/vibe-acp.exe"
fi
done
- name: Attach binaries to release

2
.vscode/launch.json vendored
View file

@ -1,5 +1,5 @@
{
"version": "1.2.1",
"version": "1.2.2",
"configurations": [
{
"name": "ACP Server",

View file

@ -5,6 +5,14 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [1.2.2] - 2025-12-22
### Fixed
- Remove dead code
- Fix artefacts automatically attached to the release
- Refactor agent post streaming
## [1.2.1] - 2025-12-18
### Fixed

View file

@ -1,7 +1,7 @@
id = "mistral-vibe"
name = "Mistral Vibe"
description = "Mistral's open-source coding assistant"
version = "1.2.1"
version = "1.2.2"
schema_version = 1
authors = ["Mistral AI"]
repository = "https://github.com/mistralai/mistral-vibe"
@ -11,25 +11,25 @@ name = "Mistral Vibe"
icon = "./icons/mistral_vibe.svg"
[agent_servers.mistral-vibe.targets.darwin-aarch64]
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.1/vibe-acp-darwin-aarch64-1.2.1.zip"
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.2/vibe-acp-darwin-aarch64-1.2.2.zip"
cmd = "./vibe-acp"
[agent_servers.mistral-vibe.targets.darwin-x86_64]
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.1/vibe-acp-darwin-x86_64-1.2.1.zip"
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.2/vibe-acp-darwin-x86_64-1.2.2.zip"
cmd = "./vibe-acp"
[agent_servers.mistral-vibe.targets.linux-aarch64]
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.1/vibe-acp-linux-aarch64-1.2.1.zip"
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.2/vibe-acp-linux-aarch64-1.2.2.zip"
cmd = "./vibe-acp"
[agent_servers.mistral-vibe.targets.linux-x86_64]
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.1/vibe-acp-linux-x86_64-1.2.1.zip"
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.2/vibe-acp-linux-x86_64-1.2.2.zip"
cmd = "./vibe-acp"
[agent_servers.mistral-vibe.targets.windows-aarch64]
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.1/vibe-acp-windows-aarch64-1.2.1.zip"
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.2/vibe-acp-windows-aarch64-1.2.2.zip"
cmd = "./vibe-acp.exe"
[agent_servers.mistral-vibe.targets.windows-x86_64]
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.1/vibe-acp-windows-x86_64-1.2.1.zip"
archive = "https://github.com/mistralai/mistral-vibe/releases/download/v1.2.2/vibe-acp-windows-x86_64-1.2.2.zip"
cmd = "./vibe-acp.exe"

View file

@ -1,6 +1,6 @@
[project]
name = "mistral-vibe"
version = "1.2.1"
version = "1.2.2"
description = "Minimal CLI coding agent by Mistral"
readme = "README.md"
requires-python = ">=3.12"

View file

@ -426,7 +426,7 @@ class TestSessionManagement:
class TestSessionUpdates:
@pytest.mark.asyncio
async def test_agent_message_chunk_structure(self, vibe_home_dir: Path) -> None:
mock_env = get_mocking_env([mock_llm_chunk(content="Hi") for _ in range(2)])
mock_env = get_mocking_env([mock_llm_chunk(content="Hi")])
async for process in get_acp_agent_process(
mock_env=mock_env, vibe_home=vibe_home_dir
):
@ -465,7 +465,6 @@ class TestSessionUpdates:
@pytest.mark.asyncio
async def test_tool_call_update_structure(self, vibe_home_dir: Path) -> None:
mock_env = get_mocking_env([
mock_llm_chunk(content="Hey"),
mock_llm_chunk(
tool_calls=[
ToolCall(
@ -475,14 +474,9 @@ class TestSessionUpdates:
type="function",
index=0,
)
],
name="bash",
finish_reason="tool_calls",
),
mock_llm_chunk(
content="The files containing the pattern 'auth' are ...",
finish_reason="stop",
]
),
mock_llm_chunk(content="The files containing the pattern 'auth' are ..."),
])
async for process in get_acp_agent_process(
mock_env=mock_env, vibe_home=vibe_home_dir
@ -567,7 +561,6 @@ class TestToolCallStructure:
self, vibe_home_grep_ask: Path
) -> None:
custom_results = [
mock_llm_chunk(content="Hey"),
mock_llm_chunk(
tool_calls=[
ToolCall(
@ -578,10 +571,8 @@ class TestToolCallStructure:
type="function",
index=0,
)
],
name="grep",
finish_reason="tool_calls",
),
]
)
]
mock_env = get_mocking_env(custom_results)
async for process in get_acp_agent_process(
@ -628,7 +619,6 @@ class TestToolCallStructure:
self, vibe_home_grep_ask: Path
) -> None:
custom_results = [
mock_llm_chunk(content="Hey"),
mock_llm_chunk(
tool_calls=[
ToolCall(
@ -639,13 +629,10 @@ class TestToolCallStructure:
type="function",
index=0,
)
],
name="grep",
finish_reason="tool_calls",
),
mock_llm_chunk(
content="The search for 'auth' has been completed", finish_reason="stop"
]
),
mock_llm_chunk(content="The search for 'auth' has been completed"),
mock_llm_chunk(content="The file test.txt has been created"),
]
mock_env = get_mocking_env(custom_results)
async for process in get_acp_agent_process(
@ -692,7 +679,6 @@ class TestToolCallStructure:
self, vibe_home_grep_ask: Path
) -> None:
custom_results = [
mock_llm_chunk(content="Hey"),
mock_llm_chunk(
tool_calls=[
ToolCall(
@ -703,14 +689,11 @@ class TestToolCallStructure:
type="function",
index=0,
)
],
name="grep",
finish_reason="tool_calls",
]
),
mock_llm_chunk(
content="The search for 'auth' has not been performed, "
"because you rejected the permission request",
finish_reason="stop",
"because you rejected the permission request"
),
]
mock_env = get_mocking_env(custom_results)
@ -759,7 +742,6 @@ class TestToolCallStructure:
self, vibe_home_grep_ask: Path
) -> None:
custom_results = [
mock_llm_chunk(content="Hey"),
mock_llm_chunk(
tool_calls=[
ToolCall(
@ -770,13 +752,10 @@ class TestToolCallStructure:
type="function",
index=0,
)
],
name="grep",
finish_reason="tool_calls",
),
mock_llm_chunk(
content="The search for 'auth' has been completed", finish_reason="stop"
]
),
mock_llm_chunk(content="The search for 'auth' has been completed"),
mock_llm_chunk(content="The command sleep 3 has been run"),
]
mock_env = get_mocking_env(custom_results)
async for process in get_acp_agent_process(
@ -820,7 +799,6 @@ class TestToolCallStructure:
self, vibe_home_grep_ask: Path
) -> None:
custom_results = [
mock_llm_chunk(content="Hey"),
mock_llm_chunk(
tool_calls=[
ToolCall(
@ -831,14 +809,11 @@ class TestToolCallStructure:
type="function",
index=0,
)
],
name="grep",
finish_reason="tool_calls",
]
),
mock_llm_chunk(
content="The search for 'auth' has failed "
"because the path does not exist",
finish_reason="stop",
"because the path does not exist"
),
]
mock_env = get_mocking_env(custom_results)
@ -894,7 +869,6 @@ class TestCancellationStructure:
self, vibe_home_dir: Path
) -> None:
custom_results = [
mock_llm_chunk(content="Hey"),
mock_llm_chunk(
tool_calls=[
ToolCall(
@ -906,14 +880,11 @@ class TestCancellationStructure:
type="function",
index=0,
)
],
name="write_file",
finish_reason="tool_calls",
]
),
mock_llm_chunk(
content="The file test.txt has not been created, "
"because you cancelled the permission request",
finish_reason="stop",
"because you cancelled the permission request"
),
]
mock_env = get_mocking_env(custom_results)

View file

@ -22,13 +22,10 @@ from vibe.core.types import LLMChunk, LLMMessage, LLMUsage, Role
@pytest.fixture
def backend() -> FakeBackend:
backend = FakeBackend(
results=[
LLMChunk(
message=LLMMessage(role=Role.assistant, content="Hi"),
finish_reason="end_turn",
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
)
]
)
return backend

View file

@ -41,7 +41,7 @@ class TestACPInitialize:
),
)
assert response.agentInfo == Implementation(
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="1.2.1"
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="1.2.2"
)
assert response.authMethods == []
@ -63,7 +63,7 @@ class TestACPInitialize:
),
)
assert response.agentInfo == Implementation(
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="1.2.1"
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="1.2.2"
)
assert response.authMethods is not None

View file

@ -116,9 +116,9 @@ class TestMultiSessionCore:
)
session2 = acp_agent.sessions[session2_response.sessionId]
backend._chunks = [
mock_llm_chunk(content="Response 1", finish_reason="stop"),
mock_llm_chunk(content="Response 2", finish_reason="stop"),
backend._streams = [
[mock_llm_chunk(content="Response 1")],
[mock_llm_chunk(content="Response 2")],
]
async def run_session1():

View file

@ -18,13 +18,10 @@ from vibe.core.types import LLMChunk, LLMMessage, LLMUsage, Role
@pytest.fixture
def backend() -> FakeBackend:
backend = FakeBackend(
results=[
LLMChunk(
message=LLMMessage(role=Role.assistant, content="Hi"),
finish_reason="end_turn",
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
)
]
)
return backend

View file

@ -17,13 +17,10 @@ from vibe.core.types import LLMChunk, LLMMessage, LLMUsage, Role
@pytest.fixture
def backend() -> FakeBackend:
backend = FakeBackend(
results=[
LLMChunk(
message=LLMMessage(role=Role.assistant, content="Hi"),
finish_reason="end_turn",
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
)
]
)
return backend

View file

@ -17,13 +17,10 @@ from vibe.core.types import LLMChunk, LLMMessage, LLMUsage, Role
@pytest.fixture
def backend() -> FakeBackend:
backend = FakeBackend(
results=[
LLMChunk(
message=LLMMessage(role=Role.assistant, content="Hi"),
finish_reason="end_turn",
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
)
]
)
return backend

View file

@ -30,7 +30,6 @@ SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
},
{
"message": "Some content",
"finish_reason": "stop",
"usage": {
"prompt_tokens": 100,
"total_tokens": 300,
@ -78,7 +77,6 @@ TOOL_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
},
{
"message": "",
"finish_reason": "tool_calls",
"tool_calls": [
{
"name": "some_tool",
@ -102,26 +100,13 @@ STREAMED_SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultDat
rb"data: [DONE]",
],
[
{
"message": "",
"finish_reason": None,
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{
"message": "",
"finish_reason": None,
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
{
"message": "Some content",
"finish_reason": None,
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{
"message": "",
"finish_reason": "stop",
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
},
{"message": "", "usage": {"prompt_tokens": 100, "completion_tokens": 200}},
],
)
]
@ -140,30 +125,19 @@ STREAMED_TOOL_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]
rb"data: [DONE]",
],
[
{
"message": "",
"finish_reason": None,
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{
"message": "",
"finish_reason": None,
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
{
"message": "Some content",
"finish_reason": None,
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{
"message": "",
"finish_reason": None,
"tool_calls": [{"name": "some_tool", "arguments": None, "index": 0}],
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{
"message": "",
"finish_reason": None,
"tool_calls": [
{
"name": None,
@ -173,11 +147,7 @@ STREAMED_TOOL_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]
],
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{
"message": "",
"finish_reason": "tool_calls",
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
},
{"message": "", "usage": {"prompt_tokens": 100, "completion_tokens": 200}},
],
)
]

View file

@ -29,7 +29,6 @@ SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
},
{
"message": "Some content",
"finish_reason": "stop",
"usage": {
"prompt_tokens": 100,
"total_tokens": 300,
@ -75,7 +74,6 @@ TOOL_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
},
{
"message": "Some content",
"finish_reason": "tool_calls",
"tool_calls": [
{
"name": "some_tool",
@ -98,21 +96,12 @@ STREAMED_SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultDat
rb"data: [DONE]",
],
[
{
"message": "",
"finish_reason": None,
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
{
"message": "Some content",
"finish_reason": None,
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{
"message": "",
"finish_reason": "stop",
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
},
{"message": "", "usage": {"prompt_tokens": 100, "completion_tokens": 200}},
],
)
]
@ -131,25 +120,18 @@ STREAMED_TOOL_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]
rb"data: [DONE]",
],
[
{
"message": "",
"finish_reason": None,
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
{
"message": "Some content",
"finish_reason": None,
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{
"message": "",
"finish_reason": None,
"tool_calls": [{"name": "some_tool", "arguments": "", "index": 0}],
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{
"message": "",
"finish_reason": None,
"tool_calls": [
{"name": "", "arguments": '{"some_argument": ', "index": 0}
],
@ -157,17 +139,12 @@ STREAMED_TOOL_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]
},
{
"message": "",
"finish_reason": None,
"tool_calls": [
{"name": "", "arguments": '"some_argument_value"}', "index": 0}
],
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
},
{
"message": "",
"finish_reason": "tool_calls",
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
},
{"message": "", "usage": {"prompt_tokens": 100, "completion_tokens": 200}},
],
)
]

View file

@ -87,7 +87,6 @@ class TestBackend:
)
assert result.message.content == result_data["message"]
assert result.finish_reason == result_data["finish_reason"]
assert result.usage is not None
assert (
result.usage.prompt_tokens == result_data["usage"]["prompt_tokens"]
@ -165,7 +164,6 @@ class TestBackend:
for result, expected_result in zip(results, result_data, strict=True):
assert result.message.content == expected_result["message"]
assert result.finish_reason == expected_result["finish_reason"]
assert result.usage is not None
assert (
result.usage.prompt_tokens

View file

@ -13,7 +13,6 @@ def mock_llm_chunk(
tool_calls: list[ToolCall] | None = None,
name: str | None = None,
tool_call_id: str | None = None,
finish_reason: str | None = None,
prompt_tokens: int = 10,
completion_tokens: int = 5,
) -> LLMChunk:
@ -29,7 +28,6 @@ def mock_llm_chunk(
usage=LLMUsage(
prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
),
finish_reason=finish_reason,
)

View file

@ -13,13 +13,11 @@ class SnapshotTestAppWithConversation(BaseSnapshotTestApp):
def __init__(self) -> None:
config = default_config()
fake_backend = FakeBackend(
results=[
mock_llm_chunk(
content="I'm the Vibe agent and I'm ready to help.",
prompt_tokens=10_000,
completion_tokens=2_500,
)
]
)
super().__init__(config=config)
self.agent = Agent(

View file

@ -1,9 +1,10 @@
from __future__ import annotations
from collections.abc import AsyncGenerator, Callable, Iterable
from typing import cast
from tests.mock.utils import mock_llm_chunk
from vibe.core.types import LLMChunk, LLMMessage
from vibe.core.types import LLMChunk, LLMMessage, Role
class FakeBackend:
@ -15,18 +16,47 @@ class FakeBackend:
def __init__(
self,
results: Iterable[LLMChunk] | None = None,
chunks: LLMChunk
| Iterable[LLMChunk]
| Iterable[Iterable[LLMChunk]]
| None = None,
*,
token_counter: Callable[[list[LLMMessage]], int] | None = None,
exception_to_raise: Exception | None = None,
) -> None:
self._chunks = list(results or [])
"""Fake backend that will output the given chunks in the order they are given.
chunks: A single chunk, a sequence of chunks, or a sequence of sequences of chunks.
A single chunk would be outputted as such in complete / complete_streaming
A sequence of chunks will is considered a single stream: a completion would output
all chunks (either streaming or in an aggregated way)
A sequence of sequences of chunks is considered a list of streams: each completion
will output a stream (either streaming or in an aggregated way)
"""
self._requests_messages: list[list[LLMMessage]] = []
self._requests_extra_headers: list[dict[str, str] | None] = []
self._count_tokens_calls: list[list[LLMMessage]] = []
self._token_counter = token_counter or self._default_token_counter
self._exception_to_raise = exception_to_raise
self._streams: list[list[LLMChunk]]
if chunks is None:
self._streams = []
return
if isinstance(chunks, LLMChunk):
self._streams = [[chunks]]
return
if all(isinstance(chunk, LLMChunk) for chunk in chunks):
self._streams = [[cast(LLMChunk, chunk) for chunk in chunks]]
return
if any(isinstance(chunk, LLMChunk) for chunk in chunks):
raise TypeError(
f"Invalid type for chunks, expected a value of type "
f"LLMChunk | Iterable[LLMChunk] | Iterable[Iterable[LLMChunk]], got {chunks!r}"
)
chunks = cast(Iterable[Iterable[LLMChunk]], chunks)
self._streams = [[chunk for chunk in stream] for stream in chunks]
@property
def requests_messages(self) -> list[list[LLMMessage]]:
return self._requests_messages
@ -61,12 +91,15 @@ class FakeBackend:
self._requests_messages.append(messages)
self._requests_extra_headers.append(extra_headers)
if self._chunks:
chunk = self._chunks.pop(0)
if not self._chunks:
chunk = chunk.model_copy(update={"finish_reason": "stop"})
return chunk
return mock_llm_chunk(content="", finish_reason="stop")
if self._streams:
stream = self._streams.pop(0)
chunk_agg = LLMChunk(message=LLMMessage(role=Role.assistant))
for chunk in stream:
chunk_agg += chunk
return chunk_agg
return mock_llm_chunk(content="")
async def complete_streaming(
self,
@ -84,22 +117,13 @@ class FakeBackend:
self._requests_messages.append(messages)
self._requests_extra_headers.append(extra_headers)
has_final_chunk = False
while self._chunks:
chunk = self._chunks.pop(0)
is_last_provided_chunk = not self._chunks
if is_last_provided_chunk:
chunk = chunk.model_copy(update={"finish_reason": "stop"})
if chunk.finish_reason is not None:
has_final_chunk = True
if self._streams:
stream = list(self._streams.pop(0))
else:
stream = [mock_llm_chunk(content="")]
for chunk in stream:
yield chunk
if has_final_chunk:
break
if not has_final_chunk:
yield mock_llm_chunk(content="", finish_reason="stop")
async def count_tokens(
self,

View file

@ -23,8 +23,8 @@ async def test_auto_compact_triggers_and_batches_observer() -> None:
observed.append((msg.role, msg.content))
backend = FakeBackend([
mock_llm_chunk(content="<summary>"),
mock_llm_chunk(content="<final>"),
[mock_llm_chunk(content="<summary>")],
[mock_llm_chunk(content="<final>")],
])
cfg = VibeConfig(
session_logging=SessionLoggingConfig(enabled=False), auto_compact_threshold=1

View file

@ -15,7 +15,7 @@ def vibe_config() -> VibeConfig:
@pytest.mark.asyncio
async def test_passes_x_affinity_header_when_asking_an_answer(vibe_config: VibeConfig):
backend = FakeBackend([mock_llm_chunk(content="Response", finish_reason="stop")])
backend = FakeBackend([mock_llm_chunk(content="Response")])
agent = Agent(vibe_config, backend=backend)
[_ async for _ in agent.act("Hello")]
@ -31,7 +31,7 @@ async def test_passes_x_affinity_header_when_asking_an_answer(vibe_config: VibeC
async def test_passes_x_affinity_header_when_asking_an_answer_streaming(
vibe_config: VibeConfig,
):
backend = FakeBackend([mock_llm_chunk(content="Response", finish_reason="stop")])
backend = FakeBackend([mock_llm_chunk(content="Response")])
agent = Agent(vibe_config, backend=backend, enable_streaming=True)
[_ async for _ in agent.act("Hello")]
@ -45,12 +45,7 @@ async def test_passes_x_affinity_header_when_asking_an_answer_streaming(
@pytest.mark.asyncio
async def test_updates_tokens_stats_based_on_backend_response(vibe_config: VibeConfig):
chunk = mock_llm_chunk(
content="Response",
finish_reason="stop",
prompt_tokens=100,
completion_tokens=50,
)
chunk = mock_llm_chunk(content="Response", prompt_tokens=100, completion_tokens=50)
backend = FakeBackend([chunk])
agent = Agent(vibe_config, backend=backend)
@ -64,10 +59,7 @@ async def test_updates_tokens_stats_based_on_backend_response_streaming(
vibe_config: VibeConfig,
):
final_chunk = mock_llm_chunk(
content="Complete",
finish_reason="stop",
prompt_tokens=200,
completion_tokens=75,
content="Complete", prompt_tokens=200, completion_tokens=75
)
backend = FakeBackend([final_chunk])
agent = Agent(vibe_config, backend=backend, enable_streaming=True)

View file

@ -10,7 +10,6 @@ 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,
@ -25,7 +24,6 @@ from vibe.core.types import (
ApprovalResponse,
AssistantEvent,
FunctionCall,
LLMChunk,
LLMMessage,
Role,
ToolCall,
@ -177,10 +175,11 @@ async def test_act_handles_streaming_with_tool_call_events_in_sequence() -> None
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."),
],
[mock_llm_chunk(content="Done reviewing todos.")],
])
agent = Agent(
make_config(
@ -222,7 +221,7 @@ async def test_act_handles_tool_call_chunk_with_content() -> None:
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"),
mock_llm_chunk(content=" complete"),
])
agent = Agent(
make_config(
@ -237,15 +236,12 @@ async def test_act_handles_tool_call_chunk_with_content() -> None:
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 events[0].content == "Preparing todo request complete"
assert any(
m.role == Role.assistant and m.content == "Preparing todo request complete"
for m in agent.messages
@ -304,35 +300,6 @@ async def test_act_merges_streamed_tool_call_arguments() -> None:
)
@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):
@ -359,7 +326,6 @@ async def test_act_handles_user_cancellation_during_streaming() -> None:
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(

View file

@ -159,9 +159,7 @@ class TestAgentStatsHelpers:
class TestReloadPreservesStats:
@pytest.mark.asyncio
async def test_reload_preserves_session_tokens(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="First response", finish_reason="stop")
])
backend = FakeBackend(mock_llm_chunk(content="First response"))
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("Hello"):
@ -185,13 +183,14 @@ class TestReloadPreservesStats:
tool_calls=[
ToolCall(
id="tc1",
index=0,
function=FunctionCall(
name="todo", arguments='{"action": "read"}'
),
)
],
),
mock_llm_chunk(content="Done", finish_reason="stop"),
mock_llm_chunk(content="Done"),
])
config = make_config(enabled_tools=["todo"])
agent = Agent(config, mode=AgentMode.AUTO_APPROVE, backend=backend)
@ -210,8 +209,8 @@ class TestReloadPreservesStats:
@pytest.mark.asyncio
async def test_reload_preserves_steps(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="R1", finish_reason="stop"),
mock_llm_chunk(content="R2", finish_reason="stop"),
[mock_llm_chunk(content="R1")],
[mock_llm_chunk(content="R2")],
])
agent = Agent(make_config(), backend=backend)
@ -231,9 +230,7 @@ class TestReloadPreservesStats:
async def test_reload_preserves_context_tokens_when_messages_preserved(
self,
) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
backend = FakeBackend(mock_llm_chunk(content="Response"))
agent = Agent(make_config(), backend=backend)
[_ async for _ in agent.act("Hello")]
assert agent.stats.context_tokens > 0
@ -258,10 +255,10 @@ class TestReloadPreservesStats:
assert agent.stats.context_tokens == 0
@pytest.mark.asyncio
async def test_reload_preserves_context_tokens_when_messages_exist(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
async def test_reload_resets_context_tokens_when_system_prompt_changes(
self,
) -> None:
backend = FakeBackend(mock_llm_chunk(content="Response"))
config1 = make_config(system_prompt_id="tests")
config2 = make_config(system_prompt_id="cli")
agent = Agent(config1, backend=backend)
@ -279,9 +276,7 @@ class TestReloadPreservesStats:
async def test_reload_updates_pricing_from_new_model(self, monkeypatch) -> None:
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
backend = FakeBackend(mock_llm_chunk(content="Response"))
config_mistral = make_config(active_model="devstral-latest")
agent = Agent(config_mistral, backend=backend)
@ -302,8 +297,8 @@ class TestReloadPreservesStats:
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
backend = FakeBackend([
mock_llm_chunk(content="First", finish_reason="stop"),
mock_llm_chunk(content="After reload", finish_reason="stop"),
[mock_llm_chunk(content="First")],
[mock_llm_chunk(content="After reload")],
])
config1 = make_config(active_model="devstral-latest")
agent = Agent(config1, backend=backend)
@ -330,9 +325,7 @@ class TestReloadPreservesStats:
class TestReloadPreservesMessages:
@pytest.mark.asyncio
async def test_reload_preserves_conversation_messages(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
backend = FakeBackend(mock_llm_chunk(content="Response"))
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("Hello"):
@ -353,9 +346,7 @@ class TestReloadPreservesMessages:
@pytest.mark.asyncio
async def test_reload_updates_system_prompt_preserves_rest(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
backend = FakeBackend(mock_llm_chunk(content="Response"))
config1 = make_config(system_prompt_id="tests")
agent = Agent(config1, backend=backend)
@ -388,9 +379,7 @@ class TestReloadPreservesMessages:
self, observer_capture
) -> None:
observed, observer = observer_capture
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
backend = FakeBackend(mock_llm_chunk(content="Response"))
agent = Agent(make_config(), message_observer=observer, backend=backend)
async for _ in agent.act("Hello"):
@ -410,8 +399,8 @@ class TestCompactStatsHandling:
@pytest.mark.asyncio
async def test_compact_preserves_cumulative_stats(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="First response", finish_reason="stop"),
mock_llm_chunk(content="<summary>", finish_reason="stop"),
[mock_llm_chunk(content="First response")],
[mock_llm_chunk(content="<summary>")],
])
agent = Agent(make_config(), backend=backend)
@ -432,8 +421,8 @@ class TestCompactStatsHandling:
@pytest.mark.asyncio
async def test_compact_updates_context_tokens(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Long response " * 100, finish_reason="stop"),
mock_llm_chunk(content="<summary>", finish_reason="stop"),
[mock_llm_chunk(content="Long response " * 100)],
[mock_llm_chunk(content="<summary>")],
])
agent = Agent(make_config(), backend=backend)
@ -449,19 +438,22 @@ class TestCompactStatsHandling:
@pytest.mark.asyncio
async def test_compact_preserves_tool_call_stats(self) -> None:
backend = FakeBackend([
[
mock_llm_chunk(
content="Using tool",
tool_calls=[
ToolCall(
id="tc1",
index=0,
function=FunctionCall(
name="todo", arguments='{"action": "read"}'
),
)
],
),
mock_llm_chunk(content="Done", finish_reason="stop"),
mock_llm_chunk(content="<summary>", finish_reason="stop"),
mock_llm_chunk(content=" todo"),
],
[mock_llm_chunk(content="<summary>")],
])
config = make_config(enabled_tools=["todo"])
agent = Agent(config, mode=AgentMode.AUTO_APPROVE, backend=backend)
@ -478,8 +470,8 @@ class TestCompactStatsHandling:
@pytest.mark.asyncio
async def test_compact_resets_session_id(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Long response " * 100, finish_reason="stop"),
mock_llm_chunk(content="<summary>", finish_reason="stop"),
[mock_llm_chunk(content="Long response " * 100)],
[mock_llm_chunk(content="<summary>")],
])
agent = Agent(make_config(disable_logging=False), backend=backend)
@ -506,8 +498,8 @@ class TestAutoCompactIntegration:
observed.append((msg.role, msg.content))
backend = FakeBackend([
mock_llm_chunk(content="<summary>", finish_reason="stop"),
mock_llm_chunk(content="<final>", finish_reason="stop"),
[mock_llm_chunk(content="<summary>")],
[mock_llm_chunk(content="<final>")],
])
cfg = VibeConfig(
session_logging=SessionLoggingConfig(enabled=False),
@ -544,9 +536,7 @@ class TestAutoCompactIntegration:
class TestClearHistoryFullReset:
@pytest.mark.asyncio
async def test_clear_history_fully_resets_stats(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
backend = FakeBackend(mock_llm_chunk(content="Response"))
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("Hello"):
@ -563,9 +553,7 @@ class TestClearHistoryFullReset:
@pytest.mark.asyncio
async def test_clear_history_preserves_pricing(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
backend = FakeBackend(mock_llm_chunk(content="Response"))
config = make_config(input_price=0.4, output_price=2.0)
agent = Agent(config, backend=backend)
@ -579,9 +567,7 @@ class TestClearHistoryFullReset:
@pytest.mark.asyncio
async def test_clear_history_removes_messages(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
backend = FakeBackend(mock_llm_chunk(content="Response"))
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("Hello"):
@ -596,9 +582,7 @@ class TestClearHistoryFullReset:
@pytest.mark.asyncio
async def test_clear_history_resets_session_id(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
backend = FakeBackend(mock_llm_chunk(content="Response"))
agent = Agent(make_config(disable_logging=False), backend=backend)
original_session_id = agent.session_id
@ -622,9 +606,7 @@ class TestStatsEdgeCases:
) -> None:
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
backend = FakeBackend(mock_llm_chunk(content="Response"))
config1 = make_config(active_model="devstral-latest")
agent = Agent(config1, backend=backend)
@ -643,9 +625,9 @@ class TestStatsEdgeCases:
@pytest.mark.asyncio
async def test_multiple_reloads_accumulate_correctly(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="R1", finish_reason="stop"),
mock_llm_chunk(content="R2", finish_reason="stop"),
mock_llm_chunk(content="R3", finish_reason="stop"),
[mock_llm_chunk(content="R1")],
[mock_llm_chunk(content="R2")],
[mock_llm_chunk(content="R3")],
])
agent = Agent(make_config(), backend=backend)
@ -668,9 +650,9 @@ class TestStatsEdgeCases:
@pytest.mark.asyncio
async def test_compact_then_reload_preserves_both(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Initial response", finish_reason="stop"),
mock_llm_chunk(content="<summary>", finish_reason="stop"),
mock_llm_chunk(content="After reload", finish_reason="stop"),
[mock_llm_chunk(content="Initial response")],
[mock_llm_chunk(content="<summary>")],
[mock_llm_chunk(content="After reload")],
])
agent = Agent(make_config(), backend=backend)

View file

@ -44,10 +44,12 @@ def make_config(todo_permission: ToolPermission = ToolPermission.ALWAYS) -> Vibe
)
def make_todo_tool_call(call_id: str, action: str = "read") -> ToolCall:
def make_todo_tool_call(
call_id: str, index: int = 0, arguments: str | None = None
) -> ToolCall:
args = arguments if arguments is not None else '{"action": "read"}'
return ToolCall(
id=call_id,
function=FunctionCall(name="todo", arguments=f'{{"action": "{action}"}}'),
id=call_id, index=index, function=FunctionCall(name="todo", arguments=args)
)
@ -72,8 +74,8 @@ 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"),
[mock_llm_chunk(content="Let me check your todos.", tool_calls=[tool_call])],
[mock_llm_chunk(content="I retrieved 0 todos.")],
])
agent = make_agent(auto_approve=True, backend=backend)
@ -108,14 +110,13 @@ async def test_tool_call_requires_approval_if_not_auto_approved() -> None:
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",
),
)
],
[mock_llm_chunk(content="I cannot execute the tool without approval.")],
]),
)
@ -148,11 +149,13 @@ async def test_tool_call_approved_by_callback() -> None:
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"),
)
],
[mock_llm_chunk(content="I retrieved 0 todos.")],
]),
)
@ -183,13 +186,13 @@ async def test_tool_call_rejected_when_auto_approve_disabled_and_rejected_by_cal
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"
),
)
],
[mock_llm_chunk(content="Understood, I won't check the todos.")],
]),
)
@ -211,11 +214,13 @@ async def test_tool_call_skipped_when_permission_is_never() -> None:
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"),
)
],
[mock_llm_chunk(content="Tool is disabled.")],
]),
)
@ -257,14 +262,20 @@ async def test_approval_always_sets_tool_permission_for_subsequent_calls() -> No
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"),
content="First check.",
tool_calls=[make_todo_tool_call("call_first")],
)
],
[mock_llm_chunk(content="First done.")],
[
mock_llm_chunk(
content="Second check.", tool_calls=[make_todo_tool_call("call_second")]
),
mock_llm_chunk(content="Second done.", finish_reason="stop"),
content="Second check.",
tool_calls=[make_todo_tool_call("call_second")],
)
],
[mock_llm_chunk(content="Second done.")],
]),
)
agent_ref = agent
@ -291,17 +302,16 @@ async def test_approval_always_sets_tool_permission_for_subsequent_calls() -> No
@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"}'),
)
tool_call = make_todo_tool_call("call_5", 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"
),
content="Let me check your todos.", tool_calls=[tool_call]
)
],
[mock_llm_chunk(content="I encountered an error with the action.")],
]),
)
@ -320,24 +330,18 @@ async def test_tool_call_with_duplicate_todo_ids() -> None:
TodoItem(id="duplicate", content="Task 1"),
TodoItem(id="duplicate", content="Task 2"),
]
tool_call = ToolCall(
id="call_6",
function=FunctionCall(
name="todo",
tool_call = make_todo_tool_call(
"call_6",
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",
),
[mock_llm_chunk(content="Let me write todos.", tool_calls=[tool_call])],
[mock_llm_chunk(content="I couldn't write todos with duplicate IDs.")],
]),
)
@ -353,23 +357,18 @@ async def test_tool_call_with_duplicate_todo_ids() -> None:
@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",
tool_call = make_todo_tool_call(
"call_7",
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"
),
[mock_llm_chunk(content="Let me write todos.", tool_calls=[tool_call])],
[mock_llm_chunk(content="I couldn't write that many todos.")],
]),
)
@ -394,7 +393,7 @@ 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="{}")
id="call_8", index=0, function=FunctionCall(name="stub_tool", arguments="{}")
)
config = VibeConfig(
session_logging=SessionLoggingConfig(enabled=False),
@ -405,8 +404,8 @@ async def test_tool_call_can_be_interrupted(
config,
mode=AgentMode.AUTO_APPROVE,
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"),
[mock_llm_chunk(content="Let me use the tool.", tool_calls=[tool_call])],
[mock_llm_chunk(content="Tool execution completed.")],
]),
)
# no dependency injection available => monkey patch
@ -433,15 +432,11 @@ async def test_fill_missing_tool_responses_inserts_placeholders() -> None:
agent = Agent(
make_config(),
mode=AgentMode.AUTO_APPROVE,
backend=FakeBackend([mock_llm_chunk(content="ok", finish_reason="stop")]),
backend=FakeBackend(mock_llm_chunk(content="ok")),
)
tool_calls_messages = [
ToolCall(
id="tc1", function=FunctionCall(name="todo", arguments='{"action": "read"}')
),
ToolCall(
id="tc2", function=FunctionCall(name="todo", arguments='{"action": "read"}')
),
make_todo_tool_call("tc1", index=0),
make_todo_tool_call("tc2", index=1),
]
assistant_msg = LLMMessage(
role=Role.assistant, content="Calling tools...", tool_calls=tool_calls_messages
@ -473,7 +468,7 @@ async def test_ensure_assistant_after_tool_appends_understood() -> None:
agent = Agent(
make_config(),
mode=AgentMode.AUTO_APPROVE,
backend=FakeBackend([mock_llm_chunk(content="ok", finish_reason="stop")]),
backend=FakeBackend(mock_llm_chunk(content="ok")),
)
tool_msg = LLMMessage(
role=Role.tool, tool_call_id="tc_z", name="todo", content="Done"

View file

@ -34,12 +34,11 @@ def test_run_programmatic_preload_streaming_is_batched(
with mock_backend_factory(
Backend.MISTRAL,
lambda provider, **kwargs: FakeBackend([
lambda provider, **kwargs: FakeBackend(
mock_llm_chunk(
content="Decorators are wrappers that modify function behavior.",
finish_reason="stop",
content="Decorators are wrappers that modify function behavior."
)
]),
),
):
cfg = VibeConfig(
session_logging=SessionLoggingConfig(enabled=False),

View file

@ -157,7 +157,6 @@ class TestAgentSwitchMode:
return FakeBackend([
LLMChunk(
message=LLMMessage(role=Role.assistant, content="Test response"),
finish_reason="stop",
usage=LLMUsage(prompt_tokens=10, completion_tokens=5),
)
])
@ -274,7 +273,6 @@ class TestPlanModeToolRestriction:
backend = FakeBackend([
LLMChunk(
message=LLMMessage(role=Role.assistant, content="ok"),
finish_reason="stop",
usage=LLMUsage(prompt_tokens=10, completion_tokens=5),
)
])
@ -298,11 +296,12 @@ class TestPlanModeToolRestriction:
async def test_plan_mode_rejects_non_plan_tool_call(self) -> None:
tool_call = ToolCall(
id="call_1",
index=0,
function=FunctionCall(name="bash", arguments='{"command": "ls"}'),
)
backend = FakeBackend([
mock_llm_chunk(content="Let me run bash", tool_calls=[tool_call]),
mock_llm_chunk(content="Tool not available", finish_reason="stop"),
mock_llm_chunk(content="Tool not available"),
])
config = VibeConfig(

2
uv.lock generated
View file

@ -661,7 +661,7 @@ wheels = [
[[package]]
name = "mistral-vibe"
version = "1.2.1"
version = "1.2.2"
source = { editable = "." }
dependencies = [
{ name = "agent-client-protocol" },

View file

@ -3,4 +3,4 @@ from __future__ import annotations
from pathlib import Path
VIBE_ROOT = Path(__file__).parent
__version__ = "1.2.1"
__version__ = "1.2.2"

View file

@ -38,7 +38,6 @@ class EventHandler:
self.get_todos_collapsed = get_todos_collapsed
self.current_tool_call: ToolCallMessage | None = None
self.current_compact: CompactMessage | None = None
self.tool_results: list[ToolResultMessage] = []
async def handle_event(
self,
@ -121,7 +120,6 @@ class EventHandler:
)
await self.mount_callback(tool_result)
self.tool_results.append(tool_result)
self.current_tool_call = None
async def _handle_assistant_message(self, event: AssistantEvent) -> None:
@ -153,6 +151,3 @@ class EventHandler:
if self.current_compact:
self.current_compact.stop_spinning(success=False)
self.current_compact = None
def get_last_tool_result(self) -> ToolResultMessage | None:
return self.tool_results[-1] if self.tool_results else None

View file

@ -1,7 +1,6 @@
from __future__ import annotations
import asyncio
from collections import OrderedDict
from collections.abc import AsyncGenerator, Callable
from enum import StrEnum, auto
import time
@ -48,9 +47,9 @@ from vibe.core.types import (
CompactStartEvent,
LLMChunk,
LLMMessage,
LLMUsage,
Role,
SyncApprovalCallback,
ToolCall,
ToolCallEvent,
ToolResultEvent,
)
@ -141,8 +140,6 @@ class Agent:
config.effective_workdir,
)
self._last_chunk: LLMChunk | None = None
@property
def mode(self) -> AgentMode:
return self._mode
@ -267,11 +264,7 @@ class Agent:
yield event
last_message = self.messages[-1]
should_break_loop = (
last_message.role != Role.tool
and self._last_chunk is not None
and self._last_chunk.finish_reason is not None
)
should_break_loop = last_message.role != Role.tool
self._flush_new_messages()
@ -293,9 +286,7 @@ class Agent:
self.messages, self.stats, self.config, self.tool_manager
)
async def _perform_llm_turn(
self,
) -> AsyncGenerator[AssistantEvent | ToolCallEvent | ToolResultEvent]:
async def _perform_llm_turn(self) -> AsyncGenerator[BaseEvent, None]:
if self.enable_streaming:
async for event in self._stream_assistant_events():
yield event
@ -305,105 +296,40 @@ class Agent:
yield assistant_event
last_message = self.messages[-1]
last_chunk = self._last_chunk
if last_chunk is None or last_chunk.usage is None:
raise LLMResponseError("LLM response missing chunk or usage data")
parsed = self.format_handler.parse_message(last_message)
resolved = self.format_handler.resolve_tool_calls(
parsed, self.tool_manager, self.config
)
if last_chunk.usage.completion_tokens > 0 and self.stats.last_turn_duration > 0:
self.stats.tokens_per_second = (
last_chunk.usage.completion_tokens / self.stats.last_turn_duration
)
if not resolved.tool_calls and not resolved.failed_calls:
return
async for event in self._handle_tool_calls(resolved):
yield event
def _create_assistant_event(
self, content: str, chunk: LLMChunk | None
) -> AssistantEvent:
return AssistantEvent(content=content)
async def _stream_assistant_events(self) -> AsyncGenerator[AssistantEvent]:
chunks: list[LLMChunk] = []
content_buffer = ""
batched_chunk = LLMChunk(message=LLMMessage(role=Role.assistant))
chunks_with_content = 0
BATCH_SIZE = 5
async for chunk in self._chat_streaming():
chunks.append(chunk)
if chunk.message.tool_calls and chunk.finish_reason is None:
if chunk.message.content:
content_buffer += chunk.message.content
chunks_with_content += 1
if content_buffer:
yield self._create_assistant_event(content_buffer, chunk)
content_buffer = ""
chunks_with_content = 0
continue
batched_chunk += chunk
if chunk.message.content:
content_buffer += chunk.message.content
chunks_with_content += 1
if chunks_with_content >= BATCH_SIZE:
yield self._create_assistant_event(content_buffer, chunk)
content_buffer = ""
yield AssistantEvent(content=cast(str, batched_chunk.message.content))
batched_chunk = LLMChunk(message=LLMMessage(role=Role.assistant))
chunks_with_content = 0
if content_buffer:
last_chunk = chunks[-1] if chunks else None
yield self._create_assistant_event(content_buffer, last_chunk)
full_content = ""
full_tool_calls_map = OrderedDict[int, ToolCall]()
for chunk in chunks:
full_content += chunk.message.content or ""
if not chunk.message.tool_calls:
continue
for tc in chunk.message.tool_calls:
if tc.index is None:
raise LLMResponseError("Tool call chunk missing index")
if tc.index not in full_tool_calls_map:
full_tool_calls_map[tc.index] = tc
else:
new_args_str = (
full_tool_calls_map[tc.index].function.arguments or ""
) + (tc.function.arguments or "")
full_tool_calls_map[tc.index].function.arguments = new_args_str
full_tool_calls = list(full_tool_calls_map.values()) or None
last_message = LLMMessage(
role=Role.assistant, content=full_content, tool_calls=full_tool_calls
)
self.messages.append(last_message)
finish_reason = next(
(c.finish_reason for c in chunks if c.finish_reason is not None), None
)
self._last_chunk = LLMChunk(
message=last_message, usage=chunks[-1].usage, finish_reason=finish_reason
)
if batched_chunk.message.content:
yield AssistantEvent(content=batched_chunk.message.content)
async def _get_assistant_event(self) -> AssistantEvent:
llm_result = await self._chat()
if llm_result.usage is None:
raise LLMResponseError(
"Usage data missing in non-streaming completion response"
)
self._last_chunk = llm_result
assistant_msg = llm_result.message
self.messages.append(assistant_msg)
return AssistantEvent(content=assistant_msg.content or "")
return AssistantEvent(content=llm_result.message.content or "")
async def _handle_tool_calls(
self, resolved: ResolvedMessage
@ -585,7 +511,6 @@ class Agent:
try:
start_time = time.perf_counter()
async with self.backend as backend:
result = await backend.complete(
model=active_model,
@ -599,31 +524,19 @@ class Agent:
},
max_tokens=max_tokens,
)
end_time = time.perf_counter()
if result.usage is None:
raise LLMResponseError(
"Usage data missing in non-streaming completion response"
)
self.stats.last_turn_duration = end_time - start_time
self.stats.last_turn_prompt_tokens = result.usage.prompt_tokens
self.stats.last_turn_completion_tokens = result.usage.completion_tokens
self.stats.session_prompt_tokens += result.usage.prompt_tokens
self.stats.session_completion_tokens += result.usage.completion_tokens
self.stats.context_tokens = (
result.usage.prompt_tokens + result.usage.completion_tokens
)
self._update_stats(usage=result.usage, time_seconds=end_time - start_time)
processed_message = self.format_handler.process_api_response_message(
result.message
)
return LLMChunk(
message=processed_message,
usage=result.usage,
finish_reason=result.finish_reason,
)
self.messages.append(processed_message)
return LLMChunk(message=processed_message, usage=result.usage)
except Exception as e:
raise RuntimeError(
@ -642,7 +555,8 @@ class Agent:
tool_choice = self.format_handler.get_tool_choice()
try:
start_time = time.perf_counter()
last_chunk = None
usage = LLMUsage()
chunk_agg = LLMChunk(message=LLMMessage(role=Role.assistant))
async with self.backend as backend:
async for chunk in backend.complete_streaming(
model=active_model,
@ -656,38 +570,40 @@ class Agent:
},
max_tokens=max_tokens,
):
last_chunk = chunk
processed_message = (
self.format_handler.process_api_response_message(chunk.message)
)
yield LLMChunk(
message=processed_message,
usage=chunk.usage,
finish_reason=chunk.finish_reason,
processed_chunk = LLMChunk(
message=processed_message, usage=chunk.usage
)
chunk_agg += processed_chunk
usage += chunk.usage or LLMUsage()
yield processed_chunk
end_time = time.perf_counter()
if last_chunk is None:
raise LLMResponseError("Streamed completion returned no chunks")
if last_chunk.usage is None:
if chunk_agg.usage is None:
raise LLMResponseError(
"Usage data missing in final chunk of streamed completion"
)
self._update_stats(usage=usage, time_seconds=end_time - start_time)
self.stats.last_turn_duration = end_time - start_time
self.stats.last_turn_prompt_tokens = last_chunk.usage.prompt_tokens
self.stats.last_turn_completion_tokens = last_chunk.usage.completion_tokens
self.stats.session_prompt_tokens += last_chunk.usage.prompt_tokens
self.stats.session_completion_tokens += last_chunk.usage.completion_tokens
self.stats.context_tokens = (
last_chunk.usage.prompt_tokens + last_chunk.usage.completion_tokens
)
self.messages.append(chunk_agg.message)
except Exception as e:
raise RuntimeError(
f"API error from {provider.name} (model: {active_model.name}): {e}"
) from e
def _update_stats(self, usage: LLMUsage, time_seconds: float) -> None:
self.stats.last_turn_duration = time_seconds
self.stats.last_turn_prompt_tokens = usage.prompt_tokens
self.stats.last_turn_completion_tokens = usage.completion_tokens
self.stats.session_prompt_tokens += usage.prompt_tokens
self.stats.session_completion_tokens += usage.completion_tokens
self.stats.context_tokens = usage.prompt_tokens + usage.completion_tokens
if time_seconds > 0 and usage.completion_tokens > 0:
self.stats.tokens_per_second = usage.completion_tokens / time_seconds
async def _should_execute_tool(
self, tool: BaseTool, args: dict[str, Any], tool_call_id: str
) -> ToolDecision:

View file

@ -143,17 +143,13 @@ class OpenAIAdapter(APIAdapter):
message = LLMMessage.model_validate(data["choices"][0]["delta"])
else:
raise ValueError("Invalid response data")
finish_reason = data["choices"][0].get("finish_reason", None)
elif "message" in data:
message = LLMMessage.model_validate(data["message"])
finish_reason = data["choices"][0].get("finish_reason", None)
elif "delta" in data:
message = LLMMessage.model_validate(data["delta"])
finish_reason = None
else:
message = LLMMessage(role=Role.assistant, content="")
finish_reason = None
usage_data = data.get("usage") or {}
usage = LLMUsage(
@ -161,7 +157,7 @@ class OpenAIAdapter(APIAdapter):
completion_tokens=usage_data.get("completion_tokens", 0),
)
return LLMChunk(message=message, usage=usage, finish_reason=finish_reason)
return LLMChunk(message=message, usage=usage)
class GenericBackend:
@ -255,7 +251,7 @@ class GenericBackend:
provider=self._provider.name,
endpoint=url,
response=e.response,
headers=dict(e.response.headers.items()),
headers=e.response.headers,
model=model.name,
messages=messages,
temperature=temperature,
@ -320,7 +316,7 @@ class GenericBackend:
provider=self._provider.name,
endpoint=url,
response=e.response,
headers=dict(e.response.headers.items()),
headers=e.response.headers,
model=model.name,
messages=messages,
temperature=temperature,
@ -369,7 +365,11 @@ class GenericBackend:
continue
DELIM_CHAR = ":"
assert f"{DELIM_CHAR} " in line, "line should look like `key: value`"
if f"{DELIM_CHAR} " not in line:
raise ValueError(
f"Stream chunk improperly formatted. "
f"Expected `key{DELIM_CHAR} value`, received `{line}`"
)
delim_index = line.find(DELIM_CHAR)
key = line[0:delim_index]
value = line[delim_index + 2 :]
@ -404,9 +404,8 @@ class GenericBackend:
tool_choice=tool_choice,
extra_headers=extra_headers,
)
assert result.usage is not None, (
"Usage should be present in non-streaming completions"
)
if result.usage is None:
raise ValueError("Missing usage in non streaming completion")
return result.usage.prompt_tokens

View file

@ -205,7 +205,6 @@ class MistralBackend:
prompt_tokens=response.usage.prompt_tokens or 0,
completion_tokens=response.usage.completion_tokens or 0,
),
finish_reason=response.choices[0].finish_reason,
)
except mistralai.SDKError as e:
@ -213,7 +212,7 @@ class MistralBackend:
provider=self._provider.name,
endpoint=self._server_url,
response=e.raw_response,
headers=dict(e.raw_response.headers.items()),
headers=e.raw_response.headers,
model=model.name,
messages=messages,
temperature=temperature,
@ -279,7 +278,6 @@ class MistralBackend:
if chunk.data.usage
else 0,
),
finish_reason=chunk.data.choices[0].finish_reason,
)
except mistralai.SDKError as e:
@ -287,7 +285,7 @@ class MistralBackend:
provider=self._provider.name,
endpoint=self._server_url,
response=e.raw_response,
headers=dict(e.raw_response.headers.items()),
headers=e.raw_response.headers,
model=model.name,
messages=messages,
temperature=temperature,
@ -325,8 +323,7 @@ class MistralBackend:
tool_choice=tool_choice,
extra_headers=extra_headers,
)
assert result.usage is not None, (
"Usage should be present in non-streaming completions"
)
if result.usage is None:
raise ValueError("Missing usage in non streaming completion")
return result.usage.prompt_tokens

View file

@ -1,7 +1,9 @@
from __future__ import annotations
from abc import ABC
from collections import OrderedDict
from collections.abc import Awaitable, Callable
import copy
from enum import StrEnum, auto
from typing import Annotated, Any, Literal
@ -185,19 +187,75 @@ class LLMMessage(BaseModel):
"tool_call_id": getattr(v, "tool_call_id", None),
}
def __add__(self, other: LLMMessage) -> LLMMessage:
"""Careful: this is not commutative!"""
if self.role != other.role:
raise ValueError("Can't accumulate messages with different roles")
if self.name != other.name:
raise ValueError("Can't accumulate messages with different names")
if self.tool_call_id != other.tool_call_id:
raise ValueError("Can't accumulate messages with different tool_call_ids")
content = (self.content or "") + (other.content or "")
if not content:
content = None
tool_calls_map = OrderedDict[int, ToolCall]()
for tool_calls in [self.tool_calls or [], other.tool_calls or []]:
for tc in tool_calls:
if tc.index is None:
raise ValueError("Tool call chunk missing index")
if tc.index not in tool_calls_map:
tool_calls_map[tc.index] = copy.deepcopy(tc)
else:
existing_name = tool_calls_map[tc.index].function.name
new_name = tc.function.name
if existing_name and new_name and existing_name != new_name:
raise ValueError(
"Can't accumulate messages with different tool call names"
)
if new_name and not existing_name:
tool_calls_map[tc.index].function.name = new_name
new_args = (tool_calls_map[tc.index].function.arguments or "") + (
tc.function.arguments or ""
)
tool_calls_map[tc.index].function.arguments = new_args
return LLMMessage(
role=self.role,
content=content,
tool_calls=list(tool_calls_map.values()) or None,
name=self.name,
tool_call_id=self.tool_call_id,
)
class LLMUsage(BaseModel):
model_config = ConfigDict(frozen=True)
prompt_tokens: int = 0
completion_tokens: int = 0
def __add__(self, other: LLMUsage) -> LLMUsage:
return LLMUsage(
prompt_tokens=self.prompt_tokens + other.prompt_tokens,
completion_tokens=self.completion_tokens + other.completion_tokens,
)
class LLMChunk(BaseModel):
model_config = ConfigDict(frozen=True)
message: LLMMessage
finish_reason: str | None = None
usage: LLMUsage | None = None
def __add__(self, other: LLMChunk) -> LLMChunk:
if self.usage is None and other.usage is None:
new_usage = None
else:
new_usage = (self.usage or LLMUsage()) + (other.usage or LLMUsage())
return LLMChunk(message=self.message + other.message, usage=new_usage)
class BaseEvent(BaseModel, ABC):
"""Abstract base class for all agent events."""
@ -209,6 +267,13 @@ class AssistantEvent(BaseEvent):
content: str
stopped_by_middleware: bool = False
def __add__(self, other: AssistantEvent) -> AssistantEvent:
return AssistantEvent(
content=self.content + other.content,
stopped_by_middleware=self.stopped_by_middleware
or other.stopped_by_middleware,
)
class ToolCallEvent(BaseEvent):
tool_name: str