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

v1.2.2

Browse source 
Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai>
Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai>
This commit is contained in:
Mathias Gesbert 2025-12-22 13:33:20 +01:00 • committed by Mathias Gesbert
parent 402e898f39
commit 2e1e15120d
32 changed files with 391 additions and 549 deletions
Download patch file Download diff file Expand all files Collapse all files

70
tests/stubs/fake_backend.py
View file

@ -1,9 +1,10 @@
from __future__ import annotations
from collections.abc import AsyncGenerator, Callable, Iterable
from typing import cast
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:
@ -15,18 +16,47 @@ class FakeBackend:
def __init__(
self,
results: Iterable[LLMChunk] | None = None,
chunks: LLMChunk
| Iterable[LLMChunk]
| Iterable[Iterable[LLMChunk]]
| None = None,
*,
token_counter: Callable[[list[LLMMessage]], int] | None = None,
exception_to_raise: Exception | 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_extra_headers: 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]]
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
def requests_messages(self) -> list[list[LLMMessage]]:
return self._requests_messages
@ -61,12 +91,15 @@ class FakeBackend:
self._requests_messages.append(messages)
self._requests_extra_headers.append(extra_headers)
if self._chunks:
chunk = self._chunks.pop(0)
if not self._chunks:
chunk = chunk.model_copy(update={"finish_reason": "stop"})
return chunk
return mock_llm_chunk(content="", finish_reason="stop")
if self._streams:
stream = self._streams.pop(0)
chunk_agg = LLMChunk(message=LLMMessage(role=Role.assistant))
for chunk in stream:
chunk_agg += chunk
return chunk_agg
return mock_llm_chunk(content="")
async def complete_streaming(
self,
@ -84,22 +117,13 @@ class FakeBackend:
self._requests_messages.append(messages)
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
if has_final_chunk:
break
if not has_final_chunk:
yield mock_llm_chunk(content="", finish_reason="stop")
async def count_tokens(
self,