# 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, List, Optional, Union import json import httpx from loggers import get_logger import asyncio import threading 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 _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 _clamp_finish_reason(value) -> str: """Coerce an upstream finish_reason into OpenAI's known chat values. Unknown values (including ``None``) become ``"stop"`` so local upstream quirks do not leak into the public API shape. """ return ( value if value in ( "stop", "length", "tool_calls", "content_filter", "function_call", ) else "stop" ) def _normalize_stop_sequences(raw): """Coerce an OpenAI/Anthropic ``stop`` value into the list-of-non-empty-strings shape llama-server expects, or ``None`` when absent. A bare string becomes a single-element list; empty strings are dropped (an empty stop sequence would terminate generation immediately at position 0).""" if isinstance(raw, str): return [raw] if raw else None if isinstance(raw, list): return [s for s in raw if isinstance(s, str) and s] or None return None def _effective_max_tokens(payload): """Resolve the generation cap, preferring OpenAI's replacement field. ``max_tokens`` is deprecated in favor of ``max_completion_tokens``; honor either for compatibility, but let the replacement field win when both are supplied. """ return ( payload.max_completion_tokens if payload.max_completion_tokens is not None else payload.max_tokens ) def _wants_multiple_choices(payload) -> bool: return (payload.n or 1) > 1 def _raise_unsupported_openai_parameter(param: str, message: str) -> None: raise HTTPException( status_code = 400, detail = openai_error_body( message, status = 400, code = "unsupported_parameter", param = param, ), ) def _raise_unsupported_n(path_label: str) -> None: _raise_unsupported_openai_parameter("n", f"n > 1 is not supported for {path_label}.") def _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_passthrough_error(status_code, text) -> "HTTPException": """HTTPException for a non-200 upstream response on the OpenAI passthrough (tools / response_format). An over-context upstream error is mapped to a 400 with code="context_length_exceeded" so these paths deliver the same signal as the non-passthrough path; any other upstream error keeps llama-server's message verbatim.""" if _classify_llama_generation_error(Exception(text)): return HTTPException( status_code = 400, detail = openai_error_body( _friendly_error(Exception(text)), status = 400, code = "context_length_exceeded", param = "messages", ), ) return HTTPException( status_code = status_code, detail = f"llama-server error: {text[:500]}", ) _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 _cap_parallel_tool_calls_sse_line(raw_line: str) -> str: """Drop tool_call deltas whose index >= 1 from one streamed OpenAI SSE ``data:`` line so only the first tool call survives (parallel_tool_calls=false, best-effort). Non-tool / unparseable payloads are returned byte-for-byte.""" payload = raw_line[len("data: ") :] if payload.strip() in ("", "[DONE]"): return raw_line try: obj = json.loads(payload) except Exception: return raw_line if not _drop_parallel_tool_call_deltas(obj): return raw_line return "data: " + json.dumps(obj, separators = (",", ":")) def _prompt_tokens_details(upstream): """Surface llama-server's real ``cached_tokens`` (KV-cache prompt hits) while keeping the full OpenAI ``prompt_tokens_details`` shape. Defaults to zero when the upstream usage doesn't carry it, so the field is always present.""" out = {"cached_tokens": 0, "audio_tokens": 0} if isinstance(upstream, dict): out.update({k: v for k, v in upstream.items() if v is not None}) return out def _wants_stream_usage(payload) -> bool: return bool((payload.stream_options or {}).get("include_usage")) def _openai_stream_usage_chunk( payload, completion_id, created, model_name, stream_usage, stream_timings ): """Build the final OpenAI-standard usage chunk (choices=[], usage populated) for a chat stream. Returns the SSE ``data:`` line, or None when the client did not opt in via ``stream_options.include_usage`` (or no usage exists).""" if not _wants_stream_usage(payload): return None if not (stream_usage or stream_timings): return None _usage = stream_usage or {} _prompt_tokens = _usage.get("prompt_tokens") or 0 _completion_tokens = _usage.get("completion_tokens") or 0 _total_tokens = _usage.get("total_tokens") or (_prompt_tokens + _completion_tokens) usage_chunk = ChatCompletionChunk( id = completion_id, created = created, model = model_name, choices = [], usage = CompletionUsage( prompt_tokens = _prompt_tokens, completion_tokens = _completion_tokens, total_tokens = _total_tokens, prompt_tokens_details = _prompt_tokens_details(_usage.get("prompt_tokens_details")), ), timings = stream_timings, ) return f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" def _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_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 _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: float = _DEFAULT_STREAM_STALL_TIMEOUT_S ) -> None: 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 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) # Async callable invoked when the client disconnects before the body # iterator is ever advanced. A generator that never started cannot run # its own try/finally, so a stream that acquires resources before its # first yield (the passthrough opens an upstream httpx stream eagerly) # passes this to release them. self._unstarted_cleanup = unstarted_cleanup async def __call__(self, scope, receive, send) -> None: # Track whether the body iterator was ever advanced: send() only emits a # body message after the generator yields its first chunk, so a failure # before then means it 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: # The generator produced at least one chunk and is suspended in # its try/finally. Throw CancelledError into it (not aclose's # GeneratorExit) so its `except asyncio.CancelledError` handler # runs and finishes any api_monitor entry; GeneratorExit would # skip it and only run `finally`. 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: # http.response.start failed before the body iterator advanced, # so its try/finally never armed and aclose()/athrow() are no-ops # on an unstarted generator. Release any resources acquired # before the first yield via the explicit cleanup hook. aclose = getattr(self.body_iterator, "aclose", None) if aclose is not None: await aclose() if self._unstarted_cleanup is not None: try: await self._unstarted_cleanup() except Exception: pass raise ClientDisconnect() if self.background is not None: await self.background() 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. 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 if iterator is not None: try: await iterator.aclose() except Exception: pass if resp is not None: try: await resp.aclose() except Exception: pass if client is not None: try: await client.aclose() except Exception: pass 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, ) -> 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 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[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 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: if ( request is not None and response is not None and post_first_item_read_timeout_s is not None and last_item_at is not None ): stall_remaining_s = post_first_item_read_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 if ( request is not None and post_first_item_read_timeout_s is not None and now - last_item_at < post_first_item_read_timeout_s ): continue raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") if ( last_item_at is None and response is not None and post_first_item_read_timeout_s is not None ): _set_stream_response_read_timeout(response, post_first_item_read_timeout_s) last_item_at = time.monotonic() yield item from models.inference import ( LoadRequest, UnloadRequest, GenerateRequest, DiffusionLoadRequest, DiffusionGenerateRequest, DiffusionGenerateResponse, DiffusionGenerateProgressResponse, DiffusionStatusResponse, DiffusionLoadProgressResponse, GalleryImage, GalleryListResponse, 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.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 = """ """ 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 /