uvos/vibe
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Initial commit

Browse source 
Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai>
Co-Authored-By: Laure Hugo <laure.hugo@mistral.ai>
Co-Authored-By: Benjamin Trom <benjamin.trom@mistral.ai>
Co-Authored-By: Mathias Gesbert <mathias.gesbert@ext.mistral.ai>
Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai>
Co-Authored-By: Clément Drouin <clement.drouin@mistral.ai>
Co-Authored-By: Vincent Guilloux <vincent.guilloux@mistral.ai>
Co-Authored-By: Valentin Berard <val@mistral.ai>
Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Quentin Torroba 2025-12-09 13:13:22 +01:00 • committed by Quentin Torroba
commit fa15fc977b
200 changed files with 30484 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

711
tests/test_agent_stats.py Normal file
View file

@ -0,0 +1,711 @@
from __future__ import annotations
from collections.abc import Callable
import pytest
from tests.mock.utils import mock_llm_chunk
from tests.stubs.fake_backend import FakeBackend
from vibe.core.agent import Agent
from vibe.core.config import (
Backend,
ModelConfig,
ProviderConfig,
SessionLoggingConfig,
VibeConfig,
)
from vibe.core.tools.base import BaseToolConfig, ToolPermission
from vibe.core.types import (
AgentStats,
AssistantEvent,
CompactEndEvent,
CompactStartEvent,
FunctionCall,
LLMMessage,
Role,
ToolCall,
)
def make_config(
*,
system_prompt_id: str = "tests",
active_model: str = "devstral-latest",
input_price: float = 0.4,
output_price: float = 2.0,
disable_logging: bool = True,
auto_compact_threshold: int = 0,
include_project_context: bool = False,
include_prompt_detail: bool = False,
enabled_tools: list[str] | None = None,
todo_permission: ToolPermission = ToolPermission.ALWAYS,
) -> VibeConfig:
models = [
ModelConfig(
name="mistral-vibe-cli-latest",
provider="mistral",
alias="devstral-latest",
input_price=input_price,
output_price=output_price,
),
ModelConfig(
name="devstral-small-latest",
provider="mistral",
alias="devstral-small",
input_price=0.1,
output_price=0.3,
),
ModelConfig(
name="strawberry",
provider="lechat",
alias="strawberry",
input_price=2.5,
output_price=10.0,
),
]
providers = [
ProviderConfig(
name="mistral",
api_base="https://api.mistral.ai/v1",
api_key_env_var="MISTRAL_API_KEY",
backend=Backend.MISTRAL,
),
ProviderConfig(
name="lechat",
api_base="https://api.mistral.ai/v1",
api_key_env_var="LECHAT_API_KEY",
backend=Backend.MISTRAL,
),
]
return VibeConfig(
session_logging=SessionLoggingConfig(enabled=not disable_logging),
auto_compact_threshold=auto_compact_threshold,
system_prompt_id=system_prompt_id,
include_project_context=include_project_context,
include_prompt_detail=include_prompt_detail,
active_model=active_model,
models=models,
providers=providers,
enabled_tools=enabled_tools or [],
tools={"todo": BaseToolConfig(permission=todo_permission)},
)
@pytest.fixture
def observer_capture() -> tuple[list[LLMMessage], Callable[[LLMMessage], None]]:
observed: list[LLMMessage] = []
def observer(msg: LLMMessage) -> None:
observed.append(msg)
return observed, observer
class TestAgentStatsHelpers:
def test_update_pricing(self) -> None:
stats = AgentStats()
stats.update_pricing(1.5, 3.0)
assert stats.input_price_per_million == 1.5
assert stats.output_price_per_million == 3.0
def test_reset_context_state_preserves_cumulative(self) -> None:
stats = AgentStats(
steps=5,
session_prompt_tokens=1000,
session_completion_tokens=500,
tool_calls_succeeded=3,
tool_calls_failed=1,
context_tokens=800,
last_turn_prompt_tokens=100,
last_turn_completion_tokens=50,
last_turn_duration=1.5,
tokens_per_second=33.3,
input_price_per_million=0.4,
output_price_per_million=2.0,
)
stats.reset_context_state()
assert stats.steps == 5
assert stats.session_prompt_tokens == 1000
assert stats.session_completion_tokens == 500
assert stats.tool_calls_succeeded == 3
assert stats.tool_calls_failed == 1
assert stats.input_price_per_million == 0.4
assert stats.output_price_per_million == 2.0
assert stats.context_tokens == 0
assert stats.last_turn_prompt_tokens == 0
assert stats.last_turn_completion_tokens == 0
assert stats.last_turn_duration == 0.0
assert stats.tokens_per_second == 0.0
def test_session_cost_computed_from_current_pricing(self) -> None:
stats = AgentStats(
session_prompt_tokens=1_000_000,
session_completion_tokens=500_000,
input_price_per_million=1.0,
output_price_per_million=2.0,
)
# Cost = 1M * $1/M + 0.5M * $2/M = $1 + $1 = $2
assert stats.session_cost == 2.0
stats.update_pricing(2.0, 4.0)
# Cost = 1M * $2/M + 0.5M * $4/M = $2 + $2 = $4
assert stats.session_cost == 4.0
class TestReloadPreservesStats:
@pytest.mark.asyncio
async def test_reload_preserves_session_tokens(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="First response", finish_reason="stop")
])
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("Hello"):
pass
old_session_prompt = agent.stats.session_prompt_tokens
old_session_completion = agent.stats.session_completion_tokens
assert old_session_prompt > 0
assert old_session_completion > 0
await agent.reload_with_initial_messages()
assert agent.stats.session_prompt_tokens == old_session_prompt
assert agent.stats.session_completion_tokens == old_session_completion
@pytest.mark.asyncio
async def test_reload_preserves_tool_call_stats(self) -> None:
backend = FakeBackend([
mock_llm_chunk(
content="Calling tool",
tool_calls=[
ToolCall(
id="tc1",
function=FunctionCall(
name="todo", arguments='{"action": "read"}'
),
)
],
),
mock_llm_chunk(content="Done", finish_reason="stop"),
])
config = make_config(enabled_tools=["todo"])
agent = Agent(config, auto_approve=True, backend=backend)
async for _ in agent.act("Check todos"):
pass
assert agent.stats.tool_calls_succeeded == 1
assert agent.stats.tool_calls_agreed == 1
await agent.reload_with_initial_messages()
assert agent.stats.tool_calls_succeeded == 1
assert agent.stats.tool_calls_agreed == 1
@pytest.mark.asyncio
async def test_reload_preserves_steps(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="R1", finish_reason="stop"),
mock_llm_chunk(content="R2", finish_reason="stop"),
])
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("First"):
pass
async for _ in agent.act("Second"):
pass
old_steps = agent.stats.steps
assert old_steps >= 2
await agent.reload_with_initial_messages()
assert agent.stats.steps == old_steps
@pytest.mark.asyncio
async def test_reload_preserves_context_tokens_when_messages_preserved(
self,
) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
agent = Agent(make_config(), backend=backend)
[_ async for _ in agent.act("Hello")]
assert agent.stats.context_tokens > 0
initial_context_tokens = agent.stats.context_tokens
assert len(agent.messages) > 1
await agent.reload_with_initial_messages()
assert len(agent.messages) > 1
assert agent.stats.context_tokens == initial_context_tokens
@pytest.mark.asyncio
async def test_reload_resets_context_tokens_when_no_messages(self) -> None:
backend = FakeBackend([])
agent = Agent(make_config(), backend=backend)
assert len(agent.messages) == 1
assert agent.stats.context_tokens == 0
await agent.reload_with_initial_messages()
assert len(agent.messages) == 1
assert agent.stats.context_tokens == 0
@pytest.mark.asyncio
async def test_reload_resets_context_tokens_when_system_prompt_changes(
self,
) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
config1 = make_config(system_prompt_id="tests")
config2 = make_config(system_prompt_id="cli")
agent = Agent(config1, backend=backend)
[_ async for _ in agent.act("Hello")]
assert agent.stats.context_tokens > 0
assert len(agent.messages) > 1
await agent.reload_with_initial_messages(config=config2)
assert len(agent.messages) > 1
assert agent.stats.context_tokens == 0
@pytest.mark.asyncio
async def test_reload_updates_pricing_from_new_model(self, monkeypatch) -> None:
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
config_mistral = make_config(active_model="devstral-latest")
agent = Agent(config_mistral, backend=backend)
async for _ in agent.act("Hello"):
pass
assert agent.stats.input_price_per_million == 0.4
assert agent.stats.output_price_per_million == 2.0
config_other = make_config(active_model="strawberry")
await agent.reload_with_initial_messages(config=config_other)
assert agent.stats.input_price_per_million == 2.5
assert agent.stats.output_price_per_million == 10.0
@pytest.mark.asyncio
async def test_reload_accumulates_tokens_across_configs(self, monkeypatch) -> None:
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
backend = FakeBackend([
mock_llm_chunk(content="First", finish_reason="stop"),
mock_llm_chunk(content="After reload", finish_reason="stop"),
])
config1 = make_config(active_model="devstral-latest")
agent = Agent(config1, backend=backend)
async for _ in agent.act("Hello"):
pass
tokens_after_first = (
agent.stats.session_prompt_tokens + agent.stats.session_completion_tokens
)
config2 = make_config(active_model="strawberry")
await agent.reload_with_initial_messages(config=config2)
async for _ in agent.act("Continue"):
pass
tokens_after_second = (
agent.stats.session_prompt_tokens + agent.stats.session_completion_tokens
)
assert tokens_after_second > tokens_after_first
class TestReloadPreservesMessages:
@pytest.mark.asyncio
async def test_reload_preserves_conversation_messages(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("Hello"):
pass
assert len(agent.messages) == 3
old_user_content = agent.messages[1].content
old_assistant_content = agent.messages[2].content
await agent.reload_with_initial_messages()
assert len(agent.messages) == 3
assert agent.messages[0].role == Role.system
assert agent.messages[1].role == Role.user
assert agent.messages[1].content == old_user_content
assert agent.messages[2].role == Role.assistant
assert agent.messages[2].content == old_assistant_content
@pytest.mark.asyncio
async def test_reload_updates_system_prompt_preserves_rest(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
config1 = make_config(system_prompt_id="tests")
agent = Agent(config1, backend=backend)
async for _ in agent.act("Hello"):
pass
old_system = agent.messages[0].content
old_user = agent.messages[1].content
config2 = make_config(system_prompt_id="cli")
await agent.reload_with_initial_messages(config=config2)
assert agent.messages[0].content != old_system
assert agent.messages[1].content == old_user
@pytest.mark.asyncio
async def test_reload_with_no_messages_stays_empty(self) -> None:
backend = FakeBackend([])
agent = Agent(make_config(), backend=backend)
assert len(agent.messages) == 1
await agent.reload_with_initial_messages()
assert len(agent.messages) == 1
assert agent.messages[0].role == Role.system
@pytest.mark.asyncio
async def test_reload_notifies_observer_with_all_messages(
self, observer_capture
) -> None:
observed, observer = observer_capture
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
agent = Agent(make_config(), message_observer=observer, backend=backend)
async for _ in agent.act("Hello"):
pass
observed.clear()
await agent.reload_with_initial_messages()
assert len(observed) == 3
assert observed[0].role == Role.system
assert observed[1].role == Role.user
assert observed[2].role == Role.assistant
class TestCompactStatsHandling:
@pytest.mark.asyncio
async def test_compact_preserves_cumulative_stats(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="First response", finish_reason="stop"),
mock_llm_chunk(content="<summary>", finish_reason="stop"),
])
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("Build something"):
pass
tokens_before_compact = agent.stats.session_prompt_tokens
completions_before = agent.stats.session_completion_tokens
steps_before = agent.stats.steps
await agent.compact()
# Cumulative stats include the compact turn
assert agent.stats.session_prompt_tokens > tokens_before_compact
assert agent.stats.session_completion_tokens > completions_before
assert agent.stats.steps > steps_before
@pytest.mark.asyncio
async def test_compact_updates_context_tokens(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Long response " * 100, finish_reason="stop"),
mock_llm_chunk(content="<summary>", finish_reason="stop"),
])
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("Do something complex"):
pass
context_before = agent.stats.context_tokens
await agent.compact()
assert agent.stats.context_tokens < context_before
@pytest.mark.asyncio
async def test_compact_preserves_tool_call_stats(self) -> None:
backend = FakeBackend([
mock_llm_chunk(
content="Using tool",
tool_calls=[
ToolCall(
id="tc1",
function=FunctionCall(
name="todo", arguments='{"action": "read"}'
),
)
],
),
mock_llm_chunk(content="Done", finish_reason="stop"),
mock_llm_chunk(content="<summary>", finish_reason="stop"),
])
config = make_config(enabled_tools=["todo"])
agent = Agent(config, auto_approve=True, backend=backend)
async for _ in agent.act("Check todos"):
pass
assert agent.stats.tool_calls_succeeded == 1
await agent.compact()
assert agent.stats.tool_calls_succeeded == 1
@pytest.mark.asyncio
async def test_compact_resets_session_id(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Long response " * 100, finish_reason="stop"),
mock_llm_chunk(content="<summary>", finish_reason="stop"),
])
agent = Agent(make_config(disable_logging=False), backend=backend)
original_session_id = agent.session_id
original_logger_session_id = agent.interaction_logger.session_id
assert agent.session_id == original_logger_session_id
async for _ in agent.act("Do something complex"):
pass
await agent.compact()
assert agent.session_id != original_session_id
assert agent.session_id == agent.interaction_logger.session_id
class TestAutoCompactIntegration:
@pytest.mark.asyncio
async def test_auto_compact_triggers_and_preserves_stats(self) -> None:
observed: list[tuple[Role, str | None]] = []
def observer(msg: LLMMessage) -> None:
observed.append((msg.role, msg.content))
backend = FakeBackend([
mock_llm_chunk(content="<summary>", finish_reason="stop"),
mock_llm_chunk(content="<final>", finish_reason="stop"),
])
cfg = VibeConfig(
session_logging=SessionLoggingConfig(enabled=False),
auto_compact_threshold=1,
)
agent = Agent(cfg, message_observer=observer, backend=backend)
agent.stats.context_tokens = 2
events = [ev async for ev in agent.act("Hello")]
assert len(events) == 3
assert isinstance(events[0], CompactStartEvent)
assert isinstance(events[1], CompactEndEvent)
assert isinstance(events[2], AssistantEvent)
start: CompactStartEvent = events[0]
end: CompactEndEvent = events[1]
final: AssistantEvent = events[2]
assert start.current_context_tokens == 2
assert start.threshold == 1
assert end.old_context_tokens == 2
assert end.new_context_tokens >= 1
assert final.content == "<final>"
roles = [r for r, _ in observed]
assert roles == [Role.system, Role.user, Role.assistant]
assert (
observed[1][1] is not None
and "Last request from user was: Hello" in observed[1][1]
)
class TestClearHistoryFullReset:
@pytest.mark.asyncio
async def test_clear_history_fully_resets_stats(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("Hello"):
pass
assert agent.stats.session_prompt_tokens > 0
assert agent.stats.steps > 0
await agent.clear_history()
assert agent.stats.session_prompt_tokens == 0
assert agent.stats.session_completion_tokens == 0
assert agent.stats.steps == 0
@pytest.mark.asyncio
async def test_clear_history_preserves_pricing(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
config = make_config(input_price=0.4, output_price=2.0)
agent = Agent(config, backend=backend)
async for _ in agent.act("Hello"):
pass
await agent.clear_history()
assert agent.stats.input_price_per_million == 0.4
assert agent.stats.output_price_per_million == 2.0
@pytest.mark.asyncio
async def test_clear_history_removes_messages(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("Hello"):
pass
assert len(agent.messages) == 3
await agent.clear_history()
assert len(agent.messages) == 1
assert agent.messages[0].role == Role.system
@pytest.mark.asyncio
async def test_clear_history_resets_session_id(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
agent = Agent(make_config(disable_logging=False), backend=backend)
original_session_id = agent.session_id
original_logger_session_id = agent.interaction_logger.session_id
assert agent.session_id == original_logger_session_id
async for _ in agent.act("Hello"):
pass
await agent.clear_history()
assert agent.session_id != original_session_id
assert agent.session_id == agent.interaction_logger.session_id
class TestStatsEdgeCases:
@pytest.mark.asyncio
async def test_session_cost_approximation_on_model_change(
self, monkeypatch
) -> None:
monkeypatch.setenv("LECHAT_API_KEY", "mock-key")
backend = FakeBackend([
mock_llm_chunk(content="Response", finish_reason="stop")
])
config1 = make_config(active_model="devstral-latest")
agent = Agent(config1, backend=backend)
async for _ in agent.act("Hello"):
pass
cost_before = agent.stats.session_cost
config2 = make_config(active_model="strawberry")
await agent.reload_with_initial_messages(config=config2)
cost_after = agent.stats.session_cost
assert cost_after > cost_before
@pytest.mark.asyncio
async def test_multiple_reloads_accumulate_correctly(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="R1", finish_reason="stop"),
mock_llm_chunk(content="R2", finish_reason="stop"),
mock_llm_chunk(content="R3", finish_reason="stop"),
])
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("First"):
pass
tokens1 = agent.stats.session_total_llm_tokens
await agent.reload_with_initial_messages()
async for _ in agent.act("Second"):
pass
tokens2 = agent.stats.session_total_llm_tokens
await agent.reload_with_initial_messages()
async for _ in agent.act("Third"):
pass
tokens3 = agent.stats.session_total_llm_tokens
assert tokens1 < tokens2 < tokens3
@pytest.mark.asyncio
async def test_compact_then_reload_preserves_both(self) -> None:
backend = FakeBackend([
mock_llm_chunk(content="Initial response", finish_reason="stop"),
mock_llm_chunk(content="<summary>", finish_reason="stop"),
mock_llm_chunk(content="After reload", finish_reason="stop"),
])
agent = Agent(make_config(), backend=backend)
async for _ in agent.act("Build something"):
pass
await agent.compact()
tokens_after_compact = agent.stats.session_prompt_tokens
await agent.reload_with_initial_messages()
assert agent.stats.session_prompt_tokens == tokens_after_compact
async for _ in agent.act("Continue"):
pass
assert agent.stats.session_prompt_tokens > tokens_after_compact
@pytest.mark.asyncio
async def test_reload_without_config_preserves_current(self) -> None:
backend = FakeBackend([])
original_config = make_config(active_model="devstral-latest")
agent = Agent(original_config, backend=backend)
await agent.reload_with_initial_messages(config=None)
assert agent.config.active_model == "devstral-latest"
@pytest.mark.asyncio
async def test_reload_with_new_config_updates_it(self) -> None:
backend = FakeBackend([])
original_config = make_config(active_model="devstral-latest")
agent = Agent(original_config, backend=backend)
new_config = make_config(active_model="devstral-small")
await agent.reload_with_initial_messages(config=new_config)
assert agent.config.active_model == "devstral-small"