v2.18.0 (#843)

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>

tests/backend/data/reasoning.py

@@ -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,
+    )