diff --git a/.github/workflows/build-and-upload.yml b/.github/workflows/build-and-upload.yml index 796c9d5..74e6055 100644 --- a/.github/workflows/build-and-upload.yml +++ b/.github/workflows/build-and-upload.yml @@ -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 diff --git a/.vscode/launch.json b/.vscode/launch.json index 97bf047..f8a7978 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -1,5 +1,5 @@ { - "version": "1.2.1", + "version": "1.2.2", "configurations": [ { "name": "ACP Server", diff --git a/CHANGELOG.md b/CHANGELOG.md index f18d340..cd014b5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/distribution/zed/extension.toml b/distribution/zed/extension.toml index ee5bb28..870a0b3 100644 --- a/distribution/zed/extension.toml +++ b/distribution/zed/extension.toml @@ -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" diff --git a/pyproject.toml b/pyproject.toml index d5bac54..27b1a6d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/tests/acp/test_acp.py b/tests/acp/test_acp.py index e758854..3bbe9ed 100644 --- a/tests/acp/test_acp.py +++ b/tests/acp/test_acp.py @@ -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) diff --git a/tests/acp/test_content.py b/tests/acp/test_content.py index 8fee387..6705455 100644 --- a/tests/acp/test_content.py +++ b/tests/acp/test_content.py @@ -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), - ) - ] + LLMChunk( + message=LLMMessage(role=Role.assistant, content="Hi"), + usage=LLMUsage(prompt_tokens=1, completion_tokens=1), + ) ) return backend diff --git a/tests/acp/test_initialize.py b/tests/acp/test_initialize.py index f690dc7..08a8b10 100644 --- a/tests/acp/test_initialize.py +++ b/tests/acp/test_initialize.py @@ -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 diff --git a/tests/acp/test_multi_session.py b/tests/acp/test_multi_session.py index 7aabc20..1e07929 100644 --- a/tests/acp/test_multi_session.py +++ b/tests/acp/test_multi_session.py @@ -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(): diff --git a/tests/acp/test_new_session.py b/tests/acp/test_new_session.py index fa4c4b6..c945a94 100644 --- a/tests/acp/test_new_session.py +++ b/tests/acp/test_new_session.py @@ -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), - ) - ] + LLMChunk( + message=LLMMessage(role=Role.assistant, content="Hi"), + usage=LLMUsage(prompt_tokens=1, completion_tokens=1), + ) ) return backend diff --git a/tests/acp/test_set_mode.py b/tests/acp/test_set_mode.py index 52681c4..91b9376 100644 --- a/tests/acp/test_set_mode.py +++ b/tests/acp/test_set_mode.py @@ -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), - ) - ] + LLMChunk( + message=LLMMessage(role=Role.assistant, content="Hi"), + usage=LLMUsage(prompt_tokens=1, completion_tokens=1), + ) ) return backend diff --git a/tests/acp/test_set_model.py b/tests/acp/test_set_model.py index fde444f..8259faf 100644 --- a/tests/acp/test_set_model.py +++ b/tests/acp/test_set_model.py @@ -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), - ) - ] + LLMChunk( + message=LLMMessage(role=Role.assistant, content="Hi"), + usage=LLMUsage(prompt_tokens=1, completion_tokens=1), + ) ) return backend diff --git a/tests/backend/data/fireworks.py b/tests/backend/data/fireworks.py index 5ecb6eb..a69aed2 100644 --- a/tests/backend/data/fireworks.py +++ b/tests/backend/data/fireworks.py @@ -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}}, ], ) ] diff --git a/tests/backend/data/mistral.py b/tests/backend/data/mistral.py index 0bcdec1..6e6504c 100644 --- a/tests/backend/data/mistral.py +++ b/tests/backend/data/mistral.py @@ -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}}, ], ) ] diff --git a/tests/backend/test_backend.py b/tests/backend/test_backend.py index a29c89b..b508b5a 100644 --- a/tests/backend/test_backend.py +++ b/tests/backend/test_backend.py @@ -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 diff --git a/tests/mock/utils.py b/tests/mock/utils.py index cfb4152..b484a9b 100644 --- a/tests/mock/utils.py +++ b/tests/mock/utils.py @@ -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, ) diff --git a/tests/snapshots/test_ui_snapshot_basic_conversation.py b/tests/snapshots/test_ui_snapshot_basic_conversation.py index bb9e60d..99a5d51 100644 --- a/tests/snapshots/test_ui_snapshot_basic_conversation.py +++ b/tests/snapshots/test_ui_snapshot_basic_conversation.py @@ -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, - ) - ] + 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( diff --git a/tests/stubs/fake_backend.py b/tests/stubs/fake_backend.py index f788d20..26b47dc 100644 --- a/tests/stubs/fake_backend.py +++ b/tests/stubs/fake_backend.py @@ -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, diff --git a/tests/test_agent_auto_compact.py b/tests/test_agent_auto_compact.py index 5403636..9efc674 100644 --- a/tests/test_agent_auto_compact.py +++ b/tests/test_agent_auto_compact.py @@ -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=""), - mock_llm_chunk(content=""), + [mock_llm_chunk(content="")], + [mock_llm_chunk(content="")], ]) cfg = VibeConfig( session_logging=SessionLoggingConfig(enabled=False), auto_compact_threshold=1 diff --git a/tests/test_agent_backend.py b/tests/test_agent_backend.py index ec04917..7460185 100644 --- a/tests/test_agent_backend.py +++ b/tests/test_agent_backend.py @@ -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) diff --git a/tests/test_agent_observer_streaming.py b/tests/test_agent_observer_streaming.py index a8ec292..ffd1e6d 100644 --- a/tests/test_agent_observer_streaming.py +++ b/tests/test_agent_observer_streaming.py @@ -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="Checking your todos."), + mock_llm_chunk(content="", tool_calls=[todo_tool_call]), + ], + [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( diff --git a/tests/test_agent_stats.py b/tests/test_agent_stats.py index fcdf891..c4cfcd3 100644 --- a/tests/test_agent_stats.py +++ b/tests/test_agent_stats.py @@ -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="", finish_reason="stop"), + [mock_llm_chunk(content="First response")], + [mock_llm_chunk(content="")], ]) 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="", finish_reason="stop"), + [mock_llm_chunk(content="Long response " * 100)], + [mock_llm_chunk(content="")], ]) 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", - function=FunctionCall( - name="todo", arguments='{"action": "read"}' - ), - ) - ], - ), - mock_llm_chunk(content="Done", finish_reason="stop"), - mock_llm_chunk(content="", finish_reason="stop"), + [ + mock_llm_chunk( + content="Using tool", + tool_calls=[ + ToolCall( + id="tc1", + index=0, + function=FunctionCall( + name="todo", arguments='{"action": "read"}' + ), + ) + ], + ), + mock_llm_chunk(content=" todo"), + ], + [mock_llm_chunk(content="")], ]) 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="", finish_reason="stop"), + [mock_llm_chunk(content="Long response " * 100)], + [mock_llm_chunk(content="")], ]) 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="", finish_reason="stop"), - mock_llm_chunk(content="", finish_reason="stop"), + [mock_llm_chunk(content="")], + [mock_llm_chunk(content="")], ]) 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="", finish_reason="stop"), - mock_llm_chunk(content="After reload", finish_reason="stop"), + [mock_llm_chunk(content="Initial response")], + [mock_llm_chunk(content="")], + [mock_llm_chunk(content="After reload")], ]) agent = Agent(make_config(), backend=backend) diff --git a/tests/test_agent_tool_call.py b/tests/test_agent_tool_call.py index 5199f0b..5cda432 100644 --- a/tests/test_agent_tool_call.py +++ b/tests/test_agent_tool_call.py @@ -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="Let me check your todos.", + tool_calls=[make_todo_tool_call("call_2")], + ) + ], + [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="Let me check your todos.", + tool_calls=[make_todo_tool_call("call_3")], + ) + ], + [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="Let me check your todos.", + tool_calls=[make_todo_tool_call("call_4")], + ) + ], + [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="Let me check your todos.", + tool_calls=[make_todo_tool_call("call_never")], + ) + ], + [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"), - mock_llm_chunk( - content="Second check.", tool_calls=[make_todo_tool_call("call_second")] - ), - mock_llm_chunk(content="Second done.", finish_reason="stop"), + [ + mock_llm_chunk( + 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.")], ]), ) 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" - ), + [ + mock_llm_chunk( + 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", - arguments=json.dumps({ - "action": "write", - "todos": [t.model_dump() for t in duplicate_todos], - }), - ), + 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", - arguments=json.dumps({ - "action": "write", - "todos": [t.model_dump() for t in many_todos], - }), - ), + 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" diff --git a/tests/test_cli_programmatic_preload.py b/tests/test_cli_programmatic_preload.py index 73279be..99d6bfe 100644 --- a/tests/test_cli_programmatic_preload.py +++ b/tests/test_cli_programmatic_preload.py @@ -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), diff --git a/tests/test_modes.py b/tests/test_modes.py index 1483c07..7ccd086 100644 --- a/tests/test_modes.py +++ b/tests/test_modes.py @@ -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( diff --git a/uv.lock b/uv.lock index fd1321c..b6b26b5 100644 --- a/uv.lock +++ b/uv.lock @@ -661,7 +661,7 @@ wheels = [ [[package]] name = "mistral-vibe" -version = "1.2.1" +version = "1.2.2" source = { editable = "." } dependencies = [ { name = "agent-client-protocol" }, diff --git a/vibe/__init__.py b/vibe/__init__.py index 4d03dae..9c7139b 100644 --- a/vibe/__init__.py +++ b/vibe/__init__.py @@ -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" diff --git a/vibe/cli/textual_ui/handlers/event_handler.py b/vibe/cli/textual_ui/handlers/event_handler.py index afd792d..1646da9 100644 --- a/vibe/cli/textual_ui/handlers/event_handler.py +++ b/vibe/cli/textual_ui/handlers/event_handler.py @@ -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 diff --git a/vibe/core/agent.py b/vibe/core/agent.py index 6319c88..a7939d1 100644 --- a/vibe/core/agent.py +++ b/vibe/core/agent.py @@ -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 = "" - chunks_with_content = 0 + if chunks_with_content >= BATCH_SIZE: + 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: diff --git a/vibe/core/llm/backend/generic.py b/vibe/core/llm/backend/generic.py index 4abe4f0..ddfc476 100644 --- a/vibe/core/llm/backend/generic.py +++ b/vibe/core/llm/backend/generic.py @@ -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 diff --git a/vibe/core/llm/backend/mistral.py b/vibe/core/llm/backend/mistral.py index 683d336..97996fb 100644 --- a/vibe/core/llm/backend/mistral.py +++ b/vibe/core/llm/backend/mistral.py @@ -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 diff --git a/vibe/core/types.py b/vibe/core/types.py index dc74ea7..fa82417 100644 --- a/vibe/core/types.py +++ b/vibe/core/types.py @@ -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