unsloth/studio/backend/routes/inference.py
Apoze 6a9b77ee37
Some checks are pending
Core / Core (HF=default + TRL=default) (push) Waiting to run
Core / Core (HF=4.57.6 + TRL<1) (push) Waiting to run
Core / Core (HF=latest + TRL=latest) (push) Waiting to run
Core / llama.cpp build + smoke (push) Waiting to run
Cross-platform parity / parity (macos-latest) (push) Waiting to run
Cross-platform parity / parity (windows-latest) (push) Waiting to run
Lint CI / Source lint (Python + shell + YAML + JSON + safety nets) (push) Waiting to run
MLX CI on Mac M1 / dispatch (push) Waiting to run
Scorecard supply-chain security / Scorecard analysis (push) Waiting to run
Security audit / pytest tests/security (push) Waiting to run
Security audit / npm provenance + new install-script diff (push) Waiting to run
Security audit / advisory audit (pip + npm + cargo) (push) Waiting to run
Security audit / pip scan-packages :: extras (push) Waiting to run
Security audit / pip scan-packages :: studio (push) Waiting to run
Security audit / pip scan-packages :: hf-stack (push) Waiting to run
Security audit / npm scan-packages (Studio frontend tarballs) (push) Waiting to run
Security audit / workflow-trigger lint (pull_request_target / cache-poisoning) (push) Waiting to run
Studio API CI / Studio API & Auth Tests (push) Waiting to run
Backend CI / (Python 3.10) (push) Waiting to run
Backend CI / (Python 3.11) (push) Waiting to run
Backend CI / (Python 3.12) (push) Waiting to run
Backend CI / (Python 3.13) (push) Waiting to run
Backend CI / Repo tests (CPU) (push) Waiting to run
Frontend CI / Frontend build + bundle sanity (push) Waiting to run
Studio GGUF CI / OpenAI, Anthropic API tests (push) Waiting to run
Studio GGUF CI / Tool calling Tests (push) Waiting to run
Studio GGUF CI / JSON, images (push) Waiting to run
Studio load-orchestrator CI / test (push) Waiting to run
Mac Studio API CI / Studio API & Auth Tests (push) Waiting to run
Mac Studio GGUF CI / OpenAI, Anthropic API tests (push) Waiting to run
Mac Studio GGUF CI / Tool calling Tests (push) Waiting to run
Mac Studio GGUF CI / JSON, images (push) Waiting to run
Mac Studio Install Matrix CI / Install + load (macos-14) (push) Waiting to run
Mac Studio Install Matrix CI / Install + load (macos-15) (push) Waiting to run
Mac Studio Install Matrix CI / Install + load (macos-26) (push) Waiting to run
Mac Studio Install Matrix CI / Install + load (macos-15-intel) (push) Waiting to run
Mac Studio Install Matrix CI / Install + load (macos-26-intel) (push) Waiting to run
Windows Studio GGUF CI / VC++ runtime detect + install round-trip (windows-latest) (push) Waiting to run
Windows Studio UI CI / Chat UI Tests (push) Waiting to run
Windows Studio Update CI / Studio Updating Tests (push) Waiting to run
Wheel CI / Wheel build + content sanity + import smoke (push) Waiting to run
Mac Studio UI CI / Chat UI Tests (push) Waiting to run
Mac Studio Update CI / Studio Updating Tests (push) Waiting to run
Studio Tauri CI / Tauri Linux debug build (no codesign) (push) Waiting to run
Studio UI CI / Chat UI Tests (push) Waiting to run
Studio Update CI / Studio Updating Tests (push) Waiting to run
Windows Studio API CI / Studio API & Auth Tests (push) Waiting to run
Windows Studio GGUF CI / OpenAI, Anthropic API tests (push) Waiting to run
Windows Studio GGUF CI / Tool calling Tests (push) Waiting to run
Windows Studio GGUF CI / JSON, images (push) Waiting to run
Windows Studio GGUF CI / Studio install + inference without Visual Studio (push) Waiting to run
Windows Studio GGUF CI / GPU prebuilt resolves without Visual Studio (push) Waiting to run
Windows Studio GGUF CI / setup.ps1 unit tests (VS 2026 / CMake guard) (push) Waiting to run
Windows Studio GGUF CI / real-VS detection (VS 2022) (push) Waiting to run
Windows Studio GGUF CI / real-VS detection (VS 2026) (push) Waiting to run
Windows Studio GGUF CI / VC++ runtime detect + install round-trip (windows-2025-vs2026) (push) Waiting to run
Studio: harden OpenAI-compatible GGUF streaming (#6950)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-07-09 12:09:08 -03:00

12958 lines
553 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Inference API routes for model loading and text generation.
"""
import os
import sys
import time
import uuid
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Request, status
from fastapi.responses import StreamingResponse, JSONResponse, Response
from starlette.requests import ClientDisconnect
from typing import Any, Callable, List, Optional, Union
import json
import httpx
from loggers import get_logger
import asyncio
import threading
import weakref
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 _positive_int_or_none(value: Any) -> Optional[int]:
if isinstance(value, bool):
return None
try:
value_int = int(value)
except (TypeError, ValueError):
return None
return value_int if value_int > 0 else None
def _nonnegative_int_or_none(value: Any) -> Optional[int]:
if isinstance(value, bool):
return None
try:
value_int = int(value)
except (TypeError, ValueError):
return None
return value_int if value_int >= 0 else None
_MLX_MPI_DISTRIBUTED_ENV_PAIRS = (
("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"),
("PMI_RANK", "PMI_SIZE"),
("PMIX_RANK", "PMIX_SIZE"),
("MPI_RANK", "MPI_WORLD_SIZE"),
("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"),
)
def _mlx_distributed_launch_detected() -> bool:
if _nonnegative_int_or_none(os.environ.get("MLX_RANK")) is not None:
world_size = _positive_int_or_none(os.environ.get("MLX_WORLD_SIZE"))
if world_size is not None and world_size > 1:
return True
return bool(
os.environ.get("MLX_HOSTFILE")
or os.environ.get("MLX_IBV_DEVICES")
or os.environ.get("MLX_JACCL_COORDINATOR")
or (os.environ.get("NCCL_HOST_IP") and os.environ.get("NCCL_PORT"))
)
return any(
_nonnegative_int_or_none(os.environ.get(rank_env)) is not None
and (_positive_int_or_none(os.environ.get(size_env)) or 0) > 1
for rank_env, size_env in _MLX_MPI_DISTRIBUTED_ENV_PAIRS
)
def _install_httpcore_asyncgen_silencer() -> None:
"""Silence benign httpx/httpcore asyncgen GC noise on Python 3.13.
When Studio proxies a llama-server stream via httpx, the innermost
``HTTP11ConnectionByteStream.__aiter__`` async generator is finalised by
the asyncgen GC hook on a task different from the one that opened it. Its
``aclose`` calls ``anyio.Lock.acquire`` → ``cancel_shielded_checkpoint``,
entering a ``CancelScope`` on the finaliser task; Python 3.13 flags the
cross-task exit as ``"Attempted to exit cancel scope in a different task"``
and prints ``"async generator ignored GeneratorExit"`` as an unraisable
warning.
Known httpx + httpcore + anyio interaction (MCP SDK python-sdk#831, agno
#3556, chainlit #2361, langchain-mcp-adapters #254). Benign: the 200
response is already delivered. The streaming pass-throughs
(``/v1/chat/completions``, ``/v1/messages``, ``/v1/responses``,
``/v1/completions``) manage their httpx lifecycle in one task with explicit
``aclose()``; we don't hold a reference to the errant generator and can't
close it ourselves.
Install one process-wide unraisable hook that swallows only this
interaction -- identified by (RuntimeError mentioning cancel scope /
GeneratorExit) + (object repr referencing HTTP11ConnectionByteStream) --
and defers to the default hook otherwise. Idempotent.
"""
prior_hook = sys.unraisablehook
if getattr(prior_hook, "_unsloth_httpcore_silencer", False):
return
def _hook(unraisable):
exc_value = getattr(unraisable, "exc_value", None)
obj = getattr(unraisable, "object", None)
obj_repr = repr(obj) if obj is not None else ""
if (
isinstance(exc_value, RuntimeError)
and "HTTP11ConnectionByteStream" in obj_repr
and (
"cancel scope" in str(exc_value)
or "GeneratorExit" in str(exc_value)
or "no running event loop" in str(exc_value)
)
):
return
prior_hook(unraisable)
_hook._unsloth_httpcore_silencer = True # type: ignore[attr-defined]
sys.unraisablehook = _hook
_install_httpcore_asyncgen_silencer()
def _loaded_chat_template() -> Optional[str]:
"""Chat template of the currently loaded GGUF model, if any."""
try:
return get_llama_cpp_backend().chat_template
except Exception:
return None
def _template_raise_message(error_text: str, chat_template: Optional[str]) -> Optional[str]:
"""A chat-template raise_exception message to surface, but only when it appears
verbatim in chat_template (simple substring check), so we never leak arbitrary
llama-server text. Anchors on llama.cpp's "Jinja Exception:" prefix."""
if not chat_template:
return None
marker = "Jinja Exception:"
idx = error_text.find(marker)
if idx == -1:
return None
candidate = error_text[idx + len(marker) :]
# llama-server appends JSON after the message; cut at the first boundary.
for stop in ('"', "\n"):
cut = candidate.find(stop)
if cut != -1:
candidate = candidate[:cut]
candidate = candidate.strip()
return candidate if candidate and candidate in chat_template else None
_LOST_CONNECTION_MSG = (
"Lost connection to the model server. It may have crashed -- try reloading the model."
)
def _friendly_error(exc: Exception) -> str:
"""Extract a user-friendly message from known llama-server errors."""
if isinstance(exc, httpx.ReadTimeout):
if "stopped producing tokens" in str(exc).lower():
return (
"The model stopped producing tokens before the response "
"completed. Try stopping and retrying, or reduce max tokens."
)
return (
"The model is still processing the prompt but did not produce a "
"first token within 20 minutes. Try reducing context length, "
"using more GPU offload, or loading a smaller model."
)
if isinstance(exc, httpx.TimeoutException):
return "Timed out communicating with the model server. Try again shortly."
# httpx transport failures from the async pass-through helpers. Any
# RequestError subclass (ConnectError, ReadError, RemoteProtocolError,
# WriteError, PoolTimeout, ...) means the llama-server subprocess is
# unreachable -- crashed or still coming up.
if isinstance(exc, httpx.RequestError):
return _LOST_CONNECTION_MSG
msg = str(exc)
m = _re.search(
r"request \((\d+) tokens?\) exceeds the available context size \((\d+) tokens?\)",
msg,
)
if m:
return (
f"Message too long: {m.group(1)} tokens exceeds the {m.group(2)}-token "
f"context window. Try increasing the Context Length in Model settings, "
f"or shorten the conversation."
)
if "Lost connection to llama-server" in msg:
return _LOST_CONNECTION_MSG
template_msg = _template_raise_message(msg, _loaded_chat_template())
if template_msg:
return f"An internal error occurred: {template_msg}"
return "An internal error occurred"
def _friendly_upstream_error(text: str) -> str:
"""Rewrite a raw llama-server error body into an actionable message where we can.
The main case is a tool-calling grammar that llama-server can't compile ("failed to
parse grammar" / "failed to initialize samplers"). This surfaces to coding agents as
a hard 400 on every tool-bearing turn. It is a llama-server limitation with some
model/quant + tool-schema combinations, and recent llama.cpp builds handle the common
coding-agent tools, so point the user at updating Studio rather than the raw body.
"""
lowered = text.lower()
if "failed to parse grammar" in lowered or "failed to initialize samplers" in lowered:
return (
"The model couldn't compile a tool-calling grammar for this request. This is a "
"llama-server limitation with some model/quant and tool-schema combinations. "
"Update Studio (it installs the latest llama.cpp, which handles the common "
"coding-agent tools) or try a different GGUF model."
)
return f"llama-server error: {text}"
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
)
_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT"
def _positive_float_env(env_name: str, default):
"""Parse a positive float from an env var. A parseable non-positive value
returns ``None`` (0 disables the guarded feature); only unparseable or unset
values fall back to ``default``."""
raw_value = os.environ.get(env_name)
if raw_value is None or not raw_value.strip():
return default
try:
value = float(raw_value.strip())
except ValueError:
return default
return value if value > 0 else None
def _effective_openai_max_tokens_from_values(max_tokens, max_completion_tokens = None):
"""Resolve the OpenAI-compatible generation cap from raw request values.
Prefers ``max_completion_tokens`` over the deprecated ``max_tokens``, and
returns ``None`` when both are omitted so callers keep their context-window
default (OpenAI treats an omitted cap as bounded only by the context
window). Explicit client caps pass through unchanged.
"""
def _validate_explicit(value, param: str):
if value is None:
return None
if isinstance(value, bool) or not isinstance(value, int):
raise HTTPException(
status_code = 400,
detail = openai_error_body(
f"'{param}' must be an integer.",
status = 400,
code = "invalid_type",
param = param,
),
)
# The legacy completions spec declares ``minimum: 0`` for max_tokens,
# so 0 is a valid (if degenerate) cap and only negatives are rejected.
# The chat fields never reach here with 0 (pydantic enforces ge=1).
if value < 0:
raise HTTPException(
status_code = 400,
detail = openai_error_body(
f"'{param}' must be at least 0.",
status = 400,
code = "invalid_value",
param = param,
),
)
return value
max_tokens = _validate_explicit(max_tokens, "max_tokens")
max_completion_tokens = _validate_explicit(max_completion_tokens, "max_completion_tokens")
return max_completion_tokens if max_completion_tokens is not None else max_tokens
def _effective_openai_max_tokens(payload):
return _effective_openai_max_tokens_from_values(
getattr(payload, "max_tokens", None),
getattr(payload, "max_completion_tokens", None),
)
def _wants_multiple_choices(payload) -> bool:
return (payload.n or 1) > 1
def _has_openai_tool_history(messages) -> bool:
for message in messages or []:
if isinstance(message, dict):
if message.get("role") == "tool" or message.get("tool_calls"):
return True
continue
if getattr(message, "role", None) == "tool" or getattr(message, "tool_calls", None):
return True
return False
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 _sse_streaming_response(content) -> StreamingResponse:
"""A ``text/event-stream`` response with the standard SSE headers used by
every streaming path here: no client/proxy caching, no proxy buffering, and
a one-shot connection. Two callers build their response inline instead: the
external-provider proxy omits ``Connection: close``, and the OpenAI
passthrough returns an empty ``keep-alive`` stream when the request is
cancelled before the upstream response starts.
Built on ``_SameTaskStreamingResponse`` (not Starlette's stock
``StreamingResponse``) so the SSE generator runs in the request task. The
legacy AnyIO task-group wrapper trips "Attempted to exit a cancel scope in a
different task" on Python 3.13 + httpx, which surfaced as a mid-stream
``response.failed``. The streaming paths that take their response inline use
``_SameTaskStreamingResponse`` directly for the same reason."""
return _SameTaskStreamingResponse(
content,
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
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_stream_error_sse(error: dict) -> str:
return f"data: {json.dumps(error)}\n\ndata: [DONE]\n\n"
def _openai_stream_error_sse_bytes(error: dict) -> bytes:
return _openai_stream_error_sse(error).encode("utf-8")
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; a tool-grammar compile failure gets the same actionable
guidance as the Anthropic passthrough; any other upstream error stays 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 = _friendly_upstream_error(text[:500]),
)
_OVERFLOW_TRUNCATE_MAX_RETRIES = 3
# Truncated-prompt share of the real window; the rest is generation headroom
# so a near-full prompt cannot cut a tool call mid-JSON at the wall.
_OVERFLOW_PROMPT_TARGET_FRACTION = 0.75
def _overflow_truncation_requested(payload) -> bool:
"""True when the request (or the UNSLOTH_CONTEXT_OVERFLOW server default,
for clients that cannot send custom fields) opted into truncation."""
requested = getattr(payload, "context_overflow", None)
if requested is not None:
return requested == "truncate_middle"
return os.environ.get("UNSLOTH_CONTEXT_OVERFLOW", "").strip().lower() == "truncate_middle"
def _parse_overflow_counts(err_text: str):
"""(n_prompt_tokens, n_ctx) from an exceed_context_size_error body, or
None. Tolerates \\" around keys (body may be a re-wrapped JSON string)."""
m_prompt = _re.search(r'n_prompt_tokens\\?"?\s*:\s*(\d+)', err_text)
m_ctx = _re.search(r'n_ctx\\?"?\s*:\s*(\d+)', err_text)
if m_prompt and m_ctx:
return int(m_prompt.group(1)), int(m_ctx.group(1))
return None
def _estimate_message_tokens(msg: dict) -> int:
try:
return max(1, len(json.dumps(msg, ensure_ascii = False)) // 4)
except Exception:
return 1
def _truncate_middle_messages(messages: list, keep_ratio: float):
"""Drop whole turn-groups from the middle of an OpenAI message list.
Always kept: leading system message(s), the first group (task anchor),
and the trailing groups. A group is a user message, or an assistant
message plus its following tool results, so surviving tool_calls stay
paired with their results as chat templates require.
Returns (new_messages, dropped_message_count).
"""
if not messages or keep_ratio >= 1.0:
return messages, 0
head: list = []
idx = 0
while idx < len(messages) and messages[idx].get("role") in ("system", "developer"):
head.append(messages[idx])
idx += 1
groups: list[list] = []
for msg in messages[idx:]:
role = msg.get("role")
if role == "tool" and groups:
groups[-1].append(msg)
elif role == "tool":
groups.append([msg]) # orphan tool result; treat as its own group
else:
groups.append([msg])
# Anchor group plus the last 3 groups stay.
protected_tail = min(3, max(1, len(groups) - 1))
if len(groups) <= 1 + protected_tail:
return messages, 0
total_est = sum(_estimate_message_tokens(m) for m in messages)
target_est = int(total_est * keep_ratio)
anchor = groups[0]
middle = groups[1:-protected_tail]
tail = groups[-protected_tail:]
current_est = total_est
kept_middle: list[list] = list(middle)
dropped = 0
# Drop oldest-first until the estimate fits the target.
while kept_middle and current_est > target_est:
victim = kept_middle.pop(0)
dropped += len(victim)
current_est -= sum(_estimate_message_tokens(m) for m in victim)
if dropped == 0:
return messages, 0
new_messages = head + anchor
for grp in kept_middle:
new_messages.extend(grp)
for grp in tail:
new_messages.extend(grp)
return new_messages, dropped
_CLIP_MARKER = "\n[... truncated by context_overflow=truncate_middle ...]\n"
# Generous head+tail first; cut harder if the estimate still misses the target.
_CLIP_KEEP_CHARS = (1500, 400)
def _clip_long_contents(messages: list, target_est: int) -> int:
"""Clip oversized string contents middle-out until ``target_est`` is met.
Tool results first, then earlier user turns, the final message last.
Message count and roles never change, so tool pairing holds even when
group-dropping could not free enough. Returns messages clipped.
"""
def _candidates():
tools = [m for m in messages if m.get("role") == "tool"]
users = [m for m in messages[:-1] if m.get("role") == "user"]
last = [messages[-1]] if messages else []
return tools + users + last
clipped = 0
for keep in _CLIP_KEEP_CHARS:
for msg in _candidates():
if sum(_estimate_message_tokens(m) for m in messages) <= target_est:
return clipped
content = msg.get("content")
if not isinstance(content, str) or len(content) <= 2 * keep + len(_CLIP_MARKER):
continue
msg["content"] = content[:keep] + _CLIP_MARKER + content[-keep:]
clipped += 1
return clipped
def _apply_overflow_truncation(body: dict, err_text: str) -> bool:
"""Shrink a passthrough body after an upstream context overflow: drop
middle turn-groups, clip still-oversized contents, clamp ``max_tokens``
to the generation headroom. Returns False when nothing could shrink."""
counts = _parse_overflow_counts(err_text)
messages = body.get("messages") or []
total_est = sum(_estimate_message_tokens(m) for m in messages)
if counts:
n_prompt, n_ctx = counts
keep_ratio = min(0.95, (_OVERFLOW_PROMPT_TARGET_FRACTION * n_ctx) / max(1, n_prompt))
else:
n_ctx = None
keep_ratio = 0.6 # no counts in the error; cut conservatively
# Scale the server-token target into char-estimate units.
target_est = int(total_est * keep_ratio)
new_messages, dropped = _truncate_middle_messages(messages, keep_ratio)
if dropped:
body["messages"] = new_messages
clipped = 0
if sum(_estimate_message_tokens(m) for m in body.get("messages") or []) > target_est:
clipped = _clip_long_contents(body.get("messages") or [], target_est)
if not dropped and not clipped:
return False
if n_ctx:
headroom = max(1024, int(n_ctx * (1.0 - _OVERFLOW_PROMPT_TARGET_FRACTION)))
cur_max = body.get("max_tokens")
body["max_tokens"] = min(cur_max, headroom) if cur_max else headroom
logger.warning(
"context_overflow=truncate_middle: dropped %d middle messages, clipped "
"%d contents (keep_ratio %.2f); retrying within the real window",
dropped,
clipped,
keep_ratio,
)
return True
def _anthropic_stream_error_event(exc, *, force: bool = False):
"""Return an Anthropic in-band stream error event when one is useful."""
_cls = _classify_llama_generation_error(exc)
if _cls is None and not force:
return None
status = 400 if _cls is not None else 500
return build_anthropic_sse_event(
"error",
anthropic_error_body(_friendly_error(exc), status = status),
)
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 _add_empty_content_to_reasoning_deltas(chunk: dict) -> bool:
"""Make reasoning-only deltas palatable to strict OpenAI adapters.
Some clients built on OpenAI-compatible streams ignore or reject chunks whose
delta only contains non-standard ``reasoning_content``. Preserve that field,
but add an empty standard ``content`` member so the chunk is still a valid
text-delta shape and downstream parsers keep the stream alive.
"""
changed = False
choices = chunk.get("choices")
if not isinstance(choices, list):
return False
for choice in choices:
if not isinstance(choice, dict):
continue
delta = choice.get("delta")
if not isinstance(delta, dict):
continue
if "reasoning_content" in delta and "content" not in delta:
delta["content"] = ""
changed = True
return changed
def _normalize_openai_passthrough_sse_line(
raw_line: str, *, cap_parallel_tool_calls: bool = False
) -> str:
"""Normalize one passthrough OpenAI SSE ``data:`` line before relaying.
The function is intentionally narrow: it leaves comments, blank events,
``[DONE]``, and unparseable upstream bytes untouched; parsed chunks are
re-serialized only when a compatibility mutation is actually required.
"""
if not raw_line.startswith("data:"):
return raw_line
# Both mutations key off JSON object keys, so a line without either quoted
# key can never change; skip the parse on the per-token common case.
if '"reasoning_content"' not in raw_line and not (
cap_parallel_tool_calls and '"tool_calls"' in raw_line
):
return raw_line
payload = raw_line[len("data:") :].lstrip()
if payload.strip() in ("", "[DONE]"):
return raw_line
try:
obj = json.loads(payload)
except Exception:
return raw_line
if not isinstance(obj, dict):
return raw_line
changed = _add_empty_content_to_reasoning_deltas(obj)
if cap_parallel_tool_calls and _drop_parallel_tool_call_deltas(obj):
changed = True
if not changed:
return raw_line
return "data: " + json.dumps(obj, separators = (",", ":"), ensure_ascii = False)
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"))
_OPENAI_PASSTHROUGH_TERMINAL_GRACE_S = 2.0
_SSE_DONE_LINE = "data: [DONE]"
def _openai_passthrough_sse_line_terminal_state(raw_line: str) -> Optional[str]:
"""Classify OpenAI-compatible chat stream terminal markers.
Some llama-server builds can emit the logical final chunk (``finish_reason``)
and optional usage chunk, then keep the HTTP stream open without sending the
OpenAI ``data: [DONE]`` sentinel. Classifying those chunks lets Studio close
the client stream promptly while preserving an optional trailing usage chunk.
"""
if not raw_line.startswith("data:"):
return None
data_str = raw_line[5:].lstrip()
if data_str == "[DONE]":
return "done"
try:
data = json.loads(data_str)
except json.JSONDecodeError:
return None
return _openai_passthrough_terminal_state_from_data(data)
def _openai_passthrough_terminal_state_from_data(data) -> Optional[str]:
"""Dict-level core of ``_openai_passthrough_sse_line_terminal_state`` for
callers that already parsed the chunk (avoids a re-parse per relayed line)."""
if not isinstance(data, dict):
return None
if _monitor_openai_error_message(data):
return "error"
choices = data.get("choices")
if isinstance(choices, list):
if not choices and isinstance(data.get("usage"), dict):
return "usage"
for choice in choices:
if isinstance(choice, dict) and choice.get("finish_reason") is not None:
return "finish"
elif isinstance(data.get("usage"), dict):
return "usage"
return None
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 _chat_chunk_sse(completion_id, created, model_name, *, delta, finish_reason) -> str:
"""One ``ChatCompletionChunk`` as an SSE ``data:`` line. The role / content /
final chunks every in-process streamer emits differ only in their ``delta``
and ``finish_reason``."""
chunk = ChatCompletionChunk(
id = completion_id,
created = created,
model = model_name,
choices = [ChunkChoice(delta = delta, finish_reason = finish_reason)],
)
return f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
def _chat_role_chunk(completion_id, created, model_name) -> str:
"""Opening assistant-role chunk for a chat stream."""
return _chat_chunk_sse(
completion_id,
created,
model_name,
delta = ChoiceDelta(role = "assistant"),
finish_reason = None,
)
def _chat_content_chunk(completion_id, created, model_name, text) -> str:
"""A content-delta chunk carrying ``text``."""
return _chat_chunk_sse(
completion_id,
created,
model_name,
delta = ChoiceDelta(content = text),
finish_reason = None,
)
def _chat_reasoning_chunk(completion_id, created, model_name, text) -> str:
"""Like ``_chat_content_chunk`` but on ``reasoning_content`` (renders the UI thinking block).
Carries ``content: ""`` alongside, like the GGUF and passthrough paths, so
strict OpenAI adapters don't drop the reasoning-only delta.
"""
return _chat_chunk_sse(
completion_id,
created,
model_name,
delta = ChoiceDelta(content = "", reasoning_content = text),
finish_reason = None,
)
def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str:
"""Terminal stop chunk (empty delta) carrying the finish reason."""
return _chat_chunk_sse(
completion_id,
created,
model_name,
delta = ChoiceDelta(),
finish_reason = finish_reason,
)
def _chat_tool_calls_chunk(completion_id, created, model_name, tool_calls) -> str:
"""Delta chunk carrying OpenAI tool-call deltas (sibling of ``_chat_content_chunk``)."""
return _chat_chunk_sse(
completion_id,
created,
model_name,
delta = ChoiceDelta(tool_calls = tool_calls),
finish_reason = None,
)
def _sf_heal_events_to_sse(
events,
completion_id,
created,
model_name,
state,
parallel_tool_calls,
monitor_id = None,
):
"""Serialize ``StreamToolCallHealer`` events into chat SSE lines.
``state["idx"]`` tracks the call index across ``feed``/``finalize``;
``parallel_tool_calls is False`` caps promotion to one call (GGUF parity).
The monitor is fed from the same events the client receives, never the
healed-away markup."""
lines = []
for kind, value in events:
if kind == "text":
if value:
lines.append(_chat_content_chunk(completion_id, created, model_name, value))
api_monitor.append_reply(monitor_id, value)
continue
if parallel_tool_calls is False and state["idx"] >= 1:
continue
lines.append(
_chat_tool_calls_chunk(
completion_id,
created,
model_name,
[
{
"index": state["idx"],
"id": value["id"],
"type": "function",
"function": value["function"],
}
],
)
)
_fn = value.get("function") or {}
api_monitor.append_reply(
monitor_id,
("[tool_calls] " if state["idx"] == 0 else "; ")
+ f"{_fn.get('name', '')}({_fn.get('arguments', '')})",
)
state["idx"] += 1
return lines
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))
try:
from core.inference import get_inference_backend
from core.inference.llama_cpp import (
LlamaCppBackend,
_DEFAULT_FIRST_TOKEN_TIMEOUT_S,
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_STREAM_STALL_TIMEOUT_S,
_canonicalize_spec_mode,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
detect_reasoning_flags,
)
from core.inference.llama_server_args import (
_effective_tensor_parallel,
_tensor_parallel_matches_loaded,
parse_split_mode_override,
resolve_tensor_parallel,
strip_shadowing_flags,
validate_extra_args,
)
from core.inference.tensor_fallback import load_with_tensor_fallback
from utils.models import ModelConfig
from utils.inference import load_inference_config
from utils.models.model_config import (
detect_mtp_file,
load_model_defaults,
)
from utils.native_path_leases import (
NativePathLeaseError,
display_label_for_native_path,
is_registered_native_path_label,
redact_native_paths,
verify_native_path_lease,
)
except ImportError:
parent_backend = backend_path.parent / "backend"
if str(parent_backend) not in sys.path:
sys.path.insert(0, str(parent_backend))
from core.inference import get_inference_backend
from core.inference.llama_cpp import (
LlamaCppBackend,
_DEFAULT_FIRST_TOKEN_TIMEOUT_S,
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_STREAM_STALL_TIMEOUT_S,
_canonicalize_spec_mode,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
detect_reasoning_flags,
)
from core.inference.llama_server_args import (
_effective_tensor_parallel,
_tensor_parallel_matches_loaded,
parse_split_mode_override,
resolve_tensor_parallel,
strip_shadowing_flags,
validate_extra_args,
)
from core.inference.tensor_fallback import load_with_tensor_fallback
from utils.models import ModelConfig
from utils.inference import load_inference_config
from utils.models.model_config import (
detect_mtp_file,
load_model_defaults,
)
from utils.native_path_leases import (
NativePathLeaseError,
display_label_for_native_path,
is_registered_native_path_label,
redact_native_paths,
verify_native_path_lease,
)
def _llama_non_streaming_generation_timeout() -> httpx.Timeout:
return httpx.Timeout(_DEFAULT_FIRST_TOKEN_TIMEOUT_S)
def _llama_streaming_generation_timeout() -> httpx.Timeout:
return httpx.Timeout(_DEFAULT_FIRST_TOKEN_TIMEOUT_S)
def _set_stream_response_read_timeout(
response: httpx.Response, read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S
) -> None:
# ``read_timeout_s = None`` clears httpx's read timeout (wait indefinitely),
# used when the stall guard is disabled so a stale first-token deadline
# can't keep timing out post-first-chunk gaps.
try:
timeout_ext = response.request.extensions.get("timeout")
if isinstance(timeout_ext, dict):
timeout_ext["read"] = read_timeout_s
except Exception:
pass
_STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25
_OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1
_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S = 5.0
_OPENAI_PASSTHROUGH_SSE_KEEPALIVE = ": keep-alive\n\n"
def _openai_compat_stream_stall_timeout():
"""Max silent gap after an OpenAI passthrough stream has produced data.
If the socket goes silent after valid SSE data, this bounds how long the
client is kept open. Defaults to the backend-wide stall timeout so this
path stalls out like every sibling stream; set the env var to tighten it
for local serving, or to 0 to disable the guard.
"""
return _positive_float_env(
_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV,
_DEFAULT_STREAM_STALL_TIMEOUT_S,
)
def _openai_passthrough_upstream_headers(*, llama_backend = None) -> dict:
headers = {}
auth_headers = getattr(llama_backend, "_auth_headers", None)
if isinstance(auth_headers, dict):
headers.update(auth_headers)
headers["Connection"] = "close"
return headers
class _CompatSameTaskTimeout:
"""Same-task timeout fallback for Python versions before asyncio.timeout."""
def __init__(self, timeout_s: float):
self.timeout_s = timeout_s
self._task = None
self._handle = None
self._timed_out = False
self._cancelling = 0
async def __aenter__(self):
self._task = asyncio.current_task()
if self._task is None:
return self
if hasattr(self._task, "cancelling"):
self._cancelling = self._task.cancelling()
loop = asyncio.get_running_loop()
self._handle = loop.call_later(max(self.timeout_s, 0), self._cancel_task)
return self
async def __aexit__(self, exc_type, exc, tb):
if self._handle is not None:
self._handle.cancel()
if exc_type is not None and issubclass(exc_type, asyncio.CancelledError):
if self._timed_out:
if self._task is not None and hasattr(self._task, "uncancel"):
if self._task.uncancel() > self._cancelling:
return None
raise asyncio.TimeoutError from exc
return None
def _cancel_task(self) -> None:
self._timed_out = True
if self._task is not None:
self._task.cancel()
def _same_task_timeout(timeout_s: float):
timeout_ctx = getattr(asyncio, "timeout", None)
if timeout_ctx is not None:
return timeout_ctx(timeout_s)
return _CompatSameTaskTimeout(timeout_s)
class _SameTaskStreamingResponse(StreamingResponse):
"""StreamingResponse without Starlette's legacy AnyIO task-group wrapper."""
def __init__(
self,
*args,
unstarted_cleanup = None,
**kwargs,
) -> None:
super().__init__(*args, **kwargs)
# Released when the client disconnects before the body iterator starts:
# its try/finally never runs, so a stream that opens resources before the
# first yield (the passthrough's upstream httpx stream) passes this.
self._unstarted_cleanup = unstarted_cleanup
async def __call__(self, scope, receive, send) -> None:
# send() emits a body message only after the first chunk, so no body
# message means the generator never entered its try/finally.
body_started = False
async def _tracking_send(message) -> None:
nonlocal body_started
if message.get("type") == "http.response.body":
body_started = True
await send(message)
try:
await self.stream_response(_tracking_send)
except OSError: # client disconnected mid-send
if body_started:
# Generator is suspended in its try/finally: throw CancelledError
# (not aclose's GeneratorExit) so its handler finishes the
# api_monitor entry. Fall back to aclose() without athrow.
athrow = getattr(self.body_iterator, "athrow", None)
if athrow is not None:
try:
await athrow(asyncio.CancelledError())
except (asyncio.CancelledError, StopAsyncIteration, RuntimeError):
pass
else:
aclose = getattr(self.body_iterator, "aclose", None)
if aclose is not None:
await aclose()
else:
# Generator never started; aclose()/athrow() are no-ops on it, so
# release eager resources via the hook. getattr guards a response
# built through __new__ without __init__ (tests, pickling).
aclose = getattr(self.body_iterator, "aclose", None)
if aclose is not None:
await aclose()
cleanup = getattr(self, "_unstarted_cleanup", None)
if cleanup is not None:
try:
await cleanup()
except Exception:
pass
raise ClientDisconnect()
if self.background is not None:
await self.background()
def _tracked_cancel_unstarted_cleanup(tracker):
"""unstarted_cleanup that exits ``tracker`` on a pre-start disconnect, when
the generator's finally (which normally exits it) never runs."""
async def _cleanup() -> None:
tracker.__exit__(None, None, None)
return _cleanup
async def _aclose_stream_resources(
*,
watchers = (),
iterator = None,
resp = None,
client = None,
) -> None:
"""Tear down an httpx streaming generator's resources in the required order:
cancel + await each watcher task, then aclose() the byte/line iterator, the
response, and the client. Each step swallows its own exceptions so teardown
always completes; a close-time CancelledError is re-raised only after every
step has run. See _anthropic_passthrough_stream for the ordering rationale."""
for watcher in watchers:
if watcher is not None:
watcher.cancel()
try:
await watcher
except (asyncio.CancelledError, Exception):
pass
close_cancelled = False
if iterator is not None:
try:
await iterator.aclose()
except asyncio.CancelledError:
close_cancelled = True
except Exception:
pass
if resp is not None:
try:
await resp.aclose()
except asyncio.CancelledError:
close_cancelled = True
except Exception:
pass
if client is not None:
try:
await client.aclose()
except asyncio.CancelledError:
close_cancelled = True
except Exception:
pass
if close_cancelled:
raise asyncio.CancelledError()
async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool:
if cancel_event is not None and cancel_event.is_set():
return True
if request is not None and await request.is_disconnected():
if cancel_event is not None:
cancel_event.set()
return True
return False
async def _wait_preheader_cancel(cancel_event = None, request: Optional[Request] = None) -> None:
while not await _preheader_cancelled(cancel_event, request):
await asyncio.sleep(0.05)
async def _send_stream_with_preheader_cancel(
client: httpx.AsyncClient,
req: httpx.Request,
cancel_event = None,
request: Optional[Request] = None,
mark_cancel_on_cancel: bool = True,
) -> Optional[httpx.Response]:
if cancel_event is None and request is None:
return await client.send(req, stream = True)
if await _preheader_cancelled(cancel_event, request):
return None
send_task = asyncio.create_task(client.send(req, stream = True))
cancel_task = asyncio.create_task(_wait_preheader_cancel(cancel_event, request))
async def _stop_send_task() -> None:
try:
await client.aclose()
except Exception:
pass
send_task.cancel()
try:
await send_task
except (asyncio.CancelledError, Exception):
pass
try:
done, _pending = await asyncio.wait(
{send_task, cancel_task},
return_when = asyncio.FIRST_COMPLETED,
)
if send_task in done:
return await send_task
await _stop_send_task()
return None
except asyncio.CancelledError:
if mark_cancel_on_cancel and cancel_event is not None:
cancel_event.set()
await _stop_send_task()
raise
finally:
cancel_task.cancel()
try:
await cancel_task
except (asyncio.CancelledError, Exception):
pass
async def _aiter_llama_stream_items(
async_iter,
*,
cancel_event = None,
request: Optional[Request] = None,
first_token_deadline: Optional[float] = None,
response: Optional[httpx.Response] = None,
post_first_item_read_timeout_s: Optional[
Union[float, Callable[[], Optional[float]]]
] = _DEFAULT_STREAM_STALL_TIMEOUT_S,
):
if first_token_deadline is None:
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
last_item_at: Optional[float] = None
def _post_first_timeout_s() -> Optional[float]:
if callable(post_first_item_read_timeout_s):
return post_first_item_read_timeout_s()
return post_first_item_read_timeout_s
while True:
if cancel_event is not None and cancel_event.is_set():
return
if request is not None and await request.is_disconnected():
if cancel_event is not None:
cancel_event.set()
return
waiting_first_item = last_item_at is None
try:
if waiting_first_item:
remaining_s = first_token_deadline - time.monotonic()
if remaining_s <= 0:
raise httpx.ReadTimeout("The model did not produce a first token in time.")
if response is not None:
_set_stream_response_read_timeout(response, remaining_s)
# Keep httpx/httpcore's AnyIO cancel scope in this task.
# asyncio.wait_for would drive __anext__ in a child task.
async with _same_task_timeout(remaining_s):
item = await async_iter.__anext__()
else:
timeout_s = _post_first_timeout_s()
if (
request is not None
and response is not None
and timeout_s is not None
and last_item_at is not None
):
stall_remaining_s = timeout_s - (time.monotonic() - last_item_at)
if stall_remaining_s <= 0:
raise httpx.ReadTimeout("The model stopped producing tokens mid-response.")
_set_stream_response_read_timeout(response, stall_remaining_s)
item = await async_iter.__anext__()
except asyncio.TimeoutError as exc:
if waiting_first_item:
raise httpx.ReadTimeout("The model did not produce a first token in time.") from exc
raise
except StopAsyncIteration:
return
except httpx.ReadTimeout:
now = time.monotonic()
if last_item_at is None:
if now >= first_token_deadline:
raise
continue
timeout_s = _post_first_timeout_s()
if request is not None and timeout_s is not None and now - last_item_at < timeout_s:
continue
raise httpx.ReadTimeout("The model stopped producing tokens mid-response.")
if last_item_at is None and response is not None:
# The first-token read deadline no longer applies once a chunk has
# arrived: switch to the stall timeout, or clear the read timeout
# entirely when the stall guard is disabled (callable returns None)
# so a long gap can't trip the stale first-token deadline.
_set_stream_response_read_timeout(response, _post_first_timeout_s())
last_item_at = time.monotonic()
yield item
from models.inference import (
LoadRequest,
UnloadRequest,
GenerateRequest,
LoadResponse,
LoadProgressResponse,
UnloadResponse,
InferenceStatusResponse,
ChatCompletionRequest,
ChatCompletionChunk,
ChatCompletion,
ToolConfirmRequest,
ChatMessage,
ChunkChoice,
ChoiceDelta,
CompletionChoice,
CompletionMessage,
CompletionUsage,
ValidateModelRequest,
ValidateModelResponse,
TextContentPart,
ImageContentPart,
ImageUrl,
ResponsesRequest,
ResponsesInputTextPart,
ResponsesInputImagePart,
ResponsesOutputTextPart,
ResponsesUnknownInputItem,
ResponsesFunctionCallInputItem,
ResponsesFunctionCallOutputInputItem,
ResponsesOutputTextContent,
ResponsesOutputMessage,
ResponsesOutputReasoning,
ResponsesOutputReasoningContent,
ResponsesOutputFunctionCall,
ResponsesUsage,
ResponsesResponse,
AnthropicMessagesRequest,
AnthropicMessagesResponse,
AnthropicResponseTextBlock,
AnthropicResponseToolUseBlock,
AnthropicUsage,
CreateOpenAIContainerBody,
DeleteOpenAIContainerBody,
ListOpenAIContainersResponse,
OpenAIContainerRequest,
OpenAIContainerSummary,
)
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,
)
from auth.authentication import get_current_subject
from state.tool_approvals import resolve_tool_decision
from core.inference.key_exchange import decrypt_api_key
from core.inference.model_ids import public_model_id
from core.inference.api_monitor import api_monitor
from core.inference.llama_http import nonstreaming_client
from core.inference.tool_call_parser import (
_strip_function_xml_calls,
_strip_gemma_wrapperless_calls,
_strip_glm_calls,
_strip_mistral_closed_calls,
)
from core.inference.tool_call_parser import TOOL_XML_SIGNALS as _PARSER_TOOL_SIGNALS
from core.inference.passthrough_healing import (
StreamToolCallHealer,
heal_gate,
heal_openai_message,
heal_openai_message_events,
nudge_enabled,
nudge_messages,
nudge_should_retry,
response_has_promotable_calls,
)
from core.inference.providers import get_base_url
from core.inference.external_provider import ExternalProviderClient
from core.inference.chat_templates import resolve_effective_chat_template_override
from storage import providers_db
from utils.utils import safe_error_detail, log_and_http_error
import io
import base64
import numpy as np
from datetime import date as _date
router = APIRouter()
# Studio-only router (not mounted on /v1 OpenAI-compat).
studio_router = APIRouter()
_ARTIFACT_PREVIEW_FRAME_ANCESTORS = "'self' tauri://localhost http://tauri.localhost"
_ARTIFACT_PREVIEW_FRAME_STRICT_CSP = (
"default-src 'none'; "
"script-src 'unsafe-inline'; "
"style-src 'unsafe-inline'; "
"img-src data: blob:; "
"font-src data:; "
"media-src data: blob:; "
"connect-src 'none'; "
"object-src 'none'; "
"base-uri 'none'; "
"form-action 'none'; "
f"frame-ancestors {_ARTIFACT_PREVIEW_FRAME_ANCESTORS}; "
"sandbox allow-scripts"
)
_ARTIFACT_PREVIEW_FRAME_NETWORK_CSP = (
"default-src http: https: data: blob:; "
"script-src 'unsafe-inline' 'unsafe-eval' http: https: data: blob:; "
"script-src-elem 'unsafe-inline' http: https: data: blob:; "
"style-src 'unsafe-inline' http: https: data: blob:; "
"style-src-elem 'unsafe-inline' http: https: data: blob:; "
"img-src http: https: data: blob:; "
"font-src http: https: data: blob:; "
"media-src http: https: data: blob:; "
"connect-src http: https: ws: wss: data: blob:; "
"worker-src http: https: blob:; "
"object-src 'none'; "
"base-uri 'none'; "
"form-action 'none'; "
f"frame-ancestors {_ARTIFACT_PREVIEW_FRAME_ANCESTORS}; "
"sandbox allow-scripts"
)
_ARTIFACT_PREVIEW_FRAME_HTML = """<!doctype html>
<html>
<head><meta charset=\"utf-8\" /></head>
<body>
<script>
(() => {
const createMemoryStorage = () => {
const data = new Map();
return {
get length() { return data.size; },
key: (index) => Array.from(data.keys())[index] ?? null,
getItem: (key) => data.has(String(key)) ? data.get(String(key)) : null,
setItem: (key, value) => data.set(String(key), String(value)),
removeItem: (key) => data.delete(String(key)),
clear: () => data.clear(),
};
};
const installStorageFallback = (name) => {
try {
void window[name];
return;
} catch {
// Opaque-origin sandboxed frames throw SecurityError for Web Storage.
}
try {
Object.defineProperty(window, name, {
value: createMemoryStorage(),
configurable: true,
});
} catch {
// Leave the sandbox failure contained in the canvas if the
// browser refuses to shadow the Web Storage accessor.
}
};
const installStorageFallbacks = () => {
installStorageFallback("localStorage");
installStorageFallback("sessionStorage");
};
const render = (html) => {
installStorageFallbacks();
document.open();
document.write(html);
document.close();
};
installStorageFallbacks();
window.addEventListener("message", (event) => {
const data = event.data;
if (!data || data.type !== "unsloth:artifact-html" || typeof data.html !== "string") return;
render(data.html);
});
})();
</script>
</body>
</html>"""
async def _authenticate_header_or_query(request: Request, token: Optional[str]) -> str:
"""Resolve the bearer token from the Authorization header or the ``?token=``
query param (needed for <img src> / <iframe>, which can't send custom
headers), validate it, and return the subject. Raises 401 when absent."""
auth_header = request.headers.get("authorization")
if auth_header and auth_header.lower().startswith("bearer "):
jwt_token = auth_header[7:]
elif token:
jwt_token = token
else:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Missing authentication token",
)
from fastapi.security import HTTPAuthorizationCredentials
creds = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = jwt_token)
return await get_current_subject(creds)
@studio_router.get("/artifact-preview-frame", include_in_schema = False)
async def artifact_preview_frame(allow_network: bool = False):
"""Serve the opaque sandbox shell for client-side HTML canvases.
No auth token by design: the URL is readable by the untrusted canvas via
location.href, and this static shell exposes no server resource (frame-ancestors
plus the sandbox already gate it), so the CSP is chosen from allow_network alone.
"""
csp = (
_ARTIFACT_PREVIEW_FRAME_NETWORK_CSP if allow_network else _ARTIFACT_PREVIEW_FRAME_STRICT_CSP
)
return Response(
content = _ARTIFACT_PREVIEW_FRAME_HTML,
media_type = "text/html; charset=utf-8",
headers = {
"Cache-Control": "no-store",
"Content-Security-Policy": csp,
"Referrer-Policy": "no-referrer",
"X-Content-Type-Options": "nosniff",
},
)
# Whitespace/escape-tolerant bare-JSON tool-template detector (matches pretty-printed and
# JSON-escaped ``{"name":`` plus the ``"function"`` alias), mirroring the parser's tolerance.
_BARE_JSON_NAME_MARKER_RE = _re.compile(r'\{\s*\\?"(?:name|function)\\?"\s*:')
def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict:
"""Classify reasoning/tool capabilities via the GGUF classifier so flags
match across backends. gpt-oss is overridden: Harmony routes reasoning and
tools through tokenizer channels, not template markup."""
model_id = getattr(backend, "active_model_name", None)
flags = detect_reasoning_flags(
chat_template,
model_identifier = model_id,
log_source = "safetensors",
)
# Markers any supported parser recognises (template advertises tools but
# uses none -> drop the pill). Reuse the parser's own signal list so this
# gate never drifts (a hand-maintained copy lost the DeepSeek variants);
# ``<arg_key>`` is GLM's unique signal, absent from the shared set. The
# bare-JSON ``{"name":`` form is matched below with the whitespace/escape-
# tolerant ``_BARE_JSON_NAME_MARKER_RE`` so pretty-printed or escaped
# templates are not mis-classified as tool-less.
_PARSER_MARKERS = (
*_PARSER_TOOL_SIGNALS,
"<arg_key>",
)
if (
flags.get("supports_tools")
and chat_template
and not any(m in chat_template for m in _PARSER_MARKERS)
and not _BARE_JSON_NAME_MARKER_RE.search(chat_template)
):
logger.info(
"safetensors: template advertises tools but uses an "
"emission format the loop cannot parse; suppressing "
"supports_tools"
)
flags["supports_tools"] = False
# gpt-oss: keep reasoning on, drop tools (Harmony channel, not the
# <tool_call> XML this loop parses).
try:
if hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model():
flags["supports_reasoning"] = True
flags["reasoning_style"] = "reasoning_effort"
flags["supports_tools"] = False
except Exception:
logger.debug("gpt_oss_check_failed", exc_info = True)
return flags
def _generation_prompt_opens_think(template: Optional[str]) -> bool:
"""True when rendering the template's generation prompt ends INSIDE an unclosed ``<think>``.
Distinguishes templates that PREFILL an open ``<think>`` in the assistant generation
prompt (DeepSeek-R1, QwQ, Qwen3-Thinking) -- where the model emits only the closing
``</think>`` and the extractor must start in reasoning mode -- from templates that merely
render PAST assistant ``<think>...</think>`` history while leaving the generation prompt
open with no ``<think>`` (e.g. Kimi-K2-Thinking), where the model self-emits its own block
and the extractor must start in normal mode. Renders a single-user-message probe with the
same sandbox transformers uses; on any failure returns True, preserving the historical
always-on prefill for templates that cannot be rendered here.
"""
if not template:
return False
try:
from jinja2.sandbox import ImmutableSandboxedEnvironment
def _raise_exception(message: str):
raise RuntimeError(message)
env = ImmutableSandboxedEnvironment(
trim_blocks = True,
lstrip_blocks = True,
extensions = ["jinja2.ext.loopcontrols"],
)
env.filters["tojson"] = lambda value, **kwargs: json.dumps(value, ensure_ascii = False)
env.globals["raise_exception"] = _raise_exception
rendered = env.from_string(template).render(
messages = [{"role": "user", "content": "hi"}],
add_generation_prompt = True,
bos_token = "",
eos_token = "",
)
except Exception:
return True
# ``<think>`` is not a substring of ``</think>`` (the ``/`` breaks it), so the last open
# tag sitting after the last close tag means the prompt ends inside an open block.
return rendered.rfind("<think>") > rendered.rfind("</think>")
def _sf_reasoning_prefill_mode(
features: dict,
enable_thinking: Optional[bool],
template: Optional[str] = None,
reasoning_effort: Optional[str] = None,
) -> bool:
"""Whether a safetensors/MLX generation begins INSIDE an unclosed ``<think>``.
``enable_thinking`` templates (Qwen3/GLM) prefill an open ``<think>`` so the model
emits only the closing ``</think>``, and the extractor must start in reasoning mode.
Gated on the STANDARD ``<think>``/``</think>`` markers: bespoke channels (gemma's
``<|think|>``) never emit ``</think>`` and would swallow the answer, so they and
gpt-oss and thinking-disabled requests return False. ``enable_thinking`` None
defaults thinking ON, so a plain request still prefills.
"""
if features.get("reasoning_style") not in ("enable_thinking", "enable_thinking_effort"):
return False
tpl = template or ""
if "</think>" not in tpl and "<think>" not in tpl:
return False
if features.get("reasoning_always_on"):
# enable_thinking_effort + always-on: the effort mechanism (not the prompt shape) keeps
# thinking on, so always-on wins over reasoning_effort and we prefill.
if features.get("reasoning_style") == "enable_thinking_effort":
return True
# ``reasoning_always_on`` fires on paired ``<think>...</think>`` anywhere in the
# template, including markup that only renders PAST assistant history (Kimi-K2-Thinking)
# while the generation prompt opens none. Prefill only when the generation prompt opens
# one, else the extractor captures a normal answer as reasoning_content and returns blank.
return _generation_prompt_opens_think(tpl)
if not features.get("supports_reasoning"):
return False
if enable_thinking is False:
return False
# Thinking-off arrives as reasoning_effort "none" on enable_thinking_effort models; honor it
# so we don't prefill and capture the answer. Plain enable_thinking models ignore effort.
if features.get("reasoning_style") == "enable_thinking_effort" and reasoning_effort == "none":
return False
return True
def _effective_enable_tools(payload) -> Optional[bool]:
"""Resolve `payload.enable_tools` against the process-level tool policy.
Returns the policy value when set (CLI hard-override from `unsloth run`),
else the per-request value.
"""
from state.tool_policy import get_tool_policy
policy = get_tool_policy()
return policy if policy is not None else payload.enable_tools
def _explicit_studio_tool_loop_requested(payload) -> bool:
"""True when the request itself asks Studio to execute local tools.
Process-wide CLI policy can default Studio's tool loop on for ordinary chat,
but it must not steal OpenAI-compatible client tools or response_format
requests from the llama-server passthrough path. A policy of ``False``
(--disable-tools) vetoes even an explicit ``enable_tools: true`` ask.
"""
from state.tool_policy import get_tool_policy
policy = get_tool_policy()
return policy is not False and (payload.enable_tools is True or bool(payload.mcp_enabled))
# Cancel registry. Proxies (e.g. Colab) can swallow client fetch aborts so
# is_disconnected() never fires. POST /inference/cancel looks up in-flight
# cancel_events here by cancel_id (per-run) or session_id / completion_id
# (fallbacks).
_CANCEL_REGISTRY: dict[str, set[threading.Event]] = {}
_CANCEL_LOCK = threading.Lock()
# Cancel POSTs arriving before registration are stashed; the next matching
# __enter__ replays set() within the TTL.
_PENDING_CANCELS: dict[str, float] = {}
_PENDING_CANCEL_TTL_S = 30.0
def _prune_pending(now: float) -> None:
for k in [k for k, ts in _PENDING_CANCELS.items() if now - ts > _PENDING_CANCEL_TTL_S]:
_PENDING_CANCELS.pop(k, None)
class _TrackedCancel:
"""Register cancel_event in _CANCEL_REGISTRY for the block's duration."""
def __init__(self, event: threading.Event, *keys):
self.event = event
self.keys = tuple(k for k in keys if k)
def __enter__(self):
# Register + consume-pending in one critical section to close the
# TOCTOU race against a concurrent cancel POST.
should_cancel = False
with _CANCEL_LOCK:
for k in self.keys:
_CANCEL_REGISTRY.setdefault(k, set()).add(self.event)
now = time.monotonic()
_prune_pending(now)
for k in self.keys:
if k and _PENDING_CANCELS.pop(k, None) is not None:
should_cancel = True
if should_cancel:
self.event.set()
return self.event
def __exit__(self, *exc):
with _CANCEL_LOCK:
for k in self.keys:
bucket = _CANCEL_REGISTRY.get(k)
if bucket is None:
continue
bucket.discard(self.event)
if not bucket:
_CANCEL_REGISTRY.pop(k, None)
return False
def _cancel_by_keys(keys) -> int:
"""Set cancel_event for matching registry entries; no stash.
session_id/completion_id are shared across runs on the same thread, so
stashing them would ghost-cancel the user's next request. Only cancel_id
is per-run unique (see _cancel_by_cancel_id_or_stash)."""
if not keys:
return 0
events: set[threading.Event] = set()
with _CANCEL_LOCK:
_prune_pending(time.monotonic())
for k in keys:
bucket = _CANCEL_REGISTRY.get(k)
if bucket:
events.update(bucket)
for ev in events:
ev.set()
return len(events)
def _cancel_by_cancel_id_or_stash(cancel_id: str) -> int:
"""Atomic lookup-or-stash; pairs with _TrackedCancel.__enter__ to
close the TOCTOU race."""
now = time.monotonic()
events: set[threading.Event] = set()
with _CANCEL_LOCK:
_prune_pending(now)
bucket = _CANCEL_REGISTRY.get(cancel_id)
if bucket:
events.update(bucket)
else:
_PENDING_CANCELS[cancel_id] = now
for ev in events:
ev.set()
return len(events)
async def _await_cancel_then_close(cancel_event, resp) -> None:
"""Watch a threading.Event from asyncio and close ``resp`` when it fires.
Used by passthrough streamers so a /cancel POST can interrupt while the
async iterator is blocked on llama-server prefill. Without it the in-loop
``cancel_event.is_set()`` check is unreachable until the first SSE chunk
arrives -- exactly the proxy/Colab case the cancel POST exists for.
Polls a threading.Event since the cancel registry is keyed by
threading.Event (so the sync /cancel handler can call .set()). The 50ms
cadence adds at most that latency to a prefill cancel; the common
streaming-cancel path still sees the event on the iterator's next chunk.
"""
try:
while not cancel_event.is_set():
await asyncio.sleep(0.05)
try:
await resp.aclose()
except Exception:
pass
except asyncio.CancelledError:
return
async def _await_disconnect_then_close(request, resp, cancel_event) -> None:
"""Close ``resp`` on client disconnect; sets ``cancel_event`` first so
the streamer's RemoteProtocolError handler treats it as cancellation.
Catches aborts the in-loop /cancel check misses during prefill. #5692.
"""
try:
while not await request.is_disconnected():
await asyncio.sleep(0.1)
cancel_event.set()
try:
await resp.aclose()
except Exception as e:
logger.debug("Failed to close response on disconnect: %s", e)
except asyncio.CancelledError:
return
async def _await_disconnect_then_cancel(request, cancel_event) -> None:
"""Set ``cancel_event`` when a same-task local stream disconnects."""
try:
while not await request.is_disconnected():
await asyncio.sleep(0.1)
cancel_event.set()
except asyncio.CancelledError:
return
def _cancelable_nonstreaming_client() -> httpx.AsyncClient:
return httpx.AsyncClient(
limits = httpx.Limits(max_connections = 1, max_keepalive_connections = 0),
trust_env = False,
)
async def _await_cancel_or_disconnect_then_close_client(
*, cancel_event, request: Optional[Request], client: httpx.AsyncClient
) -> None:
"""Close a dedicated non-streaming upstream client on cancel/disconnect.
The shared ``nonstreaming_client()`` is pooled, so cancelable generation calls
use a per-request client. Closing it interrupts a blocked llama-server
request without affecting unrelated pooled non-streaming calls.
"""
try:
while True:
if cancel_event is not None and cancel_event.is_set():
break
if request is not None and await request.is_disconnected():
if cancel_event is not None:
cancel_event.set()
break
await asyncio.sleep(0.1)
try:
await client.aclose()
except Exception:
pass
except asyncio.CancelledError:
return
async def _stop_local_disconnect_cancel_watcher(watcher) -> None:
watcher.cancel()
try:
await watcher
except (asyncio.CancelledError, Exception):
pass
# Centralized local/server tool nudge. Keep render_html guidance gated to turns
# where the canvas tool is actually present in the tool schema; otherwise
# small local models can hallucinate a missing tool call instead of following
# the fenced-HTML fallback prompt.
_TOOL_BASE_NUDGE = (
"Tools are available when they materially improve the answer. Use an enabled "
"tool for current facts, calculations, code execution, or canvases when it "
"materially helps; otherwise answer normally and follow the user's requested "
"format."
)
_TOOL_WEB_COMPACT_TIP = "When using web_search, do not repeat the same search query."
_TOOL_WEB_EXPANDED_TIP = (
"When using web_search and a result URL is relevant, fetch its full content "
"by calling web_search with the url parameter. Do not repeat the same search "
"query. If a search returns no useful results, try rephrasing or fetching a "
"result URL directly."
)
_TOOL_CODE_TIP = (
"Use code execution for math, calculations, data processing, or to parse "
"and analyze information from tool results."
)
_TOOL_ARTIFACT_TIP = (
"For HTML, CSS, or JavaScript canvas requests, call render_html once when "
"it is available with one complete self-contained HTML document in the code "
"argument. After render_html succeeds, do not call it again in the same "
"response unless the user asks for changes. Future user requests for new "
"canvases may call render_html once."
)
def _build_tool_action_nudge(*, tools: list[dict], model_name: str) -> str:
tool_names = {
(tool.get("function") or {}).get("name")
for tool in tools
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
}
has_web = "web_search" in tool_names
has_code = "python" in tool_names or "terminal" in tool_names
has_artifact = "render_html" in tool_names
if not (has_web or has_code or has_artifact):
return ""
model_size_b = _extract_model_size_b(model_name)
compact_web_tip = model_size_b is not None and model_size_b < 9
tool_tip_parts: list[str] = []
if has_web:
tool_tip_parts.append(_TOOL_WEB_COMPACT_TIP if compact_web_tip else _TOOL_WEB_EXPANDED_TIP)
if has_code:
tool_tip_parts.append(_TOOL_CODE_TIP)
if has_artifact:
tool_tip_parts.append(_TOOL_ARTIFACT_TIP)
return (
f"The current date is {_date.today().isoformat()}. "
+ _TOOL_BASE_NUDGE
+ " "
+ " ".join(tool_tip_parts)
)
# Nudge appended when the RAG knowledge-base tool is active: ground answers in
# the attached documents instead of model memory.
_RAG_GROUNDING_NUDGE = (
"The user has attached documents to this conversation. Relevant "
"passages are retrieved and provided to you automatically; base "
"your answer on them and cite them. You can also call "
"search_knowledge_base to look for more. Do not answer from "
"memory when the attached documents are relevant."
)
async def _select_request_tools(
payload: ChatCompletionRequest, *, tools_on: bool, mcp_allowed: bool
) -> list[dict]:
"""Resolve the tool list for a chat request: built-ins filtered by the
caller's opt-in (empty when MCP-only), the RAG tool dropped without a
retrieval scope, then enabled MCP tools appended. An empty result means the
caller should skip the tool loop, so a model-emitted built-in call can't
piggy-back on the empty allow-list."""
from core.inference.tools import ALL_TOOLS, get_enabled_mcp_tools
if not tools_on:
# MCP-only request: skip built-ins, leave room for MCP tools.
tools = []
elif payload.enabled_tools is not None:
tools = [t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools]
else:
# Copy so the shared module-global tool list can't be mutated by callers.
tools = list(ALL_TOOLS)
# Drop the RAG tool without a scope: nothing to search over.
if not payload.rag_scope:
tools = [t for t in tools if t["function"]["name"] != "search_knowledge_base"]
if mcp_allowed:
tools = tools + await get_enabled_mcp_tools()
return tools
def _apply_rag_nudge(nudge: str, tools: list[dict], *, rag_scope) -> str:
"""Append the RAG grounding nudge to ``nudge`` when the knowledge-base tool
is active (search_knowledge_base present and a retrieval scope is set). The
date is prefixed when the tool nudge is empty (RAG-only tool set). Returns
``nudge`` unchanged when RAG isn't active."""
tool_names = {(t.get("function") or {}).get("name") for t in (tools or [])}
if "search_knowledge_base" not in tool_names or not rag_scope:
return nudge
if not nudge:
date_line = f"The current date is {_date.today().isoformat()}."
return date_line + " " + _RAG_GROUNDING_NUDGE
return nudge + " " + _RAG_GROUNDING_NUDGE
# Strip leaked tool-call markup: every shared-parser format plus the leak shapes
# llama_cpp.py's speculative buffer splits across the visible/DRAIN boundary:
# 1. well-formed `<tool_call>...</tool_call>` / `<function=...>...</function>`
# 2. orphan opening to EOF (close was DRAINED)
# 3. bare orphan close (open was DRAINED)
# 4. tail-only `</parameter>` (outer close truncated by EOS); anchored to
# `\Z` so mid-text `<parameter>` in user code samples survives.
# 5. Mistral `[TOOL_CALLS]name{json}` / rehearsal `name[ARGS]{json}`: the balanced
# scan removes the whole call (a non-greedy regex would truncate nested JSON).
# DeepSeek/GLM/Kimi envelopes are covered by the parser's own arms/scans, so a signal
# we parse is never left un-stripped; the DeepSeek opener alternation is the parser's own.
from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC
_TOOL_XML_RE = _re.compile(
# Arm order/notes: the closed ``<function=...>`` arm runs first and extends
# to the call's REAL close so a literal ``</function>`` in a value does not
# leak the tail; the combined arm still catches ``<tool_call>`` and orphan
# tails. The python_tag arm bounds only on REAL Llama control sentinels
# (stopping at any ``<|`` truncated on literal ``<|x|>`` tokens in values).
# The last arms cover DeepSeek envelopes (all opener variants), Kimi section
# blocks, and bare Kimi calls. Name class ``[\w.\-]`` mirrors the parser.
# Those three arms carry a call-shaped lookahead (matching the parser's
# ``_TOOL_ALL_PATS``): a prose answer that merely mentions a marker
# (``See <|tool_call_begin|> in the docs``) is only stripped when a real
# call actually follows the marker, or the marker is a bare fragment at EOF.
r'<function(?:=[\w.\-]+|\s+name="[\w.\-]+")>(?:(?!<function(?:=[\w.\-]+|\s+name="[\w.\-]+")>).)*</function>'
r'|<(?:tool_call|function(?:=[\w.\-]+|\s+name="[\w.\-]+"))>.*?(?:</(?:tool_call|function)>|\Z)'
r"|<\|tool_call>.*?(?:<tool_call\|>|\Z)"
r"|</(?:tool_call|function)>"
r"|<tool_call\|>"
r"|<\|python_tag\|>(?:[^<]|<(?!\|(?:eot_id|eom_id|python_tag|start_header_id|end_header_id|begin_of_text|finetune_right_pad_id)\|))*"
r"|\[/TOOL_CALLS\]"
# Truncated canonical array (closing ``]`` lost to EOS): the balanced scan cannot remove
# it, so strip its tail here.
r"|\[TOOL_CALLS\]\s*\[.*\Z"
# Named / v11 forms and bare rehearsal; arms aligned with the parser regexes.
r"|\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|.*?\Z)"
# Rehearsal: balanced/truncated body or bare marker at EOS only (prose ``foo[ARGS]``
# survives); NAME captured as ``reh`` for the inactive-name display gate.
r"|(?<!\[CALL_ID\])\b(?P<reh>[\w-]+)\[ARGS\]\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|\{.*\Z|\Z)"
# DeepSeek envelopes (all opener variants), Kimi section blocks, and bare Kimi calls;
# each arm carries a call-shaped lookahead so prose merely mentioning a marker survives.
r"|"
+ _DS_OPEN_SRC
+ r"(?=\s*(?:<tool▁call▁begin>|function)|\s*$).*?(?:<tool▁calls▁end>|\Z)"
r"|<\|tool_calls_section_begin\|>(?=\s*<\|tool_call_begin\|>|\s*$).*?(?:<\|tool_calls_section_end\|>|\Z)"
r"|<\|tool_call_begin\|>(?=\s*[A-Za-z_][\w.\-]*:\d|\s*$).*?(?:<\|tool_call_end\|>|\Z)"
# ``</param>`` is the attribute-form alias of ``</parameter>`` (the parser accepts
# both); strip a tail-only orphan close of either spelling.
r"|</(?:parameter|param)>\s*\Z",
_re.DOTALL,
)
# Closed-only variant for segments before the last think block: the ``\Z``-anchored arms
# would treat a segment boundary as EOS and strip prose ``foo[ARGS]``.
_TOOL_XML_CLOSED_RE = _re.compile(
r"<(?:tool_call|function=[\w-]+)>.*?</(?:tool_call|function)>"
r"|<\|tool_call>.*?<tool_call\|>"
r"|</(?:tool_call|function)>"
r"|<tool_call\|>"
r"|\[/TOOL_CALLS\]",
_re.DOTALL,
)
def _gemma_strip_gate(tools) -> set:
"""Enabled tool NAMES gating the wrapper-less Gemma strip (mirrors the
parser/loop gate: only an enabled ``call:foo{...}`` is a call). With NO tools
enabled this returns an EMPTY set, not ``None``: every ``call:NAME{...}`` is
then prose, and ``None`` would strip-all and delete a legitimate answer."""
names = {
(t.get("function") or {}).get("name")
for t in (tools or [])
if isinstance(t, dict) and isinstance(t.get("function"), dict)
}
names.discard(None)
return names
def _display_tool_name_gate(active_tools):
"""Active tool NAMES for gating the rehearsal display strip, or None when no tools
are enabled. ``None`` keeps the legacy strip-all behavior, mirroring the loop gate:
a bare ``NAME[ARGS]`` is a call only when NAME is active; without a tool list every
identifier stays ambiguous, so strip."""
names = {
(t.get("function") or {}).get("name")
for t in (active_tools or [])
if isinstance(t, dict) and isinstance(t.get("function"), dict)
}
names.discard(None)
return names or None
def _strip_tool_xml_for_display(
text: str,
*,
auto_heal_tool_calls: bool,
enabled_tool_names: Optional[set] = None,
) -> str:
"""Apply route-level XML leak cleanup only when Auto-Heal is enabled.
Mirrors the parser-side segment scan: balanced strips first (Mistral, gated Gemma
wrapper-less, GLM real-close, guarded function-XML close at each call's REAL terminator
so literal markup inside a value is data), then the ``_TOOL_XML_RE`` arms cover the
DeepSeek / Kimi / orphan forms. ``<think>`` blocks are preserved verbatim and the
``\\Z``-anchored tail arms run only on the last segment (prose ``foo[ARGS]`` before a
block survives). ``enabled_tool_names`` (when not None) gates the ambiguous bare-rehearsal
``NAME[ARGS]{...}`` and wrapper-less Gemma ``call:NAME{...}`` strips on the active tool
list; an inactive NAME is prose and is kept. The ``[TOOL_CALLS]`` control-token arms strip
unconditionally regardless of NAME."""
if not auto_heal_tool_calls:
return text
from core.tool_healing import _strip_bracket_tag_calls, strip_outside_think
def _keep_inactive_rehearsal(m) -> str:
# Only the bare-rehearsal arm captures ``reh``; with a tool list an inactive
# NAME[ARGS]{...} is prose -- keep it.
if enabled_tool_names is not None:
name = m.groupdict().get("reh")
if name is not None and name not in enabled_tool_names:
return m.group(0)
return ""
def _strip_segment(seg: str, is_last: bool) -> str:
# Scan strips close at each call's REAL terminator (a literal ``</function>`` or a
# nested marker quoted inside a value cannot truncate the strip); the regex arms below
# cover the attribute form and the DeepSeek / Kimi / orphan families.
seg = _strip_mistral_closed_calls(seg)
seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names)
if is_last:
seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names)
seg = _strip_glm_calls(seg, final = is_last)
seg = _strip_function_xml_calls(seg, final = is_last)
if is_last:
return _TOOL_XML_RE.sub(_keep_inactive_rehearsal, seg)
return _TOOL_XML_CLOSED_RE.sub("", seg)
return strip_outside_think(text, _strip_segment)
def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str:
# Mistral balanced-brace pre-strip (kept explicit so the regression guards see it), then
# the shared think-aware display strip -- the one raw _TOOL_XML_RE.sub lives inside
# _strip_tool_xml_for_display, so every route cleanup site shares it. ``enabled_tool_names``
# gates the Gemma wrapper-less strip; ``None`` strips every closed call.
text = _strip_mistral_closed_calls(text)
return _strip_tool_xml_for_display(
text, auto_heal_tool_calls = True, enabled_tool_names = enabled_tool_names
)
logger = get_logger(__name__)
def _monitor_content_text(content) -> str:
if content is None:
return ""
if isinstance(content, str):
return content
if isinstance(content, list):
parts: list[str] = []
for part in content:
if isinstance(part, dict):
ptype = part.get("type")
if ptype in ("text", "input_text", "output_text"):
text = part.get("text")
if isinstance(text, str):
parts.append(text)
elif ptype in ("image_url", "input_image", "image"):
parts.append("[image]")
else:
parts.append(f"[{ptype or 'content'}]")
else:
ptype = getattr(part, "type", None)
text = getattr(part, "text", None)
if isinstance(text, str):
parts.append(text)
elif ptype in ("image_url", "input_image", "image"):
parts.append("[image]")
elif ptype:
parts.append(f"[{ptype}]")
return "\n".join(parts)
return str(content)
def _monitor_prompt_from_messages(messages) -> str:
lines: list[str] = []
for msg in messages or []:
role = msg.get("role") if isinstance(msg, dict) else getattr(msg, "role", "")
content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", "")
tool_calls = (
msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None)
)
text = _monitor_content_text(content)
if tool_calls and not text:
text = "[tool calls]"
if text:
lines.append(f"{role or 'message'}: {text}")
return "\n\n".join(lines)
def _monitor_usage(
monitor_id: Optional[str],
usage: Optional[dict],
context_length = None,
):
if not usage:
return
api_monitor.set_usage(
monitor_id,
prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens"),
completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens"),
total_tokens = usage.get("total_tokens"),
context_length = context_length,
)
def _monitor_call_text(name: Any, arguments: Any = None) -> str:
call_name = str(name or "tool")
if arguments is None or arguments == "":
return f"Tool call: {call_name}"
if not isinstance(arguments, str):
args_text = json.dumps(arguments, default = str)
else:
args_text = arguments
if len(args_text) > 500:
args_text = args_text[:497] + "..."
return f"Tool call: {call_name}({args_text})"
def _monitor_tool_calls_text(tool_calls: Any) -> str:
if not isinstance(tool_calls, list):
return ""
parts: list[str] = []
for tool_call in tool_calls:
if not isinstance(tool_call, dict):
continue
fn = tool_call.get("function") or {}
if not isinstance(fn, dict):
fn = {}
name = fn.get("name") or tool_call.get("name") or "tool"
args = fn.get("arguments")
if args is None:
args = tool_call.get("arguments")
parts.append(_monitor_call_text(name, args))
return "\n".join(parts)
def _monitor_openai_chunk(
monitor_id: Optional[str],
data: dict,
context_length = None,
):
if not monitor_id:
return
_monitor_usage(monitor_id, data.get("usage"), context_length)
# Defensive: ignore malformed shapes so the helper never raises into the
# streaming generator and aborts the user's response.
choices = data.get("choices")
if not isinstance(choices, list) or not choices:
return
reply_parts: list[tuple[int, str]] = []
for idx, choice in enumerate(choices):
if not isinstance(choice, dict):
continue
delta = choice.get("delta") or {}
message = choice.get("message") or {}
content = delta.get("content") if isinstance(delta, dict) else None
if content:
api_monitor.append_reply(monitor_id, content)
continue
if isinstance(delta, dict):
tool_text = _monitor_tool_calls_text(delta.get("tool_calls"))
if tool_text:
api_monitor.append_reply(monitor_id, tool_text)
continue
if isinstance(choice.get("text"), str):
reply_parts.append((idx, choice["text"]))
elif isinstance(message, dict):
text = message.get("content")
if isinstance(text, str):
reply_parts.append((idx, text))
else:
tool_text = _monitor_tool_calls_text(message.get("tool_calls"))
if tool_text:
reply_parts.append((idx, tool_text))
if not reply_parts:
return
if len(choices) == 1:
api_monitor.append_reply(monitor_id, reply_parts[0][1])
return
api_monitor.append_reply(
monitor_id,
"\n\n".join(f"Choice {idx + 1}:\n{text}" for idx, text in reply_parts),
)
def _monitor_openai_error_message(data: dict) -> Optional[str]:
error = data.get("error")
if isinstance(error, dict):
message = error.get("message")
if isinstance(message, str) and message:
return message
return json.dumps(error)
if isinstance(error, str) and error:
return error
return None
def _monitor_openai_sse_line(
monitor_id: Optional[str],
raw_line: str,
context_length = None,
) -> Optional[str]:
if not monitor_id:
return None
# SSE spec allows `data:value` and `data: value`; accept both.
if not raw_line.startswith("data:"):
return None
data_str = raw_line[5:].lstrip()
if data_str == "[DONE]":
api_monitor.finish(monitor_id)
return "done"
try:
data = json.loads(data_str)
except json.JSONDecodeError:
return None
if isinstance(data, dict):
error_message = _monitor_openai_error_message(data)
if error_message:
api_monitor.fail(monitor_id, error_message)
return "error"
_monitor_openai_chunk(monitor_id, data, context_length)
return None
def _monitor_openai_sse_event(
monitor_id: Optional[str],
event: bytes,
context_length = None,
) -> None:
for line in event.decode("utf-8", errors = "ignore").splitlines():
_monitor_openai_sse_line(monitor_id, line.strip(), context_length)
def _monitor_anthropic_usage(
monitor_id: Optional[str],
usage: Optional[dict],
context_length = None,
) -> None:
if not usage:
return
_monitor_usage(
monitor_id,
{
"prompt_tokens": usage.get("input_tokens") or usage.get("prompt_tokens"),
"completion_tokens": usage.get("output_tokens") or usage.get("completion_tokens"),
"total_tokens": usage.get("total_tokens"),
},
context_length,
)
_ANTHROPIC_MONITOR_TOOL_BLOCKS: dict[str, dict[int, bool]] = {}
def _monitor_anthropic_index(data: dict) -> int:
try:
return int(data.get("index") or 0)
except (TypeError, ValueError):
return 0
def _monitor_anthropic_payload(
monitor_id: Optional[str],
data: dict,
context_length = None,
) -> Optional[str]:
if not monitor_id or not isinstance(data, dict):
return None
event_type = data.get("type")
if event_type == "message_start":
message = data.get("message") or {}
if isinstance(message, dict):
_monitor_anthropic_usage(monitor_id, message.get("usage"), context_length)
return None
if event_type == "content_block_start":
content_block = data.get("content_block") or {}
if isinstance(content_block, dict) and content_block.get("type") == "tool_use":
index = _monitor_anthropic_index(data)
_ANTHROPIC_MONITOR_TOOL_BLOCKS.setdefault(monitor_id, {})[index] = False
api_monitor.append_reply(monitor_id, _monitor_call_text(content_block.get("name")))
return None
if event_type == "content_block_delta":
delta = data.get("delta") or {}
text = delta.get("text") if isinstance(delta, dict) else None
if isinstance(text, str) and text:
api_monitor.append_reply(monitor_id, text)
elif isinstance(delta, dict) and delta.get("type") == "input_json_delta":
index = _monitor_anthropic_index(data)
tool_blocks = _ANTHROPIC_MONITOR_TOOL_BLOCKS.get(monitor_id) or {}
if index in tool_blocks:
if not tool_blocks[index]:
api_monitor.append_reply(monitor_id, "\nInput: ")
tool_blocks[index] = True
partial_json = delta.get("partial_json")
if isinstance(partial_json, str) and partial_json:
api_monitor.append_reply(monitor_id, partial_json)
return None
if event_type == "content_block_stop":
index = _monitor_anthropic_index(data)
tool_blocks = _ANTHROPIC_MONITOR_TOOL_BLOCKS.get(monitor_id)
if tool_blocks is not None:
tool_blocks.pop(index, None)
if not tool_blocks:
_ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None)
return None
if event_type == "message_delta":
_monitor_anthropic_usage(monitor_id, data.get("usage"), context_length)
return None
if event_type == "error":
error = data.get("error") or {}
if isinstance(error, dict):
message = error.get("message") or json.dumps(error, default = str)
else:
message = str(error)
api_monitor.fail(monitor_id, message)
return "error"
return None
def _monitor_anthropic_sse_line(
monitor_id: Optional[str],
raw_line: str,
context_length = None,
) -> Optional[str]:
if not monitor_id or not raw_line.startswith("data:"):
return None
data_str = raw_line[5:].lstrip()
try:
data = json.loads(data_str)
except json.JSONDecodeError:
return None
return _monitor_anthropic_payload(monitor_id, data, context_length)
def _monitor_anthropic_content_blocks(content: Any) -> str:
if not isinstance(content, list):
return ""
parts: list[str] = []
for block in content:
if not isinstance(block, dict):
continue
if block.get("type") == "text" and isinstance(block.get("text"), str):
parts.append(block["text"])
elif block.get("type") == "tool_use":
parts.append(_monitor_call_text(block.get("name"), block.get("input")))
return "".join(parts)
def _monitor_anthropic_json_response(
response,
monitor_id: Optional[str],
context_length = None,
) -> None:
if not monitor_id:
return
body = getattr(response, "body", b"")
try:
data = json.loads(body.decode("utf-8") if isinstance(body, bytes) else body)
except Exception:
api_monitor.finish(monitor_id)
return
if not isinstance(data, dict):
api_monitor.finish(monitor_id)
return
text = _monitor_anthropic_content_blocks(data.get("content"))
if text:
api_monitor.set_reply(monitor_id, text)
_monitor_anthropic_usage(monitor_id, data.get("usage"), context_length)
api_monitor.finish(monitor_id)
def _monitor_anthropic_response(
response,
monitor_id,
context_length = None,
cancel_event = None,
):
if not monitor_id:
return response
body_iterator = getattr(response, "body_iterator", None)
if body_iterator is None:
_monitor_anthropic_json_response(response, monitor_id, context_length)
return response
async def _monitored_body():
terminal = False
try:
async for chunk in body_iterator:
text = (
chunk.decode("utf-8", errors = "ignore")
if isinstance(chunk, (bytes, bytearray))
else str(chunk)
)
for line in text.splitlines():
if (
_monitor_anthropic_sse_line(
monitor_id,
line.strip(),
context_length,
)
== "error"
):
terminal = True
yield chunk
if not terminal:
api_monitor.finish(
monitor_id,
"cancelled"
if cancel_event is not None and cancel_event.is_set()
else "completed",
)
except asyncio.CancelledError:
if cancel_event is not None:
cancel_event.set()
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as exc:
api_monitor.fail(monitor_id, _friendly_error(exc))
raise
finally:
_ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None)
response.body_iterator = _monitored_body()
return response
def _monitor_context_length() -> Optional[int]:
llama_backend = get_llama_cpp_backend()
if getattr(llama_backend, "is_loaded", False):
context_length = _positive_int_or_none(getattr(llama_backend, "context_length", None))
if context_length is not None:
return context_length
backend = get_inference_backend()
if not backend.active_model_name:
return None
models = getattr(backend, "models", {}) or {}
model_info = models.get(backend.active_model_name, {}) if isinstance(models, dict) else {}
context_length = _positive_int_or_none(model_info.get("context_length"))
if context_length is not None:
return context_length
for candidate in (
getattr(backend, "context_length", None),
getattr(backend, "max_seq_length", None),
):
context_length = _positive_int_or_none(candidate)
if context_length is not None:
return context_length
return None
def _monitor_active_model() -> Optional[str]:
llama_backend = get_llama_cpp_backend()
if getattr(llama_backend, "is_loaded", False):
return getattr(llama_backend, "model_identifier", None)
backend = get_inference_backend()
return backend.active_model_name
def _validate_native_gguf_companion(
companion_path: str | None, gguf_path: str | None, label: str
) -> None:
"""Reject a companion GGUF (mmproj / MTP drafter) that a native-lease load
would otherwise hand to llama-server: must be a regular file (no symlink
escaping the leased directory) living next to the selected GGUF."""
if not companion_path or not gguf_path:
return
import stat as _stat_module
companion = Path(companion_path)
gguf = Path(gguf_path)
try:
companion_lstat = os.lstat(companion)
except OSError as exc:
raise HTTPException(
status_code = 400,
detail = f"Native {label} is no longer accessible.",
) from exc
if _stat_module.S_ISLNK(companion_lstat.st_mode) or not _stat_module.S_ISREG(
companion_lstat.st_mode
):
raise HTTPException(
status_code = 400,
detail = f"Native {label} must be a regular file.",
)
try:
if companion.resolve(strict = True).parent != gguf.resolve(strict = True).parent:
raise HTTPException(
status_code = 400,
detail = f"Native {label} must live next to the selected GGUF.",
)
except OSError as exc:
raise HTTPException(
status_code = 400,
detail = f"Native {label} is no longer accessible.",
) from exc
def _normalise_settings_str(value: Optional[str]) -> Optional[str]:
"""Lowercase + strip a settings string, mapping blank/None to None."""
if value is None:
return None
if isinstance(value, str):
stripped = value.strip().lower()
return stripped or None
return value
def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[str]]) -> bool:
"""Whether an inherited --split-mode should be stripped on reload.
The binary Tensor Parallelism toggle can't carry --split-mode's row/none/
layer modes, so only strip when the toggle overrides it: tensor being turned
on, or the inherited mode is tensor (toggle turning it off). Non-tensor modes
survive. Shared by the inheritance strip and the already-loaded stale check
so they agree on what reload would do.
"""
fields_set = getattr(request, "model_fields_set", set())
return "tensor_parallel" in fields_set and (
request.tensor_parallel or resolve_tensor_parallel(backend_extra, False)
)
def _carry_preserved_tensor_intent(
*, preserved: bool, same_model: bool, explicit_drop: bool
) -> bool:
"""Carry a preserved multi-GPU layer fallback forward only for a reload of the
SAME loaded model that doesn't explicitly drop tensor intent, so a fitting model
isn't collapsed to one GPU on a ctx-only change -- but an unrelated model switch
(without /unload) or an explicit tensor-off doesn't inherit it (#6659)."""
return preserved and same_model and not explicit_drop
def _is_explicit_tensor_drop(request: LoadRequest) -> bool:
"""True only when the request explicitly selects a non-tensor --split-mode (e.g.
layer/row/none), a deliberate departure from a preserved tensor->layer fallback.
A bare tensor_parallel field is NOT a drop: the Studio UI always sends it and echoes
the /load response's resolved value back, so after a fallback every reload carries
tensor_parallel=false even though the user never changed it -- treating that as a drop
would collapse the preserved multi-GPU placement on the next ctx/settings reload. An
empty clear is not a drop either (a fallback always stores --split-mode layer, never a
tensor split mode, so a clear never wipes tensor intent), nor is an unrelated extra
(--top-k) or inherit (None). tensor_parallel=true / --split-mode tensor re-engage
tensor. Shared by the already-loaded dedup and the load carry-forward (#6659)."""
override = parse_split_mode_override(request.llama_extra_args)
return override is not None and override.strip().lower() != "tensor"
def _request_matches_loaded_settings(
request: LoadRequest,
llama_backend: LlamaCppBackend,
effective_chat_template_override: Optional[str] = None,
) -> bool:
"""True iff every runtime setting on the request matches the loaded server.
Caller has already checked model+variant+is_loaded. See #5401.
``effective_chat_template_override`` is the resolved template that will be
launched (user override, else a bundled family template such as the
gemma-4 override), so the dedup compares against what the backend actually
holds rather than the raw request field. Defaults to the request field for
callers that do not resolve a bundled override."""
# Compare requested n_ctx (not effective) so VRAM-cap doesn't mask an
# Auto-vs-explicit slider flip.
if request.max_seq_length != llama_backend.requested_n_ctx:
return False
if _normalise_settings_str(request.cache_type_kv) != _normalise_settings_str(
llama_backend.cache_type_kv
):
return False
# Reconcile a user --split-mode in extras into the effective tensor state.
# When the request omits llama_extra_args ("inherit"), compare using the
# stored extras stripped the way the reload strips them, so an extras-driven
# tensor load isn't seen as a mismatch that needlessly reloads the server.
backend_extra = list(llama_backend.extra_args) if llama_backend.extra_args else []
effective_extra = (
request.llama_extra_args
if request.llama_extra_args is not None
else strip_shadowing_flags(
backend_extra,
strip_split_mode = _should_strip_split_mode(request, backend_extra),
)
)
if not _tensor_parallel_matches_loaded(
effective_extra, request.tensor_parallel, llama_backend.tensor_parallel
):
return False
# Preserved tensor->layer fallback (both report tensor=off, so the check above
# matches): if the user now explicitly drops tensor intent, reload so placement
# re-selects instead of keeping the all-GPU mask (#6659). The effective check
# includes the env, so an env-only tensor (LLAMA_ARG_SPLIT_MODE=tensor) that
# can't actually be dropped falls through to the env-downgrade match, not a loop.
if llama_backend.layer_preserves_tensor_intent and _is_explicit_tensor_drop(request):
return False
# Spec decoding works on vision models too (MTP is mmproj-compatible,
# llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare
# the real requested mode -- coercing vision to ``off`` here used to
# swallow every spec-mode change on a vision model as already_loaded.
req_mode = _canonicalize_spec_mode(request.speculative_type) or "auto"
backend_mode = llama_backend.requested_spec_mode or "auto"
if req_mode != backend_mode:
return False
# Prior HF load fell back with drafter_not_found: a same-settings reload must
# retry the download, not dedupe to the stale fallback. HF only (hf_repo set);
# local/native loads have no download to retry (handled by the path compare).
if (
llama_backend.hf_repo
and llama_backend.spec_fallback_reason == "drafter_not_found"
and req_mode in ("auto", "mtp", "mtp+ngram")
and not _extra_args_set_spec_type(effective_extra)
):
return False
# spec_draft_n_max only matters with an MTP variant; None means "platform
# default" and matches whatever the backend chose.
if backend_mode in ("mtp", "mtp+ngram") and request.spec_draft_n_max is not None:
if int(request.spec_draft_n_max) != (llama_backend.spec_draft_n_max or 0):
return False
_effective_cto = (
effective_chat_template_override
if effective_chat_template_override is not None
else request.chat_template_override
)
if (_effective_cto or None) != (llama_backend.chat_template_override or None):
return False
# llama_extra_args=None means "inherit"; only an explicit differing list
# forces a reload. On the inherit path, refuse to match if stored extras
# contain any shadow flag, so the reload path strips them rather than
# leaving a stale override in effect. (backend_extra computed above.)
if request.llama_extra_args is None:
# Mirror the reload's conditional split-mode strip, so a preserved
# non-tensor mode (row/none/layer) isn't seen as stale and doesn't
# trigger a needless reload of a healthy server.
if (
backend_extra
and strip_shadowing_flags(
backend_extra,
strip_split_mode = _should_strip_split_mode(request, backend_extra),
)
!= backend_extra
):
return False
else:
if list(request.llama_extra_args) != backend_extra:
return False
# A separate drafter (Gemma's root mtp-*.gguf) appearing or disappearing
# next to the loaded weights changes the launch command (--model-draft),
# so a duplicate /load must reload rather than dedupe. Always compare the
# detected vs stored drafter when the mode can use one and the user does
# not own --spec-type: the resolved-path compare is cheap and handles all
# four cases (both None -> match; one None -> reload; equal -> match;
# different -> reload), including a drafter deleted out from under a
# running server. Runs last: it stats the filesystem, so every pure-memory
# comparison above short-circuits first. Resolve both sides since the
# stored launch path may be a snapshot symlink while detect_mtp_file
# returns the resolved blob.
if req_mode in ("auto", "mtp", "mtp+ngram") and llama_backend.gguf_path:
effective_extras = (
request.llama_extra_args
if request.llama_extra_args is not None
else llama_backend.extra_args
)
if not _extra_args_set_spec_type(effective_extras):
detected = detect_mtp_file(llama_backend.gguf_path)
stored = llama_backend.mtp_draft_path
try:
detected_resolved = Path(detected).resolve() if detected else None
stored_resolved = Path(stored).resolve() if stored else None
except OSError:
return False
if detected_resolved != stored_resolved:
return False
return True
def _resolve_model_identifier_for_request(
request: LoadRequest | ValidateModelRequest, *, operation: str
) -> tuple[str, str, bool]:
if not request.native_path_lease:
return request.model_path, request.model_path, False
try:
grant = verify_native_path_lease(
request.native_path_lease,
operation = operation,
expected_kind = "model",
expected_path_type = "file",
allowed_suffixes = (".gguf",),
)
except NativePathLeaseError as exc:
# Curated, client-correctable lease error (expired / wrong type / re-select);
# keep the actionable message, just redact paths.
logger.warning("inference.native_path_lease_failed: %s", exc)
raise HTTPException(
status_code = 400,
detail = redact_native_paths(str(exc)),
) from exc
display_label = grant.display_label or Path(request.model_path).name or "Native model"
return str(grant.canonical_path), display_label, True
# GGUF inference backend (llama-server)
_llama_cpp_backend = LlamaCppBackend()
def get_llama_cpp_backend() -> LlamaCppBackend:
return _llama_cpp_backend
# Serializes opt-in auto-switch loads so two requests can't race a swap. One
# lock per running loop, since a module-level asyncio.Lock binds to a single
# loop and breaks multi-loop runners (e.g. pytest's per-test loops on pre-3.10).
_auto_switch_locks: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
_auto_switch_locks_guard = threading.Lock()
def _auto_switch_lock() -> asyncio.Lock:
loop = asyncio.get_running_loop()
# WeakKeyDictionary mutation isn't thread-safe; guard get-or-create so two
# loops on different threads can't race it.
with _auto_switch_locks_guard:
lock = _auto_switch_locks.get(loop)
if lock is None:
lock = _auto_switch_locks[loop] = asyncio.Lock()
return lock
# Process-wide gate so a swap on another event loop in this process can't race
# this one for the single model slot: the asyncio lock above is per loop, but the
# backend slot and _load_model_impl are process-wide. threading.Lock so it serializes
# across loops/threads; released from the loop thread (Lock allows cross-thread release).
_auto_switch_process_lock = threading.Lock()
async def _acquire_swap_gate() -> None:
# Non-blocking first for the common single-loop case; otherwise poll off a
# short sleep rather than awaiting to_thread(acquire). A cancelled to_thread
# (client disconnect mid-wait) leaves its worker thread still acquiring, so the
# gate gets taken but the finally that releases it never runs -- deadlocking
# later swaps. Polling keeps the wait off this loop AND cancellation-safe: a
# cancel lands during the sleep, when the gate is not held.
while not _auto_switch_process_lock.acquire(blocking = False):
await asyncio.sleep(0.02)
# Counts in-flight auto-switch requests per (target, variant). The busy guard
# subtracts same-target waiters so concurrent requests for one model load once
# instead of each 409-ing the other.
_auto_switch_waiters: dict[tuple[str, str], int] = {}
_auto_switch_waiters_guard = threading.Lock()
def _switch_key(override_id: str, variant: Optional[str]) -> tuple[str, str]:
return (override_id.lower(), (variant or "").lower())
def _note_switch_waiter(key: tuple[str, str], delta: int) -> None:
with _auto_switch_waiters_guard:
n = _auto_switch_waiters.get(key, 0) + delta
if n > 0:
_auto_switch_waiters[key] = n
else:
_auto_switch_waiters.pop(key, None)
def _same_target_waiters(key: tuple[str, str]) -> int:
with _auto_switch_waiters_guard:
return _auto_switch_waiters.get(key, 0)
# A second waiter map keyed by the raw requested model, registered before the
# (slow) resolve. The middleware counts a concurrent same-model request as
# in-flight before it resolves and joins _auto_switch_waiters, so without this
# the first request would see it as an unrelated request and 409.
_auto_switch_request_waiters: dict[str, int] = {}
_auto_switch_request_waiters_guard = threading.Lock()
def _request_waiter_key(requested_model: str) -> str:
return requested_model.strip().lower()
def _note_request_waiter(key: str, delta: int) -> None:
with _auto_switch_request_waiters_guard:
n = _auto_switch_request_waiters.get(key, 0) + delta
if n > 0:
_auto_switch_request_waiters[key] = n
else:
_auto_switch_request_waiters.pop(key, None)
def _same_request_waiters(key: str) -> int:
with _auto_switch_request_waiters_guard:
return _auto_switch_request_waiters.get(key, 0)
def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]:
"""The id to report for the loaded GGUF in API responses: the advertised repo
id from an auto-switch load, else the cleaned public id, never the on-disk
.gguf path (see core.inference.model_ids.public_model_id)."""
return (
getattr(llama_backend, "_openai_advertised_id", None)
or public_model_id(getattr(llama_backend, "model_identifier", None))
or public_model_id(fallback)
or fallback
)
_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY = "_unsloth_disable_openai_auto_switch"
# Sentinel a raw-body endpoint passes when the request omits ``model``: it must
# only restore an idle-freed model, never run the resolver (so a downloaded GGUF
# literally named "default" can't be swapped to). The NUL keeps it off any index.
_RELOAD_ONLY_MODEL = "\x00reload-only"
def _switch_model_for_payload(payload) -> str:
# A pydantic request fills an omitted ``model`` with "default"; only an
# explicitly set model may switch, else reload-only so a GGUF named "default"
# is never matched (mirrors the raw-body sentinel path).
return payload.model if "model" in payload.model_fields_set else _RELOAD_ONLY_MODEL
def _target_is_vision(load_path: str) -> bool:
# A local GGUF's vision capability is its companion mmproj, a filesystem check
# (no model load). Matches the loaded backend's is_vision, so rejecting a swap
# here can't differ from the post-load guard. Thread the ambient HF token so the
# probe keeps the capability-probe invariant (the resolver only yields local
# paths, where the token is unused, but the rule requires it regardless).
from utils.models.model_config import is_vision_model
try:
return bool(is_vision_model(load_path, hf_token = os.environ.get("HF_TOKEN")))
except Exception as exc:
# Detection failure: don't block the swap, let the load decide.
logger.debug("auto-switch: vision probe failed for %s: %s", load_path, exc)
return True
def _messages_have_image(messages) -> bool:
return any(
isinstance(m.content, list) and any(isinstance(p, ImageContentPart) for p in m.content)
for m in messages
)
def _request_has_image(payload) -> bool:
if getattr(payload, "image_base64", None):
return True
return _messages_have_image(payload.messages)
def _anthropic_request_has_image(payload) -> bool:
# Mirror anthropic_messages_to_openai: an Anthropic image block carries
# ``type == "image"`` (typed AnthropicImageBlock or a raw dict).
for msg in getattr(payload, "messages", None) or []:
content = getattr(msg, "content", None)
if not isinstance(content, list):
continue
for block in content:
bt = block.get("type") if isinstance(block, dict) else getattr(block, "type", None)
if bt == "image":
return True
return False
def disable_openai_auto_switch_for_request(scope) -> None:
"""Opt a request out of OpenAI auto-switch. The public preview route uses this:
it always serves its pinned checkpoint, so a caller-supplied model must never
swap the loaded model."""
if isinstance(scope, dict):
scope[_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY] = True
def _automatic_model_load_may_run() -> bool:
"""True when a request can trigger an automatic load: either resolver-based
auto-switch is on, or a standalone idle TTL can reload an idle-freed model. The
validate-before-switch guards key off this so an invalid request never loads."""
from utils.openai_auto_switch_settings import (
get_openai_auto_switch_enabled,
get_auto_unload_idle_seconds,
)
return get_openai_auto_switch_enabled() or get_auto_unload_idle_seconds() > 0
async def _maybe_auto_switch_model(
requested_model: Optional[str],
fastapi_request: Request,
current_subject: str,
*,
require_vision: bool = False,
) -> None:
"""Load a downloaded local GGUF named by an OpenAI request when auto-switch is on.
No-op unless enabled and ``requested_model`` resolves to a downloaded local
model different from the loaded one. Unknown names fall through (drop-in
compat) and no remote download is triggered. ``require_vision`` rejects a swap
to a text-only target before it runs, so an image request can't evict the
resident vision model only to 400 afterwards.
"""
from utils.openai_auto_switch_settings import (
get_openai_auto_switch_enabled,
get_auto_unload_idle_seconds,
get_model_override,
)
from core.inference.local_model_resolver import resolve_local_gguf
from core.inference.llama_keepwarm import (
get_last_unloaded_model,
other_inference_request_count,
inference_lifecycle_gate,
)
# Treat a non-string model (e.g. {"model": 123} on a raw-body endpoint) as
# absent so it falls through instead of raising in the membership checks below.
if not isinstance(requested_model, str) or not requested_model:
return
# The public preview route opts out so a caller cannot switch away from the
# pinned preview checkpoint it just loaded.
scope = getattr(fastapi_request, "scope", None)
if isinstance(scope, dict) and scope.get(_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY):
return
auto_switch_on = get_openai_auto_switch_enabled()
# The reload-stash path also runs when idle-unload is active on its own (a
# standalone UNSLOTH_MODEL_IDLE_TTL with auto-switch off), so a model the idle
# loop freed is restored on the next request. The resolver-based switch still
# requires the auto-switch toggle.
if not auto_switch_on and get_auto_unload_idle_seconds() <= 0:
return
# Register by the raw requested model before resolving (which can be slow):
# the middleware already counts a concurrent same-model request as in-flight,
# so the busy guard must know it shares this target even while it resolves.
request_key = _request_waiter_key(requested_model)
_note_request_waiter(request_key, 1)
try:
# Off the loop: a cold-cache rebuild walks several model dirs + HF caches.
# With auto-switch off (or an omitted-model reload-only request), skip the
# resolve so only the reload-stash path runs and no name is ever matched.
reload_only = requested_model == _RELOAD_ONLY_MODEL
resolved = (
await asyncio.to_thread(resolve_local_gguf, requested_model)
if auto_switch_on and not reload_only
else None
)
if resolved is None:
# Idle-unload may have freed the model; reload exactly what it freed
# (path + quant + advertised id) so an alias/unknown name stays servable
# and keeps the override keyed by the advertised id, not the load path.
last = get_last_unloaded_model()
# A non-GGUF (Unsloth/Transformers) model loaded after the idle-unload
# leaves the GGUF slot empty but is the live model, so don't resurrect
# the stale GGUF over it (that load would tear the active model down).
if (
not last
or get_llama_cpp_backend().is_loaded
or getattr(get_inference_backend(), "active_model_name", None)
):
return
if len(last) == 3:
target_id, variant, override_id = last
else: # pre-3-tuple stash: fall back to the path as the override key
target_id, variant = last
override_id = target_id
else:
# load_path is a concrete local path (never the bare repo id), so /load
# takes the local branch and cannot trigger a download. override_id is the
# advertised repo id, the launch-override key and the public model id.
target_id, variant, override_id = resolved
backend = get_llama_cpp_backend()
# A bare model id (no :VARIANT) is satisfied by any loaded quant of that
# repo, so it never reloads a different local quant that already serves it.
bare = ":" not in requested_model
def _already_serving() -> bool:
# Match against both the concrete load path and the advertised repo id,
# so a model loaded manually by repo id (identifier = repo id) and one
# loaded by auto-switch (identifier = path, advertised = repo id) both
# count as already serving rather than triggering a needless reswap.
if not backend.is_loaded or not backend.model_identifier:
return False
loaded_keys = {backend.model_identifier.lower()}
advertised = getattr(backend, "_openai_advertised_id", None)
if advertised:
loaded_keys.add(advertised.lower())
if loaded_keys.isdisjoint({target_id.lower(), override_id.lower()}):
return False
if bare:
return True
if variant:
loaded_variant = (getattr(backend, "hf_variant", None) or "").lower()
return loaded_variant == variant.lower()
return True
def _record_serving_alias() -> None:
# When an advertised alias already resolves to the loaded model (e.g. a
# model loaded by local path, requested by its repo/LM Studio id), record
# the alias as the public id so /v1/models and responses report it (and
# mark it loaded) instead of the path-derived basename. Resolver branch
# only: the reload-stash override_id can be the bare path, not a repo id.
# Lock-free is safe here: an in-flight request blocks any concurrent swap
# (single-slot busy guard), so the loaded model can't change under this.
if resolved is None or not override_id:
return
b = get_llama_cpp_backend()
if getattr(b, "_openai_advertised_id", None) != override_id:
b._openai_advertised_id = override_id
if _already_serving():
_record_serving_alias()
return
# An image/audio request naming a different text-only GGUF would load it
# here and only 400 below, evicting the working model. Reject before the
# swap. Only the resolver branch (an explicit new target); the reload-stash
# path just restores the model the request was already using. Both vision and
# audio input come from a companion mmproj (a filesystem probe) -- run it off
# the loop, like the resolver above.
if (
require_vision
and resolved is not None
and not await asyncio.to_thread(_target_is_vision, target_id)
):
raise HTTPException(
status_code = 400,
detail = openai_error_body(
"The requested model does not support the image or audio input in this request.",
status = 400,
code = "invalid_value",
param = "model",
),
)
key = _switch_key(override_id, variant)
_note_switch_waiter(key, 1)
try:
async with _auto_switch_lock():
# The asyncio lock is per loop; add a process-wide gate so a swap on
# another loop in this process can't race the single slot.
await _acquire_swap_gate()
try:
# Hold the keep-warm gate across the swap so no new inference can
# start on the model while it is being torn down and replaced.
async with inference_lifecycle_gate():
if _already_serving():
_record_serving_alias()
return
# Single slot: refuse a cross-model swap while another inference
# request is active rather than killing its response. Requests
# heading to this same target (by resolved id or raw name) are
# excluded, so concurrent requests for one model load once. A
# pending request is still in the middleware, not generating, so
# it is not counted here.
same_others = max(
_same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0
)
others = other_inference_request_count(
current_request_counted = True, include_pending = False
)
# Not gated on the GGUF being loaded: _load_model_impl also
# tears down an active Unsloth backend before loading a GGUF,
# so refuse whenever any other inference request is in flight.
if others > same_others:
raise HTTPException(
status_code = 409,
detail = openai_error_body(
"Cannot switch models while another inference request is in progress.",
status = 409,
code = "model_switch_busy",
param = "model",
),
)
# Apply this model's saved launch flags so the swap honors the config.
override = get_model_override(override_id)
load_kwargs = {"model_path": target_id, "gguf_variant": variant}
if override.get("llama_extra_args") is not None:
load_kwargs["llama_extra_args"] = override["llama_extra_args"]
if override.get("max_seq_length") is not None:
load_kwargs["max_seq_length"] = override["max_seq_length"]
# Reuse the load impl so its dedup, tensor fallback, and threading
# apply. Call the impl directly: we already hold the lifecycle gate
# the /load route would otherwise take, so the route would deadlock.
await _load_model_impl(
LoadRequest(**load_kwargs),
fastapi_request,
current_subject,
)
# Advertise the repo id (not the concrete load path) as the loaded
# model's public id and override key for /v1/models and idle stash.
get_llama_cpp_backend()._openai_advertised_id = override_id
finally:
_auto_switch_process_lock.release()
finally:
_note_switch_waiter(key, -1)
finally:
_note_request_waiter(request_key, -1)
async def _auto_switch_from_request_body(request: Request, current_subject: str):
"""Run auto-switch from a raw-body endpoint's ``model`` without changing its
pre-feature status codes: a malformed/non-dict body yields no model (so an
unloaded backend still 503s, not 500), and the caller re-reads to surface the
original parse error after the loaded-state check. Returns the parsed body, or
None if it could not be parsed."""
try:
body = await request.json()
except (json.JSONDecodeError, ValueError):
return None
if isinstance(body, dict):
# A raw-body client may omit ``model`` and rely on the loaded backend. Pass
# a reload-only sentinel so the idle-stash reload still runs (an idle-freed
# model is restored) without the resolver ever matching a real name.
model = body.get("model") or _RELOAD_ONLY_MODEL
else:
model = None
await _maybe_auto_switch_model(model, request, current_subject)
return body
def _effective_load_in_4bit(config: ModelConfig, requested: bool) -> bool:
"""Effective quantization the loader will use: a LoRA adapter can flip 4-bit to
16-bit via adapter_config.json, so the guard sizes this, not the raw request."""
load_in_4bit = requested
if not getattr(config, "is_lora", False) or not getattr(config, "path", None):
return load_in_4bit
adapter_cfg_path = Path(config.path) / "adapter_config.json"
if not adapter_cfg_path.exists():
return load_in_4bit
try:
with open(adapter_cfg_path) as f:
adapter_cfg = json.load(f)
if not isinstance(adapter_cfg, dict): # malformed -> keep requested
return load_in_4bit
except Exception as e:
logger.warning(f"Could not read adapter_config.json: {e}")
return load_in_4bit
training_method = adapter_cfg.get("unsloth_training_method")
if training_method == "lora":
return False
if training_method == "qlora":
return True
if not training_method and config.base_model and "-bnb-4bit" not in config.base_model.lower():
return False
return load_in_4bit
def _remote_gguf_companion_bytes(
repo: str, *, hf_token: Optional[str], include_mmproj: bool
) -> int:
"""Bytes of MTP/mmproj companion GGUFs llama-server auto-downloads. 0 on error,
so it can only add headroom, never refuse a load by itself."""
try:
from huggingface_hub import model_info
info = model_info(repo, token = hf_token, files_metadata = True)
total = 0
for sibling in info.siblings or []:
name = sibling.rfilename or ""
base = Path(name).name.lower()
if not base.endswith(".gguf"):
continue
# Root-level mtp- only: -hf auto-fetches the repo-root drafter, not
# the MTP/ subdir copies (which now share the mtp- prefix too).
is_root_mtp = "/" not in name and base.startswith("mtp-")
if is_root_mtp or (include_mmproj and "mmproj" in base):
total += getattr(sibling, "size", 0) or 0
return total
except Exception as e:
logger.warning(f"Could not size GGUF companions for {repo}: {e}")
return 0
def _estimate_gguf_kv_gb(
gguf_path: str,
max_seq_length: int,
llama_extra_args: Optional[list[str]] = None,
n_parallel: int = 1,
) -> float:
"""KV-cache VRAM (GB) at the larger of max_seq_length and any `--ctx-size`/`-c`
override, over n_parallel slots, with the default f16 cache so the estimate is
never below what the server allocates. 0 if metadata is unreadable."""
try:
from core.inference.llama_server_args import parse_ctx_override
probe = LlamaCppBackend()
probe._read_gguf_metadata(gguf_path)
if not probe._can_estimate_kv():
return 0.0
try:
ctx_override = parse_ctx_override(llama_extra_args) or 0
except Exception:
ctx_override = 0 # malformed extras are rejected upstream; fall back
ctx = max(max_seq_length or 0, ctx_override) or (probe._context_length or 0)
if ctx <= 0:
return 0.0
kv = probe._estimate_kv_cache_bytes(ctx, n_parallel = max(1, n_parallel or 1))
return kv / (1024**3)
except Exception as e:
logger.warning(f"Could not size GGUF KV cache for training guard: {e}")
return 0.0
def _estimate_gguf_required_gb(
config: ModelConfig,
hf_token: Optional[str] = None,
max_seq_length: int = 0,
llama_extra_args: Optional[list[str]] = None,
n_parallel: int = 1,
) -> Optional[float]:
"""Approximate GGUF VRAM (GB): quantized weights + companions, plus the KV
cache for local files (unreadable pre-download for remote). None when nothing
resolves so the caller default-denies."""
try:
total_bytes = 0
main = getattr(config, "gguf_file", None)
if main and Path(main).is_file():
total_bytes += LlamaCppBackend._get_gguf_size_bytes(str(main))
for attr in ("gguf_mmproj_file", "gguf_mtp_file"):
f = getattr(config, attr, None)
if f and Path(f).is_file():
total_bytes += Path(f).stat().st_size
if total_bytes > 0:
return total_bytes / (1024**3) + _estimate_gguf_kv_gb(
main, max_seq_length, llama_extra_args, n_parallel
)
repo = getattr(config, "gguf_hf_repo", None)
variant = getattr(config, "gguf_variant", None)
if repo and variant:
from utils.models.model_config import list_gguf_variants
variants, has_vision = list_gguf_variants(repo, hf_token = hf_token)
main_bytes = next(
(v.size_bytes for v in variants if v.quant.lower() == variant.lower()), None
)
if main_bytes is None:
return None
companions = _remote_gguf_companion_bytes(
repo, hf_token = hf_token, include_mmproj = bool(has_vision)
)
return (main_bytes + companions) / (1024**3)
return None
except Exception as e:
logger.warning(f"Could not size GGUF model for training guard: {e}")
return None
def _guard_chat_load_against_training(
config: ModelConfig,
*,
model_identifier: str,
hf_token: Optional[str],
load_in_4bit: bool,
max_seq_length: int,
requested_gpu_ids: Optional[List[int]],
llama_extra_args: Optional[list[str]] = None,
n_parallel: int = 1,
) -> None:
"""Refuse loading a local chat model that would OOM an active training run.
No-op when training is inactive or unknown. `load_in_4bit` must be the
effective quantization (see _effective_load_in_4bit). Raises HTTP 409 when the
model would not fit alongside training."""
from core.training import get_training_backend
from routes.training_vram import can_load_chat_during_training
try:
if not get_training_backend().is_training_active():
return
except Exception as e:
logger.warning("Could not check training state for chat-load guard: %s", e)
return
is_gguf = bool(getattr(config, "is_gguf", False))
required_override_gb = (
_estimate_gguf_required_gb(
config,
hf_token = hf_token,
max_seq_length = max_seq_length,
llama_extra_args = llama_extra_args,
n_parallel = n_parallel,
)
if is_gguf
else None
)
ok, info = can_load_chat_during_training(
model_name = model_identifier,
hf_token = hf_token,
load_in_4bit = load_in_4bit,
max_seq_length = max_seq_length,
requested_gpu_ids = requested_gpu_ids,
is_gguf = is_gguf,
required_override_gb = required_override_gb,
)
if ok:
return
usable = info.get("usable_gb")
needed = info.get("needed_gb")
if needed is None:
needed = info.get("required_gb")
if needed is not None and usable is not None:
detail = (
f"Not enough free GPU memory to load this model while training is "
f"running (needs ~{needed:.0f} GB including safety headroom, "
f"~{usable:.0f} GB free). Training was left untouched. Use an external "
f"provider, a smaller or more quantized model, or try again after "
f"training finishes."
)
else:
detail = (
"Can't load this model while training is running: its GPU memory use "
"could not be verified, so the load was refused to protect the "
"training run. Use an external provider or try again after training "
"finishes."
)
logger.info("Refusing chat-model load during training: %s", info)
raise HTTPException(status_code = 409, detail = detail)
def _model_json_response(model, status_code: int = 200) -> Response:
"""Serialize a pydantic response once via pydantic-core.
Equivalent body to ``JSONResponse(content = model.model_dump())`` but
avoids the dict round-trip plus Starlette's second ``json.dumps``.
"""
return Response(
content = model.model_dump_json(),
media_type = "application/json",
status_code = status_code,
)
_NOT_SUPPORTED_HINTS = (
"No config file found",
"not yet supported",
"is not supported",
"does not support",
)
def _maybe_unsupported_message(msg: str) -> str:
"""Rewrite a load/validate error into the friendly "not supported yet"
message when it matches a known unsupported-model signature; otherwise
return ``msg`` unchanged."""
if any(h.lower() in msg.lower() for h in _NOT_SUPPORTED_HINTS):
return f"This model is not supported yet. Try a different model. (Original error: {msg})"
return msg
@router.post("/load", response_model = LoadResponse)
async def load_model(
request: LoadRequest,
fastapi_request: Request,
current_subject: str = Depends(get_current_subject),
):
"""
Load a model for inference.
model_path is a clean identifier from GET /models/list. Returns inference
config (temperature, top_p, top_k, min_p) from the model's YAML, falling
back to default.yaml for missing values.
GGUF models load via llama-server (llama.cpp) instead of Unsloth.
"""
# Hold the lifecycle gate across the load so idle auto-unload can't unload the
# model mid-load. Auto-switch calls _load_model_impl directly since it already
# holds this gate.
from core.inference.llama_keepwarm import inference_lifecycle_gate
async with inference_lifecycle_gate():
return await _load_model_impl(request, fastapi_request, current_subject)
async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str):
from core.inference.llama_cpp import LlamaServerNotFoundError
native_grant_backed = False
model_log_label = request.model_path
try:
# Validate user pass-through args up front so a managed-flag collision
# returns 400 before any model work.
try:
extra_llama_args = validate_extra_args(request.llama_extra_args)
except ValueError as exc:
# Keep the curated validation message (names the flag); just strip paths.
logger.warning("inference.validate_extra_args_failed: %s", exc)
raise HTTPException(
status_code = 400,
detail = redact_native_paths(str(exc)),
)
# Re-narrow []-from-None back to None so the inheritance path below can
# tell "caller omitted" from "caller explicit []".
extra_llama_args: Optional[list[str]] = (
None if request.llama_extra_args is None else extra_llama_args
)
model_identifier, model_log_label, native_grant_backed = (
_resolve_model_identifier_for_request(request, operation = "load-model")
)
# Version switching is handled by the subprocess-based inference
# backend -- no ensure_transformers_version() needed here.
# Resolve the effective chat-template override once, up front: an
# explicit user override, else a bundled family template (e.g. the
# gemma-4 override that ships preserve_thinking without re-downloading
# quants), else None. Used for both the reload-dedup check below and the
# load_model calls, so the live backend state and the incoming request
# compare against the same template text.
effective_chat_template_override = resolve_effective_chat_template_override(
model_identifier = model_identifier,
user_override = request.chat_template_override,
)
# ── Already-loaded check: skip reload if the exact model is active ──
backend = get_inference_backend()
llama_backend = get_llama_cpp_backend()
is_direct_gguf_request = model_identifier.lower().endswith(".gguf")
if request.gguf_variant or is_direct_gguf_request:
gguf_variant_matches = is_direct_gguf_request or bool(
llama_backend.hf_variant
and request.gguf_variant
and llama_backend.hf_variant.lower() == request.gguf_variant.lower()
)
if (
llama_backend.is_loaded
and gguf_variant_matches
and llama_backend.model_identifier
and llama_backend.model_identifier.lower() == model_identifier.lower()
# Match runtime settings so Apply isn't dropped (#5401).
and _request_matches_loaded_settings(
request,
llama_backend,
effective_chat_template_override,
)
# Skip if a prior audio probe failed -- let load_model retry.
and getattr(llama_backend, "_audio_probed", True)
):
logger.info(
"Model already loaded (GGUF): "
f"{model_log_label} variant={request.gguf_variant or llama_backend.hf_variant}, skipping reload"
)
inference_config = load_inference_config(llama_backend.model_identifier)
_gguf_audio = getattr(llama_backend, "_audio_type", None)
_gguf_is_audio = getattr(llama_backend, "_is_audio", False)
return LoadResponse(
status = "already_loaded",
model = model_log_label
if native_grant_backed
else llama_backend.model_identifier,
display_name = model_log_label
if native_grant_backed
else llama_backend.model_identifier,
is_vision = llama_backend._is_vision,
is_lora = False,
is_gguf = True,
is_diffusion = llama_backend.is_diffusion,
is_audio = _gguf_is_audio,
audio_type = _gguf_audio,
has_audio_input = getattr(llama_backend, "_has_audio_input", False),
inference = inference_config,
# GGUF loads via llama.cpp: auto_map never executes, so inert (matches validate_model).
requires_trust_remote_code = False,
context_length = llama_backend.context_length,
max_context_length = llama_backend.max_context_length,
native_context_length = llama_backend.native_context_length,
supports_reasoning = llama_backend.supports_reasoning,
reasoning_style = llama_backend.reasoning_style,
reasoning_effort_levels = llama_backend.reasoning_effort_levels,
reasoning_always_on = llama_backend.reasoning_always_on,
supports_preserve_thinking = llama_backend.supports_preserve_thinking,
supports_tools = llama_backend.supports_tools,
chat_template = llama_backend.chat_template,
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
tensor_parallel = llama_backend.tensor_parallel,
)
else:
if (
backend.active_model_name
and backend.active_model_name.lower() == model_identifier.lower()
):
logger.info(f"Model already loaded (Unsloth): {model_log_label}, skipping reload")
inference_config = load_inference_config(backend.active_model_name)
_model_info = backend.models.get(backend.active_model_name, {})
_chat_template = None
try:
_tpl_info = _model_info.get("chat_template_info", {})
_chat_template = _tpl_info.get("template")
except Exception as e:
logger.warning(
f"Could not retrieve chat template for {backend.active_model_name}: {e}"
)
# Classify via the same path as GGUF.
_sf_flags = _detect_safetensors_features(backend, _chat_template)
_sf_supports_reasoning = _sf_flags["supports_reasoning"]
_sf_reasoning_style = _sf_flags["reasoning_style"]
return LoadResponse(
status = "already_loaded",
model = model_log_label if native_grant_backed else backend.active_model_name,
display_name = model_log_label
if native_grant_backed
else backend.active_model_name,
is_vision = _model_info.get("is_vision", False),
is_lora = _model_info.get("is_lora", False),
is_gguf = False,
is_audio = _model_info.get("is_audio", False),
audio_type = _model_info.get("audio_type"),
has_audio_input = _model_info.get("has_audio_input", False),
inference = inference_config,
requires_trust_remote_code = _resolve_loaded_trust_remote_code(
backend.active_model_name, _model_info, inference_config
),
supports_reasoning = _sf_supports_reasoning,
reasoning_style = _sf_reasoning_style,
reasoning_effort_levels = _sf_flags.get("reasoning_effort_levels", []),
reasoning_always_on = _sf_flags["reasoning_always_on"],
supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
supports_tools = _sf_flags["supports_tools"],
context_length = _positive_int_or_none(_model_info.get("context_length")),
chat_template = _chat_template,
)
# is_lora auto-detected from adapter_config.json on disk/HF.
# DNS-probe wrap so offline loads skip 30-60s of soft-failed network
# checks before the worker starts.
with _hf_offline_if_dns_dead():
config = ModelConfig.from_identifier(
model_id = model_identifier,
hf_token = request.hf_token,
gguf_variant = request.gguf_variant,
)
if not config:
raise HTTPException(
status_code = 400,
detail = f"Invalid model identifier: {model_log_label}",
)
# Normalize gpu_ids: empty list means auto-selection, same as None
effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
# Reject GGUF + gpu_ids first so the guard can't mask it with a VRAM 409.
if config.is_gguf and effective_gpu_ids is not None:
raise HTTPException(
status_code = 400,
detail = "gpu_ids is not supported for GGUF models yet.",
)
if not config.is_gguf and _mlx_distributed_launch_detected():
raise HTTPException(
status_code = 400,
detail = (
"Studio does not support distributed MLX inference under "
"mlx.launch. Use `mlx.launch ... unsloth chat` or run Studio "
"without the distributed launcher."
),
)
# Effective quantization (LoRA can flip 4-bit -> 16-bit); guard + load reuse it.
effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit)
if effective_load_in_4bit != request.load_in_4bit:
logger.info(
f"Resolved load_in_4bit={effective_load_in_4bit} for '{model_log_label}' "
f"from adapter_config.json / base model (requested {request.load_in_4bit})"
)
# Refuse a load that would OOM active training, before the unload step below
# frees the resident model. Off-loop: guard does sync nvidia-smi / HF work.
await asyncio.to_thread(
_guard_chat_load_against_training,
config,
model_identifier = model_identifier,
hf_token = request.hf_token,
load_in_4bit = effective_load_in_4bit,
max_seq_length = request.max_seq_length,
requested_gpu_ids = effective_gpu_ids,
llama_extra_args = extra_llama_args,
n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1),
)
# ── GGUF path: load via llama-server ──────────────────────
if config.is_gguf:
llama_backend = get_llama_cpp_backend()
unsloth_backend = get_inference_backend()
# Unload any active Unsloth model to free VRAM (off the event loop:
# unload takes _gen_lock and can wait on an in-flight stream).
if unsloth_backend.active_model_name:
logger.info(
f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF"
)
await asyncio.to_thread(
unsloth_backend.unload_model, unsloth_backend.active_model_name
)
# Inherit llama_extra_args from the previous load when the request
# omits the field (the chat-settings Apply path doesn't round-trip
# them; explicit [] still clears). Gated on (model_identifier,
# hf_variant) to refuse cross-model pickup, and shadowing flags are
# stripped so an inherited override can't win the last-wins CLI
# parse against a freshly-supplied first-class field.
if request.llama_extra_args is None and llama_backend.extra_args:
source = llama_backend.extra_args_source
# Compare against the resolved variant, not the request
# field: callers commonly omit gguf_variant for local
# ``.gguf`` paths and HF auto-pick flows. ``config.gguf_
# variant`` is the variant load_model was actually
# invoked with (see the HF / local branches below), so
# both sides of the comparison key off the same string.
resolved_variant = (config.gguf_variant or "").lower()
request_variant = (request.gguf_variant or "").lower()
stored_variant = (source[1] or "").lower() if source else ""
same_model = bool(
source and source[0] and source[0].lower() == model_identifier.lower()
)
if request.gguf_variant:
variant_mismatch = request_variant != stored_variant
else:
variant_mismatch = bool(stored_variant and resolved_variant != stored_variant)
same_source = same_model and not variant_mismatch
if not same_source:
logger.info(
"Not inheriting llama_extra_args: stored args came from %s, loading %s",
source,
(model_identifier, resolved_variant),
)
# Cross-model: clear explicitly so the backend doesn't
# inherit via "no opinion" semantics.
extra_llama_args = []
else:
# Strip only the groups whose first-class field was set by
# the caller, so an inherited --chat-template-file survives
# an Apply that omits chat_template_override. A bundled family
# template (e.g. the gemma-4 override) is an effective
# first-class template setting even when the raw request
# omits chat_template_override, so strip the inherited
# --chat-template-file in that case too -- otherwise the stale
# extra arg (appended last) shadows the bundled template while
# Studio reports the bundled template's capabilities.
fields_set = getattr(request, "model_fields_set", set())
stripped = strip_shadowing_flags(
llama_backend.extra_args,
strip_context = "max_seq_length" in fields_set,
strip_cache = "cache_type_kv" in fields_set,
strip_spec = (
"speculative_type" in fields_set or "spec_draft_n_max" in fields_set
),
strip_template = (
"chat_template_override" in fields_set
or effective_chat_template_override is not None
),
strip_split_mode = _should_strip_split_mode(
request, llama_backend.extra_args
),
)
try:
extra_llama_args = validate_extra_args(stripped)
except ValueError:
# Shouldn't happen on already-validated args; degrade to
# no-extras rather than 400 if managed flags changed.
logger.warning(
"Stored llama_extra_args failed revalidation; loading without them: %s",
stripped,
)
extra_llama_args = []
else:
if extra_llama_args:
logger.info(
"Inheriting llama_extra_args from previous "
"load (same model, shadow-stripped): %s",
extra_llama_args,
)
# Route to HF or local mode based on config. Run in a thread so the
# event loop stays free for progress polling and other requests
# during the (potentially long) GGUF download + llama-server start.
_n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1)
# Load kwargs common to HF and local modes; the two differ only by
# the model-source args (hf_repo/-token vs gguf_path/mmproj).
_common_load_kwargs = dict(
model_identifier = config.identifier,
is_vision = config.is_vision,
n_ctx = request.max_seq_length,
chat_template_override = effective_chat_template_override,
cache_type_kv = request.cache_type_kv,
speculative_type = request.speculative_type,
spec_draft_n_max = request.spec_draft_n_max,
n_parallel = _n_parallel,
)
if config.gguf_hf_repo:
# HF mode: download via huggingface_hub then start llama-server
_source_load_kwargs = dict(
hf_repo = config.gguf_hf_repo,
hf_variant = config.gguf_variant,
hf_token = request.hf_token,
)
else:
# Local mode: llama-server loads via -m <path>
if native_grant_backed:
if config.gguf_mmproj_file:
_validate_native_gguf_companion(
config.gguf_mmproj_file, config.gguf_file, "vision companion"
)
if config.gguf_mtp_file:
# The drafter is optional (unlike mmproj for a vision
# model): drop it rather than fail the load.
try:
_validate_native_gguf_companion(
config.gguf_mtp_file, config.gguf_file, "MTP drafter"
)
except HTTPException as exc:
logger.warning("Dropping MTP drafter for native load: %s", exc.detail)
config.gguf_mtp_file = None
_source_load_kwargs = dict(
gguf_path = config.gguf_file,
mmproj_path = config.gguf_mmproj_file,
mtp_draft_path = config.gguf_mtp_file,
# Pass the resolved variant so _extra_args_source keys off
# the same string the inheritance check at the top of /load
# uses (#5401 followup).
hf_variant = config.gguf_variant,
)
# Tensor intent for this load: the request itself, or a preserved
# multi-GPU layer fallback carried across a reload of the SAME model that
# doesn't drop it (e.g. a ctx-only change), so a fitting model doesn't
# silently collapse to one GPU. Only an explicit non-tensor --split-mode
# override counts as the drop -- the tensor field echo / unrelated extras keep
# the preserved placement; the same-model guard stops a switch-without-unload
# inheriting the prior model's intent.
_explicit_tensor_drop = _is_explicit_tensor_drop(request)
# Compare the resolved config.identifier (what load_model stores), not the
# raw request id: from_identifier normalizes shorthands (adds unsloth/, fixes
# case), so a reload with the shorthand would otherwise miss the match and
# drop the carry-forward. #6659
_same_model_loaded = (
llama_backend.is_loaded
and (llama_backend.model_identifier or "").lower()
== (config.identifier or "").lower()
)
# model_identifier is variant-agnostic for HF repos and dir-level for a
# local multi-variant directory, so also require the loaded quant to match
# (path else variant, mirroring _already_in_target_state) -- otherwise a
# different variant inherits the prior one's preserved intent. #6659
if _same_model_loaded:
if config.gguf_file and llama_backend.gguf_path:
try:
_same_model_loaded = (
Path(llama_backend.gguf_path).resolve()
== Path(config.gguf_file).resolve()
)
except OSError:
_same_model_loaded = False
else:
_same_model_loaded = (llama_backend.hf_variant or "").lower() == (
config.gguf_variant or ""
).lower()
_tensor_intent_overall = _effective_tensor_parallel(
extra_llama_args, request.tensor_parallel
) or _carry_preserved_tensor_intent(
preserved = llama_backend.layer_preserves_tensor_intent,
same_model = _same_model_loaded,
explicit_drop = _explicit_tensor_drop,
)
# Run a single load attempt with the given tensor flag + extras.
async def _attempt_gguf_load(
tensor_parallel: bool, attempt_extra_args: Optional[list[str]]
) -> bool:
attempt_kwargs = {
**_common_load_kwargs,
"extra_args": attempt_extra_args,
}
return await asyncio.to_thread(
llama_backend.load_model,
**_source_load_kwargs,
**attempt_kwargs,
tensor_parallel = tensor_parallel,
# True on the layer fallback retry (tensor wanted overall but not on
# this attempt): keep multi-GPU. Mirrors the fallback's key.
preserve_multi_gpu_on_layer = bool(
_tensor_intent_overall
and not _effective_tensor_parallel(attempt_extra_args, tensor_parallel)
),
)
# Tensor parallelism is arch-gated in llama.cpp and crashes some loads
# outright (e.g. Gemma 3n aborts with a GGML_ASSERT). The helper auto-
# falls back to layer split so the checkbox never blocks a model from
# loading; the response reports the backend's actual tensor_parallel
# state so the UI toggle reflects the fallback.
success = await load_with_tensor_fallback(
_attempt_gguf_load,
requested_tensor = request.tensor_parallel,
extra_args = extra_llama_args,
label = config.identifier,
cancelled = llama_backend.load_cancelled,
)
if not success:
raise HTTPException(
status_code = 500,
detail = f"Failed to load GGUF model: {model_log_label if native_grant_backed else config.display_name}",
)
logger.info(
f"Loaded GGUF model via llama-server: {model_log_label if native_grant_backed else config.identifier}"
)
# Clear any idle-unload reload stash now, not only on the next poll.
from core.inference.llama_keepwarm import note_model_loaded
note_model_loaded()
# A plain load advertises its own identifier; auto-switch overwrites
# this with the repo id right after _load_model_impl returns.
llama_backend._openai_advertised_id = None
# Audio detection moved into load_model under _serial_load_lock (#5642).
_gguf_audio = llama_backend._audio_type
_gguf_is_audio = llama_backend._is_audio
llama_backend._native_display_label = model_log_label if native_grant_backed else None
llama_backend._native_grant_backed = bool(native_grant_backed)
if _gguf_is_audio:
logger.info(f"GGUF model detected as audio: audio_type={_gguf_audio}")
inference_config = load_inference_config(config.identifier)
return LoadResponse(
status = "loaded",
model = model_log_label if native_grant_backed else config.identifier,
display_name = model_log_label if native_grant_backed else config.display_name,
is_vision = llama_backend.is_vision,
is_lora = False,
is_gguf = True,
is_diffusion = llama_backend.is_diffusion,
is_audio = _gguf_is_audio,
audio_type = _gguf_audio,
has_audio_input = llama_backend._has_audio_input,
inference = inference_config,
# GGUF loads via llama.cpp: auto_map never executes, so inert (matches validate_model).
requires_trust_remote_code = False,
context_length = llama_backend.context_length,
max_context_length = llama_backend.max_context_length,
native_context_length = llama_backend.native_context_length,
supports_reasoning = llama_backend.supports_reasoning,
reasoning_style = llama_backend.reasoning_style,
reasoning_effort_levels = llama_backend.reasoning_effort_levels,
reasoning_always_on = llama_backend.reasoning_always_on,
supports_preserve_thinking = llama_backend.supports_preserve_thinking,
supports_tools = llama_backend.supports_tools,
cache_type_kv = llama_backend.cache_type_kv,
chat_template = llama_backend.chat_template,
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
tensor_parallel = llama_backend.tensor_parallel,
)
# ── Standard path: load via Unsloth/transformers ──────────
backend = get_inference_backend()
# Unload any active GGUF model first
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded:
logger.info("Unloading GGUF model before loading Unsloth model")
llama_backend.unload_model()
# Shut down any export subprocess to free VRAM
try:
from core.export import get_export_backend
exp_backend = get_export_backend()
if exp_backend.current_checkpoint:
logger.info("Shutting down export subprocess to free GPU memory for inference")
exp_backend._shutdown_subprocess()
exp_backend.current_checkpoint = None
exp_backend.is_vision = False
exp_backend.is_peft = False
except Exception as e:
logger.warning("Could not shut down export subprocess: %s", e)
# Resolved before the guard so both size the same load.
load_in_4bit = effective_load_in_4bit
# Load in a thread so the event loop stays free for download progress
# polling and other requests.
success = await asyncio.to_thread(
backend.load_model,
config = config,
max_seq_length = request.max_seq_length,
load_in_4bit = load_in_4bit,
hf_token = request.hf_token,
trust_remote_code = request.trust_remote_code,
approved_remote_code_fingerprint = request.approved_remote_code_fingerprint,
gpu_ids = effective_gpu_ids,
subject = current_subject,
)
if not success:
# Check if YAML says this model needs trust_remote_code.
if not request.trust_remote_code:
model_defaults = load_model_defaults(config.identifier)
yaml_trust = model_defaults.get("inference", {}).get("trust_remote_code", False)
if yaml_trust:
raise HTTPException(
status_code = 400,
detail = (
f"Model '{config.display_name}' requires trust_remote_code to be enabled. "
f"Please enable 'Trust remote code' in Chat Settings and try again."
),
)
raise HTTPException(
status_code = 500,
detail = f"Failed to load model: {model_log_label if native_grant_backed else config.display_name}",
)
logger.info(
f"Loaded model: {model_log_label if native_grant_backed else config.identifier}"
)
# Clear any idle-unload reload stash: a manual load supersedes an idle-freed
# GGUF, so the next /v1 request must not resurrect it. Mirror the GGUF branch
# above; without this a non-GGUF load leaves a stale stash until the idle
# poll clears it (and never, while idle-unload is off).
from core.inference.llama_keepwarm import note_model_loaded
note_model_loaded()
# Load inference configuration parameters
inference_config = load_inference_config(config.identifier)
# Get chat template from tokenizer
_chat_template = None
try:
_model_info = backend.models.get(config.identifier, {})
_tpl_info = _model_info.get("chat_template_info", {})
_chat_template = _tpl_info.get("template")
except Exception:
pass
# Classify reasoning/tool flags via the GGUF sniffer.
_sf_flags = _detect_safetensors_features(backend, _chat_template)
# Report validate_model's requirement (raw auto_map OR YAML) plus the value the
# load used, and persist it, so a later retry/rollback doesn't send
# trust_remote_code=false for a custom-code model (and status reports it too).
_requires_rc = _resolve_loaded_trust_remote_code(
config.identifier,
None,
inference_config,
request.hf_token,
trust_remote_code_used = bool(getattr(request, "trust_remote_code", False)),
)
try:
backend.models.setdefault(config.identifier, {})["requires_trust_remote_code"] = (
_requires_rc
)
except Exception:
pass
return LoadResponse(
status = "loaded",
model = model_log_label if native_grant_backed else config.identifier,
display_name = model_log_label if native_grant_backed else config.display_name,
is_vision = config.is_vision,
is_lora = config.is_lora,
is_gguf = False,
is_audio = config.is_audio,
audio_type = config.audio_type,
has_audio_input = config.has_audio_input,
inference = inference_config,
requires_trust_remote_code = _requires_rc,
supports_reasoning = _sf_flags["supports_reasoning"],
reasoning_style = _sf_flags["reasoning_style"],
reasoning_effort_levels = _sf_flags.get("reasoning_effort_levels", []),
reasoning_always_on = _sf_flags["reasoning_always_on"],
supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
supports_tools = _sf_flags["supports_tools"],
context_length = _positive_int_or_none(_model_info.get("context_length")),
chat_template = _chat_template,
)
except HTTPException:
raise
except ValueError as e:
if native_grant_backed:
redacted_msg = redact_native_paths(str(e))
logger.warning(
"Rejected inference selection for native model %s: %s",
model_log_label,
redacted_msg,
)
raise HTTPException(status_code = 400, detail = redacted_msg)
logger.warning("Rejected inference GPU selection: %s", e)
# User-facing validation (e.g. "Invalid gpu_ids [99]"): redact paths, keep detail.
raise HTTPException(status_code = 400, detail = redact_native_paths(str(e)))
except LlamaServerNotFoundError as e:
# Missing GGUF runtime: 400 with the install message, not a generic 500.
logger.warning("GGUF runtime missing while loading '%s': %s", model_log_label, e)
raise HTTPException(status_code = 400, detail = str(e))
except Exception as e:
# Friendlier message for models Unsloth cannot load.
if native_grant_backed:
redacted_msg = redact_native_paths(str(e))
logger.error(
"Error loading native model %s: %s",
model_log_label,
redacted_msg,
)
msg = _maybe_unsupported_message(redacted_msg)
raise HTTPException(
status_code = 500,
detail = f"Failed to load native model {model_log_label}: {msg}",
)
logger.error(f"Error loading model: {e}", exc_info = True)
msg = _maybe_unsupported_message(redact_native_paths(str(e)))
raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}")
def _requires_trust_remote_code_for_model(
model_identifier: str, hf_token: Optional[str] = None
) -> bool:
"""Whether loading this model would execute custom repo code, so the consent
dialog must run first. True if the Studio YAML default enables
``trust_remote_code`` OR the raw config declares an ``auto_map`` (Hub/local,
config.json or tokenizer_config.json). Reads raw JSON only; never imports
model code."""
from utils.inference import load_inference_config
try:
if bool(load_inference_config(model_identifier).get("trust_remote_code", False)):
return True
except Exception:
pass
try:
from utils.security.consent import _config_has_auto_map
return _config_has_auto_map(model_identifier, hf_token) is True
except Exception:
return False
def _resolve_loaded_trust_remote_code(
model_id,
model_info,
inference_config,
hf_token = None,
trust_remote_code_used = False,
) -> bool:
"""TRC requirement to report for an ALREADY-LOADED model, consistent with
``validate_model``.
``validate_model`` reports ``requires_trust_remote_code`` from
``_requires_trust_remote_code_for_model`` (YAML default OR raw ``auto_map``), but
the load / already-loaded / status responses historically reported only the YAML
default. That dropped raw-``auto_map`` models: after approving and loading one, the
response said ``false``, so the frontend stored ``false`` and a later retry/rollback
sent ``trust_remote_code=false`` and failed.
Resolution order: a value stored on the model at load time (so a status refresh does
not re-derive it) -> the trust_remote_code the load actually used -> the YAML default
-> the raw ``auto_map`` check (reads the loaded model's cached config; no network)."""
stored = (model_info or {}).get("requires_trust_remote_code")
if stored is not None:
return bool(stored)
if trust_remote_code_used or bool((inference_config or {}).get("trust_remote_code", False)):
return True
try:
return bool(_requires_trust_remote_code_for_model(model_id, hf_token))
except Exception:
return False
def _requires_security_review_for_model(
model_identifier: str, hf_token: Optional[str] = None
) -> bool:
"""Whether Hugging Face's security scan flagged unsafe files for this repo, so
the consent dialog must open as a hard block before loading. Metadata-only;
never downloads the flagged files. Fails open (False) on any error."""
try:
from utils.security import evaluate_file_security, security_load_subdirs
return evaluate_file_security(
model_identifier,
hf_token,
load_subdirs = security_load_subdirs(model_identifier, hf_token),
).blocked
except Exception:
return False
@router.post("/validate", response_model = ValidateModelResponse)
async def validate_model(
request: ValidateModelRequest, current_subject: str = Depends(get_current_subject)
):
"""
Lightweight validation endpoint for model identifiers.
Checks that ModelConfig.from_identifier() can resolve model_path, but does
NOT load model weights into GPU memory.
"""
from core.inference.llama_cpp import LlamaServerNotFoundError
native_grant_backed = False
model_log_label = request.model_path
try:
model_identifier, model_log_label, native_grant_backed = (
_resolve_model_identifier_for_request(request, operation = "validate-model")
)
config = ModelConfig.from_identifier(
model_id = model_identifier,
hf_token = request.hf_token,
gguf_variant = request.gguf_variant,
)
if not config:
raise HTTPException(
status_code = 400,
detail = f"Invalid model identifier: {model_log_label}",
)
# Refuse early (before the frontend unloads to load this) if it can't fit
# alongside training, using the same settings /load uses so they agree.
effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
# Mirror /load: reject GGUF + gpu_ids before the guard so both return 400.
if config.is_gguf and effective_gpu_ids is not None:
raise HTTPException(
status_code = 400,
detail = "gpu_ids is not supported for GGUF models yet.",
)
effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit)
# Off-loop: guard does sync nvidia-smi / HF work.
await asyncio.to_thread(
_guard_chat_load_against_training,
config,
model_identifier = model_identifier,
hf_token = request.hf_token,
load_in_4bit = effective_load_in_4bit,
max_seq_length = request.max_seq_length,
requested_gpu_ids = effective_gpu_ids,
)
# Both checks cover the [adapter, base] set (matching the scan route and workers):
# either repo can ship auto_map code or a poisoned pickle.
security_targets = [config.identifier]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so its code/weights are reviewed too.
_base = get_base_model_from_lora_identifier(model_identifier, request.hf_token)
if _base:
security_targets.append(_base)
except Exception:
pass
security_targets = list(dict.fromkeys(security_targets))
is_gguf = getattr(config, "is_gguf", False)
# A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a
# mixed repo are inert for this load, so gating on them is a false positive. Only
# run the remote-code/security preflight for non-GGUF loads.
requires_trust_remote_code = False
requires_security_review = False
if not is_gguf:
requires_trust_remote_code = any(
_requires_trust_remote_code_for_model(_t, request.hf_token)
for _t in security_targets
)
requires_security_review = any(
_requires_security_review_for_model(_t, request.hf_token) for _t in security_targets
)
# Native context length, read from the local GGUF header when present.
# Lets the staged ("Load on selection" off) flow populate the context
# slider before the GPU load; None until the file is downloaded.
context_length: Optional[int] = None
if request.include_context_length and is_gguf:
from hub.utils.gguf import resolve_local_gguf_path
from utils.models.gguf_metadata import read_gguf_context_length
# Best-effort: a header-read failure must never fail validation of an
# otherwise-valid model (the outer except turns it into a 400).
try:
if native_grant_backed:
# model_identifier is the resolved canonical .gguf path.
local_gguf = model_identifier
else:
# Local folder / exported GGUFs already have their file
# resolved on the config (gguf_file is None for HF repos, so
# those fall back to the HF-cache lookup).
local_gguf = config.gguf_file or resolve_local_gguf_path(
model_identifier, request.gguf_variant
)
if local_gguf:
context_length = read_gguf_context_length(local_gguf)
except Exception as e:
logger.debug("Context-length probe failed for %s: %s", model_log_label, e)
return ValidateModelResponse(
valid = True,
message = "Model identifier is valid.",
identifier = model_log_label if native_grant_backed else config.identifier,
display_name = model_log_label
if native_grant_backed
else getattr(config, "display_name", config.identifier),
is_gguf = is_gguf,
is_lora = getattr(config, "is_lora", False),
is_vision = getattr(config, "is_vision", False),
requires_trust_remote_code = requires_trust_remote_code,
requires_security_review = requires_security_review,
context_length = context_length,
)
except HTTPException:
raise
except LlamaServerNotFoundError as e:
# Missing GGUF runtime: 400 with the install message, not a generic "Invalid model".
logger.warning("GGUF runtime missing while validating '%s': %s", request.model_path, e)
raise HTTPException(status_code = 400, detail = str(e))
except Exception as e:
if native_grant_backed:
redacted_msg = redact_native_paths(str(e))
logger.error(
"Error validating native model %s: %s",
model_log_label,
redacted_msg,
)
msg = _maybe_unsupported_message(redacted_msg)
raise HTTPException(
status_code = 400,
detail = f"Invalid native model {model_log_label}: {msg}",
)
logger.error(
f"Error validating model identifier '{request.model_path}': {e}",
exc_info = True,
)
# RuntimeError / ValueError carry intentional, actionable messages here
# (e.g. "llama-server binary not found - cannot load GGUF models. Run
# setup.sh ..."), so surface them instead of a blank "Invalid model".
# Path-redact for safety and keep any other exception type generic so an
# unexpected internal error never leaks its details to the client.
if isinstance(e, (RuntimeError, ValueError)):
msg = redact_native_paths(str(e)).strip()
if msg:
msg = _maybe_unsupported_message(msg)
raise HTTPException(
status_code = 400,
detail = msg,
)
raise HTTPException(
status_code = 400,
detail = "Invalid model",
)
@router.post("/unload", response_model = UnloadResponse)
async def unload_model(request: UnloadRequest, current_subject: str = Depends(get_current_subject)):
"""
Unload a model from memory.
Routes to the correct backend (llama-server for GGUF, Unsloth otherwise).
"""
# A deliberate unload means "stay unloaded": drop any idle reload stash so the
# next /v1 request can't resurrect this model. The idle loop unloads via the
# backend directly (not this route), so clearing here never fights keep-warm.
from core.inference.llama_keepwarm import inference_lifecycle_gate, note_model_unloaded
try:
# "Stop loading" (frontend cancelLoading -> /unload) must abort a still-loading
# model promptly. /load holds the lifecycle gate for the whole (multi-minute) load,
# so gating first would make the cancel wait it out. cancel_load only tears the
# loading subprocess down (no unload command), so it is safe off-gate.
backend = get_inference_backend()
loading = getattr(backend, "get_loading_model", lambda: None)()
if (
loading is not None
and hasattr(backend, "cancel_load")
and (request.model_path == loading or request.model_path.lower() == loading.lower())
):
if await asyncio.to_thread(backend.cancel_load, request.model_path):
note_model_unloaded()
logger.info(f"Cancelled in-flight load: {request.model_path}")
return UnloadResponse(status = "unloaded", model = request.model_path)
# Same "stop loading" fast path for a still-loading GGUF (llama-server spawned,
# health check not yet passed). A gated unload would wait out the multi-minute
# load; unload_model() sets the cancel_event load_model polls off its own lock and
# kills the child, sending no worker command, so it is safe off-gate like
# cancel_load. The gated GGUF branch below handles the already-loaded case. Gate on
# the loading model (identifier or native label): the single llama-server loads one
# GGUF at a time, so an unload for a different model must not cancel this load.
llama_backend = get_llama_cpp_backend()
if (
llama_backend.is_active
and not llama_backend.is_loaded
and (
llama_backend.model_identifier == request.model_path
or is_registered_native_path_label(
llama_backend.model_identifier, request.model_path
)
)
):
await asyncio.to_thread(llama_backend.unload_model)
note_model_unloaded()
logger.info(f"Cancelled in-flight GGUF load: {request.model_path}")
return UnloadResponse(status = "unloaded", model = request.model_path)
# Serialize with /load under the same lifecycle gate: the Unsloth unload now runs
# off the event loop (asyncio.to_thread), so without this a concurrent /load could
# swap in a fresh subprocess mid-unload and the unload command would land on the
# new worker. The gate makes load and unload exclusive.
async with inference_lifecycle_gate():
# Check if the GGUF backend has this model loaded or is loading it.
llama_backend = get_llama_cpp_backend()
if llama_backend.is_active and (
llama_backend.model_identifier == request.model_path
or is_registered_native_path_label(
llama_backend.model_identifier, request.model_path
)
or not llama_backend.is_loaded
):
# A manual unload is a deliberate user action: tear down now even if a
# request is mid-stream (only the automatic idle loop defers to it).
llama_backend.unload_model()
note_model_unloaded()
logger.info(f"Unloaded GGUF model: {request.model_path}")
return UnloadResponse(status = "unloaded", model = request.model_path)
# Unload from Unsloth backend off the event loop: unload takes _gen_lock, which
# a slow SSE stream paused between tokens still holds, so a sync call would block
# the loop that drives the stream's next token and the lock release.
backend = get_inference_backend()
await asyncio.to_thread(backend.unload_model, request.model_path)
note_model_unloaded()
logger.info(f"Unloaded model: {request.model_path}")
return UnloadResponse(status = "unloaded", model = request.model_path)
except Exception as e:
logger.error(f"Error unloading model: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = "Failed to unload model")
@studio_router.post("/cancel")
async def cancel_inference(request: Request, current_subject: str = Depends(get_current_subject)):
"""Cancel in-flight inference requests.
Body (JSON, at least one key required):
cancel_id - preferred: per-run UUID, matched exclusively.
session_id - fallback when cancel_id is absent.
completion_id - fallback when cancel_id is absent.
A cancel_id arriving before its stream registers is stashed briefly and
replayed on registration. Returns {"cancelled": N}.
"""
try:
body = await request.json()
if not isinstance(body, dict):
body = {}
except Exception as e:
logger.debug("Failed to parse cancel request body: %s", e)
body = {}
cancel_id = body.get("cancel_id")
if isinstance(cancel_id, str) and cancel_id:
return {"cancelled": _cancel_by_cancel_id_or_stash(cancel_id)}
keys = []
# `message_id` is the Anthropic passthrough's per-run identifier, so
# /v1/messages clients can cancel by their native id.
for k in ("completion_id", "session_id", "message_id"):
v = body.get(k)
if isinstance(v, str) and v:
keys.append(v)
if not keys:
return {"cancelled": 0}
n = _cancel_by_keys(keys)
return {"cancelled": n}
@studio_router.post("/tool-confirm")
async def confirm_tool_call(
request: ToolConfirmRequest, current_subject: str = Depends(get_current_subject)
):
matched = resolve_tool_decision(
request.approval_id,
request.decision,
session_id = request.session_id,
)
if not matched:
raise HTTPException(status_code = 404, detail = "No pending tool call confirmation")
return {"resolved": True}
@studio_router.get("/monitor")
async def get_api_monitor(current_subject: str = Depends(get_current_subject)):
"""Return recent OpenAI-compatible API activity for Studio."""
active_model = _monitor_active_model()
active_requests = api_monitor.active_count(subject = current_subject)
if active_requests:
operating_status = "generating"
elif active_model:
operating_status = "ready"
else:
operating_status = "idle"
return {
"status": operating_status,
"active_model": active_model,
"context_length": _monitor_context_length(),
"active_requests": active_requests,
"entries": api_monitor.snapshot(include_details = False, subject = current_subject),
}
@studio_router.get("/monitor/{entry_id}")
async def get_api_monitor_entry(entry_id: str, current_subject: str = Depends(get_current_subject)):
"""Return full prompt/reply details for one OpenAI-compatible API request."""
entry = api_monitor.get(entry_id, subject = current_subject)
if entry is None:
raise HTTPException(status_code = 404, detail = "Monitor entry not found")
return entry
@router.post("/generate/stream")
async def generate_stream(
request: GenerateRequest,
fastapi_request: Request,
current_subject: str = Depends(get_current_subject),
):
"""
Generate a chat response with Server-Sent Events (SSE) streaming.
For vision models, provide image_base64 (base64-encoded image).
"""
backend = get_inference_backend()
if not backend.active_model_name:
raise HTTPException(
status_code = 400, detail = "No model loaded. Call POST /inference/load first."
)
# Decode image if provided (vision models)
image = None
if request.image_base64:
try:
import base64
from PIL import Image
from io import BytesIO
# Check current model supports vision
model_info = backend.models.get(backend.active_model_name, {})
if not model_info.get("is_vision"):
raise HTTPException(
status_code = 400,
detail = "Image provided but current model is text-only. Load a vision model.",
)
image_data = base64.b64decode(request.image_base64)
image = Image.open(BytesIO(image_data))
image = backend.resize_image(image)
except HTTPException:
raise
except Exception as e:
raise log_and_http_error(
e,
400,
"Failed to decode image",
event = "inference.decode_image_failed",
log = logger,
)
cancel_event = threading.Event()
async def stream():
gen = None
completed = False
# Cancel the generation when the client disconnects. The generator only
# awaits asyncio.to_thread(next, gen, ...), so without a concurrent
# watcher a disconnect during a long prefill/generation would go
# unnoticed until the next send and the backend would keep generating.
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(fastapi_request, cancel_event)
)
try:
gen = backend.generate_chat_response(
messages = request.messages,
system_prompt = request.system_prompt,
image = image,
temperature = request.temperature,
top_p = request.top_p,
top_k = request.top_k,
min_p = request.min_p,
max_new_tokens = request.max_new_tokens,
repetition_penalty = request.repetition_penalty,
presence_penalty = request.presence_penalty,
cancel_event = cancel_event,
)
_DONE = object()
while True:
if cancel_event.is_set():
# Watcher set cancel_event between chunks. Reset here: closing
# the generator does not signal a subprocess backend, so it would
# keep decoding. The finally's reset is guarded, so no double-run.
backend.reset_generation_state()
break
chunk = await asyncio.to_thread(next, gen, _DONE)
if chunk is _DONE:
completed = True
break
yield f"data: {json.dumps({'content': chunk})}\n\n"
if completed:
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
backend.reset_generation_state()
raise
except Exception as e:
cancel_event.set()
backend.reset_generation_state()
logger.error(f"Error during generation: {e}", exc_info = True)
yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n"
finally:
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
if not completed and not cancel_event.is_set():
cancel_event.set()
backend.reset_generation_state()
if gen is not None:
try:
await asyncio.to_thread(gen.close)
except (RuntimeError, ValueError):
pass
return _sse_streaming_response(stream())
@router.get("/status", response_model = InferenceStatusResponse)
async def get_status(current_subject: str = Depends(get_current_subject)):
"""
Get current inference backend status.
Reports whichever backend (Unsloth or llama-server) is active.
"""
try:
llama_backend = get_llama_cpp_backend()
# MTP probe + freshness check (both cached); drive the UI banner.
try:
_bin = type(llama_backend)._find_llama_server_binary()
_caps = type(llama_backend).probe_server_capabilities(_bin)
_supports_mtp = bool(_caps.get("supports_mtp", False))
except Exception:
_bin = None
_supports_mtp = True # fail open
try:
from utils.llama_cpp_freshness import check_prebuilt_freshness
_freshness = check_prebuilt_freshness(_bin)
except Exception:
_freshness = {}
_stale = bool(_freshness.get("stale"))
_installed_tag = _freshness.get("installed_tag")
_latest_tag = _freshness.get("latest_tag")
# If a GGUF model is loaded via llama-server, report that
if llama_backend.is_loaded:
_model_id = llama_backend.model_identifier
_native_grant_backed = getattr(llama_backend, "_native_grant_backed", False)
_display_model_id = getattr(
llama_backend, "_native_display_label", None
) or display_label_for_native_path(_model_id)
if (
_native_grant_backed
and _model_id
and _display_model_id == _model_id
and os.path.isabs(_model_id)
):
_display_model_id = os.path.basename(_model_id)
_inference_cfg = load_inference_config(_model_id) if _model_id else None
_audio_type = getattr(llama_backend, "_audio_type", None)
# Don't surface Studio's auto-applied bundled family template (e.g. the
# gemma-4 override) as a user-authored override: the frontend adopts
# status.chat_template_override as editable state and would otherwise
# re-send it as an explicit override for a later, unrelated model. Only
# expose a genuine user override.
_reported_chat_template_override = llama_backend.chat_template_override
_auto_chat_template_override = resolve_effective_chat_template_override(
model_identifier = _model_id,
user_override = None,
)
if (
_auto_chat_template_override is not None
and _reported_chat_template_override == _auto_chat_template_override
):
_reported_chat_template_override = None
return InferenceStatusResponse(
active_model = _display_model_id,
model_identifier = None if _native_grant_backed else _model_id,
is_vision = llama_backend.is_vision,
is_gguf = True,
is_diffusion = llama_backend.is_diffusion,
gguf_variant = llama_backend.hf_variant,
is_audio = getattr(llama_backend, "_is_audio", False),
audio_type = _audio_type,
has_audio_input = getattr(llama_backend, "_has_audio_input", False),
loading = [],
loaded = [_display_model_id] if _display_model_id else [],
inference = _inference_cfg,
# GGUF status: auto_map never executes, so inert (matches validate_model).
requires_trust_remote_code = False,
supports_reasoning = llama_backend.supports_reasoning,
reasoning_style = llama_backend.reasoning_style,
reasoning_effort_levels = llama_backend.reasoning_effort_levels,
reasoning_always_on = llama_backend.reasoning_always_on,
supports_preserve_thinking = llama_backend.supports_preserve_thinking,
supports_tools = llama_backend.supports_tools,
chat_template = llama_backend.chat_template,
context_length = llama_backend.context_length,
max_context_length = llama_backend.max_context_length,
native_context_length = llama_backend.native_context_length,
cache_type_kv = llama_backend.cache_type_kv,
chat_template_override = _reported_chat_template_override,
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
tensor_parallel = llama_backend.tensor_parallel,
llama_cpp_supports_mtp = _supports_mtp,
spec_fallback_reason = llama_backend.spec_fallback_reason,
llama_cpp_prebuilt_stale = _stale,
llama_cpp_installed_tag = _installed_tag,
llama_cpp_latest_tag = _latest_tag,
)
# Otherwise, report Unsloth backend status
backend = get_inference_backend()
is_vision = False
is_audio = False
audio_type = None
has_audio_input = False
model_info = {}
if backend.active_model_name:
model_info = backend.models.get(backend.active_model_name, {})
is_vision = model_info.get("is_vision", False)
is_audio = model_info.get("is_audio", False)
audio_type = model_info.get("audio_type")
has_audio_input = model_info.get("has_audio_input", False)
chat_template_info = model_info.get("chat_template_info", {})
chat_template = (
chat_template_info.get("template") if isinstance(chat_template_info, dict) else None
)
# Non-GGUF: classify from the loaded template.
_sf_flags = _detect_safetensors_features(backend, chat_template)
inference_config = (
load_inference_config(backend.active_model_name) if backend.active_model_name else None
)
return InferenceStatusResponse(
active_model = backend.active_model_name,
model_identifier = backend.active_model_name,
is_vision = is_vision,
is_gguf = False,
is_audio = is_audio,
audio_type = audio_type,
has_audio_input = has_audio_input,
loading = list(getattr(backend, "loading_models", set())),
loaded = list(backend.models.keys()),
inference = inference_config,
requires_trust_remote_code = _resolve_loaded_trust_remote_code(
backend.active_model_name, model_info, inference_config
),
supports_reasoning = _sf_flags["supports_reasoning"],
reasoning_style = _sf_flags["reasoning_style"],
reasoning_effort_levels = _sf_flags.get("reasoning_effort_levels", []),
reasoning_always_on = _sf_flags["reasoning_always_on"],
supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
supports_tools = _sf_flags["supports_tools"],
context_length = _positive_int_or_none(model_info.get("context_length")),
chat_template = chat_template,
llama_cpp_supports_mtp = _supports_mtp,
llama_cpp_prebuilt_stale = _stale,
llama_cpp_installed_tag = _installed_tag,
llama_cpp_latest_tag = _latest_tag,
)
except Exception as e:
logger.error(f"Error getting status: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = "Failed to get status")
@router.get("/load-progress", response_model = LoadProgressResponse)
async def get_load_progress(current_subject: str = Depends(get_current_subject)):
"""
Return the active GGUF load's mmap/upload progress.
During the warmup window after a GGUF download -- when llama-server pages
~tens-to-hundreds of GB of shards into the page cache before pushing layers
to VRAM -- ``/api/inference/status`` only shows a generic spinner. This
exposes sampled progress so the UI can render a real bar plus rate/ETA.
Returns an empty payload (``phase=null, bytes=0``) when no load is in
flight. The frontend should stop polling once ``phase`` becomes ``ready``.
"""
try:
llama_backend = get_llama_cpp_backend()
progress = llama_backend.load_progress()
if progress is None:
return LoadProgressResponse()
return LoadProgressResponse(**progress)
except Exception as e:
logger.warning(f"Error sampling load progress: {e}")
return LoadProgressResponse()
# =====================================================================
# Audio (TTS) Generation (/audio/generate)
# =====================================================================
@router.post("/audio/generate")
async def generate_audio(
payload: ChatCompletionRequest,
request: Request,
current_subject: str = Depends(get_current_subject),
):
"""
Generate audio (TTS) from the latest user message.
Returns JSON with base64-encoded WAV audio.
Works with both GGUF (llama-server) and Unsloth/transformers backends.
"""
import base64
# Extract text from the last user message
_, chat_messages, _ = _extract_content_parts(payload.messages)
if not chat_messages:
raise HTTPException(status_code = 400, detail = "No messages provided.")
last_user_msg = next((m for m in reversed(chat_messages) if m["role"] == "user"), None)
if not last_user_msg:
raise HTTPException(status_code = 400, detail = "No user message found.")
text = last_user_msg["content"]
# Restore an idle-evicted GGUF before selecting a backend: this path is
# keep-warm-tracked but had no reload hook, so a standalone idle TTL could
# unload an audio GGUF the next request then failed to restore. Validation
# above ran first, so an invalid request never triggers a reload.
#
# Reload-only on purpose: a local GGUF's audio-input capability is not a cheap
# pre-load probe (the companion mmproj signal can't tell an audio projector
# from a vision one, and codec-based TTS ships no projector at all), so passing
# the client model through the resolver could load a text- or vision-only target
# and evict the working audio model before the audio backend check fails. Only
# the idle-stash restore runs here; switching TTS models is an explicit /load.
await _maybe_auto_switch_model(_RELOAD_ONLY_MODEL, request, current_subject)
# Pick backend — both return (wav_bytes, sample_rate)
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded and getattr(llama_backend, "_is_audio", False):
# Advertised repo id after an auto-switch load, else a clean public id,
# never the absolute .gguf path.
model_name = _llama_public_model_id(llama_backend)
gen = lambda: llama_backend.generate_audio_response(
text = text,
audio_type = llama_backend._audio_type,
temperature = payload.temperature,
top_p = payload.top_p,
top_k = payload.top_k,
min_p = payload.min_p,
max_new_tokens = _effective_max_tokens(payload) or 2048,
repetition_penalty = payload.repetition_penalty,
)
else:
backend = get_inference_backend()
if not backend.active_model_name:
raise HTTPException(status_code = 400, detail = "No model loaded.")
model_info = backend.models.get(backend.active_model_name, {})
if not model_info.get("is_audio"):
raise HTTPException(status_code = 400, detail = "Active model is not an audio model.")
model_name = public_model_id(backend.active_model_name)
gen = lambda: backend.generate_audio_response(
text = text,
temperature = payload.temperature,
top_p = payload.top_p,
top_k = payload.top_k,
min_p = payload.min_p,
max_new_tokens = _effective_max_tokens(payload) or 2048,
repetition_penalty = payload.repetition_penalty,
use_adapter = payload.use_adapter,
)
try:
wav_bytes, sample_rate = await asyncio.to_thread(gen)
except Exception as e:
logger.error(f"Audio generation error: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = safe_error_detail(e))
audio_b64 = base64.b64encode(wav_bytes).decode("ascii")
return JSONResponse(
content = {
"id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
"object": "chat.completion.audio",
"model": model_name,
"audio": {"data": audio_b64, "format": "wav", "sample_rate": sample_rate},
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": f'[Generated audio from: "{text[:100]}"]',
},
"finish_reason": "stop",
}
],
}
)
# =====================================================================
# OpenAI-Compatible Chat Completions (/chat/completions)
# =====================================================================
def _decode_audio_base64(b64: str) -> np.ndarray:
"""Decode base64 audio (any format) → float32 numpy array at 16kHz."""
import torchaudio
import tempfile
import os
from utils.paths import ensure_dir, tmp_root
raw = base64.b64decode(b64)
# torchaudio.load needs a path or file-like with a format hint; write a
# temp file so it can auto-detect the format.
with tempfile.NamedTemporaryFile(
suffix = ".audio",
delete = False,
dir = str(ensure_dir(tmp_root())),
) as tmp:
tmp.write(raw)
tmp_path = tmp.name
try:
waveform, sr = torchaudio.load(tmp_path)
finally:
os.unlink(tmp_path)
if waveform.shape[0] > 1:
waveform = waveform.mean(dim = 0, keepdim = True)
if sr != 16000:
resampler = torchaudio.transforms.Resample(orig_freq = sr, new_freq = 16000)
waveform = resampler(waveform)
return waveform.squeeze(0).numpy()
# Reject oversized audio before decoding. base64 inflates raw bytes by ~4/3, so
# cap the encoded length to bound the upload. _MAX_AUDIO_SECONDS additionally
# bounds the *decoded* length, since a small compressed file (opus/flac/etc.)
# can expand to a far larger PCM array than the encoded-size cap implies.
_MAX_AUDIO_RAW_BYTES = 25 * 1024 * 1024
_MAX_AUDIO_B64_CHARS = _MAX_AUDIO_RAW_BYTES * 4 // 3
_MAX_AUDIO_SECONDS = 30 * 60
_WAV_HEADER_BYTES = 44
_MIN_TRANSCODE_AUDIO_SAMPLE_RATE = 8000
def _sniff_audio_container(raw: bytes) -> Optional[str]:
"""Return 'wav' or 'mp3' if the bytes are a container llama-server accepts
directly (so we can forward them untouched), else None (needs transcoding)."""
if len(raw) >= 12 and raw[:4] == b"RIFF" and raw[8:12] == b"WAVE":
return "wav"
# mp3: ID3 tag, or an MPEG audio frame sync (no other accepted format leads
# with 0xFF, so the simple sync check doesn't collide).
if raw[:3] == b"ID3" or (len(raw) >= 2 and raw[0] == 0xFF and (raw[1] & 0xE0) == 0xE0):
return "mp3"
return None
def _mono_f32_to_wav_bytes(arr: np.ndarray, sample_rate: int) -> bytes:
"""Encode a mono float32 array as 16-bit PCM WAV bytes.
Torch-free (numpy + stdlib only) so it works on no-torch GGUF-only installs;
the shared audio_codecs helper pulls in torch at import time.
"""
import io
import wave
arr = np.nan_to_num(np.asarray(arr, dtype = np.float32).flatten(), posinf = 0.0, neginf = 0.0)
if arr.size == 0:
raise ValueError("decoded audio is empty")
peak = float(np.abs(arr).max())
if peak > 1.0:
arr = arr / peak
pcm = (arr * 32767.0).astype(np.int16)
buf = io.BytesIO()
with wave.open(buf, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2)
wf.setframerate(int(sample_rate))
wf.writeframes(pcm.tobytes())
return buf.getvalue()
def _resample_mono_linear(arr: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray:
"""Small numpy-only resampler for upload size limiting."""
if source_rate <= 0 or target_rate <= 0 or source_rate == target_rate:
return arr
duration = len(arr) / float(source_rate)
target_len = max(1, int(round(duration * target_rate)))
if target_len == len(arr):
return arr
source_x = np.linspace(0.0, duration, num = len(arr), endpoint = False)
target_x = np.linspace(0.0, duration, num = target_len, endpoint = False)
return np.interp(target_x, source_x, arr).astype(np.float32)
def _fit_transcoded_audio_to_wav_cap(arr: np.ndarray, sample_rate: int) -> tuple[np.ndarray, int]:
"""Downsample only when needed so transcoded WAV stays within the upload cap."""
if sample_rate <= 0:
raise ValueError("decoded audio has an invalid sample rate")
wav_bytes = _WAV_HEADER_BYTES + len(arr) * 2
if wav_bytes <= _MAX_AUDIO_RAW_BYTES:
return arr, sample_rate
duration = len(arr) / float(sample_rate)
max_samples = max(1, (_MAX_AUDIO_RAW_BYTES - _WAV_HEADER_BYTES) // 2)
target_rate = int(max_samples // duration)
if target_rate < _MIN_TRANSCODE_AUDIO_SAMPLE_RATE:
raise ValueError("decoded audio exceeds the transcoded WAV size limit")
target_rate = min(sample_rate, target_rate)
fitted = _resample_mono_linear(arr, sample_rate, target_rate)
if _WAV_HEADER_BYTES + len(fitted) * 2 > _MAX_AUDIO_RAW_BYTES:
raise ValueError("decoded audio exceeds the transcoded WAV size limit")
return fitted, target_rate
def _decode_audio_mono(raw: bytes) -> tuple[np.ndarray, int]:
"""Decode audio bytes to (mono float32 array, native sample_rate).
soundfile (libsndfile) reads wav/mp3/ogg/flac straight from memory. librosa
(ffmpeg-backed) additionally covers m4a/webm but needs a real path and is
absent on no-torch GGUF-only installs. Both imports are inside the fallback
so a missing decoder degrades to the next one (and finally a clear error)
rather than crashing.
"""
import io
try:
import soundfile as sf
arr, sr = sf.read(io.BytesIO(raw), dtype = "float32")
except Exception:
try:
import librosa
except ModuleNotFoundError as e:
raise RuntimeError(
"this audio format needs librosa, which is not installed in "
"GGUF-only environments; use wav, mp3, ogg or flac"
) from e
import os
import tempfile
from utils.paths import ensure_dir, tmp_root
with tempfile.NamedTemporaryFile(
suffix = ".audio",
delete = False,
dir = str(ensure_dir(tmp_root())),
) as tmp:
tmp.write(raw)
tmp_path = tmp.name
try:
arr, sr = librosa.load(tmp_path, sr = None, mono = True)
finally:
os.unlink(tmp_path)
if arr.ndim > 1:
arr = arr.mean(axis = 1)
if sr > 0 and len(arr) > sr * _MAX_AUDIO_SECONDS:
raise ValueError(f"decoded audio exceeds the {_MAX_AUDIO_SECONDS // 60}-minute limit")
return arr, sr
def _prepare_audio_for_llama(b64: str) -> tuple[str, str]:
"""Return (base64, format) ready for llama-server's input_audio part.
llama-server's API only accepts wav/mp3, and decodes/resamples/down-mixes
them itself, so wav and mp3 uploads are forwarded untouched (no decode, no
PCM payload inflation). Other containers (m4a/ogg/webm/flac) are decoded to
a mono WAV. Blocking; call via a thread from async paths.
"""
if b64.startswith("data:"):
b64 = b64.split(",", 1)[1] if "," in b64 else ""
raw = base64.b64decode(b64)
passthrough = _sniff_audio_container(raw)
if passthrough is not None:
return b64, passthrough
arr, sr = _decode_audio_mono(raw)
arr, sr = _fit_transcoded_audio_to_wav_cap(arr, sr)
return base64.b64encode(_mono_f32_to_wav_bytes(arr, sr)).decode("ascii"), "wav"
def _inject_audio_part(messages: list[dict], audio_b64: str, audio_format: str) -> None:
"""Append an input_audio part to the last user message, in place.
Audio rides in the message list like image_url parts do, so it flows through
both the plain and tool-calling generation paths.
"""
part = {
"type": "input_audio",
"input_audio": {"data": audio_b64, "format": audio_format},
}
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content")
if isinstance(content, list):
content.append(part)
else:
msg["content"] = [{"type": "text", "text": content or ""}, part]
return
def _extract_content_parts(messages: list) -> tuple[str, list[dict], "Optional[str]"]:
"""
Parse OpenAI-format messages into components the inference backend expects.
Handles both plain-string ``content`` and multimodal content-part arrays
(``[{type: "text", ...}, {type: "image_url", ...}]``).
Returns:
system_prompt: System message text (empty string if none).
chat_messages: Non-system messages with content flattened to strings.
image_base64: Base64 of the *first* image found, or ``None``.
"""
system_parts: list[str] = []
chat_messages: list[dict] = []
first_image_b64: Optional[str] = None
for msg in messages:
# ── System / developer messages → extract as system_prompt ────────
if msg.role in ("system", "developer"):
if isinstance(msg.content, str):
system_parts.append(msg.content)
elif isinstance(msg.content, list):
# Unlikely but handle: join text parts
system_parts.append("\n".join(p.text for p in msg.content if p.type == "text"))
continue
# ── User / assistant messages ─────────────────────────
if isinstance(msg.content, str):
# Plain string content — pass through
chat_messages.append({"role": msg.role, "content": msg.content})
elif isinstance(msg.content, list):
# Multimodal content parts
text_parts: list[str] = []
for part in msg.content:
if part.type == "text":
text_parts.append(part.text)
elif part.type == "image_url" and first_image_b64 is None:
url = part.image_url.url
if url.startswith("data:"):
# data:image/png;base64,<DATA> -> extract <DATA>
first_image_b64 = url.split(",", 1)[1] if "," in url else None
else:
logger.warning(f"Remote image URLs not yet supported: {url[:80]}...")
combined_text = "\n".join(text_parts) if text_parts else ""
chat_messages.append({"role": msg.role, "content": combined_text})
return "\n\n".join(p for p in system_parts if p), chat_messages, first_image_b64
# ── External provider proxy ──────────────────────────────────────
# Providers whose stream helper translates `input_document` parts into a
# native attachment block on the wire. Anthropic: `_stream_anthropic` ->
# {type:"document", source:...}; OpenAI: `_stream_openai_responses` ->
# {type:"input_file", file_data|file_url}. Every other provider (gemini /
# mistral / kimi / openrouter / deepseek / custom OpenAI-compat) goes through
# the generic /chat/completions passthrough that forwards messages verbatim,
# so handing them an `input_document` part would 400 with an unknown
# content_part type.
_INPUT_DOCUMENT_PROVIDERS = frozenset({"anthropic", "openai"})
def _build_external_messages(
messages: list,
supports_vision: bool,
provider_type: Optional[str] = None,
base_url: Optional[str] = None,
) -> list[dict]:
"""
Convert ChatMessage list to OpenAI-compatible dicts for external providers.
Behaviour per content-part type:
- `text`: always preserved.
- `image_url`: preserved on vision providers; stripped on non-vision.
- `input_document`: preserved ONLY when the provider's stream helper has
explicit translation logic (Anthropic + OpenAI today, see
``_INPUT_DOCUMENT_PROVIDERS``). Stripped for every other provider so the
unknown type doesn't reach generic /chat/completions and 400.
- `reasoning`: OpenAI-only Responses reasoning item paired with a prior
tool output. Forwarded ONLY when provider_type=="openai" so follow-up
image edits can replay the required reasoning item.
- `image_generation_call`: OpenAI-only Responses image reference. Forwarded
ONLY when provider_type=="openai" so follow-up image edits can reference
prior generated images.
- `compaction`: Anthropic-only synthetic part (round-trips server-side
compaction state). Forwarded ONLY when provider_type=="anthropic";
stripped elsewhere so the unknown part doesn't reach generic
/chat/completions and 400 (DeepSeek, Mistral, Gemini, Kimi, OpenRouter).
"""
document_provider = provider_type in _INPUT_DOCUMENT_PROVIDERS
anthropic = provider_type == "anthropic"
openai = provider_type == "openai"
# `extra_content` carries the assistant's text-part `thoughtSignature`
# round-trip on Gemini's native streamGenerateContent endpoint. Custom
# Gemini OpenAI-compat gateways (LiteLLM etc.) route through
# /chat/completions where the field is unknown and can be rejected -- gate
# strictly on the Google-hosted Gemini base.
_native_gemini = False
if provider_type == "gemini" and base_url:
try:
from urllib.parse import urlparse as _urlparse
_host = (_urlparse(base_url).hostname or "").lower()
_native_gemini = _host == "generativelanguage.googleapis.com"
except Exception:
_native_gemini = False
emit_extra_content = _native_gemini
_SERVER_BUILTIN_TOOL_NAMES = frozenset(
{"web_search", "web_fetch", "code_execution", "image_generation"}
)
def _is_marked_server_builtin_tool_call(tc: Any) -> bool:
"""Return True iff `tc` is a synthetic provider-side tool card with a
canonical builtin name and either:
- the `args._server_tool` marker stamped by the backend, or
- a Gemini `args.google.native_part` payload (durable replay signal
for code_execution / image_generation that predates the marker).
Such cards must not be forwarded to non-native providers: they aren't
real user functions, so the receiving API rejects the orphan tool
history. Real user functions with these names normally have neither
signal.
"""
if not isinstance(tc, dict):
return False
fn = tc.get("function")
if not isinstance(fn, dict):
return False
name = (fn.get("name") or "").lower()
if name not in _SERVER_BUILTIN_TOOL_NAMES:
return False
raw_args = fn.get("arguments") or ""
try:
args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
except Exception:
return False
if not isinstance(args, dict):
return False
if args.get("_server_tool") is True:
return True
google = args.get("google")
return isinstance(google, dict) and isinstance(google.get("native_part"), dict)
# When we drop a server-side builtin tool_call, the matching `role="tool"`
# follow-up must also be dropped -- else the provider gets an orphan
# tool_call_id with no matching assistant call, which OpenAI Responses and
# Anthropic both reject.
dropped_server_builtin_tool_call_ids: set[str] = set()
def _filter_tool_calls(tool_calls: Any) -> Optional[list]:
"""Sanitize assistant `tool_calls` for non-native-Gemini providers.
Two concerns:
1. `tool_calls[i].extra_content` carries Gemini-only thoughtSignature
metadata; strip it for providers that can't parse the unknown key.
2. Marked server-side builtin cards (`_server_tool: true` on a
canonical builtin name, or a Gemini `native_part` payload) are
Studio-internal tool cards from a prior native Gemini turn;
forwarding them to OpenAI / Anthropic / custom OAI-compat gateways
sends an orphan `tool_calls` entry (no matching tool declaration,
often no matching `role="tool"` reply) that can be rejected. We
record the dropped call_ids so the matching role=tool message is
skipped below.
Native Gemini keeps both untouched so the translator can replay them
via `native_part`.
"""
if not tool_calls:
return None
if not isinstance(tool_calls, list):
return tool_calls
if emit_extra_content:
return tool_calls
cleaned: list = []
for _tc in tool_calls:
if _is_marked_server_builtin_tool_call(_tc):
_tc_id = _tc.get("id") if isinstance(_tc, dict) else None
if isinstance(_tc_id, str) and _tc_id:
dropped_server_builtin_tool_call_ids.add(_tc_id)
continue
if not isinstance(_tc, dict):
cleaned.append(_tc)
continue
if "extra_content" not in _tc:
cleaned.append(_tc)
continue
_stripped = {k: v for k, v in _tc.items() if k != "extra_content"}
cleaned.append(_stripped)
return cleaned
def _openai_responses_part(item: Any) -> Optional[dict[str, Any]]:
"""Rebuild a forwarded OpenAI Responses assistant part (`reasoning` or
`image_generation_call`); returns None for any other part type."""
if item.type == "reasoning":
reasoning: dict[str, Any] = {
"type": "reasoning",
"id": item.id,
"summary": item.summary,
}
if item.status:
reasoning["status"] = item.status
return reasoning
if item.type == "image_generation_call":
image_ref: dict[str, Any] = {"type": "image_generation_call", "id": item.id}
if getattr(item, "response_id", None):
image_ref["response_id"] = item.response_id
return image_ref
return None
result = []
for msg in messages:
# Drop role=tool messages whose matching server-builtin tool_call was
# filtered above. An orphan tool_result with no matching tool_call is
# rejected by OpenAI Responses and Anthropic.
if (
msg.role == "tool"
and isinstance(msg.tool_call_id, str)
and msg.tool_call_id in dropped_server_builtin_tool_call_ids
):
continue
if isinstance(msg.content, str):
# Drop bare assistant messages with no content AND no tool_calls
# (some providers reject empty assistant turns). Preserve assistant
# turns whose only payload is tool_calls so multi-turn
# function-call loops round-trip.
if msg.role == "assistant" and not msg.content.strip() and not msg.tool_calls:
continue
out: dict[str, Any] = {"role": msg.role, "content": msg.content}
if msg.role == "assistant" and msg.tool_calls:
_tcs = _filter_tool_calls(msg.tool_calls)
if _tcs:
out["tool_calls"] = _tcs
elif not msg.content.strip():
# Every tool_call was a dropped synthetic provider card;
# the turn would be an empty
# `{"role":"assistant","content":""}` that some providers
# reject. Skip it entirely.
continue
if msg.role == "tool":
if msg.tool_call_id:
out["tool_call_id"] = msg.tool_call_id
if msg.name:
out["name"] = msg.name
if emit_extra_content and msg.role == "assistant" and msg.extra_content:
out["extra_content"] = msg.extra_content
result.append(out)
continue
# Assistant messages with content=None but populated tool_calls are
# valid (post-tool-call turn). Forward them so the provider helper can
# rebuild the functionCall part.
if msg.content is None and msg.role == "assistant" and msg.tool_calls:
_filtered_tcs = _filter_tool_calls(msg.tool_calls)
if not _filtered_tcs:
# Every tool_call was provider-side synthetic and dropped;
# skip the whole message to avoid an empty assistant turn.
continue
_assistant_only: dict[str, Any] = {
"role": "assistant",
"content": "",
"tool_calls": _filtered_tcs,
}
if emit_extra_content and msg.extra_content:
_assistant_only["extra_content"] = msg.extra_content
result.append(_assistant_only)
continue
if isinstance(msg.content, list):
if supports_vision:
parts = []
for part in msg.content:
if part.type == "text":
parts.append({"type": "text", "text": part.text})
elif part.type == "image_url":
parts.append(
{
"type": "image_url",
"image_url": {"url": part.image_url.url},
}
)
elif (
openai
and msg.role == "assistant"
and (_rp := _openai_responses_part(part)) is not None
):
# ExternalProviderClient maps image_generation_call onto a
# top-level Responses input item after the current user
# prompt, or onto `previous_response_id` when response_id
# is available from the prior turn.
parts.append(_rp)
elif part.type == "input_document" and document_provider:
# ExternalProviderClient maps this onto Anthropic's
# `document` or OpenAI Responses' `input_file` block;
# every other provider would 400 on the unknown part.
doc: dict[str, Any] = {"type": "input_document"}
if part.file_data:
doc["file_data"] = part.file_data
if part.file_url:
doc["file_url"] = part.file_url
if part.filename:
doc["filename"] = part.filename
if part.media_type:
doc["media_type"] = part.media_type
parts.append(doc)
elif part.type == "compaction" and anthropic:
# Anthropic stream helper forwards this as a native
# `compaction` block; every other provider would 400 on
# the unknown part, so gate by provider_type.
parts.append({"type": "compaction", "content": part.content})
entry: dict[str, Any] = {"role": msg.role, "content": parts}
if msg.role == "assistant" and msg.tool_calls:
_tcs = _filter_tool_calls(msg.tool_calls)
if _tcs:
entry["tool_calls"] = _tcs
elif not parts:
# All tool_calls were synthetic and dropped, and no
# content parts survived. Skip rather than forward an
# empty assistant turn that downstream providers reject.
continue
elif msg.role == "assistant" and not parts:
continue
if msg.role == "tool":
if msg.tool_call_id:
entry["tool_call_id"] = msg.tool_call_id
if msg.name:
entry["name"] = msg.name
if emit_extra_content and msg.role == "assistant" and msg.extra_content:
entry["extra_content"] = msg.extra_content
result.append(entry)
else:
# Non-vision provider: strip images / documents, keep text,
# optionally keep compaction (Anthropic only --
# compaction-capable Anthropic models all report
# supports_vision=True today, but gate here for safety).
preserved = []
for p in msg.content:
if p.type == "text":
preserved.append({"type": "text", "text": p.text})
elif (
openai
and msg.role == "assistant"
and (_rp := _openai_responses_part(p)) is not None
):
preserved.append(_rp)
elif p.type == "compaction" and anthropic:
preserved.append({"type": "compaction", "content": p.content})
if msg.role == "assistant" and not preserved:
continue
if len(preserved) == 1 and preserved[0]["type"] == "text":
# Single text part collapses to a string for providers that
# don't accept content arrays.
entry = {"role": msg.role, "content": preserved[0]["text"]}
else:
entry = {"role": msg.role, "content": preserved}
if msg.role == "assistant" and msg.tool_calls:
_tcs = _filter_tool_calls(msg.tool_calls)
if _tcs:
entry["tool_calls"] = _tcs
else:
# All tool_calls were synthetic and dropped; skip if no
# content survived either.
_entry_content = entry.get("content")
_has_text = (
isinstance(_entry_content, str) and _entry_content.strip()
) or (isinstance(_entry_content, list) and len(_entry_content) > 0)
if not _has_text:
continue
if msg.role == "tool":
if msg.tool_call_id:
entry["tool_call_id"] = msg.tool_call_id
if msg.name:
entry["name"] = msg.name
if emit_extra_content and msg.role == "assistant" and msg.extra_content:
entry["extra_content"] = msg.extra_content
result.append(entry)
return result
async def _proxy_to_external_provider(
payload: ChatCompletionRequest,
request: Request,
current_subject: Optional[str] = None,
) -> StreamingResponse:
"""
Proxy a chat completion request to an external LLM provider.
Resolves provider config (DB or registry), decrypts the API key, and
streams the response back in OpenAI SSE format.
"""
# Resolve provider type and base URL
provider_type = payload.provider_type
base_url = payload.provider_base_url
if payload.provider_id:
config = providers_db.get_provider(payload.provider_id)
if config is None:
raise HTTPException(
status_code = 404,
detail = f"Provider config not found: {payload.provider_id}",
)
if not config["is_enabled"]:
raise HTTPException(
status_code = 400,
detail = f"Provider '{config['display_name']}' is disabled.",
)
provider_type = provider_type or config["provider_type"]
base_url = base_url or config["base_url"]
if not provider_type:
raise HTTPException(
status_code = 400,
detail = "Either provider_id or provider_type is required for external provider routing.",
)
# Fall back to registry default base URL
if not base_url:
base_url = get_base_url(provider_type)
if not base_url:
raise HTTPException(
status_code = 400,
detail = f"Unknown provider type: {provider_type}",
)
api_key = ""
if payload.encrypted_api_key:
try:
api_key = decrypt_api_key(payload.encrypted_api_key)
except Exception as exc:
logger.warning("external_provider.decrypt_failed", error = str(exc))
raise HTTPException(
status_code = 400,
detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.",
)
model = payload.external_model or payload.model
if model == "default":
raise HTTPException(
status_code = 400,
detail = "external_model is required when using an external provider.",
)
# Build messages, preserving multimodal content for vision providers
from core.inference.providers import get_provider_info as _get_provider_info
_pinfo = _get_provider_info(provider_type) or {}
_supports_vision = _pinfo.get("supports_vision", False)
chat_messages = _build_external_messages(
payload.messages,
_supports_vision,
provider_type = provider_type,
base_url = base_url,
)
monitor_id = None
if not getattr(request.state, "skip_api_monitor", False):
monitor_id = api_monitor.start(
endpoint = request.url.path,
method = request.method,
model = model,
prompt = _monitor_prompt_from_messages(payload.messages),
context_length = None,
subject = current_subject,
)
client = ExternalProviderClient(
provider_type = provider_type,
base_url = base_url,
api_key = api_key,
)
# `top_k` defaults to 20 in ChatCompletionRequest because the local path
# expects an int, but the external-provider path treats "field omitted from
# JSON" as "use provider default" so callers sending only model/messages
# don't silently get different sampling than before this PR. Pydantic's
# `model_fields_set` tracks explicit-vs-default per request.
_top_k_explicit = payload.top_k if "top_k" in payload.model_fields_set else None
async def _stream():
gen = client.stream_chat_completion(
messages = chat_messages,
model = model,
temperature = payload.temperature,
top_p = payload.top_p,
# 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,
reasoning_effort = payload.reasoning_effort,
enabled_tools = payload.enabled_tools,
enable_prompt_caching = payload.enable_prompt_caching,
openai_code_exec_container_id = payload.openai_code_exec_container_id,
anthropic_code_exec_container_id = payload.anthropic_code_exec_container_id,
prompt_cache_ttl = payload.prompt_cache_ttl,
compaction_threshold = payload.compaction_threshold,
tools = payload.tools,
tool_choice = payload.tool_choice,
fast_mode = payload.fast_mode,
stream = payload.stream,
)
try:
sent_done = False
stream_failed = False
async for line in gen:
monitor_event = _monitor_openai_sse_line(monitor_id, line)
if monitor_event is None:
try:
_monitor_openai_chunk(monitor_id, json.loads(line))
except Exception:
pass
if monitor_event == "error":
stream_failed = True
yield f"{line}\n\n"
if monitor_event == "done":
sent_done = True
if not sent_done:
if not stream_failed:
api_monitor.finish(monitor_id)
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as exc:
logger.error("external_provider.stream_error", error = str(exc))
api_monitor.fail(monitor_id, _friendly_error(exc))
# Surface the failure: a bare EOF (e.g. after a read timeout) is treated
# by the chat client as success, saving a partial answer with no error.
yield (
"data: "
+ json.dumps({"error": {"message": _friendly_error(exc), "type": "server_error"}})
+ "\n\n"
)
yield "data: [DONE]\n\n"
finally:
try:
await gen.aclose()
except RuntimeError:
pass # suppress httpcore asyncgen cleanup error (Python 3.13 + httpcore 1.0.x)
await client.close()
return StreamingResponse(
_stream(),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"X-Accel-Buffering": "no",
},
)
# ── OpenAI shell-tool container management ───────────────────────
def _resolve_openai_cloud_client(body: OpenAIContainerRequest) -> ExternalProviderClient:
"""
Decrypt the API key + validate the base URL points at OpenAI cloud, then
build an ExternalProviderClient for the three container CRUD endpoints
below. The shell tool only exists on api.openai.com, so rejecting non-cloud
bases up front prevents confusing 404s on ollama / llama.cpp / vLLM /
custom presets.
"""
base_url = body.provider_base_url or get_base_url("openai")
if not base_url or "api.openai.com" not in base_url:
raise HTTPException(
status_code = 400,
detail = (
"OpenAI container management is only available on the "
"managed cloud (api.openai.com). The provider's base URL "
f"points at {base_url!r}."
),
)
try:
api_key = decrypt_api_key(body.encrypted_api_key)
except Exception as exc:
logger.warning("external_provider.decrypt_failed", error = str(exc))
raise HTTPException(
status_code = 400,
detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.",
)
return ExternalProviderClient(
provider_type = "openai",
base_url = base_url,
api_key = api_key,
)
def _summarize_container(raw: dict) -> OpenAIContainerSummary:
expires = raw.get("expires_after")
expires_minutes: Optional[int] = None
if isinstance(expires, dict):
minutes = expires.get("minutes")
if isinstance(minutes, int):
expires_minutes = minutes
return OpenAIContainerSummary(
id = str(raw.get("id") or ""),
name = raw.get("name"),
created_at = raw.get("created_at") if isinstance(raw.get("created_at"), int) else None,
last_active_at = raw.get("last_active_at")
if isinstance(raw.get("last_active_at"), int)
else None,
expires_after_minutes = expires_minutes,
status = raw.get("status") if isinstance(raw.get("status"), str) else None,
)
@router.post(
"/external/openai/containers/list",
response_model = ListOpenAIContainersResponse,
)
async def list_openai_containers(
body: OpenAIContainerRequest, current_subject: str = Depends(get_current_subject)
) -> ListOpenAIContainersResponse:
"""List the user's OpenAI shell-tool containers."""
client = _resolve_openai_cloud_client(body)
try:
try:
raw = await client.list_openai_containers()
except httpx.HTTPStatusError as exc:
detail = exc.response.text[:500] if exc.response is not None else str(exc)
raise HTTPException(
status_code = exc.response.status_code if exc.response else 502,
detail = f"OpenAI rejected /containers list: {detail}",
)
except httpx.HTTPError as exc:
raise log_and_http_error(
exc,
502,
"Could not reach OpenAI.",
event = "openai_container_list.transport_error",
log = logger,
)
# OpenAI keeps expired containers in /v1/containers indefinitely with
# status="expired" -- dead but still listed. Hide them so the picker
# only shows usable containers.
return ListOpenAIContainersResponse(
containers = [
_summarize_container(c)
for c in raw
if isinstance(c, dict) and c.get("status") != "expired"
],
)
finally:
await client.close()
@router.post(
"/external/openai/containers/create",
response_model = OpenAIContainerSummary,
)
async def create_openai_container(
body: CreateOpenAIContainerBody, current_subject: str = Depends(get_current_subject)
) -> OpenAIContainerSummary:
"""Create a named container with the user-chosen idle TTL."""
client = _resolve_openai_cloud_client(body)
try:
try:
raw = await client.create_openai_container(
name = body.name,
ttl_minutes = body.ttl_minutes,
)
except httpx.HTTPStatusError as exc:
detail = exc.response.text[:500] if exc.response is not None else str(exc)
raise HTTPException(
status_code = exc.response.status_code if exc.response else 502,
detail = f"OpenAI rejected /containers create: {detail}",
)
except httpx.HTTPError as exc:
raise log_and_http_error(
exc,
502,
"Could not reach OpenAI.",
event = "openai_container_create.transport_error",
log = logger,
)
if not isinstance(raw, dict):
raise HTTPException(
status_code = 502,
detail = "OpenAI returned an unexpected container payload.",
)
return _summarize_container(raw)
finally:
await client.close()
@router.post("/external/openai/containers/delete", status_code = 204)
async def delete_openai_container(
body: DeleteOpenAIContainerBody, current_subject: str = Depends(get_current_subject)
) -> None:
"""Delete a named container by id."""
logger.info(
"openai_container_delete.request subject=%s container_id=%s base_url=%s",
current_subject,
body.container_id,
body.provider_base_url,
)
client = _resolve_openai_cloud_client(body)
try:
try:
await client.delete_openai_container(body.container_id)
logger.info(
"openai_container_delete.success container_id=%s",
body.container_id,
)
except httpx.HTTPStatusError as exc:
detail = exc.response.text[:500] if exc.response is not None else str(exc)
logger.warning(
"openai_container_delete.openai_rejected container_id=%s status=%s body=%s",
body.container_id,
exc.response.status_code if exc.response else None,
detail,
)
raise HTTPException(
status_code = exc.response.status_code if exc.response else 502,
detail = f"OpenAI rejected /containers delete: {detail}",
)
except httpx.HTTPError as exc:
raise log_and_http_error(
exc,
502,
"Could not reach OpenAI.",
event = "openai_container_delete.transport_error",
log = logger,
)
finally:
await client.close()
@router.post("/chat/completions")
async def openai_chat_completions(
payload: ChatCompletionRequest,
request: Request,
current_subject: str = Depends(get_current_subject),
):
"""
OpenAI-compatible chat completions endpoint.
Supports multimodal messages: ``content`` may be a plain string or a list
of content parts (``text`` / ``image_url``).
Non-streaming (default): returns a single ChatCompletion JSON object.
Streaming: returns SSE chunks matching OpenAI's format.
``stream`` defaults to ``false`` per OpenAI's spec; clients opt into SSE by
sending ``stream: true``.
Routes to the correct backend automatically:
- 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:
# External provider: this request won't touch the local GGUF, so drop it
# from the keep-warm count or its in-flight stream would falsely block a
# concurrent local auto-switch with model_switch_busy.
from core.inference.llama_keepwarm import untrack_current_request
untrack_current_request(request.scope)
# Bypass Permissions suppresses the confirm gate, so do not reject a
# request that sets both flags (effective confirm is then False).
if (
payload.confirm_tool_calls
and not payload.bypass_permissions
and (
payload.enable_tools is True
or bool(payload.enabled_tools)
or bool(payload.tools)
or bool(payload.openai_code_exec_container_id)
or bool(payload.anthropic_code_exec_container_id)
)
):
raise HTTPException(
status_code = 400,
detail = openai_error_body(
"confirm_tool_calls is only supported for local streaming tools.",
status = 400,
code = "invalid_request_error",
param = "confirm_tool_calls",
),
)
if _wants_multiple_choices(payload):
_raise_unsupported_n("external provider chat completions")
return await _proxy_to_external_provider(payload, request, current_subject)
# 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",
),
)
# Reject a system-only chat before any automatic load so an invalid request
# never swaps or reloads the resident model (as /responses and /messages
# already validate before switching). Gate on every automatic-load trigger,
# not just auto-switch, since a standalone idle TTL can also reload here.
# Parse once and reuse below.
_pre_parsed = None
_needs_vision = False
if _automatic_model_load_may_run():
_pre_parsed = _extract_content_parts(payload.messages)
if not _pre_parsed[1]:
raise HTTPException(
status_code = 400, detail = "At least one non-system message is required."
)
# Reject confirm-without-stream local tool requests before the switch: the
# local tool path requires stream=true for the confirm gate, so this shape
# is invalid and must not evict the resident model first. Mirror that path's
# enablement exactly (_effective_enable_tools honors a CLI --enable-tools
# policy hard-override; mcp_enabled opens the tool loop on its own but still
# defers to a CLI --disable-tools policy), or an mcp_enabled/policy-forced
# request would slip past this guard and only 400 after the swap.
from state.tool_policy import get_tool_policy as _get_confirm_tool_policy
_confirm_cli_policy = _get_confirm_tool_policy()
if (
payload.confirm_tool_calls
and not payload.bypass_permissions
and not payload.stream
and (
_effective_enable_tools(payload)
or (bool(payload.mcp_enabled) and _confirm_cli_policy is not False)
or bool(payload.enabled_tools)
or bool(payload.tools)
or bool(payload.openai_code_exec_container_id)
or bool(payload.anthropic_code_exec_container_id)
)
):
raise HTTPException(
status_code = 400,
detail = openai_error_body(
"confirm_tool_calls requires stream=true for local tool execution.",
status = 400,
code = "invalid_request_error",
param = "confirm_tool_calls",
),
)
# Reject a malformed tool_choice forcing object before the switch: a
# {"type": "function", "function": {}} with no name would otherwise be
# forwarded to llama-server and rejected only after the model swapped.
_tc = payload.tool_choice
if isinstance(_tc, dict) and _tc.get("type") == "function":
_tc_fn = _tc.get("function")
_tc_name = _tc_fn.get("name") if isinstance(_tc_fn, dict) else None
if not isinstance(_tc_name, str) or not _tc_name.strip():
raise HTTPException(
status_code = 400,
detail = openai_error_body(
"Invalid 'tool_choice': the forced function must have a 'name'.",
status = 400,
code = "invalid_value",
param = "tool_choice",
),
)
# Reject an oversized audio upload before the switch: the size cap is a
# cheap, target-independent length check, so a too-large payload must not
# load a GGUF only to 413 afterward (the decode itself stays post-switch to
# avoid decoding a valid upload twice).
if payload.audio_base64 and len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS:
raise HTTPException(status_code = 413, detail = "Audio file is too large (max ~25 MB).")
# Reject streaming n>1 before the switch: only the non-streaming GGUF path
# returns multiple choices, so stream=true + n>1 is invalid on every local
# serving path (the external path already rejected it before its early
# return). Both fields are known here, so a bad shape must not load model B
# only to 400. The non-streaming n>1 cases stay post-switch, where the
# serving path decides whether the shape is supported.
if payload.stream and _wants_multiple_choices(payload):
_raise_unsupported_n("streaming chat completions")
# Audio input rides the same companion-mmproj projector as vision, so a
# text-only target can't serve it either; guard both before the switch.
_needs_vision = (
bool(_pre_parsed[2]) or _request_has_image(payload) or bool(payload.audio_base64)
)
await _maybe_auto_switch_model(
_switch_model_for_payload(payload),
request,
current_subject,
require_vision = _needs_vision,
)
llama_backend = get_llama_cpp_backend()
using_gguf = llama_backend.is_loaded
# OpenAI-SDK clients send ``chat_template_kwargs`` via ``extra_body``, which
# the SDK spreads into the request body at the top level. Studio's
# ChatCompletionRequest has ``extra="allow"`` so pydantic stashes them in
# ``model_extra``, but downstream generators consume the typed
# ``payload.enable_thinking``. Lift ``enable_thinking`` from the extra-body
# chat_template_kwargs onto the typed field so clients that only know the
# OpenAI shape (data_designer recipe runs, etc.) can still control the
# reasoning preamble.
_extra = getattr(payload, "model_extra", None)
if payload.enable_thinking is None and isinstance(_extra, dict):
_tpl_kw = _extra.get("chat_template_kwargs")
if isinstance(_tpl_kw, dict) and "enable_thinking" in _tpl_kw:
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.
monitor_id = None
async def _monitored_generate_audio(model_label: str, context_length: Optional[int] = None):
tts_monitor_id = None
if not getattr(request.state, "skip_api_monitor", False):
tts_monitor_id = api_monitor.start(
endpoint = request.url.path,
method = request.method,
model = model_label,
prompt = _monitor_prompt_from_messages(payload.messages),
context_length = context_length,
subject = current_subject,
)
try:
response = await generate_audio(payload, request)
except asyncio.CancelledError:
api_monitor.finish(tts_monitor_id, "cancelled")
raise
except Exception as e:
api_monitor.fail(tts_monitor_id, _friendly_error(e))
raise
if isinstance(response, JSONResponse):
try:
body = json.loads(response.body.decode())
choices = body.get("choices") or []
message = (choices[0].get("message") or {}) if choices else {}
content = message.get("content")
if isinstance(content, str):
api_monitor.set_reply(tts_monitor_id, content)
except Exception:
pass
api_monitor.finish(tts_monitor_id)
return response
if using_gguf:
# Advertised repo id after an auto-switch load, else a clean public id,
# never the absolute .gguf path.
model_name = _llama_public_model_id(llama_backend, payload.model)
if getattr(llama_backend, "_is_audio", False):
if _wants_multiple_choices(payload):
_raise_unsupported_n("GGUF audio chat completions")
return await _monitored_generate_audio(
model_name,
context_length = llama_backend.context_length,
)
else:
backend = get_inference_backend()
if not backend.active_model_name:
raise HTTPException(
status_code = 400,
detail = "No model loaded. Call POST /inference/load first.",
)
# Clean public id so the response never echoes a local path; the audio
# branch below receives this sanitized label too.
model_name = public_model_id(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)
model_info = backend.models.get(backend.active_model_name, {})
if model_info.get("is_audio") and model_info.get("audio_type") != "whisper":
return await _monitored_generate_audio(model_name)
# ── Whisper without audio: return clear error ──
if model_info.get("audio_type") == "whisper" and not payload.audio_base64:
raise HTTPException(
status_code = 400,
detail = "Whisper models require audio input. Please upload an audio file.",
)
if not getattr(request.state, "skip_api_monitor", False):
monitor_id = api_monitor.start(
endpoint = request.url.path,
method = request.method,
model = model_name,
prompt = _monitor_prompt_from_messages(payload.messages),
context_length = _monitor_context_length(),
subject = current_subject,
)
# ── Audio INPUT path: decode WAV and route to audio input generation ──
if payload.audio_base64 and model_info.get("has_audio_input"):
try:
audio_array = _decode_audio_base64(payload.audio_base64)
system_prompt, chat_messages, _ = _extract_content_parts(payload.messages)
except Exception as e:
api_monitor.fail(monitor_id, _friendly_error(e))
raise
cancel_event = threading.Event()
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created = int(time.time())
def audio_input_generate():
if model_info.get("audio_type") == "whisper":
return backend.generate_whisper_response(
audio_array = audio_array,
cancel_event = cancel_event,
)
return backend.generate_audio_input_response(
messages = chat_messages,
system_prompt = system_prompt,
audio_array = audio_array,
temperature = payload.temperature,
top_p = payload.top_p,
top_k = payload.top_k,
min_p = payload.min_p,
max_new_tokens = _effective_max_tokens(payload) or 2048,
repetition_penalty = payload.repetition_penalty,
cancel_event = cancel_event,
)
if payload.stream:
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_tracker = _TrackedCancel(cancel_event, *_cancel_keys)
_tracker.__enter__()
async def audio_input_stream():
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
try:
yield _chat_role_chunk(completion_id, created, model_name)
gen = audio_input_generate()
_DONE = object()
cancelled = False
while True:
if cancel_event.is_set():
cancelled = True
break
if await request.is_disconnected():
cancel_event.set()
api_monitor.finish(monitor_id, "cancelled")
return
chunk_text = await asyncio.to_thread(next, gen, _DONE)
if chunk_text is _DONE:
break
if chunk_text:
api_monitor.append_reply(monitor_id, chunk_text)
yield _chat_content_chunk(
completion_id, created, model_name, chunk_text
)
api_monitor.finish(monitor_id, "cancelled" if cancelled else "completed")
yield _chat_final_chunk(completion_id, created, model_name, "stop")
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as e:
logger.error(f"Error during audio input streaming: {e}", exc_info = True)
api_monitor.fail(monitor_id, _friendly_error(e))
yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n"
finally:
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
_tracker.__exit__(None, None, None)
return _SameTaskStreamingResponse(
audio_input_stream(),
unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
else:
try:
full_text = "".join(audio_input_generate())
except Exception as e:
api_monitor.fail(monitor_id, _friendly_error(e))
raise
api_monitor.set_reply(monitor_id, full_text)
api_monitor.finish(monitor_id)
response = ChatCompletion(
id = completion_id,
created = created,
model = model_name,
choices = [
CompletionChoice(
message = CompletionMessage(content = full_text),
finish_reason = "stop",
)
],
)
return _model_json_response(response)
if monitor_id is None and not getattr(request.state, "skip_api_monitor", False):
monitor_id = api_monitor.start(
endpoint = request.url.path,
method = request.method,
model = model_name,
prompt = _monitor_prompt_from_messages(payload.messages),
context_length = _monitor_context_length(),
subject = current_subject,
)
# Finalize the monitor entry on validation rejection before raising.
def _reject(status_code: int, detail: Any) -> "HTTPException":
if monitor_id is not None:
fail_detail = detail if isinstance(detail, str) else json.dumps(detail, default = str)
api_monitor.fail(monitor_id, fail_detail)
return HTTPException(status_code = status_code, detail = detail)
def _reject_unsupported_n(path_label: str) -> "HTTPException":
return _reject(
400,
openai_error_body(
f"n > 1 is not supported for {path_label}.",
status = 400,
code = "unsupported_parameter",
param = "n",
),
)
# ── 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 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_openai_max_tokens(payload)
normalized_stop = _normalize_stop_sequences(payload.stop)
_has_tool_messages = _has_openai_tool_history(payload.messages)
# Route guided-decoding requests through the verbatim passthrough so
# ``response_format`` (JSON schema) reaches llama-server and the model's
# GBNF-constrained output comes back unmodified. The non-passthrough GGUF
# path below calls ``generate_chat_completion`` which has no response_format
# kwarg, so the schema gets silently dropped and data_designer falls back to
# free-form sampling. Guided decoding does not require ``supports_tools`` --
# the grammar machinery is independent of tool-call parsing.
_has_response_format = _extract_response_format(payload) is not None
_has_tool_catalog = bool(payload.tools and len(payload.tools) > 0)
_has_active_tool_catalog = _has_tool_catalog and payload.tool_choice != "none"
_has_client_tool_contract = _has_active_tool_catalog or _has_tool_messages
# The Studio tool loop needs a tool-capable backend, so a request that asks
# for it on a backend that can't run it (DiffusionGemma forces supports_tools
# off) must not steal client tools from the passthrough (#6851).
_studio_tool_loop_requested = (
_explicit_studio_tool_loop_requested(payload) and llama_backend.supports_tools
)
_client_disabled_tool_calls = payload.tool_choice == "none" and not _studio_tool_loop_requested
_supports_tool_passthrough = getattr(
llama_backend, "supports_tool_passthrough", llama_backend.supports_tools
)
_tools_passthrough = _supports_tool_passthrough and _has_client_tool_contract
if (
using_gguf
and not _studio_tool_loop_requested
and _has_client_tool_contract
and not _supports_tool_passthrough
):
raise _reject(
400,
openai_error_body(
(
"Client-supplied tools or tool-call history require a GGUF chat template "
"with tool-call support; the current model/template does not advertise tools."
),
status = 400,
code = "unsupported_parameter",
param = "tools" if payload.tools else "messages",
),
)
if (
using_gguf
and not _studio_tool_loop_requested
and (_tools_passthrough or _has_response_format)
):
if _wants_multiple_choices(payload):
raise _reject_unsupported_n("GGUF tool or response_format passthrough")
if payload.audio_base64:
# This path forwards the request verbatim, so the transcoded audio
# never gets injected. (The agentic tool loop below does support
# audio.)
raise _reject(
400,
"Audio input is not supported together with guided decoding or client-supplied tools yet.",
)
# Preserve the vision guard from the non-passthrough path below:
# text-only tool-capable GGUFs should return a clear 400 here rather
# than forwarding the image to llama-server and surfacing an opaque
# upstream error.
if not llama_backend.is_vision and (
payload.image_base64
or any(
isinstance(m.content, list)
and any(isinstance(p, ImageContentPart) for p in m.content)
for m in payload.messages
)
):
raise _reject(
400,
"Image provided but current GGUF model does not support vision.",
)
cancel_event = threading.Event()
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
# `stream` defaults to False on ChatCompletionRequest (OpenAI spec
# parity). Naive curl / .NET / System.Text.Json clients omitting the
# field used to get SSE here and choke on deserialization (#5047).
if payload.stream:
return await _openai_passthrough_stream(
request,
cancel_event,
llama_backend,
payload,
model_name,
completion_id,
monitor_id = monitor_id,
)
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_tracker = _TrackedCancel(cancel_event, *_cancel_keys)
_tracker.__enter__()
try:
return await _openai_passthrough_non_streaming(
llama_backend,
payload,
model_name,
monitor_id = monitor_id,
request = request,
cancel_event = cancel_event,
)
finally:
_tracker.__exit__(None, None, None)
# ── Parse messages (handles multimodal content parts) ─────
# Reuse the pre-hook parse when auto-switch did it, else parse now.
if _pre_parsed is not None:
system_prompt, chat_messages, extracted_image_b64 = _pre_parsed
else:
system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages)
if not chat_messages:
raise _reject(400, "At least one non-system message is required.")
# ── GGUF path: proxy to llama-server /v1/chat/completions ──
if using_gguf:
# Forward uploaded audio as an input_audio part. wav/mp3 pass through
# untouched (llama-server decodes and resamples them via the mmproj
# audio encoder); other containers are transcoded to WAV here. The part
# is injected into the message list below so it rides through both the
# plain and tool-calling paths, exactly like image_url parts.
audio_b64 = None
audio_format = "wav"
if payload.audio_base64:
if not getattr(llama_backend, "_has_audio_input", False):
raise _reject(
400,
"Audio provided but current GGUF model does not support audio input.",
)
if len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS:
raise _reject(413, "Audio file is too large (max ~25 MB).")
try:
audio_b64, audio_format = await asyncio.to_thread(
_prepare_audio_for_llama, payload.audio_base64
)
except Exception as e:
logger.warning("Audio decode failed: %s", e, exc_info = True)
raise _reject(400, "Could not decode the provided audio file.")
gguf_messages, _ = _openai_messages_for_gguf_chat(
payload,
llama_backend.is_vision,
)
gguf_messages = _set_or_prepend_system_message(gguf_messages, system_prompt)
image_b64 = None
if audio_b64:
_inject_audio_part(gguf_messages, audio_b64, audio_format)
cancel_event = threading.Event()
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created = int(time.time())
def _new_chat_reasoning_extractor():
return _ResponsesReasoningExtractor(
parse_think_markers = _responses_should_parse_think_markers(
payload,
llama_backend,
)
)
def _gguf_chat_delta_line(delta: ChoiceDelta, finish_reason = None) -> str:
if delta.reasoning_content is not None and delta.content is None:
delta = delta.model_copy(update = {"content": ""})
chunk = ChatCompletionChunk(
id = completion_id,
created = created,
model = model_name,
choices = [
ChunkChoice(
delta = delta,
finish_reason = finish_reason,
)
],
)
return f"data: {chunk.model_dump_json(exclude_none = True)}\n\n"
# ── Tool-calling path (agentic loop) ──────────────────
# `_effective_enable_tools` lets `unsloth run --enable-tools/--disable-tools`
# hard-override the per-request value, else falls back to
# `payload.enable_tools`. `mcp_enabled=true` also opens the tool loop so
# MCP-only callers needn't flip a second flag, BUT must still honor a
# CLI `--disable-tools` policy -- checking the raw policy here keeps
# `mcp_enabled` from re-enabling tools the operator explicitly forbade.
from state.tool_policy import get_tool_policy as _get_tool_policy_g
_cli_policy = _get_tool_policy_g()
_tools_on = False if _client_disabled_tool_calls else _effective_enable_tools(payload)
_mcp_allowed = (
not _client_disabled_tool_calls
and bool(payload.mcp_enabled)
and _cli_policy is not False
)
use_tools = (_tools_on or _mcp_allowed) and llama_backend.supports_tools
if use_tools:
tools_to_use = await _select_request_tools(
payload, tools_on = _tools_on, mcp_allowed = _mcp_allowed
)
# Skip the tool loop when no tool survived, so the safetensors
# loop's "empty = allow all" semantic can't reach built-in tools
# the caller didn't opt into. Callers who omit enabled_tools still
# get ALL_TOOLS here, so this only suppresses the loop when
# discovery + opt-in left it genuinely empty.
if not tools_to_use:
use_tools = False
if use_tools:
# Bypass Permissions suppresses confirm, so the stream requirement
# (the gate needs streaming to prompt) no longer applies.
if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream:
raise _reject(
400,
openai_error_body(
"confirm_tool_calls requires stream=true for local tool execution.",
status = 400,
code = "invalid_request_error",
param = "confirm_tool_calls",
),
)
if _wants_multiple_choices(payload):
raise _reject_unsupported_n("GGUF tool chat completions")
# ── Tool-use system prompt nudge ──────────────────────
_nudge = _build_tool_action_nudge(
tools = tools_to_use,
model_name = model_name,
)
# Nudge the model to ground in attached documents instead of memory.
_nudge = _apply_rag_nudge(_nudge, tools_to_use, rag_scope = payload.rag_scope)
if _nudge:
# Append nudge to system prompt (preserve user's prompt)
if system_prompt:
system_prompt = system_prompt.rstrip() + "\n\n" + _nudge
else:
system_prompt = _nudge
gguf_messages = _set_or_prepend_system_message(gguf_messages, system_prompt)
_gguf_auto_heal_tool_calls = (
payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True
)
# Active tool names gating the bare-rehearsal strip, matching the loop gate.
_gguf_display_tool_names = _display_tool_name_gate(tools_to_use)
# ── Strip stale tool-call XML from conversation history ─
for _msg in gguf_messages:
if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str):
# Gate on enabled tool names, like the live strip, so a documented inactive
# ``foo[ARGS]{...}`` survives in the replayed prompt context.
_msg["content"] = _strip_tool_xml_for_display(
_msg["content"],
auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
enabled_tool_names = _gguf_display_tool_names,
).strip()
def gguf_generate_with_tools():
return llama_backend.generate_chat_completion_with_tools(
messages = gguf_messages,
tools = tools_to_use,
temperature = payload.temperature,
top_p = payload.top_p,
top_k = payload.top_k,
min_p = payload.min_p,
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,
auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
nudge_tool_calls = payload.nudge_tool_calls,
max_tool_iterations = payload.max_tool_calls_per_message
if payload.max_tool_calls_per_message is not None
else 25,
tool_call_timeout = payload.tool_call_timeout
if payload.tool_call_timeout is not None
else 300,
session_id = payload.session_id,
rag_scope = payload.rag_scope,
disable_parallel_tool_use = payload.parallel_tool_calls is False,
# Bypass Permissions takes precedence over the confirm gate:
# never prompt while bypassing.
confirm_tool_calls = bool(payload.confirm_tool_calls)
and not bool(payload.bypass_permissions),
bypass_permissions = bool(payload.bypass_permissions),
)
_tool_sentinel = object()
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_tracker = _TrackedCancel(cancel_event, *_cancel_keys)
_tracker.__enter__()
async def gguf_tool_stream():
gen = None
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
try:
yield _chat_role_chunk(completion_id, created, model_name)
# Iterate the sync generator in a thread so the event loop
# stays free for disconnect detection.
gen = gguf_generate_with_tools()
prev_text = ""
reasoning_extractor = _new_chat_reasoning_extractor()
_stream_usage = None
_stream_timings = None
_stream_finish = None
def _flush_reasoning_extractor():
final_reasoning, final_visible = reasoning_extractor.finish()
chunks = []
if final_reasoning:
chunks.append(
_gguf_chat_delta_line(
ChoiceDelta(reasoning_content = final_reasoning)
)
)
if final_visible:
api_monitor.append_reply(monitor_id, final_visible)
chunks.append(_gguf_chat_delta_line(ChoiceDelta(content = final_visible)))
return chunks
while True:
if cancel_event.is_set():
break
if await request.is_disconnected():
cancel_event.set()
api_monitor.finish(monitor_id, "cancelled")
return
event = await asyncio.to_thread(next, gen, _tool_sentinel)
if event is _tool_sentinel:
break
if event["type"] == "status":
# Empty status marks an iteration boundary in the
# GGUF tool loop (e.g. after a re-prompt). Reset the
# cumulative cursor so the next assistant turn
# streams cleanly.
if not event["text"]:
for chunk in _flush_reasoning_extractor():
yield chunk
prev_text = ""
reasoning_extractor = _new_chat_reasoning_extractor()
# Emit tool status as a custom SSE event (including
# empty ones to clear UI badges)
status_data = json.dumps(
{
"type": "tool_status",
"content": event["text"],
}
)
yield f"data: {status_data}\n\n"
continue
if event["type"] in ("tool_start", "tool_end"):
if event["type"] == "tool_start":
for chunk in _flush_reasoning_extractor():
yield chunk
prev_text = ""
reasoning_extractor = _new_chat_reasoning_extractor()
yield f"data: {json.dumps(event)}\n\n"
continue
if event["type"] == "metadata":
_stream_usage = event.get("usage")
_stream_timings = event.get("timings")
_stream_finish = event.get("finish_reason")
continue
if event["type"] == "reasoning_summary":
# Forward server-side reasoning timing to the UI.
yield f"data: {json.dumps(event)}\n\n"
continue
# "content" type -- cumulative text. Sanitize the full
# cumulative then diff against the last sanitized
# snapshot so cross-chunk XML tags are handled correctly.
raw_cumulative = event.get("text", "")
clean_cumulative = _strip_tool_xml_for_display(
raw_cumulative,
auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
enabled_tool_names = _gguf_display_tool_names,
)
new_text = clean_cumulative[len(prev_text) :]
prev_text = clean_cumulative
if not new_text:
continue
reasoning_delta, visible_delta = reasoning_extractor.feed(new_text)
if reasoning_delta:
yield _gguf_chat_delta_line(
ChoiceDelta(reasoning_content = reasoning_delta)
)
if visible_delta:
api_monitor.append_reply(monitor_id, visible_delta)
yield _gguf_chat_delta_line(ChoiceDelta(content = visible_delta))
for chunk in _flush_reasoning_extractor():
yield chunk
final_chunk = ChatCompletionChunk(
id = completion_id,
created = created,
model = model_name,
choices = [
ChunkChoice(
delta = ChoiceDelta(),
finish_reason = _clamp_finish_reason(_stream_finish),
)
],
)
# Emit the terminal chunk carrying finish_reason before the
# optional usage chunk and [DONE], so OpenAI-compatible
# clients can detect stop/length/tool_calls.
yield f"data: {final_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
_monitor_usage(monitor_id, _stream_usage, _monitor_context_length())
api_monitor.finish(
monitor_id, "cancelled" if cancel_event.is_set() else "completed"
)
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as e:
logger.error(f"Error during GGUF tool streaming: {e}", exc_info = True)
api_monitor.fail(monitor_id, _friendly_error(e))
# Recover if an MTP+tensor crash killed the server mid-stream.
get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
error_chunk = _openai_stream_error_chunk(e)
yield _openai_stream_error_sse(error_chunk)
finally:
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
if gen is not None:
try:
gen.close()
except (RuntimeError, ValueError):
pass
_tracker.__exit__(None, None, None)
if payload.stream:
return _SameTaskStreamingResponse(
gguf_tool_stream(),
unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
# Non-streaming JSON: drain the agentic generator into one
# ChatCompletion, like the standard GGUF `else` branch. stream:false
# with tools enabled used to return an SSE body, breaking
# non-streaming clients; `unsloth studio run --model` forces tools on
# process-wide, so plain requests reach this path (#6570).
def _drain_gguf_tool_loop():
full_text = ""
usage = None
finish = None
gen = gguf_generate_with_tools()
try:
for event in gen:
if cancel_event.is_set():
break
if event.get("type") == "metadata":
usage = event.get("usage")
finish = event.get("finish_reason")
elif event.get("type") == "content":
# Content is cumulative within a turn and resets
# between turns, so the last event holds the final
# turn's text. As in the safetensors drain, a visible
# preamble emitted before a tool call (its own earlier
# turn) isn't carried -- only the final turn is.
full_text = _strip_tool_xml_for_display(
event.get("text", ""),
auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
enabled_tool_names = _gguf_display_tool_names,
)
return full_text, usage, finish
finally:
# Close the generator on early break/cancel so the underlying
# llama-server stream socket is released, like the SSE path.
try:
gen.close()
except (RuntimeError, ValueError):
pass
try:
full_text, completion_usage, completion_finish = await asyncio.to_thread(
_drain_gguf_tool_loop
)
reasoning_text, visible_text = _extract_responses_reasoning(
full_text,
parse_think_markers = _responses_should_parse_think_markers(
payload, llama_backend
),
)
message_kwargs = {"content": visible_text}
if reasoning_text:
message_kwargs["reasoning_content"] = reasoning_text
_usage = completion_usage or {}
_prompt_tokens = _usage.get("prompt_tokens") or 0
_completion_tokens = _usage.get("completion_tokens") or 0
response = ChatCompletion(
id = completion_id,
created = created,
model = model_name,
choices = [
CompletionChoice(
message = CompletionMessage(**message_kwargs),
finish_reason = _clamp_finish_reason(completion_finish),
)
],
usage = CompletionUsage(
prompt_tokens = _prompt_tokens,
completion_tokens = _completion_tokens,
total_tokens = _prompt_tokens + _completion_tokens,
prompt_tokens_details = _prompt_tokens_details(
_usage.get("prompt_tokens_details")
),
),
)
api_monitor.set_reply(monitor_id, visible_text)
_monitor_usage(
monitor_id,
{
"prompt_tokens": _prompt_tokens,
"completion_tokens": _completion_tokens,
"total_tokens": _prompt_tokens + _completion_tokens,
},
_monitor_context_length(),
)
api_monitor.finish(
monitor_id, "cancelled" if cancel_event.is_set() else "completed"
)
return _model_json_response(response)
except Exception as e:
logger.error(f"Error during GGUF tool completion: {e}", exc_info = True)
api_monitor.fail(monitor_id, _friendly_error(e))
# Recover if an MTP+tensor crash killed the server.
get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
# 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))
finally:
_tracker.__exit__(None, None, None)
# ── Standard GGUF path (no tools) ─────────────────────
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,
temperature = payload.temperature,
top_p = payload.top_p,
top_k = payload.top_k,
min_p = payload.min_p,
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 _reject_unsupported_n("streaming GGUF chat completions")
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_tracker = _TrackedCancel(cancel_event, *_cancel_keys)
_tracker.__enter__()
async def gguf_stream_chunks():
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
gen = None
next_task = None
stream_completed = False
try:
yield _chat_role_chunk(completion_id, created, model_name)
# Iterate the sync generator in a thread so the event loop
# stays free for disconnect detection.
gen = gguf_generate()
prev_text = ""
reasoning_extractor = _new_chat_reasoning_extractor()
_stream_usage = None
_stream_timings = None
_stream_finish = None
while True:
if cancel_event.is_set():
break
if await request.is_disconnected():
cancel_event.set()
api_monitor.finish(monitor_id, "cancelled")
return
next_task = asyncio.create_task(
asyncio.to_thread(next, gen, _gguf_sentinel)
)
try:
cumulative = await asyncio.shield(next_task)
finally:
if next_task.done():
next_task = None
if cumulative is _gguf_sentinel:
break
# Capture server metadata for the final usage chunk
if isinstance(cumulative, dict):
if cumulative.get("type") == "metadata":
_stream_usage = cumulative.get("usage")
_stream_timings = cumulative.get("timings")
_stream_finish = cumulative.get("finish_reason")
elif cumulative.get("type") == "diffusion_frame":
# Diffusion frame (per-step canvas): pass through as a raw SSE line on the
# tool_status channel. No assistant text, so it never enters the cumulative diff.
yield f"data: {json.dumps(cumulative)}\n\n"
else:
logger.warning(
"gguf_stream_chunks: unexpected dict event: %s",
{k: v for k, v in cumulative.items() if k != "timings"},
)
continue
new_text = cumulative[len(prev_text) :]
prev_text = cumulative
if not new_text:
continue
reasoning_delta, visible_delta = reasoning_extractor.feed(new_text)
if reasoning_delta:
yield _gguf_chat_delta_line(
ChoiceDelta(reasoning_content = reasoning_delta)
)
if visible_delta:
api_monitor.append_reply(monitor_id, visible_delta)
yield _gguf_chat_delta_line(ChoiceDelta(content = visible_delta))
final_reasoning, final_visible = reasoning_extractor.finish()
if final_reasoning:
yield _gguf_chat_delta_line(ChoiceDelta(reasoning_content = final_reasoning))
if final_visible:
api_monitor.append_reply(monitor_id, final_visible)
yield _gguf_chat_delta_line(ChoiceDelta(content = final_visible))
# Final chunk
final_chunk = ChatCompletionChunk(
id = completion_id,
created = created,
model = model_name,
choices = [
ChunkChoice(
delta = ChoiceDelta(),
finish_reason = _clamp_finish_reason(_stream_finish),
)
],
)
# Emit the terminal chunk carrying finish_reason before the
# optional usage chunk and [DONE], so OpenAI-compatible
# clients can detect stop/length/tool_calls.
yield f"data: {final_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
_monitor_usage(monitor_id, _stream_usage, _monitor_context_length())
api_monitor.finish(
monitor_id, "cancelled" if cancel_event.is_set() else "completed"
)
stream_completed = True
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as e:
logger.error(f"Error during GGUF streaming: {e}", exc_info = True)
api_monitor.fail(monitor_id, _friendly_error(e))
error_chunk = _openai_stream_error_chunk(e)
yield _openai_stream_error_sse(error_chunk)
finally:
try:
if not stream_completed:
cancel_event.set()
task_to_drain = next_task
next_task = None
while task_to_drain is not None and not task_to_drain.done():
try:
await asyncio.shield(task_to_drain)
except asyncio.CancelledError:
cancel_event.set()
continue
except Exception:
break
if task_to_drain is not None and task_to_drain.done():
try:
task_to_drain.exception()
except (asyncio.CancelledError, Exception):
pass
if gen is not None and not stream_completed:
try:
await asyncio.to_thread(gen.close)
except (RuntimeError, ValueError):
pass
except Exception:
logger.debug(
"Error closing GGUF stream generator during cleanup",
exc_info = True,
)
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
finally:
_tracker.__exit__(None, None, None)
return _SameTaskStreamingResponse(
gguf_stream_chunks(),
unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
else:
try:
# ``n`` requests several independent completions; the single
# decode slot yields one at a time, so loop sequentially.
drain_task = None
async def _drain_cancelled_gguf_task():
if drain_task is None:
return
while not drain_task.done():
try:
await asyncio.shield(drain_task)
except asyncio.CancelledError:
cancel_event.set()
continue
except Exception:
break
if drain_task.done():
try:
drain_task.exception()
except (asyncio.CancelledError, Exception):
pass
def _drain_gguf_choices():
_n = payload.n or 1
_choices = []
_monitor_replies = []
_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
reasoning_text, visible_text = _extract_responses_reasoning(
full_text,
parse_think_markers = _responses_should_parse_think_markers(
payload,
llama_backend,
),
)
message_kwargs = {"content": visible_text}
if reasoning_text:
message_kwargs["reasoning_content"] = reasoning_text
_choices.append(
CompletionChoice(
index = _idx,
message = CompletionMessage(**message_kwargs),
finish_reason = _clamp_finish_reason(completion_finish),
)
)
_monitor_replies.append(visible_text)
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")
return (
_n,
_choices,
_monitor_replies,
_prompt_tokens,
_sum_completion,
_prompt_details,
)
drain_task = asyncio.create_task(asyncio.to_thread(_drain_gguf_choices))
(
_n,
_choices,
_monitor_replies,
_prompt_tokens,
_sum_completion,
_prompt_details,
) = await asyncio.shield(drain_task)
response = ChatCompletion(
id = completion_id,
created = created,
model = model_name,
choices = _choices,
usage = CompletionUsage(
prompt_tokens = _prompt_tokens,
completion_tokens = _sum_completion,
total_tokens = _prompt_tokens + _sum_completion,
prompt_tokens_details = _prompt_tokens_details(_prompt_details),
),
)
monitor_reply = _monitor_replies[-1] if _monitor_replies else ""
if _n > 1:
monitor_reply = "\n\n".join(
f"Choice {_idx + 1}:\n{text}" for _idx, text in enumerate(_monitor_replies)
)
api_monitor.set_reply(monitor_id, monitor_reply)
_monitor_usage(
monitor_id,
{
"prompt_tokens": _prompt_tokens,
"completion_tokens": _sum_completion,
"total_tokens": _prompt_tokens + _sum_completion,
},
_monitor_context_length(),
)
api_monitor.finish(monitor_id)
return _model_json_response(response)
except asyncio.CancelledError:
cancel_event.set()
await _drain_cancelled_gguf_task()
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as e:
logger.error(f"Error during GGUF completion: {e}", exc_info = True)
api_monitor.fail(monitor_id, _friendly_error(e))
# Recover if an MTP+tensor crash killed the server.
get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
# 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 ─────────────────────────────────
# Decode image (from content parts OR legacy field)
image_b64 = extracted_image_b64 or payload.image_base64
image = None
if image_b64:
try:
import base64
from PIL import Image
from io import BytesIO
model_info = backend.models.get(backend.active_model_name, {})
if not model_info.get("is_vision"):
raise HTTPException(
status_code = 400,
detail = "Image provided but current model is text-only. Load a vision model.",
)
image_data = base64.b64decode(image_b64)
image = Image.open(BytesIO(image_data))
image = backend.resize_image(image)
except HTTPException:
raise
except Exception as e:
raise log_and_http_error(
e,
400,
"Failed to decode image",
event = "inference.decode_image_failed",
log = logger,
)
# Classify capability flags from the loaded template.
_sf_model_info = backend.models.get(backend.active_model_name, {})
_sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template")
_sf_features = _detect_safetensors_features(backend, _sf_tpl)
# GGUF parity: enable_thinking templates prefill an unclosed <think>; split into
# reasoning_content deltas so the UI renders the block for safetensors and MLX.
_sf_parse_think = bool(
_sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on")
)
# Prefilled-open only for prefill styles with thinking on; gpt-oss uses the normal mode.
_sf_reasoning_prefilled = _sf_reasoning_prefill_mode(
_sf_features,
payload.enable_thinking,
_sf_tpl,
reasoning_effort = payload.reasoning_effort,
)
def _new_sf_reasoning_extractor():
return _ResponsesReasoningExtractor(
parse_think_markers = _sf_parse_think,
reasoning_prefilled = _sf_reasoning_prefilled,
)
cancel_event = threading.Event()
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created = int(time.time())
# ── Safetensors tool-calling path ─────────────────────────
# Mirrors the GGUF agentic loop's event shape. Disabled for vision turns
# (untested overlap with image render slot) and for gpt-oss (Harmony uses
# dedicated channels, not <tool_call> XML -- gpt-oss tools still work via
# the GGUF path).
_sf_is_gptoss = False
try:
_sf_is_gptoss = bool(hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model())
except Exception:
_sf_is_gptoss = False
_sf_tool_budget = (
payload.max_tool_calls_per_message if payload.max_tool_calls_per_message is not None else 25
)
# Match the GGUF path: mcp_enabled also opens the tool loop on its own
# but must still honor a CLI `--disable-tools` policy.
from state.tool_policy import get_tool_policy as _get_tool_policy_sf
_sf_cli_policy = _get_tool_policy_sf()
_sf_tools_on = _effective_enable_tools(payload)
_sf_mcp_allowed = bool(payload.mcp_enabled) and _sf_cli_policy is not False
_sf_use_tools = (
(_sf_tools_on or _sf_mcp_allowed)
and _sf_features.get("supports_tools", False)
and image is None
and not _sf_is_gptoss
and _sf_tool_budget > 0
)
if _sf_use_tools:
_sf_tools_to_use = await _select_request_tools(
payload, tools_on = _sf_tools_on, mcp_allowed = _sf_mcp_allowed
)
# Mirror the GGUF path: refuse to enter the tool loop when nothing
# survived, so a model-emitted built-in call can't piggy-back on the
# empty allow-list.
if not _sf_tools_to_use:
_sf_use_tools = False
if _sf_use_tools:
# Bypass Permissions suppresses confirm, so the stream requirement
# (the gate needs streaming to prompt) no longer applies.
if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream:
raise _reject(
400,
openai_error_body(
"confirm_tool_calls requires stream=true for local tool execution.",
status = 400,
code = "invalid_request_error",
param = "confirm_tool_calls",
),
)
_sf_nudge = _build_tool_action_nudge(
tools = _sf_tools_to_use,
model_name = model_name,
)
# RAG nudge, mirroring the GGUF path.
_sf_nudge = _apply_rag_nudge(_sf_nudge, _sf_tools_to_use, rag_scope = payload.rag_scope)
_sf_system_prompt = system_prompt
if _sf_nudge:
if _sf_system_prompt:
_sf_system_prompt = _sf_system_prompt.rstrip() + "\n\n" + _sf_nudge
else:
_sf_system_prompt = _sf_nudge
_sf_auto_heal_tool_calls = (
payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True
)
# Active tool names gating the bare-rehearsal strip, matching the loop gate.
_sf_display_tool_names = _display_tool_name_gate(_sf_tools_to_use)
# Strip stale tool-call XML from prior assistant turns.
_sf_chat_messages = []
for _msg in chat_messages:
if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str):
_sf_chat_messages.append(
{
**_msg,
"content": _strip_tool_xml_for_display(
_msg["content"],
auto_heal_tool_calls = _sf_auto_heal_tool_calls,
enabled_tool_names = _sf_display_tool_names,
).strip(),
}
)
else:
_sf_chat_messages.append(_msg)
# Request-scoped usage/timings receptacle (filled at gen_done).
_sf_stats_holder: dict = {}
def sf_generate_with_tools():
return backend.generate_chat_completion_with_tools(
messages = _sf_chat_messages,
tools = _sf_tools_to_use,
system_prompt = _sf_system_prompt or "",
temperature = payload.temperature,
top_p = payload.top_p,
top_k = payload.top_k,
min_p = payload.min_p,
max_tokens = effective_max_tokens,
repetition_penalty = payload.repetition_penalty,
presence_penalty = payload.presence_penalty,
cancel_event = cancel_event,
enable_thinking = payload.enable_thinking,
reasoning_effort = payload.reasoning_effort,
preserve_thinking = payload.preserve_thinking,
auto_heal_tool_calls = _sf_auto_heal_tool_calls,
nudge_tool_calls = payload.nudge_tool_calls,
max_tool_iterations = _sf_tool_budget,
tool_call_timeout = payload.tool_call_timeout
if payload.tool_call_timeout is not None
else 300,
session_id = payload.session_id,
rag_scope = payload.rag_scope,
# Bypass Permissions takes precedence over the confirm gate:
# never prompt while bypassing.
confirm_tool_calls = bool(payload.confirm_tool_calls)
and not bool(payload.bypass_permissions),
bypass_permissions = bool(payload.bypass_permissions),
use_adapter = payload.use_adapter,
stats_holder = _sf_stats_holder,
)
_sf_tool_sentinel = object()
_sf_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_sf_tracker = _TrackedCancel(cancel_event, *_sf_cancel_keys)
_sf_tracker.__enter__()
async def sf_tool_stream():
gen = None
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
try:
yield _chat_role_chunk(completion_id, created, model_name)
gen = sf_generate_with_tools()
prev_text = ""
reasoning_extractor = _new_sf_reasoning_extractor()
def _sf_flush_reasoning():
# Drain the extractor at turn/stream end (mirrors GGUF); only visible text hits the monitor.
fr, fv = reasoning_extractor.finish()
out = []
if fr:
out.append(_chat_reasoning_chunk(completion_id, created, model_name, fr))
if fv:
api_monitor.append_reply(monitor_id, fv)
out.append(_chat_content_chunk(completion_id, created, model_name, fv))
return out
while True:
if cancel_event.is_set():
backend.reset_generation_state()
break
if await request.is_disconnected():
cancel_event.set()
backend.reset_generation_state()
api_monitor.finish(monitor_id, "cancelled")
return
event = await asyncio.to_thread(next, gen, _sf_tool_sentinel)
if event is _sf_tool_sentinel:
break
if event["type"] == "status":
if not event["text"]:
# Iteration boundary: flush reasoning, then a fresh prefilled extractor for the next turn.
for _c in _sf_flush_reasoning():
yield _c
prev_text = ""
reasoning_extractor = _new_sf_reasoning_extractor()
status_data = json.dumps(
{
"type": "tool_status",
"content": event["text"],
}
)
yield f"data: {status_data}\n\n"
continue
if event["type"] in ("tool_start", "tool_end"):
if event["type"] == "tool_start":
# Flush reasoning before tool_start so the thinking block closes ahead of the card.
for _c in _sf_flush_reasoning():
yield _c
prev_text = ""
reasoning_extractor = _new_sf_reasoning_extractor()
yield f"data: {json.dumps(event)}\n\n"
continue
# Diff cumulative cleaned text against last snapshot.
raw_cumulative = event.get("text", "")
clean_cumulative = _strip_tool_xml_for_display(
raw_cumulative,
auto_heal_tool_calls = _sf_auto_heal_tool_calls,
enabled_tool_names = _sf_display_tool_names,
)
new_text = clean_cumulative[len(prev_text) :]
prev_text = clean_cumulative
if not new_text:
continue
# Split reasoning vs visible; only visible reaches the monitor.
reasoning_delta, visible_delta = reasoning_extractor.feed(new_text)
if reasoning_delta:
yield _chat_reasoning_chunk(
completion_id, created, model_name, reasoning_delta
)
if visible_delta:
api_monitor.append_reply(monitor_id, visible_delta)
yield _chat_content_chunk(completion_id, created, model_name, visible_delta)
for _c in _sf_flush_reasoning():
yield _c
yield _chat_final_chunk(completion_id, created, model_name, "stop")
# Usage chunk from the last turn, same shape as the
# GGUF tool loop's metadata. Request-scoped holder, so
# concurrent streams cannot read each other's stats.
_stats = _sf_stats_holder.get("stats")
if _stats:
usage_line = _openai_stream_usage_chunk(
payload,
completion_id,
created,
model_name,
_stats.get("usage"),
_stats.get("timings"),
)
if usage_line is not None:
yield usage_line
_monitor_usage(monitor_id, _stats.get("usage"))
api_monitor.finish(
monitor_id, "cancelled" if cancel_event.is_set() else "completed"
)
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
backend.reset_generation_state()
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception:
backend.reset_generation_state()
# Generic wire message; full trace stays in the log (CWE-209:
# transformers/torch errors may leak paths).
logger.exception("safetensors tool stream error")
api_monitor.fail(monitor_id, "An internal error occurred.")
error_chunk = {
"error": {
"message": "An internal error occurred.",
"type": "server_error",
},
}
yield _openai_stream_error_sse(error_chunk)
finally:
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
if gen is not None:
try:
gen.close()
except (RuntimeError, ValueError):
pass
_sf_tracker.__exit__(None, None, None)
if payload.stream:
return _SameTaskStreamingResponse(
sf_tool_stream(),
unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_sf_tracker),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
# Non-streaming JSON: drain the loop, build one ChatCompletion.
try:
def _drain_to_text():
full_text = ""
gen = sf_generate_with_tools()
for event in gen:
if cancel_event.is_set():
break
if event.get("type") == "content":
full_text = _strip_tool_xml_for_display(
event.get("text", ""),
auto_heal_tool_calls = _sf_auto_heal_tool_calls,
enabled_tool_names = _sf_display_tool_names,
)
return full_text
content_text = await asyncio.to_thread(_drain_to_text)
# Split prefilled <think> out of the visible answer (GGUF parity); the monitor gets visible text only.
_reasoning_text, _visible_text = _extract_responses_reasoning(
content_text,
parse_think_markers = _sf_parse_think,
reasoning_prefilled = _sf_reasoning_prefilled,
)
api_monitor.set_reply(monitor_id, _visible_text)
_stats = _sf_stats_holder.get("stats")
if _stats:
_monitor_usage(monitor_id, _stats.get("usage"))
api_monitor.finish(monitor_id, "cancelled" if cancel_event.is_set() else "completed")
_sf_msg_kwargs = {"content": _visible_text}
if _reasoning_text:
_sf_msg_kwargs["reasoning_content"] = _reasoning_text
response = ChatCompletion(
id = completion_id,
created = created,
model = model_name,
choices = [
CompletionChoice(
message = CompletionMessage(**_sf_msg_kwargs),
finish_reason = "stop",
)
],
)
return _model_json_response(response)
except asyncio.CancelledError:
cancel_event.set()
backend.reset_generation_state()
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception:
backend.reset_generation_state()
# CWE-209: generic detail; full trace in log.
logger.exception("safetensors tool completion error")
api_monitor.fail(monitor_id, "An internal error occurred.")
raise HTTPException(
status_code = 500,
detail = "An internal error occurred.",
)
finally:
_sf_tracker.__exit__(None, None, None)
# Shared generation kwargs
gen_kwargs = dict(
messages = chat_messages,
system_prompt = system_prompt,
image = image,
temperature = payload.temperature,
top_p = payload.top_p,
top_k = payload.top_k,
min_p = payload.min_p,
max_new_tokens = effective_max_tokens or 2048,
repetition_penalty = payload.repetition_penalty,
presence_penalty = payload.presence_penalty,
)
# Forward reasoning kwargs; the worker/template wrapper peels off any the
# template doesn't accept.
if payload.enable_thinking is not None:
gen_kwargs["enable_thinking"] = payload.enable_thinking
if payload.reasoning_effort is not None:
gen_kwargs["reasoning_effort"] = payload.reasoning_effort
if payload.preserve_thinking is not None:
gen_kwargs["preserve_thinking"] = payload.preserve_thinking
# ── Client-tool passthrough (safetensors + MLX) ──────────────
# Client tools (or tool-result history) without server-side tools: render
# tools into the template, generate one turn, heal text-form calls (#6801).
# supports_tools=False falls through to plain relay (GGUF gate parity).
_sf_has_tool_msgs = any(m.role == "tool" or m.tool_calls for m in payload.messages)
# Gate on _sf_use_tools (did the server-side path claim the request?), not
# raw mcp_enabled: an empty MCP registry must not silently drop client tools.
_sf_client_tools = (
not _effective_enable_tools(payload)
and not _sf_use_tools
and image is None
and not _sf_is_gptoss
and _sf_features.get("supports_tools", False)
and ((payload.tools and len(payload.tools) > 0) or _sf_has_tool_msgs)
)
_sf_heal = (
heal_gate(payload.auto_heal_tool_calls, payload.tools, payload.tool_choice)
if _sf_client_tools
else None
)
if _sf_client_tools:
# Re-derive from payload.messages so tool_calls / role="tool" history
# survives templating; fold system/developer into one leading system
# message (templates reject "developer") and clear prompt to avoid a dup.
gen_kwargs["messages"] = _set_or_prepend_system_message(
_structured_tool_history_for_local_template(
_flatten_content_parts_for_local_template(_openai_messages_for_passthrough(payload))
),
system_prompt,
)
gen_kwargs["system_prompt"] = ""
# tool_choice="none": keep history templating but advertise no tools
# (heal_gate is off, markup would relay as prose). A forced function
# narrows templating to that one schema. Both mirror the GGUF path,
# where llama-server honors tool_choice itself.
_sf_tc = payload.tool_choice
_sf_forced = None
if isinstance(_sf_tc, dict) and isinstance(_sf_tc.get("function"), dict):
_sf_forced = _sf_tc["function"].get("name")
if _sf_tc == "none":
gen_kwargs["tools"] = None
elif isinstance(_sf_forced, str):
gen_kwargs["tools"] = [
t
for t in payload.tools or []
if isinstance(t, dict)
and isinstance(t.get("function"), dict)
and t["function"].get("name") == _sf_forced
] or None
else:
gen_kwargs["tools"] = payload.tools
# Request-scoped usage/timings receptacle (filled at gen_done).
stats_holder: dict = {}
if payload.use_adapter is not None:
def generate(messages_override = None):
kw = (
gen_kwargs
if messages_override is None
else {**gen_kwargs, "messages": messages_override}
)
return backend.generate_with_adapter_control(
use_adapter = payload.use_adapter,
cancel_event = cancel_event,
stats_holder = stats_holder,
**kw,
)
else:
def generate(messages_override = None):
kw = (
gen_kwargs
if messages_override is None
else {**gen_kwargs, "messages": messages_override}
)
return backend.generate_chat_response(
cancel_event = cancel_event,
stats_holder = stats_holder,
**kw,
)
# ── Streaming response ────────────────────────────────────────
if payload.stream:
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_tracker = _TrackedCancel(cancel_event, *_cancel_keys)
_tracker.__enter__()
async def stream_chunks():
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
try:
yield _chat_role_chunk(completion_id, created, model_name)
# Client-tool passthrough: heal text-form calls on the fly
# (None => relay verbatim).
healer = StreamToolCallHealer(_sf_heal, payload.tools) if _sf_heal else None
heal_state = {"idx": 0}
prev_text = ""
# Split prefilled <think> into reasoning_content deltas (GGUF parity); single turn, serves MLX.
reasoning_extractor = _new_sf_reasoning_extractor()
# Run the sync generator in a thread pool to avoid blocking the
# event loop. Critical for compare mode: two SSE requests arrive
# concurrently but the orchestrator serializes them via
# _gen_lock; without run_in_executor the second request's
# blocking lock acquisition would freeze the entire event loop,
# stalling both streams.
_DONE = object() # sentinel for generator exhaustion
loop = asyncio.get_event_loop()
gen = generate()
while True:
if cancel_event.is_set():
backend.reset_generation_state()
break
# next(gen, _DONE) returns _DONE instead of raising
# StopIteration -- StopIteration can't propagate through
# asyncio futures (Python limitation).
cumulative = await loop.run_in_executor(None, next, gen, _DONE)
if cumulative is _DONE:
break
if await request.is_disconnected():
cancel_event.set()
backend.reset_generation_state()
api_monitor.finish(monitor_id, "cancelled")
return
new_text = cumulative[len(prev_text) :]
prev_text = cumulative
if not new_text:
continue
# Split prefilled <think> reasoning first (GGUF/MLX parity),
# then route only the visible text through the client-tool
# healer so tool markup inside a reasoning block is not promoted.
reasoning_delta, visible_delta = reasoning_extractor.feed(new_text)
if reasoning_delta:
yield _chat_reasoning_chunk(
completion_id, created, model_name, reasoning_delta
)
if visible_delta:
if healer is None:
# Monitor mirrors the verbatim relay; with healing on,
# _sf_heal_events_to_sse records the healed events instead.
api_monitor.append_reply(monitor_id, visible_delta)
yield _chat_content_chunk(
completion_id, created, model_name, visible_delta
)
else:
for line in _sf_heal_events_to_sse(
healer.feed(visible_delta),
completion_id,
created,
model_name,
heal_state,
payload.parallel_tool_calls,
monitor_id,
):
yield line
final_reasoning, final_visible = reasoning_extractor.finish()
if final_reasoning:
yield _chat_reasoning_chunk(completion_id, created, model_name, final_reasoning)
if final_visible:
if healer is None:
api_monitor.append_reply(monitor_id, final_visible)
yield _chat_content_chunk(completion_id, created, model_name, final_visible)
else:
for line in _sf_heal_events_to_sse(
healer.feed(final_visible),
completion_id,
created,
model_name,
heal_state,
payload.parallel_tool_calls,
monitor_id,
):
yield line
# A cancelled stream must not promote buffered-but-incomplete
# markup: finalize()'s allow_incomplete heal would execute a tool
# the user just cancelled. Disconnect returns earlier; "Stop" only
# sets cancel_event, so guard on it here too.
_cancelled = cancel_event.is_set()
if healer is not None and not _cancelled:
for line in _sf_heal_events_to_sse(
healer.finalize(),
completion_id,
created,
model_name,
heal_state,
payload.parallel_tool_calls,
monitor_id,
):
yield line
_finish = (
"tool_calls"
if (healer is not None and not _cancelled and healer.healed)
else "stop"
)
yield _chat_final_chunk(completion_id, created, model_name, _finish)
# Usage chunk (choices=[], usage set), same shape as the
# GGUF path so the speed popover works for MLX too.
# Request-scoped holder, so concurrent streams cannot
# read each other's stats.
_stats = stats_holder.get("stats")
if _stats:
usage_line = _openai_stream_usage_chunk(
payload,
completion_id,
created,
model_name,
_stats.get("usage"),
_stats.get("timings"),
)
if usage_line is not None:
yield usage_line
_monitor_usage(monitor_id, _stats.get("usage"))
api_monitor.finish(
monitor_id, "cancelled" if cancel_event.is_set() else "completed"
)
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
backend.reset_generation_state()
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as e:
backend.reset_generation_state()
logger.error(f"Error during OpenAI streaming: {e}", exc_info = True)
_msg = _friendly_error(e)
api_monitor.fail(monitor_id, _msg)
error_chunk = {
"error": {
"message": _msg,
"type": "server_error",
},
}
yield _openai_stream_error_sse(error_chunk)
finally:
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
_tracker.__exit__(None, None, None)
return _SameTaskStreamingResponse(
stream_chunks(),
unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
# ── Non-streaming response ────────────────────────────────────
else:
try:
full_text = ""
for token in generate():
full_text = token
# Split prefilled <think> reasoning (GGUF parity); also covers MLX via
# the shared generate(). Client-tool healing then runs on the visible
# text so tool markup inside a reasoning block is never promoted.
_reasoning_text, _visible_text = _extract_responses_reasoning(
full_text,
parse_think_markers = _sf_parse_think,
reasoning_prefilled = _sf_reasoning_prefilled,
)
# Client-tool passthrough: promote text-form calls; opt-in single
# nudge retry on unparseable tool markup.
_msg = {"role": "assistant", "content": _visible_text}
if _reasoning_text:
_msg["reasoning_content"] = _reasoning_text
_finish = "stop"
if _sf_heal:
if heal_openai_message(_msg, _sf_heal, payload.tools):
_finish = "tool_calls"
elif nudge_enabled(payload.nudge_tool_calls):
_data = {
"choices": [{"message": {"role": "assistant", "content": _visible_text}}]
}
if nudge_should_retry(_data, _sf_heal, payload.tools):
# A failed retry must not 500 the request; keep the first
# response (GGUF nudge parity). The retry's generate()
# overwrites stats_holder, so save the first attempt's stats
# and restore them if the retry is discarded.
_first_stats = stats_holder.get("stats")
try:
retry_text = ""
for token in generate(
[*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)]
):
retry_text = token
# Re-split reasoning on the retry so its visible text is
# what heals into a call (and reaches the monitor).
_retry_reasoning, _retry_visible = _extract_responses_reasoning(
retry_text,
parse_think_markers = _sf_parse_think,
reasoning_prefilled = _sf_reasoning_prefilled,
)
retry_msg = {"role": "assistant", "content": _retry_visible}
if _retry_reasoning:
retry_msg["reasoning_content"] = _retry_reasoning
if heal_openai_message(retry_msg, _sf_heal, payload.tools):
_visible_text, _msg, _finish = (
_retry_visible,
retry_msg,
"tool_calls",
)
else:
# Retry produced no healable call -> first response wins.
stats_holder["stats"] = _first_stats
except Exception as retry_exc:
logger.debug(
"Nudge retry failed; keeping first response: %s", retry_exc
)
stats_holder["stats"] = _first_stats
# parallel_tool_calls=false: cap to one call (GGUF parity).
if payload.parallel_tool_calls is False:
_tcs = _msg.get("tool_calls")
if isinstance(_tcs, list) and len(_tcs) > 1:
_msg["tool_calls"] = _tcs[:1]
response = ChatCompletion(
id = completion_id,
created = created,
model = model_name,
choices = [
CompletionChoice(
message = CompletionMessage(
content = _msg["content"],
reasoning_content = _msg.get("reasoning_content"),
tool_calls = _msg.get("tool_calls"),
),
finish_reason = _finish,
)
],
)
_monitor_reply = _msg.get("content") or ""
if _finish == "tool_calls":
_tcs = _msg.get("tool_calls") or []
_calls_text = "; ".join(
f"{(tc.get('function') or {}).get('name', '')}"
f"({(tc.get('function') or {}).get('arguments', '')})"
for tc in _tcs
)
_monitor_reply = (_msg.get("content") or "") + (
f"[tool_calls] {_calls_text}" if _calls_text else ""
)
api_monitor.set_reply(monitor_id, _monitor_reply)
_stats = stats_holder.get("stats")
if _stats:
_monitor_usage(monitor_id, _stats.get("usage"))
api_monitor.finish(monitor_id)
return _model_json_response(response)
except Exception as e:
backend.reset_generation_state()
logger.error(f"Error during OpenAI completion: {e}", exc_info = True)
api_monitor.fail(monitor_id, _friendly_error(e))
raise HTTPException(status_code = 500, detail = safe_error_detail(e))
# =====================================================================
# Sandbox file serving (/sandbox/{session_id}/{filename})
# =====================================================================
_SANDBOX_MEDIA_TYPES = {
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".gif": "image/gif",
".webp": "image/webp",
".bmp": "image/bmp",
}
@router.get("/sandbox/{session_id}/{filename}")
async def serve_sandbox_file(
session_id: str,
filename: str,
request: Request,
token: Optional[str] = None,
):
"""
Serve image files created by Python tool execution.
Accepts auth via Authorization header OR ?token= query param (needed
because <img src> cannot send custom headers).
"""
from fastapi.responses import FileResponse
# ── Authentication (header or query param) ──────────────────
await _authenticate_header_or_query(request, token)
# ── Filename sanitization ───────────────────────────────────
safe_filename = os.path.basename(filename)
if not safe_filename or safe_filename in (".", ".."):
raise HTTPException(status_code = 404, detail = "Not found")
# Defense-in-depth allowlist (clears CodeQL py/path-injection), still allowing
# names like "loss curve.png"; basename + extension + realpath below are the guards.
if not _re.fullmatch(r"[^/\\\x00-\x1f]{1,255}", safe_filename):
raise HTTPException(status_code = 404, detail = "Not found")
# ── Extension allowlist ─────────────────────────────────────
ext = os.path.splitext(safe_filename)[1].lower()
media_type = _SANDBOX_MEDIA_TYPES.get(ext)
if not media_type:
raise HTTPException(
status_code = status.HTTP_403_FORBIDDEN,
detail = "File type not allowed",
)
# ── Path containment check ──────────────────────────────────
from core.inference.tools import get_sandbox_workdir
sandbox_dir = os.path.realpath(get_sandbox_workdir(session_id))
file_path = os.path.realpath(os.path.join(sandbox_dir, safe_filename))
if file_path != sandbox_dir and not file_path.startswith(sandbox_dir + os.sep):
raise HTTPException(
status_code = status.HTTP_403_FORBIDDEN,
detail = "Access denied",
)
if not os.path.isfile(file_path):
raise HTTPException(status_code = 404, detail = "Not found")
return FileResponse(
path = file_path,
media_type = media_type,
headers = {
"Cache-Control": "private, no-store",
"X-Content-Type-Options": "nosniff",
},
)
# =====================================================================
# OpenAI-Compatible Models Listing (/models → /v1/models)
# =====================================================================
# `owned_by` marker on every /v1/models entry (loaded and available alike).
_OWNED_BY = "unsloth-studio"
def _openai_model_objects() -> list[dict]:
"""The model objects GET /v1/models exposes (one per loaded local backend).
Shared by the LIST and RETRIEVE handlers so both report the same ids and
field shape.
"""
models: list[dict] = []
_created = int(time.time())
# Check GGUF backend
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded:
# Advertise the repo id an auto-switch load recorded, not the concrete
# on-disk load path, so /v1/models never leaks a host path or lists a
# model twice (path plus repo id).
entry = {
# Advertised repo id after an auto-switch load, else a clean public id,
# never the absolute .gguf path (which leaks the host filesystem layout).
"id": _llama_public_model_id(llama_backend),
"object": "model",
"created": _created,
"owned_by": _OWNED_BY,
}
_ctx = _positive_int_or_none(getattr(llama_backend, "context_length", None))
if _ctx is not None:
entry["context_length"] = _ctx
_max_ctx = _positive_int_or_none(getattr(llama_backend, "max_context_length", None))
if _max_ctx is not None:
entry["max_context_length"] = _max_ctx
_native_ctx = _positive_int_or_none(getattr(llama_backend, "native_context_length", None))
if _native_ctx is not None:
entry["native_context_length"] = _native_ctx
models.append(entry)
# Check Unsloth backend
backend = get_inference_backend()
if backend.active_model_name:
model_info = backend.models.get(backend.active_model_name, {})
entry = {
"id": public_model_id(backend.active_model_name),
"object": "model",
"created": _created,
"owned_by": _OWNED_BY,
}
_ctx = _positive_int_or_none(model_info.get("context_length"))
if _ctx is None:
for _candidate in (
getattr(backend, "context_length", None),
getattr(backend, "max_seq_length", None),
):
_ctx = _positive_int_or_none(_candidate)
if _ctx is not None:
break
if _ctx is not None:
entry["context_length"] = _ctx
models.append(entry)
return models
# Brief cache for the local-model filesystem scan so repeated /v1/models calls
# don't rescan the HF cache and models dirs on every request.
_CATALOG_CACHE: dict = {"at": 0.0, "models": []}
_CATALOG_TTL_S = 30.0
# Per-loop lock (like _auto_switch_lock): a module-level asyncio.Lock ties its
# waiters to the loop that first awaited it, so a second event loop awaiting it
# in a multi-loop ASGI process can hang. The cache double-check keeps correctness
# even when two loops each scan once.
_catalog_locks: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
_catalog_locks_guard = threading.Lock()
def _catalog_lock() -> asyncio.Lock:
loop = asyncio.get_running_loop()
with _catalog_locks_guard:
lock = _catalog_locks.get(loop)
if lock is None:
lock = _catalog_locks[loop] = asyncio.Lock()
return lock
async def _cached_local_catalog() -> list:
"""Locally available models (models dir + HF caches + LM Studio + scan
folders), cached for a few seconds. Returns a list of LocalModelInfo.
The scan walks several directories and stats many files, so it runs in a
worker thread (asyncio.to_thread) -- calling it inline would block the event
loop and stall every concurrent request and in-flight inference stream. A
lock with a double-check collapses a burst of simultaneous /v1/models calls
into a single scan instead of one per request."""
# Validity is keyed on "at" (set only after a scan), not on list contents, so
# an empty/errored scan is still cached instead of rescanning on every poll.
now = time.monotonic()
if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S:
return _CATALOG_CACHE["models"]
async with _catalog_lock():
now = time.monotonic()
if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S:
return _CATALOG_CACHE["models"]
try:
from routes.models import collect_local_models
_CATALOG_CACHE["models"] = await asyncio.to_thread(
collect_local_models, Path("./models").resolve()
)
except Exception as exc:
logger.debug("model catalog scan failed: %s", exc)
_CATALOG_CACHE["models"] = []
# Stamp after the scan, not the pre-scan "now": a scan slower than the TTL
# would otherwise leave the cache already expired, so every waiter rescans.
_CATALOG_CACHE["at"] = time.monotonic()
return _CATALOG_CACHE["models"]
async def _openai_catalog_objects() -> list[dict]:
"""Every model the server knows about for ``GET /v1/models``: the loaded
model(s) plus locally available (downloaded/cached) models discovered by
scanning. Loaded entries keep their context fields and are marked
``loaded: true``. All ids are clean public ids (never absolute paths)."""
_created = int(time.time())
# Loaded models first (clean ids + context fields), marked loaded.
by_id: dict[str, dict] = {}
for entry in _openai_model_objects():
by_id[entry["id"]] = {**entry, "loaded": True}
# Locally available (downloaded/cached) models that are not already loaded.
# Advertise only GGUF models /v1 can actually serve (llama.cpp). GGUF-ness is
# read from the on-disk files, not model_format: the HF-cache scanner leaves
# model_format unset for GGUF snapshots, so a model_format filter would drop
# every cached GGUF. The file checks run off the loop.
from core.inference.local_model_resolver import info_has_local_gguf
catalog = await _cached_local_catalog()
servable = await asyncio.to_thread(lambda: [i for i in catalog if info_has_local_gguf(i)])
for info in servable:
cid = getattr(info, "model_id", None) or public_model_id(getattr(info, "id", None))
if not cid or cid in by_id:
continue
obj = {
"id": cid,
"object": "model",
"created": _created,
"owned_by": _OWNED_BY,
"loaded": False,
}
display = getattr(info, "display_name", None)
if display:
obj["display_name"] = display
by_id[cid] = obj
return list(by_id.values())
@router.get("/models")
async def openai_list_models(current_subject: str = Depends(get_current_subject)):
"""
OpenAI-compatible model listing endpoint (``GET /v1/models``).
Lists every model available on this server -- the loaded model(s) plus
locally available (downloaded/cached) models -- not only what is resident in
memory. Each entry carries a clean public id and a ``loaded`` flag.
"""
return {"object": "list", "data": await _openai_catalog_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 known model
(loaded or locally available), 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.
"""
from core.inference.model_ids import model_id_matches
# Loaded models resolve without a catalog scan (the common case); only build
# the full catalog -- which may hit the filesystem -- for unloaded ids. Match
# case-insensitively, like the catalog loop below and the resolver's index.
_loaded = _openai_model_objects()
for entry in _loaded:
eid = entry["id"]
if isinstance(eid, str) and eid.lower() == model_id.lower():
return {**entry, "loaded": True}
objects = await _openai_catalog_objects()
for model in objects:
# Case-insensitive to match the resolver, which lowercases its index.
mid = model.get("id")
if isinstance(mid, str) and mid.lower() == model_id.lower():
return model
# Backward compatibility: a client may still send the legacy raw identifier
# (e.g. an absolute .gguf path cached from an older /v1/models). Map it to the
# loaded model's object so it keeps working, without ever echoing the path back.
# Key each raw id to the SAME public id its /v1/models entry uses: an
# auto-switch load advertises a repo id while its identifier is the snapshot
# path, so public_model_id(path) would miss the advertised entry and 404 a
# model that is in fact loaded.
llama_backend = get_llama_cpp_backend()
backend = get_inference_backend()
raw_to_public: list[tuple[str, Optional[str]]] = []
if llama_backend.is_loaded and llama_backend.model_identifier:
raw_to_public.append(
(llama_backend.model_identifier, _llama_public_model_id(llama_backend))
)
if backend.active_model_name:
raw_to_public.append(
(backend.active_model_name, public_model_id(backend.active_model_name))
)
for raw, clean in raw_to_public:
if model_id_matches(model_id, raw):
for entry in _loaded:
if entry["id"] == clean:
return {**entry, "loaded": True}
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",
),
)
# =====================================================================
# OpenAI-Compatible Completions Proxy (/completions → /v1/completions)
# =====================================================================
def _flatten_monitor_prompt(value) -> str:
"""Flatten an OpenAI prompt/input field (str or list) into the single
string the api_monitor prompt preview expects."""
if isinstance(value, list):
return "\n".join(str(part) for part in value)
return str(value)
def _completions_prompt_present(body: dict) -> bool:
"""Whether a completions body carries a usable ``prompt`` (non-empty)."""
prompt = body.get("prompt")
if isinstance(prompt, str):
return prompt != ""
if isinstance(prompt, (list, tuple)):
return len(prompt) > 0
return prompt is not None
@router.post("/completions")
async def openai_completions(request: Request, current_subject: str = Depends(get_current_subject)):
"""
OpenAI-compatible text completions endpoint (non-chat).
Proxies to the running llama-server's ``/v1/completions``. Only available
when a GGUF model is loaded.
"""
llama_backend = get_llama_cpp_backend()
# Reject a request with no prompt before any automatic load so an invalid
# request never swaps or reloads the resident model (as chat/embeddings already
# validate before switching). Gate on every automatic-load trigger.
if _automatic_model_load_may_run():
try:
_pre = await request.json()
except (json.JSONDecodeError, ValueError):
_pre = None
if isinstance(_pre, dict):
_pre_prompt = _pre.get("prompt")
if _pre_prompt is not None and not isinstance(_pre_prompt, (str, list, tuple)):
# An object/number prompt is a deterministic client error (only a
# string or array is valid); reject it before the switch so a bad
# shape can't load a GGUF only to be rejected by llama-server after.
raise HTTPException(status_code = 400, detail = "'prompt' must be a string or array.")
if not _completions_prompt_present(_pre):
raise HTTPException(status_code = 400, detail = "'prompt' is required for completions.")
# Opt-in: load the requested local GGUF before the loaded-state check.
body = await _auto_switch_from_request_body(request, current_subject)
if not llama_backend.is_loaded:
raise HTTPException(
status_code = 503,
detail = "No GGUF model loaded. Load a GGUF model first.",
)
if not isinstance(body, dict):
# Re-read to re-raise a malformed-body error (post-503, pre-feature behavior);
# a valid non-dict body such as a list is a clean 400 rather than a 500.
body = await request.json()
if not isinstance(body, dict):
raise HTTPException(status_code = 400, detail = "Request body must be a JSON object")
_resolved_max_tokens = _effective_openai_max_tokens_from_values(body.get("max_tokens"))
body["max_tokens"] = (
_resolved_max_tokens
if _resolved_max_tokens is not None
else (llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR)
)
target_url = f"{llama_backend.base_url}/v1/completions"
is_stream = body.get("stream", False)
prompt_text = _flatten_monitor_prompt(body.get("prompt", ""))
monitor_id = api_monitor.start(
endpoint = request.url.path,
method = request.method,
model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"),
prompt = prompt_text,
context_length = llama_backend.context_length,
subject = current_subject,
)
if is_stream:
async def _stream():
# 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:
# <async_generator>" / 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 = _llama_streaming_generation_timeout(),
trust_env = False,
)
resp = None
bytes_iter = None
disconnect_event = threading.Event()
disconnect_watcher = None
try:
req = client.build_request(
"POST", target_url, json = body, headers = {"Connection": "close"}
)
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
resp = await _send_stream_with_preheader_cancel(client, req, request = request)
if resp is None:
api_monitor.finish(monitor_id, "cancelled")
return
if resp.status_code != 200:
err_bytes = await resp.aread()
err_text = err_bytes.decode("utf-8", errors = "replace")
api_monitor.fail(monitor_id, err_text[:500])
raise RuntimeError(f"llama-server returned {resp.status_code}: {err_text}")
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_close(request, resp, disconnect_event)
)
bytes_iter = resp.aiter_bytes()
buffer = b""
async for chunk in _aiter_llama_stream_items(
bytes_iter,
cancel_event = disconnect_event,
request = request,
first_token_deadline = first_token_deadline,
response = resp,
):
buffer += chunk
while b"\n\n" in buffer:
event, buffer = buffer.split(b"\n\n", 1)
_monitor_openai_sse_event(
monitor_id,
event,
llama_backend.context_length,
)
out = _cmpl_stream_event_out(event, _include_usage)
if out is not None:
yield out + b"\n\n"
if not disconnect_event.is_set() and buffer:
_monitor_openai_sse_event(
monitor_id,
buffer,
llama_backend.context_length,
)
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"
if disconnect_event.is_set():
api_monitor.finish(monitor_id, "cancelled")
return
api_monitor.finish(monitor_id)
except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e:
if not disconnect_event.is_set():
logger.error("openai_completions stream error: %s", e)
api_monitor.fail(monitor_id, _friendly_error(e))
error_chunk = _openai_stream_error_chunk(e)
yield _openai_stream_error_sse_bytes(error_chunk)
return
api_monitor.finish(monitor_id, "cancelled")
return
except asyncio.CancelledError:
disconnect_event.set()
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as e:
if disconnect_event.is_set():
api_monitor.finish(monitor_id, "cancelled")
return
logger.error("openai_completions stream error: %s", e)
api_monitor.fail(monitor_id, _friendly_error(e))
error_chunk = _openai_stream_error_chunk(e)
yield _openai_stream_error_sse_bytes(error_chunk)
return
finally:
await _aclose_stream_resources(
watchers = (disconnect_watcher,),
iterator = bytes_iter,
resp = resp,
client = client,
)
return _sse_streaming_response(_stream())
else:
try:
resp = await nonstreaming_client().post(
target_url,
json = body,
timeout = _llama_non_streaming_generation_timeout(),
)
except asyncio.CancelledError:
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as e:
api_monitor.fail(monitor_id, _friendly_error(e))
raise
if resp.status_code != 200:
api_monitor.fail(monitor_id, resp.text[:500])
raise _openai_passthrough_error(resp.status_code, resp.text)
try:
_monitor_openai_chunk(monitor_id, resp.json(), llama_backend.context_length)
except Exception:
pass
api_monitor.finish(monitor_id)
return Response(
content = _rewrite_cmpl_id(resp.content),
status_code = resp.status_code,
media_type = "application/json",
)
# =====================================================================
# OpenAI-Compatible Embeddings Proxy (/embeddings → /v1/embeddings)
# =====================================================================
def _embeddings_input_present(body: dict) -> bool:
"""Whether an embeddings body carries a usable ``input`` (non-empty)."""
inp = body.get("input")
if isinstance(inp, str):
return inp != ""
if isinstance(inp, (list, tuple)):
return len(inp) > 0
return inp is not None
@router.post("/embeddings")
async def openai_embeddings(request: Request, current_subject: str = Depends(get_current_subject)):
"""
OpenAI-compatible embeddings endpoint.
Proxies to the running llama-server's ``/v1/embeddings``. Only available
when a GGUF model is loaded.
Note: the loaded model must support pooling, else llama-server returns an
error (expected).
"""
llama_backend = get_llama_cpp_backend()
# Reject a request with no input before any automatic load so an invalid
# request never swaps or reloads the resident model (as chat/responses/messages
# already validate before switching). Gate on every automatic-load trigger,
# not just auto-switch, since a standalone idle TTL can also reload here.
if _automatic_model_load_may_run():
try:
_pre = await request.json()
except (json.JSONDecodeError, ValueError):
_pre = None
if isinstance(_pre, dict):
_pre_input = _pre.get("input")
if _pre_input is not None and not isinstance(_pre_input, (str, list, tuple)):
# An object/number input is a deterministic client error (only a
# string or array is valid); reject it before the switch so a bad
# shape can't load a GGUF only to be rejected by llama-server after.
raise HTTPException(status_code = 400, detail = "'input' must be a string or array.")
if not _embeddings_input_present(_pre):
raise HTTPException(status_code = 400, detail = "'input' is required for embeddings.")
# Embeddings is a model-bearing inference path too, so honor auto-switch. Unlike
# vision (cheaply pre-checked via a companion mmproj), GGUF pooling capability has
# no reliable pre-load probe -- is_embedding_model keys on a sentence-transformers
# modules.json a bare .gguf never has -- so embeddings auto-switch is best-effort:
# a non-embedding target switches, then llama-server returns a no-pooling error.
body = await _auto_switch_from_request_body(request, current_subject)
if not llama_backend.is_loaded:
raise HTTPException(
status_code = 503,
detail = "No GGUF model loaded. Load a GGUF model first.",
)
if not isinstance(body, dict):
# Re-read to re-raise a malformed-body error (post-503, pre-feature behavior);
# a valid non-dict body such as a list is a clean 400 rather than a 500.
body = await request.json()
if not isinstance(body, dict):
raise HTTPException(status_code = 400, detail = "Request body must be a JSON object")
target_url = f"{llama_backend.base_url}/v1/embeddings"
prompt_text = _flatten_monitor_prompt(body.get("input", ""))
monitor_id = None
if not getattr(request.state, "skip_api_monitor", False):
monitor_id = api_monitor.start(
endpoint = request.url.path,
method = request.method,
model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"),
prompt = prompt_text,
context_length = llama_backend.context_length,
subject = current_subject,
)
try:
resp = await nonstreaming_client().post(
target_url,
json = body,
timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S,
)
except asyncio.CancelledError:
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as exc:
api_monitor.fail(monitor_id, _friendly_error(exc))
raise
if resp.status_code != 200:
api_monitor.fail(monitor_id, resp.text[:500])
else:
try:
_monitor_usage(monitor_id, resp.json().get("usage"), _monitor_context_length())
except Exception:
pass
api_monitor.finish(monitor_id)
return Response(
content = resp.content,
status_code = resp.status_code,
media_type = "application/json",
)
# =====================================================================
# OpenAI Responses API (/responses → /v1/responses)
# =====================================================================
def _translate_responses_tools_to_chat(tools: Optional[list[dict]]) -> Optional[list[dict]]:
"""Translate Responses-shape function tools to the Chat Completions nested shape.
Responses uses a flat shape per tool entry::
{"type": "function", "name": "...", "description": "...",
"parameters": {...}, "strict": true}
The Chat Completions / llama-server passthrough expects the nested shape::
{"type": "function",
"function": {"name": "...", "description": "...",
"parameters": {...}, "strict": true}}
Only ``type=="function"`` entries are forwarded. Built-in Responses tools
(``web_search``, ``file_search``, ``mcp``, ...) are dropped: llama-server
doesn't implement them server-side, so keeping them would produce an opaque
upstream 400.
"""
if not tools:
return None
out: list[dict] = []
for tool in tools:
if not isinstance(tool, dict):
continue
if tool.get("type") != "function":
continue
fn: dict = {}
if "name" in tool:
fn["name"] = tool["name"]
if tool.get("description") is not None:
fn["description"] = tool["description"]
if tool.get("parameters") is not None:
fn["parameters"] = tool["parameters"]
if tool.get("strict") is not None:
fn["strict"] = tool["strict"]
out.append({"type": "function", "function": fn})
return out or None
def _translate_responses_tool_choice_to_chat(tool_choice: Any) -> Any:
"""Translate a Responses-shape ``tool_choice`` to the Chat Completions shape.
String values (``"auto"``/``"none"``/``"required"``) pass through unchanged.
The Responses forcing object ``{"type": "function", "name": "X"}`` becomes
Chat Completions' ``{"type": "function", "function": {"name": "X"}}``.
Unknown / built-in tool choices are forwarded as-is; llama-server ignores
what it doesn't recognise.
"""
if tool_choice is None:
return None
if isinstance(tool_choice, str):
return tool_choice
if (
isinstance(tool_choice, dict)
and tool_choice.get("type") == "function"
and "name" in tool_choice
and "function" not in tool_choice
):
return {"type": "function", "function": {"name": tool_choice["name"]}}
return tool_choice
def _responses_message_text(content: Union[str, list]) -> str:
"""Flatten a ResponsesInputMessage ``content`` into a plain text string.
Used for system/developer message hoisting and for assistant-replay
(``output_text``) messages when images/unknown parts are irrelevant.
Returns an empty string for empty input.
"""
if isinstance(content, str):
return content
parts: list[str] = []
for part in content or []:
if isinstance(part, (ResponsesInputTextPart, ResponsesOutputTextPart)):
parts.append(part.text)
return "\n".join(parts)
def _responses_tool_output_content(output: Union[str, list]) -> Union[str, list]:
"""Return Chat Completions-safe content for a Responses tool result."""
if isinstance(output, str):
return output if output.strip() else "(no output)"
if not output:
return "(no output)"
text_parts: list[str] = []
chat_parts: list = []
has_multimodal = False
for part in output:
if not isinstance(part, dict):
return json.dumps(output)
part_type = part.get("type")
if part_type in ("input_text", "output_text", "text"):
text = part.get("text")
if text is None:
_raise_unsupported_openai_parameter(
"input",
"Responses function_call_output.output text parts require a text field.",
)
text = str(text)
text_parts.append(text)
chat_parts.append(TextContentPart(type = "text", text = text))
continue
if part_type == "input_image":
image_url = part.get("image_url")
if not isinstance(image_url, str) or not image_url:
if part.get("file_id"):
_raise_unsupported_openai_parameter(
"input",
"Responses function_call_output.output input_image parts with file_id are not supported by the local adapter. Use image_url instead.",
)
_raise_unsupported_openai_parameter(
"input",
"Responses function_call_output.output input_image parts require an image_url string.",
)
detail = part.get("detail", "auto")
if detail is None:
detail = "auto"
if detail not in ("auto", "low", "high", "original"):
_raise_unsupported_openai_parameter(
"input",
"Responses function_call_output.output input_image detail must be auto, low, high, or original.",
)
chat_parts.append(
ImageContentPart(
type = "image_url",
image_url = ImageUrl(url = image_url, detail = detail),
)
)
has_multimodal = True
continue
if part_type == "input_file":
_raise_unsupported_openai_parameter(
"input",
"Responses function_call_output.output input_file parts are not supported by the local adapter.",
)
return json.dumps(output)
if has_multimodal:
return chat_parts
text = "\n".join(text_parts)
return text if text.strip() else "(no output)"
_RESPONSES_THINK_OPEN = "<think>"
_RESPONSES_THINK_CLOSE = "</think>"
_RESPONSES_REASONING_EFFORTS = {"none", "minimal", "low", "medium", "high", "max", "xhigh"}
def _coerce_responses_reasoning_text(value: Any) -> str:
if value is None:
return ""
if isinstance(value, str):
return value
if isinstance(value, list):
return "".join(_coerce_responses_reasoning_text(part) for part in value)
if isinstance(value, dict):
for key in ("text", "reasoning_text", "content"):
text = _coerce_responses_reasoning_text(value.get(key))
if text:
return text
return ""
return json.dumps(value)
def _responses_marker_holdback(text: str, markers: tuple[str, ...]) -> int:
"""Number of trailing chars to retain because they may start a marker."""
for size in range(min(len(text), max(len(m) for m in markers) - 1), 0, -1):
suffix = text[-size:]
if any(marker.startswith(suffix) for marker in markers):
return size
return 0
class _ResponsesReasoningExtractor:
"""Split local <think> markup into Responses reasoning and visible text."""
def __init__(
self,
*,
parse_think_markers: bool = False,
reasoning_prefilled: bool = False,
) -> None:
self._buffer = ""
# reasoning_prefilled: the template inserts an unclosed <think>, so output begins inside
# the block; start in reasoning until the first close tag. Existing callers pass False.
self._in_reasoning = reasoning_prefilled
# Splitting requires marker parsing; a prefilled open implies it.
self._parse_think_markers = parse_think_markers or reasoning_prefilled
def feed(
self,
text: str = "",
reasoning_content: Any = None,
) -> tuple[str, str]:
reasoning_parts: list[str] = []
visible_parts: list[str] = []
structured_reasoning = _coerce_responses_reasoning_text(reasoning_content)
if structured_reasoning:
reasoning_parts.append(structured_reasoning)
if text:
self._buffer += text
if not self._parse_think_markers:
visible_parts.append(self._buffer)
self._buffer = ""
return "".join(reasoning_parts), "".join(visible_parts)
while self._buffer:
if self._in_reasoning:
close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE)
if close_idx != -1:
reasoning_parts.append(
self._buffer[:close_idx].replace(_RESPONSES_THINK_OPEN, "")
)
self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :]
self._in_reasoning = False
continue
# Hold back a trailing partial of either marker: the close (clean split across chunks)
# and a stray open (a re-emitted <think> is suppressed, not leaked).
keep = _responses_marker_holdback(
self._buffer, (_RESPONSES_THINK_CLOSE, _RESPONSES_THINK_OPEN)
)
if keep == len(self._buffer):
break
emit = self._buffer[:-keep] if keep else self._buffer
reasoning_parts.append(emit.replace(_RESPONSES_THINK_OPEN, ""))
self._buffer = self._buffer[-keep:] if keep else ""
break
open_idx = self._buffer.find(_RESPONSES_THINK_OPEN)
close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE)
if close_idx != -1 and (open_idx == -1 or close_idx < open_idx):
visible_parts.append(self._buffer[:close_idx])
self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :]
continue
if open_idx != -1:
visible_parts.append(self._buffer[:open_idx])
self._buffer = self._buffer[open_idx + len(_RESPONSES_THINK_OPEN) :]
self._in_reasoning = True
continue
keep = _responses_marker_holdback(
self._buffer,
(_RESPONSES_THINK_OPEN, _RESPONSES_THINK_CLOSE),
)
if keep == len(self._buffer):
break
visible_parts.append(self._buffer[:-keep] if keep else self._buffer)
self._buffer = self._buffer[-keep:] if keep else ""
break
return "".join(reasoning_parts), "".join(visible_parts)
def finish(self) -> tuple[str, str]:
if not self._buffer:
return "", ""
remaining = self._buffer
self._buffer = ""
if not self._parse_think_markers:
return "", remaining
if self._in_reasoning:
self._in_reasoning = False
return remaining.replace(_RESPONSES_THINK_OPEN, ""), ""
return "", remaining.replace(_RESPONSES_THINK_CLOSE, "")
def _extract_responses_reasoning(
text: str = "",
reasoning_content: Any = None,
*,
parse_think_markers: bool = False,
reasoning_prefilled: bool = False,
) -> tuple[str, str]:
extractor = _ResponsesReasoningExtractor(
parse_think_markers = parse_think_markers,
reasoning_prefilled = reasoning_prefilled,
)
reasoning, visible = extractor.feed(text, reasoning_content)
final_reasoning, final_visible = extractor.finish()
return reasoning + final_reasoning, visible + final_visible
def _responses_should_parse_think_markers(
chat_req: ChatCompletionRequest, llama_backend: Any = None
) -> bool:
if llama_backend is not None and getattr(llama_backend, "is_loaded", False):
if getattr(llama_backend, "reasoning_always_on", False):
return True
if getattr(llama_backend, "supports_reasoning", False):
return True
return False
if chat_req.enable_thinking is True:
return True
return chat_req.enable_thinking is None and chat_req.reasoning_effort not in (None, "none")
def _responses_reasoning_output_item(reasoning_text: str, item_id: Optional[str] = None) -> dict:
kwargs: dict[str, Any] = {
"status": "completed",
"summary": [],
"content": [ResponsesOutputReasoningContent(text = reasoning_text)],
}
if item_id is not None:
kwargs["id"] = item_id
return ResponsesOutputReasoning(**kwargs).model_dump()
def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]:
"""Convert a ResponsesRequest's ``input`` into a Chat-format ``ChatMessage`` list.
Handles the three input item shapes allowed by the Responses API:
- ``ResponsesInputMessage`` -- regular chat messages (text or multimodal).
- ``ResponsesFunctionCallInputItem`` -- a prior assistant tool call
replayed on a follow-up turn. Becomes an assistant message carrying a
Chat Completions ``tool_calls`` entry keyed by ``call_id``.
- ``ResponsesFunctionCallOutputInputItem`` -- a tool result the client is
returning. Becomes a ``role="tool"`` message with ``tool_call_id`` set to
the originating ``call_id`` so llama-server can reconcile call with result.
System / developer content is collected from ``instructions`` *and* any
``role="system"`` / ``role="developer"`` entries in ``input``, then merged
into a single top-of-list ``role="system"`` message. This satisfies strict
chat templates (harmony / gpt-oss, Qwen3, ...) whose Jinja raises
``"System message must be at the beginning."`` when more than one system
message is present or a system message follows a user turn -- the exact
pattern the OpenAI Codex CLI hits, since Codex sets ``instructions`` *and*
also sends a developer message in ``input``.
"""
system_parts: list[str] = []
messages: list[ChatMessage] = []
if payload.instructions:
system_parts.append(payload.instructions)
def _with_system(msgs: list[ChatMessage]) -> list[ChatMessage]:
if not system_parts:
return msgs
merged = "\n\n".join(p for p in system_parts if p)
return [ChatMessage(role = "system", content = merged), *msgs]
# Simple string input
if isinstance(payload.input, str):
if payload.input:
messages.append(ChatMessage(role = "user", content = payload.input))
return _with_system(messages)
for item in payload.input:
if isinstance(item, ResponsesFunctionCallInputItem):
messages.append(
ChatMessage(
role = "assistant",
content = None,
tool_calls = [
{
"id": item.call_id,
"type": "function",
"function": {
"name": item.name,
"arguments": item.arguments,
},
}
],
)
)
continue
if isinstance(item, ResponsesFunctionCallOutputInputItem):
# Flatten pure text arrays for broad template compatibility, and
# forward image URL outputs as real multimodal parts for vision models.
output = _responses_tool_output_content(item.output)
messages.append(
ChatMessage(
role = "tool",
tool_call_id = item.call_id,
content = output,
)
)
continue
if isinstance(item, ResponsesUnknownInputItem):
# Reasoning items and other unmodelled top-level Responses item
# types are silently dropped -- llama-server-backed GGUFs can't
# consume them; lenient validation lets them in so unrelated turns
# don't 422.
continue
# ResponsesInputMessage -- hoist system/developer to the top, merge.
if item.role in ("system", "developer"):
hoisted = _responses_message_text(item.content)
if hoisted:
system_parts.append(hoisted)
continue
if isinstance(item.content, str):
messages.append(ChatMessage(role = item.role, content = item.content))
continue
# Assistant-replay turns come back as content = [output_text, ...].
# Chat Completions' assistant role expects a plain string, not a
# multimodal array, so flatten output_text (and any stray input_text /
# unknown text) to a single string.
if item.role == "assistant":
text = _responses_message_text(item.content)
if text:
messages.append(ChatMessage(role = "assistant", content = text))
continue
# User (and any other remaining roles) -- keep multimodal when present,
# drop unknown content parts silently.
parts: list = []
for part in item.content:
if isinstance(part, (ResponsesInputTextPart, ResponsesOutputTextPart)):
parts.append(TextContentPart(type = "text", text = part.text))
elif isinstance(part, ResponsesInputImagePart):
parts.append(
ImageContentPart(
type = "image_url",
image_url = ImageUrl(url = part.image_url, detail = part.detail),
)
)
# ResponsesUnknownContentPart and anything else: drop.
if parts:
# Collapse single-text-part content to a plain string so roles that
# reject multimodal arrays (e.g. legacy templates) still accept it.
if len(parts) == 1 and isinstance(parts[0], TextContentPart):
messages.append(ChatMessage(role = item.role, content = parts[0].text))
else:
messages.append(ChatMessage(role = item.role, content = parts))
return _with_system(messages)
def _build_chat_request(
payload: ResponsesRequest, messages: list[ChatMessage], stream: bool
) -> ChatCompletionRequest:
"""Build a ChatCompletionRequest from a ResponsesRequest.
Tools and ``tool_choice`` are translated from the flat Responses shape to
the nested Chat Completions shape here so the existing #5099
``/v1/chat/completions`` client-side pass-through picks them up unchanged.
"""
chat_kwargs: dict = dict(
messages = messages,
stream = stream,
)
# Only forward an explicitly set model so an omitted Responses model stays
# reload-only when openai_chat_completions re-checks on the non-streaming path.
if "model" in payload.model_fields_set:
chat_kwargs["model"] = payload.model
if payload.temperature is not None:
chat_kwargs["temperature"] = payload.temperature
if payload.top_p is not None:
chat_kwargs["top_p"] = payload.top_p
if payload.max_output_tokens is not None:
chat_kwargs["max_tokens"] = payload.max_output_tokens
chat_tools = _translate_responses_tools_to_chat(payload.tools)
if chat_tools is not None:
chat_kwargs["tools"] = chat_tools
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
# ``chat_template_kwargs`` (e.g. ``{"enable_thinking": true}``) arrives via
# the Responses extra-body: ResponsesRequest has ``extra="allow"``, so the
# OpenAI SDK's ``extra_body`` spread lands the dict in ``model_extra``. The
# downstream Chat Completions paths consume the typed ``enable_thinking``
# field -- the non-streaming path lifts it in ``openai_chat_completions``
# only when it is still ``None``, and the streaming pass-through reads
# ``payload.enable_thinking`` directly -- so lift it here, mirroring that
# handler, to cover both Responses paths.
explicit_enable_thinking = False
_extra = getattr(payload, "model_extra", None)
if isinstance(_extra, dict):
_tpl_kw = _extra.get("chat_template_kwargs")
if isinstance(_tpl_kw, dict) and "enable_thinking" in _tpl_kw:
chat_kwargs["enable_thinking"] = bool(_tpl_kw["enable_thinking"])
explicit_enable_thinking = True
# auto_heal_tool_calls / nudge_tool_calls are not typed on
# ResponsesRequest; lift them from the extra-body so passthrough
# healing (and the opt-in nudge) honor them on both paths.
if isinstance(_extra.get("auto_heal_tool_calls"), bool):
chat_kwargs["auto_heal_tool_calls"] = _extra["auto_heal_tool_calls"]
if isinstance(_extra.get("nudge_tool_calls"), bool):
chat_kwargs["nudge_tool_calls"] = _extra["nudge_tool_calls"]
if isinstance(payload.reasoning, dict):
effort = payload.reasoning.get("effort")
if isinstance(effort, str) and effort in _RESPONSES_REASONING_EFFORTS:
if not explicit_enable_thinking:
chat_kwargs["reasoning_effort"] = effort
chat_kwargs["enable_thinking"] = effort != "none"
elif chat_kwargs.get("enable_thinking") is False:
chat_kwargs["reasoning_effort"] = "none"
elif effort != "none":
chat_kwargs["reasoning_effort"] = effort
return ChatCompletionRequest(**chat_kwargs)
def _chat_tool_calls_to_responses_output(tool_calls: list[dict]) -> list[dict]:
"""Map Chat Completions ``tool_calls`` into Responses ``function_call`` output items.
The Chat Completions id (``call_xxx``) is the shared correlation key across
turns in the Responses API -- stored as ``call_id`` on the output item and
echoed back by the client as ``function_call_output.call_id`` next turn.
"""
items: list[dict] = []
for tc in tool_calls:
if tc.get("type") != "function":
continue
fn = tc.get("function") or {}
items.append(
ResponsesOutputFunctionCall(
call_id = tc.get("id", ""),
name = fn.get("name", ""),
arguments = fn.get("arguments", "") or "",
status = "completed",
).model_dump()
)
return items
async def _responses_non_streaming(
payload: ResponsesRequest,
messages: list[ChatMessage],
request: Request,
current_subject: Optional[str] = None,
) -> JSONResponse:
"""Handle a non-streaming Responses API call."""
chat_req = _build_chat_request(payload, messages, stream = False)
request_state = getattr(request, "state", None)
if request_state is None:
request_state = type("_RequestState", (), {})()
try:
setattr(request, "state", request_state)
except Exception:
request_state = None
previous_skip_monitor = (
bool(getattr(request_state, "skip_api_monitor", False))
if request_state is not None
else False
)
monitor_id = None
if not previous_skip_monitor:
monitor_id = api_monitor.start(
endpoint = getattr(getattr(request, "url", None), "path", "/v1/responses"),
method = getattr(request, "method", "POST"),
model = payload.model,
prompt = _monitor_prompt_from_messages(messages),
context_length = _monitor_context_length(),
subject = current_subject,
)
if request_state is not None:
request_state.skip_api_monitor = True
try:
result = await openai_chat_completions(chat_req, request)
# openai_chat_completions returns a JSONResponse for non-streaming.
if isinstance(result, Response):
body = json.loads(result.body.decode())
else:
body = result
choices = body.get("choices", [])
text = ""
reasoning_text = ""
tool_calls: list[dict] = []
if choices:
msg = choices[0].get("message", {}) or {}
raw_content = msg.get("content", "") or ""
raw_text = raw_content if isinstance(raw_content, str) else json.dumps(raw_content)
llama_backend = get_llama_cpp_backend()
reasoning_text, text = _extract_responses_reasoning(
raw_text,
msg.get("reasoning_content"),
parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend),
)
tool_calls = msg.get("tool_calls") or []
usage_data = body.get("usage", {})
input_tokens = usage_data.get("prompt_tokens", 0)
output_tokens = usage_data.get("completion_tokens", 0)
resp_id = f"resp_{uuid.uuid4().hex[:12]}"
# Responses API emits each tool call as its own top-level output item,
# plus an optional assistant text message. Emit the text message only when
# the model produced content, so clients expecting a pure tool-call turn
# (finish_reason="tool_calls") don't see a spurious empty message item.
output_items: list[dict] = []
if reasoning_text:
output_items.append(_responses_reasoning_output_item(reasoning_text))
if text:
msg_id = f"msg_{uuid.uuid4().hex[:12]}"
output_items.append(
ResponsesOutputMessage(
id = msg_id,
status = "completed",
role = "assistant",
content = [ResponsesOutputTextContent(text = text)],
).model_dump()
)
output_items.extend(_chat_tool_calls_to_responses_output(tool_calls))
response = ResponsesResponse(
id = resp_id,
created_at = int(time.time()),
status = "completed",
model = body.get("model", payload.model),
output = output_items,
usage = ResponsesUsage(
input_tokens = input_tokens,
output_tokens = output_tokens,
total_tokens = input_tokens + output_tokens,
),
temperature = payload.temperature,
top_p = payload.top_p,
max_output_tokens = payload.max_output_tokens,
instructions = payload.instructions,
)
api_monitor.set_reply(monitor_id, text or _monitor_tool_calls_text(tool_calls))
_monitor_usage(monitor_id, usage_data, _monitor_context_length())
api_monitor.finish(monitor_id)
return _model_json_response(response)
except asyncio.CancelledError:
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as exc:
api_monitor.fail(monitor_id, _friendly_error(exc))
raise
finally:
if request_state is not None:
request_state.skip_api_monitor = previous_skip_monitor
async def _responses_stream(
payload: ResponsesRequest,
messages: list[ChatMessage],
request: Request,
monitor_id: Optional[str] = None,
):
"""Handle a streaming Responses API call, emitting named SSE events.
For GGUF models the request goes directly to llama-server's
``/v1/chat/completions`` from inside the StreamingResponse child task -- one
httpx lifecycle, one async generator. Wrapping the existing
``openai_chat_completions`` pass-through (which has its own httpx lifecycle)
stacks two generators: Python 3.13 + httpcore 1.0.x then loses the
close-propagation chain on the innermost ``HTTP11ConnectionByteStream`` at
asyncgen finalisation, tripping "Attempted to exit cancel scope in a
different task" / "async generator ignored GeneratorExit". The direct path
avoids that. Non-GGUF falls back to the wrapper (which doesn't use httpx, so
the issue doesn't apply).
Output items are allocated as upstream deltas appear. Reasoning/text deltas
open top-level ``reasoning`` / ``message`` items; each tool call from
``delta.tool_calls[]`` is promoted to its own top-level ``function_call``
item (one per distinct ``tool_calls[].index``) and relayed as
``response.function_call_arguments.delta`` / ``.done`` events so clients
(Codex, OpenAI Python SDK) can reconstruct the call incrementally and reply
with a ``function_call_output`` item next turn.
"""
resp_id = f"resp_{uuid.uuid4().hex[:12]}"
created_at = int(time.time())
chat_req = _build_chat_request(payload, messages, stream = True)
llama_backend = get_llama_cpp_backend()
if not llama_backend.is_loaded:
# The direct pass-through is GGUF-only. Non-GGUF /v1/responses streaming
# isn't a Codex-compatible path today, and wrapping the transformers
# backend's streaming generator here would re-introduce the
# double-layer asyncgen close pattern that produces "Attempted to exit
# cancel scope in a different task" on Python 3.13. Surface a typed 400
# so the client sees a useful error instead of a dangling stream.
raise HTTPException(
status_code = 400,
detail = (
"Streaming /v1/responses requires a GGUF model loaded via "
"llama-server. Use non-streaming /v1/responses, "
"/v1/chat/completions, or load a GGUF model."
),
)
# Direct pass-through bypasses the openai_chat_completions image gate.
if not llama_backend.is_vision and any(
isinstance(m.content, list) and any(isinstance(p, ImageContentPart) for p in m.content)
for m in messages
):
raise HTTPException(
status_code = 400,
detail = "Image provided but current GGUF model does not support vision.",
)
body = _build_openai_passthrough_body(
chat_req, backend_ctx = llama_backend.context_length, llama_backend = llama_backend
)
body["stream_options"] = {"include_usage": True}
target_url = f"{llama_backend.base_url}/v1/chat/completions"
async def event_generator():
# Clean public id for every response envelope. Prefer the loaded model's
# id so the stream agrees with /v1/models, chat/completions and the
# non-streaming twin; fall back to a sanitized payload.model (a legacy
# raw .gguf path is stripped, never echoed back). Use the advertised-id
# helper, not the raw identifier: after an auto-switch to a cached HF GGUF
# the identifier is the snapshot path while the repo id lives in
# _openai_advertised_id, so the raw form would stream a snapshot basename.
_clean_model = _llama_public_model_id(llama_backend, payload.model) or payload.model
full_text = ""
full_reasoning = ""
input_tokens = 0
output_tokens = 0
extractor = _ResponsesReasoningExtractor(
parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend)
)
reasoning_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False}
message_state: dict[str, Any] = {
"output_index": None,
"item_id": None,
"opened": False,
"text": "",
}
# Message items already closed mid-stream (a healed tool call splits
# the assistant text into separate message items, as native Responses
# streams do). Kept for the final response.completed snapshot.
closed_message_states: list[dict] = []
# Per-tool-call state keyed by Chat Completions `tool_calls[].index`,
# stable across chunks for the same call. Values:
# {output_index, item_id, call_id, name, arguments, opened}
tool_call_state: dict[int, dict] = {}
next_output_index = 0
# Text-form tool calls promoted back to structured calls (declared
# client tools only); dormant once grammar-mode structured deltas appear.
_allowed_tools = heal_gate(
getattr(chat_req, "auto_heal_tool_calls", None),
body.get("tools"),
body.get("tool_choice"),
)
healer = StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None
healed_tc_index = 0
def _healed_tc(call: dict):
# Chat-delta shape for a healed call. Indexes live in a disjoint
# range so a healed call can never merge into a structured call's
# state slot; parallel_tool_calls=false caps healed calls too (the
# upstream cap ran before injection).
nonlocal healed_tc_index
if payload.parallel_tool_calls is False and healed_tc_index >= 1:
return None
tc = {
"index": 1_000_000 + healed_tc_index,
"id": call["id"],
"type": "function",
"function": call["function"],
}
healed_tc_index += 1
return tc
def _sse(event_name: str, payload: dict) -> str:
return f"event: {event_name}\ndata: {json.dumps(payload)}\n\n"
def _tool_call_delta_events(tc: dict) -> list:
# One Chat Completions tool_calls delta -> Responses SSE events,
# allocating/merging per-call state (shared by the structured loop
# and the healer's promoted calls).
events = []
idx = tc.get("index", 0)
st = tool_call_state.get(idx)
fn = tc.get("function") or {}
if st is None:
# First chunk for this tool call -- allocate an
# output_index and emit output_item.added.
st = {
"output_index": _claim_output_index(),
"item_id": f"fc_{uuid.uuid4().hex[:12]}",
"call_id": tc.get("id") or "",
"name": fn.get("name") or "",
"arguments": "",
"opened": False,
}
tool_call_state[idx] = st
else:
# Later chunks sometimes carry id/name only once; merge
# when present.
if tc.get("id") and not st["call_id"]:
st["call_id"] = tc["id"]
if fn.get("name") and not st["name"]:
st["name"] = fn["name"]
if not st["opened"] and st["call_id"] and st["name"]:
item_added = {
"type": "response.output_item.added",
"output_index": st["output_index"],
"item": {
"type": "function_call",
"id": st["item_id"],
"status": "in_progress",
"call_id": st["call_id"],
"name": st["name"],
"arguments": "",
},
}
events.append(_sse("response.output_item.added", item_added))
st["opened"] = True
arg_delta = fn.get("arguments") or ""
if arg_delta and st["opened"]:
st["arguments"] += arg_delta
args_delta_event = {
"type": "response.function_call_arguments.delta",
"item_id": st["item_id"],
"output_index": st["output_index"],
"delta": arg_delta,
}
events.append(_sse("response.function_call_arguments.delta", args_delta_event))
elif arg_delta:
# Buffer args until we can open the item (some models
# send id/name in the same chunk as the first arg delta;
# if not, stash).
st["arguments"] += arg_delta
return events
def _claim_output_index() -> int:
nonlocal next_output_index
output_index = next_output_index
next_output_index += 1
return output_index
def _apply_usage(u) -> None:
nonlocal input_tokens, output_tokens
if not u:
return
input_tokens = u.get("prompt_tokens", input_tokens)
output_tokens = u.get("completion_tokens", output_tokens)
_monitor_usage(monitor_id, u, llama_backend.context_length)
def _ensure_reasoning_open() -> list[str]:
if reasoning_state["opened"]:
return []
reasoning_state["output_index"] = _claim_output_index()
reasoning_state["item_id"] = f"rs_{uuid.uuid4().hex[:12]}"
reasoning_state["opened"] = True
output_index = reasoning_state["output_index"]
item_id = reasoning_state["item_id"]
return [
_sse(
"response.output_item.added",
{
"type": "response.output_item.added",
"output_index": output_index,
"item": {
"type": "reasoning",
"id": item_id,
"status": "in_progress",
"summary": [],
"content": [],
},
},
),
_sse(
"response.content_part.added",
{
"type": "response.content_part.added",
"item_id": item_id,
"output_index": output_index,
"content_index": 0,
"part": {"type": "reasoning_text", "text": ""},
},
),
]
def _ensure_message_open() -> list[str]:
if message_state["opened"]:
return []
message_state["output_index"] = _claim_output_index()
message_state["item_id"] = f"msg_{uuid.uuid4().hex[:12]}"
message_state["opened"] = True
output_index = message_state["output_index"]
item_id = message_state["item_id"]
return [
_sse(
"response.output_item.added",
{
"type": "response.output_item.added",
"output_index": output_index,
"item": {
"type": "message",
"id": item_id,
"status": "in_progress",
"role": "assistant",
"content": [],
},
},
),
_sse(
"response.content_part.added",
{
"type": "response.content_part.added",
"item_id": item_id,
"output_index": output_index,
"content_index": 0,
"part": {"type": "output_text", "text": "", "annotations": []},
},
),
]
def _close_message_item() -> list[str]:
"""Close the open message item so later text opens a fresh one.
Emits the same done-event triplet the end-of-stream close loop
would, records the item for the final snapshot, and resets the
state in place. No-op when no message item is open.
"""
if not message_state["opened"]:
return []
text = message_state["text"]
events = [
_sse(
"response.output_text.done",
{
"type": "response.output_text.done",
"item_id": message_state["item_id"],
"output_index": message_state["output_index"],
"content_index": 0,
"text": text,
},
),
_sse(
"response.content_part.done",
{
"type": "response.content_part.done",
"item_id": message_state["item_id"],
"output_index": message_state["output_index"],
"content_index": 0,
"part": {"type": "output_text", "text": text, "annotations": []},
},
),
_sse(
"response.output_item.done",
{
"type": "response.output_item.done",
"output_index": message_state["output_index"],
"item": {
"type": "message",
"id": message_state["item_id"],
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": text, "annotations": []}],
},
},
),
]
closed_message_states.append(dict(message_state))
message_state.update(
{"output_index": None, "item_id": None, "opened": False, "text": ""}
)
return events
def _healed_event_sse(events) -> list[str]:
"""Serialize healer events preserving their order.
Text around a healed call must keep its position relative to the
function_call item (output indexes are claimed in emission order),
so never split an event list into all-text-then-all-calls. A healed
call also CLOSES any open message item, so trailing text opens a
fresh message with a later output index, exactly like a native
Responses stream that interleaves messages and calls.
"""
nonlocal full_text
out: list[str] = []
for kind, value in events:
if kind == "text":
if not value:
continue
out.extend(_ensure_message_open())
full_text += value
message_state["text"] += value
api_monitor.append_reply(monitor_id, value)
out.append(
_sse(
"response.output_text.delta",
{
"type": "response.output_text.delta",
"item_id": message_state["item_id"],
"output_index": message_state["output_index"],
"content_index": 0,
"delta": value,
},
)
)
else:
tc = _healed_tc(value)
if tc is None:
continue
out.extend(_close_message_item())
out.extend(_tool_call_delta_events(tc))
return out
def _snapshot_output() -> list[dict]:
"""Snapshot of all completed output items for response.completed."""
indexed_items: list[tuple[int, dict]] = []
if reasoning_state["opened"]:
indexed_items.append(
(
reasoning_state["output_index"],
{
"type": "reasoning",
"id": reasoning_state["item_id"],
"status": "completed",
"summary": [],
"content": [{"type": "reasoning_text", "text": full_reasoning}],
},
)
)
# Closed copies keep opened=True (snapshotted before reset); the
# live state contributes only when a message is currently open.
for msg_st in [*closed_message_states, message_state]:
if not msg_st["opened"]:
continue
indexed_items.append(
(
msg_st["output_index"],
{
"type": "message",
"id": msg_st["item_id"],
"status": "completed",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": msg_st["text"],
"annotations": [],
}
],
},
)
)
for st in tool_call_state.values():
indexed_items.append(
(
st["output_index"],
{
"type": "function_call",
"id": st["item_id"],
"status": "completed",
"call_id": st["call_id"],
"name": st["name"],
"arguments": st["arguments"],
},
)
)
return [item for _, item in sorted(indexed_items, key = lambda pair: pair[0])]
def _failed_response_payload(exc: Exception, status_code: int) -> dict:
return {
"type": "response.failed",
"response": {
"id": resp_id,
"object": "response",
"created_at": created_at,
"status": "failed",
"model": _clean_model,
"output": _snapshot_output(),
"usage": {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens,
},
"error": {
"code": status_code,
"message": _friendly_error(exc),
},
},
}
# ── Preamble events ──
yield _sse(
"response.created",
{
"type": "response.created",
"response": {
"id": resp_id,
"object": "response",
"created_at": created_at,
"status": "in_progress",
"model": _clean_model,
"output": [],
"usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
},
},
)
# ── Direct httpx lifecycle to llama-server ──
# Full same-task open + close, same pattern as
# _openai_passthrough_stream and _anthropic_passthrough_stream: no
# `async with`, explicit aclose of lines_iter BEFORE resp / client so
# the innermost httpcore byte stream is finalised in this task (not via
# the asyncgen GC in a sibling task).
client = httpx.AsyncClient(
timeout = _llama_streaming_generation_timeout(),
trust_env = False,
)
resp = None
lines_iter = None
disconnect_watcher = None
disconnect_event = threading.Event()
try:
req = client.build_request(
"POST", target_url, json = body, headers = {"Connection": "close"}
)
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
try:
resp = await _send_stream_with_preheader_cancel(client, req, request = request)
if resp is None:
api_monitor.finish(monitor_id, "cancelled")
return
except httpx.RequestError as e:
logger.error("responses stream: upstream unreachable: %s", e)
api_monitor.fail(monitor_id, _friendly_error(e))
yield _sse(
"response.failed",
{
"type": "response.failed",
"response": {
"id": resp_id,
"object": "response",
"created_at": created_at,
"status": "failed",
"model": _clean_model,
"output": [],
"error": {"code": 502, "message": _friendly_error(e)},
},
},
)
return
if resp.status_code != 200:
err_bytes = await resp.aread()
err_text = err_bytes.decode("utf-8", errors = "replace")
logger.error(
"responses stream upstream error: status=%s body=%s",
resp.status_code,
err_text[:500],
)
api_monitor.fail(monitor_id, err_text[:500])
yield _sse(
"response.failed",
{
"type": "response.failed",
"response": {
"id": resp_id,
"object": "response",
"created_at": created_at,
"status": "failed",
"model": _clean_model,
"output": [],
"error": {
"code": resp.status_code,
"message": _friendly_upstream_error(err_text[:500]),
},
},
},
)
return
lines_iter = resp.aiter_lines()
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_close(request, resp, disconnect_event)
)
async for raw_line in _aiter_llama_stream_items(
lines_iter,
cancel_event = disconnect_event,
request = request,
first_token_deadline = first_token_deadline,
response = resp,
):
if not raw_line:
continue
if not raw_line.startswith("data: "):
continue
data_str = raw_line[6:]
if data_str.strip() == "[DONE]":
break
try:
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:
_apply_usage(chunk_data.get("usage"))
continue
delta = choices[0].get("delta", {}) or {}
reasoning_delta, visible_delta = extractor.feed(
delta.get("content") or "",
delta.get("reasoning_content"),
)
if reasoning_delta:
for event in _ensure_reasoning_open():
yield event
full_reasoning += reasoning_delta
yield _sse(
"response.reasoning_text.delta",
{
"type": "response.reasoning_text.delta",
"item_id": reasoning_state["item_id"],
"output_index": reasoning_state["output_index"],
"content_index": 0,
"delta": reasoning_delta,
},
)
# Heal text-form tool calls in the visible stream (never in
# reasoning text): promoted calls join the structured tc loop
# below through the same state machinery, and healer events are
# emitted IN ORDER so text after a healed call never jumps ahead
# of the function_call item. Once a structured delta arrives,
# grammar mode worked and the healer goes dormant.
if healer is not None and not healer.dormant:
healed_events = []
if delta.get("tool_calls"):
# Held text preceded the structured call; the call's own
# deltas follow in the structured loop below.
healed_events = healer.structured_tool_call_seen()
if visible_delta:
healed_events.append(("text", visible_delta))
elif visible_delta:
healed_events = healer.feed(visible_delta)
visible_delta = ""
for event in _healed_event_sse(healed_events):
yield event
if visible_delta:
for event in _ensure_message_open():
yield event
full_text += visible_delta
message_state["text"] += visible_delta
api_monitor.append_reply(monitor_id, visible_delta)
yield _sse(
"response.output_text.delta",
{
"type": "response.output_text.delta",
"item_id": message_state["item_id"],
"output_index": message_state["output_index"],
"content_index": 0,
"delta": visible_delta,
},
)
for tc in delta.get("tool_calls") or []:
if (
payload.parallel_tool_calls is False
and healed_tc_index >= 1
and tc.get("index", 0) not in tool_call_state
):
# A healed call already consumed the single allowed slot;
# _drop_parallel_tool_call_deltas only sees native indexes,
# so a native index-0 call would still open a second
# function_call item. Skip it (and its later argument
# deltas, which never allocate a state either).
continue
for event in _tool_call_delta_events(tc):
yield event
_apply_usage(chunk_data.get("usage"))
except asyncio.CancelledError:
disconnect_event.set()
api_monitor.finish(monitor_id, "cancelled")
raise
except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e:
if not disconnect_event.is_set():
logger.error("responses stream error: %s", e)
api_monitor.fail(monitor_id, _friendly_error(e))
status_code = 400 if _classify_llama_generation_error(e) is not None else 500
yield _sse(
"response.failed",
_failed_response_payload(e, status_code),
)
return
except Exception as e:
if disconnect_event.is_set():
api_monitor.finish(monitor_id, "cancelled")
return
logger.error("responses stream error: %s", e)
api_monitor.fail(monitor_id, _friendly_error(e))
status_code = 400 if _classify_llama_generation_error(e) is not None else 500
yield _sse(
"response.failed",
_failed_response_payload(e, status_code),
)
return
finally:
await _aclose_stream_resources(
watchers = (disconnect_watcher,),
iterator = lines_iter,
resp = resp,
client = client,
)
if disconnect_event.is_set():
api_monitor.finish(monitor_id, "cancelled")
return
final_reasoning, final_visible = extractor.finish()
if final_reasoning:
for event in _ensure_reasoning_open():
yield event
full_reasoning += final_reasoning
yield _sse(
"response.reasoning_text.delta",
{
"type": "response.reasoning_text.delta",
"item_id": reasoning_state["item_id"],
"output_index": reasoning_state["output_index"],
"content_index": 0,
"delta": final_reasoning,
},
)
# Last-chance heal of any held residue (e.g. a tool block the model
# never closed) before the trailing visible text is flushed; events
# keep healer order so trailing text stays behind a healed call.
if healer is not None:
events = (healer.feed(final_visible) if final_visible else []) + healer.finalize()
final_visible = ""
for event in _healed_event_sse(events):
yield event
if final_visible:
for event in _ensure_message_open():
yield event
full_text += final_visible
message_state["text"] += final_visible
api_monitor.append_reply(monitor_id, final_visible)
yield _sse(
"response.output_text.delta",
{
"type": "response.output_text.delta",
"item_id": message_state["item_id"],
"output_index": message_state["output_index"],
"content_index": 0,
"delta": final_visible,
},
)
close_items: list[tuple[int, str, dict[str, Any]]] = []
if reasoning_state["opened"]:
close_items.append((reasoning_state["output_index"], "reasoning", reasoning_state))
if message_state["opened"]:
close_items.append((message_state["output_index"], "message", message_state))
close_items.extend((st["output_index"], "tool", st) for st in tool_call_state.values())
for _, kind, st in sorted(close_items, key = lambda item: item[0]):
if kind == "reasoning":
yield _sse(
"response.reasoning_text.done",
{
"type": "response.reasoning_text.done",
"item_id": st["item_id"],
"output_index": st["output_index"],
"content_index": 0,
"text": full_reasoning,
},
)
yield _sse(
"response.content_part.done",
{
"type": "response.content_part.done",
"item_id": st["item_id"],
"output_index": st["output_index"],
"content_index": 0,
"part": {"type": "reasoning_text", "text": full_reasoning},
},
)
yield _sse(
"response.output_item.done",
{
"type": "response.output_item.done",
"output_index": st["output_index"],
"item": {
"type": "reasoning",
"id": st["item_id"],
"status": "completed",
"summary": [],
"content": [{"type": "reasoning_text", "text": full_reasoning}],
},
},
)
continue
if kind == "message":
# Per-item text: message items closed mid-stream (healed-call
# rotation) already emitted their done events, so this state
# carries only its own text, not the whole stream's.
_msg_text = st["text"]
yield _sse(
"response.output_text.done",
{
"type": "response.output_text.done",
"item_id": st["item_id"],
"output_index": st["output_index"],
"content_index": 0,
"text": _msg_text,
},
)
yield _sse(
"response.content_part.done",
{
"type": "response.content_part.done",
"item_id": st["item_id"],
"output_index": st["output_index"],
"content_index": 0,
"part": {"type": "output_text", "text": _msg_text, "annotations": []},
},
)
yield _sse(
"response.output_item.done",
{
"type": "response.output_item.done",
"output_index": st["output_index"],
"item": {
"type": "message",
"id": st["item_id"],
"status": "completed",
"role": "assistant",
"content": [
{"type": "output_text", "text": _msg_text, "annotations": []}
],
},
},
)
continue
# If id/name never arrived (malformed upstream), synthesise so the
# client still sees a coherent frame sequence.
if not st["opened"]:
if not st["call_id"]:
st["call_id"] = f"call_{uuid.uuid4().hex[:12]}"
item_added = {
"type": "response.output_item.added",
"output_index": st["output_index"],
"item": {
"type": "function_call",
"id": st["item_id"],
"status": "in_progress",
"call_id": st["call_id"],
"name": st["name"],
"arguments": "",
},
}
yield _sse("response.output_item.added", item_added)
if st["arguments"]:
yield _sse(
"response.function_call_arguments.delta",
{
"type": "response.function_call_arguments.delta",
"item_id": st["item_id"],
"output_index": st["output_index"],
"delta": st["arguments"],
},
)
st["opened"] = True
args_done = {
"type": "response.function_call_arguments.done",
"item_id": st["item_id"],
"output_index": st["output_index"],
"name": st["name"],
"arguments": st["arguments"],
}
yield _sse("response.function_call_arguments.done", args_done)
item_done = {
"type": "response.output_item.done",
"output_index": st["output_index"],
"item": {
"type": "function_call",
"id": st["item_id"],
"status": "completed",
"call_id": st["call_id"],
"name": st["name"],
"arguments": st["arguments"],
},
}
api_monitor.append_reply(monitor_id, _monitor_call_text(st["name"], st["arguments"]))
yield _sse("response.output_item.done", item_done)
# response.completed
total_tokens = input_tokens + output_tokens
completed_response = {
"type": "response.completed",
"response": {
"id": resp_id,
"object": "response",
"created_at": created_at,
"status": "completed",
"model": _clean_model,
"output": _snapshot_output(),
"usage": {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": total_tokens,
},
},
}
api_monitor.finish(monitor_id)
yield _sse("response.completed", completed_response)
return _SameTaskStreamingResponse(
event_generator(),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
@router.post("/responses")
async def openai_responses(
payload: ResponsesRequest,
request: Request,
current_subject: str = Depends(get_current_subject),
):
"""
OpenAI Responses API endpoint.
Accepts a Responses-format request, converts it to a ChatCompletionRequest
internally, and returns a response matching the Responses API schema
(output array, input_tokens/output_tokens, named SSE events for streaming).
"""
messages = _normalise_responses_input(payload)
if not messages:
raise HTTPException(status_code = 400, detail = "No input provided.")
# System/developer-only input normalises to a non-empty list, so reject it
# before the switch (mirror chat) or an invalid request evicts the resident
# model only for the chat handler to 400 it as having no non-system message.
if not any(m.role not in ("system", "developer") for m in messages):
raise HTTPException(status_code = 400, detail = "At least one non-system message is required.")
# Reject a malformed function tool before any model load, mirroring the
# /v1/chat/completions check, so an invalid request never switches the model.
# Built-in tools (web_search, mcp, ...) carry no name and are dropped later.
for _tool in payload.tools or []:
if not isinstance(_tool, dict) or _tool.get("type") != "function":
continue
_name = _tool.get("name")
if not isinstance(_name, str) or not _name.strip():
raise HTTPException(
status_code = 400,
detail = openai_error_body(
"Invalid 'tools': each function tool must have a 'name'.",
status = 400,
code = "invalid_value",
param = "tools",
),
)
# Reject a forcing-function tool_choice with no name before the switch (mirror
# chat), so a malformed request can't evict the model. Responses forces with
# {"type": "function", "name": "X"}; the streaming path would otherwise forward
# the bad choice and the non-streaming path only 400s after the swap.
_tc = payload.tool_choice
if isinstance(_tc, dict) and _tc.get("type") == "function":
_tc_name = _tc.get("name")
if not isinstance(_tc_name, str) or not _tc_name.strip():
raise HTTPException(
status_code = 400,
detail = openai_error_body(
"Invalid 'tool_choice': the forced function must have a 'name'.",
status = 400,
code = "invalid_value",
param = "tool_choice",
),
)
# After input validation so a 400 never triggers a load. Switches the
# streaming path; non-streaming re-checks via the idempotent chat handler.
# require_vision rejects a swap to a text-only target before it runs, so an
# image request can't evict the resident vision model only to 400 afterwards
# (the non-streaming chat re-check short-circuits on _already_serving).
await _maybe_auto_switch_model(
_switch_model_for_payload(payload),
request,
current_subject,
require_vision = _messages_have_image(messages),
)
if payload.stream:
monitor_id = None
if not getattr(request.state, "skip_api_monitor", False):
monitor_id = api_monitor.start(
endpoint = request.url.path,
method = request.method,
model = payload.model,
prompt = _monitor_prompt_from_messages(messages),
context_length = _monitor_context_length(),
subject = current_subject,
)
try:
return await _responses_stream(payload, messages, request, monitor_id)
except HTTPException as exc:
detail = exc.detail
if not isinstance(detail, str):
detail = json.dumps(detail, default = str)
api_monitor.fail(monitor_id, detail)
raise
except Exception as exc:
api_monitor.fail(monitor_id, _friendly_error(exc))
raise
return await _responses_non_streaming(payload, messages, request, current_subject)
# =====================================================================
# Anthropic-Compatible Messages API (/messages → /v1/messages)
# =====================================================================
_STUDIO_ANTHROPIC_TOOL_ALIASES = {
"web_search": "web_search",
"web_search_20250305": "web_search",
"web_fetch": "web_search",
"web_fetch_20250910": "web_search",
"web_fetch_20260209": "web_search",
"python": "python",
"terminal": "terminal",
}
def _anthropic_requested_studio_tools(tools: Optional[list]) -> set[str]:
requested: set[str] = set()
for tool in tools or []:
td = tool if isinstance(tool, dict) else tool.model_dump()
# Client tools always carry input_schema; server tools never do.
if td.get("input_schema") is not None:
continue
# Anthropic dispatches server tools by `type`, not bare `name`; matching
# name too would let a malformed client tool like `{"name": "python"}`
# silently flip into server-execution mode.
type_ = td.get("type")
if isinstance(type_, str) and type_ in _STUDIO_ANTHROPIC_TOOL_ALIASES:
requested.add(_STUDIO_ANTHROPIC_TOOL_ALIASES[type_])
return requested
def _select_anthropic_server_tools(
all_tools: list[dict], requested_studio_tools: set[str], enabled_tools: Optional[list[str]]
) -> list[dict]:
"""Select Studio tools requested through Anthropic tools and extensions."""
if not requested_studio_tools and enabled_tools is None:
return all_tools
selected_names = set(requested_studio_tools)
if enabled_tools is not None:
selected_names.update(enabled_tools)
return [tool for tool in all_tools if tool["function"]["name"] in selected_names]
def _image_bytes_to_png_b64(raw: bytes) -> str:
"""Decode raw image bytes and re-encode to a base64-ascii PNG string.
llama-server's stb_image only handles a few formats (JPEG/PNG/BMP/...); re-
encoding to PNG keeps JPEG/WebP/... inputs loadable. Raises on undecodable
input; callers wrap the call in ``try`` -> HTTPException(400)."""
from PIL import Image
img = Image.open(io.BytesIO(raw)).convert("RGB")
buf = io.BytesIO()
img.save(buf, format = "PNG")
return base64.b64encode(buf.getvalue()).decode("ascii")
def _normalize_anthropic_openai_images(openai_messages: list[dict], is_vision: bool) -> bool:
"""Enforce the vision guard on translated Anthropic messages and normalize
any base64-data-URL ``image_url`` parts to PNG.
llama-server's stb_image only handles a few formats (JPEG/PNG/BMP/…);
Anthropic clients commonly send JPEG or WebP, and Claude Code sends WebP.
Re-encoding everything to PNG mirrors `_openai_messages_for_passthrough` /
the GGUF branch of `/v1/chat/completions` so the two endpoints agree.
Mutates ``openai_messages`` in place. Returns ``True`` when any image part
was seen (so the caller can skip a second scan). Raises HTTPException(400)
when images are present but the active model isn't a vision model, or when
an image cannot be decoded.
"""
has_image = False
for msg in openai_messages:
content = msg.get("content")
if not isinstance(content, list):
continue
for part in content:
if part.get("type") != "image_url":
continue
has_image = True
if not is_vision:
raise HTTPException(
status_code = 400,
detail = "Image provided but current GGUF model does not support vision.",
)
url = (part.get("image_url") or {}).get("url", "")
if not url.startswith("data:"):
# Remote URLs are forwarded as-is; llama-server will
# fetch (or fail) per its own support matrix.
continue
try:
_, b64data = url.split(",", 1)
raw = base64.b64decode(b64data)
png_b64 = _image_bytes_to_png_b64(raw)
except Exception:
raise HTTPException(
status_code = 400,
detail = "Failed to process image.",
)
part["image_url"] = {"url": f"data:image/png;base64,{png_b64}"}
return has_image
def _validate_anthropic_client_tools(tools) -> None:
# Reject malformed client tools before any model load, so an invalid request
# never evicts the loaded model. AnthropicTool relaxed name/input_schema to
# Optional for server tools, so the converter silently drops incomplete
# entries; surface them as 400 here. A `type` field marks a server-tool
# declaration (unrecognized server tools are no-ops); anything else without
# input_schema or name is malformed.
for tool in tools or []:
td = tool if isinstance(tool, dict) else tool.model_dump()
name, type_, schema = td.get("name"), td.get("type"), td.get("input_schema")
if schema is None and not isinstance(type_, str):
raise HTTPException(
status_code = 400,
detail = f"Tool {name!r} is missing required field 'input_schema'.",
)
if schema is not None and (not isinstance(name, str) or not name):
raise HTTPException(
status_code = 400,
detail = "Client tool is missing required field 'name'.",
)
@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.
"""
# Reject malformed tools before the switch, like /messages, so an invalid
# count request can't evict the loaded model.
_validate_anthropic_client_tools(payload.tools)
# Count with the requested model's tokenizer, like the sibling /messages.
# Carry the vision guard too: an image count naming a text-only GGUF must not
# evict a loaded vision model for a swap that can't serve the request.
await _maybe_auto_switch_model(
_switch_model_for_payload(payload),
request,
current_subject,
require_vision = _anthropic_request_has_image(payload),
)
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).
# Coalesce adjacent user turns left behind by dropping an empty / null assistant
# turn, so a strict GGUF chat template does not 400 on non-alternating roles
# (mirrors the GGUF chat path); a no-op for already-alternating histories.
openai_messages = _coalesce_consecutive_user_turns(
_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]:
"""Return messages with a single leading system prompt, preserving multimodal parts."""
safe_messages = messages or []
if not system_prompt:
return safe_messages
# Drop existing system/developer turns so the backend never sees duplicate
# or conflicting system instructions, then prepend the resolved prompt.
others = [dict(msg) for msg in safe_messages if msg.get("role") not in ("system", "developer")]
return [{"role": "system", "content": system_prompt}, *others]
@router.post("/messages")
async def anthropic_messages(
payload: AnthropicMessagesRequest,
request: Request,
current_subject: str = Depends(get_current_subject),
):
"""
Anthropic-compatible Messages API endpoint.
Translates Anthropic message format to internal OpenAI format, runs through
the existing agentic tool loop when tools are provided, and returns
responses in Anthropic Messages API format (streaming SSE or non-streaming
JSON).
"""
llama_backend = get_llama_cpp_backend()
# Default-off parity: with no automatic load possible and nothing loaded, 503
# before any request-shape check, exactly as the pre-feature endpoint did. When
# an automatic load can run (auto-switch or a standalone idle TTL), fall through
# so validation runs before the reload hook gets a chance to restore the model.
if not llama_backend.is_loaded and not _automatic_model_load_may_run():
raise HTTPException(
status_code = 503,
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. Validate before
# auto-switch so a rejected request never triggers a model load.
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",
),
)
# Reject malformed client tools before any model load (see helper), so an
# invalid request never evicts the loaded model.
_validate_anthropic_client_tools(payload.tools)
# Mixing Anthropic server tools with custom client tools is unsupported (the
# server-tool loop can't relay client functions back to the caller). Reject
# before the switch too -- it depends only on the payload -- so an invalid
# request never evicts the loaded model. Reused below for tool routing.
requested_studio_tools = _anthropic_requested_studio_tools(payload.tools)
_has_client_tool = any(
(t if isinstance(t, dict) else t.model_dump()).get("input_schema") is not None
for t in payload.tools or []
)
if requested_studio_tools and _has_client_tool:
raise HTTPException(
status_code = 400,
detail = (
"Mixing Anthropic server tools (e.g. web_search_20250305) "
"with custom client tools in a single request is not "
"supported. Send them in separate requests."
),
)
# require_vision rejects a swap to a text-only target before it runs, so an
# image request can't evict the resident vision model only to hit the vision
# guard (_normalize_anthropic_openai_images) below after the load.
await _maybe_auto_switch_model(
_switch_model_for_payload(payload),
request,
current_subject,
require_vision = _anthropic_request_has_image(payload),
)
if not llama_backend.is_loaded:
raise HTTPException(
status_code = 503,
detail = "No GGUF model loaded. Load a GGUF model first.",
)
# Advertised repo id after an auto-switch load, else a clean public id, never
# the local .gguf path (and a legacy raw path in payload.model is sanitized).
model_name = _llama_public_model_id(llama_backend, payload.model)
message_id = f"msg_{uuid.uuid4().hex[:24]}"
# ── Translate Anthropic → OpenAI ──────────────────────────
openai_messages = anthropic_messages_to_openai(
[m.model_dump() for m in payload.messages],
payload.system,
)
# Strip synthetic provider-side builtin tool history (web_search,
# web_fetch, code_execution, image_generation cards tagged with
# _server_tool or extra_content.google.native_part) before handing off to
# local llama-server. The local /v1/chat/completions and GGUF passthrough
# builders apply the same strip; without it an Anthropic /v1/messages caller
# replaying a prior provider-side tool_use forwards fake builtin tool
# history to a backend with no matching function declarations.
# Coalesce adjacent user turns left behind by dropping an empty / null assistant
# turn, so a strict GGUF chat template does not 400 on non-alternating roles
# (mirrors the GGUF chat path); a no-op for already-alternating histories.
openai_messages = _coalesce_consecutive_user_turns(
_strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages))
)
# Enforce vision guard + re-encode embedded images to PNG so the Anthropic
# endpoint matches /v1/chat/completions.
_has_image = _normalize_anthropic_openai_images(openai_messages, llama_backend.is_vision)
temperature = payload.temperature if payload.temperature is not None else 0.6
top_p = payload.top_p if payload.top_p is not None else 0.95
top_k = payload.top_k if payload.top_k is not None else 20
min_p = payload.min_p if payload.min_p is not None else 0.01
repetition_penalty = (
payload.repetition_penalty if payload.repetition_penalty is not None else 1.0
)
presence_penalty = payload.presence_penalty if payload.presence_penalty is not None else 0.0
stop = payload.stop_sequences or None
# Translate Anthropic tool_choice to OpenAI format for llama-server. Falls
# back to "auto" when unset or unrecognized (prior hardcoded behavior).
openai_tool_choice = anthropic_tool_choice_to_openai(payload.tool_choice)
if openai_tool_choice is None:
openai_tool_choice = "auto"
cancel_event = threading.Event()
# ── Tool routing ──────────────────────────────────────────
# Three paths:
# 1. enable_tools=true → server-side execution of built-in tools (Unsloth shorthand)
# 2. tools=[...] only → client-side pass-through (standard Anthropic behavior)
# 3. neither → plain chat
# The server-side agentic loop doesn't support multimodal input -- matches
# the `not image_b64` gate in /v1/chat/completions. requested_studio_tools and
# the mixed-mode rejection were computed before the switch above.
openai_client_tools = [
tool
for tool in anthropic_tools_to_openai(payload.tools or [])
if tool.get("function", {}).get("name") not in requested_studio_tools
]
# An Anthropic server-tool declaration implies server-tool mode, but only
# when tools aren't explicitly disabled (CLI --disable-tools or per-request
# enable_tools=false). Explicit False always wins.
_enable = _effective_enable_tools(payload)
server_tools = (
(_enable or (_enable is None and bool(requested_studio_tools)))
and llama_backend.supports_tools
and not _has_image
)
client_tools = (
not server_tools
and len(openai_client_tools) > 0
and getattr(llama_backend, "supports_tool_passthrough", 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")
)
monitor_id = None
monitor_context_length = _monitor_context_length()
request_state = getattr(request, "state", None)
if not getattr(request_state, "skip_api_monitor", False):
request_url = getattr(request, "url", None)
monitor_id = api_monitor.start(
endpoint = getattr(request_url, "path", "/v1/messages"),
method = getattr(request, "method", "POST"),
model = model_name,
prompt = _monitor_prompt_from_messages(openai_messages),
context_length = monitor_context_length,
subject = current_subject,
)
async def _monitored_anthropic(coro):
try:
response = await coro
except asyncio.CancelledError:
cancel_event.set()
api_monitor.finish(monitor_id, "cancelled")
raise
except Exception as exc:
api_monitor.fail(monitor_id, _friendly_error(exc))
raise
return _monitor_anthropic_response(
response,
monitor_id,
monitor_context_length,
cancel_event,
)
# ── Client-side pass-through path ─────────────────────────
if client_tools:
openai_tools = openai_client_tools
if payload.stream:
return await _monitored_anthropic(
_anthropic_passthrough_stream(
request,
cancel_event,
llama_backend,
openai_messages,
openai_tools,
temperature,
top_p,
top_k,
payload.max_tokens,
message_id,
model_name,
stop = stop,
min_p = min_p,
repetition_penalty = repetition_penalty,
presence_penalty = presence_penalty,
tool_choice = openai_tool_choice,
session_id = payload.session_id,
cancel_id = payload.cancel_id,
disable_parallel_tool_use = _disable_parallel,
auto_heal_tool_calls = payload.auto_heal_tool_calls,
)
)
return await _monitored_anthropic(
_anthropic_passthrough_non_streaming(
llama_backend,
openai_messages,
openai_tools,
temperature,
top_p,
top_k,
payload.max_tokens,
message_id,
model_name,
stop = stop,
min_p = min_p,
repetition_penalty = repetition_penalty,
presence_penalty = presence_penalty,
tool_choice = openai_tool_choice,
disable_parallel_tool_use = _disable_parallel,
auto_heal_tool_calls = payload.auto_heal_tool_calls,
nudge_tool_calls = payload.nudge_tool_calls,
)
)
if server_tools:
# Bypass Permissions suppresses confirm, so both flags together is fine.
if bool(getattr(payload, "confirm_tool_calls", False)) and not bool(
getattr(payload, "bypass_permissions", False)
):
api_monitor.fail(
monitor_id,
"confirm_tool_calls is not supported for Anthropic Messages server tools.",
)
raise HTTPException(
status_code = 400,
detail = anthropic_error_body(
"confirm_tool_calls is not supported for Anthropic Messages server tools.",
status = 400,
err_type = "invalid_request_error",
),
)
from core.inference.tools import ALL_TOOLS
openai_tools = _select_anthropic_server_tools(
ALL_TOOLS,
requested_studio_tools,
payload.enabled_tools,
)
# Build tool-use system prompt nudge (same logic as /chat/completions)
_nudge = _build_tool_action_nudge(
tools = openai_tools,
model_name = model_name,
)
if _nudge:
# Inject into system prompt
if openai_messages and openai_messages[0].get("role") == "system":
openai_messages[0]["content"] = (
openai_messages[0]["content"].rstrip() + "\n\n" + _nudge
)
else:
openai_messages.insert(0, {"role": "system", "content": _nudge})
# Strip stale tool-call XML via the protected display helper (think rehearsal and [TOOL_CALLS]
# prose survive), gated on enabled tool names so documented inactive examples are kept.
_anthropic_history_gate = _display_tool_name_gate(openai_tools)
for _msg in openai_messages:
if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str):
_msg["content"] = _strip_tool_xml_for_display(
_msg["content"],
auto_heal_tool_calls = True,
enabled_tool_names = _anthropic_history_gate,
).strip()
def _run_tool_gen():
return llama_backend.generate_chat_completion_with_tools(
messages = openai_messages,
tools = openai_tools,
temperature = temperature,
top_p = top_p,
top_k = top_k,
min_p = min_p,
repetition_penalty = repetition_penalty,
presence_penalty = presence_penalty,
max_tokens = payload.max_tokens,
stop = stop,
cancel_event = cancel_event,
max_tool_iterations = 25,
auto_heal_tool_calls = True,
nudge_tool_calls = payload.nudge_tool_calls,
tool_call_timeout = 300,
session_id = payload.session_id,
# Anthropic passthrough has no rag_scope field (RAG is local-only).
rag_scope = getattr(payload, "rag_scope", None),
disable_parallel_tool_use = _disable_parallel,
bypass_permissions = bool(payload.bypass_permissions),
)
if payload.stream:
return await _monitored_anthropic(
_anthropic_tool_stream(
request,
cancel_event,
_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 _monitored_anthropic(
_anthropic_tool_non_streaming(
_run_tool_gen,
message_id,
model_name,
disable_parallel_tool_use = _disable_parallel,
openai_tools = openai_tools,
)
)
# ── No-tool path ──────────────────────────────────────────
def _run_plain_gen():
return llama_backend.generate_chat_completion(
messages = openai_messages,
temperature = temperature,
top_p = top_p,
top_k = top_k,
min_p = min_p,
repetition_penalty = repetition_penalty,
presence_penalty = presence_penalty,
max_tokens = payload.max_tokens,
stop = stop,
cancel_event = cancel_event,
)
if payload.stream:
return await _monitored_anthropic(
_anthropic_plain_stream(
request,
cancel_event,
_run_plain_gen,
message_id,
model_name,
llama_backend = llama_backend,
openai_messages = openai_messages,
)
)
return await _monitored_anthropic(
_anthropic_plain_non_streaming(
_run_plain_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()
# Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final
# answer is prose and must survive in the delivered text.
_display_names = _display_tool_name_gate(openai_tools)
# 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, 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()
# Watcher to cancel on disconnect: the in-loop poll fires only between
# events, so a mid-prefill disconnect would otherwise hold the decode slot.
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
try:
while True:
if cancel_event.is_set() or await request.is_disconnected():
cancel_event.set()
return
event = await asyncio.to_thread(next, gen, _sentinel)
if event is _sentinel:
break
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.
# Protected helper preserves <think> rehearsal and balanced
# [TOOL_CALLS] trailing prose (raw _TOOL_XML_RE.sub corrupts both).
if etype == "content":
event = dict(event)
event["text"] = _strip_tool_xml_for_display(
event["text"],
auto_heal_tool_calls = True,
enabled_tool_names = _display_names,
)
# 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
finally:
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
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 _sse_streaming_response(_stream())
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, input_tokens = input_tokens):
yield line
captured_finish_reason = None
gen = run_gen()
# Watcher to cancel on disconnect: the in-loop poll fires only between
# chunks, so a mid-prefill disconnect would otherwise hold the decode slot.
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_cancel(request, cancel_event)
)
try:
while True:
if cancel_event.is_set() or await request.is_disconnected():
cancel_event.set()
return
cumulative = await asyncio.to_thread(next, gen, _sentinel)
if cumulative is _sentinel:
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
# Plain generator yields cumulative text strings
for line in emitter.feed({"type": "content", "text": cumulative}):
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
finally:
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
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 _sse_streaming_response(_stream())
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)
def _anthropic_message_json_response(
message_id, model_name, content_blocks, stop_reason, usage
) -> Response:
"""Assemble the terminal Anthropic non-streaming JSON response shared by the
tool / plain / passthrough paths."""
return _model_json_response(
AnthropicMessagesResponse(
id = message_id,
model = model_name,
content = content_blocks,
stop_reason = stop_reason,
usage = AnthropicUsage(
input_tokens = usage.get("prompt_tokens", 0),
output_tokens = usage.get("completion_tokens", 0),
),
)
)
async def _anthropic_tool_non_streaming(
run_gen,
message_id,
model_name,
disable_parallel_tool_use = False,
openai_tools = None,
):
"""Non-streaming response for the tool-calling path.
Builds ``content_blocks`` in generation order (text → tool_use → text →
tool_use → ...), mirroring the streaming emitter. Deltas within one
synthesis turn merge into the trailing text block; tool_use blocks interrupt
the text sequence and open a new text block on the next content event.
``prev_text`` is reset on ``tool_end`` because
``generate_chat_completion_with_tools`` yields cumulative content *per
turn* -- the first content event of turn N+1 must diff against an empty
baseline, not turn N's final length.
"""
content_blocks: list = []
tool_blocks_by_id: dict[str, AnthropicResponseToolUseBlock] = {}
usage = {}
prev_text = ""
captured_finish_reason = None
# Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final
# answer is prose and must survive in the delivered text.
_display_names = _display_tool_name_gate(openai_tools)
# Pending client tool_use; cleared by tool_end (server execution) or
# trailing text. See the stop_reason mapping below.
ends_on_tool_use = False
events = _collect_anthropic_events(run_gen)
for event in events:
etype = event.get("type", "")
if etype == "content":
# Strip leaked tool XML (protected helper keeps think rehearsal and trailing prose).
clean = _strip_tool_xml_for_display(
event["text"], auto_heal_tool_calls = True, enabled_tool_names = _display_names
)
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:
content_blocks.append(AnthropicResponseTextBlock(text = new))
elif etype == "tool_start":
tool_call_id = event["tool_call_id"]
arguments = event.get("arguments", {})
existing_tool_block = tool_blocks_by_id.get(tool_call_id) if tool_call_id else None
if existing_tool_block is not None:
if arguments or not existing_tool_block.input:
existing_tool_block.input = arguments
if event.get("tool_name") and not existing_tool_block.name:
existing_tool_block.name = event["tool_name"]
else:
tool_block = AnthropicResponseToolUseBlock(
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
)
return _anthropic_message_json_response(
message_id, model_name, content_blocks, stop_reason, usage
)
async def _anthropic_plain_non_streaming(run_gen, message_id, model_name):
"""Non-streaming response for the no-tool path."""
text_parts = []
usage = {}
prev_text = ""
captured_finish_reason = None
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
if new:
text_parts.append(new)
full_text = "".join(text_parts)
content_blocks = []
if full_text:
content_blocks.append(AnthropicResponseTextBlock(text = full_text))
stop_reason = openai_finish_to_anthropic_stop(captured_finish_reason, had_tool_calls = False)
return _anthropic_message_json_response(
message_id, model_name, content_blocks, stop_reason, usage
)
# =====================================================================
# Client-side tool pass-through (Anthropic-native tools field)
# =====================================================================
def _build_passthrough_payload(
openai_messages,
openai_tools,
temperature,
top_p,
top_k,
max_tokens,
stream,
stop = None,
min_p = None,
repetition_penalty = None,
presence_penalty = None,
tool_choice = "auto",
response_format = None,
chat_template_kwargs = None,
backend_ctx = None,
seed = None,
stream_options = None,
):
body = {
"messages": openai_messages,
"temperature": temperature,
"top_p": top_p,
"top_k": top_k,
"stream": stream,
}
if openai_tools:
body["tools"] = openai_tools
if tool_choice is not None:
body["tool_choice"] = tool_choice
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)
)
# 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:
# llama-server's field is "repeat_penalty", not "repetition_penalty".
body["repeat_penalty"] = repetition_penalty
if presence_penalty is not None:
body["presence_penalty"] = presence_penalty
if response_format is not None:
# llama-server applies a GBNF grammar derived from the JSON schema when
# response_format is present. The field is documented flat at the
# request root (tools/server/README.md), which is also what the OpenAI
# SDK produces by spreading extra_body into the body top.
body["response_format"] = response_format
if chat_template_kwargs is not None:
# Propagate reasoning / template overrides (e.g. enable_thinking) so
# llama-server renders the Jinja template in the caller's mode instead
# of the model's load-time default.
body["chat_template_kwargs"] = chat_template_kwargs
return body
async def _anthropic_passthrough_stream(
request,
cancel_event,
llama_backend,
openai_messages,
openai_tools,
temperature,
top_p,
top_k,
max_tokens,
message_id,
model_name,
stop = None,
min_p = None,
repetition_penalty = None,
presence_penalty = None,
tool_choice = "auto",
session_id = None,
cancel_id = None,
disable_parallel_tool_use = False,
auto_heal_tool_calls = None,
):
"""Streaming client-side pass-through: forward tools to llama-server and
translate its stream to Anthropic SSE without executing anything."""
target_url = f"{llama_backend.base_url}/v1/chat/completions"
body = _build_passthrough_payload(
openai_messages,
openai_tools,
temperature,
top_p,
top_k,
max_tokens,
True,
stop = stop,
min_p = min_p,
repetition_penalty = repetition_penalty,
presence_penalty = presence_penalty,
tool_choice = tool_choice,
backend_ctx = llama_backend.context_length,
stream_options = {"include_usage": True},
)
# 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()
# Promote text-form tool calls (declared client tools only) into
# tool_use blocks; verbatim behavior when healing is off or no tools.
# tool_choice arrives here already converted to the OpenAI shape.
_allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice)
if _allowed_tools:
emitter.enable_healing(
_allowed_tools,
openai_tools,
disable_parallel_tool_use = disable_parallel_tool_use,
)
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
# generator MANUALLY -- no `async with`, no anonymous iterator.
#
# On Python 3.13 + httpcore 1.0.x, `async for raw_line in
# resp.aiter_lines():` creates an anonymous async generator. When the
# loop exits via `break` (or the generator is orphaned by a mid-stream
# client disconnect), `async for` does NOT auto-close the iterator like
# a sync `for` would. The iterator stays reachable only from the current
# coroutine frame; once `_stream()` returns, the frame is GC'd and the
# iterator becomes unreachable. The asyncgen finalizer then runs aclose()
# on a LATER GC pass in a DIFFERENT asyncio task, where httpcore's
# `HTTP11ConnectionByteStream.aclose()` enters `anyio.CancelScope.__exit__`
# with a mismatched task and prints `RuntimeError: Attempted to exit
# cancel scope in a different task` / `RuntimeError: async generator
# ignored GeneratorExit` as "Exception ignored in:" unraisable warnings.
#
# Fix: save `resp.aiter_lines()` as `lines_iter`, and in finally
# explicitly `await lines_iter.aclose()` BEFORE `resp.aclose()` /
# `client.aclose()`. This closes the iterator in our own task's event
# loop, cleaning up the httpcore byte-stream before the asyncgen
# finalizer has anything orphaned to finalize. Each aclose is wrapped in
# `try: ... except Exception: pass` so nested anyio cleanup noise can't
# bubble out.
client = httpx.AsyncClient(
timeout = _llama_streaming_generation_timeout(),
limits = httpx.Limits(max_keepalive_connections = 0),
trust_env = False,
)
resp = None
lines_iter = None
cancel_watcher = None
disconnect_watcher = None
try:
req = client.build_request(
"POST", target_url, json = body, headers = {"Connection": "close"}
)
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
resp = await _send_stream_with_preheader_cancel(
client, req, cancel_event, request = request
)
if resp is None:
return
# 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(
_friendly_upstream_error(_err_text),
status = resp.status_code,
),
)
return
# Watchers unblock aiter_lines() during prefill, before in-loop
# cancel/disconnect checks can run.
cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp))
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_close(request, resp, cancel_event)
)
lines_iter = resp.aiter_lines()
async for raw_line in _aiter_llama_stream_items(
lines_iter,
cancel_event = cancel_event,
request = request,
first_token_deadline = first_token_deadline,
response = resp,
):
if not raw_line or not raw_line.startswith("data: "):
continue
data_str = raw_line[6:]
if data_str.strip() == "[DONE]":
break
try:
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 Exception as e:
if not cancel_event.is_set():
logger.error("anthropic_messages passthrough stream error: %s", e)
get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
event = _anthropic_stream_error_event(
e,
force = True,
)
if event is not None:
yield event
return
finally:
# Same shape as the OpenAI passthrough: a close-time CancelledError
# re-raised by _aclose_stream_resources must not skip the tracker exit.
try:
await _aclose_stream_resources(
watchers = (cancel_watcher, disconnect_watcher),
iterator = lines_iter,
resp = resp,
client = client,
)
finally:
_tracker.__exit__(None, None, None)
for line in emitter.finish():
yield line
return _sse_streaming_response(_stream())
async def _anthropic_passthrough_non_streaming(
llama_backend,
openai_messages,
openai_tools,
temperature,
top_p,
top_k,
max_tokens,
message_id,
model_name,
stop = None,
min_p = None,
repetition_penalty = None,
presence_penalty = None,
tool_choice = "auto",
disable_parallel_tool_use = False,
auto_heal_tool_calls = None,
nudge_tool_calls = None,
):
"""Non-streaming client-side pass-through."""
target_url = f"{llama_backend.base_url}/v1/chat/completions"
body = _build_passthrough_payload(
openai_messages,
openai_tools,
temperature,
top_p,
top_k,
max_tokens,
False,
stop = stop,
min_p = min_p,
repetition_penalty = repetition_penalty,
presence_penalty = presence_penalty,
tool_choice = tool_choice,
backend_ctx = llama_backend.context_length,
)
resp = await nonstreaming_client().post(
target_url,
json = body,
timeout = _llama_non_streaming_generation_timeout(),
)
if resp.status_code != 200:
raise HTTPException(
status_code = resp.status_code,
detail = _friendly_upstream_error(resp.text[:500]),
)
data = resp.json()
# tool_choice arrives here already converted to the OpenAI shape.
_allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice)
# Opt-in single-retry nudge (mirrors the OpenAI passthrough): the model
# tried to call a tool but nothing usable came out; re-ask once with the
# prompt prefix intact so llama-server's KV cache is reused.
if (
_allowed_tools
and nudge_enabled(nudge_tool_calls)
and nudge_should_retry(data, _allowed_tools, openai_tools)
):
retry_body = {
**body,
"messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)],
}
try:
retry_resp = await nonstreaming_client().post(
target_url,
json = retry_body,
timeout = _llama_non_streaming_generation_timeout(),
)
if retry_resp.status_code == 200:
retry_data = retry_resp.json()
if response_has_promotable_calls(retry_data, _allowed_tools, openai_tools):
data = retry_data
except (httpx.RequestError, ValueError) as exc:
logger.warning("tool-call nudge retry failed; keeping original: %s", exc)
choice = (data.get("choices") or [{}])[0]
message = choice.get("message") or {}
finish_reason = choice.get("finish_reason")
healing_active = bool(_allowed_tools)
healed_events = (
heal_openai_message_events(message, _allowed_tools, openai_tools)
if healing_active
else None
)
content_blocks = []
tool_calls = []
if healed_events:
emitted_tool_uses = 0
for kind, value in healed_events:
if kind == "text":
text = str(value).strip()
if text:
content_blocks.append(AnthropicResponseTextBlock(text = text))
continue
if disable_parallel_tool_use and emitted_tool_uses >= 1:
continue
fn = value.get("function") or {}
try:
args = json.loads(fn.get("arguments", "{}"))
except json.JSONDecodeError:
args = {}
tool_calls.append(value)
emitted_tool_uses += 1
content_blocks.append(
AnthropicResponseToolUseBlock(
id = anthropic_tool_use_id(value.get("id")),
name = fn.get("name", ""),
input = args,
)
)
else:
text = message.get("content") or ""
if text:
# Keep unpromoted bytes when healing is active; legacy stripping is
# only for opted-out or no-client-tool requests. Protected helper (not
# raw _TOOL_XML_RE.sub): preserves <think> rehearsal and balanced
# [TOOL_CALLS] trailing prose, gated on the declared tools so an
# inactive NAME[ARGS]{...} example in the final text is kept.
if not healing_active:
text = _strip_tool_xml_for_display(
text,
auto_heal_tool_calls = True,
enabled_tool_names = _display_tool_name_gate(openai_tools),
)
text = text.strip()
if text:
content_blocks.append(AnthropicResponseTextBlock(text = text))
tool_calls = message.get("tool_calls") or []
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:
args = json.loads(fn.get("arguments", "{}"))
except json.JSONDecodeError:
args = {}
content_blocks.append(
AnthropicResponseToolUseBlock(
id = anthropic_tool_use_id(tc.get("id")),
name = fn.get("name", ""),
input = args,
)
)
stop_reason = openai_finish_to_anthropic_stop(finish_reason, had_tool_calls = bool(tool_calls))
usage = data.get("usage") or {}
return _anthropic_message_json_response(
message_id, model_name, content_blocks, stop_reason, usage
)
# =====================================================================
# Client-side tool pass-through (OpenAI-native /v1/chat/completions)
# =====================================================================
def _drop_empty_assistant_sentinels(messages: list[dict]) -> list[dict]:
"""Drop bare ``{"role":"assistant"}`` Stop-button sentinels; passthrough backends reject them."""
out: list[dict] = []
for m in messages:
if m.get("role") == "assistant":
has_content = bool(m.get("content"))
has_tool_calls = bool(m.get("tool_calls"))
if not has_content and not has_tool_calls:
continue
out.append(m)
return out
def _merge_user_content(a: Any, b: Any) -> Any:
"""Join two user ``content`` values: strings with a blank line, else as concatenated parts."""
if isinstance(a, str) and isinstance(b, str):
if not a:
return b
if not b:
return a
return a + "\n\n" + b
def _parts(c: Any) -> list:
if c is None:
return []
if isinstance(c, str):
return [{"type": "text", "text": c}] if c else []
if isinstance(c, list):
return list(c)
return [{"type": "text", "text": str(c)}]
return _parts(a) + _parts(b)
def _coalesce_consecutive_user_turns(messages: list[dict]) -> list[dict]:
"""Merge adjacent user turns so the GGUF history stays alternating.
Dropping an empty assistant turn (0-token reply or Stop-button sentinel) can
leave two user turns in a row, which makes strict templates (Gemma 3, ...)
raise "Conversation roles must alternate" -> llama-server 400. Only user turns
merge (assistant/tool turns may carry tool_calls/tool_call_id); multimodal
parts are preserved; no-op for already-alternating histories.
"""
out: list[dict] = []
for m in messages:
if m.get("role") == "user" and out and out[-1].get("role") == "user":
prev = dict(out[-1])
prev["content"] = _merge_user_content(prev.get("content"), m.get("content"))
out[-1] = prev
continue
out.append(m)
return out
_LOCAL_SERVER_BUILTIN_TOOL_NAMES = frozenset(
{"web_search", "web_fetch", "code_execution", "image_generation"}
)
def _strip_provider_synthetic_tool_history(messages: list[dict]) -> list[dict]:
"""Drop synthetic provider-side tool_calls + matching role=tool replies on
the local-backend (llama-server / GGUF) dispatch path.
A Gemini chat that ran code_execution / image_generation persists the
server-side tool card into history as an assistant tool_calls entry tagged
with ``args._server_tool`` (or a Gemini ``args.google.native_part`` payload)
plus a follow-up role=tool reply. When the user switches the SAME thread to
a local GGUF model, those synthetic tool_calls aren't real user functions,
llama-server has no matching declaration, and Gemini-only ``extra_content``
/ ``native_part`` payloads are meaningless. Forward only ordinary user
function calls; strip the matched role=tool replies too so the backend never
sees an orphan tool_call_id.
"""
dropped_ids: set[str] = set()
sanitized_assistant: list[dict] = []
for m in messages:
if m.get("role") != "assistant":
sanitized_assistant.append(m)
continue
tool_calls = m.get("tool_calls")
if not isinstance(tool_calls, list) or not tool_calls:
# Plain text Gemini reply: still strip message-level
# `extra_content` (carries `google.thought_signature` replay
# metadata) so a text-only Gemini turn switched to a local GGUF
# backend doesn't leak Gemini-only fields to llama-server.
# ChatMessage didn't used to have `extra_content` (implicitly
# dropped); round-22 added it, which made this leak possible.
if "extra_content" in m:
m = {k: v for k, v in m.items() if k != "extra_content"}
sanitized_assistant.append(m)
continue
cleaned: list[dict] = []
for tc in tool_calls:
if not isinstance(tc, dict):
cleaned.append(tc)
continue
fn = tc.get("function")
name = ""
if isinstance(fn, dict):
name = (fn.get("name") or "").lower()
if name in _LOCAL_SERVER_BUILTIN_TOOL_NAMES:
raw_args = fn.get("arguments") if isinstance(fn, dict) else None
args_obj: Any = None
if isinstance(raw_args, str):
try:
args_obj = json.loads(raw_args) if raw_args else None
except Exception:
args_obj = None
elif isinstance(raw_args, dict):
args_obj = raw_args
is_synthetic = False
if isinstance(args_obj, dict):
if args_obj.get("_server_tool") is True:
is_synthetic = True
google = args_obj.get("google")
if isinstance(google, dict) and isinstance(google.get("native_part"), dict):
is_synthetic = True
if is_synthetic:
tc_id = tc.get("id")
if isinstance(tc_id, str) and tc_id:
dropped_ids.add(tc_id)
continue
# Strip Gemini-only `extra_content` on real user tool_calls too --
# llama-server has no use for it and may pass it to the model
# unchanged.
if "extra_content" in tc:
tc = {k: v for k, v in tc.items() if k != "extra_content"}
cleaned.append(tc)
# Drop message-level `extra_content` (Gemini thoughtSignature replay
# metadata) on local dispatch.
m_clean = {k: v for k, v in m.items() if k != "extra_content"}
if cleaned:
m_clean["tool_calls"] = cleaned
else:
m_clean.pop("tool_calls", None)
if not m_clean.get("content") and not m_clean.get("tool_calls"):
continue # assistant turn now empty, drop
sanitized_assistant.append(m_clean)
if not dropped_ids:
return sanitized_assistant
out: list[dict] = []
for m in sanitized_assistant:
if (
m.get("role") == "tool"
and isinstance(m.get("tool_call_id"), str)
and m["tool_call_id"] in dropped_ids
):
continue
out.append(m)
return out
def _splice_image_into_last_user(messages: list[dict], image_part: dict) -> None:
"""Splice an image content part into the last user message, in place.
String content becomes a text part plus the image; an existing content-part
list gets the image appended; any other shape is replaced by the lone image.
With no user message present, a new user turn carrying the image is appended."""
for msg in reversed(messages):
if msg.get("role") != "user":
continue
existing = msg.get("content")
if isinstance(existing, str):
msg["content"] = [{"type": "text", "text": existing}, image_part]
elif isinstance(existing, list):
existing.append(image_part)
else:
msg["content"] = [image_part]
break
else:
messages.append({"role": "user", "content": [image_part]})
def _openai_messages_for_passthrough(payload) -> list[dict]:
"""Build OpenAI-format message dicts for the /v1/chat/completions
passthrough path.
``payload.messages`` are dumped through Pydantic (dropping unset optional
fields), so they're already standard OpenAI format -- including
``role="tool"`` tool-result messages and assistant messages carrying
structured ``tool_calls``. Content-parts images already in the list are
left untouched.
When a client uses Studio's legacy ``image_base64`` top-level field, the
image is re-encoded to PNG (llama-server's stb_image has limited format
support) and spliced into the last user message as an OpenAI ``image_url``
content part so vision + function-calling requests work transparently.
"""
messages = _strip_provider_synthetic_tool_history(
_drop_empty_assistant_sentinels([m.model_dump(exclude_none = True) for m in payload.messages])
)
if not payload.image_base64:
return messages
try:
raw = base64.b64decode(payload.image_base64)
png_b64 = _image_bytes_to_png_b64(raw)
except Exception:
raise HTTPException(
status_code = 400,
detail = "Failed to process image.",
)
data_url = f"data:image/png;base64,{png_b64}"
image_part = {"type": "image_url", "image_url": {"url": data_url}}
_splice_image_into_last_user(messages, image_part)
return messages
def _flatten_content_parts_for_local_template(messages: list[dict]) -> list[dict]:
"""Flatten OpenAI content-part lists to plain strings.
Local text templates take string content and raise on part lists (e.g. a
remote ``image_url`` that leaves ``image is None``): keep the text parts,
drop the rest, like the plain non-GGUF path. GGUF keeps the parts."""
out = []
for msg in messages:
content = msg.get("content")
if isinstance(content, list):
text_parts = [
part.get("text", "")
for part in content
if isinstance(part, dict) and part.get("type") == "text"
]
msg = {**msg, "content": "\n".join(text_parts) if text_parts else ""}
out.append(msg)
return out
def _structured_tool_history_for_local_template(messages: list[dict]) -> list[dict]:
"""Deserialize assistant ``tool_calls[].function.arguments`` JSON strings to
mappings for local templating.
Clients send prior-turn arguments as JSON strings, but local templates take
mappings (some raise on strings). Only the internal messages copy is
rewritten; the HTTP response stays OpenAI-shaped and unparseable strings
are left untouched."""
out = []
for msg in messages:
tool_calls = msg.get("tool_calls")
if isinstance(tool_calls, list) and tool_calls:
new_calls = []
for tc in tool_calls:
fn = tc.get("function") if isinstance(tc, dict) else None
args = fn.get("arguments") if isinstance(fn, dict) else None
if isinstance(args, str):
try:
parsed = json.loads(args)
except ValueError:
parsed = None
if isinstance(parsed, dict):
tc = {**tc, "function": {**fn, "arguments": parsed}}
new_calls.append(tc)
msg = {**msg, "tool_calls": new_calls}
out.append(msg)
return out
def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict], bool]:
"""Build llama-server messages for the standard GGUF chat path.
llama-server accepts OpenAI multimodal content parts directly. Preserve all
per-turn ``image_url`` parts so multi-image chat history keeps each image
attached to its original turn.
"""
# Coalesce only on the GGUF chat path (strict Jinja template); the tool path
# reuses this via _set_or_prepend_system_message. Passthrough forwards verbatim.
messages = _coalesce_consecutive_user_turns(
_strip_provider_synthetic_tool_history(
_drop_empty_assistant_sentinels(
[m.model_dump(exclude_none = True) for m in payload.messages]
)
)
)
has_message_image = any(
isinstance(msg.get("content"), list)
and any(part.get("type") == "image_url" for part in msg["content"])
for msg in messages
)
if payload.image_base64 and not has_message_image:
# Legacy bytes can be any format; the normalizer below sniffs and
# re-encodes to PNG, so the declared mime is rewritten anyway.
image_part = {
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{payload.image_base64}",
},
}
_splice_image_into_last_user(messages, image_part)
has_image = _normalize_anthropic_openai_images(messages, is_vision)
return messages, has_image
def _extract_response_format(payload):
"""Return the ``response_format`` field on an incoming ChatCompletionRequest
(or None). The model uses ``extra="allow"`` so pydantic stashes unknown
top-level fields in ``model_extra``; OpenAI-SDK clients spread ``extra_body``
into the request body top level, where guided-decoding recipes park their
JSON-schema response_format.
"""
extra = getattr(payload, "model_extra", None)
if not isinstance(extra, dict):
return None
rf = extra.get("response_format")
return rf if isinstance(rf, dict) else None
def _build_openai_passthrough_body(
payload,
backend_ctx = None,
llama_backend = None,
) -> dict:
"""Assemble the llama-server request body from a ChatCompletionRequest.
Only known OpenAI / llama-server fields are forwarded, so Studio-specific
extensions (``enable_tools``, ``enabled_tools``, ``session_id``, ...) never
leak to the backend.
"""
messages = _openai_messages_for_passthrough(payload)
system_prompt, _, _ = _extract_content_parts(payload.messages)
messages = _set_or_prepend_system_message(messages, system_prompt)
tool_choice = payload.tool_choice if payload.tool_choice is not None else "auto"
tools = payload.tools
if payload.tool_choice == "none" and not _has_openai_tool_history(payload.messages):
tools = None
# Forward per-request reasoning fields (enable_thinking / reasoning_effort /
# preserve_thinking) via chat_template_kwargs so the Jinja template renders
# in the caller's mode, gated on the active template's capabilities exactly
# like the non-passthrough paths.
tpl_kwargs = (
llama_backend._request_reasoning_kwargs(
payload.enable_thinking,
payload.reasoning_effort,
payload.preserve_thinking,
)
if llama_backend is not None
else None
)
return _build_passthrough_payload(
messages,
tools,
payload.temperature,
payload.top_p,
payload.top_k,
# Honor max_completion_tokens on the tools/response_format passthrough too.
_effective_openai_max_tokens(payload),
payload.stream,
stop = payload.stop,
min_p = payload.min_p,
repetition_penalty = payload.repetition_penalty,
presence_penalty = payload.presence_penalty,
tool_choice = tool_choice,
response_format = _extract_response_format(payload),
chat_template_kwargs = tpl_kwargs,
backend_ctx = backend_ctx,
seed = payload.seed,
stream_options = payload.stream_options,
)
async def _openai_passthrough_stream(
request,
cancel_event,
llama_backend,
payload,
model_name,
completion_id,
monitor_id: Optional[str] = None,
):
"""Streaming client-side pass-through for /v1/chat/completions.
Forwards the client's OpenAI function-calling request to llama-server and
relays the SSE stream back with minimal normalization (reasoning-only
deltas gain ``content: ""``; errors and missing terminal markers get a
closing ``[DONE]``), preserving llama-server's native response ``id``,
``finish_reason`` (including ``"tool_calls"``), ``delta.tool_calls``, and
any client-requested trailing ``usage`` chunk so the client sees a
standard OpenAI response.
Reasoning/tool-call splitting is delegated to llama-server (``--jinja
--reasoning-format auto``), so ``delta.content`` carries no raw markup and is
deliberately not re-parsed locally, unlike the ``/completion`` paths.
"""
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_tracker = _TrackedCancel(cancel_event, *_cancel_keys)
_tracker.__enter__()
target_url = f"{llama_backend.base_url}/v1/chat/completions"
upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend)
client = None
resp = None
send_task: Optional[asyncio.Task[Optional[httpx.Response]]] = None
async def _aclose_send_task(task: Optional[asyncio.Task[Optional[httpx.Response]]]) -> None:
if task is None:
return
if not task.done():
task.cancel()
try:
task_resp = await task
if task_resp is not None:
try:
await task_resp.aclose()
except Exception:
pass
except (asyncio.CancelledError, Exception):
pass
# Keep tracker cleanup paired if pre-header dispatch is cancelled.
try:
body = _build_openai_passthrough_body(
payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend
)
# Text-form tool calls from small models get promoted to structured calls on
# the way back (declared client tools only); requests without tools or with
# auto_heal_tool_calls=false keep the unhealed relay. tool_choice constrains
# the allowlist ("none" disables, a forced function narrows to it).
_allowed_tools = heal_gate(
payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice")
)
# Keep the pre-header window short so accepted SSE clients receive
# immediate headers in the common timeout-reduced stall.
client = httpx.AsyncClient(
timeout = _llama_streaming_generation_timeout(),
limits = httpx.Limits(max_keepalive_connections = 0),
trust_env = False,
)
_truncate_budget = (
_OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0
)
while True:
try:
req = client.build_request("POST", target_url, json = body, headers = upstream_headers)
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
send_task = asyncio.create_task(
_send_stream_with_preheader_cancel(
client,
req,
cancel_event,
request = request,
mark_cancel_on_cancel = False,
)
)
done, _ = await asyncio.wait(
{send_task},
timeout = _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S,
return_when = asyncio.FIRST_COMPLETED,
)
if send_task not in done:
break
# Dispatch returned quickly enough to preserve pre-header status.
resp = await send_task
send_task = None
except httpx.RequestError as e:
# llama-server subprocess crashed / starting / unreachable.
logger.error("openai passthrough stream: upstream unreachable: %s", e)
api_monitor.fail(monitor_id, _friendly_error(e))
await _aclose_send_task(send_task)
await _aclose_stream_resources(resp = resp, client = client)
raise HTTPException(
status_code = 502,
detail = _friendly_error(e),
)
if resp is None and send_task is not None and not send_task.done():
break
if resp is None:
if cancel_event is not None:
cancel_event.set()
api_monitor.finish(monitor_id, "cancelled")
try:
await _aclose_send_task(send_task)
await _aclose_stream_resources(client = client)
finally:
_tracker.__exit__(None, None, None)
return _SameTaskStreamingResponse(
iter(()),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
if resp.status_code == 200:
break
err_bytes = await resp.aread()
err_text = err_bytes.decode("utf-8", errors = "replace")
logger.error(
"openai passthrough upstream error: status=%s body=%s",
resp.status_code,
err_text[:500],
)
upstream_status = resp.status_code
try:
await resp.aclose()
except Exception:
pass
resp = None
# Opt-in overflow policy: shrink and retry instead of a fatal 400.
if (
_truncate_budget > 0
and _classify_llama_generation_error(Exception(err_text))
and _apply_overflow_truncation(body, err_text)
):
_truncate_budget -= 1
continue
try:
await client.aclose()
except Exception:
pass
api_monitor.fail(monitor_id, err_text[:500])
raise _openai_passthrough_error(upstream_status, err_text)
# Keep tracker cleanup paired if pre-header dispatch is cancelled after we
# have already committed headers.
async def _stream():
# Same httpx lifecycle pattern as _anthropic_passthrough_stream:
# save resp.aiter_lines() so the finally block can aclose() it on
# our task. See that function for full rationale.
lines_iter = None
# Watchers unblock aiter_lines() during prefill, before in-loop
# cancel/disconnect checks can run.
cancel_watcher = None
disconnect_watcher = None
nonlocal resp, send_task, first_token_deadline, _truncate_budget
nonlocal client
monitor_done = False
saw_finish_reason = False
saw_done = False
saw_stream_error = False
saw_stream_item = False
saw_tool_call_delta = False
terminal_seen = False
last_chunk_id = completion_id
last_chunk_model = model_name
last_chunk_created = int(time.time())
healer = (
StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None
)
healed_call_index = 0
def _synthetic_finish_line() -> str:
healed = healer is not None and healer.healed
finish_reason = "tool_calls" if (saw_tool_call_delta or healed) else "stop"
chunk = ChatCompletionChunk(
id = last_chunk_id,
created = last_chunk_created,
model = last_chunk_model,
choices = [
ChunkChoice(
delta = ChoiceDelta(),
finish_reason = finish_reason,
)
],
)
return f"data: {chunk.model_dump_json(exclude_none = True)}"
def _healer_sse_lines(events) -> list:
# Serialize healer events as chunks matching the upstream stream's
# id/model/created so clients see one coherent completion.
nonlocal healed_call_index
lines = []
for kind, value in events:
if kind == "text":
if not value:
continue
delta = {"content": value}
else:
# parallel_tool_calls=false caps healed calls too (the SSE
# line cap only sees structured upstream deltas).
if payload.parallel_tool_calls is False and healed_call_index >= 1:
continue
delta = {
"tool_calls": [
{
"index": healed_call_index,
"id": value["id"],
"type": "function",
"function": value["function"],
}
]
}
healed_call_index += 1
chunk = {
"id": last_chunk_id,
"object": "chat.completion.chunk",
"created": last_chunk_created,
"model": last_chunk_model,
"choices": [{"index": 0, "delta": delta, "finish_reason": None}],
}
lines.append("data: " + json.dumps(chunk, ensure_ascii = False))
return lines
stall_timeout_s = _openai_compat_stream_stall_timeout()
def _terminal_read_timeout_s() -> Optional[float]:
if terminal_seen:
return _OPENAI_PASSTHROUGH_TERMINAL_GRACE_S
return stall_timeout_s
def _heal_transform(chunk_data: dict, raw_line: str) -> list:
"""SSE lines to emit in place of one upstream line (healing on)."""
choices = chunk_data.get("choices")
if not (isinstance(choices, list) and choices and isinstance(choices[0], dict)):
return [raw_line]
choice = choices[0]
delta = choice.get("delta")
delta = delta if isinstance(delta, dict) else {}
if delta.get("tool_calls"):
# Structured call streamed: grammar mode worked. Flush any held
# text (it preceded the call) and relay verbatim from here on.
lines = _healer_sse_lines(healer.structured_tool_call_seen())
if healed_call_index:
if payload.parallel_tool_calls is False:
# A healed call already consumed the single allowed
# slot; the upstream SSE cap keeps native index 0, so
# drop the native call here or the client gets two.
del delta["tool_calls"]
if delta or choice.get("finish_reason") or chunk_data.get("usage"):
lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False))
return lines
# A healed call already went out on index 0..n-1; OpenAI
# clients merge tool-call deltas by index, so shift the
# native calls into the next indexes or they would merge
# into the healed call.
for tc in delta["tool_calls"]:
if isinstance(tc, dict) and isinstance(tc.get("index"), int):
tc["index"] += healed_call_index
return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)]
return lines + [raw_line]
content = delta.get("content")
finish = choice.get("finish_reason")
if not isinstance(content, str) or not content:
if not finish:
return [raw_line]
# Finish chunk: last-chance heal of the residue, and rewrite a
# "stop" into "tool_calls" when text-form calls were promoted.
lines = _healer_sse_lines(healer.finalize())
if healer.healed and finish == "stop":
choice["finish_reason"] = "tool_calls"
return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)]
return lines + [raw_line]
events = healer.feed(content)
if finish:
events += healer.finalize()
if not finish and events == [("text", content)]:
# Nothing held or promoted: the healer passed the chunk
# through whole, so keep the verbatim upstream bytes.
return [raw_line]
del delta["content"]
prefix_lines = []
if delta:
prefix_chunk = {k: v for k, v in chunk_data.items() if k != "usage"}
prefix_choice = dict(choice)
prefix_choice["delta"] = dict(delta)
prefix_choice["finish_reason"] = None
prefix_chunk["choices"] = [prefix_choice]
prefix_lines.append("data: " + json.dumps(prefix_chunk, ensure_ascii = False))
delta.clear()
lines = prefix_lines + _healer_sse_lines(events)
if delta or finish or chunk_data.get("usage"):
if healer.healed and finish == "stop":
choice["finish_reason"] = "tool_calls"
lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False))
return lines
try:
while True:
if send_task is not None:
last_keepalive_at = time.monotonic()
while not send_task.done():
# Wake often enough that _preheader_cancelled keeps
# cancel/disconnect latency sub-second during prefill;
# keepalives still pace off last_keepalive_at.
wait_timeout = min(
_STREAM_DISCONNECT_POLL_TIMEOUT_S,
_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S,
)
done, _ = await asyncio.wait(
{send_task},
timeout = wait_timeout,
return_when = asyncio.FIRST_COMPLETED,
)
if send_task in done:
break
if await _preheader_cancelled(cancel_event, request):
api_monitor.finish(monitor_id, "cancelled")
return
# The downstream SSE response is already committed;
# keep strict clients and proxies from treating a long
# llama-server prefill/header wait as a dead stream.
now = time.monotonic()
if (
now - last_keepalive_at
>= _OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S
):
last_keepalive_at = now
yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE
if resp is None:
try:
resp = send_task.result()
except httpx.RequestError as e:
logger.error(
"openai passthrough stream: upstream unreachable: %s", e
)
api_monitor.fail(monitor_id, _friendly_error(e))
yield _openai_stream_error_sse(_openai_stream_error_chunk(e))
return
send_task = None
if resp is None:
api_monitor.finish(monitor_id, "cancelled")
return
if resp.status_code == 200:
break
err_bytes = await resp.aread()
err_text = err_bytes.decode("utf-8", errors = "replace")
logger.error(
"openai passthrough upstream error: status=%s body=%s",
resp.status_code,
err_text[:500],
)
upstream_status = resp.status_code
try:
await resp.aclose()
except Exception:
pass
resp = None
if (
_truncate_budget > 0
and _classify_llama_generation_error(Exception(err_text))
and _apply_overflow_truncation(body, err_text)
):
_truncate_budget -= 1
req = client.build_request(
"POST", target_url, json = body, headers = upstream_headers
)
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
send_task = asyncio.create_task(
_send_stream_with_preheader_cancel(
client,
req,
cancel_event,
request = request,
mark_cancel_on_cancel = False,
)
)
continue
upstream_error = _openai_passthrough_error(upstream_status, err_text)
error_payload = (
upstream_error.detail
if isinstance(upstream_error.detail, dict)
else openai_error_body(
str(upstream_error.detail),
status = upstream_status,
)
)
api_monitor.fail(monitor_id, err_text[:500])
yield _openai_stream_error_sse(error_payload)
return
cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp))
disconnect_watcher = asyncio.create_task(
_await_disconnect_then_close(request, resp, cancel_event)
)
lines_iter = resp.aiter_lines()
async for raw_line in _aiter_llama_stream_items(
lines_iter,
cancel_event = cancel_event,
request = request,
first_token_deadline = first_token_deadline,
response = resp,
post_first_item_read_timeout_s = _terminal_read_timeout_s,
):
if not raw_line:
continue
if not raw_line.startswith("data:"):
continue
saw_stream_item = True
data_text = raw_line[5:].strip()
if data_text == "[DONE]":
saw_done = True
# Upstream ended without a finish chunk: heal the residue
# first so the synthetic finish sees healer.healed.
if healer is not None and not saw_stream_error:
for held_line in _healer_sse_lines(healer.finalize()):
_monitor_openai_sse_line(
monitor_id, held_line, llama_backend.context_length
)
yield held_line + "\n\n"
if (
not saw_finish_reason
and not saw_stream_error
and not cancel_event.is_set()
):
finish_line = _synthetic_finish_line()
_monitor_openai_sse_line(
monitor_id,
finish_line,
llama_backend.context_length,
)
yield finish_line + "\n\n"
saw_finish_reason = True
_monitor_openai_sse_line(
monitor_id,
raw_line,
llama_backend.context_length,
)
yield raw_line + "\n\n"
monitor_done = True
break
raw_line = _normalize_openai_passthrough_sse_line(
raw_line,
cap_parallel_tool_calls = payload.parallel_tool_calls is False,
)
data_text = raw_line[5:].strip()
try:
chunk_data = json.loads(data_text)
except json.JSONDecodeError:
chunk_data = None
if isinstance(chunk_data, dict):
if isinstance(chunk_data.get("id"), str):
last_chunk_id = chunk_data["id"]
if isinstance(chunk_data.get("model"), str):
last_chunk_model = chunk_data["model"]
if isinstance(chunk_data.get("created"), int):
last_chunk_created = chunk_data["created"]
choices = chunk_data.get("choices")
if isinstance(choices, list) and choices:
choice = choices[0]
if isinstance(choice, dict):
if choice.get("finish_reason"):
saw_finish_reason = True
delta = choice.get("delta")
if isinstance(delta, dict) and delta.get("tool_calls"):
saw_tool_call_delta = True
# Detect an error chunk independently of API monitoring
# (skip_api_monitor returns early), else the synthetic
# finish would fire after a failed stream.
if _monitor_openai_error_message(chunk_data):
saw_stream_error = True
# With healing active, a content-bearing line may be replaced by
# held/promoted chunks; otherwise the single (already
# normalized) line relays unchanged (monitored exactly as
# emitted either way).
if (
healer is not None
and not healer.dormant
and isinstance(chunk_data, dict)
and not saw_stream_error
):
out_lines = _heal_transform(chunk_data, raw_line)
else:
out_lines = [raw_line]
# If a trailing usage-only chunk (include_usage) arrives before
# any finish chunk, emit the synthetic finish first so the order
# stays finish -> usage -> [DONE], matching the other streams.
if (
isinstance(chunk_data, dict)
and chunk_data.get("usage")
and not (
isinstance(chunk_data.get("choices"), list) and chunk_data["choices"]
)
and not saw_finish_reason
and not saw_stream_error
and not cancel_event.is_set()
):
if healer is not None:
# Residue must precede the finish it may upgrade.
held = _healer_sse_lines(healer.finalize())
for held_line in held:
_monitor_openai_sse_line(
monitor_id, held_line, llama_backend.context_length
)
yield held_line + "\n\n"
finish_line = _synthetic_finish_line()
_monitor_openai_sse_line(
monitor_id, finish_line, llama_backend.context_length
)
yield finish_line + "\n\n"
saw_finish_reason = True
for out_line in out_lines:
monitor_event = _monitor_openai_sse_line(
monitor_id,
out_line,
llama_backend.context_length,
)
if monitor_event == "error":
saw_stream_error = True
# Relay to preserve llama-server's native id,
# finish_reason, delta.tool_calls, and usage chunks.
yield out_line + "\n\n"
if monitor_event == "done":
monitor_done = True
break
terminal_state = (
_openai_passthrough_terminal_state_from_data(chunk_data)
if out_line is raw_line
else _openai_passthrough_sse_line_terminal_state(out_line)
)
if terminal_state == "usage" or (
terminal_state == "finish" and not _wants_stream_usage(payload)
):
done_line = _SSE_DONE_LINE
_monitor_openai_sse_line(
monitor_id,
done_line,
llama_backend.context_length,
)
yield done_line + "\n\n"
saw_done = True
monitor_done = True
break
if terminal_state == "finish":
terminal_seen = True
if monitor_done:
break
if not saw_done and not saw_stream_error and not cancel_event.is_set():
# Synthesize a finish chunk only if one was not already
# emitted (e.g. before a trailing usage-only chunk), but
# always close with [DONE] whenever the upstream omitted it,
# so the stream ends on the [DONE] sentinel either way.
if healer is not None:
for held_line in _healer_sse_lines(healer.finalize()):
_monitor_openai_sse_line(
monitor_id, held_line, llama_backend.context_length
)
yield held_line + "\n\n"
if not saw_finish_reason:
finish_line = _synthetic_finish_line()
_monitor_openai_sse_line(
monitor_id,
finish_line,
llama_backend.context_length,
)
yield finish_line + "\n\n"
done_line = _SSE_DONE_LINE
_monitor_openai_sse_line(
monitor_id,
done_line,
llama_backend.context_length,
)
yield done_line + "\n\n"
monitor_done = True
if not monitor_done:
api_monitor.finish(
monitor_id,
"cancelled" if cancel_event.is_set() else "completed",
)
except asyncio.CancelledError:
api_monitor.finish(monitor_id, "cancelled")
raise
except httpx.ReadTimeout as e:
if terminal_seen and not saw_stream_error and not cancel_event.is_set():
done_line = _SSE_DONE_LINE
_monitor_openai_sse_line(
monitor_id,
done_line,
llama_backend.context_length,
)
yield done_line + "\n\n"
api_monitor.finish(monitor_id)
return
if cancel_event.is_set():
api_monitor.finish(monitor_id, "cancelled")
return
logger.error(
"openai passthrough stream %s: %s",
"stalled mid-response" if saw_stream_item else "timeout",
e,
)
api_monitor.fail(monitor_id, _friendly_error(e))
get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
err = _openai_stream_error_chunk(e)
yield _openai_stream_error_sse(err)
except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e:
# Watcher closed resp on cancel. Emit nothing extra; the client
# initiated the cancel or already disconnected.
if not cancel_event.is_set():
api_monitor.fail(monitor_id, "Stream interrupted")
get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
raise
api_monitor.finish(monitor_id, "cancelled")
except HTTPException as exc:
status_code = getattr(exc, "status_code", 500) or 500
detail = exc.detail
error_payload = (
detail
if isinstance(detail, dict) and "error" in detail
else openai_error_body(str(detail), status = status_code)
)
api_monitor.fail(monitor_id, str(detail))
yield _openai_stream_error_sse(error_payload)
except Exception as e:
if cancel_event.is_set():
api_monitor.finish(monitor_id, "cancelled")
return
# 200 headers already flushed; errors must go in the SSE body.
logger.error("openai passthrough stream error: %s", e)
api_monitor.fail(monitor_id, _friendly_error(e))
get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
err = _openai_stream_error_chunk(e)
yield _openai_stream_error_sse(err)
finally:
# _aclose_stream_resources re-raises a close-time CancelledError
# only after finishing teardown, and the tracker exits either way.
try:
await _aclose_send_task(send_task)
await _aclose_stream_resources(
watchers = (cancel_watcher, disconnect_watcher),
iterator = lines_iter,
resp = resp,
client = client,
)
finally:
_tracker.__exit__(None, None, None)
async def _unstarted_cleanup() -> None:
# Client disconnected before the body stream started, so _stream()'s
# finally never ran. Release the eagerly-opened upstream resp/client
# and the cancel-registry entry here; the watchers and line iterator
# are created inside _stream(), so there is nothing else to close.
try:
await _aclose_send_task(send_task)
await _aclose_stream_resources(resp = resp, client = client)
finally:
_tracker.__exit__(None, None, None)
return _SameTaskStreamingResponse(
_stream(),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
unstarted_cleanup = _unstarted_cleanup,
)
except BaseException as exc:
if isinstance(exc, asyncio.CancelledError):
if cancel_event is not None:
cancel_event.set()
api_monitor.finish(monitor_id, "cancelled")
else:
detail = exc.detail if isinstance(exc, HTTPException) else _friendly_error(exc)
api_monitor.fail(monitor_id, str(detail))
try:
await _aclose_send_task(send_task)
await _aclose_stream_resources(resp = resp, client = client)
finally:
_tracker.__exit__(None, None, None)
raise
async def _openai_passthrough_non_streaming(
llama_backend,
payload,
model_name,
monitor_id: Optional[str] = None,
*,
request: Optional[Request] = None,
cancel_event = None,
):
"""Non-streaming client-side pass-through for /v1/chat/completions.
Returns llama-server's JSON response verbatim so the client sees the native
response ``id``, ``finish_reason`` (including ``"tool_calls"``), structured
``tool_calls``, and accurate ``usage`` token counts.
"""
target_url = f"{llama_backend.base_url}/v1/chat/completions"
upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend)
body = _build_openai_passthrough_body(
payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend
)
body["stream"] = False
body.pop("stream_options", None)
_truncate_budget = (
_OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0
)
async def _post(body_to_send):
if cancel_event is None and request is None:
return await nonstreaming_client().post(
target_url,
json = body_to_send,
headers = upstream_headers,
timeout = _llama_non_streaming_generation_timeout(),
)
if cancel_event is None:
cancel = threading.Event()
else:
cancel = cancel_event
client = _cancelable_nonstreaming_client()
watcher = asyncio.create_task(
_await_cancel_or_disconnect_then_close_client(
cancel_event = cancel,
request = request,
client = client,
)
)
try:
try:
response = await client.post(
target_url,
json = body_to_send,
headers = upstream_headers,
timeout = _llama_non_streaming_generation_timeout(),
)
except httpx.RequestError:
if cancel.is_set():
raise asyncio.CancelledError()
raise
if cancel.is_set():
raise asyncio.CancelledError()
return response
finally:
watcher.cancel()
try:
await watcher
except (asyncio.CancelledError, Exception):
pass
try:
await client.aclose()
except Exception:
pass
while True:
try:
resp = await _post(body)
except asyncio.CancelledError:
api_monitor.finish(monitor_id, "cancelled")
raise
except httpx.RequestError as e:
# llama-server subprocess crashed / starting / unreachable. Surface the
# same friendly message the sync chat path emits so operators don't see
# a bare 500 with no diagnostic.
logger.error("openai passthrough non-streaming: upstream unreachable: %s", e)
api_monitor.fail(monitor_id, _friendly_error(e))
get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
raise HTTPException(
status_code = 502,
detail = _friendly_error(e),
)
if resp.status_code == 200:
break
# Opt-in overflow policy: shrink and retry instead of a fatal 400.
if (
_truncate_budget > 0
and _classify_llama_generation_error(Exception(resp.text))
and _apply_overflow_truncation(body, resp.text)
):
_truncate_budget -= 1
continue
api_monitor.fail(monitor_id, 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
_allowed_tools = heal_gate(
payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice")
)
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,
)
api_monitor.finish(monitor_id)
return Response(content = resp.content, media_type = "application/json")
# Opt-in single-retry nudge: the model clearly tried to call a tool (signal
# present) but nothing parseable/declared came out, so re-ask once with the
# original prompt prefix intact (llama-server reuses the slot's KV cache)
# plus a two-message nudge suffix. The retry replaces the original response
# only when it actually yields a usable call.
if (
_allowed_tools
and nudge_enabled(payload.nudge_tool_calls)
and nudge_should_retry(data, _allowed_tools, body.get("tools"))
):
retry_body = {
**body,
"messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)],
}
try:
retry_resp = await _post(retry_body)
if retry_resp.status_code == 200:
retry_data = retry_resp.json()
if response_has_promotable_calls(retry_data, _allowed_tools, body.get("tools")):
resp, data = retry_resp, retry_data
except asyncio.CancelledError:
api_monitor.finish(monitor_id, "cancelled")
raise
except (httpx.RequestError, ValueError) as exc:
logger.warning("tool-call nudge retry failed; keeping original: %s", exc)
changed = False
for choice in data.get("choices", []):
if not isinstance(choice, dict):
continue
msg = choice.get("message")
if not isinstance(msg, dict):
continue
# Small models emit tool calls as text instead of structured tool_calls;
# promote them (declared client tools only) so the agent sees a real call.
# Truncation wins over the upgrade (same rule as the streaming and
# Anthropic paths): a call cut off at max_tokens keeps
# finish_reason="length" so the client knows the arguments may be
# incomplete, while the healed call itself stays attached.
if _allowed_tools and heal_openai_message(msg, _allowed_tools, body.get("tools")):
if choice.get("finish_reason") == "stop":
choice["finish_reason"] = "tool_calls"
changed = True
# 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
# 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
_monitor_openai_chunk(monitor_id, data, llama_backend.context_length)
api_monitor.finish(monitor_id)
if not changed:
# Nothing mutated: relay the upstream bytes verbatim, skipping a
# redundant parse + re-serialize round-trip.
return Response(content = resp.content, media_type = "application/json")
return JSONResponse(content = data)