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>
65 lines
2 KiB
Python
65 lines
2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from tests import constants as c
|
|
from tests.agent_loop.e2e.providers.base import ProviderAPI
|
|
from tests.backend.data import Chunk, JsonResponse
|
|
from tests.backend.data.reasoning import (
|
|
reasoning_message,
|
|
reasoning_text_stream,
|
|
reasoning_thinking_message,
|
|
reasoning_thinking_tool_use_stream,
|
|
reasoning_tool_use,
|
|
)
|
|
from tests.conftest import build_test_vibe_config
|
|
from vibe.core.config import ModelConfig, ProviderConfig, ThinkingLevel, VibeConfig
|
|
from vibe.core.types import Backend
|
|
|
|
|
|
def e2e_config(*, thinking: ThinkingLevel = "off", **overrides: Any) -> VibeConfig:
|
|
provider = ProviderConfig(
|
|
name="reasoning",
|
|
api_base=c.REASONING_BASE_URL,
|
|
api_key_env_var="REASONING_API_KEY",
|
|
api_style="reasoning",
|
|
backend=Backend.GENERIC,
|
|
)
|
|
models = [
|
|
ModelConfig(
|
|
name="reasoning-test",
|
|
provider="reasoning",
|
|
alias="reasoning",
|
|
thinking=thinking,
|
|
)
|
|
]
|
|
return build_test_vibe_config(
|
|
active_model="reasoning", models=models, providers=[provider], **overrides
|
|
)
|
|
|
|
|
|
class API(ProviderAPI):
|
|
base_url = c.REASONING_BASE_URL
|
|
post_path = c.REASONING_COMPLETIONS_PATH
|
|
|
|
|
|
class Mocks:
|
|
def answer(self, text: str) -> JsonResponse:
|
|
return reasoning_message(text)
|
|
|
|
def text_stream(self, text: str) -> list[Chunk]:
|
|
return reasoning_text_stream(text)
|
|
|
|
def tool_call(self, name: str, arguments: dict[str, Any]) -> JsonResponse:
|
|
return reasoning_tool_use(name, json.dumps(arguments))
|
|
|
|
def reasoning_answer(self, text: str, reasoning: str) -> JsonResponse:
|
|
return reasoning_thinking_message(text, reasoning)
|
|
|
|
def reasoning_tool_call_stream(
|
|
self, name: str, arguments: dict[str, Any], *, reasoning: str
|
|
) -> list[Chunk]:
|
|
return reasoning_thinking_tool_use_stream(
|
|
name, json.dumps(arguments), reasoning=reasoning
|
|
)
|