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>
651 lines
20 KiB
Python
651 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from tests.backend.data import (
|
|
ANSWER_COMPLETION_TOKENS,
|
|
ANSWER_PROMPT_TOKENS,
|
|
Chunk,
|
|
JsonResponse,
|
|
ResultData,
|
|
Url,
|
|
)
|
|
from tests.constants import OPENAI_BASE_URL
|
|
|
|
|
|
def _sse_event(data: dict[str, Any] | str) -> Chunk:
|
|
if data == "[DONE]":
|
|
return b"data: [DONE]"
|
|
return f"data: {json.dumps(data, separators=(',', ':'))}".encode()
|
|
|
|
|
|
def _usage(prompt_tokens: int = 0, completion_tokens: int = 0) -> dict[str, int]:
|
|
return {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens}
|
|
|
|
|
|
def _result(
|
|
message: str = "",
|
|
*,
|
|
prompt_tokens: int = 0,
|
|
completion_tokens: int = 0,
|
|
reasoning_content: str | None = None,
|
|
tool_calls: list[dict[str, Any]] | None = None,
|
|
) -> ResultData:
|
|
result: ResultData = {
|
|
"message": message,
|
|
"usage": _usage(prompt_tokens, completion_tokens),
|
|
}
|
|
if reasoning_content is not None:
|
|
result["reasoning_content"] = reasoning_content
|
|
if tool_calls is not None:
|
|
result["tool_calls"] = tool_calls
|
|
return result
|
|
|
|
|
|
def _tool_call_result(
|
|
name: str | None, arguments: str, index: int | None
|
|
) -> dict[str, Any]:
|
|
return {"name": name, "arguments": arguments, "index": index}
|
|
|
|
|
|
def _output_text(text: str) -> dict[str, Any]:
|
|
return {"type": "output_text", "text": text, "annotations": [], "logprobs": []}
|
|
|
|
|
|
def _message_output_item(
|
|
message_id: str, text: str, *, phase: str | None = None, status: str = "completed"
|
|
) -> dict[str, Any]:
|
|
item: dict[str, Any] = {
|
|
"id": message_id,
|
|
"type": "message",
|
|
"status": status,
|
|
"content": [_output_text(text)],
|
|
"role": "assistant",
|
|
}
|
|
if phase is not None:
|
|
item["phase"] = phase
|
|
return item
|
|
|
|
|
|
def _stream_message_item(
|
|
message_id: str, *, phase: str | None = None, status: str = "in_progress"
|
|
) -> dict[str, Any]:
|
|
item: dict[str, Any] = {
|
|
"id": message_id,
|
|
"type": "message",
|
|
"status": status,
|
|
"content": [],
|
|
"role": "assistant",
|
|
}
|
|
if phase is not None:
|
|
item["phase"] = phase
|
|
return item
|
|
|
|
|
|
def _function_call_item(
|
|
item_id: str, call_id: str, name: str, arguments: str, *, status: str = "completed"
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"id": item_id,
|
|
"type": "function_call",
|
|
"call_id": call_id,
|
|
"name": name,
|
|
"arguments": arguments,
|
|
"status": status,
|
|
}
|
|
|
|
|
|
def openai_message_item(
|
|
text: str, *, phase: str | None = None, message_id: str = "msg_1"
|
|
) -> dict[str, Any]:
|
|
return _message_output_item(message_id, text, phase=phase)
|
|
|
|
|
|
def openai_function_call_item(
|
|
name: str, arguments: str, *, item_id: str = "fc_1", call_id: str = "call_1"
|
|
) -> dict[str, Any]:
|
|
return _function_call_item(item_id, call_id, name, arguments)
|
|
|
|
|
|
def openai_response(
|
|
output: list[dict[str, Any]],
|
|
*,
|
|
input_tokens: int = ANSWER_PROMPT_TOKENS,
|
|
output_tokens: int = ANSWER_COMPLETION_TOKENS,
|
|
response_id: str = "resp_1",
|
|
model: str = "gpt-test",
|
|
) -> JsonResponse:
|
|
return {
|
|
"id": response_id,
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": model,
|
|
"output": output,
|
|
"usage": {
|
|
"input_tokens": input_tokens,
|
|
"output_tokens": output_tokens,
|
|
"total_tokens": input_tokens + output_tokens,
|
|
},
|
|
}
|
|
|
|
|
|
def openai_sse(*events: dict[str, Any] | str) -> list[Chunk]:
|
|
return [_sse_event(event) for event in events]
|
|
|
|
|
|
def openai_reasoning_tool_call_stream(
|
|
name: str,
|
|
arguments: str,
|
|
*,
|
|
reasoning: str = "thinking...",
|
|
call_id: str = "call_1",
|
|
item_id: str = "fc_1",
|
|
input_tokens: int = 20,
|
|
output_tokens: int = 5,
|
|
) -> list[Chunk]:
|
|
return openai_sse(
|
|
{"type": "response.created", "response": {"id": "resp_1", "output": []}},
|
|
{
|
|
"type": "response.output_item.added",
|
|
"output_index": 0,
|
|
"item": _stream_message_item("msg_r", phase="commentary"),
|
|
},
|
|
{
|
|
"type": "response.reasoning_summary_text.delta",
|
|
"output_index": 0,
|
|
"delta": reasoning,
|
|
},
|
|
{
|
|
"type": "response.output_item.added",
|
|
"output_index": 1,
|
|
"item": _function_call_item(
|
|
item_id, call_id, name, "", status="in_progress"
|
|
),
|
|
},
|
|
{
|
|
"type": "response.function_call_arguments.delta",
|
|
"output_index": 1,
|
|
"call_id": call_id,
|
|
"delta": arguments,
|
|
},
|
|
{
|
|
"type": "response.function_call_arguments.done",
|
|
"output_index": 1,
|
|
"call_id": call_id,
|
|
"name": name,
|
|
"arguments": arguments,
|
|
},
|
|
{
|
|
"type": "response.output_item.done",
|
|
"output_index": 1,
|
|
"item": _function_call_item(item_id, call_id, name, arguments),
|
|
},
|
|
{
|
|
"type": "response.completed",
|
|
"response": openai_response(
|
|
[_function_call_item(item_id, call_id, name, arguments)],
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
),
|
|
},
|
|
"[DONE]",
|
|
)
|
|
|
|
|
|
def openai_text_stream(
|
|
text: str, *, input_tokens: int = 10, output_tokens: int = 2
|
|
) -> list[Chunk]:
|
|
return openai_sse(
|
|
{"type": "response.created", "response": {"id": "resp_1", "output": []}},
|
|
{
|
|
"type": "response.output_item.added",
|
|
"output_index": 0,
|
|
"item": _stream_message_item("msg_1"),
|
|
},
|
|
{
|
|
"type": "response.output_text.delta",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"delta": text,
|
|
},
|
|
{
|
|
"type": "response.completed",
|
|
"response": openai_response(
|
|
[openai_message_item(text)],
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
),
|
|
},
|
|
"[DONE]",
|
|
)
|
|
|
|
|
|
SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
|
|
(
|
|
OPENAI_BASE_URL,
|
|
{
|
|
"id": "resp_fake_id_1234",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": "gpt-4o-2024-08-06",
|
|
"output": [
|
|
_message_output_item(
|
|
"msg_fake_id_5678", "Hello! How can I help you today?"
|
|
)
|
|
],
|
|
"usage": {"input_tokens": 100, "output_tokens": 200, "total_tokens": 300},
|
|
},
|
|
_result(
|
|
"Hello! How can I help you today?", prompt_tokens=100, completion_tokens=200
|
|
),
|
|
)
|
|
]
|
|
|
|
TOOL_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
|
|
(
|
|
OPENAI_BASE_URL,
|
|
{
|
|
"id": "resp_fake_id_9012",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": "gpt-4o-2024-08-06",
|
|
"output": [
|
|
_message_output_item("msg_fake_id_3456", ""),
|
|
_function_call_item(
|
|
"fc_fake_id_7890",
|
|
"call_fake_id_1111",
|
|
"some_tool",
|
|
'{"some_argument": "some_argument_value"}',
|
|
),
|
|
],
|
|
"usage": {"input_tokens": 100, "output_tokens": 200, "total_tokens": 300},
|
|
},
|
|
_result(
|
|
prompt_tokens=100,
|
|
completion_tokens=200,
|
|
tool_calls=[
|
|
_tool_call_result(
|
|
"some_tool", '{"some_argument": "some_argument_value"}', 1
|
|
)
|
|
],
|
|
),
|
|
)
|
|
]
|
|
|
|
STREAMED_SIMPLE_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]]] = [
|
|
(
|
|
OPENAI_BASE_URL,
|
|
[
|
|
_sse_event({
|
|
"type": "response.created",
|
|
"response": {
|
|
"id": "resp_fake_id_1234",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": "gpt-4o-2024-08-06",
|
|
"output": [],
|
|
"usage": None,
|
|
},
|
|
}),
|
|
_sse_event({
|
|
"type": "response.in_progress",
|
|
"response": {"id": "resp_fake_id_1234"},
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_item.added",
|
|
"output_index": 0,
|
|
"item": _stream_message_item("msg_fake_id_5678"),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.content_part.added",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"part": _output_text(""),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_text.delta",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"delta": "Hello",
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_text.delta",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"delta": "!",
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_text.done",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"text": "Hello!",
|
|
}),
|
|
_sse_event({
|
|
"type": "response.content_part.done",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"part": _output_text("Hello!"),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_item.done",
|
|
"output_index": 0,
|
|
"item": _message_output_item("msg_fake_id_5678", "Hello!"),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.completed",
|
|
"response": {
|
|
"id": "resp_fake_id_1234",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": "gpt-4o-2024-08-06",
|
|
"output": [_message_output_item("msg_fake_id_5678", "Hello!")],
|
|
"usage": {
|
|
"input_tokens": 100,
|
|
"output_tokens": 200,
|
|
"total_tokens": 300,
|
|
},
|
|
},
|
|
}),
|
|
_sse_event("[DONE]"),
|
|
],
|
|
[
|
|
_result(),
|
|
_result(),
|
|
_result(),
|
|
_result(),
|
|
_result("Hello"),
|
|
_result("!"),
|
|
_result(),
|
|
_result(),
|
|
_result(),
|
|
_result(prompt_tokens=100, completion_tokens=200),
|
|
],
|
|
)
|
|
]
|
|
|
|
COMMENTARY_CONVERSATION_PARAMS: list[tuple[Url, JsonResponse, ResultData]] = [
|
|
(
|
|
OPENAI_BASE_URL,
|
|
{
|
|
"id": "resp_thinking_1234",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": "gpt-5.4-2025-04-14",
|
|
"output": [
|
|
_message_output_item(
|
|
"msg_commentary_5678",
|
|
"The user said hello, I should respond warmly.",
|
|
phase="commentary",
|
|
),
|
|
_message_output_item(
|
|
"msg_final_9012",
|
|
"Hello! How can I help you today?",
|
|
phase="final_answer",
|
|
),
|
|
],
|
|
"usage": {"input_tokens": 150, "output_tokens": 250, "total_tokens": 400},
|
|
},
|
|
_result(
|
|
"Hello! How can I help you today?",
|
|
prompt_tokens=150,
|
|
completion_tokens=250,
|
|
reasoning_content="The user said hello, I should respond warmly.",
|
|
),
|
|
)
|
|
]
|
|
|
|
STREAMED_COMMENTARY_CONVERSATION_PARAMS: list[
|
|
tuple[Url, list[Chunk], list[ResultData]]
|
|
] = [
|
|
(
|
|
OPENAI_BASE_URL,
|
|
[
|
|
_sse_event({
|
|
"type": "response.created",
|
|
"response": {
|
|
"id": "resp_thinking_1234",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": "gpt-5.4-2025-04-14",
|
|
"output": [],
|
|
"usage": None,
|
|
},
|
|
}),
|
|
_sse_event({
|
|
"type": "response.in_progress",
|
|
"response": {"id": "resp_thinking_1234"},
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_item.added",
|
|
"output_index": 0,
|
|
"item": _stream_message_item("msg_commentary_5678", phase="commentary"),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.content_part.added",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"part": _output_text(""),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_text.delta",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"delta": "Thinking",
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_text.delta",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"delta": " about it...",
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_text.done",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"text": "Thinking about it...",
|
|
}),
|
|
_sse_event({
|
|
"type": "response.content_part.done",
|
|
"output_index": 0,
|
|
"content_index": 0,
|
|
"part": _output_text("Thinking about it..."),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_item.done",
|
|
"output_index": 0,
|
|
"item": _message_output_item(
|
|
"msg_commentary_5678", "Thinking about it...", phase="commentary"
|
|
),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_item.added",
|
|
"output_index": 1,
|
|
"item": _stream_message_item("msg_final_9012", phase="final_answer"),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.content_part.added",
|
|
"output_index": 1,
|
|
"content_index": 0,
|
|
"part": _output_text(""),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_text.delta",
|
|
"output_index": 1,
|
|
"content_index": 0,
|
|
"delta": "Hello",
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_text.delta",
|
|
"output_index": 1,
|
|
"content_index": 0,
|
|
"delta": "!",
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_text.done",
|
|
"output_index": 1,
|
|
"content_index": 0,
|
|
"text": "Hello!",
|
|
}),
|
|
_sse_event({
|
|
"type": "response.content_part.done",
|
|
"output_index": 1,
|
|
"content_index": 0,
|
|
"part": _output_text("Hello!"),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_item.done",
|
|
"output_index": 1,
|
|
"item": _message_output_item(
|
|
"msg_final_9012", "Hello!", phase="final_answer"
|
|
),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.completed",
|
|
"response": {
|
|
"id": "resp_thinking_1234",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": "gpt-5.4-2025-04-14",
|
|
"output": [
|
|
_message_output_item(
|
|
"msg_commentary_5678",
|
|
"Thinking about it...",
|
|
phase="commentary",
|
|
),
|
|
_message_output_item(
|
|
"msg_final_9012", "Hello!", phase="final_answer"
|
|
),
|
|
],
|
|
"usage": {
|
|
"input_tokens": 150,
|
|
"output_tokens": 250,
|
|
"total_tokens": 400,
|
|
},
|
|
},
|
|
}),
|
|
_sse_event("[DONE]"),
|
|
],
|
|
[
|
|
_result(),
|
|
_result(),
|
|
_result(),
|
|
_result(),
|
|
_result(reasoning_content="Thinking"),
|
|
_result(reasoning_content=" about it..."),
|
|
_result(),
|
|
_result(),
|
|
_result(),
|
|
_result(),
|
|
_result(),
|
|
_result("Hello"),
|
|
_result("!"),
|
|
_result(),
|
|
_result(),
|
|
_result(),
|
|
_result(prompt_tokens=150, completion_tokens=250),
|
|
],
|
|
)
|
|
]
|
|
|
|
STREAMED_TOOL_CONVERSATION_PARAMS: list[tuple[Url, list[Chunk], list[ResultData]]] = [
|
|
(
|
|
OPENAI_BASE_URL,
|
|
[
|
|
_sse_event({
|
|
"type": "response.created",
|
|
"response": {
|
|
"id": "resp_fake_id_9012",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": "gpt-4o-2024-08-06",
|
|
"output": [],
|
|
"usage": None,
|
|
},
|
|
}),
|
|
_sse_event({
|
|
"type": "response.in_progress",
|
|
"response": {"id": "resp_fake_id_9012"},
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_item.added",
|
|
"output_index": 0,
|
|
"item": _function_call_item(
|
|
"fc_fake_id_7890",
|
|
"call_fake_id_1111",
|
|
"some_tool",
|
|
"",
|
|
status="in_progress",
|
|
),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.function_call_arguments.delta",
|
|
"output_index": 0,
|
|
"call_id": "call_fake_id_1111",
|
|
"delta": '{"some_argument": ',
|
|
}),
|
|
_sse_event({
|
|
"type": "response.function_call_arguments.delta",
|
|
"output_index": 0,
|
|
"call_id": "call_fake_id_1111",
|
|
"delta": '"some_argument_value"}',
|
|
}),
|
|
_sse_event({
|
|
"type": "response.function_call_arguments.done",
|
|
"output_index": 0,
|
|
"call_id": "call_fake_id_1111",
|
|
"name": "some_tool",
|
|
"arguments": '{"some_argument": "some_argument_value"}',
|
|
}),
|
|
_sse_event({
|
|
"type": "response.output_item.done",
|
|
"output_index": 0,
|
|
"item": _function_call_item(
|
|
"fc_fake_id_7890",
|
|
"call_fake_id_1111",
|
|
"some_tool",
|
|
'{"some_argument": "some_argument_value"}',
|
|
),
|
|
}),
|
|
_sse_event({
|
|
"type": "response.completed",
|
|
"response": {
|
|
"id": "resp_fake_id_9012",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"model": "gpt-4o-2024-08-06",
|
|
"output": [
|
|
_function_call_item(
|
|
"fc_fake_id_7890",
|
|
"call_fake_id_1111",
|
|
"some_tool",
|
|
'{"some_argument": "some_argument_value"}',
|
|
)
|
|
],
|
|
"usage": {
|
|
"input_tokens": 100,
|
|
"output_tokens": 200,
|
|
"total_tokens": 300,
|
|
},
|
|
},
|
|
}),
|
|
_sse_event("[DONE]"),
|
|
],
|
|
[
|
|
_result(),
|
|
_result(),
|
|
_result(tool_calls=[_tool_call_result("some_tool", "", 0)]),
|
|
_result(),
|
|
_result(),
|
|
_result(
|
|
tool_calls=[
|
|
_tool_call_result(
|
|
"some_tool", '{"some_argument": "some_argument_value"}', 0
|
|
)
|
|
]
|
|
),
|
|
_result(),
|
|
_result(prompt_tokens=100, completion_tokens=200),
|
|
],
|
|
)
|
|
]
|