From 57be5868f9a4826b6be2b63646ce192afae24f50 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 9 Jun 2026 12:13:25 -0300 Subject: [PATCH] Studio: improve OpenAI- and Anthropic-compatible API spec compliance (#6010) * Studio: fix OpenAI- and Anthropic-compatible API spec compliance * Studio: fix API spec-compliance gaps on passthrough and streaming paths * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: carry context_length_exceeded through the OpenAI passthrough error path * Studio: count tool-schema tokens in the Anthropic server-tool stream, and small stream-handling guards * Studio: guard message_delta usage against None and normalize developer role before proxying * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: honor max_completion_tokens on the external-provider proxy path * Studio: forward llama-server cached_tokens into OpenAI prompt_tokens_details * Studio: sanitize messages in count_tokens to match the /v1/messages prompt * Studio: report max_tokens for truncated tool calls and guard null usage in metadata events * Studio: drop the request-id middleware (headers aren't declared in either spec) * Studio: include the required request_id field in Anthropic error bodies * Studio: honor max_completion_tokens on the audio (TTS / audio-input) paths * Studio: add the _effective_max_tokens helper and route all max-token sites through it * Studio: align API compatibility edge cases * Studio: clarify multi-choice chat support * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: clarify logprobs chat support * Studio: opt the local chat UI into the streaming usage chunk so the context bar and tok/s repopulate * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: forward seed to llama-server, and fix Anthropic server-tool stop_reason, tool_result id correlation, and parallel-tool execution cap * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: align OpenAI chat completion spec edge cases * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: align backend API compatibility tests * Studio: honor tool caps and internal stream usage * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: coerce nullable stream usage counts * Studio: preserve system prompts with developer messages --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: wasimysaid --- .../core/inference/anthropic_compat.py | 112 +- studio/backend/core/inference/llama_cpp.py | 178 ++- studio/backend/main.py | 12 +- studio/backend/models/inference.py | 56 +- studio/backend/routes/data_recipe/jobs.py | 6 + studio/backend/routes/inference.py | 1143 +++++++++++++---- .../backend/tests/test_anthropic_messages.py | 173 ++- .../tests/test_gguf_completion_usage.py | 28 +- .../tests/test_openai_tool_passthrough.py | 348 ++++- .../tests/test_responses_tool_passthrough.py | 149 +++ studio/backend/utils/api_errors.py | 252 ++++ .../src/features/chat/api/chat-adapter.ts | 3 + .../frontend/src/features/chat/types/api.ts | 7 + 13 files changed, 2174 insertions(+), 293 deletions(-) create mode 100644 studio/backend/utils/api_errors.py diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index c1f1c6b81e..0307336dde 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -10,9 +10,37 @@ Pure functions plus stateful stream emitters; no FastAPI, no I/O. from __future__ import annotations import json +import uuid from typing import Any, Optional, Union +def openai_finish_to_anthropic_stop(finish_reason, had_tool_calls = False) -> str: + """Map an OpenAI finish_reason to an Anthropic stop_reason. + 'length' -> 'max_tokens' (truncation wins even mid tool call, so a cut-off + tool call isn't mislabeled tool_use); tool_calls / had_tool_calls -> 'tool_use'; + 'stop_sequence' -> 'stop_sequence'; 'stop'/None/unknown -> 'end_turn'.""" + # Truncation takes precedence: a tool call cut off at max_tokens has possibly + # incomplete arguments, so report max_tokens rather than telling the client to + # run the tool. + if finish_reason == "length": + return "max_tokens" + if finish_reason == "tool_calls" or had_tool_calls: + return "tool_use" + if finish_reason == "stop_sequence": + return "stop_sequence" + # "stop", None, and any unknown value collapse to end_turn. + return "end_turn" + + +def anthropic_tool_use_id(upstream_id = None) -> str: + """Return an Anthropic-style tool_use id (prefix 'toolu_'). Reuses an + upstream id only if it already starts with 'toolu_'; otherwise mints a fresh + 'toolu_<24 hex>'.""" + if upstream_id and isinstance(upstream_id, str) and upstream_id.startswith("toolu_"): + return upstream_id + return f"toolu_{uuid.uuid4().hex[:24]}" + + def _anthropic_image_block_to_openai_part(block: dict) -> Optional[dict]: """Translate one Anthropic ``image`` block to an OpenAI ``image_url`` part. @@ -204,6 +232,19 @@ def build_anthropic_sse_event(event_type: str, data: dict) -> str: return f"event: {event_type}\ndata: {json.dumps(data)}\n\n" +def _message_delta_usage(usage: Optional[dict]) -> dict: + """Usage block for a message_delta event (cumulative token counts). Cache + fields are always 0 — no prompt caching backend. ``usage`` may be None when a + metadata event carried usage=None (e.g. only finish_reason set).""" + usage = usage or {} + return { + "input_tokens": usage.get("prompt_tokens", 0), + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": usage.get("completion_tokens", 0), + } + + class AnthropicStreamEmitter: """Converts generate_chat_completion_with_tools() events into Anthropic Messages SSE strings.""" @@ -212,11 +253,19 @@ class AnthropicStreamEmitter: self.block_index: int = 0 self._text_block_open: bool = False self._open_tool_call_id: Optional[str] = None + # The mapped Anthropic ``toolu_*`` id published in content_block_start, + # reused for the paired tool_result so consumers can correlate them. + self._open_tool_use_id: Optional[str] = None self._open_tool_args_sent: bool = False self._prev_text: str = "" self._usage: dict = {} - def start(self, message_id: str, model: str) -> list[str]: + def start( + self, + message_id: str, + model: str, + input_tokens: int = 0, + ) -> list[str]: """Emit message_start and open the first text content block.""" events = [] events.append( @@ -232,7 +281,12 @@ class AnthropicStreamEmitter: "model": model, "stop_reason": None, "stop_sequence": None, - "usage": {"input_tokens": 0, "output_tokens": 0}, + "usage": { + "input_tokens": input_tokens, + "output_tokens": 0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + }, }, }, ) @@ -255,22 +309,28 @@ class AnthropicStreamEmitter: # status events — no Anthropic equivalent return [] - def finish(self, stop_reason: str = "end_turn") -> list[str]: + def finish( + self, + stop_reason: str = "end_turn", + stop_sequence = None, + ) -> list[str]: """Close any open block and emit message_delta + message_stop.""" events = [] if self._text_block_open or self._open_tool_call_id is not None: events.append(self._close_block()) self._open_tool_call_id = None + self._open_tool_use_id = None self._open_tool_args_sent = False events.append( build_anthropic_sse_event( "message_delta", { "type": "message_delta", - "delta": {"stop_reason": stop_reason, "stop_sequence": None}, - "usage": { - "output_tokens": self._usage.get("completion_tokens", 0), + "delta": { + "stop_reason": stop_reason, + "stop_sequence": stop_sequence, }, + "usage": _message_delta_usage(self._usage), }, ) ) @@ -319,11 +379,13 @@ class AnthropicStreamEmitter: elif self._open_tool_call_id is not None: events.append(self._close_block()) self._open_tool_call_id = None + self._open_tool_use_id = None self._open_tool_args_sent = False # Open a tool_use block. self.block_index += 1 self._open_tool_call_id = tool_call_id + self._open_tool_use_id = anthropic_tool_use_id(tool_call_id) self._open_tool_args_sent = False events.append( build_anthropic_sse_event( @@ -333,7 +395,7 @@ class AnthropicStreamEmitter: "index": self.block_index, "content_block": { "type": "tool_use", - "id": tool_call_id, + "id": self._open_tool_use_id, "name": event.get("tool_name", ""), "input": {}, }, @@ -368,7 +430,11 @@ class AnthropicStreamEmitter: # Close the tool_use block. if self._open_tool_call_id is not None or self._text_block_open: events.append(self._close_block()) + # Reuse the id published in content_block_start; fall back to mapping + # the raw id only if no tool_start preceded this end. + tool_use_id = self._open_tool_use_id or anthropic_tool_use_id(event.get("tool_call_id", "")) self._open_tool_call_id = None + self._open_tool_use_id = None self._open_tool_args_sent = False # Emit custom tool_result event (non-standard, ignored by SDKs) events.append( @@ -376,7 +442,7 @@ class AnthropicStreamEmitter: "tool_result", { "type": "tool_result", - "tool_use_id": event.get("tool_call_id", ""), + "tool_use_id": tool_use_id, "content": event.get("result", ""), }, ) @@ -427,8 +493,14 @@ class AnthropicPassthroughEmitter: self._tool_call_states: dict = {} # delta index -> {block_index, id, name} self._usage: dict = {} self._stop_reason: str = "end_turn" + self._stop_sequence: Optional[str] = None - def start(self, message_id: str, model: str) -> list[str]: + def start( + self, + message_id: str, + model: str, + input_tokens: int = 0, + ) -> list[str]: return [ build_anthropic_sse_event( "message_start", @@ -442,7 +514,12 @@ class AnthropicPassthroughEmitter: "model": model, "stop_reason": None, "stop_sequence": None, - "usage": {"input_tokens": 0, "output_tokens": 0}, + "usage": { + "input_tokens": input_tokens, + "output_tokens": 0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + }, }, }, ) @@ -492,7 +569,7 @@ class AnthropicPassthroughEmitter: # New tool call — close prior block, open tool_use block if self._current_block_type is not None: events.append(self._close_current_block()) - tc_id = tc.get("id", "") + tc_id = anthropic_tool_use_id(tc.get("id", "")) tc_name = fn.get("name", "") self.block_index += 1 self._current_block_type = "tool_use" @@ -535,12 +612,7 @@ class AnthropicPassthroughEmitter: # ── Finish reason ── if finish_reason: - if finish_reason == "tool_calls": - self._stop_reason = "tool_use" - elif finish_reason == "length": - self._stop_reason = "max_tokens" - else: - self._stop_reason = "end_turn" + self._stop_reason = openai_finish_to_anthropic_stop(finish_reason) return events @@ -555,11 +627,9 @@ class AnthropicPassthroughEmitter: "type": "message_delta", "delta": { "stop_reason": self._stop_reason, - "stop_sequence": None, - }, - "usage": { - "output_tokens": self._usage.get("completion_tokens", 0), + "stop_sequence": self._stop_sequence, }, + "usage": _message_delta_usage(self._usage), }, ) ) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index d12eb6e29e..ec3a216258 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4116,6 +4116,7 @@ class LlamaCppBackend: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + seed: Optional[int] = None, ) -> Generator[str | dict, None, None]: """ Send a chat completion to llama-server and stream tokens back. @@ -4156,6 +4157,8 @@ class LlamaCppBackend: payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: payload["stop"] = stop + if seed is not None: + payload["seed"] = seed payload["stream_options"] = {"include_usage": True} url = f"{self.base_url}/v1/chat/completions" @@ -4164,6 +4167,7 @@ class LlamaCppBackend: _stream_done = False _metadata_usage = None _metadata_timings = None + _metadata_finish_reason = None try: # _stream_with_retry uses a 120 s read timeout so prefill can @@ -4228,6 +4232,9 @@ class LlamaCppBackend: choices = data.get("choices", []) if choices: delta = choices[0].get("delta", {}) + _fr = choices[0].get("finish_reason") + if _fr: + _metadata_finish_reason = _fr # Reasoning/thinking tokens: llama-server # sends these as "reasoning_content"; wrap @@ -4253,14 +4260,18 @@ class LlamaCppBackend: logger.debug(f"Skipping malformed SSE line: {line[:100]}") if _stream_done: break # exit outer for - if _metadata_usage or _metadata_timings: + if _metadata_usage or _metadata_timings or _metadata_finish_reason: _metadata_usage = _backfill_usage_from_timings( _metadata_usage, _metadata_timings ) yield { "type": "metadata", - "usage": _metadata_usage, + # Never None: a finish-only metadata event (no usage, + # no timings) would otherwise crash consumers that do + # usage.get(...) on the non-streaming paths. + "usage": _metadata_usage or {}, "timings": _metadata_timings, + "finish_reason": _metadata_finish_reason, } except httpx.ConnectError: @@ -4292,6 +4303,8 @@ class LlamaCppBackend: auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, + seed: Optional[int] = None, + disable_parallel_tool_use: bool = False, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -4394,6 +4407,8 @@ class LlamaCppBackend: payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: payload["stop"] = stop + if seed is not None: + payload["seed"] = seed try: _auth_headers = ( @@ -4419,6 +4434,7 @@ class LlamaCppBackend: has_structured_tc = False _iter_usage = None _iter_timings = None + _iter_finish_reason = None _stream_done = False _last_emitted = "" provisional_render_html_tool_call_ids = set() @@ -4498,6 +4514,9 @@ class LlamaCppBackend: continue delta = choices[0].get("delta", {}) + _fr = choices[0].get("finish_reason") + if _fr: + _iter_finish_reason = _fr # ── Structured tool_calls ── tc_deltas = delta.get("tool_calls") @@ -4840,6 +4859,7 @@ class LlamaCppBackend: "total_tokens": _fp + _tc, }, "timings": _mt, + "finish_reason": _iter_finish_reason, } return @@ -4915,6 +4935,7 @@ class LlamaCppBackend: "total_tokens": _fp + _tc, }, "timings": _mt, + "finish_reason": _iter_finish_reason, } return @@ -4926,6 +4947,12 @@ class LlamaCppBackend: _accumulated_predicted_ms += _it.get("predicted_ms", 0) _accumulated_predicted_n += _it.get("predicted_n", 0) + # disable_parallel_tool_use: execute only the first tool call + # this turn. Truncate before building assistant_msg so the + # conversation stays consistent and extra calls are never executed. + if disable_parallel_tool_use and tool_calls and len(tool_calls) > 1: + tool_calls = tool_calls[:1] + assistant_msg: dict = {"role": "assistant", "content": content_text} assistant_appended = False @@ -5040,6 +5067,8 @@ class LlamaCppBackend: stream_payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: stream_payload["stop"] = stop + if seed is not None: + stream_payload["seed"] = seed stream_payload["stream_options"] = {"include_usage": True} cumulative = "" @@ -5049,6 +5078,7 @@ class LlamaCppBackend: reasoning_text = "" _metadata_usage = None _metadata_timings = None + _metadata_finish_reason = None _stream_done = False try: @@ -5107,6 +5137,9 @@ class LlamaCppBackend: choices = chunk_data.get("choices", []) if choices: delta = choices[0].get("delta", {}) + _fr = choices[0].get("finish_reason") + if _fr: + _metadata_finish_reason = _fr reasoning = delta.get("reasoning_content", "") if reasoning: @@ -5137,7 +5170,7 @@ class LlamaCppBackend: _final_completion = _final_usage.get("completion_tokens", 0) _final_prompt = _final_usage.get("prompt_tokens", 0) _total_completion = _final_completion + _accumulated_completion_tokens - if _metadata_usage or _metadata_timings: + if _metadata_usage or _metadata_timings or _metadata_finish_reason: _merged_timings = dict(_metadata_timings) if _metadata_timings else {} if _accumulated_predicted_ms or _accumulated_predicted_n: _merged_timings["predicted_ms"] = ( @@ -5160,6 +5193,7 @@ class LlamaCppBackend: "total_tokens": _final_prompt + _total_completion, }, "timings": _merged_timings, + "finish_reason": _metadata_finish_reason, } except httpx.ConnectError: @@ -5169,6 +5203,144 @@ class LlamaCppBackend: return raise + # ── Prompt token counting ────────────────────────────────── + + def count_chat_tokens( + self, + messages, + system = None, + tools = None, + strict: bool = False, + ) -> int: + """Count prompt tokens for a chat request via llama-server. + + Non-strict callers keep the historical best-effort behavior and receive + 0 when a count cannot be determined. Strict callers (public count_tokens + endpoints) get an exception instead of a successful-looking zero when + tokenizer/template calls fail or a multimodal prompt would fall back to a + text-only approximation. + """ + if not self.is_loaded: + if strict: + raise RuntimeError("llama-server is not loaded") + return 0 + + def _has_non_text_content(content) -> bool: + if isinstance(content, list): + for block in content: + if isinstance(block, str): + continue + if not isinstance(block, dict): + return True + if block.get("type") == "text" and isinstance(block.get("text"), str): + continue + if isinstance(block.get("text"), str): + continue + return True + return False + + def _has_non_text_prompt_parts() -> bool: + if _has_non_text_content(system): + return True + for msg in messages or []: + if isinstance(msg, dict) and _has_non_text_content(msg.get("content", "")): + return True + return False + + def _block_text(content) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [] + for block in content: + if isinstance(block, dict): + if block.get("type") == "text" and isinstance(block.get("text"), str): + parts.append(block["text"]) + elif isinstance(block.get("text"), str): + parts.append(block["text"]) + elif isinstance(block, str): + parts.append(block) + return "".join(parts) + return "" + + # Normalize system into a leading message / plain text. + system_text = "" + if isinstance(system, str): + system_text = system + elif isinstance(system, list): + system_text = _block_text(system) + + try: + _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + with httpx.Client(timeout = 10, headers = _auth_headers) as client: + + def _tokenize(text: str) -> int: + r = client.post( + f"{self.base_url}/tokenize", + json = {"content": text, "add_special": True}, + ) + if r.status_code != 200: + if strict: + raise RuntimeError("llama-server tokenizer failed") + return 0 + tokens = r.json().get("tokens", []) + if not isinstance(tokens, list): + if strict: + raise RuntimeError("llama-server tokenizer returned invalid tokens") + return 0 + return len(tokens) + + # 1. Try /apply-template to render the real chat prompt. + template_messages = list(messages) if messages else [] + if system_text: + template_messages = [ + {"role": "system", "content": system_text} + ] + template_messages + apply_template_failed = False + try: + # llama-server's /apply-template renders tool declarations + # into the prompt when ``tools`` is supplied, so pass them + # through — otherwise tool-schema tokens go uncounted. + template_body = {"messages": template_messages} + if tools: + template_body["tools"] = tools + resp = client.post( + f"{self.base_url}/apply-template", + json = template_body, + ) + if resp.status_code == 200: + prompt = resp.json().get("prompt", "") + if isinstance(prompt, str): + return _tokenize(prompt) + apply_template_failed = True + except Exception: + apply_template_failed = True + + if strict and apply_template_failed and _has_non_text_prompt_parts(): + raise RuntimeError( + "cannot fall back to text-only token counting for multimodal messages" + ) + + # 2. Fallback: concatenate plain text and tokenize. Append a + # serialized form of the tools so they still contribute to the + # count when /apply-template is unavailable. + parts = [] + if system_text: + parts.append(system_text) + for msg in messages or []: + if isinstance(msg, dict): + parts.append(_block_text(msg.get("content", ""))) + if tools: + try: + parts.append(json.dumps(tools, ensure_ascii = False)) + except Exception: + pass + return _tokenize("\n".join(p for p in parts if p)) + except Exception: + if strict: + raise + return 0 + # ── TTS support ──────────────────────────────────────────── def detect_audio_type(self) -> Optional[str]: diff --git a/studio/backend/main.py b/studio/backend/main.py index 321995885c..1794e0a9af 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -247,6 +247,7 @@ from utils.update_status import ( get_studio_update_status, ) from utils.studio_version import get_studio_version +from utils.api_errors import install_api_error_handlers def get_unsloth_version() -> str: @@ -719,6 +720,7 @@ app.add_middleware( allow_headers = ["*"], ) + # ============ Register API Routes ============ # Register routers @@ -743,6 +745,10 @@ app.include_router(training_history_router, prefix = "/api/train", tags = ["trai app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"]) app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"]) +# Re-wrap client-error responses on the /v1/* surface into OpenAI/Anthropic +# error envelopes; non-/v1 paths keep FastAPI's default {"detail": ...} shape. +install_api_error_handlers(app) + # ============ Health and System Endpoints ============ @@ -1055,8 +1061,12 @@ def setup_frontend(app: FastAPI, build_path: Path): @app.get("/{full_path:path}") async def serve_frontend(request: Request, full_path: str): + # Unknown API paths: raise a real 404 so the api_errors handlers can + # render the correct envelope for /v1/* (and {"detail":...} for /api/*). + # This handler only sees paths NOT matched by a real route. The full + # request path is "/" + full_path. if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")): - return {"error": "API endpoint not found"} + raise HTTPException(status_code = 404, detail = "API endpoint not found") file_path = (build_path / full_path).resolve() diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 9005f087be..3482303eeb 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -495,7 +495,9 @@ class ChatMessage(BaseModel): ``role="tool"`` is resolved at the ``ChatCompletionRequest`` layer. """ - role: Literal["system", "user", "assistant", "tool"] = Field(..., description = "Message role") + role: Literal["system", "user", "assistant", "tool", "developer"] = Field( + ..., description = "Message role" + ) content: Optional[Union[str, list[ContentPart]]] = Field( None, description = "Message content (string or multimodal parts)" ) @@ -590,6 +592,34 @@ class ChatCompletionRequest(BaseModel): "{'type': 'function', 'function': {'name': ...}}" ), ) + max_completion_tokens: Optional[int] = Field( + None, + ge = 1, + description = "OpenAI upper bound on generated tokens (supersedes the deprecated max_tokens).", + ) + n: Optional[int] = Field( + None, + ge = 1, + le = 128, + description = "Number of chat completion choices to generate.", + ) + logprobs: Optional[bool] = Field( + None, description = "Whether to return log probabilities of the output tokens." + ) + top_logprobs: Optional[int] = Field( + None, + ge = 0, + le = 20, + description = "Number of most likely tokens (0-20) to return per position; requires logprobs=true.", + ) + parallel_tool_calls: Optional[bool] = Field( + None, description = "Whether to enable parallel function calling during tool use." + ) + seed: Optional[int] = Field(None, description = "Best-effort deterministic sampling seed.") + stream_options: Optional[dict] = Field( + None, + description = 'Streaming options, e.g. {"include_usage": true} to emit a final usage chunk.', + ) # ── Unsloth extensions (ignored by standard OpenAI clients) ── top_k: int = Field(20, ge = -1, le = 100, description = "[x-unsloth] Top-k sampling") @@ -933,12 +963,16 @@ class ChoiceDelta(BaseModel): content: Optional[str] = None +OpenAIFinishReason = Literal["stop", "length", "tool_calls", "content_filter", "function_call"] + + class ChunkChoice(BaseModel): """A single choice in a streaming chunk.""" index: int = 0 delta: ChoiceDelta - finish_reason: Optional[Literal["stop", "length"]] = None + finish_reason: Optional[OpenAIFinishReason] = None + logprobs: Optional[dict] = None class ChatCompletionChunk(BaseModel): @@ -961,6 +995,7 @@ class CompletionMessage(BaseModel): role: Literal["assistant"] = "assistant" content: str + refusal: Optional[str] = None class CompletionChoice(BaseModel): @@ -968,7 +1003,8 @@ class CompletionChoice(BaseModel): index: int = 0 message: CompletionMessage - finish_reason: Literal["stop", "length"] = "stop" + finish_reason: OpenAIFinishReason = "stop" + logprobs: Optional[dict] = None class CompletionUsage(BaseModel): @@ -977,6 +1013,17 @@ class CompletionUsage(BaseModel): prompt_tokens: int = 0 completion_tokens: int = 0 total_tokens: int = 0 + prompt_tokens_details: Optional[dict] = Field( + default_factory = lambda: {"cached_tokens": 0, "audio_tokens": 0} + ) + completion_tokens_details: Optional[dict] = Field( + default_factory = lambda: { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + } + ) class ChatCompletion(BaseModel): @@ -988,6 +1035,7 @@ class ChatCompletion(BaseModel): model: str = "default" choices: list[CompletionChoice] usage: CompletionUsage = Field(default_factory = CompletionUsage) + system_fingerprint: Optional[str] = None # ===================================================================== @@ -1434,6 +1482,8 @@ class AnthropicMessagesRequest(BaseModel): class AnthropicUsage(BaseModel): input_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 output_tokens: int = 0 diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py index 3874bd587c..59714380da 100644 --- a/studio/backend/routes/data_recipe/jobs.py +++ b/studio/backend/routes/data_recipe/jobs.py @@ -243,6 +243,12 @@ def _inject_local_structured_response_format( "type": "json_schema", "schema": output_format, } + # The OpenAI chat endpoint now returns raw JSON by default for + # response_format requests (spec compliance for public clients). This + # internal opt-in flag rides through the OpenAI SDK's extra_body + # passthrough alongside response_format and re-enables the ```json + # markdown fence that data_designer's structured-output parser expects. + extra_body["_unsloth_guided_fence"] = True params["extra_body"] = extra_body new_configs.append(clone) column["model_alias"] = clone_alias diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 28634396f5..a2e0e5a811 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -26,6 +26,8 @@ import re as _re # Model size extraction (shared with core/inference/llama_cpp.py) from utils.models import extract_model_size_b as _extract_model_size_b +from utils.api_errors import openai_error_body, anthropic_error_body + def _install_httpcore_asyncgen_silencer() -> None: """Silence benign httpx/httpcore asyncgen GC noise on Python 3.13. @@ -107,6 +109,273 @@ def _friendly_error(exc: Exception) -> str: return "An internal error occurred" +def _clamp_finish_reason(value) -> str: + """Coerce an upstream finish_reason into OpenAI's known chat values. + + Unknown values (including ``None``) become ``"stop"`` so local upstream + quirks do not leak into the public API shape. + """ + return ( + value + if value + in ( + "stop", + "length", + "tool_calls", + "content_filter", + "function_call", + ) + else "stop" + ) + + +def _normalize_stop_sequences(raw): + """Coerce an OpenAI/Anthropic ``stop`` value into the list-of-non-empty-strings + shape llama-server expects, or ``None`` when absent. A bare string becomes a + single-element list; empty strings are dropped (an empty stop sequence would + terminate generation immediately at position 0).""" + if isinstance(raw, str): + return [raw] if raw else None + if isinstance(raw, list): + return [s for s in raw if isinstance(s, str) and s] or None + return None + + +def _effective_max_tokens(payload): + """Resolve the generation cap, preferring OpenAI's replacement field. + + ``max_tokens`` is deprecated in favor of ``max_completion_tokens``; honor + either for compatibility, but let the replacement field win when both are + supplied. + """ + return ( + payload.max_completion_tokens + if payload.max_completion_tokens is not None + else payload.max_tokens + ) + + +def _wants_multiple_choices(payload) -> bool: + return (payload.n or 1) > 1 + + +def _raise_unsupported_openai_parameter(param: str, message: str) -> None: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + message, + status = 400, + code = "unsupported_parameter", + param = param, + ), + ) + + +def _raise_unsupported_n(path_label: str) -> None: + _raise_unsupported_openai_parameter("n", f"n > 1 is not supported for {path_label}.") + + +def _openai_stream_error_chunk(exc) -> dict: + """Build an in-band OpenAI error chunk for a mid-stream failure. Once the + stream's 200 headers are flushed the status can't change, so the error must + ride in the SSE body. An upstream context-window overflow is mapped to + code=context_length_exceeded so client compaction/trim loops can detect it + (a code-less error hides it).""" + _cls = _classify_llama_generation_error(exc) + if _cls: + return openai_error_body(_friendly_error(exc), status = 400, code = "context_length_exceeded") + if _cls is False: + return openai_error_body(_friendly_error(exc), status = 400) + return openai_error_body(_friendly_error(exc), status = 500) + + +def _openai_passthrough_error(status_code, text) -> "HTTPException": + """HTTPException for a non-200 upstream response on the OpenAI passthrough + (tools / response_format). An over-context upstream error is mapped to a 400 + with code="context_length_exceeded" so these paths deliver the same signal as + the non-passthrough path; any other upstream error keeps llama-server's + message verbatim.""" + if _classify_llama_generation_error(Exception(text)): + return HTTPException( + status_code = 400, + detail = openai_error_body( + _friendly_error(Exception(text)), + status = 400, + code = "context_length_exceeded", + param = "messages", + ), + ) + return HTTPException( + status_code = status_code, + detail = f"llama-server error: {text[:500]}", + ) + + +def _anthropic_stream_error_event(exc): + """Anthropic in-band SSE ``error`` event for a mid-stream failure, or ``None`` + to fall through to a normal message_delta finish. Returns an event only for a + classifiable upstream client error (context overflow / 4xx) so a streaming + over-context request surfaces a real error instead of a silent empty + end_turn message.""" + if _classify_llama_generation_error(exc) is None: + return None + return build_anthropic_sse_event( + "error", + anthropic_error_body(_friendly_error(exc), status = 400), + ) + + +def _drop_parallel_tool_call_deltas(chunk) -> bool: + """In-place: drop tool_call deltas whose index >= 1 from a parsed OpenAI + streaming chunk so only the first tool call survives (parallel_tool_calls=false + / disable_parallel_tool_use, best-effort). Returns True if anything changed.""" + if not isinstance(chunk, dict): + return False + changed = False + for ch in chunk.get("choices") or []: + delta = ch.get("delta") or {} + tcs = delta.get("tool_calls") + if isinstance(tcs, list): + kept = [tc for tc in tcs if isinstance(tc, dict) and (tc.get("index") or 0) == 0] + if len(kept) != len(tcs): + delta["tool_calls"] = kept + changed = True + return changed + + +def _cap_parallel_tool_calls_sse_line(raw_line: str) -> str: + """Drop tool_call deltas whose index >= 1 from one streamed OpenAI SSE + ``data:`` line so only the first tool call survives (parallel_tool_calls=false, + best-effort). Non-tool / unparseable payloads are returned byte-for-byte.""" + payload = raw_line[len("data: ") :] + if payload.strip() in ("", "[DONE]"): + return raw_line + try: + obj = json.loads(payload) + except Exception: + return raw_line + if not _drop_parallel_tool_call_deltas(obj): + return raw_line + return "data: " + json.dumps(obj, separators = (",", ":")) + + +def _prompt_tokens_details(upstream): + """Surface llama-server's real ``cached_tokens`` (KV-cache prompt hits) while + keeping the full OpenAI ``prompt_tokens_details`` shape. Defaults to zero when + the upstream usage doesn't carry it, so the field is always present.""" + out = {"cached_tokens": 0, "audio_tokens": 0} + if isinstance(upstream, dict): + out.update({k: v for k, v in upstream.items() if v is not None}) + return out + + +def _wants_stream_usage(payload) -> bool: + return bool((payload.stream_options or {}).get("include_usage")) + + +def _openai_stream_usage_chunk( + payload, completion_id, created, model_name, stream_usage, stream_timings +): + """Build the final OpenAI-standard usage chunk (choices=[], usage populated) + for a chat stream. Returns the SSE ``data:`` line, or None when the client + did not opt in via ``stream_options.include_usage`` (or no usage exists).""" + if not _wants_stream_usage(payload): + return None + if not (stream_usage or stream_timings): + return None + _usage = stream_usage or {} + _prompt_tokens = _usage.get("prompt_tokens") or 0 + _completion_tokens = _usage.get("completion_tokens") or 0 + _total_tokens = _usage.get("total_tokens") or (_prompt_tokens + _completion_tokens) + usage_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [], + usage = CompletionUsage( + prompt_tokens = _prompt_tokens, + completion_tokens = _completion_tokens, + total_tokens = _total_tokens, + prompt_tokens_details = _prompt_tokens_details(_usage.get("prompt_tokens_details")), + ), + timings = stream_timings, + ) + return f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" + + +def _rewrite_cmpl_id(raw: bytes) -> bytes: + """Rewrite llama-server's chat-style ``chatcmpl-`` ids to the ``cmpl-`` + prefix OpenAI's legacy /v1/completions use. Anchored on the ``"id":`` key + (both spacing variants) so the rest of the body stays byte-exact.""" + return raw.replace(b'"id":"chatcmpl-', b'"id":"cmpl-').replace( + b'"id": "chatcmpl-', b'"id": "cmpl-' + ) + + +def _cmpl_stream_event_out(event: bytes, include_usage: bool) -> Optional[bytes]: + """Process one legacy /v1/completions SSE event (text between blank-line + separators). + + Always rewrites the ``chatcmpl-`` -> ``cmpl-`` id prefix. When the client + did NOT request ``stream_options.include_usage``, also removes the usage + statistics so the stream matches OpenAI's contract. + + Shape note: on /v1/completions, llama-server attaches ``usage`` to the + FINAL content chunk (the ``finish_reason`` chunk, which has a populated + ``choices`` array) -- unlike the chat stream, which emits a standalone + ``choices: []`` usage chunk. Both shapes are handled: a standalone + usage-only chunk is dropped; an inline ``usage`` field is stripped from a + content chunk while keeping ``choices``/``finish_reason`` intact. + + Returns the event bytes to emit, or ``None`` to drop the event. Only a + usage-bearing event is re-serialized; every other event keeps exact bytes. + """ + if include_usage: + return _rewrite_cmpl_id(event) + lines = event.split(b"\n") + changed = False + for i, ln in enumerate(lines): + if not ln.startswith(b"data:"): + continue + payload = ln[len(b"data:") :].strip() + if not payload or payload == b"[DONE]": + continue + try: + obj = json.loads(payload) + except Exception: + continue + if not isinstance(obj, dict) or obj.get("usage") is None: + continue + # Standalone usage-only chunk (chat-style) -> drop the whole event. + if obj.get("choices") == []: + return None + # Usage on a content/finish chunk (completions-style) -> strip it. + obj.pop("usage", None) + lines[i] = b"data: " + json.dumps(obj, separators = (",", ":")).encode("utf-8") + changed = True + return _rewrite_cmpl_id(b"\n".join(lines) if changed else event) + + +def _classify_llama_generation_error(exc: Exception) -> Optional[bool]: + """Classify an error raised while consuming the GGUF generator. + + Returns True for a context-window overflow, False for any other upstream + 4xx (a client error), or None when it should stay a 500. Distinguishes a + real client error from a genuine crash by the explicit "llama-server + returned 4xx" marker, not a bare "tokens"/"exceed" substring. + """ + msg = str(exc) + msg_l = msg.lower() + if "n_ctx" in msg_l or ( + "context" in msg_l and any(t in msg_l for t in ("exceed", "length", "window", "too long")) + ): + return True + if _re.search(r"llama-server returned (4\d\d)", msg): + return False + return None + + +# Add backend directory to path backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) @@ -214,6 +483,9 @@ from core.inference.anthropic_compat import ( anthropic_messages_to_openai, anthropic_tools_to_openai, anthropic_tool_choice_to_openai, + openai_finish_to_anthropic_stop, + anthropic_tool_use_id, + build_anthropic_sse_event, AnthropicStreamEmitter, AnthropicPassthroughEmitter, ) @@ -1672,7 +1944,7 @@ async def generate_audio( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_new_tokens = payload.max_tokens or 2048, + max_new_tokens = _effective_max_tokens(payload) or 2048, repetition_penalty = payload.repetition_penalty, ) else: @@ -1689,7 +1961,7 @@ async def generate_audio( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_new_tokens = payload.max_tokens or 2048, + max_new_tokens = _effective_max_tokens(payload) or 2048, repetition_penalty = payload.repetition_penalty, use_adapter = payload.use_adapter, ) @@ -1773,7 +2045,7 @@ def _extract_content_parts(messages: list) -> tuple[str, list[dict], "Optional[s chat_messages: Non-system messages with content flattened to strings. image_base64: Base64 of the *first* image found, or ``None``. """ - system_prompt = "" + system_parts: list[str] = [] chat_messages: list[dict] = [] first_image_b64: Optional[str] = None @@ -1781,10 +2053,10 @@ def _extract_content_parts(messages: list) -> tuple[str, list[dict], "Optional[s # ── System messages → extract as system_prompt ──────── if msg.role == "system": if isinstance(msg.content, str): - system_prompt = msg.content + system_parts.append(msg.content) elif isinstance(msg.content, list): # Unlikely but handle: join text parts - system_prompt = "\n".join(p.text for p in msg.content if p.type == "text") + system_parts.append("\n".join(p.text for p in msg.content if p.type == "text")) continue # ── User / assistant messages ───────────────────────── @@ -1807,7 +2079,7 @@ def _extract_content_parts(messages: list) -> tuple[str, list[dict], "Optional[s combined_text = "\n".join(text_parts) if text_parts else "" chat_messages.append({"role": msg.role, "content": combined_text}) - return system_prompt, chat_messages, first_image_b64 + return "\n\n".join(p for p in system_parts if p), chat_messages, first_image_b64 # ── External provider proxy ────────────────────────────────────── @@ -2227,7 +2499,10 @@ async def _proxy_to_external_provider( model = model, temperature = payload.temperature, top_p = payload.top_p, - max_tokens = payload.max_tokens, + # Honor max_completion_tokens when max_tokens is absent, so a + # provider-routed request capped only by the newer field still gets + # a limit instead of falling back to the provider default. + max_tokens = _effective_max_tokens(payload), presence_penalty = payload.presence_penalty, top_k = _top_k_explicit, enable_thinking = payload.enable_thinking, @@ -2469,11 +2744,56 @@ async def openai_chat_completions( - GGUF models → llama-server via LlamaCppBackend - Other models → Unsloth/transformers via InferenceBackend """ + # OpenAI's newer "developer" role is equivalent to "system". Normalize it + # before provider routing so external providers (which may not accept the + # "developer" role) get "system" too, matching the local path. + for _m in payload.messages: + if _m.role == "developer": + _m.role = "system" + + if payload.logprobs: + _raise_unsupported_openai_parameter( + "logprobs", "logprobs is not supported for chat completions." + ) + if payload.top_logprobs is not None: + _raise_unsupported_openai_parameter( + "top_logprobs", "top_logprobs is not supported for chat completions." + ) + # ── External provider routing ──────────────────────────────── # encrypted_api_key is optional -- local providers (llama.cpp / vLLM / Ollama) may run without auth. if payload.provider_id or payload.provider_type: + if _wants_multiple_choices(payload): + _raise_unsupported_n("external provider chat completions") return await _proxy_to_external_provider(payload, request) + # Reject a malformed function tool here: it would otherwise reach + # llama-server and surface as an opaque 500 "Failed to parse tools". + if payload.tools: + for _tool in payload.tools: + if not isinstance(_tool, dict): + continue + # llama-server 500s ("Failed to parse tools: Missing tool type") when + # a function tool omits "type". Default it to "function" so a + # well-formed tool isn't rejected over a missing discriminator (and a + # malformed one still surfaces as a clean 400 below, not a 500). + if _tool.get("type") is None and isinstance(_tool.get("function"), dict): + _tool["type"] = "function" + if _tool.get("type") != "function": + continue + _fn = _tool.get("function") + _name = _fn.get("name") if isinstance(_fn, dict) else None + if not isinstance(_name, str) or not _name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tools': each tool must have a 'function' with a 'name'.", + status = 400, + code = "invalid_value", + param = "tools", + ), + ) + llama_backend = get_llama_cpp_backend() using_gguf = llama_backend.is_loaded @@ -2492,9 +2812,13 @@ async def openai_chat_completions( payload.enable_thinking = bool(_tpl_kw["enable_thinking"]) # ── Determine which backend is active ───────────────────── + # Single-model server: any model name serves the loaded model (drop-in + # OpenAI compat), so payload.model is only a fallback label here. if using_gguf: model_name = llama_backend.model_identifier or payload.model if getattr(llama_backend, "_is_audio", False): + if _wants_multiple_choices(payload): + _raise_unsupported_n("GGUF audio chat completions") return await generate_audio(payload, request) else: backend = get_inference_backend() @@ -2504,6 +2828,8 @@ async def openai_chat_completions( detail = "No model loaded. Call POST /inference/load first.", ) model_name = backend.active_model_name or payload.model + if _wants_multiple_choices(payload): + _raise_unsupported_n("non-GGUF chat completions") # ── Audio TTS path: auto-route to audio generation ──── # (Whisper is ASR not TTS -- handled below in audio input path) @@ -2540,7 +2866,7 @@ async def openai_chat_completions( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_new_tokens = payload.max_tokens or 2048, + max_new_tokens = _effective_max_tokens(payload) or 2048, repetition_penalty = payload.repetition_penalty, cancel_event = cancel_event, ) @@ -2634,11 +2960,16 @@ async def openai_chat_completions( # ── Standard OpenAI function-calling pass-through (GGUF only) ──── # When a client (opencode / Claude Code via OpenAI compat / Cursor / # Continue / ...) sends standard OpenAI `tools` without Studio's - # `enable_tools` shorthand, forward to llama-server verbatim so structured - # `tool_calls` flow back. Runs BEFORE `_extract_content_parts` because that - # helper is unaware of `role="tool"` messages and assistant messages that - # only carry `tool_calls` (content=None) -- both valid in multi-turn - # client-side tool loops. + # `enable_tools` shorthand, forward the request to llama-server + # verbatim so structured `tool_calls` flow back to the client. This + # branch runs BEFORE `_extract_content_parts` because that helper is + # unaware of `role="tool"` messages and assistant messages that only + # carry `tool_calls` (content=None) — both of which are valid in + # multi-turn client-side tool loops. + effective_max_tokens = _effective_max_tokens(payload) + + normalized_stop = _normalize_stop_sequences(payload.stop) + _has_tool_messages = any(m.role == "tool" or m.tool_calls for m in payload.messages) # Route guided-decoding requests through the verbatim passthrough so # ``response_format`` (JSON schema) reaches llama-server and the model's @@ -2656,6 +2987,8 @@ async def openai_chat_completions( and not _effective_enable_tools(payload) and (_tools_passthrough or _has_response_format) ): + if _wants_multiple_choices(payload): + _raise_unsupported_n("GGUF tool or response_format passthrough") if payload.audio_base64: raise HTTPException( status_code = 400, @@ -2766,6 +3099,8 @@ async def openai_chat_completions( use_tools = False if use_tools: + if _wants_multiple_choices(payload): + _raise_unsupported_n("GGUF tool chat completions") # ── Tool-use system prompt nudge ────────────────────── _nudge = _build_tool_action_nudge( tools = tools_to_use, @@ -2800,10 +3135,12 @@ async def openai_chat_completions( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_tokens = payload.max_tokens, + max_tokens = effective_max_tokens, repetition_penalty = payload.repetition_penalty, presence_penalty = payload.presence_penalty, + stop = normalized_stop, cancel_event = cancel_event, + seed = payload.seed, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, @@ -2815,6 +3152,7 @@ async def openai_chat_completions( if payload.tool_call_timeout is not None else 300, session_id = payload.session_id, + disable_parallel_tool_use = payload.parallel_tool_calls is False, ) _tool_sentinel = object() @@ -2844,6 +3182,7 @@ async def openai_chat_completions( prev_text = "" _stream_usage = None _stream_timings = None + _stream_finish = None while True: if cancel_event.is_set(): break @@ -2882,6 +3221,7 @@ async def openai_chat_completions( if event["type"] == "metadata": _stream_usage = event.get("usage") _stream_timings = event.get("timings") + _stream_finish = event.get("finish_reason") continue # "content" type -- cumulative text. Sanitize the full @@ -2916,27 +3256,21 @@ async def openai_chat_completions( choices = [ ChunkChoice( delta = ChoiceDelta(), - finish_reason = "stop", + finish_reason = _clamp_finish_reason(_stream_finish), ) ], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" - # Usage chunk (OpenAI standard: choices=[], usage populated) - if _stream_usage or _stream_timings: - usage_obj = CompletionUsage( - prompt_tokens = (_stream_usage or {}).get("prompt_tokens", 0), - completion_tokens = (_stream_usage or {}).get("completion_tokens", 0), - total_tokens = (_stream_usage or {}).get("total_tokens", 0), - ) - usage_chunk = ChatCompletionChunk( - id = completion_id, - created = created, - model = model_name, - choices = [], - usage = usage_obj, - timings = _stream_timings, - ) - yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" + usage_line = _openai_stream_usage_chunk( + payload, + completion_id, + created, + model_name, + _stream_usage, + _stream_timings, + ) + if usage_line is not None: + yield usage_line yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -2947,12 +3281,7 @@ async def openai_chat_completions( tb = traceback.format_exc() logger.error(f"Error during GGUF tool streaming: {e}\n{tb}") - error_chunk = { - "error": { - "message": _friendly_error(e), - "type": "server_error", - }, - } + error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: _tracker.__exit__(None, None, None) @@ -2969,7 +3298,10 @@ async def openai_chat_completions( # ── Standard GGUF path (no tools) ───────────────────── - def gguf_generate(): + def gguf_generate(choice_index: int = 0): + _seed = payload.seed + if _seed is not None and _seed >= 0 and choice_index: + _seed += choice_index return llama_backend.generate_chat_completion( messages = gguf_messages, image_b64 = image_b64, @@ -2977,18 +3309,22 @@ async def openai_chat_completions( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_tokens = payload.max_tokens, + max_tokens = effective_max_tokens, repetition_penalty = payload.repetition_penalty, presence_penalty = payload.presence_penalty, + stop = normalized_stop, cancel_event = cancel_event, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, + seed = _seed, ) _gguf_sentinel = object() if payload.stream: + if _wants_multiple_choices(payload): + _raise_unsupported_n("streaming GGUF chat completions") _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() @@ -3015,6 +3351,7 @@ async def openai_chat_completions( prev_text = "" _stream_usage = None _stream_timings = None + _stream_finish = None while True: if cancel_event.is_set(): break @@ -3029,6 +3366,7 @@ async def openai_chat_completions( if cumulative.get("type") == "metadata": _stream_usage = cumulative.get("usage") _stream_timings = cumulative.get("timings") + _stream_finish = cumulative.get("finish_reason") else: logger.warning( "gguf_stream_chunks: unexpected dict event: %s", @@ -3060,27 +3398,21 @@ async def openai_chat_completions( choices = [ ChunkChoice( delta = ChoiceDelta(), - finish_reason = "stop", + finish_reason = _clamp_finish_reason(_stream_finish), ) ], ) yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" - # Usage chunk (OpenAI standard: choices=[], usage populated) - if _stream_usage or _stream_timings: - usage_obj = CompletionUsage( - prompt_tokens = (_stream_usage or {}).get("prompt_tokens", 0), - completion_tokens = (_stream_usage or {}).get("completion_tokens", 0), - total_tokens = (_stream_usage or {}).get("total_tokens", 0), - ) - usage_chunk = ChatCompletionChunk( - id = completion_id, - created = created, - model = model_name, - choices = [], - usage = usage_obj, - timings = _stream_timings, - ) - yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" + usage_line = _openai_stream_usage_chunk( + payload, + completion_id, + created, + model_name, + _stream_usage, + _stream_timings, + ) + if usage_line is not None: + yield usage_line yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -3088,12 +3420,7 @@ async def openai_chat_completions( raise except Exception as e: logger.error(f"Error during GGUF streaming: {e}", exc_info = True) - error_chunk = { - "error": { - "message": _friendly_error(e), - "type": "server_error", - }, - } + error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: _tracker.__exit__(None, None, None) @@ -3109,35 +3436,74 @@ async def openai_chat_completions( ) else: try: - full_text = "" - completion_usage = None - for token in gguf_generate(): - if isinstance(token, dict): - if token.get("type") == "metadata": - completion_usage = token.get("usage") - continue - full_text = token + # ``n`` requests several independent completions; the single + # decode slot yields one at a time, so loop sequentially. + _n = payload.n or 1 + + _choices = [] + _prompt_tokens = 0 + _sum_completion = 0 + _prompt_details = None + for _idx in range(_n): + # Stop spawning the remaining choices once cancelled. + if cancel_event.is_set(): + break + full_text = "" + completion_usage = None + completion_finish = None + for token in gguf_generate(_idx): + if isinstance(token, dict): + if token.get("type") == "metadata": + completion_usage = token.get("usage") + completion_finish = token.get("finish_reason") + continue + full_text = token + + _choices.append( + CompletionChoice( + index = _idx, + message = CompletionMessage(content = full_text), + finish_reason = _clamp_finish_reason(completion_finish), + ) + ) + if completion_usage: + # The prompt is shared across all n choices, so count its + # tokens ONCE (OpenAI bills only generated tokens for each + # extra choice). Only completion_tokens accumulates. + _prompt_tokens = completion_usage.get("prompt_tokens") or _prompt_tokens + _sum_completion += completion_usage.get("completion_tokens") or 0 + if _prompt_details is None: + _prompt_details = completion_usage.get("prompt_tokens_details") response = ChatCompletion( id = completion_id, created = created, model = model_name, - choices = [ - CompletionChoice( - message = CompletionMessage(content = full_text), - finish_reason = "stop", - ) - ], + choices = _choices, usage = CompletionUsage( - prompt_tokens = (completion_usage or {}).get("prompt_tokens") or 0, - completion_tokens = (completion_usage or {}).get("completion_tokens") or 0, - total_tokens = (completion_usage or {}).get("total_tokens") or 0, + prompt_tokens = _prompt_tokens, + completion_tokens = _sum_completion, + total_tokens = _prompt_tokens + _sum_completion, + prompt_tokens_details = _prompt_tokens_details(_prompt_details), ), ) return JSONResponse(content = response.model_dump()) except Exception as e: logger.error(f"Error during GGUF completion: {e}", exc_info = True) + # An over-context prompt makes llama-server return 400; map any + # upstream 4xx to a 400 client error rather than leaking a 500. + _cls = _classify_llama_generation_error(e) + if _cls is not None: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + _friendly_error(e), + status = 400, + code = "context_length_exceeded" if _cls else None, + param = "messages", + ), + ) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) # ── Standard Unsloth path ───────────────────────────────── @@ -3279,7 +3645,7 @@ async def openai_chat_completions( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_tokens = payload.max_tokens, + max_tokens = effective_max_tokens, repetition_penalty = payload.repetition_penalty, cancel_event = cancel_event, enable_thinking = payload.enable_thinking, @@ -3388,20 +3754,16 @@ async def openai_chat_completions( # concurrent streams cannot read each other's stats. _stats = _sf_stats_holder.get("stats") if _stats: - _stream_usage = _stats.get("usage") or {} - usage_chunk = ChatCompletionChunk( - id = completion_id, - created = created, - model = model_name, - choices = [], - usage = CompletionUsage( - prompt_tokens = _stream_usage.get("prompt_tokens", 0), - completion_tokens = _stream_usage.get("completion_tokens", 0), - total_tokens = _stream_usage.get("total_tokens", 0), - ), - timings = _stats.get("timings"), + usage_line = _openai_stream_usage_chunk( + payload, + completion_id, + created, + model_name, + _stats.get("usage"), + _stats.get("timings"), ) - yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" + if usage_line is not None: + yield usage_line yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -3483,7 +3845,7 @@ async def openai_chat_completions( top_p = payload.top_p, top_k = payload.top_k, min_p = payload.min_p, - max_new_tokens = payload.max_tokens or 2048, + max_new_tokens = effective_max_tokens or 2048, repetition_penalty = payload.repetition_penalty, ) # Forward reasoning kwargs; the worker/template wrapper peels off any the @@ -3596,20 +3958,16 @@ async def openai_chat_completions( # read each other's stats. _stats = stats_holder.get("stats") if _stats: - _stream_usage = _stats.get("usage") or {} - usage_chunk = ChatCompletionChunk( - id = completion_id, - created = created, - model = model_name, - choices = [], - usage = CompletionUsage( - prompt_tokens = _stream_usage.get("prompt_tokens", 0), - completion_tokens = _stream_usage.get("completion_tokens", 0), - total_tokens = _stream_usage.get("total_tokens", 0), - ), - timings = _stats.get("timings"), + usage_line = _openai_stream_usage_chunk( + payload, + completion_id, + created, + model_name, + _stats.get("usage"), + _stats.get("timings"), ) - yield f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" + if usage_line is not None: + yield usage_line yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -3757,15 +4115,14 @@ async def serve_sandbox_file( # ===================================================================== -@router.get("/models") -async def openai_list_models(current_subject: str = Depends(get_current_subject)): - """ - OpenAI-compatible model listing endpoint. +def _openai_model_objects() -> list[dict]: + """The model objects GET /v1/models exposes (one per loaded local backend). - Returns the currently loaded model in the format expected by - OpenAI-compatible clients (``GET /v1/models``). + Shared by the LIST and RETRIEVE handlers so both report the same ids and + field shape. """ - models = [] + models: list[dict] = [] + _created = int(time.time()) # Check GGUF backend llama_backend = get_llama_cpp_backend() @@ -3774,6 +4131,7 @@ async def openai_list_models(current_subject: str = Depends(get_current_subject) { "id": llama_backend.model_identifier, "object": "model", + "created": _created, "owned_by": "local", } ) @@ -3785,11 +4143,46 @@ async def openai_list_models(current_subject: str = Depends(get_current_subject) { "id": backend.active_model_name, "object": "model", + "created": _created, "owned_by": "local", } ) - return {"object": "list", "data": models} + return models + + +@router.get("/models") +async def openai_list_models(current_subject: str = Depends(get_current_subject)): + """ + OpenAI-compatible model listing endpoint. + + Returns the currently loaded model in the format expected by + OpenAI-compatible clients (``GET /v1/models``). + """ + return {"object": "list", "data": _openai_model_objects()} + + +@router.get("/models/{model_id:path}") +async def openai_retrieve_model(model_id: str, current_subject: str = Depends(get_current_subject)): + """ + OpenAI-compatible single-model retrieval endpoint (``GET /v1/models/{id}``). + + Returns the bare model object when ``model_id`` matches a loaded local + model, or 404 model_not_found otherwise. Defined after the LIST route so + it does not shadow it; ``{model_id:path}`` keeps ids with slashes intact. + """ + for model in _openai_model_objects(): + if model["id"] == model_id: + return model + raise HTTPException( + status_code = 404, + detail = openai_error_body( + f"The model '{model_id}' does not exist", + status = 404, + code = "model_not_found", + param = "id", + ), + ) # ===================================================================== @@ -3819,14 +4212,17 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge if is_stream: async def _stream(): - # Manual httpx client/response lifecycle AND explicit aiter_bytes() - # iterator close -- see _anthropic_passthrough_stream for the full - # rationale. Saving `bytes_iter = resp.aiter_bytes()` and `await - # bytes_iter.aclose()` in finally is what avoids the Python 3.13 + - # httpcore 1.0.x "Exception ignored in: " / anyio - # cancel-scope trace: an anonymous async for leaves the iterator - # unclosed, so the asyncgen GC finalizer runs cleanup later in a - # different asyncio task. + # Manual httpx client/response lifecycle AND explicit iterator + # close — see _anthropic_passthrough_stream for the full rationale. + # Saving the iterator and closing it in the finally block avoids the + # Python 3.13 + httpcore 1.0.x "Exception ignored in: + # " / anyio cancel-scope trace. + # + # Buffer the relay into whole SSE events (split on the blank-line + # separator) so _cmpl_stream_event_out can rewrite the cmpl- id and + # honor stream_options.include_usage per event, while keeping SSE + # framing and token bytes intact. + _include_usage = bool((body.get("stream_options") or {}).get("include_usage")) client = httpx.AsyncClient(timeout = 600) resp = None bytes_iter = None @@ -3834,8 +4230,21 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge req = client.build_request("POST", target_url, json = body) resp = await client.send(req, stream = True) bytes_iter = resp.aiter_bytes() + buffer = b"" async for chunk in bytes_iter: - yield chunk + buffer += chunk + while b"\n\n" in buffer: + event, buffer = buffer.split(b"\n\n", 1) + out = _cmpl_stream_event_out(event, _include_usage) + if out is not None: + yield out + b"\n\n" + if buffer: + out = _cmpl_stream_event_out(buffer, _include_usage) + if out is not None: + # Re-add the SSE separator the split consumed, so a final + # event arriving without a trailing blank line is still + # terminated for the client's parser. + yield out + b"\n\n" except Exception as e: logger.error("openai_completions stream error: %s", e) finally: @@ -3858,8 +4267,12 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge else: async with httpx.AsyncClient() as client: resp = await client.post(target_url, json = body, timeout = 600) + + if resp.status_code != 200: + raise _openai_passthrough_error(resp.status_code, resp.text) + return Response( - content = resp.content, + content = _rewrite_cmpl_id(resp.content), status_code = resp.status_code, media_type = "application/json", ) @@ -4139,14 +4552,10 @@ def _build_chat_request( chat_tool_choice = _translate_responses_tool_choice_to_chat(payload.tool_choice) if chat_tool_choice is not None: chat_kwargs["tool_choice"] = chat_tool_choice + if payload.parallel_tool_calls is not None: + chat_kwargs["parallel_tool_calls"] = payload.parallel_tool_calls - req = ChatCompletionRequest(**chat_kwargs) - # `parallel_tool_calls` isn't a first-class field on ChatCompletionRequest, - # but the model allows extras and _build_openai_passthrough_body forwards - # only known fields. Llama-server doesn't implement parallel_tool_calls - # semantics, so accept-and-ignore it on the Responses side to avoid breaking - # SDK clients that always send it. - return req + return ChatCompletionRequest(**chat_kwargs) def _chat_tool_calls_to_responses_output(tool_calls: list[dict]) -> list[dict]: @@ -4295,6 +4704,7 @@ async def _responses_stream( ) body = _build_openai_passthrough_body(chat_req, backend_ctx = llama_backend.context_length) + body["stream_options"] = {"include_usage": True} target_url = f"{llama_backend.base_url}/v1/chat/completions" async def event_generator(): @@ -4399,6 +4809,8 @@ async def _responses_stream( chunk_data = json.loads(data_str) except json.JSONDecodeError: continue + if payload.parallel_tool_calls is False: + _drop_parallel_tool_call_deltas(chunk_data) choices = chunk_data.get("choices", []) if not choices: @@ -4720,6 +5132,56 @@ def _normalize_anthropic_openai_images(openai_messages: list[dict], is_vision: b return has_image +@router.post("/messages/count_tokens") +async def anthropic_count_tokens( + payload: AnthropicMessagesRequest, + request: Request, + current_subject: str = Depends(get_current_subject), +): + """Anthropic-compatible token-counting endpoint (POST /v1/messages/count_tokens). + + Translates the Anthropic request to OpenAI form (the same translation the + /messages handler uses), counts prompt tokens with the loaded GGUF model's + tokenizer, and returns ``{"input_tokens": int}`` only. Unlike /messages, + max_tokens is NOT required here. + """ + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + # Same Anthropic → OpenAI translation as anthropic_messages: system is + # folded into the messages list, so pass system=None to the counter. + openai_messages = anthropic_messages_to_openai( + [m.model_dump() for m in payload.messages], + payload.system, + ) + # Apply the same sanitization /messages does before generation, so the count + # matches the prompt the real request would build (otherwise empty-assistant + # sentinels / synthetic tool history inflate the count or hit the fallback). + openai_messages = _strip_provider_synthetic_tool_history( + _drop_empty_assistant_sentinels(openai_messages) + ) + openai_tools = anthropic_tools_to_openai(payload.tools or []) or None + + try: + count = await asyncio.to_thread( + llama_backend.count_chat_tokens, + openai_messages, + None, + openai_tools, + strict = True, + ) + except Exception: + raise HTTPException( + status_code = 503, + detail = "Unable to count tokens with the loaded model tokenizer.", + ) + return JSONResponse(content = {"input_tokens": int(count)}) + + def _set_or_prepend_system_message( messages: Optional[list[dict]], system_prompt: str ) -> list[dict]: @@ -4755,6 +5217,18 @@ async def anthropic_messages( detail = "No GGUF model loaded. Load a GGUF model first.", ) + # max_tokens is a required field on the Anthropic Messages API; real + # Anthropic returns a 400 invalid_request_error when it is omitted. + if payload.max_tokens is None: + raise HTTPException( + status_code = 400, + detail = anthropic_error_body( + "max_tokens: field required", + status = 400, + err_type = "invalid_request_error", + ), + ) + model_name = getattr(llama_backend, "model_identifier", None) or payload.model message_id = f"msg_{uuid.uuid4().hex[:24]}" @@ -4866,6 +5340,14 @@ async def anthropic_messages( not server_tools and len(openai_client_tools) > 0 and llama_backend.supports_tools ) + # Anthropic tool_choice.disable_parallel_tool_use caps the response to a + # single tool_use block. Computed here so BOTH the client-tool passthrough + # and the server-tool path honor it. + _disable_parallel = bool( + isinstance(payload.tool_choice, dict) + and payload.tool_choice.get("disable_parallel_tool_use") + ) + # ── Client-side pass-through path ───────────────────────── if client_tools: openai_tools = openai_client_tools @@ -4890,6 +5372,7 @@ async def anthropic_messages( tool_choice = openai_tool_choice, session_id = payload.session_id, cancel_id = payload.cancel_id, + disable_parallel_tool_use = _disable_parallel, ) return await _anthropic_passthrough_non_streaming( llama_backend, @@ -4906,6 +5389,7 @@ async def anthropic_messages( repetition_penalty = repetition_penalty, presence_penalty = presence_penalty, tool_choice = openai_tool_choice, + disable_parallel_tool_use = _disable_parallel, ) if server_tools: @@ -4954,6 +5438,7 @@ async def anthropic_messages( auto_heal_tool_calls = True, tool_call_timeout = 300, session_id = payload.session_id, + disable_parallel_tool_use = _disable_parallel, ) if payload.stream: @@ -4963,11 +5448,16 @@ async def anthropic_messages( _run_tool_gen, message_id, model_name, + llama_backend = llama_backend, + openai_messages = openai_messages, + openai_tools = openai_tools, + disable_parallel_tool_use = _disable_parallel, ) return await _anthropic_tool_non_streaming( _run_tool_gen, message_id, model_name, + disable_parallel_tool_use = _disable_parallel, ) # ── No-tool path ────────────────────────────────────────── @@ -4992,6 +5482,8 @@ async def anthropic_messages( _run_plain_gen, message_id, model_name, + llama_backend = llama_backend, + openai_messages = openai_messages, ) return await _anthropic_plain_non_streaming( _run_plain_gen, @@ -5000,15 +5492,44 @@ async def anthropic_messages( ) -async def _anthropic_tool_stream(request, cancel_event, run_gen, message_id, model_name): +async def _anthropic_tool_stream( + request, + cancel_event, + run_gen, + message_id, + model_name, + llama_backend = None, + openai_messages = None, + openai_tools = None, + disable_parallel_tool_use = False, +): """Streaming response for the tool-calling path.""" _sentinel = object() + # Prompt-token count for message_start.usage.input_tokens. count_chat_tokens + # makes blocking HTTP calls to llama-server, so run it off the event loop. + # Pass the tools so tool-schema tokens are counted (the generator renders + # them too), matching the non-stream / count_tokens / passthrough paths. + input_tokens = 0 + if llama_backend is not None and openai_messages is not None: + input_tokens = await asyncio.to_thread( + llama_backend.count_chat_tokens, openai_messages, None, openai_tools + ) + async def _stream(): emitter = AnthropicStreamEmitter() - for line in emitter.start(message_id, model_name): + for line in emitter.start(message_id, model_name, input_tokens = input_tokens): yield line + captured_finish_reason = None + # Whether the response currently ends on a pending tool_use block (the + # client must act → stop_reason "tool_use") as opposed to final text. + # The server may run a tool and then keep generating, which flips this + # back to False — that is an end_turn (or max_tokens) response. + ends_on_tool_use = False + tool_blocks_emitted = 0 + drop_until_tool_end = False + gen = run_gen() try: while True: @@ -5018,16 +5539,53 @@ async def _anthropic_tool_stream(request, cancel_event, run_gen, message_id, mod event = await asyncio.to_thread(next, gen, _sentinel) if event is _sentinel: break - # Strip leaked tool-call XML from content events - if event.get("type") == "content": + etype = event.get("type") + if drop_until_tool_end: + # disable_parallel_tool_use: a later tool call is being + # dropped — skip every event until (and including) its tool_end. + if etype == "tool_end": + drop_until_tool_end = False + continue + if etype == "metadata": + _fr = event.get("finish_reason") + if _fr is not None: + captured_finish_reason = _fr + # Strip leaked tool-call XML from content events first, so a + # content event that was purely tool XML doesn't count as text. + if etype == "content": event = dict(event) event["text"] = _TOOL_XML_RE.sub("", event["text"]) + # disable_parallel_tool_use: keep only the first tool_use block, + # dropping every later tool_start and its paired tool_end (robust + # to empty tool-call ids — tracked by state, not id matching). + if etype == "tool_start": + if disable_parallel_tool_use and tool_blocks_emitted >= 1: + drop_until_tool_end = True + continue + ends_on_tool_use = True + elif etype == "tool_end": + tool_blocks_emitted += 1 + # A tool_end means Studio executed the tool server-side, so + # the response no longer ends on a pending client action. + # Without this, a server tool that produces no trailing text + # would be mislabeled stop_reason "tool_use", telling the + # client to run a tool Studio already ran. + ends_on_tool_use = False + elif etype == "content" and event.get("text"): + ends_on_tool_use = False for line in emitter.feed(event): yield line except Exception as e: logger.error("anthropic_messages stream error: %s", e) + _error_event = _anthropic_stream_error_event(e) + if _error_event is not None: + yield _error_event + return - for line in emitter.finish("end_turn"): + stop_reason = openai_finish_to_anthropic_stop( + captured_finish_reason, had_tool_calls = ends_on_tool_use + ) + for line in emitter.finish(stop_reason = stop_reason, stop_sequence = None): yield line return StreamingResponse( @@ -5041,15 +5599,31 @@ async def _anthropic_tool_stream(request, cancel_event, run_gen, message_id, mod ) -async def _anthropic_plain_stream(request, cancel_event, run_gen, message_id, model_name): +async def _anthropic_plain_stream( + request, + cancel_event, + run_gen, + message_id, + model_name, + llama_backend = None, + openai_messages = None, +): """Streaming response for the no-tool path.""" _sentinel = object() + # Prompt-token count for message_start.usage.input_tokens. count_chat_tokens + # makes blocking HTTP calls to llama-server, so run it off the event loop. + input_tokens = 0 + if llama_backend is not None and openai_messages is not None: + input_tokens = await asyncio.to_thread(llama_backend.count_chat_tokens, openai_messages) + async def _stream(): emitter = AnthropicStreamEmitter() - for line in emitter.start(message_id, model_name): + for line in emitter.start(message_id, model_name, input_tokens = input_tokens): yield line + captured_finish_reason = None + gen = run_gen() try: while True: @@ -5061,6 +5635,9 @@ async def _anthropic_plain_stream(request, cancel_event, run_gen, message_id, mo break if isinstance(cumulative, dict): if cumulative.get("type") == "metadata": + _fr = cumulative.get("finish_reason") + if _fr is not None: + captured_finish_reason = _fr for line in emitter.feed(cumulative): yield line continue @@ -5069,8 +5646,13 @@ async def _anthropic_plain_stream(request, cancel_event, run_gen, message_id, mo yield line except Exception as e: logger.error("anthropic_messages stream error: %s", e) + _error_event = _anthropic_stream_error_event(e) + if _error_event is not None: + yield _error_event + return - for line in emitter.finish("end_turn"): + stop_reason = openai_finish_to_anthropic_stop(captured_finish_reason, had_tool_calls = False) + for line in emitter.finish(stop_reason = stop_reason, stop_sequence = None): yield line return StreamingResponse( @@ -5084,7 +5666,38 @@ async def _anthropic_plain_stream(request, cancel_event, run_gen, message_id, mo ) -async def _anthropic_tool_non_streaming(run_gen, message_id, model_name): +def _anthropic_map_generation_error(e: Exception) -> HTTPException: + """Map an upstream 4xx / context-overflow generation error to a clean + Anthropic 400 invalid_request_error. Genuine 5xx errors stay 500.""" + if _classify_llama_generation_error(e) is not None: + return HTTPException( + status_code = 400, + detail = anthropic_error_body( + _friendly_error(e), + status = 400, + err_type = "invalid_request_error", + ), + ) + return HTTPException(status_code = 500, detail = _friendly_error(e)) + + +def _collect_anthropic_events(run_gen) -> list: + """Drain the generator into a list, mapping an upstream 4xx / context + overflow to a clean Anthropic 400 instead of leaking a 500.""" + try: + return list(run_gen()) + except HTTPException: + raise + except Exception as e: + raise _anthropic_map_generation_error(e) + + +async def _anthropic_tool_non_streaming( + run_gen, + message_id, + model_name, + disable_parallel_tool_use = False, +): """Non-streaming response for the tool-calling path. Builds ``content_blocks`` in generation order (text → tool_use → text → @@ -5101,8 +5714,14 @@ async def _anthropic_tool_non_streaming(run_gen, message_id, model_name): tool_blocks_by_id: dict[str, AnthropicResponseToolUseBlock] = {} usage = {} prev_text = "" + captured_finish_reason = None + # Pending client tool_use; cleared by tool_end (server execution) or + # trailing text. See the stop_reason mapping below. + ends_on_tool_use = False - for event in run_gen(): + events = _collect_anthropic_events(run_gen) + + for event in events: etype = event.get("type", "") if etype == "content": # Strip leaked tool-call XML @@ -5110,6 +5729,7 @@ async def _anthropic_tool_non_streaming(run_gen, message_id, model_name): new = clean[len(prev_text) :] prev_text = clean if new: + ends_on_tool_use = False if content_blocks and isinstance(content_blocks[-1], AnthropicResponseTextBlock): content_blocks[-1].text += new else: @@ -5125,23 +5745,50 @@ async def _anthropic_tool_non_streaming(run_gen, message_id, model_name): existing_tool_block.name = event["tool_name"] else: tool_block = AnthropicResponseToolUseBlock( - id = tool_call_id, + id = anthropic_tool_use_id(tool_call_id), name = event["tool_name"], input = arguments, ) if tool_call_id: tool_blocks_by_id[tool_call_id] = tool_block content_blocks.append(tool_block) + ends_on_tool_use = True elif etype == "tool_end": prev_text = "" + # Server-executed: no longer pending a client action (see above). + ends_on_tool_use = False elif etype == "metadata": usage = event.get("usage", {}) + _fr = event.get("finish_reason") + if _fr is not None: + captured_finish_reason = _fr + + # disable_parallel_tool_use: cap the response to at most one tool_use + # block. Keep the first tool_use and drop any later ones. + if disable_parallel_tool_use: + _seen_tool_use = False + _capped: list = [] + for block in content_blocks: + if isinstance(block, AnthropicResponseToolUseBlock): + if _seen_tool_use: + continue + _seen_tool_use = True + _capped.append(block) + content_blocks = _capped + + # stop_reason "tool_use" only when the response still ends on a pending + # tool_use (client must act). `ends_on_tool_use` is tracked through the + # event stream above: it is True only if the last tool_start had no + # following tool_end (server execution) or trailing text. + stop_reason = openai_finish_to_anthropic_stop( + captured_finish_reason, had_tool_calls = ends_on_tool_use + ) resp = AnthropicMessagesResponse( id = message_id, model = model_name, content = content_blocks, - stop_reason = "end_turn", + stop_reason = stop_reason, usage = AnthropicUsage( input_tokens = usage.get("prompt_tokens", 0), output_tokens = usage.get("completion_tokens", 0), @@ -5155,11 +5802,17 @@ async def _anthropic_plain_non_streaming(run_gen, message_id, model_name): text_parts = [] usage = {} prev_text = "" + captured_finish_reason = None - for cumulative in run_gen(): + events = _collect_anthropic_events(run_gen) + + for cumulative in events: if isinstance(cumulative, dict): if cumulative.get("type") == "metadata": usage = cumulative.get("usage", {}) + _fr = cumulative.get("finish_reason") + if _fr is not None: + captured_finish_reason = _fr continue new = cumulative[len(prev_text) :] prev_text = cumulative @@ -5171,11 +5824,13 @@ async def _anthropic_plain_non_streaming(run_gen, message_id, model_name): if full_text: content_blocks.append(AnthropicResponseTextBlock(text = full_text)) + stop_reason = openai_finish_to_anthropic_stop(captured_finish_reason, had_tool_calls = False) + resp = AnthropicMessagesResponse( id = message_id, model = model_name, content = content_blocks, - stop_reason = "end_turn", + stop_reason = stop_reason, usage = AnthropicUsage( input_tokens = usage.get("prompt_tokens", 0), output_tokens = usage.get("completion_tokens", 0), @@ -5205,6 +5860,8 @@ def _build_passthrough_payload( response_format = None, chat_template_kwargs = None, backend_ctx = None, + seed = None, + stream_options = None, ): body = { "messages": openai_messages, @@ -5215,14 +5872,19 @@ def _build_passthrough_payload( "top_k": top_k, "stream": stream, } - if stream: - body["stream_options"] = {"include_usage": True} + if seed is not None: + body["seed"] = seed + if stream and stream_options is not None: + body["stream_options"] = stream_options body["max_tokens"] = ( max_tokens if max_tokens is not None else (backend_ctx or _DEFAULT_MAX_TOKENS_FLOOR) ) body["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS - if stop: - body["stop"] = stop + # Normalize stop the same way the non-passthrough path does (the passthrough + # was previously the one path that forwarded an empty stop string verbatim). + _stop = _normalize_stop_sequences(stop) + if _stop: + body["stop"] = _stop if min_p is not None: body["min_p"] = min_p if repetition_penalty is not None: @@ -5263,6 +5925,7 @@ async def _anthropic_passthrough_stream( tool_choice = "auto", session_id = None, cancel_id = None, + disable_parallel_tool_use = False, ): """Streaming client-side pass-through: forward tools to llama-server and translate its stream to Anthropic SSE without executing anything.""" @@ -5281,16 +5944,25 @@ async def _anthropic_passthrough_stream( presence_penalty = presence_penalty, tool_choice = tool_choice, backend_ctx = llama_backend.context_length, + stream_options = {"include_usage": True}, ) - # cancel_id mirrors the OpenAI passthrough so a per-run cancel POST works - # without the caller knowing the local message_id. + # Prompt-token count for message_start.usage.input_tokens. count_chat_tokens + # makes blocking HTTP calls to llama-server, so run it off the event loop. + # Pass the tools through so tool-schema tokens are counted (otherwise the + # streaming input_tokens undercounts vs the non-stream / count_tokens paths). + input_tokens = await asyncio.to_thread( + llama_backend.count_chat_tokens, openai_messages, None, openai_tools + ) + + # cancel_id mirrors the OpenAI passthrough so a per-run cancel POST + # works without the caller having to know the local message_id. _tracker = _TrackedCancel(cancel_event, cancel_id, session_id, message_id) _tracker.__enter__() async def _stream(): emitter = AnthropicPassthroughEmitter() - for line in emitter.start(message_id, model_name): + for line in emitter.start(message_id, model_name, input_tokens = input_tokens): yield line # Manage the httpx client, response, AND the aiter_lines() async @@ -5327,10 +5999,31 @@ async def _anthropic_passthrough_stream( req = client.build_request("POST", target_url, json = body) resp = await client.send(req, stream = True) - # See _openai_passthrough_stream for rationale: aiter_lines() blocks - # during llama-server prefill, so the in-loop cancel check is - # unreachable until the first SSE chunk arrives. The watcher closes - # `resp` on cancel, raising in aiter_lines. + # Upstream client error (e.g. over-context 400) arrives before any + # SSE. The 200 stream headers are already flushed, so surface it as + # an in-band Anthropic ``error`` event instead of silently finishing + # with an empty end_turn message. + if resp.status_code != 200: + _err_bytes = await resp.aread() + _err_text = _err_bytes.decode("utf-8", "replace")[:500] + logger.error( + "anthropic passthrough upstream error: status=%s body=%s", + resp.status_code, + _err_text, + ) + yield build_anthropic_sse_event( + "error", + anthropic_error_body( + f"llama-server error: {_err_text}", + status = resp.status_code, + ), + ) + return + + # See _openai_passthrough_stream for rationale: aiter_lines() + # blocks during llama-server prefill, so the in-loop cancel + # check is unreachable until the first SSE chunk arrives. + # The watcher closes `resp` on cancel, raising in aiter_lines. cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) lines_iter = resp.aiter_lines() async for raw_line in lines_iter: @@ -5348,6 +6041,8 @@ async def _anthropic_passthrough_stream( chunk = json.loads(data_str) except json.JSONDecodeError: continue + if disable_parallel_tool_use: + _drop_parallel_tool_call_deltas(chunk) for line in emitter.feed_chunk(chunk): yield line except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError): @@ -5407,6 +6102,7 @@ async def _anthropic_passthrough_non_streaming( repetition_penalty = None, presence_penalty = None, tool_choice = "auto", + disable_parallel_tool_use = False, ): """Non-streaming client-side pass-through.""" target_url = f"{llama_backend.base_url}/v1/chat/completions" @@ -5448,6 +6144,9 @@ async def _anthropic_passthrough_non_streaming( content_blocks.append(AnthropicResponseTextBlock(text = text)) tool_calls = message.get("tool_calls") or [] + # disable_parallel_tool_use: keep only the first tool_use block. + if disable_parallel_tool_use and len(tool_calls) > 1: + tool_calls = tool_calls[:1] for tc in tool_calls: fn = tc.get("function") or {} try: @@ -5456,18 +6155,13 @@ async def _anthropic_passthrough_non_streaming( args = {} content_blocks.append( AnthropicResponseToolUseBlock( - id = tc.get("id", ""), + id = anthropic_tool_use_id(tc.get("id")), name = fn.get("name", ""), input = args, ) ) - if tool_calls: - stop_reason = "tool_use" - elif finish_reason == "length": - stop_reason = "max_tokens" - else: - stop_reason = "end_turn" + stop_reason = openai_finish_to_anthropic_stop(finish_reason, had_tool_calls = bool(tool_calls)) usage = data.get("usage") or {} resp_obj = AnthropicMessagesResponse( @@ -5734,7 +6428,8 @@ def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: payload.temperature, payload.top_p, payload.top_k, - payload.max_tokens, + # Honor max_completion_tokens on the tools/response_format passthrough too. + _effective_max_tokens(payload), payload.stream, stop = payload.stop, min_p = payload.min_p, @@ -5744,6 +6439,8 @@ def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: response_format = _extract_response_format(payload), chat_template_kwargs = tpl_kwargs, backend_ctx = backend_ctx, + seed = payload.seed, + stream_options = payload.stream_options, ) @@ -5755,8 +6452,8 @@ async def _openai_passthrough_stream( Forwards the client's OpenAI function-calling request to llama-server and relays the SSE stream back verbatim, preserving llama-server's native response ``id``, ``finish_reason`` (including ``"tool_calls"``), - ``delta.tool_calls``, and the trailing ``usage`` chunk so the client sees a - standard OpenAI response. + ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so + the client sees a standard OpenAI response. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" body = _build_openai_passthrough_body(payload, backend_ctx = llama_backend.context_length) @@ -5814,10 +6511,7 @@ async def _openai_passthrough_stream( await client.aclose() except Exception: pass - raise HTTPException( - status_code = upstream_status, - detail = f"llama-server error: {err_text[:500]}", - ) + raise _openai_passthrough_error(upstream_status, err_text) async def _stream(): # Same httpx lifecycle pattern as _anthropic_passthrough_stream: @@ -5843,6 +6537,12 @@ async def _openai_passthrough_stream( continue if not raw_line.startswith("data: "): continue + # Honor parallel_tool_calls=false (best-effort): drop tool_call + # deltas with index>=1 so only the first call streams. Only + # lines carrying tool_calls are reparsed; everything else is + # relayed byte-for-byte. + if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: + raw_line = _cap_parallel_tool_calls_sse_line(raw_line) # Relay verbatim to preserve llama-server's native id, # finish_reason, delta.tool_calls, and usage chunks. yield raw_line + "\n\n" @@ -5856,12 +6556,7 @@ async def _openai_passthrough_stream( except Exception as e: # 200 headers already flushed; errors must go in the SSE body. logger.error("openai passthrough stream error: %s", e) - err = { - "error": { - "message": _friendly_error(e), - "type": "server_error", - }, - } + err = _openai_stream_error_chunk(e) yield f"data: {json.dumps(err)}\n\n" finally: cancel_watcher.cancel() @@ -5922,46 +6617,62 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): ) if resp.status_code != 200: - raise HTTPException( - status_code = resp.status_code, - detail = f"llama-server error: {resp.text[:500]}", + raise _openai_passthrough_error(resp.status_code, resp.text) + + # The guided-decoding fence wraps each choice's JSON content in a + # ```json ... ``` markdown fence that data_designer's structured parser + # requires but which CORRUPTS output for standard OpenAI clients doing + # ``json.loads(content)``. It is therefore opt-in: only the internal + # data-recipe path sets ``_unsloth_guided_fence``; public response_format + # clients get the raw upstream JSON verbatim. + _guided_fence = bool((payload.model_extra or {}).get("_unsloth_guided_fence")) + _do_fence = _guided_fence and _extract_response_format(payload) is not None + _cap_parallel = payload.parallel_tool_calls is False + + try: + data = resp.json() + except Exception as exc: + # Non-JSON / unparseable upstream body: relay verbatim as before. + logger.warning( + "openai passthrough non-streaming: response not JSON, relaying raw: %s", + exc, ) + return Response(content = resp.content, media_type = "application/json") - # Guided-decoding fence wrap. llama-server returns raw JSON matching the - # schema (no markdown) since the GBNF grammar emits only the JSON object. - # data_designer's llm-structured parser looks for a ```json ... ``` fence - # and discards unfenced output, collapsing a 100%-valid guided-decoding run - # to 0/N. Wrap each choice's content in the expected fence when the caller - # asked for guided decoding, leaving already-fenced content alone. - if _extract_response_format(payload) is not None: - try: - data = resp.json() - changed = False - for choice in data.get("choices", []): - if not isinstance(choice, dict): - continue - msg = choice.get("message") - if not isinstance(msg, dict): - continue - content = msg.get("content") - if not isinstance(content, str): - continue - stripped = content.strip() - if not stripped or stripped.startswith("```"): - continue - msg["content"] = f"```json\n{stripped}\n```" + changed = False + for choice in data.get("choices", []): + if not isinstance(choice, dict): + continue + msg = choice.get("message") + if not isinstance(msg, dict): + continue + + # OpenAI requires content=null on a pure tool-call turn; llama-server + # emits content="". + if msg.get("tool_calls") and msg.get("content") == "": + msg["content"] = None + changed = True + + # Honor parallel_tool_calls=false (best-effort) by capping to one call. + if _cap_parallel: + _tcs = msg.get("tool_calls") + if isinstance(_tcs, list) and len(_tcs) > 1: + msg["tool_calls"] = _tcs[:1] changed = True - if changed: - return JSONResponse(content = data) - except Exception as exc: - # Wrap is best-effort; fall through to the verbatim body if the - # response isn't JSON-shaped or the structure is unusual. - logger.warning( - "response_format fence wrap skipped: %s", - exc, - ) - # Pass the upstream body through as raw bytes -- skips a redundant - # parse+re-serialize round-trip and keeps the response truly verbatim - # (matches the docstring). Status is guaranteed 200 by the check above. - return Response(content = resp.content, media_type = "application/json") + # Guided-decoding fence wrap (opt-in via _unsloth_guided_fence). + if _do_fence: + content = msg.get("content") + if not isinstance(content, str): + continue + stripped = content.strip() + if not stripped or stripped.startswith("```"): + continue + msg["content"] = f"```json\n{stripped}\n```" + changed = True + + # Nothing mutated: relay the upstream bytes verbatim, skipping a redundant + # parse + re-serialize round-trip. + if not changed: + return Response(content = resp.content, media_type = "application/json") + return JSONResponse(content = data) diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 207c4b3805..e1230ae113 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -6,7 +6,9 @@ import sys import os import json +import threading +import httpx import pytest _backend = os.path.join(os.path.dirname(__file__), "..") @@ -36,6 +38,7 @@ from routes.inference import ( _normalize_anthropic_openai_images, _select_anthropic_server_tools, _anthropic_requested_studio_tools, + _anthropic_passthrough_stream, _anthropic_tool_non_streaming, anthropic_messages, ) @@ -671,7 +674,7 @@ class TestAnthropicStreamEmitter: and payload["content_block"]["type"] == "tool_use" ] assert len(tool_starts) == 1 - assert tool_starts[0]["content_block"]["id"] == "call_0" + assert tool_starts[0]["content_block"]["id"].startswith("toolu_") assert second_payloads == [ { "type": "content_block_delta", @@ -686,7 +689,7 @@ class TestAnthropicStreamEmitter: def test_tool_end_closes_tool_opens_new_text_block(self): e = AnthropicStreamEmitter() e.start("msg_1", "m") - e.feed( + start_events = e.feed( { "type": "tool_start", "tool_name": "t", @@ -694,6 +697,13 @@ class TestAnthropicStreamEmitter: "arguments": {}, } ) + start_payload = next( + json.loads(event.split("data: ")[1]) + for event in start_events + if "content_block_start" in event + ) + tool_use_id = start_payload["content_block"]["id"] + assert tool_use_id.startswith("toolu_") events = e.feed( { "type": "tool_end", @@ -708,7 +718,7 @@ class TestAnthropicStreamEmitter: assert "tool_result" in events[1] parsed = json.loads(events[1].split("data: ")[1]) assert parsed["content"] == "done" - assert parsed["tool_use_id"] == "tc_1" + assert parsed["tool_use_id"] == tool_use_id assert "content_block_start" in events[2] assert '"type": "text"' in events[2] @@ -837,14 +847,11 @@ class TestAnthropicToolNonStreaming: body = json.loads(response.body) tool_blocks = [block for block in body["content"] if block["type"] == "tool_use"] - assert tool_blocks == [ - { - "type": "tool_use", - "id": "call_0", - "name": "render_html", - "input": {"code": ""}, - } - ] + assert len(tool_blocks) == 1 + assert tool_blocks[0]["type"] == "tool_use" + assert tool_blocks[0]["id"].startswith("toolu_") + assert tool_blocks[0]["name"] == "render_html" + assert tool_blocks[0]["input"] == {"code": ""} # ===================================================================== @@ -912,7 +919,7 @@ class TestAnthropicPassthroughEmitter: parsed = self._parse(events[0]) assert parsed["type"] == "content_block_start" assert parsed["content_block"]["type"] == "tool_use" - assert parsed["content_block"]["id"] == "call_1" + assert parsed["content_block"]["id"].startswith("toolu_") assert parsed["content_block"]["name"] == "Bash" def test_tool_call_arguments_streamed_as_input_json_delta(self): @@ -1117,7 +1124,102 @@ class TestAnthropicPassthroughEmitter: assert "content_block_start" in events[1] parsed = self._parse(events[1]) assert parsed["content_block"]["name"] == "Read" - assert parsed["content_block"]["id"] == "c2" + assert parsed["content_block"]["id"].startswith("toolu_") + + +class TestAnthropicPassthroughStreamAdapter: + class _Request: + async def is_disconnected(self): + return False + + @staticmethod + async def _collect(response): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk.decode() if isinstance(chunk, bytes) else chunk) + return chunks + + @staticmethod + def _payloads(lines, event_name): + prefix = f"event: {event_name}\n" + return [ + json.loads(line.split("data: ", 1)[1].strip()) + for line in lines + if line.startswith(prefix) + ] + + def test_stream_requests_usage_for_final_message_delta(self, monkeypatch): + import routes.inference as inf_mod + + captured = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode()) + chunks = [ + {"choices": [{"delta": {"content": "hi"}}]}, + { + "choices": [], + "usage": { + "prompt_tokens": 2, + "completion_tokens": 4, + "total_tokens": 6, + }, + }, + ] + content = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + content += "data: [DONE]\n\n" + return httpx.Response( + 200, + content = content.encode(), + headers = {"content-type": "text/event-stream"}, + ) + + transport = httpx.MockTransport(handler) + real_async_client = httpx.AsyncClient + + def _client(*args, **kwargs): + return real_async_client( + transport = transport, + timeout = kwargs.get("timeout", 600), + ) + + monkeypatch.setattr(inf_mod.httpx, "AsyncClient", _client) + backend = SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + count_chat_tokens = lambda *args, **kwargs: 2, + ) + + async def run(): + response = await _anthropic_passthrough_stream( + self._Request(), + threading.Event(), + backend, + [{"role": "user", "content": "hi"}], + [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + 0.7, + 0.95, + 20, + 16, + "msg_1", + "test-model", + ) + return await self._collect(response) + + lines = asyncio.run(run()) + + assert captured["body"]["stream_options"] == {"include_usage": True} + message_delta = self._payloads(lines, "message_delta")[0] + assert message_delta["usage"]["input_tokens"] == 2 + assert message_delta["usage"]["output_tokens"] == 4 # ===================================================================== @@ -1271,27 +1373,23 @@ class TestAnthropicRequestedStudioTools: # ===================================================================== -class _PlainPathCalled(Exception): - pass - - -class _ToolPathCalled(Exception): - pass - - def _mock_backend(monkeypatch, **overrides): """Install a minimal stub backend on routes.inference. - Generation methods raise sentinels so the caller can assert which path - the route entered. + Generation methods record which path the route entered, then yield one + content event so the route can complete normally. """ import routes.inference as inf_mod + calls = [] + def _gen_plain(**kwargs): - raise _PlainPathCalled() + calls.append(("plain", kwargs)) + yield {"type": "content", "text": "ok"} def _gen_tools(**kwargs): - raise _ToolPathCalled() + calls.append(("tools", kwargs)) + yield {"type": "content", "text": "ok"} backend = SimpleNamespace( is_loaded = True, @@ -1300,6 +1398,7 @@ def _mock_backend(monkeypatch, **overrides): model_identifier = "test-model", generate_chat_completion = _gen_plain, generate_chat_completion_with_tools = _gen_tools, + calls = calls, ) backend.__dict__.update(overrides) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) @@ -1414,42 +1513,42 @@ class TestAnthropicMessagesToolRouting: assert "input_schema" in exc.value.detail def test_unrecognized_server_tool_accepted_as_noop(self, monkeypatch): - _mock_backend(monkeypatch) + backend = _mock_backend(monkeypatch) payload = _basic_payload( tools = [{"type": "code_execution_20250825", "name": "code_execution"}], ) - with pytest.raises(_PlainPathCalled): - _drive(anthropic_messages(payload, request = None, current_subject = "t")) + _drive(anthropic_messages(payload, request = None, current_subject = "t")) + assert backend.calls[0][0] == "plain" def test_disable_tools_policy_overrides_server_tool_alias(self, monkeypatch): # CLI `unsloth run --disable-tools` sets policy=False. A request with # a Studio server-tool alias must NOT enter the agentic loop then. - _mock_backend(monkeypatch) + backend = _mock_backend(monkeypatch) set_tool_policy(False) payload = _basic_payload( tools = [{"type": "web_search_20250305", "name": "web_search"}], ) - with pytest.raises(_PlainPathCalled): - _drive(anthropic_messages(payload, request = None, current_subject = "t")) + _drive(anthropic_messages(payload, request = None, current_subject = "t")) + assert backend.calls[0][0] == "plain" def test_server_tool_alias_enters_tool_path_when_policy_unset(self, monkeypatch): # Mirror of the previous test for the default (None) policy. - _mock_backend(monkeypatch) + backend = _mock_backend(monkeypatch) payload = _basic_payload( tools = [{"type": "web_search_20250305", "name": "web_search"}], ) - with pytest.raises(_ToolPathCalled): - _drive(anthropic_messages(payload, request = None, current_subject = "t")) + _drive(anthropic_messages(payload, request = None, current_subject = "t")) + assert backend.calls[0][0] == "tools" def test_per_request_enable_tools_false_blocks_server_tool_alias(self, monkeypatch): - _mock_backend(monkeypatch) + backend = _mock_backend(monkeypatch) payload = _basic_payload( enable_tools = False, tools = [{"type": "web_search_20250305", "name": "web_search"}], ) - with pytest.raises(_PlainPathCalled): - _drive(anthropic_messages(payload, request = None, current_subject = "t")) + _drive(anthropic_messages(payload, request = None, current_subject = "t")) + assert backend.calls[0][0] == "plain" diff --git a/studio/backend/tests/test_gguf_completion_usage.py b/studio/backend/tests/test_gguf_completion_usage.py index 9f9de80a84..d1e05f3e0e 100644 --- a/studio/backend/tests/test_gguf_completion_usage.py +++ b/studio/backend/tests/test_gguf_completion_usage.py @@ -53,10 +53,16 @@ def test_non_streaming_gguf_completion_includes_generated_usage(monkeypatch): ) assert response.status_code == 200 - assert response.json()["usage"] == { - "prompt_tokens": 23, - "completion_tokens": 1283, - "total_tokens": 1306, + usage = response.json()["usage"] + assert usage["prompt_tokens"] == 23 + assert usage["completion_tokens"] == 1283 + assert usage["total_tokens"] == 1306 + assert usage["prompt_tokens_details"] == {"cached_tokens": 0, "audio_tokens": 0} + assert usage["completion_tokens_details"] == { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, } @@ -67,8 +73,14 @@ def test_non_streaming_gguf_completion_defaults_nullable_usage_to_zero(monkeypat ) assert response.status_code == 200 - assert response.json()["usage"] == { - "prompt_tokens": 0, - "completion_tokens": 1283, - "total_tokens": 0, + usage = response.json()["usage"] + assert usage["prompt_tokens"] == 0 + assert usage["completion_tokens"] == 1283 + assert usage["total_tokens"] == 1283 + assert usage["prompt_tokens_details"] == {"cached_tokens": 0, "audio_tokens": 0} + assert usage["completion_tokens_details"] == { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, } diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 4437bf2212..495ebc434d 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -12,6 +12,7 @@ mapping. No server or GPU required. import os import sys import asyncio +import json from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") @@ -25,13 +26,19 @@ from pydantic import ValidationError from models.inference import ( ChatCompletionRequest, ChatMessage, + CompletionChoice, + CompletionMessage, ) from core.inference.anthropic_compat import ( anthropic_tool_choice_to_openai, ) from routes.inference import ( _build_passthrough_payload, + _clamp_finish_reason, + _effective_max_tokens, + _extract_content_parts, _friendly_error, + _openai_stream_usage_chunk, _set_or_prepend_system_message, openai_chat_completions, ) @@ -271,17 +278,19 @@ class TestChatCompletionRequestToolFields: assert req.stop is None def test_extra_fields_accepted(self): - # `frequency_penalty`, `seed`, `response_format` aren't explicitly - # declared but must survive Pydantic parsing now that extra="allow". + # `frequency_penalty` and `response_format` are not yet explicitly + # declared but must survive Pydantic parsing now that extra="allow" is + # set. `seed` is declared and should land on the typed field instead. req = self._make( frequency_penalty = 0.5, seed = 42, response_format = {"type": "json_object"}, ) + assert req.seed == 42 # Extras land in model_extra assert req.model_extra is not None assert req.model_extra.get("frequency_penalty") == 0.5 - assert req.model_extra.get("seed") == 42 + assert "seed" not in req.model_extra assert req.model_extra.get("response_format") == {"type": "json_object"} def test_unsloth_extensions_still_work(self): @@ -337,6 +346,153 @@ class TestChatCompletionRequestToolFields: assert "text/event-stream" not in resp.headers["content-type"] assert captured["stream"] is False + def _v1_client( + self, + monkeypatch, + llama_backend, + inference_backend = None, + ): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + import routes.inference as inference_route + from auth.authentication import get_current_subject + from utils.api_errors import install_api_error_handlers + + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama_backend) + if inference_backend is not None: + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: inference_backend) + + app = FastAPI() + app.include_router(inference_route.router, prefix = "/v1") + install_api_error_handlers(app) + app.dependency_overrides[get_current_subject] = lambda: "test-user" + return TestClient(app) + + def _assert_unsupported_param(self, response, param): + assert response.status_code == 400 + body = response.json() + assert body["error"]["param"] == param + assert body["error"]["code"] == "unsupported_parameter" + + def _assert_unsupported_n(self, response): + self._assert_unsupported_param(response, "n") + + def test_n_allows_openai_chat_completion_range(self): + req = self._make(n = 128) + assert req.n == 128 + with pytest.raises(ValidationError): + self._make(n = 129) + + def test_n_rejected_for_external_provider_path(self, monkeypatch): + class _UnusedBackend: + is_loaded = False + + client = self._v1_client(monkeypatch, _UnusedBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "provider_type": "openai", + "n": 2, + }, + ) + self._assert_unsupported_n(resp) + + def test_logprobs_rejected_until_supported(self, monkeypatch): + class _UnusedBackend: + is_loaded = False + + client = self._v1_client(monkeypatch, _UnusedBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "provider_type": "openai", + "logprobs": True, + }, + ) + self._assert_unsupported_param(resp, "logprobs") + + def test_top_logprobs_rejected_until_supported(self, monkeypatch): + class _UnusedBackend: + is_loaded = False + + client = self._v1_client(monkeypatch, _UnusedBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "provider_type": "openai", + "top_logprobs": 3, + }, + ) + self._assert_unsupported_param(resp, "top_logprobs") + + def test_n_rejected_for_gguf_streaming_path(self, monkeypatch): + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "n": 2, + }, + ) + self._assert_unsupported_n(resp) + + def test_n_rejected_for_gguf_tools_passthrough_path(self, monkeypatch): + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = True + is_vision = False + _is_audio = False + + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + "n": 2, + }, + ) + self._assert_unsupported_n(resp) + + def test_n_rejected_for_non_gguf_path(self, monkeypatch): + class _NoGGUFBackend: + is_loaded = False + + class _InferenceBackend: + active_model_name = "test-model" + models = {"test-model": {}} + + client = self._v1_client(monkeypatch, _NoGGUFBackend(), _InferenceBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "n": 2, + }, + ) + self._assert_unsupported_n(resp) + def test_multiturn_tool_loop_messages(self): req = ChatCompletionRequest( messages = [ @@ -449,18 +605,115 @@ class TestBuildPassthroughPayloadToolChoice: body = _build_passthrough_payload(**self._args(), tool_choice = tc) assert body["tool_choice"] == tc - def test_stream_adds_include_usage(self): + def test_stream_omits_usage_options_when_client_did_not_request_them(self): args = self._args() args["stream"] = True body = _build_passthrough_payload(**args) + assert "stream_options" not in body + + def test_stream_forwards_include_usage_when_client_requests_it(self): + args = self._args() + args["stream"] = True + body = _build_passthrough_payload( + **args, + stream_options = {"include_usage": True}, + ) assert body.get("stream_options") == {"include_usage": True} + def test_stream_forwards_include_usage_false_when_client_requests_it(self): + args = self._args() + args["stream"] = True + body = _build_passthrough_payload( + **args, + stream_options = {"include_usage": False}, + ) + assert body.get("stream_options") == {"include_usage": False} + def test_repetition_penalty_renamed(self): body = _build_passthrough_payload(**self._args(), repetition_penalty = 1.1) assert body.get("repeat_penalty") == 1.1 assert "repetition_penalty" not in body +# ===================================================================== +# OpenAI API compatibility helpers — verified spec edge cases +# ===================================================================== + + +class TestOpenAICompatibilityHelpers: + def test_max_completion_tokens_wins_over_deprecated_max_tokens(self): + payload = SimpleNamespace(max_tokens = 128, max_completion_tokens = 64) + assert _effective_max_tokens(payload) == 64 + + @pytest.mark.parametrize( + "finish_reason", + ["stop", "length", "tool_calls", "content_filter", "function_call"], + ) + def test_clamp_finish_reason_preserves_openai_finish_reasons(self, finish_reason): + assert _clamp_finish_reason(finish_reason) == finish_reason + + def test_clamp_finish_reason_defaults_unknown_to_stop(self): + assert _clamp_finish_reason(None) == "stop" + assert _clamp_finish_reason("unexpected") == "stop" + + def test_non_streaming_completion_choice_accepts_tool_calls_finish_reason(self): + choice = CompletionChoice( + index = 0, + message = CompletionMessage(content = ""), + finish_reason = "tool_calls", + ) + assert choice.finish_reason == "tool_calls" + + def test_stream_usage_chunk_requires_include_usage(self): + usage = {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5} + payload = SimpleNamespace(stream_options = None) + assert ( + _openai_stream_usage_chunk(payload, "chatcmpl-test", 123, "model", usage, None) is None + ) + + payload.stream_options = {"include_usage": True} + line = _openai_stream_usage_chunk(payload, "chatcmpl-test", 123, "model", usage, None) + assert line is not None + assert '"choices":[]' in line + assert '"usage"' in line + + def test_stream_usage_chunk_coerces_nullable_counts(self): + payload = SimpleNamespace(stream_options = {"include_usage": True}) + line = _openai_stream_usage_chunk( + payload, + "chatcmpl-test", + 123, + "model", + {"prompt_tokens": None, "completion_tokens": 7, "total_tokens": None}, + None, + ) + + assert line is not None + parsed = json.loads(line.removeprefix("data: ")) + usage = parsed["usage"] + assert usage["prompt_tokens"] == 0 + assert usage["completion_tokens"] == 7 + assert usage["total_tokens"] == 7 + + def test_developer_message_preserves_existing_system_prompt(self): + payload = ChatCompletionRequest( + messages = [ + {"role": "system", "content": "original system"}, + {"role": "developer", "content": "developer rules"}, + {"role": "user", "content": "hi"}, + ] + ) + for message in payload.messages: + if message.role == "developer": + message.role = "system" + + system_prompt, chat_messages, image_b64 = _extract_content_parts(payload.messages) + + assert system_prompt == "original system\n\ndeveloper rules" + assert chat_messages == [{"role": "user", "content": "hi"}] + assert image_b64 is None + + # ===================================================================== # _friendly_error — httpx transport failures # ===================================================================== @@ -800,3 +1053,90 @@ class TestGgufVisionToolRouting: assert tool_messages[0]["role"] == "system" assert tool_messages[1]["role"] == "user" assert tool_messages[1]["content"][1]["type"] == "image_url" + + def test_parallel_tool_calls_false_reaches_gguf_tool_loop(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + captured = {} + + def _plain(**kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**kwargs): + captured["kwargs"] = kwargs + yield {"type": "content", "text": "done"} + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + model_identifier = "test-gguf", + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + enable_tools = True, + enabled_tools = ["web_search"], + parallel_tool_calls = False, + messages = [{"role": "user", "content": "search once"}], + ) + + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + self._consume_response(response) + + assert captured["kwargs"]["disable_parallel_tool_use"] is True + + @pytest.mark.parametrize( + ("seed", "expected"), + [ + (41, [41, 42, 43]), + (-1, [-1, -1, -1]), + ], + ) + def test_gguf_n_choices_vary_explicit_non_negative_seed(self, monkeypatch, seed, expected): + import routes.inference as inf_mod + + seen_seeds = [] + + def _generate(**kwargs): + seen_seeds.append(kwargs.get("seed")) + yield f"choice-{len(seen_seeds)}" + yield { + "type": "metadata", + "usage": { + "prompt_tokens": 5, + "completion_tokens": 7, + "total_tokens": 12, + }, + "finish_reason": "stop", + } + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + model_identifier = "test-gguf", + generate_chat_completion = _generate, + ) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + n = 3, + seed = seed, + ) + + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + body = json.loads(response.body) + + assert seen_seeds == expected + assert [choice["index"] for choice in body["choices"]] == [0, 1, 2] diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index c182e3d6d0..f0f2714214 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -26,12 +26,15 @@ No running server or GPU required. import os import sys +import asyncio +from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") sys.path.insert(0, _backend) import json +import httpx import pytest from pydantic import ValidationError @@ -52,8 +55,10 @@ from models.inference import ( ResponsesUsage, ) from routes.inference import ( + _build_chat_request, _chat_tool_calls_to_responses_output, _normalise_responses_input, + _responses_stream, _translate_responses_tool_choice_to_chat, _translate_responses_tools_to_chat, ) @@ -259,6 +264,26 @@ class TestToolChoiceTranslation: assert _translate_responses_tool_choice_to_chat(obj) == obj +class TestBuildChatRequest: + def test_parallel_tool_calls_false_is_preserved_for_passthrough_caps(self): + payload = ResponsesRequest( + input = "hi", + tools = [ + { + "type": "function", + "name": "lookup", + "parameters": {"type": "object"}, + } + ], + parallel_tool_calls = False, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = True) + + assert chat_req.parallel_tool_calls is False + + # ===================================================================== # _normalise_responses_input — multi-turn tool mapping # ===================================================================== @@ -424,6 +449,130 @@ class TestChatToolCallsToResponsesOutput: assert items[0]["arguments"] == "" +# ===================================================================== +# Streaming Responses adapter +# ===================================================================== + + +class TestResponsesStreamAdapter: + class _Request: + async def is_disconnected(self): + return False + + @staticmethod + async def _collect(response): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk.decode() if isinstance(chunk, bytes) else chunk) + return chunks + + @staticmethod + def _payloads(lines, event_name): + prefix = f"event: {event_name}\n" + return [ + json.loads(line.split("data: ", 1)[1].strip()) + for line in lines + if line.startswith(prefix) + ] + + def test_requests_usage_and_caps_parallel_tool_calls(self, monkeypatch): + import routes.inference as inf_mod + + captured = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode()) + chunks = [ + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "type": "function", + "function": {"name": "first", "arguments": "{}"}, + }, + { + "index": 1, + "id": "call_1", + "type": "function", + "function": {"name": "second", "arguments": "{}"}, + }, + ] + } + } + ] + }, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + content = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + content += "data: [DONE]\n\n" + return httpx.Response( + 200, + content = content.encode(), + headers = {"content-type": "text/event-stream"}, + ) + + transport = httpx.MockTransport(handler) + real_async_client = httpx.AsyncClient + + def _client(*args, **kwargs): + return real_async_client( + transport = transport, + timeout = kwargs.get("timeout", 600), + ) + + monkeypatch.setattr(inf_mod.httpx, "AsyncClient", _client) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + is_vision = False, + context_length = 4096, + base_url = "http://llama.test", + ), + ) + + payload = ResponsesRequest( + input = "hi", + stream = True, + parallel_tool_calls = False, + tools = [ + { + "type": "function", + "name": "first", + "parameters": {"type": "object"}, + }, + { + "type": "function", + "name": "second", + "parameters": {"type": "object"}, + }, + ], + ) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + assert captured["body"]["stream_options"] == {"include_usage": True} + joined = "".join(lines) + assert "call_0" in joined + assert "call_1" not in joined + completed = self._payloads(lines, "response.completed")[0] + assert completed["response"]["usage"] == { + "input_tokens": 2, + "output_tokens": 3, + "total_tokens": 5, + } + + # ===================================================================== # Response model — ResponsesOutputFunctionCall / mixed output # ===================================================================== diff --git a/studio/backend/utils/api_errors.py b/studio/backend/utils/api_errors.py new file mode 100644 index 0000000000..b1c55b61b9 --- /dev/null +++ b/studio/backend/utils/api_errors.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Error-envelope helpers for the OpenAI/Anthropic-compatible ``/v1/*`` API surface. + +FastAPI's defaults emit ``{"detail": ...}`` bodies (status 422 for validation, +``exc.status_code`` for ``HTTPException``). Real OpenAI/Anthropic clients expect +provider-specific error envelopes instead, so this module re-wraps Unsloth's own +client-error responses on the ``/v1/*`` surface: + +- OpenAI surface (``/v1/chat/completions``, ``/v1/completions``, ``/v1/models``, + ``/v1/responses``, ``/v1/embeddings``, ...):: + + {"error": {"message": str, "type": str, "param": None|str, "code": None|str}} + +- Anthropic surface (any path starting with ``/v1/messages``):: + + {"type": "error", "error": {"type": str, "message": str}} + +CRITICAL: the exception handlers installed by :func:`install_api_error_handlers` +are global, but they ONLY transform responses for paths that start with ``/v1/``. +For every other path (``/api/...``, frontend routes) they reproduce FastAPI's +default behavior byte-for-byte, because the Studio frontend depends on the +``{"detail": ...}`` shape for ``/api/*``. + +Public contract (other modules depend on these): + +- ``OPENAI_TYPE_BY_STATUS`` / ``ANTHROPIC_TYPE_BY_STATUS``: status -> type maps. +- ``openai_error_body(message, *, status=400, err_type=None, code=None, param=None)`` +- ``anthropic_error_body(message, *, status=400, err_type=None)`` +- ``is_anthropic_path(path)`` +- ``error_body_for_path(path, message, *, status, err_type=None, code=None, param=None)`` +- ``install_api_error_handlers(app)`` +""" + +from fastapi.encoders import jsonable_encoder +from fastapi.responses import JSONResponse, Response +from fastapi.exceptions import RequestValidationError +from fastapi.utils import is_body_allowed_for_status_code +from starlette.exceptions import HTTPException as StarletteHTTPException + + +# Status-code -> error ``type`` string for the OpenAI error envelope. +OPENAI_TYPE_BY_STATUS = { + 400: "invalid_request_error", + 401: "authentication_error", + 403: "permission_error", + 404: "not_found_error", + 409: "conflict_error", + 413: "invalid_request_error", + 422: "invalid_request_error", + 429: "rate_limit_error", + 500: "api_error", + 502: "api_error", + 503: "api_error", +} + +# Status-code -> error ``type`` string for the Anthropic error envelope. +ANTHROPIC_TYPE_BY_STATUS = { + 400: "invalid_request_error", + 401: "authentication_error", + 403: "permission_error", + 404: "not_found_error", + 409: "conflict_error", + 413: "request_too_large", + 422: "invalid_request_error", + 429: "rate_limit_error", + 500: "api_error", + 502: "api_error", + 503: "api_error", + 529: "overloaded_error", +} + + +def openai_error_body( + message, + *, + status = 400, + err_type = None, + code = None, + param = None, +) -> dict: + """Build an OpenAI-style error envelope. + + Returns ``{"error": {"message", "type", "param", "code"}}``. The ``param`` + and ``code`` keys are always present (value may be ``None``). ``err_type`` + defaults to :data:`OPENAI_TYPE_BY_STATUS` for ``status`` (``"api_error"`` + fallback). + """ + return { + "error": { + "message": str(message), + "type": err_type or OPENAI_TYPE_BY_STATUS.get(status, "api_error"), + "param": param, + "code": code, + } + } + + +def anthropic_error_body( + message, + *, + status = 400, + err_type = None, +) -> dict: + """Build an Anthropic-style error envelope. + + Returns ``{"type": "error", "request_id": None, "error": {"type", "message"}}``. + ``request_id`` is a required (nullable) field on the spec's ErrorResponse; + Studio has no request-id system, so it is null. ``err_type`` defaults to + :data:`ANTHROPIC_TYPE_BY_STATUS` for ``status`` (``"api_error"`` fallback). + """ + return { + "type": "error", + "request_id": None, + "error": { + "type": err_type or ANTHROPIC_TYPE_BY_STATUS.get(status, "api_error"), + "message": str(message), + }, + } + + +def is_anthropic_path(path: str) -> bool: + """True iff ``path`` belongs to the Anthropic surface (``/v1/messages*``).""" + return path.startswith("/v1/messages") + + +def error_body_for_path( + path, + message, + *, + status, + err_type = None, + code = None, + param = None, +) -> dict: + """Dispatch to the correct envelope builder based on ``path``. + + Anthropic surface paths use :func:`anthropic_error_body` (``code``/``param`` + are not part of that envelope and are ignored); all other ``/v1/*`` paths use + :func:`openai_error_body`. + """ + if is_anthropic_path(path): + return anthropic_error_body(message, status = status, err_type = err_type) + return openai_error_body(message, status = status, err_type = err_type, code = code, param = param) + + +def _summarize_validation_errors(errors) -> tuple: + """Derive a readable one-line message and (optional) body param from ``exc.errors()``. + + Returns ``(summary, param)``. ``summary`` is a human-readable string like + ``"messages: Field required"``. ``param`` is the offending body field name when + one can be extracted (used as the OpenAI envelope ``param``), else ``None``. + + Malformed-JSON bodies surface here as ``type == "json_invalid"`` and get a + dedicated message. + """ + if not errors: + return "Invalid request", None + + first = errors[0] + if first.get("type") == "json_invalid": + return "Invalid JSON in request body", None + + loc = first.get("loc", ()) or () + msg = first.get("msg", "Invalid request") + + # Extract the body field name (the loc element after a leading "body"). + param = None + loc_parts = [p for p in loc if p not in ("body",)] + if loc and loc[0] == "body" and loc_parts: + # First non-"body" element that is a field name (string). + for part in loc_parts: + if isinstance(part, str): + param = part + break + + label = ".".join(str(p) for p in loc_parts) if loc_parts else ".".join(str(p) for p in loc) + summary = f"{label}: {msg}" if label else str(msg) + return summary, param + + +def install_api_error_handlers(app) -> None: + """Register validation + HTTPException handlers that emit ``/v1/*`` envelopes. + + Both handlers are global but only transform responses for paths starting with + ``/v1/``. Non-``/v1/`` paths reproduce FastAPI's default ``{"detail": ...}`` + behavior exactly so the Studio frontend keeps working. + """ + + @app.exception_handler(RequestValidationError) + async def _handle_validation_error(request, exc): + path = request.url.path + if path.startswith("/v1/"): + summary, param = _summarize_validation_errors(exc.errors()) + return JSONResponse( + status_code = 400, + content = error_body_for_path(path, summary, status = 400, param = param), + ) + # Default FastAPI behavior for every other path. + return JSONResponse( + status_code = 422, + content = {"detail": jsonable_encoder(exc.errors())}, + ) + + @app.exception_handler(StarletteHTTPException) + async def _handle_http_exception(request, exc): + path = request.url.path + headers = getattr(exc, "headers", None) + # Statuses like 204/304/1xx must not carry a body — mirror FastAPI's + # default http_exception_handler, which returns a bodiless Response. + if not is_body_allowed_for_status_code(exc.status_code): + return Response(status_code = exc.status_code, headers = headers) + if path.startswith("/v1/"): + detail = exc.detail + # Already a fully-formed envelope: pass through untouched. + if isinstance(detail, dict) and ("error" in detail or detail.get("type") == "error"): + return JSONResponse( + status_code = exc.status_code, + content = detail, + headers = headers, + ) + # A dict carrying our individual fields. + if isinstance(detail, dict): + message = detail.get("message", detail) + err_type = detail.get("type") + code = detail.get("code") + param = detail.get("param") + else: + # Plain message string (the common HTTPException case). + message = detail + err_type = None + code = None + param = None + return JSONResponse( + status_code = exc.status_code, + content = error_body_for_path( + path, + message, + status = exc.status_code, + err_type = err_type, + code = code, + param = param, + ), + headers = headers, + ) + # Default FastAPI behavior for every other path. + return JSONResponse( + status_code = exc.status_code, + content = {"detail": exc.detail}, + headers = headers, + ) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 34cae518e3..91a9bffa36 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -2339,6 +2339,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { model: params.checkpoint, messages: outboundMessages, stream: true, + // Opt into the trailing usage chunk so the context-usage bar + // and tok/s readout populate (backend gates it on include_usage). + stream_options: { include_usage: true }, temperature: params.temperature, top_p: params.topP, max_tokens: params.maxTokens, diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 48efdbba2e..a69c216891 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -309,6 +309,13 @@ export interface OpenAIChatCompletionsRequest { * See https://platform.claude.com/docs/en/build-with-claude/fast-mode */ fast_mode?: boolean | null; + /** + * Opt into the OpenAI-standard trailing usage chunk on streams + * (`choices: []` with `usage` + llama-server `timings` populated). The + * backend only emits it when `include_usage` is set; the local chat UI + * sends it so the context-usage bar and tok/s readout populate. + */ + stream_options?: { include_usage?: boolean } | null; } export interface OpenAIChatDelta {