e607ccbb00
Co-authored-by: Clément Drouin <clement.drouin@mistral.ai> Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai> Co-authored-by: Cyprien <courtot.c@gmail.com> Co-authored-by: Guillaume LE GOFF <guillaume.lgf@gmail.com> Co-authored-by: Jean Burellier <sheplu@users.noreply.github.com> Co-authored-by: Kim-Adeline Miguel <51720070+kimadeline@users.noreply.github.com> Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: Nelson PROIA <144663685+Nelson-PROIA@users.noreply.github.com> Co-authored-by: Paul VEZIA <166131032+le-codeur-rapide@users.noreply.github.com> Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai> Co-authored-by: Quentin <quentin.torroba@mistral.ai> Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com> Co-authored-by: josephine-delas <57808586+josephine-delas@users.noreply.github.com> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
import respx
|
|
|
|
from tests import constants as c
|
|
from tests.agent_loop.e2e.providers.anthropic import Mocks
|
|
from tests.agent_loop.e2e.providers.base import ProviderAPI
|
|
from tests.conftest import build_test_vibe_config
|
|
from vibe.core.config import ModelConfig, ProviderConfig, VibeConfig
|
|
from vibe.core.types import Backend
|
|
|
|
__all__ = ["API", "Mocks", "e2e_config"]
|
|
|
|
|
|
def e2e_config(**overrides: Any) -> VibeConfig:
|
|
provider = ProviderConfig(
|
|
name="vertex",
|
|
api_base="",
|
|
api_key_env_var="VERTEX_API_KEY",
|
|
api_style="vertex-anthropic",
|
|
backend=Backend.GENERIC,
|
|
project_id=c.VERTEX_PROJECT_ID,
|
|
region=c.VERTEX_REGION,
|
|
)
|
|
models = [ModelConfig(name=c.VERTEX_MODEL, provider="vertex", alias="vertex")]
|
|
return build_test_vibe_config(
|
|
active_model="vertex", models=models, providers=[provider], **overrides
|
|
)
|
|
|
|
|
|
class API(ProviderAPI):
|
|
"""Vertex splits rawPredict and streamRawPredict across two wire paths."""
|
|
|
|
base_url = c.VERTEX_BASE_URL
|
|
post_path = c.VERTEX_RAW_PREDICT_PATH
|
|
|
|
def setup_router(self, router: respx.MockRouter) -> None:
|
|
super().setup_router(router)
|
|
self._stream_route = router.post(c.VERTEX_STREAM_PREDICT_PATH)
|
|
|
|
@classmethod
|
|
def setup_monkeypatch(cls, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
# google.auth is the only external dependency Vertex pulls in; stub the
|
|
# credential lookup so the real token-refresh logic runs without network.
|
|
creds = MagicMock()
|
|
creds.valid = True
|
|
creds.token = "fake-vertex-token"
|
|
monkeypatch.setattr(
|
|
"vibe.core.llm.backend.vertex.google.auth.default",
|
|
lambda **_kwargs: (creds, c.VERTEX_PROJECT_ID),
|
|
)
|
|
|
|
def reply_stream(self, chunks: list[bytes]) -> None:
|
|
self._stream_route.mock(return_value=self._stream_response(chunks))
|
|
|
|
def reply_streams(self, *chunk_lists: list[bytes]) -> None:
|
|
self._stream_route.mock(
|
|
side_effect=[self._stream_response(chunks) for chunks in chunk_lists]
|
|
)
|
|
|
|
@property
|
|
def request_json(self) -> dict[str, Any]:
|
|
route = self._stream_route if self._stream_route.calls else self.route
|
|
return json.loads(route.calls.last.request.content)
|