d33db9fff8
Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai> Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai> Co-Authored-By: Clément Drouin <clement.drouin@mistral.ai> Co-Authored-By: Vincent Guilloux <vincent.guilloux@mistral.ai> Co-Authored-By: Clément Siriex <clement.sirieix@mistral.ai> Co-Authored-By: Kim-Adeline Miguel <kimadeline.miguel@mistral.ai> Co-Authored-By: Thaddee Tyl <thaddee.tyl@gmail.com> Co-Authored-By: David Brochart <david.brochart@gmail.com> Co-Authored-By: Joseph Guhlin <joseph.guhlin@gmail.com> Co-Authored-By: Thomas Kenbeek <thomaskenbeek@gmail.com> Co-Authored-By: Remenby31 <baptiste.cruvellier31@gmail.com>
69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from tests.mock.utils import mock_llm_chunk
|
|
from tests.stubs.fake_backend import FakeBackend
|
|
from vibe.core.agent_loop import AgentLoop
|
|
from vibe.core.config import SessionLoggingConfig, VibeConfig
|
|
|
|
|
|
@pytest.fixture
|
|
def vibe_config() -> VibeConfig:
|
|
return VibeConfig(session_logging=SessionLoggingConfig(enabled=False))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_passes_x_affinity_header_when_asking_an_answer(vibe_config: VibeConfig):
|
|
backend = FakeBackend([mock_llm_chunk(content="Response")])
|
|
agent = AgentLoop(vibe_config, backend=backend)
|
|
|
|
[_ async for _ in agent.act("Hello")]
|
|
|
|
assert len(backend.requests_extra_headers) > 0
|
|
headers = backend.requests_extra_headers[0]
|
|
assert headers is not None
|
|
assert "x-affinity" in headers
|
|
assert headers["x-affinity"] == agent.session_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_passes_x_affinity_header_when_asking_an_answer_streaming(
|
|
vibe_config: VibeConfig,
|
|
):
|
|
backend = FakeBackend([mock_llm_chunk(content="Response")])
|
|
agent = AgentLoop(vibe_config, backend=backend, enable_streaming=True)
|
|
|
|
[_ async for _ in agent.act("Hello")]
|
|
|
|
assert len(backend.requests_extra_headers) > 0
|
|
headers = backend.requests_extra_headers[0]
|
|
assert headers is not None
|
|
assert "x-affinity" in headers
|
|
assert headers["x-affinity"] == agent.session_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_updates_tokens_stats_based_on_backend_response(vibe_config: VibeConfig):
|
|
chunk = mock_llm_chunk(content="Response", prompt_tokens=100, completion_tokens=50)
|
|
backend = FakeBackend([chunk])
|
|
agent = AgentLoop(vibe_config, backend=backend)
|
|
|
|
[_ async for _ in agent.act("Hello")]
|
|
|
|
assert agent.stats.context_tokens == 150
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_updates_tokens_stats_based_on_backend_response_streaming(
|
|
vibe_config: VibeConfig,
|
|
):
|
|
final_chunk = mock_llm_chunk(
|
|
content="Complete", prompt_tokens=200, completion_tokens=75
|
|
)
|
|
backend = FakeBackend([final_chunk])
|
|
agent = AgentLoop(vibe_config, backend=backend, enable_streaming=True)
|
|
|
|
[_ async for _ in agent.act("Hello")]
|
|
|
|
assert agent.stats.context_tokens == 275
|