v1.2.2
Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai> Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai>
This commit is contained in:
parent
402e898f39
commit
2e1e15120d
32 changed files with 391 additions and 549 deletions
9
.github/workflows/build-and-upload.yml
vendored
9
.github/workflows/build-and-upload.yml
vendored
|
|
@ -108,12 +108,17 @@ jobs:
|
||||||
with:
|
with:
|
||||||
path: artifacts
|
path: artifacts
|
||||||
|
|
||||||
- name: Zip artifacts like GitHub UI
|
- name: Zip artifacts
|
||||||
run: |
|
run: |
|
||||||
mkdir release-assets
|
mkdir release-assets
|
||||||
for dir in artifacts/*; do
|
for dir in artifacts/*; do
|
||||||
name=$(basename "$dir")
|
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
|
done
|
||||||
|
|
||||||
- name: Attach binaries to release
|
- name: Attach binaries to release
|
||||||
|
|
|
||||||
2
.vscode/launch.json
vendored
2
.vscode/launch.json
vendored
|
|
@ -1,5 +1,5 @@
|
||||||
{
|
{
|
||||||
"version": "1.2.1",
|
"version": "1.2.2",
|
||||||
"configurations": [
|
"configurations": [
|
||||||
{
|
{
|
||||||
"name": "ACP Server",
|
"name": "ACP Server",
|
||||||
|
|
|
||||||
|
|
@ -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/),
|
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).
|
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
|
## [1.2.1] - 2025-12-18
|
||||||
|
|
||||||
### Fixed
|
### Fixed
|
||||||
|
|
|
||||||
|
|
@ -1,7 +1,7 @@
|
||||||
id = "mistral-vibe"
|
id = "mistral-vibe"
|
||||||
name = "Mistral Vibe"
|
name = "Mistral Vibe"
|
||||||
description = "Mistral's open-source coding assistant"
|
description = "Mistral's open-source coding assistant"
|
||||||
version = "1.2.1"
|
version = "1.2.2"
|
||||||
schema_version = 1
|
schema_version = 1
|
||||||
authors = ["Mistral AI"]
|
authors = ["Mistral AI"]
|
||||||
repository = "https://github.com/mistralai/mistral-vibe"
|
repository = "https://github.com/mistralai/mistral-vibe"
|
||||||
|
|
@ -11,25 +11,25 @@ name = "Mistral Vibe"
|
||||||
icon = "./icons/mistral_vibe.svg"
|
icon = "./icons/mistral_vibe.svg"
|
||||||
|
|
||||||
[agent_servers.mistral-vibe.targets.darwin-aarch64]
|
[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"
|
cmd = "./vibe-acp"
|
||||||
|
|
||||||
[agent_servers.mistral-vibe.targets.darwin-x86_64]
|
[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"
|
cmd = "./vibe-acp"
|
||||||
|
|
||||||
[agent_servers.mistral-vibe.targets.linux-aarch64]
|
[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"
|
cmd = "./vibe-acp"
|
||||||
|
|
||||||
[agent_servers.mistral-vibe.targets.linux-x86_64]
|
[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"
|
cmd = "./vibe-acp"
|
||||||
|
|
||||||
[agent_servers.mistral-vibe.targets.windows-aarch64]
|
[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"
|
cmd = "./vibe-acp.exe"
|
||||||
|
|
||||||
[agent_servers.mistral-vibe.targets.windows-x86_64]
|
[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"
|
cmd = "./vibe-acp.exe"
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,6 @@
|
||||||
[project]
|
[project]
|
||||||
name = "mistral-vibe"
|
name = "mistral-vibe"
|
||||||
version = "1.2.1"
|
version = "1.2.2"
|
||||||
description = "Minimal CLI coding agent by Mistral"
|
description = "Minimal CLI coding agent by Mistral"
|
||||||
readme = "README.md"
|
readme = "README.md"
|
||||||
requires-python = ">=3.12"
|
requires-python = ">=3.12"
|
||||||
|
|
|
||||||
|
|
@ -426,7 +426,7 @@ class TestSessionManagement:
|
||||||
class TestSessionUpdates:
|
class TestSessionUpdates:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_agent_message_chunk_structure(self, vibe_home_dir: Path) -> None:
|
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(
|
async for process in get_acp_agent_process(
|
||||||
mock_env=mock_env, vibe_home=vibe_home_dir
|
mock_env=mock_env, vibe_home=vibe_home_dir
|
||||||
):
|
):
|
||||||
|
|
@ -465,7 +465,6 @@ class TestSessionUpdates:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_tool_call_update_structure(self, vibe_home_dir: Path) -> None:
|
async def test_tool_call_update_structure(self, vibe_home_dir: Path) -> None:
|
||||||
mock_env = get_mocking_env([
|
mock_env = get_mocking_env([
|
||||||
mock_llm_chunk(content="Hey"),
|
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
tool_calls=[
|
tool_calls=[
|
||||||
ToolCall(
|
ToolCall(
|
||||||
|
|
@ -475,14 +474,9 @@ class TestSessionUpdates:
|
||||||
type="function",
|
type="function",
|
||||||
index=0,
|
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(
|
async for process in get_acp_agent_process(
|
||||||
mock_env=mock_env, vibe_home=vibe_home_dir
|
mock_env=mock_env, vibe_home=vibe_home_dir
|
||||||
|
|
@ -567,7 +561,6 @@ class TestToolCallStructure:
|
||||||
self, vibe_home_grep_ask: Path
|
self, vibe_home_grep_ask: Path
|
||||||
) -> None:
|
) -> None:
|
||||||
custom_results = [
|
custom_results = [
|
||||||
mock_llm_chunk(content="Hey"),
|
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
tool_calls=[
|
tool_calls=[
|
||||||
ToolCall(
|
ToolCall(
|
||||||
|
|
@ -578,10 +571,8 @@ class TestToolCallStructure:
|
||||||
type="function",
|
type="function",
|
||||||
index=0,
|
index=0,
|
||||||
)
|
)
|
||||||
],
|
]
|
||||||
name="grep",
|
)
|
||||||
finish_reason="tool_calls",
|
|
||||||
),
|
|
||||||
]
|
]
|
||||||
mock_env = get_mocking_env(custom_results)
|
mock_env = get_mocking_env(custom_results)
|
||||||
async for process in get_acp_agent_process(
|
async for process in get_acp_agent_process(
|
||||||
|
|
@ -628,7 +619,6 @@ class TestToolCallStructure:
|
||||||
self, vibe_home_grep_ask: Path
|
self, vibe_home_grep_ask: Path
|
||||||
) -> None:
|
) -> None:
|
||||||
custom_results = [
|
custom_results = [
|
||||||
mock_llm_chunk(content="Hey"),
|
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
tool_calls=[
|
tool_calls=[
|
||||||
ToolCall(
|
ToolCall(
|
||||||
|
|
@ -639,13 +629,10 @@ class TestToolCallStructure:
|
||||||
type="function",
|
type="function",
|
||||||
index=0,
|
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)
|
mock_env = get_mocking_env(custom_results)
|
||||||
async for process in get_acp_agent_process(
|
async for process in get_acp_agent_process(
|
||||||
|
|
@ -692,7 +679,6 @@ class TestToolCallStructure:
|
||||||
self, vibe_home_grep_ask: Path
|
self, vibe_home_grep_ask: Path
|
||||||
) -> None:
|
) -> None:
|
||||||
custom_results = [
|
custom_results = [
|
||||||
mock_llm_chunk(content="Hey"),
|
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
tool_calls=[
|
tool_calls=[
|
||||||
ToolCall(
|
ToolCall(
|
||||||
|
|
@ -703,14 +689,11 @@ class TestToolCallStructure:
|
||||||
type="function",
|
type="function",
|
||||||
index=0,
|
index=0,
|
||||||
)
|
)
|
||||||
],
|
]
|
||||||
name="grep",
|
|
||||||
finish_reason="tool_calls",
|
|
||||||
),
|
),
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="The search for 'auth' has not been performed, "
|
content="The search for 'auth' has not been performed, "
|
||||||
"because you rejected the permission request",
|
"because you rejected the permission request"
|
||||||
finish_reason="stop",
|
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
mock_env = get_mocking_env(custom_results)
|
mock_env = get_mocking_env(custom_results)
|
||||||
|
|
@ -759,7 +742,6 @@ class TestToolCallStructure:
|
||||||
self, vibe_home_grep_ask: Path
|
self, vibe_home_grep_ask: Path
|
||||||
) -> None:
|
) -> None:
|
||||||
custom_results = [
|
custom_results = [
|
||||||
mock_llm_chunk(content="Hey"),
|
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
tool_calls=[
|
tool_calls=[
|
||||||
ToolCall(
|
ToolCall(
|
||||||
|
|
@ -770,13 +752,10 @@ class TestToolCallStructure:
|
||||||
type="function",
|
type="function",
|
||||||
index=0,
|
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)
|
mock_env = get_mocking_env(custom_results)
|
||||||
async for process in get_acp_agent_process(
|
async for process in get_acp_agent_process(
|
||||||
|
|
@ -820,7 +799,6 @@ class TestToolCallStructure:
|
||||||
self, vibe_home_grep_ask: Path
|
self, vibe_home_grep_ask: Path
|
||||||
) -> None:
|
) -> None:
|
||||||
custom_results = [
|
custom_results = [
|
||||||
mock_llm_chunk(content="Hey"),
|
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
tool_calls=[
|
tool_calls=[
|
||||||
ToolCall(
|
ToolCall(
|
||||||
|
|
@ -831,14 +809,11 @@ class TestToolCallStructure:
|
||||||
type="function",
|
type="function",
|
||||||
index=0,
|
index=0,
|
||||||
)
|
)
|
||||||
],
|
]
|
||||||
name="grep",
|
|
||||||
finish_reason="tool_calls",
|
|
||||||
),
|
),
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="The search for 'auth' has failed "
|
content="The search for 'auth' has failed "
|
||||||
"because the path does not exist",
|
"because the path does not exist"
|
||||||
finish_reason="stop",
|
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
mock_env = get_mocking_env(custom_results)
|
mock_env = get_mocking_env(custom_results)
|
||||||
|
|
@ -894,7 +869,6 @@ class TestCancellationStructure:
|
||||||
self, vibe_home_dir: Path
|
self, vibe_home_dir: Path
|
||||||
) -> None:
|
) -> None:
|
||||||
custom_results = [
|
custom_results = [
|
||||||
mock_llm_chunk(content="Hey"),
|
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
tool_calls=[
|
tool_calls=[
|
||||||
ToolCall(
|
ToolCall(
|
||||||
|
|
@ -906,14 +880,11 @@ class TestCancellationStructure:
|
||||||
type="function",
|
type="function",
|
||||||
index=0,
|
index=0,
|
||||||
)
|
)
|
||||||
],
|
]
|
||||||
name="write_file",
|
|
||||||
finish_reason="tool_calls",
|
|
||||||
),
|
),
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="The file test.txt has not been created, "
|
content="The file test.txt has not been created, "
|
||||||
"because you cancelled the permission request",
|
"because you cancelled the permission request"
|
||||||
finish_reason="stop",
|
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
mock_env = get_mocking_env(custom_results)
|
mock_env = get_mocking_env(custom_results)
|
||||||
|
|
|
||||||
|
|
@ -22,13 +22,10 @@ from vibe.core.types import LLMChunk, LLMMessage, LLMUsage, Role
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def backend() -> FakeBackend:
|
def backend() -> FakeBackend:
|
||||||
backend = FakeBackend(
|
backend = FakeBackend(
|
||||||
results=[
|
|
||||||
LLMChunk(
|
LLMChunk(
|
||||||
message=LLMMessage(role=Role.assistant, content="Hi"),
|
message=LLMMessage(role=Role.assistant, content="Hi"),
|
||||||
finish_reason="end_turn",
|
|
||||||
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
|
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
|
||||||
)
|
)
|
||||||
]
|
|
||||||
)
|
)
|
||||||
return backend
|
return backend
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -41,7 +41,7 @@ class TestACPInitialize:
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
assert response.agentInfo == Implementation(
|
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 == []
|
assert response.authMethods == []
|
||||||
|
|
@ -63,7 +63,7 @@ class TestACPInitialize:
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
assert response.agentInfo == Implementation(
|
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
|
assert response.authMethods is not None
|
||||||
|
|
|
||||||
|
|
@ -116,9 +116,9 @@ class TestMultiSessionCore:
|
||||||
)
|
)
|
||||||
session2 = acp_agent.sessions[session2_response.sessionId]
|
session2 = acp_agent.sessions[session2_response.sessionId]
|
||||||
|
|
||||||
backend._chunks = [
|
backend._streams = [
|
||||||
mock_llm_chunk(content="Response 1", finish_reason="stop"),
|
[mock_llm_chunk(content="Response 1")],
|
||||||
mock_llm_chunk(content="Response 2", finish_reason="stop"),
|
[mock_llm_chunk(content="Response 2")],
|
||||||
]
|
]
|
||||||
|
|
||||||
async def run_session1():
|
async def run_session1():
|
||||||
|
|
|
||||||
|
|
@ -18,13 +18,10 @@ from vibe.core.types import LLMChunk, LLMMessage, LLMUsage, Role
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def backend() -> FakeBackend:
|
def backend() -> FakeBackend:
|
||||||
backend = FakeBackend(
|
backend = FakeBackend(
|
||||||
results=[
|
|
||||||
LLMChunk(
|
LLMChunk(
|
||||||
message=LLMMessage(role=Role.assistant, content="Hi"),
|
message=LLMMessage(role=Role.assistant, content="Hi"),
|
||||||
finish_reason="end_turn",
|
|
||||||
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
|
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
|
||||||
)
|
)
|
||||||
]
|
|
||||||
)
|
)
|
||||||
return backend
|
return backend
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -17,13 +17,10 @@ from vibe.core.types import LLMChunk, LLMMessage, LLMUsage, Role
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def backend() -> FakeBackend:
|
def backend() -> FakeBackend:
|
||||||
backend = FakeBackend(
|
backend = FakeBackend(
|
||||||
results=[
|
|
||||||
LLMChunk(
|
LLMChunk(
|
||||||
message=LLMMessage(role=Role.assistant, content="Hi"),
|
message=LLMMessage(role=Role.assistant, content="Hi"),
|
||||||
finish_reason="end_turn",
|
|
||||||
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
|
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
|
||||||
)
|
)
|
||||||
]
|
|
||||||
)
|
)
|
||||||
return backend
|
return backend
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -17,13 +17,10 @@ from vibe.core.types import LLMChunk, LLMMessage, LLMUsage, Role
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def backend() -> FakeBackend:
|
def backend() -> FakeBackend:
|
||||||
backend = FakeBackend(
|
backend = FakeBackend(
|
||||||
results=[
|
|
||||||
LLMChunk(
|
LLMChunk(
|
||||||
message=LLMMessage(role=Role.assistant, content="Hi"),
|
message=LLMMessage(role=Role.assistant, content="Hi"),
|
||||||
finish_reason="end_turn",
|
|
||||||
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
|
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
|
||||||
)
|
)
|
||||||
]
|
|
||||||
)
|
)
|
||||||
return backend
|
return backend
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -30,7 +30,6 @@ SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"message": "Some content",
|
"message": "Some content",
|
||||||
"finish_reason": "stop",
|
|
||||||
"usage": {
|
"usage": {
|
||||||
"prompt_tokens": 100,
|
"prompt_tokens": 100,
|
||||||
"total_tokens": 300,
|
"total_tokens": 300,
|
||||||
|
|
@ -78,7 +77,6 @@ TOOL_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"message": "",
|
"message": "",
|
||||||
"finish_reason": "tool_calls",
|
|
||||||
"tool_calls": [
|
"tool_calls": [
|
||||||
{
|
{
|
||||||
"name": "some_tool",
|
"name": "some_tool",
|
||||||
|
|
@ -102,26 +100,13 @@ STREAMED_SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultDat
|
||||||
rb"data: [DONE]",
|
rb"data: [DONE]",
|
||||||
],
|
],
|
||||||
[
|
[
|
||||||
{
|
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
|
||||||
"message": "",
|
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
|
||||||
"finish_reason": None,
|
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"message": "",
|
|
||||||
"finish_reason": None,
|
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
|
||||||
},
|
|
||||||
{
|
{
|
||||||
"message": "Some content",
|
"message": "Some content",
|
||||||
"finish_reason": None,
|
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||||
},
|
},
|
||||||
{
|
{"message": "", "usage": {"prompt_tokens": 100, "completion_tokens": 200}},
|
||||||
"message": "",
|
|
||||||
"finish_reason": "stop",
|
|
||||||
"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]",
|
rb"data: [DONE]",
|
||||||
],
|
],
|
||||||
[
|
[
|
||||||
{
|
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
|
||||||
"message": "",
|
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
|
||||||
"finish_reason": None,
|
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
|
||||||
},
|
|
||||||
{
|
|
||||||
"message": "",
|
|
||||||
"finish_reason": None,
|
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
|
||||||
},
|
|
||||||
{
|
{
|
||||||
"message": "Some content",
|
"message": "Some content",
|
||||||
"finish_reason": None,
|
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"message": "",
|
"message": "",
|
||||||
"finish_reason": None,
|
|
||||||
"tool_calls": [{"name": "some_tool", "arguments": None, "index": 0}],
|
"tool_calls": [{"name": "some_tool", "arguments": None, "index": 0}],
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"message": "",
|
"message": "",
|
||||||
"finish_reason": None,
|
|
||||||
"tool_calls": [
|
"tool_calls": [
|
||||||
{
|
{
|
||||||
"name": None,
|
"name": None,
|
||||||
|
|
@ -173,11 +147,7 @@ STREAMED_TOOL_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]
|
||||||
],
|
],
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||||
},
|
},
|
||||||
{
|
{"message": "", "usage": {"prompt_tokens": 100, "completion_tokens": 200}},
|
||||||
"message": "",
|
|
||||||
"finish_reason": "tool_calls",
|
|
||||||
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
|
|
||||||
},
|
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -29,7 +29,6 @@ SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"message": "Some content",
|
"message": "Some content",
|
||||||
"finish_reason": "stop",
|
|
||||||
"usage": {
|
"usage": {
|
||||||
"prompt_tokens": 100,
|
"prompt_tokens": 100,
|
||||||
"total_tokens": 300,
|
"total_tokens": 300,
|
||||||
|
|
@ -75,7 +74,6 @@ TOOL_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"message": "Some content",
|
"message": "Some content",
|
||||||
"finish_reason": "tool_calls",
|
|
||||||
"tool_calls": [
|
"tool_calls": [
|
||||||
{
|
{
|
||||||
"name": "some_tool",
|
"name": "some_tool",
|
||||||
|
|
@ -98,21 +96,12 @@ STREAMED_SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultDat
|
||||||
rb"data: [DONE]",
|
rb"data: [DONE]",
|
||||||
],
|
],
|
||||||
[
|
[
|
||||||
{
|
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
|
||||||
"message": "",
|
|
||||||
"finish_reason": None,
|
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
|
||||||
},
|
|
||||||
{
|
{
|
||||||
"message": "Some content",
|
"message": "Some content",
|
||||||
"finish_reason": None,
|
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||||
},
|
},
|
||||||
{
|
{"message": "", "usage": {"prompt_tokens": 100, "completion_tokens": 200}},
|
||||||
"message": "",
|
|
||||||
"finish_reason": "stop",
|
|
||||||
"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]",
|
rb"data: [DONE]",
|
||||||
],
|
],
|
||||||
[
|
[
|
||||||
{
|
{"message": "", "usage": {"prompt_tokens": 0, "completion_tokens": 0}},
|
||||||
"message": "",
|
|
||||||
"finish_reason": None,
|
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
|
||||||
},
|
|
||||||
{
|
{
|
||||||
"message": "Some content",
|
"message": "Some content",
|
||||||
"finish_reason": None,
|
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"message": "",
|
"message": "",
|
||||||
"finish_reason": None,
|
|
||||||
"tool_calls": [{"name": "some_tool", "arguments": "", "index": 0}],
|
"tool_calls": [{"name": "some_tool", "arguments": "", "index": 0}],
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"message": "",
|
"message": "",
|
||||||
"finish_reason": None,
|
|
||||||
"tool_calls": [
|
"tool_calls": [
|
||||||
{"name": "", "arguments": '{"some_argument": ', "index": 0}
|
{"name": "", "arguments": '{"some_argument": ', "index": 0}
|
||||||
],
|
],
|
||||||
|
|
@ -157,17 +139,12 @@ STREAMED_TOOL_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"message": "",
|
"message": "",
|
||||||
"finish_reason": None,
|
|
||||||
"tool_calls": [
|
"tool_calls": [
|
||||||
{"name": "", "arguments": '"some_argument_value"}', "index": 0}
|
{"name": "", "arguments": '"some_argument_value"}', "index": 0}
|
||||||
],
|
],
|
||||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||||
},
|
},
|
||||||
{
|
{"message": "", "usage": {"prompt_tokens": 100, "completion_tokens": 200}},
|
||||||
"message": "",
|
|
||||||
"finish_reason": "tool_calls",
|
|
||||||
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
|
|
||||||
},
|
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -87,7 +87,6 @@ class TestBackend:
|
||||||
)
|
)
|
||||||
|
|
||||||
assert result.message.content == result_data["message"]
|
assert result.message.content == result_data["message"]
|
||||||
assert result.finish_reason == result_data["finish_reason"]
|
|
||||||
assert result.usage is not None
|
assert result.usage is not None
|
||||||
assert (
|
assert (
|
||||||
result.usage.prompt_tokens == result_data["usage"]["prompt_tokens"]
|
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):
|
for result, expected_result in zip(results, result_data, strict=True):
|
||||||
assert result.message.content == expected_result["message"]
|
assert result.message.content == expected_result["message"]
|
||||||
assert result.finish_reason == expected_result["finish_reason"]
|
|
||||||
assert result.usage is not None
|
assert result.usage is not None
|
||||||
assert (
|
assert (
|
||||||
result.usage.prompt_tokens
|
result.usage.prompt_tokens
|
||||||
|
|
|
||||||
|
|
@ -13,7 +13,6 @@ def mock_llm_chunk(
|
||||||
tool_calls: list[ToolCall] | None = None,
|
tool_calls: list[ToolCall] | None = None,
|
||||||
name: str | None = None,
|
name: str | None = None,
|
||||||
tool_call_id: str | None = None,
|
tool_call_id: str | None = None,
|
||||||
finish_reason: str | None = None,
|
|
||||||
prompt_tokens: int = 10,
|
prompt_tokens: int = 10,
|
||||||
completion_tokens: int = 5,
|
completion_tokens: int = 5,
|
||||||
) -> LLMChunk:
|
) -> LLMChunk:
|
||||||
|
|
@ -29,7 +28,6 @@ def mock_llm_chunk(
|
||||||
usage=LLMUsage(
|
usage=LLMUsage(
|
||||||
prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
|
prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
|
||||||
),
|
),
|
||||||
finish_reason=finish_reason,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -13,13 +13,11 @@ class SnapshotTestAppWithConversation(BaseSnapshotTestApp):
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
config = default_config()
|
config = default_config()
|
||||||
fake_backend = FakeBackend(
|
fake_backend = FakeBackend(
|
||||||
results=[
|
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="I'm the Vibe agent and I'm ready to help.",
|
content="I'm the Vibe agent and I'm ready to help.",
|
||||||
prompt_tokens=10_000,
|
prompt_tokens=10_000,
|
||||||
completion_tokens=2_500,
|
completion_tokens=2_500,
|
||||||
)
|
)
|
||||||
]
|
|
||||||
)
|
)
|
||||||
super().__init__(config=config)
|
super().__init__(config=config)
|
||||||
self.agent = Agent(
|
self.agent = Agent(
|
||||||
|
|
|
||||||
|
|
@ -1,9 +1,10 @@
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from collections.abc import AsyncGenerator, Callable, Iterable
|
from collections.abc import AsyncGenerator, Callable, Iterable
|
||||||
|
from typing import cast
|
||||||
|
|
||||||
from tests.mock.utils import mock_llm_chunk
|
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:
|
class FakeBackend:
|
||||||
|
|
@ -15,18 +16,47 @@ class FakeBackend:
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
results: Iterable[LLMChunk] | None = None,
|
chunks: LLMChunk
|
||||||
|
| Iterable[LLMChunk]
|
||||||
|
| Iterable[Iterable[LLMChunk]]
|
||||||
|
| None = None,
|
||||||
*,
|
*,
|
||||||
token_counter: Callable[[list[LLMMessage]], int] | None = None,
|
token_counter: Callable[[list[LLMMessage]], int] | None = None,
|
||||||
exception_to_raise: Exception | None = None,
|
exception_to_raise: Exception | None = 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_messages: list[list[LLMMessage]] = []
|
||||||
self._requests_extra_headers: list[dict[str, str] | None] = []
|
self._requests_extra_headers: list[dict[str, str] | None] = []
|
||||||
self._count_tokens_calls: list[list[LLMMessage]] = []
|
self._count_tokens_calls: list[list[LLMMessage]] = []
|
||||||
self._token_counter = token_counter or self._default_token_counter
|
self._token_counter = token_counter or self._default_token_counter
|
||||||
self._exception_to_raise = exception_to_raise
|
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
|
@property
|
||||||
def requests_messages(self) -> list[list[LLMMessage]]:
|
def requests_messages(self) -> list[list[LLMMessage]]:
|
||||||
return self._requests_messages
|
return self._requests_messages
|
||||||
|
|
@ -61,12 +91,15 @@ class FakeBackend:
|
||||||
|
|
||||||
self._requests_messages.append(messages)
|
self._requests_messages.append(messages)
|
||||||
self._requests_extra_headers.append(extra_headers)
|
self._requests_extra_headers.append(extra_headers)
|
||||||
if self._chunks:
|
|
||||||
chunk = self._chunks.pop(0)
|
if self._streams:
|
||||||
if not self._chunks:
|
stream = self._streams.pop(0)
|
||||||
chunk = chunk.model_copy(update={"finish_reason": "stop"})
|
chunk_agg = LLMChunk(message=LLMMessage(role=Role.assistant))
|
||||||
return chunk
|
for chunk in stream:
|
||||||
return mock_llm_chunk(content="", finish_reason="stop")
|
chunk_agg += chunk
|
||||||
|
return chunk_agg
|
||||||
|
|
||||||
|
return mock_llm_chunk(content="")
|
||||||
|
|
||||||
async def complete_streaming(
|
async def complete_streaming(
|
||||||
self,
|
self,
|
||||||
|
|
@ -84,22 +117,13 @@ class FakeBackend:
|
||||||
|
|
||||||
self._requests_messages.append(messages)
|
self._requests_messages.append(messages)
|
||||||
self._requests_extra_headers.append(extra_headers)
|
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
|
yield chunk
|
||||||
if has_final_chunk:
|
|
||||||
break
|
|
||||||
|
|
||||||
if not has_final_chunk:
|
|
||||||
yield mock_llm_chunk(content="", finish_reason="stop")
|
|
||||||
|
|
||||||
async def count_tokens(
|
async def count_tokens(
|
||||||
self,
|
self,
|
||||||
|
|
|
||||||
|
|
@ -23,8 +23,8 @@ async def test_auto_compact_triggers_and_batches_observer() -> None:
|
||||||
observed.append((msg.role, msg.content))
|
observed.append((msg.role, msg.content))
|
||||||
|
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="<summary>"),
|
[mock_llm_chunk(content="<summary>")],
|
||||||
mock_llm_chunk(content="<final>"),
|
[mock_llm_chunk(content="<final>")],
|
||||||
])
|
])
|
||||||
cfg = VibeConfig(
|
cfg = VibeConfig(
|
||||||
session_logging=SessionLoggingConfig(enabled=False), auto_compact_threshold=1
|
session_logging=SessionLoggingConfig(enabled=False), auto_compact_threshold=1
|
||||||
|
|
|
||||||
|
|
@ -15,7 +15,7 @@ def vibe_config() -> VibeConfig:
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_passes_x_affinity_header_when_asking_an_answer(vibe_config: VibeConfig):
|
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)
|
agent = Agent(vibe_config, backend=backend)
|
||||||
|
|
||||||
[_ async for _ in agent.act("Hello")]
|
[_ 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(
|
async def test_passes_x_affinity_header_when_asking_an_answer_streaming(
|
||||||
vibe_config: VibeConfig,
|
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)
|
agent = Agent(vibe_config, backend=backend, enable_streaming=True)
|
||||||
|
|
||||||
[_ async for _ in agent.act("Hello")]
|
[_ 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
|
@pytest.mark.asyncio
|
||||||
async def test_updates_tokens_stats_based_on_backend_response(vibe_config: VibeConfig):
|
async def test_updates_tokens_stats_based_on_backend_response(vibe_config: VibeConfig):
|
||||||
chunk = mock_llm_chunk(
|
chunk = mock_llm_chunk(content="Response", prompt_tokens=100, completion_tokens=50)
|
||||||
content="Response",
|
|
||||||
finish_reason="stop",
|
|
||||||
prompt_tokens=100,
|
|
||||||
completion_tokens=50,
|
|
||||||
)
|
|
||||||
backend = FakeBackend([chunk])
|
backend = FakeBackend([chunk])
|
||||||
agent = Agent(vibe_config, backend=backend)
|
agent = Agent(vibe_config, backend=backend)
|
||||||
|
|
||||||
|
|
@ -64,10 +59,7 @@ async def test_updates_tokens_stats_based_on_backend_response_streaming(
|
||||||
vibe_config: VibeConfig,
|
vibe_config: VibeConfig,
|
||||||
):
|
):
|
||||||
final_chunk = mock_llm_chunk(
|
final_chunk = mock_llm_chunk(
|
||||||
content="Complete",
|
content="Complete", prompt_tokens=200, completion_tokens=75
|
||||||
finish_reason="stop",
|
|
||||||
prompt_tokens=200,
|
|
||||||
completion_tokens=75,
|
|
||||||
)
|
)
|
||||||
backend = FakeBackend([final_chunk])
|
backend = FakeBackend([final_chunk])
|
||||||
agent = Agent(vibe_config, backend=backend, enable_streaming=True)
|
agent = Agent(vibe_config, backend=backend, enable_streaming=True)
|
||||||
|
|
|
||||||
|
|
@ -10,7 +10,6 @@ from tests.mock.utils import mock_llm_chunk
|
||||||
from tests.stubs.fake_backend import FakeBackend
|
from tests.stubs.fake_backend import FakeBackend
|
||||||
from vibe.core.agent import Agent
|
from vibe.core.agent import Agent
|
||||||
from vibe.core.config import SessionLoggingConfig, VibeConfig
|
from vibe.core.config import SessionLoggingConfig, VibeConfig
|
||||||
from vibe.core.llm.types import BackendLike
|
|
||||||
from vibe.core.middleware import (
|
from vibe.core.middleware import (
|
||||||
ConversationContext,
|
ConversationContext,
|
||||||
MiddlewareAction,
|
MiddlewareAction,
|
||||||
|
|
@ -25,7 +24,6 @@ from vibe.core.types import (
|
||||||
ApprovalResponse,
|
ApprovalResponse,
|
||||||
AssistantEvent,
|
AssistantEvent,
|
||||||
FunctionCall,
|
FunctionCall,
|
||||||
LLMChunk,
|
|
||||||
LLMMessage,
|
LLMMessage,
|
||||||
Role,
|
Role,
|
||||||
ToolCall,
|
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"}'),
|
function=FunctionCall(name="todo", arguments='{"action": "read"}'),
|
||||||
)
|
)
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
|
[
|
||||||
mock_llm_chunk(content="Checking your todos."),
|
mock_llm_chunk(content="Checking your todos."),
|
||||||
mock_llm_chunk(content="", tool_calls=[todo_tool_call]),
|
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(
|
agent = Agent(
|
||||||
make_config(
|
make_config(
|
||||||
|
|
@ -222,7 +221,7 @@ async def test_act_handles_tool_call_chunk_with_content() -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="Preparing "),
|
mock_llm_chunk(content="Preparing "),
|
||||||
mock_llm_chunk(content="todo request", tool_calls=[todo_tool_call]),
|
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(
|
agent = Agent(
|
||||||
make_config(
|
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.")]
|
events = [event async for event in agent.act("Check todos with content.")]
|
||||||
|
|
||||||
assert [type(event) for event in events] == [
|
assert [type(event) for event in events] == [
|
||||||
AssistantEvent,
|
|
||||||
AssistantEvent,
|
AssistantEvent,
|
||||||
ToolCallEvent,
|
ToolCallEvent,
|
||||||
ToolResultEvent,
|
ToolResultEvent,
|
||||||
]
|
]
|
||||||
assert isinstance(events[0], AssistantEvent)
|
assert isinstance(events[0], AssistantEvent)
|
||||||
assert events[0].content == "Preparing todo request"
|
assert events[0].content == "Preparing todo request complete"
|
||||||
assert isinstance(events[1], AssistantEvent)
|
|
||||||
assert events[1].content == " complete"
|
|
||||||
assert any(
|
assert any(
|
||||||
m.role == Role.assistant and m.content == "Preparing todo request complete"
|
m.role == Role.assistant and m.content == "Preparing todo request complete"
|
||||||
for m in agent.messages
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_act_handles_user_cancellation_during_streaming() -> None:
|
async def test_act_handles_user_cancellation_during_streaming() -> None:
|
||||||
class CountingMiddleware(MiddlewarePipeline):
|
class CountingMiddleware(MiddlewarePipeline):
|
||||||
|
|
@ -359,7 +326,6 @@ async def test_act_handles_user_cancellation_during_streaming() -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="Preparing "),
|
mock_llm_chunk(content="Preparing "),
|
||||||
mock_llm_chunk(content="todo request", tool_calls=[todo_tool_call]),
|
mock_llm_chunk(content="todo request", tool_calls=[todo_tool_call]),
|
||||||
mock_llm_chunk(content="", finish_reason="stop"),
|
|
||||||
])
|
])
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
make_config(
|
make_config(
|
||||||
|
|
|
||||||
|
|
@ -159,9 +159,7 @@ class TestAgentStatsHelpers:
|
||||||
class TestReloadPreservesStats:
|
class TestReloadPreservesStats:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_reload_preserves_session_tokens(self) -> None:
|
async def test_reload_preserves_session_tokens(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend(mock_llm_chunk(content="First response"))
|
||||||
mock_llm_chunk(content="First response", finish_reason="stop")
|
|
||||||
])
|
|
||||||
agent = Agent(make_config(), backend=backend)
|
agent = Agent(make_config(), backend=backend)
|
||||||
|
|
||||||
async for _ in agent.act("Hello"):
|
async for _ in agent.act("Hello"):
|
||||||
|
|
@ -185,13 +183,14 @@ class TestReloadPreservesStats:
|
||||||
tool_calls=[
|
tool_calls=[
|
||||||
ToolCall(
|
ToolCall(
|
||||||
id="tc1",
|
id="tc1",
|
||||||
|
index=0,
|
||||||
function=FunctionCall(
|
function=FunctionCall(
|
||||||
name="todo", arguments='{"action": "read"}'
|
name="todo", arguments='{"action": "read"}'
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
],
|
],
|
||||||
),
|
),
|
||||||
mock_llm_chunk(content="Done", finish_reason="stop"),
|
mock_llm_chunk(content="Done"),
|
||||||
])
|
])
|
||||||
config = make_config(enabled_tools=["todo"])
|
config = make_config(enabled_tools=["todo"])
|
||||||
agent = Agent(config, mode=AgentMode.AUTO_APPROVE, backend=backend)
|
agent = Agent(config, mode=AgentMode.AUTO_APPROVE, backend=backend)
|
||||||
|
|
@ -210,8 +209,8 @@ class TestReloadPreservesStats:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_reload_preserves_steps(self) -> None:
|
async def test_reload_preserves_steps(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="R1", finish_reason="stop"),
|
[mock_llm_chunk(content="R1")],
|
||||||
mock_llm_chunk(content="R2", finish_reason="stop"),
|
[mock_llm_chunk(content="R2")],
|
||||||
])
|
])
|
||||||
agent = Agent(make_config(), backend=backend)
|
agent = Agent(make_config(), backend=backend)
|
||||||
|
|
||||||
|
|
@ -231,9 +230,7 @@ class TestReloadPreservesStats:
|
||||||
async def test_reload_preserves_context_tokens_when_messages_preserved(
|
async def test_reload_preserves_context_tokens_when_messages_preserved(
|
||||||
self,
|
self,
|
||||||
) -> None:
|
) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend(mock_llm_chunk(content="Response"))
|
||||||
mock_llm_chunk(content="Response", finish_reason="stop")
|
|
||||||
])
|
|
||||||
agent = Agent(make_config(), backend=backend)
|
agent = Agent(make_config(), backend=backend)
|
||||||
[_ async for _ in agent.act("Hello")]
|
[_ async for _ in agent.act("Hello")]
|
||||||
assert agent.stats.context_tokens > 0
|
assert agent.stats.context_tokens > 0
|
||||||
|
|
@ -258,10 +255,10 @@ class TestReloadPreservesStats:
|
||||||
assert agent.stats.context_tokens == 0
|
assert agent.stats.context_tokens == 0
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_reload_preserves_context_tokens_when_messages_exist(self) -> None:
|
async def test_reload_resets_context_tokens_when_system_prompt_changes(
|
||||||
backend = FakeBackend([
|
self,
|
||||||
mock_llm_chunk(content="Response", finish_reason="stop")
|
) -> None:
|
||||||
])
|
backend = FakeBackend(mock_llm_chunk(content="Response"))
|
||||||
config1 = make_config(system_prompt_id="tests")
|
config1 = make_config(system_prompt_id="tests")
|
||||||
config2 = make_config(system_prompt_id="cli")
|
config2 = make_config(system_prompt_id="cli")
|
||||||
agent = Agent(config1, backend=backend)
|
agent = Agent(config1, backend=backend)
|
||||||
|
|
@ -279,9 +276,7 @@ class TestReloadPreservesStats:
|
||||||
async def test_reload_updates_pricing_from_new_model(self, monkeypatch) -> None:
|
async def test_reload_updates_pricing_from_new_model(self, monkeypatch) -> None:
|
||||||
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
|
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
|
||||||
|
|
||||||
backend = FakeBackend([
|
backend = FakeBackend(mock_llm_chunk(content="Response"))
|
||||||
mock_llm_chunk(content="Response", finish_reason="stop")
|
|
||||||
])
|
|
||||||
config_mistral = make_config(active_model="devstral-latest")
|
config_mistral = make_config(active_model="devstral-latest")
|
||||||
agent = Agent(config_mistral, backend=backend)
|
agent = Agent(config_mistral, backend=backend)
|
||||||
|
|
||||||
|
|
@ -302,8 +297,8 @@ class TestReloadPreservesStats:
|
||||||
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
|
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
|
||||||
|
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="First", finish_reason="stop"),
|
[mock_llm_chunk(content="First")],
|
||||||
mock_llm_chunk(content="After reload", finish_reason="stop"),
|
[mock_llm_chunk(content="After reload")],
|
||||||
])
|
])
|
||||||
config1 = make_config(active_model="devstral-latest")
|
config1 = make_config(active_model="devstral-latest")
|
||||||
agent = Agent(config1, backend=backend)
|
agent = Agent(config1, backend=backend)
|
||||||
|
|
@ -330,9 +325,7 @@ class TestReloadPreservesStats:
|
||||||
class TestReloadPreservesMessages:
|
class TestReloadPreservesMessages:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_reload_preserves_conversation_messages(self) -> None:
|
async def test_reload_preserves_conversation_messages(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend(mock_llm_chunk(content="Response"))
|
||||||
mock_llm_chunk(content="Response", finish_reason="stop")
|
|
||||||
])
|
|
||||||
agent = Agent(make_config(), backend=backend)
|
agent = Agent(make_config(), backend=backend)
|
||||||
|
|
||||||
async for _ in agent.act("Hello"):
|
async for _ in agent.act("Hello"):
|
||||||
|
|
@ -353,9 +346,7 @@ class TestReloadPreservesMessages:
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_reload_updates_system_prompt_preserves_rest(self) -> None:
|
async def test_reload_updates_system_prompt_preserves_rest(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend(mock_llm_chunk(content="Response"))
|
||||||
mock_llm_chunk(content="Response", finish_reason="stop")
|
|
||||||
])
|
|
||||||
config1 = make_config(system_prompt_id="tests")
|
config1 = make_config(system_prompt_id="tests")
|
||||||
agent = Agent(config1, backend=backend)
|
agent = Agent(config1, backend=backend)
|
||||||
|
|
||||||
|
|
@ -388,9 +379,7 @@ class TestReloadPreservesMessages:
|
||||||
self, observer_capture
|
self, observer_capture
|
||||||
) -> None:
|
) -> None:
|
||||||
observed, observer = observer_capture
|
observed, observer = observer_capture
|
||||||
backend = FakeBackend([
|
backend = FakeBackend(mock_llm_chunk(content="Response"))
|
||||||
mock_llm_chunk(content="Response", finish_reason="stop")
|
|
||||||
])
|
|
||||||
agent = Agent(make_config(), message_observer=observer, backend=backend)
|
agent = Agent(make_config(), message_observer=observer, backend=backend)
|
||||||
|
|
||||||
async for _ in agent.act("Hello"):
|
async for _ in agent.act("Hello"):
|
||||||
|
|
@ -410,8 +399,8 @@ class TestCompactStatsHandling:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_compact_preserves_cumulative_stats(self) -> None:
|
async def test_compact_preserves_cumulative_stats(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="First response", finish_reason="stop"),
|
[mock_llm_chunk(content="First response")],
|
||||||
mock_llm_chunk(content="<summary>", finish_reason="stop"),
|
[mock_llm_chunk(content="<summary>")],
|
||||||
])
|
])
|
||||||
agent = Agent(make_config(), backend=backend)
|
agent = Agent(make_config(), backend=backend)
|
||||||
|
|
||||||
|
|
@ -432,8 +421,8 @@ class TestCompactStatsHandling:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_compact_updates_context_tokens(self) -> None:
|
async def test_compact_updates_context_tokens(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="Long response " * 100, finish_reason="stop"),
|
[mock_llm_chunk(content="Long response " * 100)],
|
||||||
mock_llm_chunk(content="<summary>", finish_reason="stop"),
|
[mock_llm_chunk(content="<summary>")],
|
||||||
])
|
])
|
||||||
agent = Agent(make_config(), backend=backend)
|
agent = Agent(make_config(), backend=backend)
|
||||||
|
|
||||||
|
|
@ -449,19 +438,22 @@ class TestCompactStatsHandling:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_compact_preserves_tool_call_stats(self) -> None:
|
async def test_compact_preserves_tool_call_stats(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
|
[
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="Using tool",
|
content="Using tool",
|
||||||
tool_calls=[
|
tool_calls=[
|
||||||
ToolCall(
|
ToolCall(
|
||||||
id="tc1",
|
id="tc1",
|
||||||
|
index=0,
|
||||||
function=FunctionCall(
|
function=FunctionCall(
|
||||||
name="todo", arguments='{"action": "read"}'
|
name="todo", arguments='{"action": "read"}'
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
],
|
],
|
||||||
),
|
),
|
||||||
mock_llm_chunk(content="Done", finish_reason="stop"),
|
mock_llm_chunk(content=" todo"),
|
||||||
mock_llm_chunk(content="<summary>", finish_reason="stop"),
|
],
|
||||||
|
[mock_llm_chunk(content="<summary>")],
|
||||||
])
|
])
|
||||||
config = make_config(enabled_tools=["todo"])
|
config = make_config(enabled_tools=["todo"])
|
||||||
agent = Agent(config, mode=AgentMode.AUTO_APPROVE, backend=backend)
|
agent = Agent(config, mode=AgentMode.AUTO_APPROVE, backend=backend)
|
||||||
|
|
@ -478,8 +470,8 @@ class TestCompactStatsHandling:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_compact_resets_session_id(self) -> None:
|
async def test_compact_resets_session_id(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="Long response " * 100, finish_reason="stop"),
|
[mock_llm_chunk(content="Long response " * 100)],
|
||||||
mock_llm_chunk(content="<summary>", finish_reason="stop"),
|
[mock_llm_chunk(content="<summary>")],
|
||||||
])
|
])
|
||||||
agent = Agent(make_config(disable_logging=False), backend=backend)
|
agent = Agent(make_config(disable_logging=False), backend=backend)
|
||||||
|
|
||||||
|
|
@ -506,8 +498,8 @@ class TestAutoCompactIntegration:
|
||||||
observed.append((msg.role, msg.content))
|
observed.append((msg.role, msg.content))
|
||||||
|
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="<summary>", finish_reason="stop"),
|
[mock_llm_chunk(content="<summary>")],
|
||||||
mock_llm_chunk(content="<final>", finish_reason="stop"),
|
[mock_llm_chunk(content="<final>")],
|
||||||
])
|
])
|
||||||
cfg = VibeConfig(
|
cfg = VibeConfig(
|
||||||
session_logging=SessionLoggingConfig(enabled=False),
|
session_logging=SessionLoggingConfig(enabled=False),
|
||||||
|
|
@ -544,9 +536,7 @@ class TestAutoCompactIntegration:
|
||||||
class TestClearHistoryFullReset:
|
class TestClearHistoryFullReset:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_clear_history_fully_resets_stats(self) -> None:
|
async def test_clear_history_fully_resets_stats(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend(mock_llm_chunk(content="Response"))
|
||||||
mock_llm_chunk(content="Response", finish_reason="stop")
|
|
||||||
])
|
|
||||||
agent = Agent(make_config(), backend=backend)
|
agent = Agent(make_config(), backend=backend)
|
||||||
|
|
||||||
async for _ in agent.act("Hello"):
|
async for _ in agent.act("Hello"):
|
||||||
|
|
@ -563,9 +553,7 @@ class TestClearHistoryFullReset:
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_clear_history_preserves_pricing(self) -> None:
|
async def test_clear_history_preserves_pricing(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend(mock_llm_chunk(content="Response"))
|
||||||
mock_llm_chunk(content="Response", finish_reason="stop")
|
|
||||||
])
|
|
||||||
config = make_config(input_price=0.4, output_price=2.0)
|
config = make_config(input_price=0.4, output_price=2.0)
|
||||||
agent = Agent(config, backend=backend)
|
agent = Agent(config, backend=backend)
|
||||||
|
|
||||||
|
|
@ -579,9 +567,7 @@ class TestClearHistoryFullReset:
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_clear_history_removes_messages(self) -> None:
|
async def test_clear_history_removes_messages(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend(mock_llm_chunk(content="Response"))
|
||||||
mock_llm_chunk(content="Response", finish_reason="stop")
|
|
||||||
])
|
|
||||||
agent = Agent(make_config(), backend=backend)
|
agent = Agent(make_config(), backend=backend)
|
||||||
|
|
||||||
async for _ in agent.act("Hello"):
|
async for _ in agent.act("Hello"):
|
||||||
|
|
@ -596,9 +582,7 @@ class TestClearHistoryFullReset:
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_clear_history_resets_session_id(self) -> None:
|
async def test_clear_history_resets_session_id(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend(mock_llm_chunk(content="Response"))
|
||||||
mock_llm_chunk(content="Response", finish_reason="stop")
|
|
||||||
])
|
|
||||||
agent = Agent(make_config(disable_logging=False), backend=backend)
|
agent = Agent(make_config(disable_logging=False), backend=backend)
|
||||||
|
|
||||||
original_session_id = agent.session_id
|
original_session_id = agent.session_id
|
||||||
|
|
@ -622,9 +606,7 @@ class TestStatsEdgeCases:
|
||||||
) -> None:
|
) -> None:
|
||||||
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
|
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
|
||||||
|
|
||||||
backend = FakeBackend([
|
backend = FakeBackend(mock_llm_chunk(content="Response"))
|
||||||
mock_llm_chunk(content="Response", finish_reason="stop")
|
|
||||||
])
|
|
||||||
config1 = make_config(active_model="devstral-latest")
|
config1 = make_config(active_model="devstral-latest")
|
||||||
agent = Agent(config1, backend=backend)
|
agent = Agent(config1, backend=backend)
|
||||||
|
|
||||||
|
|
@ -643,9 +625,9 @@ class TestStatsEdgeCases:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_multiple_reloads_accumulate_correctly(self) -> None:
|
async def test_multiple_reloads_accumulate_correctly(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="R1", finish_reason="stop"),
|
[mock_llm_chunk(content="R1")],
|
||||||
mock_llm_chunk(content="R2", finish_reason="stop"),
|
[mock_llm_chunk(content="R2")],
|
||||||
mock_llm_chunk(content="R3", finish_reason="stop"),
|
[mock_llm_chunk(content="R3")],
|
||||||
])
|
])
|
||||||
agent = Agent(make_config(), backend=backend)
|
agent = Agent(make_config(), backend=backend)
|
||||||
|
|
||||||
|
|
@ -668,9 +650,9 @@ class TestStatsEdgeCases:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_compact_then_reload_preserves_both(self) -> None:
|
async def test_compact_then_reload_preserves_both(self) -> None:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="Initial response", finish_reason="stop"),
|
[mock_llm_chunk(content="Initial response")],
|
||||||
mock_llm_chunk(content="<summary>", finish_reason="stop"),
|
[mock_llm_chunk(content="<summary>")],
|
||||||
mock_llm_chunk(content="After reload", finish_reason="stop"),
|
[mock_llm_chunk(content="After reload")],
|
||||||
])
|
])
|
||||||
agent = Agent(make_config(), backend=backend)
|
agent = Agent(make_config(), backend=backend)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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(
|
return ToolCall(
|
||||||
id=call_id,
|
id=call_id, index=index, function=FunctionCall(name="todo", arguments=args)
|
||||||
function=FunctionCall(name="todo", arguments=f'{{"action": "{action}"}}'),
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -72,8 +74,8 @@ async def test_single_tool_call_executes_under_auto_approve() -> None:
|
||||||
mocked_tool_call_id = "call_1"
|
mocked_tool_call_id = "call_1"
|
||||||
tool_call = make_todo_tool_call(mocked_tool_call_id)
|
tool_call = make_todo_tool_call(mocked_tool_call_id)
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="Let me check your todos.", tool_calls=[tool_call]),
|
[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="I retrieved 0 todos.")],
|
||||||
])
|
])
|
||||||
agent = make_agent(auto_approve=True, backend=backend)
|
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,
|
auto_approve=False,
|
||||||
todo_permission=ToolPermission.ASK,
|
todo_permission=ToolPermission.ASK,
|
||||||
backend=FakeBackend([
|
backend=FakeBackend([
|
||||||
|
[
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="Let me check your todos.",
|
content="Let me check your todos.",
|
||||||
tool_calls=[make_todo_tool_call("call_2")],
|
tool_calls=[make_todo_tool_call("call_2")],
|
||||||
),
|
)
|
||||||
mock_llm_chunk(
|
],
|
||||||
content="I cannot execute the tool without approval.",
|
[mock_llm_chunk(content="I cannot execute the tool without approval.")],
|
||||||
finish_reason="stop",
|
|
||||||
),
|
|
||||||
]),
|
]),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -148,11 +149,13 @@ async def test_tool_call_approved_by_callback() -> None:
|
||||||
todo_permission=ToolPermission.ASK,
|
todo_permission=ToolPermission.ASK,
|
||||||
approval_callback=approval_callback,
|
approval_callback=approval_callback,
|
||||||
backend=FakeBackend([
|
backend=FakeBackend([
|
||||||
|
[
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="Let me check your todos.",
|
content="Let me check your todos.",
|
||||||
tool_calls=[make_todo_tool_call("call_3")],
|
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,
|
todo_permission=ToolPermission.ASK,
|
||||||
approval_callback=approval_callback,
|
approval_callback=approval_callback,
|
||||||
backend=FakeBackend([
|
backend=FakeBackend([
|
||||||
|
[
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="Let me check your todos.",
|
content="Let me check your todos.",
|
||||||
tool_calls=[make_todo_tool_call("call_4")],
|
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,
|
auto_approve=False,
|
||||||
todo_permission=ToolPermission.NEVER,
|
todo_permission=ToolPermission.NEVER,
|
||||||
backend=FakeBackend([
|
backend=FakeBackend([
|
||||||
|
[
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="Let me check your todos.",
|
content="Let me check your todos.",
|
||||||
tool_calls=[make_todo_tool_call("call_never")],
|
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,
|
todo_permission=ToolPermission.ASK,
|
||||||
approval_callback=approval_callback,
|
approval_callback=approval_callback,
|
||||||
backend=FakeBackend([
|
backend=FakeBackend([
|
||||||
|
[
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="First check.", tool_calls=[make_todo_tool_call("call_first")]
|
content="First check.",
|
||||||
),
|
tool_calls=[make_todo_tool_call("call_first")],
|
||||||
mock_llm_chunk(content="First done.", finish_reason="stop"),
|
)
|
||||||
|
],
|
||||||
|
[mock_llm_chunk(content="First done.")],
|
||||||
|
[
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="Second check.", tool_calls=[make_todo_tool_call("call_second")]
|
content="Second check.",
|
||||||
),
|
tool_calls=[make_todo_tool_call("call_second")],
|
||||||
mock_llm_chunk(content="Second done.", finish_reason="stop"),
|
)
|
||||||
|
],
|
||||||
|
[mock_llm_chunk(content="Second done.")],
|
||||||
]),
|
]),
|
||||||
)
|
)
|
||||||
agent_ref = agent
|
agent_ref = agent
|
||||||
|
|
@ -291,17 +302,16 @@ async def test_approval_always_sets_tool_permission_for_subsequent_calls() -> No
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_tool_call_with_invalid_action() -> None:
|
async def test_tool_call_with_invalid_action() -> None:
|
||||||
tool_call = ToolCall(
|
tool_call = make_todo_tool_call("call_5", arguments='{"action": "invalid_action"}')
|
||||||
id="call_5",
|
|
||||||
function=FunctionCall(name="todo", arguments='{"action": "invalid_action"}'),
|
|
||||||
)
|
|
||||||
agent = make_agent(
|
agent = make_agent(
|
||||||
auto_approve=True,
|
auto_approve=True,
|
||||||
backend=FakeBackend([
|
backend=FakeBackend([
|
||||||
mock_llm_chunk(content="Let me check your todos.", tool_calls=[tool_call]),
|
[
|
||||||
mock_llm_chunk(
|
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 1"),
|
||||||
TodoItem(id="duplicate", content="Task 2"),
|
TodoItem(id="duplicate", content="Task 2"),
|
||||||
]
|
]
|
||||||
tool_call = ToolCall(
|
tool_call = make_todo_tool_call(
|
||||||
id="call_6",
|
"call_6",
|
||||||
function=FunctionCall(
|
|
||||||
name="todo",
|
|
||||||
arguments=json.dumps({
|
arguments=json.dumps({
|
||||||
"action": "write",
|
"action": "write",
|
||||||
"todos": [t.model_dump() for t in duplicate_todos],
|
"todos": [t.model_dump() for t in duplicate_todos],
|
||||||
}),
|
}),
|
||||||
),
|
|
||||||
)
|
)
|
||||||
agent = make_agent(
|
agent = make_agent(
|
||||||
auto_approve=True,
|
auto_approve=True,
|
||||||
backend=FakeBackend([
|
backend=FakeBackend([
|
||||||
mock_llm_chunk(content="Let me write todos.", tool_calls=[tool_call]),
|
[mock_llm_chunk(content="Let me write todos.", tool_calls=[tool_call])],
|
||||||
mock_llm_chunk(
|
[mock_llm_chunk(content="I couldn't write todos with duplicate IDs.")],
|
||||||
content="I couldn't write todos with duplicate IDs.",
|
|
||||||
finish_reason="stop",
|
|
||||||
),
|
|
||||||
]),
|
]),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -353,23 +357,18 @@ async def test_tool_call_with_duplicate_todo_ids() -> None:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_tool_call_with_exceeding_max_todos() -> None:
|
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)]
|
many_todos = [TodoItem(id=f"todo_{i}", content=f"Task {i}") for i in range(150)]
|
||||||
tool_call = ToolCall(
|
tool_call = make_todo_tool_call(
|
||||||
id="call_7",
|
"call_7",
|
||||||
function=FunctionCall(
|
|
||||||
name="todo",
|
|
||||||
arguments=json.dumps({
|
arguments=json.dumps({
|
||||||
"action": "write",
|
"action": "write",
|
||||||
"todos": [t.model_dump() for t in many_todos],
|
"todos": [t.model_dump() for t in many_todos],
|
||||||
}),
|
}),
|
||||||
),
|
|
||||||
)
|
)
|
||||||
agent = make_agent(
|
agent = make_agent(
|
||||||
auto_approve=True,
|
auto_approve=True,
|
||||||
backend=FakeBackend([
|
backend=FakeBackend([
|
||||||
mock_llm_chunk(content="Let me write todos.", tool_calls=[tool_call]),
|
[mock_llm_chunk(content="Let me write todos.", tool_calls=[tool_call])],
|
||||||
mock_llm_chunk(
|
[mock_llm_chunk(content="I couldn't write that many todos.")],
|
||||||
content="I couldn't write that many todos.", finish_reason="stop"
|
|
||||||
),
|
|
||||||
]),
|
]),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -394,7 +393,7 @@ async def test_tool_call_can_be_interrupted(
|
||||||
exception_class: type[BaseException],
|
exception_class: type[BaseException],
|
||||||
) -> None:
|
) -> None:
|
||||||
tool_call = ToolCall(
|
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(
|
config = VibeConfig(
|
||||||
session_logging=SessionLoggingConfig(enabled=False),
|
session_logging=SessionLoggingConfig(enabled=False),
|
||||||
|
|
@ -405,8 +404,8 @@ async def test_tool_call_can_be_interrupted(
|
||||||
config,
|
config,
|
||||||
mode=AgentMode.AUTO_APPROVE,
|
mode=AgentMode.AUTO_APPROVE,
|
||||||
backend=FakeBackend([
|
backend=FakeBackend([
|
||||||
mock_llm_chunk(content="Let me use the tool.", tool_calls=[tool_call]),
|
[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="Tool execution completed.")],
|
||||||
]),
|
]),
|
||||||
)
|
)
|
||||||
# no dependency injection available => monkey patch
|
# no dependency injection available => monkey patch
|
||||||
|
|
@ -433,15 +432,11 @@ async def test_fill_missing_tool_responses_inserts_placeholders() -> None:
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
make_config(),
|
make_config(),
|
||||||
mode=AgentMode.AUTO_APPROVE,
|
mode=AgentMode.AUTO_APPROVE,
|
||||||
backend=FakeBackend([mock_llm_chunk(content="ok", finish_reason="stop")]),
|
backend=FakeBackend(mock_llm_chunk(content="ok")),
|
||||||
)
|
)
|
||||||
tool_calls_messages = [
|
tool_calls_messages = [
|
||||||
ToolCall(
|
make_todo_tool_call("tc1", index=0),
|
||||||
id="tc1", function=FunctionCall(name="todo", arguments='{"action": "read"}')
|
make_todo_tool_call("tc2", index=1),
|
||||||
),
|
|
||||||
ToolCall(
|
|
||||||
id="tc2", function=FunctionCall(name="todo", arguments='{"action": "read"}')
|
|
||||||
),
|
|
||||||
]
|
]
|
||||||
assistant_msg = LLMMessage(
|
assistant_msg = LLMMessage(
|
||||||
role=Role.assistant, content="Calling tools...", tool_calls=tool_calls_messages
|
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(
|
agent = Agent(
|
||||||
make_config(),
|
make_config(),
|
||||||
mode=AgentMode.AUTO_APPROVE,
|
mode=AgentMode.AUTO_APPROVE,
|
||||||
backend=FakeBackend([mock_llm_chunk(content="ok", finish_reason="stop")]),
|
backend=FakeBackend(mock_llm_chunk(content="ok")),
|
||||||
)
|
)
|
||||||
tool_msg = LLMMessage(
|
tool_msg = LLMMessage(
|
||||||
role=Role.tool, tool_call_id="tc_z", name="todo", content="Done"
|
role=Role.tool, tool_call_id="tc_z", name="todo", content="Done"
|
||||||
|
|
|
||||||
|
|
@ -34,12 +34,11 @@ def test_run_programmatic_preload_streaming_is_batched(
|
||||||
|
|
||||||
with mock_backend_factory(
|
with mock_backend_factory(
|
||||||
Backend.MISTRAL,
|
Backend.MISTRAL,
|
||||||
lambda provider, **kwargs: FakeBackend([
|
lambda provider, **kwargs: FakeBackend(
|
||||||
mock_llm_chunk(
|
mock_llm_chunk(
|
||||||
content="Decorators are wrappers that modify function behavior.",
|
content="Decorators are wrappers that modify function behavior."
|
||||||
finish_reason="stop",
|
|
||||||
)
|
)
|
||||||
]),
|
),
|
||||||
):
|
):
|
||||||
cfg = VibeConfig(
|
cfg = VibeConfig(
|
||||||
session_logging=SessionLoggingConfig(enabled=False),
|
session_logging=SessionLoggingConfig(enabled=False),
|
||||||
|
|
|
||||||
|
|
@ -157,7 +157,6 @@ class TestAgentSwitchMode:
|
||||||
return FakeBackend([
|
return FakeBackend([
|
||||||
LLMChunk(
|
LLMChunk(
|
||||||
message=LLMMessage(role=Role.assistant, content="Test response"),
|
message=LLMMessage(role=Role.assistant, content="Test response"),
|
||||||
finish_reason="stop",
|
|
||||||
usage=LLMUsage(prompt_tokens=10, completion_tokens=5),
|
usage=LLMUsage(prompt_tokens=10, completion_tokens=5),
|
||||||
)
|
)
|
||||||
])
|
])
|
||||||
|
|
@ -274,7 +273,6 @@ class TestPlanModeToolRestriction:
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
LLMChunk(
|
LLMChunk(
|
||||||
message=LLMMessage(role=Role.assistant, content="ok"),
|
message=LLMMessage(role=Role.assistant, content="ok"),
|
||||||
finish_reason="stop",
|
|
||||||
usage=LLMUsage(prompt_tokens=10, completion_tokens=5),
|
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:
|
async def test_plan_mode_rejects_non_plan_tool_call(self) -> None:
|
||||||
tool_call = ToolCall(
|
tool_call = ToolCall(
|
||||||
id="call_1",
|
id="call_1",
|
||||||
|
index=0,
|
||||||
function=FunctionCall(name="bash", arguments='{"command": "ls"}'),
|
function=FunctionCall(name="bash", arguments='{"command": "ls"}'),
|
||||||
)
|
)
|
||||||
backend = FakeBackend([
|
backend = FakeBackend([
|
||||||
mock_llm_chunk(content="Let me run bash", tool_calls=[tool_call]),
|
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(
|
config = VibeConfig(
|
||||||
|
|
|
||||||
2
uv.lock
generated
2
uv.lock
generated
|
|
@ -661,7 +661,7 @@ wheels = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "mistral-vibe"
|
name = "mistral-vibe"
|
||||||
version = "1.2.1"
|
version = "1.2.2"
|
||||||
source = { editable = "." }
|
source = { editable = "." }
|
||||||
dependencies = [
|
dependencies = [
|
||||||
{ name = "agent-client-protocol" },
|
{ name = "agent-client-protocol" },
|
||||||
|
|
|
||||||
|
|
@ -3,4 +3,4 @@ from __future__ import annotations
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
VIBE_ROOT = Path(__file__).parent
|
VIBE_ROOT = Path(__file__).parent
|
||||||
__version__ = "1.2.1"
|
__version__ = "1.2.2"
|
||||||
|
|
|
||||||
|
|
@ -38,7 +38,6 @@ class EventHandler:
|
||||||
self.get_todos_collapsed = get_todos_collapsed
|
self.get_todos_collapsed = get_todos_collapsed
|
||||||
self.current_tool_call: ToolCallMessage | None = None
|
self.current_tool_call: ToolCallMessage | None = None
|
||||||
self.current_compact: CompactMessage | None = None
|
self.current_compact: CompactMessage | None = None
|
||||||
self.tool_results: list[ToolResultMessage] = []
|
|
||||||
|
|
||||||
async def handle_event(
|
async def handle_event(
|
||||||
self,
|
self,
|
||||||
|
|
@ -121,7 +120,6 @@ class EventHandler:
|
||||||
)
|
)
|
||||||
await self.mount_callback(tool_result)
|
await self.mount_callback(tool_result)
|
||||||
|
|
||||||
self.tool_results.append(tool_result)
|
|
||||||
self.current_tool_call = None
|
self.current_tool_call = None
|
||||||
|
|
||||||
async def _handle_assistant_message(self, event: AssistantEvent) -> None:
|
async def _handle_assistant_message(self, event: AssistantEvent) -> None:
|
||||||
|
|
@ -153,6 +151,3 @@ class EventHandler:
|
||||||
if self.current_compact:
|
if self.current_compact:
|
||||||
self.current_compact.stop_spinning(success=False)
|
self.current_compact.stop_spinning(success=False)
|
||||||
self.current_compact = None
|
self.current_compact = None
|
||||||
|
|
||||||
def get_last_tool_result(self) -> ToolResultMessage | None:
|
|
||||||
return self.tool_results[-1] if self.tool_results else None
|
|
||||||
|
|
|
||||||
|
|
@ -1,7 +1,6 @@
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
from collections import OrderedDict
|
|
||||||
from collections.abc import AsyncGenerator, Callable
|
from collections.abc import AsyncGenerator, Callable
|
||||||
from enum import StrEnum, auto
|
from enum import StrEnum, auto
|
||||||
import time
|
import time
|
||||||
|
|
@ -48,9 +47,9 @@ from vibe.core.types import (
|
||||||
CompactStartEvent,
|
CompactStartEvent,
|
||||||
LLMChunk,
|
LLMChunk,
|
||||||
LLMMessage,
|
LLMMessage,
|
||||||
|
LLMUsage,
|
||||||
Role,
|
Role,
|
||||||
SyncApprovalCallback,
|
SyncApprovalCallback,
|
||||||
ToolCall,
|
|
||||||
ToolCallEvent,
|
ToolCallEvent,
|
||||||
ToolResultEvent,
|
ToolResultEvent,
|
||||||
)
|
)
|
||||||
|
|
@ -141,8 +140,6 @@ class Agent:
|
||||||
config.effective_workdir,
|
config.effective_workdir,
|
||||||
)
|
)
|
||||||
|
|
||||||
self._last_chunk: LLMChunk | None = None
|
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def mode(self) -> AgentMode:
|
def mode(self) -> AgentMode:
|
||||||
return self._mode
|
return self._mode
|
||||||
|
|
@ -267,11 +264,7 @@ class Agent:
|
||||||
yield event
|
yield event
|
||||||
|
|
||||||
last_message = self.messages[-1]
|
last_message = self.messages[-1]
|
||||||
should_break_loop = (
|
should_break_loop = last_message.role != Role.tool
|
||||||
last_message.role != Role.tool
|
|
||||||
and self._last_chunk is not None
|
|
||||||
and self._last_chunk.finish_reason is not None
|
|
||||||
)
|
|
||||||
|
|
||||||
self._flush_new_messages()
|
self._flush_new_messages()
|
||||||
|
|
||||||
|
|
@ -293,9 +286,7 @@ class Agent:
|
||||||
self.messages, self.stats, self.config, self.tool_manager
|
self.messages, self.stats, self.config, self.tool_manager
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _perform_llm_turn(
|
async def _perform_llm_turn(self) -> AsyncGenerator[BaseEvent, None]:
|
||||||
self,
|
|
||||||
) -> AsyncGenerator[AssistantEvent | ToolCallEvent | ToolResultEvent]:
|
|
||||||
if self.enable_streaming:
|
if self.enable_streaming:
|
||||||
async for event in self._stream_assistant_events():
|
async for event in self._stream_assistant_events():
|
||||||
yield event
|
yield event
|
||||||
|
|
@ -305,105 +296,40 @@ class Agent:
|
||||||
yield assistant_event
|
yield assistant_event
|
||||||
|
|
||||||
last_message = self.messages[-1]
|
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)
|
parsed = self.format_handler.parse_message(last_message)
|
||||||
resolved = self.format_handler.resolve_tool_calls(
|
resolved = self.format_handler.resolve_tool_calls(
|
||||||
parsed, self.tool_manager, self.config
|
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:
|
if not resolved.tool_calls and not resolved.failed_calls:
|
||||||
return
|
return
|
||||||
|
|
||||||
async for event in self._handle_tool_calls(resolved):
|
async for event in self._handle_tool_calls(resolved):
|
||||||
yield event
|
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]:
|
async def _stream_assistant_events(self) -> AsyncGenerator[AssistantEvent]:
|
||||||
chunks: list[LLMChunk] = []
|
batched_chunk = LLMChunk(message=LLMMessage(role=Role.assistant))
|
||||||
content_buffer = ""
|
|
||||||
chunks_with_content = 0
|
chunks_with_content = 0
|
||||||
BATCH_SIZE = 5
|
BATCH_SIZE = 5
|
||||||
|
|
||||||
async for chunk in self._chat_streaming():
|
async for chunk in self._chat_streaming():
|
||||||
chunks.append(chunk)
|
batched_chunk += 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
|
|
||||||
|
|
||||||
if chunk.message.content:
|
if chunk.message.content:
|
||||||
content_buffer += chunk.message.content
|
|
||||||
chunks_with_content += 1
|
chunks_with_content += 1
|
||||||
|
|
||||||
if chunks_with_content >= BATCH_SIZE:
|
if chunks_with_content >= BATCH_SIZE:
|
||||||
yield self._create_assistant_event(content_buffer, chunk)
|
yield AssistantEvent(content=cast(str, batched_chunk.message.content))
|
||||||
content_buffer = ""
|
batched_chunk = LLMChunk(message=LLMMessage(role=Role.assistant))
|
||||||
chunks_with_content = 0
|
chunks_with_content = 0
|
||||||
|
|
||||||
if content_buffer:
|
if batched_chunk.message.content:
|
||||||
last_chunk = chunks[-1] if chunks else None
|
yield AssistantEvent(content=batched_chunk.message.content)
|
||||||
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
|
|
||||||
)
|
|
||||||
|
|
||||||
async def _get_assistant_event(self) -> AssistantEvent:
|
async def _get_assistant_event(self) -> AssistantEvent:
|
||||||
llm_result = await self._chat()
|
llm_result = await self._chat()
|
||||||
if llm_result.usage is None:
|
return AssistantEvent(content=llm_result.message.content or "")
|
||||||
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 "")
|
|
||||||
|
|
||||||
async def _handle_tool_calls(
|
async def _handle_tool_calls(
|
||||||
self, resolved: ResolvedMessage
|
self, resolved: ResolvedMessage
|
||||||
|
|
@ -585,7 +511,6 @@ class Agent:
|
||||||
|
|
||||||
try:
|
try:
|
||||||
start_time = time.perf_counter()
|
start_time = time.perf_counter()
|
||||||
|
|
||||||
async with self.backend as backend:
|
async with self.backend as backend:
|
||||||
result = await backend.complete(
|
result = await backend.complete(
|
||||||
model=active_model,
|
model=active_model,
|
||||||
|
|
@ -599,31 +524,19 @@ class Agent:
|
||||||
},
|
},
|
||||||
max_tokens=max_tokens,
|
max_tokens=max_tokens,
|
||||||
)
|
)
|
||||||
|
|
||||||
end_time = time.perf_counter()
|
end_time = time.perf_counter()
|
||||||
|
|
||||||
if result.usage is None:
|
if result.usage is None:
|
||||||
raise LLMResponseError(
|
raise LLMResponseError(
|
||||||
"Usage data missing in non-streaming completion response"
|
"Usage data missing in non-streaming completion response"
|
||||||
)
|
)
|
||||||
|
self._update_stats(usage=result.usage, time_seconds=end_time - start_time)
|
||||||
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
|
|
||||||
)
|
|
||||||
|
|
||||||
processed_message = self.format_handler.process_api_response_message(
|
processed_message = self.format_handler.process_api_response_message(
|
||||||
result.message
|
result.message
|
||||||
)
|
)
|
||||||
|
self.messages.append(processed_message)
|
||||||
return LLMChunk(
|
return LLMChunk(message=processed_message, usage=result.usage)
|
||||||
message=processed_message,
|
|
||||||
usage=result.usage,
|
|
||||||
finish_reason=result.finish_reason,
|
|
||||||
)
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
|
|
@ -642,7 +555,8 @@ class Agent:
|
||||||
tool_choice = self.format_handler.get_tool_choice()
|
tool_choice = self.format_handler.get_tool_choice()
|
||||||
try:
|
try:
|
||||||
start_time = time.perf_counter()
|
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 with self.backend as backend:
|
||||||
async for chunk in backend.complete_streaming(
|
async for chunk in backend.complete_streaming(
|
||||||
model=active_model,
|
model=active_model,
|
||||||
|
|
@ -656,38 +570,40 @@ class Agent:
|
||||||
},
|
},
|
||||||
max_tokens=max_tokens,
|
max_tokens=max_tokens,
|
||||||
):
|
):
|
||||||
last_chunk = chunk
|
|
||||||
processed_message = (
|
processed_message = (
|
||||||
self.format_handler.process_api_response_message(chunk.message)
|
self.format_handler.process_api_response_message(chunk.message)
|
||||||
)
|
)
|
||||||
yield LLMChunk(
|
processed_chunk = LLMChunk(
|
||||||
message=processed_message,
|
message=processed_message, usage=chunk.usage
|
||||||
usage=chunk.usage,
|
|
||||||
finish_reason=chunk.finish_reason,
|
|
||||||
)
|
)
|
||||||
|
chunk_agg += processed_chunk
|
||||||
|
usage += chunk.usage or LLMUsage()
|
||||||
|
yield processed_chunk
|
||||||
end_time = time.perf_counter()
|
end_time = time.perf_counter()
|
||||||
if last_chunk is None:
|
|
||||||
raise LLMResponseError("Streamed completion returned no chunks")
|
if chunk_agg.usage is None:
|
||||||
if last_chunk.usage is None:
|
|
||||||
raise LLMResponseError(
|
raise LLMResponseError(
|
||||||
"Usage data missing in final chunk of streamed completion"
|
"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.messages.append(chunk_agg.message)
|
||||||
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
|
|
||||||
)
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
f"API error from {provider.name} (model: {active_model.name}): {e}"
|
f"API error from {provider.name} (model: {active_model.name}): {e}"
|
||||||
) from 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(
|
async def _should_execute_tool(
|
||||||
self, tool: BaseTool, args: dict[str, Any], tool_call_id: str
|
self, tool: BaseTool, args: dict[str, Any], tool_call_id: str
|
||||||
) -> ToolDecision:
|
) -> ToolDecision:
|
||||||
|
|
|
||||||
|
|
@ -143,17 +143,13 @@ class OpenAIAdapter(APIAdapter):
|
||||||
message = LLMMessage.model_validate(data["choices"][0]["delta"])
|
message = LLMMessage.model_validate(data["choices"][0]["delta"])
|
||||||
else:
|
else:
|
||||||
raise ValueError("Invalid response data")
|
raise ValueError("Invalid response data")
|
||||||
finish_reason = data["choices"][0].get("finish_reason", None)
|
|
||||||
|
|
||||||
elif "message" in data:
|
elif "message" in data:
|
||||||
message = LLMMessage.model_validate(data["message"])
|
message = LLMMessage.model_validate(data["message"])
|
||||||
finish_reason = data["choices"][0].get("finish_reason", None)
|
|
||||||
elif "delta" in data:
|
elif "delta" in data:
|
||||||
message = LLMMessage.model_validate(data["delta"])
|
message = LLMMessage.model_validate(data["delta"])
|
||||||
finish_reason = None
|
|
||||||
else:
|
else:
|
||||||
message = LLMMessage(role=Role.assistant, content="")
|
message = LLMMessage(role=Role.assistant, content="")
|
||||||
finish_reason = None
|
|
||||||
|
|
||||||
usage_data = data.get("usage") or {}
|
usage_data = data.get("usage") or {}
|
||||||
usage = LLMUsage(
|
usage = LLMUsage(
|
||||||
|
|
@ -161,7 +157,7 @@ class OpenAIAdapter(APIAdapter):
|
||||||
completion_tokens=usage_data.get("completion_tokens", 0),
|
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:
|
class GenericBackend:
|
||||||
|
|
@ -255,7 +251,7 @@ class GenericBackend:
|
||||||
provider=self._provider.name,
|
provider=self._provider.name,
|
||||||
endpoint=url,
|
endpoint=url,
|
||||||
response=e.response,
|
response=e.response,
|
||||||
headers=dict(e.response.headers.items()),
|
headers=e.response.headers,
|
||||||
model=model.name,
|
model=model.name,
|
||||||
messages=messages,
|
messages=messages,
|
||||||
temperature=temperature,
|
temperature=temperature,
|
||||||
|
|
@ -320,7 +316,7 @@ class GenericBackend:
|
||||||
provider=self._provider.name,
|
provider=self._provider.name,
|
||||||
endpoint=url,
|
endpoint=url,
|
||||||
response=e.response,
|
response=e.response,
|
||||||
headers=dict(e.response.headers.items()),
|
headers=e.response.headers,
|
||||||
model=model.name,
|
model=model.name,
|
||||||
messages=messages,
|
messages=messages,
|
||||||
temperature=temperature,
|
temperature=temperature,
|
||||||
|
|
@ -369,7 +365,11 @@ class GenericBackend:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
DELIM_CHAR = ":"
|
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)
|
delim_index = line.find(DELIM_CHAR)
|
||||||
key = line[0:delim_index]
|
key = line[0:delim_index]
|
||||||
value = line[delim_index + 2 :]
|
value = line[delim_index + 2 :]
|
||||||
|
|
@ -404,9 +404,8 @@ class GenericBackend:
|
||||||
tool_choice=tool_choice,
|
tool_choice=tool_choice,
|
||||||
extra_headers=extra_headers,
|
extra_headers=extra_headers,
|
||||||
)
|
)
|
||||||
assert result.usage is not None, (
|
if result.usage is None:
|
||||||
"Usage should be present in non-streaming completions"
|
raise ValueError("Missing usage in non streaming completion")
|
||||||
)
|
|
||||||
|
|
||||||
return result.usage.prompt_tokens
|
return result.usage.prompt_tokens
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -205,7 +205,6 @@ class MistralBackend:
|
||||||
prompt_tokens=response.usage.prompt_tokens or 0,
|
prompt_tokens=response.usage.prompt_tokens or 0,
|
||||||
completion_tokens=response.usage.completion_tokens or 0,
|
completion_tokens=response.usage.completion_tokens or 0,
|
||||||
),
|
),
|
||||||
finish_reason=response.choices[0].finish_reason,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
except mistralai.SDKError as e:
|
except mistralai.SDKError as e:
|
||||||
|
|
@ -213,7 +212,7 @@ class MistralBackend:
|
||||||
provider=self._provider.name,
|
provider=self._provider.name,
|
||||||
endpoint=self._server_url,
|
endpoint=self._server_url,
|
||||||
response=e.raw_response,
|
response=e.raw_response,
|
||||||
headers=dict(e.raw_response.headers.items()),
|
headers=e.raw_response.headers,
|
||||||
model=model.name,
|
model=model.name,
|
||||||
messages=messages,
|
messages=messages,
|
||||||
temperature=temperature,
|
temperature=temperature,
|
||||||
|
|
@ -279,7 +278,6 @@ class MistralBackend:
|
||||||
if chunk.data.usage
|
if chunk.data.usage
|
||||||
else 0,
|
else 0,
|
||||||
),
|
),
|
||||||
finish_reason=chunk.data.choices[0].finish_reason,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
except mistralai.SDKError as e:
|
except mistralai.SDKError as e:
|
||||||
|
|
@ -287,7 +285,7 @@ class MistralBackend:
|
||||||
provider=self._provider.name,
|
provider=self._provider.name,
|
||||||
endpoint=self._server_url,
|
endpoint=self._server_url,
|
||||||
response=e.raw_response,
|
response=e.raw_response,
|
||||||
headers=dict(e.raw_response.headers.items()),
|
headers=e.raw_response.headers,
|
||||||
model=model.name,
|
model=model.name,
|
||||||
messages=messages,
|
messages=messages,
|
||||||
temperature=temperature,
|
temperature=temperature,
|
||||||
|
|
@ -325,8 +323,7 @@ class MistralBackend:
|
||||||
tool_choice=tool_choice,
|
tool_choice=tool_choice,
|
||||||
extra_headers=extra_headers,
|
extra_headers=extra_headers,
|
||||||
)
|
)
|
||||||
assert result.usage is not None, (
|
if result.usage is None:
|
||||||
"Usage should be present in non-streaming completions"
|
raise ValueError("Missing usage in non streaming completion")
|
||||||
)
|
|
||||||
|
|
||||||
return result.usage.prompt_tokens
|
return result.usage.prompt_tokens
|
||||||
|
|
|
||||||
|
|
@ -1,7 +1,9 @@
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from abc import ABC
|
from abc import ABC
|
||||||
|
from collections import OrderedDict
|
||||||
from collections.abc import Awaitable, Callable
|
from collections.abc import Awaitable, Callable
|
||||||
|
import copy
|
||||||
from enum import StrEnum, auto
|
from enum import StrEnum, auto
|
||||||
from typing import Annotated, Any, Literal
|
from typing import Annotated, Any, Literal
|
||||||
|
|
||||||
|
|
@ -185,19 +187,75 @@ class LLMMessage(BaseModel):
|
||||||
"tool_call_id": getattr(v, "tool_call_id", None),
|
"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):
|
class LLMUsage(BaseModel):
|
||||||
model_config = ConfigDict(frozen=True)
|
model_config = ConfigDict(frozen=True)
|
||||||
prompt_tokens: int = 0
|
prompt_tokens: int = 0
|
||||||
completion_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):
|
class LLMChunk(BaseModel):
|
||||||
model_config = ConfigDict(frozen=True)
|
model_config = ConfigDict(frozen=True)
|
||||||
message: LLMMessage
|
message: LLMMessage
|
||||||
finish_reason: str | None = None
|
|
||||||
usage: LLMUsage | 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):
|
class BaseEvent(BaseModel, ABC):
|
||||||
"""Abstract base class for all agent events."""
|
"""Abstract base class for all agent events."""
|
||||||
|
|
@ -209,6 +267,13 @@ class AssistantEvent(BaseEvent):
|
||||||
content: str
|
content: str
|
||||||
stopped_by_middleware: bool = False
|
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):
|
class ToolCallEvent(BaseEvent):
|
||||||
tool_name: str
|
tool_name: str
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue