Co-authored-by: Maxime Dolores <maxime.dolores@ext.mistral.ai>
Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com>
Co-authored-by: Peter Evers <peter.evers@mistral.ai>
Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai>
Co-authored-by: Quentin <quentin.torroba@mistral.ai>
Co-authored-by: Val <102326092+vdeva@users.noreply.github.com>
Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com>
Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Clément Drouin 2026-05-29 16:17:54 +02:00 committed by GitHub
parent 198277af3f
commit ad0d5c9520
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
61 changed files with 969 additions and 461 deletions

View file

@ -1,6 +1,6 @@
from __future__ import annotations
from collections.abc import AsyncGenerator, Callable, Iterable, Sequence
from collections.abc import AsyncGenerator, Iterable
from typing import cast
from tests.mock.utils import mock_llm_chunk
@ -21,7 +21,6 @@ class FakeBackend:
| Iterable[Iterable[LLMChunk]]
| 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.
@ -36,8 +35,6 @@ 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
self._streams: list[list[LLMChunk]]
@ -70,10 +67,6 @@ class FakeBackend:
def requests_metadata(self) -> list[dict[str, str] | None]:
return self._requests_metadata
@staticmethod
def _default_token_counter(messages: Sequence[LLMMessage]) -> int:
return 1
async def __aenter__(self):
return self
@ -133,20 +126,3 @@ class FakeBackend:
stream = [mock_llm_chunk(content="")]
for chunk in stream:
yield chunk
async def count_tokens(
self,
*,
model,
messages,
temperature=0.0,
tools,
tool_choice=None,
extra_headers,
metadata=None,
) -> int:
self._requests_messages.append(list(messages))
self._requests_extra_headers.append(extra_headers)
self._requests_metadata.append(metadata)
self._count_tokens_calls.append(list(messages))
return self._token_counter(messages)