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>
This commit is contained in:
Laure Hugo 2026-06-25 16:08:45 +02:00 committed by GitHub
parent 725d3a56ce
commit e607ccbb00
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
242 changed files with 7372 additions and 1974 deletions

View file

@ -4,3 +4,8 @@ Url = str
JsonResponse = dict
ResultData = dict
Chunk = bytes
# Shared usage every provider's `answer()` mock reports
ANSWER_PROMPT_TOKENS = 10
ANSWER_COMPLETION_TOKENS = 5
ANSWER_CONTEXT_TOKENS = ANSWER_PROMPT_TOKENS + ANSWER_COMPLETION_TOKENS

View file

@ -3,13 +3,42 @@ from __future__ import annotations
import json
from typing import Any
from tests.backend.data import Chunk, JsonResponse
from tests.backend.data import (
ANSWER_COMPLETION_TOKENS,
ANSWER_PROMPT_TOKENS,
Chunk,
JsonResponse,
)
def _sse_event(data: dict[str, Any]) -> Chunk:
return f"data: {json.dumps(data, separators=(',', ':'))}".encode()
def _block_start(index: int, block: dict[str, Any]) -> dict[str, Any]:
return {"type": "content_block_start", "index": index, "content_block": block}
def _block_delta(index: int, delta: dict[str, Any]) -> dict[str, Any]:
return {"type": "content_block_delta", "index": index, "delta": delta}
def _content_stream(
blocks: list[dict[str, Any]], *, input_tokens: int, output_tokens: int, stop: str
) -> list[Chunk]:
events = [
{"type": "message_start", "message": {"usage": {"input_tokens": input_tokens}}},
*blocks,
{
"type": "message_delta",
"delta": {"stop_reason": stop},
"usage": {"output_tokens": output_tokens},
},
{"type": "message_stop"},
]
return [_sse_event(e) for e in events]
def anthropic_request_content_blocks(payload: dict[str, Any]) -> list[dict[str, Any]]:
# Flatten every structured content block across a request's messages
# (text/thinking/tool_use/tool_result blocks).
@ -24,8 +53,8 @@ def anthropic_request_content_blocks(payload: dict[str, Any]) -> list[dict[str,
def anthropic_message(
text: str,
*,
input_tokens: int = 12,
output_tokens: int = 3,
input_tokens: int = ANSWER_PROMPT_TOKENS,
output_tokens: int = ANSWER_COMPLETION_TOKENS,
stop_reason: str = "end_turn",
) -> JsonResponse:
return {
@ -62,71 +91,31 @@ def anthropic_reasoning_tool_use_stream(
input_tokens: int = 20,
output_tokens: int = 5,
) -> list[Chunk]:
return [
_sse_event(e)
for e in (
{
"type": "message_start",
"message": {"usage": {"input_tokens": input_tokens}},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "thinking", "thinking": reasoning},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "signature_delta", "signature": signature},
},
return _content_stream(
[
_block_start(0, {"type": "thinking", "thinking": reasoning}),
_block_delta(0, {"type": "signature_delta", "signature": signature}),
{"type": "content_block_stop", "index": 0},
{
"type": "content_block_start",
"index": 1,
"content_block": {"type": "tool_use", "id": tool_id, "name": name},
},
{
"type": "content_block_delta",
"index": 1,
"delta": {"type": "input_json_delta", "partial_json": arguments},
},
_block_start(1, {"type": "tool_use", "id": tool_id, "name": name}),
_block_delta(1, {"type": "input_json_delta", "partial_json": arguments}),
{"type": "content_block_stop", "index": 1},
{
"type": "message_delta",
"delta": {"stop_reason": "tool_use"},
"usage": {"output_tokens": output_tokens},
},
{"type": "message_stop"},
)
]
],
input_tokens=input_tokens,
output_tokens=output_tokens,
stop="tool_use",
)
def anthropic_text_stream(
text: str, *, input_tokens: int = 12, output_tokens: int = 3
) -> list[Chunk]:
return [
_sse_event(e)
for e in (
{
"type": "message_start",
"message": {"usage": {"input_tokens": input_tokens}},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": text},
},
return _content_stream(
[
_block_start(0, {"type": "text", "text": ""}),
_block_delta(0, {"type": "text_delta", "text": text}),
{"type": "content_block_stop", "index": 0},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": output_tokens},
},
{"type": "message_stop"},
)
]
],
input_tokens=input_tokens,
output_tokens=output_tokens,
stop="end_turn",
)

View file

@ -3,7 +3,14 @@ from __future__ import annotations
import json
from typing import Any
from tests.backend.data import Chunk, JsonResponse, ResultData, Url
from tests.backend.data import (
ANSWER_COMPLETION_TOKENS,
ANSWER_PROMPT_TOKENS,
Chunk,
JsonResponse,
ResultData,
Url,
)
from tests.constants import OPENAI_BASE_URL
@ -104,8 +111,8 @@ def openai_function_call_item(
def openai_response(
output: list[dict[str, Any]],
*,
input_tokens: int = 10,
output_tokens: int = 2,
input_tokens: int = ANSWER_PROMPT_TOKENS,
output_tokens: int = ANSWER_COMPLETION_TOKENS,
response_id: str = "resp_1",
model: str = "gpt-test",
) -> JsonResponse:

View file

@ -0,0 +1,135 @@
from __future__ import annotations
import json
from typing import Any
from tests.backend.data import (
ANSWER_COMPLETION_TOKENS,
ANSWER_PROMPT_TOKENS,
Chunk,
JsonResponse,
)
def _sse_event(data: dict[str, Any]) -> Chunk:
return f"data: {json.dumps(data, separators=(',', ':'))}".encode()
def _thinking_block(reasoning: str) -> dict[str, Any]:
return {"type": "thinking", "thinking": [{"type": "text", "text": reasoning}]}
def reasoning_message(
text: str,
*,
prompt_tokens: int = ANSWER_PROMPT_TOKENS,
completion_tokens: int = ANSWER_COMPLETION_TOKENS,
) -> JsonResponse:
return {
"choices": [{"message": {"role": "assistant", "content": text}}],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
},
}
def reasoning_thinking_message(
text: str, reasoning: str, *, prompt_tokens: int = 12, completion_tokens: int = 5
) -> JsonResponse:
return {
"choices": [
{
"message": {
"role": "assistant",
"content": [
_thinking_block(reasoning),
{"type": "text", "text": text},
],
}
}
],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
},
}
def reasoning_tool_use(
name: str,
arguments: str,
*,
reasoning: str | None = None,
tool_id: str = "call_1",
prompt_tokens: int = 20,
completion_tokens: int = 5,
) -> JsonResponse:
message: dict[str, Any] = {
"role": "assistant",
"tool_calls": [
{
"id": tool_id,
"type": "function",
"function": {"name": name, "arguments": arguments},
}
],
}
if reasoning is not None:
message["content"] = [_thinking_block(reasoning)]
return {
"choices": [{"message": message}],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
},
}
def _delta_stream(
deltas: list[dict[str, Any]], *, prompt_tokens: int, completion_tokens: int
) -> list[Chunk]:
events = [
{"choices": [{"delta": {"role": "assistant"}}]},
*({"choices": [{"delta": delta}]} for delta in deltas),
{
"choices": [{"delta": {}}],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
},
},
]
return [_sse_event(e) for e in events] + [b"data: [DONE]"]
def reasoning_text_stream(
text: str, *, prompt_tokens: int = 10, completion_tokens: int = 3
) -> list[Chunk]:
return _delta_stream(
[{"content": text}],
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
)
def reasoning_thinking_tool_use_stream(
name: str,
arguments: str,
*,
reasoning: str = "thinking...",
tool_id: str = "call_1",
prompt_tokens: int = 20,
completion_tokens: int = 5,
) -> list[Chunk]:
tool_call = {
"id": tool_id,
"type": "function",
"index": 0,
"function": {"name": name, "arguments": arguments},
}
return _delta_stream(
[{"content": [_thinking_block(reasoning)]}, {"tool_calls": [tool_call]}],
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
)