Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai>
Co-authored-by: Guillaume LE GOFF <guillaume.lgf@gmail.com>
Co-authored-by: Hdandria <henri.dandria@mistral.ai>
Co-authored-by: Ivana Dunisijevic <ivana.dunisijevic@mistral.ai>
Co-authored-by: Jean Burellier <sheplu@users.noreply.github.com>
Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai>
Co-authored-by: Mert Unsal <mert.unsal@mistral.ai>
Co-authored-by: Michel Thomazo <51709227+michelTho@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: Val <102326092+vdeva@users.noreply.github.com>
Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com>
Co-authored-by: renovate-mistral[bot] <253709520+renovate-mistral[bot]@users.noreply.github.com>
Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Clément Drouin 2026-06-19 11:01:24 +02:00 committed by GitHub
parent 564a14365e
commit 6bedf271ce
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
223 changed files with 10533 additions and 6947 deletions

View file

@ -8,6 +8,8 @@ import threading
import time
from typing import TypedDict, cast
from tests.constants import CHAT_COMPLETIONS_PATH
class StreamOptionsPayload(TypedDict, total=False):
include_usage: bool
@ -85,6 +87,104 @@ class StreamingMockServer:
),
]
@staticmethod
def _merge_tool_call_delta(
tool_calls_by_index: dict[int, dict[str, object]],
delta_tool_call: object,
*,
fallback_index: int | None,
) -> int | None:
if not isinstance(delta_tool_call, dict):
return fallback_index
index = delta_tool_call.get("index")
if not isinstance(index, int):
index = fallback_index
if index is None:
index = len(tool_calls_by_index)
tool_call = tool_calls_by_index.setdefault(index, {"index": index})
for key in ("id", "type"):
if value := delta_tool_call.get(key):
tool_call[key] = value
function_delta = delta_tool_call.get("function")
if not isinstance(function_delta, dict):
return index
function = tool_call.setdefault("function", {})
if not isinstance(function, dict):
function = {}
tool_call["function"] = function
if name := function_delta.get("name"):
function["name"] = name
if arguments := function_delta.get("arguments"):
function["arguments"] = f"{function.get('arguments', '')}{arguments}"
return index
@staticmethod
def _completion_response_from_chunks(
chunks: list[StreamChunk],
) -> dict[str, object]:
content_parts: list[str] = []
tool_calls_by_index: dict[int, dict[str, object]] = {}
active_tool_call_index: int | None = None
finish_reason: object = "stop"
usage: object = {"prompt_tokens": 3, "completion_tokens": 4}
created = 123
for chunk in chunks:
chunk_created = chunk.get("created")
if isinstance(chunk_created, int):
created = chunk_created
if chunk_usage := chunk.get("usage"):
usage = chunk_usage
choices = chunk.get("choices")
if not isinstance(choices, list) or not choices:
continue
choice = choices[0]
if not isinstance(choice, dict):
continue
if choice.get("finish_reason") is not None:
finish_reason = choice.get("finish_reason")
delta = choice.get("delta")
if not isinstance(delta, dict):
continue
if content := delta.get("content"):
content_parts.append(str(content))
if delta_tool_calls := delta.get("tool_calls"):
if isinstance(delta_tool_calls, list):
for delta_tool_call in delta_tool_calls:
active_tool_call_index = (
StreamingMockServer._merge_tool_call_delta(
tool_calls_by_index,
delta_tool_call,
fallback_index=active_tool_call_index,
)
)
message: dict[str, object] = {
"role": "assistant",
"content": "".join(content_parts),
}
tool_calls = [
tool_calls_by_index[index] for index in sorted(tool_calls_by_index)
]
if tool_calls:
message["tool_calls"] = tool_calls
return {
"id": "mock-id",
"object": "chat.completion",
"created": created,
"model": "mock-model",
"choices": [
{"index": 0, "message": message, "finish_reason": finish_reason}
],
"usage": usage,
}
def __init__(
self,
*,
@ -110,7 +210,7 @@ class StreamingMockServer:
return
def do_POST(self) -> None:
if self.path != "/v1/chat/completions":
if self.path != CHAT_COMPLETIONS_PATH:
self.send_response(404)
self.end_headers()
return
@ -125,17 +225,28 @@ class StreamingMockServer:
parent.requests.append(payload)
request_index = len(parent.requests) - 1
self.send_response(200)
self.send_header("Content-Type", "text/event-stream")
self.send_header("Cache-Control", "no-cache")
self.end_headers()
chunks = (
parent._chunk_factory(request_index, payload)
if parent._chunk_factory is not None
else parent._stream_chunks()
)
if not payload.get("stream"):
response = parent._completion_response_from_chunks(chunks)
response_body = json.dumps(response, ensure_ascii=False).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(response_body)))
self.end_headers()
self.wfile.write(response_body)
self.wfile.flush()
return
self.send_response(200)
self.send_header("Content-Type", "text/event-stream")
self.send_header("Cache-Control", "no-cache")
self.end_headers()
for chunk in chunks:
data = json.dumps(chunk, ensure_ascii=False)
self.wfile.write(f"data: {data}\n\n".encode())