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>
This commit is contained in:
Quentin Torroba
Quentin Torroba
commit
fa15fc977b
200 changed files with 30484 additions and 0 deletions
0
tests/backend/__init__.py
Normal file
0
tests/backend/__init__.py
Normal file
6
tests/backend/data/__init__.py
Normal file
6
tests/backend/data/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
Url = str
|
||||
JsonResponse = dict
|
||||
ResultData = dict
|
||||
Chunk = bytes
|
||||
183
tests/backend/data/fireworks.py
Normal file
183
tests/backend/data/fireworks.py
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from tests.backend.data import Chunk, JsonResponse, ResultData, Url
|
||||
|
||||
SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
|
||||
(
|
||||
"https://api.fireworks.ai",
|
||||
{
|
||||
"id": "fake_id_1234",
|
||||
"object": "chat.completion",
|
||||
"created": 1234567890,
|
||||
"model": "accounts/fireworks/models/glm-4p5",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "Some content",
|
||||
"reasoning_content": "Some reasoning content",
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"total_tokens": 300,
|
||||
"completion_tokens": 200,
|
||||
"prompt_tokens_details": {"cached_tokens": 0},
|
||||
},
|
||||
},
|
||||
{
|
||||
"message": "Some content",
|
||||
"finish_reason": "stop",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"total_tokens": 300,
|
||||
"completion_tokens": 200,
|
||||
},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
TOOL_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
|
||||
(
|
||||
"https://api.fireworks.ai",
|
||||
{
|
||||
"id": "fake_id_1234",
|
||||
"object": "chat.completion",
|
||||
"created": 1234567890,
|
||||
"model": "accounts/fireworks/models/glm-4p5",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"reasoning_content": "Some reasoning content",
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "fake_id_5678",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "some_tool",
|
||||
"arguments": '{"some_argument": "some_argument_value"}',
|
||||
},
|
||||
"name": None,
|
||||
}
|
||||
],
|
||||
},
|
||||
"finish_reason": "tool_calls",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 200,
|
||||
"prompt_tokens_details": {"cached_tokens": 0},
|
||||
},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": "tool_calls",
|
||||
"tool_calls": [
|
||||
{
|
||||
"name": "some_tool",
|
||||
"arguments": '{"some_argument": "some_argument_value"}',
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
STREAMED_SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]]] = [
|
||||
(
|
||||
"https://api.fireworks.ai",
|
||||
[
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"accounts/fireworks/models/glm-4p5","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}],"usage":null}',
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"accounts/fireworks/models/glm-4p5","choices":[{"index":0,"delta":{"reasoning_content":"Some reasoning content"},"finish_reason":null}],"usage":null}',
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"accounts/fireworks/models/glm-4p5","choices":[{"index":0,"delta":{"content":"Some content"},"finish_reason":null}],"usage":null}',
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"accounts/fireworks/models/glm-4p5","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":100,"total_tokens":300,"completion_tokens":200,"prompt_tokens_details":{"cached_tokens":0}}}',
|
||||
rb"data: [DONE]",
|
||||
],
|
||||
[
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": None,
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": None,
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "Some content",
|
||||
"finish_reason": None,
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": "stop",
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
|
||||
},
|
||||
],
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
STREAMED_TOOL_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]]] = [
|
||||
(
|
||||
"https://api.fireworks.ai",
|
||||
[
|
||||
rb'data: {"id": "fake_id_1234","object": "chat.completion.chunk","created": 1234567890,"model": "accounts/fireworks/models/glm-4p5","choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": null}],"usage": null}',
|
||||
rb'data: {"id": "fake_id_1234","object": "chat.completion.chunk","created": 1234567890,"model": "accounts/fireworks/models/glm-4p5","choices": [{"index": 0,"delta": {"reasoning_content": "Some reasoning content"},"finish_reason": null}],"usage": null}',
|
||||
rb'data: {"id": "fake_id_1234","object": "chat.completion.chunk","created": 1234567890,"model": "accounts/fireworks/models/glm-4p5","choices": [{"index": 0,"delta": {"content": "Some content"},"finish_reason": null}],"usage": null}',
|
||||
rb'data: {"id": "fake_id_1234","object": "chat.completion.chunk","created": 1234567890,"model": "accounts/fireworks/models/glm-4p5","choices": [{"index": 0,"delta": {"tool_calls": [{"index": 0,"id": "fake_id_151617","type": "function","function": {"name": "some_tool"}}]},"finish_reason": null}],"usage": null}',
|
||||
rb'data: {"id": "fake_id_1234","object": "chat.completion.chunk","created": 1234567890,"model": "accounts/fireworks/models/glm-4p5","choices": [{"index": 0,"delta": {"tool_calls": [{"index": 0,"id": null,"type": "function","function": {"arguments": "{\"some_argument\": \"some_arguments_value\"}"}}]},"finish_reason": null}],"usage": null}',
|
||||
rb'data: {"id": "fake_id_1234","object": "chat.completion.chunk","created": 1234567890,"model": "accounts/fireworks/models/glm-4p5","choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}],"usage": {"prompt_tokens": 100,"total_tokens": 300,"completion_tokens": 200,"prompt_tokens_details": {"cached_tokens": 190}}}',
|
||||
rb"data: [DONE]",
|
||||
],
|
||||
[
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": None,
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": None,
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "Some content",
|
||||
"finish_reason": None,
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": None,
|
||||
"tool_calls": [{"name": "some_tool", "arguments": None, "index": 0}],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"name": None,
|
||||
"arguments": '{"some_argument": "some_arguments_value"}',
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": "tool_calls",
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
|
||||
},
|
||||
],
|
||||
)
|
||||
]
|
||||
173
tests/backend/data/mistral.py
Normal file
173
tests/backend/data/mistral.py
Normal file
|
|
@ -0,0 +1,173 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from tests.backend.data import Chunk, JsonResponse, ResultData, Url
|
||||
|
||||
SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
|
||||
(
|
||||
"https://api.mistral.ai",
|
||||
{
|
||||
"id": "fake_id_1234",
|
||||
"created": 1234567890,
|
||||
"model": "devstral-latest",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"total_tokens": 300,
|
||||
"completion_tokens": 200,
|
||||
},
|
||||
"object": "chat.completion",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"tool_calls": None,
|
||||
"content": "Some content",
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"message": "Some content",
|
||||
"finish_reason": "stop",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"total_tokens": 300,
|
||||
"completion_tokens": 200,
|
||||
},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
TOOL_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
|
||||
(
|
||||
"https://api.mistral.ai",
|
||||
{
|
||||
"id": "fake_id_1234",
|
||||
"created": 1234567890,
|
||||
"model": "devstral-latest",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"total_tokens": 300,
|
||||
"completion_tokens": 200,
|
||||
},
|
||||
"object": "chat.completion",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "tool_calls",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "fake_id_5678",
|
||||
"function": {
|
||||
"name": "some_tool",
|
||||
"arguments": '{"some_argument": "some_argument_value"}',
|
||||
},
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
"content": "Some content",
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"message": "Some content",
|
||||
"finish_reason": "tool_calls",
|
||||
"tool_calls": [
|
||||
{
|
||||
"name": "some_tool",
|
||||
"arguments": '{"some_argument": "some_argument_value"}',
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
STREAMED_SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]]] = [
|
||||
(
|
||||
"https://api.mistral.ai",
|
||||
[
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"devstral-latest","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}',
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"devstral-latest","choices":[{"index":0,"delta":{"content":"Some content"},"finish_reason":null}],"p":"abcde"}',
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"devstral-latest","choices":[{"index":0,"delta":{"content":""},"finish_reason":"stop"}],"usage":{"prompt_tokens":100,"total_tokens":300,"completion_tokens":200},"p":"abcdefghijklmnopq"}',
|
||||
rb"data: [DONE]",
|
||||
],
|
||||
[
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": None,
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "Some content",
|
||||
"finish_reason": None,
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": "stop",
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
|
||||
},
|
||||
],
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
STREAMED_TOOL_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]]] = [
|
||||
(
|
||||
"https://api.mistral.ai",
|
||||
[
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"devstral-latest","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}',
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"devstral-latest","choices":[{"index":0,"delta":{"content":"Some content"},"finish_reason":null}],"p":"a"}',
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"devstral-latest","choices":[{"index":0,"delta":{"tool_calls":[{"id":"fake_id_1234","function":{"name":"some_tool","arguments":""},"index":0}]},"finish_reason":null}],"p":"abcdef"}',
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"devstral-latest","choices":[{"index":0,"delta":{"tool_calls":[{"function":{"name":"","arguments":"{\"some_argument\": "},"index":0}]},"finish_reason":null}],"p":"abcdefghijklmnopq"}',
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"devstral-latest","choices":[{"index":0,"delta":{"tool_calls":[{"id":"null","function":{"name":"","arguments":"\"some_argument_value\"}"},"index":0}]},"finish_reason":null}],"p":"abcdefghijklmnopqrstuvwxyz0123456"}',
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"devstral-latest","choices":[{"index":0,"delta":{"content":""},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":100,"total_tokens":300,"completion_tokens":200},"p":"abcdefghijklmnopq"}',
|
||||
rb"data: [DONE]",
|
||||
],
|
||||
[
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": None,
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "Some content",
|
||||
"finish_reason": None,
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": None,
|
||||
"tool_calls": [{"name": "some_tool", "arguments": "", "index": 0}],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": None,
|
||||
"tool_calls": [
|
||||
{"name": "", "arguments": '{"some_argument": ', "index": 0}
|
||||
],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": None,
|
||||
"tool_calls": [
|
||||
{"name": "", "arguments": '"some_argument_value"}', "index": 0}
|
||||
],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
},
|
||||
{
|
||||
"message": "",
|
||||
"finish_reason": "tool_calls",
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 200},
|
||||
},
|
||||
],
|
||||
)
|
||||
]
|
||||
248
tests/backend/test_backend.py
Normal file
248
tests/backend/test_backend.py
Normal file
|
|
@ -0,0 +1,248 @@
|
|||
"""Test data for this module was generated using real LLM provider API responses,
|
||||
with responses simplified and formatted to make them readable and maintainable.
|
||||
|
||||
To update or modify test parameters:
|
||||
1. Make actual API calls to the target providers
|
||||
2. Use the raw API responses as a base for updating test data
|
||||
3. Simplify only where necessary for readability while preserving core structure
|
||||
|
||||
The closer test data remains to real API responses, the more reliable and accurate
|
||||
the tests will be. Always prefer real API data over manually constructed examples.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from tests.backend.data import Chunk, JsonResponse, ResultData, Url
|
||||
from tests.backend.data.fireworks import (
|
||||
SIMPLE_CONVERSATION_PARAMS as FIREWORKS_SIMPLE_CONVERSATION_PARAMS,
|
||||
STREAMED_SIMPLE_CONVERSATION_PARAMS as FIREWORKS_STREAMED_SIMPLE_CONVERSATION_PARAMS,
|
||||
STREAMED_TOOL_CONVERSATION_PARAMS as FIREWORKS_STREAMED_TOOL_CONVERSATION_PARAMS,
|
||||
TOOL_CONVERSATION_PARAMS as FIREWORKS_TOOL_CONVERSATION_PARAMS,
|
||||
)
|
||||
from tests.backend.data.mistral import (
|
||||
SIMPLE_CONVERSATION_PARAMS as MISTRAL_SIMPLE_CONVERSATION_PARAMS,
|
||||
STREAMED_SIMPLE_CONVERSATION_PARAMS as MISTRAL_STREAMED_SIMPLE_CONVERSATION_PARAMS,
|
||||
STREAMED_TOOL_CONVERSATION_PARAMS as MISTRAL_STREAMED_TOOL_CONVERSATION_PARAMS,
|
||||
TOOL_CONVERSATION_PARAMS as MISTRAL_TOOL_CONVERSATION_PARAMS,
|
||||
)
|
||||
from vibe.core.config import ModelConfig, ProviderConfig
|
||||
from vibe.core.llm.backend.generic import GenericBackend
|
||||
from vibe.core.llm.backend.mistral import MistralBackend
|
||||
from vibe.core.llm.exceptions import BackendError
|
||||
from vibe.core.llm.types import BackendLike
|
||||
from vibe.core.types import LLMChunk, LLMMessage, Role, ToolCall
|
||||
|
||||
|
||||
class TestBackend:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"base_url,json_response,result_data",
|
||||
[
|
||||
*FIREWORKS_SIMPLE_CONVERSATION_PARAMS,
|
||||
*FIREWORKS_TOOL_CONVERSATION_PARAMS,
|
||||
*MISTRAL_SIMPLE_CONVERSATION_PARAMS,
|
||||
*MISTRAL_TOOL_CONVERSATION_PARAMS,
|
||||
],
|
||||
)
|
||||
async def test_backend_complete(
|
||||
self, base_url: Url, json_response: JsonResponse, result_data: ResultData
|
||||
):
|
||||
with respx.mock(base_url=base_url) as mock_api:
|
||||
mock_api.post("/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(status_code=200, json=json_response)
|
||||
)
|
||||
provider = ProviderConfig(
|
||||
name="provider_name",
|
||||
api_base=f"{base_url}/v1",
|
||||
api_key_env_var="API_KEY",
|
||||
)
|
||||
|
||||
BackendClasses = [
|
||||
GenericBackend,
|
||||
*([MistralBackend] if base_url == "https://api.mistral.ai" else []),
|
||||
]
|
||||
for BackendClass in BackendClasses:
|
||||
backend: BackendLike = BackendClass(provider=provider)
|
||||
model = ModelConfig(
|
||||
name="model_name", provider="provider_name", alias="model_alias"
|
||||
)
|
||||
messages = [LLMMessage(role=Role.user, content="Just say hi")]
|
||||
|
||||
result = await backend.complete(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.2,
|
||||
tools=None,
|
||||
max_tokens=None,
|
||||
tool_choice=None,
|
||||
extra_headers=None,
|
||||
)
|
||||
|
||||
assert result.message.content == result_data["message"]
|
||||
assert result.finish_reason == result_data["finish_reason"]
|
||||
assert result.usage is not None
|
||||
assert (
|
||||
result.usage.prompt_tokens == result_data["usage"]["prompt_tokens"]
|
||||
)
|
||||
assert (
|
||||
result.usage.completion_tokens
|
||||
== result_data["usage"]["completion_tokens"]
|
||||
)
|
||||
|
||||
if result.message.tool_calls is None:
|
||||
return
|
||||
|
||||
assert len(result.message.tool_calls) == len(result_data["tool_calls"])
|
||||
for i, tool_call in enumerate[ToolCall](result.message.tool_calls):
|
||||
assert (
|
||||
tool_call.function.name == result_data["tool_calls"][i]["name"]
|
||||
)
|
||||
assert (
|
||||
tool_call.function.arguments
|
||||
== result_data["tool_calls"][i]["arguments"]
|
||||
)
|
||||
assert tool_call.index == result_data["tool_calls"][i]["index"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"base_url,chunks,result_data",
|
||||
[
|
||||
*FIREWORKS_STREAMED_SIMPLE_CONVERSATION_PARAMS,
|
||||
*FIREWORKS_STREAMED_TOOL_CONVERSATION_PARAMS,
|
||||
*MISTRAL_STREAMED_SIMPLE_CONVERSATION_PARAMS,
|
||||
*MISTRAL_STREAMED_TOOL_CONVERSATION_PARAMS,
|
||||
],
|
||||
)
|
||||
async def test_backend_complete_streaming(
|
||||
self, base_url: Url, chunks: list[Chunk], result_data: list[ResultData]
|
||||
):
|
||||
with respx.mock(base_url=base_url) as mock_api:
|
||||
mock_api.post("/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(
|
||||
status_code=200,
|
||||
stream=httpx.ByteStream(stream=b"\n\n".join(chunks)),
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
)
|
||||
)
|
||||
provider = ProviderConfig(
|
||||
name="provider_name",
|
||||
api_base=f"{base_url}/v1",
|
||||
api_key_env_var="API_KEY",
|
||||
)
|
||||
BackendClasses = [
|
||||
GenericBackend,
|
||||
*([MistralBackend] if base_url == "https://api.mistral.ai" else []),
|
||||
]
|
||||
for BackendClass in BackendClasses:
|
||||
backend: BackendLike = BackendClass(provider=provider)
|
||||
model = ModelConfig(
|
||||
name="model_name", provider="provider_name", alias="model_alias"
|
||||
)
|
||||
|
||||
messages = [
|
||||
LLMMessage(role=Role.user, content="List files in current dir")
|
||||
]
|
||||
|
||||
results: list[LLMChunk] = []
|
||||
async for result in backend.complete_streaming(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.2,
|
||||
tools=None,
|
||||
max_tokens=None,
|
||||
tool_choice=None,
|
||||
extra_headers=None,
|
||||
):
|
||||
results.append(result)
|
||||
|
||||
for result, expected_result in zip(results, result_data, strict=True):
|
||||
assert result.message.content == expected_result["message"]
|
||||
assert result.finish_reason == expected_result["finish_reason"]
|
||||
assert result.usage is not None
|
||||
assert (
|
||||
result.usage.prompt_tokens
|
||||
== expected_result["usage"]["prompt_tokens"]
|
||||
)
|
||||
assert (
|
||||
result.usage.completion_tokens
|
||||
== expected_result["usage"]["completion_tokens"]
|
||||
)
|
||||
|
||||
if result.message.tool_calls is None:
|
||||
continue
|
||||
|
||||
for i, tool_call in enumerate(result.message.tool_calls):
|
||||
assert (
|
||||
tool_call.function.name
|
||||
== expected_result["tool_calls"][i]["name"]
|
||||
)
|
||||
assert (
|
||||
tool_call.function.arguments
|
||||
== expected_result["tool_calls"][i]["arguments"]
|
||||
)
|
||||
assert (
|
||||
tool_call.index == expected_result["tool_calls"][i]["index"]
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"base_url,backend_class,response",
|
||||
[
|
||||
(
|
||||
"https://api.fireworks.ai",
|
||||
GenericBackend,
|
||||
httpx.Response(status_code=500, text="Internal Server Error"),
|
||||
),
|
||||
(
|
||||
"https://api.fireworks.ai",
|
||||
GenericBackend,
|
||||
httpx.Response(status_code=429, text="Rate Limit Exceeded"),
|
||||
),
|
||||
(
|
||||
"https://api.mistral.ai",
|
||||
MistralBackend,
|
||||
httpx.Response(status_code=500, text="Internal Server Error"),
|
||||
),
|
||||
(
|
||||
"https://api.mistral.ai",
|
||||
MistralBackend,
|
||||
httpx.Response(status_code=429, text="Rate Limit Exceeded"),
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_backend_complete_streaming_error(
|
||||
self,
|
||||
base_url: Url,
|
||||
backend_class: type[MistralBackend | GenericBackend],
|
||||
response: httpx.Response,
|
||||
):
|
||||
with respx.mock(base_url=base_url) as mock_api:
|
||||
mock_api.post("/v1/chat/completions").mock(return_value=response)
|
||||
provider = ProviderConfig(
|
||||
name="provider_name",
|
||||
api_base=f"{base_url}/v1",
|
||||
api_key_env_var="API_KEY",
|
||||
)
|
||||
backend = backend_class(provider=provider)
|
||||
model = ModelConfig(
|
||||
name="model_name", provider="provider_name", alias="model_alias"
|
||||
)
|
||||
messages = [LLMMessage(role=Role.user, content="Just say hi")]
|
||||
with pytest.raises(BackendError) as e:
|
||||
async for _ in backend.complete_streaming(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.2,
|
||||
tools=None,
|
||||
max_tokens=None,
|
||||
tool_choice=None,
|
||||
extra_headers=None,
|
||||
):
|
||||
pass
|
||||
assert e.value.status == response.status_code
|
||||
assert e.value.reason == response.reason_phrase
|
||||
assert e.value.parsed_error is None
|
||||
Loading…
Add table
Add a link
Reference in a new issue