Initial commit

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>

tests/stubs/fake_backend.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator, Callable, Iterable
+
+from tests.mock.utils import mock_llm_chunk
+from vibe.core.types import LLMChunk, LLMMessage
+
+
+class FakeBackend:
+    """Minimal async backend stub to drive Agent.act without network.
+
+    Provide a finite sequence of LLMResult objects to be returned by
+    `complete`. When exhausted, returns an empty assistant message.
+    """
+
+    def __init__(
+        self,
+        results: Iterable[LLMChunk] | None = None,
+        *,
+        token_counter: Callable[[list[LLMMessage]], int] | None = None,
+        exception_to_raise: Exception | None = None,
+    ) -> None:
+        self._chunks = list(results or [])
+        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
+
+    @property
+    def requests_messages(self) -> list[list[LLMMessage]]:
+        return self._requests_messages
+
+    @property
+    def requests_extra_headers(self) -> list[dict[str, str] | None]:
+        return self._requests_extra_headers
+
+    @staticmethod
+    def _default_token_counter(messages: list[LLMMessage]) -> int:
+        return 1
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        return None
+
+    async def complete(
+        self,
+        *,
+        model,
+        messages,
+        temperature,
+        tools,
+        tool_choice,
+        extra_headers,
+        max_tokens,
+    ) -> LLMChunk:
+        if self._exception_to_raise:
+            raise self._exception_to_raise
+
+        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")
+
+    async def complete_streaming(
+        self,
+        *,
+        model,
+        messages,
+        temperature,
+        tools,
+        tool_choice,
+        extra_headers,
+        max_tokens,
+    ) -> AsyncGenerator[LLMChunk]:
+        if self._exception_to_raise:
+            raise self._exception_to_raise
+
+        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
+
+            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,
+        *,
+        model,
+        messages,
+        temperature=0.0,
+        tools,
+        tool_choice=None,
+        extra_headers,
+    ) -> int:
+        self._count_tokens_calls.append(list(messages))
+        return self._token_counter(messages)