Co-authored-by: Carlo <carloantonio.patti@mistral.ai>
Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai>
Co-authored-by: Clement Sirieix <clem.sirieix@gmail.com>
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: Thomas Kenbeek <thomas.kenbeek@mistral.ai>
Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Quentin 2026-02-27 17:31:58 +01:00 committed by GitHub
parent a560a47ce8
commit 5d2e01a6d7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
139 changed files with 7152 additions and 1457 deletions

View file

@ -1,6 +1,6 @@
from __future__ import annotations
from collections.abc import AsyncGenerator, Callable, Iterable
from collections.abc import AsyncGenerator, Callable, Iterable, Sequence
from typing import cast
from tests.mock.utils import mock_llm_chunk
@ -21,7 +21,7 @@ class FakeBackend:
| Iterable[Iterable[LLMChunk]]
| None = None,
*,
token_counter: Callable[[list[LLMMessage]], int] | None = None,
token_counter: Callable[[Sequence[LLMMessage]], int] | None = None,
exception_to_raise: Exception | None = None,
) -> None:
"""Fake backend that will output the given chunks in the order they are given.
@ -35,6 +35,7 @@ class FakeBackend:
"""
self._requests_messages: list[list[LLMMessage]] = []
self._requests_extra_headers: list[dict[str, str] | None] = []
self._requests_metadata: list[dict[str, str] | None] = []
self._count_tokens_calls: list[list[LLMMessage]] = []
self._token_counter = token_counter or self._default_token_counter
self._exception_to_raise = exception_to_raise
@ -65,8 +66,12 @@ class FakeBackend:
def requests_extra_headers(self) -> list[dict[str, str] | None]:
return self._requests_extra_headers
@property
def requests_metadata(self) -> list[dict[str, str] | None]:
return self._requests_metadata
@staticmethod
def _default_token_counter(messages: list[LLMMessage]) -> int:
def _default_token_counter(messages: Sequence[LLMMessage]) -> int:
return 1
async def __aenter__(self):
@ -85,12 +90,14 @@ class FakeBackend:
tool_choice,
extra_headers,
max_tokens,
metadata=None,
) -> LLMChunk:
if self._exception_to_raise:
raise self._exception_to_raise
self._requests_messages.append(messages)
self._requests_extra_headers.append(extra_headers)
self._requests_metadata.append(metadata)
if self._streams:
stream = self._streams.pop(0)
@ -111,12 +118,14 @@ class FakeBackend:
tool_choice,
extra_headers,
max_tokens,
metadata=None,
) -> AsyncGenerator[LLMChunk]:
if self._exception_to_raise:
raise self._exception_to_raise
self._requests_messages.append(messages)
self._requests_extra_headers.append(extra_headers)
self._requests_metadata.append(metadata)
if self._streams:
stream = list(self._streams.pop(0))
@ -134,6 +143,7 @@ class FakeBackend:
tools,
tool_choice=None,
extra_headers,
metadata=None,
) -> int:
self._count_tokens_calls.append(list(messages))
return self._token_counter(messages)