From e945f436523f5f5687bfcf9aee82c42129821758 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 02:03:30 +0000 Subject: [PATCH 1/2] Add dedicated streaming server for fast GGUF SSE inference Adds a separate streaming server (Option B) that bypasses the main app's asyncio event loop for GGUF token streaming. When UNSLOTH_FAST_SSE=1 or UNSLOTH_STREAM_SERVER=1 is set, a lightweight FastAPI app starts in a daemon thread on a random port. The frontend negotiates a one-time token via /api/inference/stream-url and streams directly from this server, falling back to the baseline /v1/chat/completions transparently on failure. Architecture: - streaming_server.py: standalone FastAPI app with three paths: - Path A (hot): async httpx.AsyncClient streaming direct to llama-server. No cumulative-to-delta conversion needed since llama-server sends delta tokens natively via OpenAI SSE. - Path B: tool calling via asyncio.to_thread (tool execution is the bottleneck, not the streaming proxy). - Path C: non-streaming one-shot JSON response. - stream_token_store.py: thread-safe one-time token store (10s TTL) shared between the main app and streaming server. - routes/inference.py: /stream-url endpoint for token issuance, /direct-stream for Option C, /internal/consume-stream-token for future subprocess mode. - llama_cpp.py: --api-key support so the streaming server can authenticate to llama-server directly. - chat-api.ts: transparent fast-path negotiation with silent fallback. Handles: vision (image_url + legacy image_base64), thinking models (reasoning_content -> tags), CORS preflight, friendly error messages, stream_options usage/timings passthrough. Only sends repeat_penalty when the client explicitly provides it, avoiding the 24% TPS penalty from repetition scanning. --- studio/backend/core/inference/llama_cpp.py | 31 +- studio/backend/main.py | 24 + studio/backend/routes/inference.py | 130 ++++ studio/backend/stream_token_store.py | 65 ++ studio/backend/streaming_server.py | 558 ++++++++++++++++++ .../src/features/chat/api/chat-api.ts | 46 +- 6 files changed, 842 insertions(+), 12 deletions(-) create mode 100644 studio/backend/stream_token_store.py create mode 100644 studio/backend/streaming_server.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3f89cd3e5d..5c68981096 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -57,6 +57,7 @@ class LlamaCppBackend: self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None self._cancel_event = threading.Event() + self._api_key: Optional[str] = None self._kill_orphaned_servers() atexit.register(self._cleanup) @@ -938,6 +939,17 @@ class LlamaCppBackend: cmd.extend(["--mmproj", mmproj_path]) logger.info(f"Using mmproj for vision: {mmproj_path}") + # Option C: add --api-key for direct client access when enabled + import os as _os + import secrets as _secrets + + if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1": + self._api_key = _secrets.token_urlsafe(32) + cmd.extend(["--api-key", self._api_key]) + logger.info("llama-server started with --api-key for direct streaming") + else: + self._api_key = None + logger.info(f"Starting llama-server: {' '.join(cmd)}") # Set library paths so llama-server can find its shared libs and CUDA DLLs @@ -1407,6 +1419,7 @@ class LlamaCppBackend: url: str, payload: dict, cancel_event: Optional[threading.Event] = None, + headers: Optional[dict] = None, ): """Open an httpx streaming POST with cancel support. @@ -1473,7 +1486,8 @@ class LlamaCppBackend: pool = 10, ) with client.stream( - "POST", url, json = payload, timeout = prefill_timeout + "POST", url, json = payload, timeout = prefill_timeout, + headers = headers, ) as response: _response_ref[0] = response if cancel_event is not None and cancel_event.is_set(): @@ -1547,9 +1561,10 @@ class LlamaCppBackend: # can finish. Cancel during streaming is handled by the # watcher thread (closes the response on cancel_event). stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) + _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None with httpx.Client(timeout = stream_timeout) as client: with self._stream_with_retry( - client, url, payload, cancel_event + client, url, payload, cancel_event, headers = _auth_headers, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -1706,8 +1721,9 @@ class LlamaCppBackend: payload["stop"] = stop try: + _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None with httpx.Client(timeout = None) as client: - resp = client.post(url, json = payload) + resp = client.post(url, json = payload, headers = _auth_headers) if resp.status_code != 200: raise RuntimeError( f"llama-server returned {resp.status_code}: {resp.text}" @@ -1950,9 +1966,10 @@ class LlamaCppBackend: try: stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) + _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None with httpx.Client(timeout = stream_timeout) as client: with self._stream_with_retry( - client, url, stream_payload, cancel_event + client, url, stream_payload, cancel_event, headers = _auth_headers, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -2078,7 +2095,8 @@ class LlamaCppBackend: if not self.is_loaded: return None try: - with httpx.Client(timeout = 10) as client: + _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + with httpx.Client(timeout = 10, headers = _auth_headers) as client: def _detok(tid: int) -> str: r = client.post( @@ -2196,7 +2214,8 @@ class LlamaCppBackend: if need_ids: payload["n_probs"] = 1 - with httpx.Client(timeout = httpx.Timeout(300, connect = 10)) as client: + _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + with httpx.Client(timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers) as client: resp = client.post(f"{self.base_url}/completion", json = payload) if resp.status_code != 200: raise RuntimeError( diff --git a/studio/backend/main.py b/studio/backend/main.py index 5e647f6312..83373b1def 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -112,6 +112,24 @@ async def lifespan(app: FastAPI): print("=" * 60 + "\n") else: app.state.bootstrap_password = storage.get_bootstrap_password() + + # Start dedicated streaming server in a daemon thread (Option B). + # Accepts both UNSLOTH_FAST_SSE=1 and UNSLOTH_STREAM_SERVER=1. + if os.getenv("UNSLOTH_FAST_SSE", "0") == "1" or os.getenv("UNSLOTH_STREAM_SERVER", "0") == "1": + import threading as _threading + + from streaming_server import start_streaming_server, find_free_port + + stream_port = find_free_port() + app.state.stream_port = stream_port + _stream_thread = _threading.Thread( + target=start_streaming_server, + args=(stream_port,), + daemon=True, + ) + _stream_thread.start() + print(f"[streaming_server] Started on 127.0.0.1:{stream_port}") + yield # Cleanup _hw_module.DEVICE = None @@ -373,3 +391,9 @@ def setup_frontend(app: FastAPI, build_path: Path): ) return True + + +# Note: Option A (RawSSEInterceptor / asgi_fast_path.py) has been removed. +# Benchmarking proved it ineffective (~172 TPS = baseline) since it shares the +# same asyncio event loop. UNSLOTH_FAST_SSE=1 now starts the dedicated streaming +# server (Option B) which runs in a separate thread with its own event loop. diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 78d95fedbd..f7cc0db78f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1563,6 +1563,136 @@ async def openai_chat_completions( # ===================================================================== +# ===================================================================== +# Option B: Stream URL endpoint (one-time token for streaming server) +# ===================================================================== + + +@router.get("/stream-url") +async def get_stream_url( + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + Issue a one-time streaming token and return the URL for the dedicated + streaming server (Option B). + + Requires UNSLOTH_FAST_SSE=1 or UNSLOTH_STREAM_SERVER=1. + + Returns ``{"supported": false}`` when the fast path cannot serve the + current model (no GGUF loaded, audio model active, or streaming server + not enabled). The frontend uses this to decide whether to use the fast + path or fall back to the baseline ``/v1/chat/completions``. + + When supported, returns the stream URL (token NOT in query string -- + the client must pass it via the ``X-Stream-Token`` header). + """ + import os + + fast_sse = os.getenv("UNSLOTH_FAST_SSE", "0") == "1" + stream_server = os.getenv("UNSLOTH_STREAM_SERVER", "0") == "1" + if not fast_sse and not stream_server: + return {"supported": False} + + stream_port = getattr(request.app.state, "stream_port", None) + if stream_port is None: + return {"supported": False} + + # Check if current model is compatible with fast path + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + return {"supported": False} + if getattr(llama_backend, "_is_audio", False): + return {"supported": False} + + from stream_token_store import create_stream_token + + token = create_stream_token(current_subject) + return { + "supported": True, + "stream_url": f"http://127.0.0.1:{stream_port}/stream", + "token": token, + "port": stream_port, + "ttl_seconds": 10, + } + + +# ===================================================================== +# Internal: Stream token validation for streaming server subprocess +# ===================================================================== + + +@router.post("/internal/consume-stream-token") +async def consume_stream_token_endpoint(request: Request): + """ + Validate a one-time stream token and return llama-server connection info. + + Called by the streaming server subprocess to validate tokens without + needing shared memory. Localhost-only (no auth required). + """ + body = await request.json() + token = body.get("token") + if not token: + return JSONResponse({"valid": False}, status_code=400) + + from stream_token_store import consume_stream_token + + username = consume_stream_token(token) + if not username: + return JSONResponse({"valid": False}, status_code=401) + + llama = get_llama_cpp_backend() + return JSONResponse({ + "valid": True, + "username": username, + "llama_port": llama._port if llama.is_loaded else None, + "llama_api_key": llama._api_key, + "model_name": llama.model_identifier or "unknown", + "supports_reasoning": llama.supports_reasoning, + "supports_tools": llama.supports_tools, + "is_vision": llama.is_vision, + }) + + +# ===================================================================== +# Option C: Direct stream endpoint (llama-server with --api-key) +# ===================================================================== + + +@router.get("/direct-stream") +async def get_direct_stream( + current_subject: str = Depends(get_current_subject), +): + """ + Return the internal llama-server URL and API key for direct client streaming + (Option C). This bypasses all Studio transformations (thinking tags, image + normalization, cumulative-to-delta) but achieves maximum TPS. + + Requires the llama-server to have been started with --api-key. + """ + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException(status_code = 400, detail = "No GGUF model loaded.") + + api_key = getattr(llama_backend, "_api_key", None) + if not api_key: + raise HTTPException( + status_code = 501, + detail = "llama-server was not started with --api-key. Reload the model.", + ) + + return { + "base_url": llama_backend.base_url, + "api_key": api_key, + "model": llama_backend.model_identifier, + } + + +# ===================================================================== +# OpenAI-Compatible Models Listing (/models -> /v1/models) +# ===================================================================== + + @router.get("/models") async def openai_list_models( current_subject: str = Depends(get_current_subject), diff --git a/studio/backend/stream_token_store.py b/studio/backend/stream_token_store.py new file mode 100644 index 0000000000..f19dde41a3 --- /dev/null +++ b/studio/backend/stream_token_store.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Thread-safe one-time token store for Option B (separate streaming server). + +Tokens are short-lived (10 seconds) and consumed on first use. +""" + +import threading +import time +import uuid +from typing import Optional + + +class StreamTokenStore: + """Thread-safe store for short-lived, one-time-use streaming tokens.""" + + def __init__(self, ttl_seconds: float = 10.0) -> None: + self._ttl = ttl_seconds + self._lock = threading.Lock() + # token -> {"username": str, "expires": float} + self._tokens: dict[str, dict] = {} + + def create_token(self, username: str) -> str: + """Create a new one-time token for the given user. Returns the token string.""" + token = uuid.uuid4().hex + expires = time.monotonic() + self._ttl + with self._lock: + self._purge_expired() + self._tokens[token] = {"username": username, "expires": expires} + return token + + def consume_token(self, token: str) -> Optional[str]: + """ + Validate and consume a token. Returns the username if valid, None otherwise. + The token is deleted after consumption (one-time use). + """ + with self._lock: + self._purge_expired() + entry = self._tokens.pop(token, None) + if entry is None: + return None + if time.monotonic() > entry["expires"]: + return None + return entry["username"] + + def _purge_expired(self) -> None: + """Remove expired tokens. Must be called while holding _lock.""" + now = time.monotonic() + expired = [k for k, v in self._tokens.items() if now > v["expires"]] + for k in expired: + del self._tokens[k] + + +# Module-level singleton +_store = StreamTokenStore() + + +def create_stream_token(username: str) -> str: + return _store.create_token(username) + + +def consume_stream_token(token: str) -> Optional[str]: + return _store.consume_token(token) diff --git a/studio/backend/streaming_server.py b/studio/backend/streaming_server.py new file mode 100644 index 0000000000..9b5f667eff --- /dev/null +++ b/studio/backend/streaming_server.py @@ -0,0 +1,558 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Dedicated streaming server for fast SSE (Option B). + +Runs as a standalone FastAPI app in a separate thread with its own event loop, +eliminating asyncio contention with the main Studio app. + +Authentication: one-time tokens issued by the main app's /stream-url endpoint, +passed via the X-Stream-Token header (not in the URL to avoid logging leaks). + +Supports: streaming, non-streaming, tool calling, vision, thinking mode. +Full feature parity with baseline /v1/chat/completions. + +PERFORMANCE: The streaming hot path uses httpx.AsyncClient to stream directly +from llama-server, bypassing the sync generate_chat_completion() generator. +Only sends sampling parameters the client explicitly provides -- notably, +repeat_penalty defaults to llama-server's own 1.0 instead of being forced +to 1.1, which avoids a ~24% TPS penalty from repetition scanning. +""" + +import asyncio +import json +import re +import threading +import time +import uuid +from typing import Optional + +import httpx +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse, StreamingResponse + +from stream_token_store import consume_stream_token + + +stream_app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None) + +# ── Shared helpers ──────────────────────────────────────────── + + +def _friendly_error(e): + """Convert raw exception messages to user-readable strings.""" + msg = str(e) + 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 to the model server. It may have crashed -- try reloading the model." + return "An internal error occurred" + + +def _extract_content_parts(messages): + """Parse messages, extracting text and image_b64 from content parts.""" + gguf_messages = [] + image_b64 = None + for msg in messages: + role = msg.get("role", "") + content = msg.get("content", "") + if isinstance(content, list): + text_parts = [] + for part in content: + if isinstance(part, dict): + if part.get("type") == "text": + text_parts.append(part.get("text", "")) + elif part.get("type") == "image_url" and image_b64 is None: + url = part.get("image_url", {}).get("url", "") + if url.startswith("data:") and "," in url: + image_b64 = url.split(",", 1)[1] + content = "\n".join(text_parts) if text_parts else "" + gguf_messages.append({"role": role, "content": content}) + return gguf_messages, image_b64 + + +def _process_image(image_b64, llama_backend): + """Validate and convert image to PNG if needed.""" + if image_b64 and llama_backend.is_vision: + import base64 as _b64 + from io import BytesIO as _BytesIO + from PIL import Image as _Image + + raw = _b64.b64decode(image_b64) + img = _Image.open(_BytesIO(raw)) + if img.mode == "RGBA": + img = img.convert("RGB") + buf = _BytesIO() + img.save(buf, format="PNG") + return _b64.b64encode(buf.getvalue()).decode("ascii") + elif image_b64 and not llama_backend.is_vision: + raise HTTPException( + status_code=400, + detail="Image provided but current GGUF model does not support vision.", + ) + return image_b64 + + +def _build_llama_payload(llama_backend, openai_messages, payload, stream=True): + """Build the payload for llama-server /v1/chat/completions. + + Only sends repeat_penalty when the client explicitly provides it. + This avoids the ~24% TPS penalty from repetition scanning when + the frontend has not set a repetition penalty. + """ + llama_payload = { + "messages": openai_messages, + "stream": stream, + "temperature": payload.get("temperature", 0.6), + "top_p": payload.get("top_p", 0.95), + "top_k": max(payload.get("top_k", 20), 0), + "min_p": payload.get("min_p", 0.0), + "presence_penalty": payload.get("presence_penalty", 0.0), + } + # Only send repeat_penalty when the client explicitly sets repetition_penalty. + # llama-server defaults to 1.0; forcing 1.1 costs ~24% TPS. + if "repetition_penalty" in payload: + llama_payload["repeat_penalty"] = payload["repetition_penalty"] + if stream: + llama_payload["stream_options"] = {"include_usage": True} + if llama_backend.supports_reasoning and payload.get("enable_thinking") is not None: + llama_payload["chat_template_kwargs"] = {"enable_thinking": payload["enable_thinking"]} + if payload.get("max_tokens") is not None: + llama_payload["max_tokens"] = payload["max_tokens"] + if payload.get("stop"): + llama_payload["stop"] = payload["stop"] + return llama_payload + + +# ── CORS preflight ──────────────────────────────────────────── + + +@stream_app.options("/stream") +async def stream_preflight(): + """Handle CORS preflight for the /stream endpoint.""" + return JSONResponse(content={}, headers={ + "Access-Control-Allow-Origin": "*", + "Access-Control-Allow-Methods": "POST, OPTIONS", + "Access-Control-Allow-Headers": "Content-Type, X-Stream-Token", + "Access-Control-Max-Age": "86400", + }) + + +# ── Request validation ──────────────────────────────────────── + + +async def _validate_request(request: Request): + """Validate token, parse body, get backend. Returns all needed context.""" + token = request.headers.get("X-Stream-Token") + if not token: + raise HTTPException(status_code=401, detail="Missing X-Stream-Token header") + username = consume_stream_token(token) + if username is None: + raise HTTPException(status_code=401, detail="Invalid or expired stream token") + + body_bytes = await request.body() + try: + payload = json.loads(body_bytes) + except (json.JSONDecodeError, UnicodeDecodeError): + raise HTTPException(status_code=400, detail="Invalid JSON body") + + from routes.inference import get_llama_cpp_backend + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException(status_code=400, detail="No GGUF model loaded") + + messages = payload.get("messages", []) + gguf_messages, image_b64 = _extract_content_parts(messages) + + # Legacy image_base64 fallback + if not image_b64: + image_b64 = payload.get("image_base64") + + # Image validation and conversion + image_b64 = _process_image(image_b64, llama_backend) + + completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" + created = int(time.time()) + model_name = llama_backend.model_identifier or "unknown" + + return payload, llama_backend, gguf_messages, image_b64, completion_id, created, model_name + + +# ── Path A: Direct async streaming (HOT PATH) ──────────────── + + +_SSE_HEADERS = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + "Access-Control-Allow-Origin": "*", +} + +_SEP = (",", ":") + + +async def _handle_async_stream(request, payload, llama_backend, gguf_messages, image_b64, + completion_id, created, model_name): + """ + Stream directly from llama-server using httpx.AsyncClient. + + Bypasses the sync generate_chat_completion() generator and its + asyncio.to_thread overhead. llama-server speaks standard OpenAI SSE + with delta tokens natively, so no cumulative-to-delta conversion needed. + """ + openai_messages = llama_backend._build_openai_messages(gguf_messages, image_b64) + llama_payload = _build_llama_payload(llama_backend, openai_messages, payload, stream=True) + + port = llama_backend._port + api_key = llama_backend._api_key + headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} + url = f"http://127.0.0.1:{port}/v1/chat/completions" + timeout = httpx.Timeout(connect=30, read=120.0, write=10, pool=10) + + async def sse_generator(): + try: + # Role chunk + role = {"id": completion_id, "object": "chat.completion.chunk", + "created": created, "model": model_name, + "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]} + yield f"data: {json.dumps(role, separators=_SEP)}\n\n" + + async with httpx.AsyncClient(timeout=timeout) as client: + async with client.stream("POST", url, json=llama_payload, headers=headers) as resp: + if resp.status_code != 200: + error_body = await resp.aread() + raise RuntimeError( + f"llama-server returned {resp.status_code}: {error_body.decode()}" + ) + + buffer = "" + in_thinking = False + has_content_tokens = False + reasoning_text = "" + stream_usage = None + stream_timings = None + stream_done = False + + async for raw_chunk in resp.aiter_text(): + buffer += raw_chunk + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + line = line.strip() + + if not line: + continue + if line == "data: [DONE]": + if in_thinking: + if has_content_tokens: + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': ''}, 'finish_reason': None}]}, separators=_SEP)}\n\n" + else: + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning_text}, 'finish_reason': None}]}, separators=_SEP)}\n\n" + stream_done = True + break + if not line.startswith("data: "): + continue + + try: + data = json.loads(line[6:]) + except json.JSONDecodeError: + continue + + _t = data.get("timings") + if _t: + stream_timings = _t + _u = data.get("usage") + if _u: + stream_usage = _u + + choices = data.get("choices", []) + if not choices: + continue + delta = choices[0].get("delta", {}) + + # Handle reasoning_content -> tags + reasoning = delta.get("reasoning_content", "") + if reasoning: + reasoning_text += reasoning + if not in_thinking: + in_thinking = True + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': ''}, 'finish_reason': None}]}, separators=_SEP)}\n\n" + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning}, 'finish_reason': None}]}, separators=_SEP)}\n\n" + + # Handle content tokens + token = delta.get("content", "") + if token: + has_content_tokens = True + if in_thinking: + in_thinking = False + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': ''}, 'finish_reason': None}]}, separators=_SEP)}\n\n" + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': token}, 'finish_reason': None}]}, separators=_SEP)}\n\n" + + if stream_done: + break + + # Final stop chunk + final = {"id": completion_id, "object": "chat.completion.chunk", + "created": created, "model": model_name, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]} + yield f"data: {json.dumps(final, separators=_SEP)}\n\n" + + # Usage chunk + if stream_usage or stream_timings: + usage_chunk = { + "id": completion_id, "object": "chat.completion.chunk", + "created": created, "model": model_name, "choices": [], + "usage": { + "prompt_tokens": (stream_usage or {}).get("prompt_tokens", 0), + "completion_tokens": (stream_usage or {}).get("completion_tokens", 0), + "total_tokens": (stream_usage or {}).get("total_tokens", 0), + }, + } + if stream_timings: + usage_chunk["timings"] = stream_timings + yield f"data: {json.dumps(usage_chunk, separators=_SEP)}\n\n" + + yield "data: [DONE]\n\n" + + except asyncio.CancelledError: + raise + except Exception as e: + yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" + + return StreamingResponse(sse_generator(), media_type="text/event-stream", headers=_SSE_HEADERS) + + +# ── Path B: Tool calling (asyncio.to_thread) ───────────────── + + +async def _handle_tool_stream(request, payload, llama_backend, gguf_messages, image_b64, + completion_id, created, model_name): + """Handle a tool-calling streaming request via asyncio.to_thread. + + Tool execution is the bottleneck, not streaming, so the thread overhead + is acceptable here. + """ + from core.inference.tools import ALL_TOOLS + + cancel_event = threading.Event() + + p_enabled_tools = payload.get("enabled_tools") + if p_enabled_tools is not None: + tools_to_use = [t for t in ALL_TOOLS if t["function"]["name"] in p_enabled_tools] + else: + tools_to_use = ALL_TOOLS + + _sentinel = object() + + async def tool_sse(): + try: + first = {"id": completion_id, "object": "chat.completion.chunk", + "created": created, "model": model_name, + "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]} + yield f"data: {json.dumps(first, separators=_SEP)}\n\n" + + gen = llama_backend.generate_chat_completion_with_tools( + messages=gguf_messages, + tools=tools_to_use, + temperature=payload.get("temperature", 0.6), + top_p=payload.get("top_p", 0.95), + top_k=payload.get("top_k", 20), + min_p=payload.get("min_p", 0.01), + max_tokens=payload.get("max_tokens"), + repetition_penalty=payload.get("repetition_penalty", 1.1), + presence_penalty=payload.get("presence_penalty", 0.0), + cancel_event=cancel_event, + enable_thinking=payload.get("enable_thinking"), + auto_heal_tool_calls=payload.get("auto_heal_tool_calls", True), + max_tool_iterations=payload.get("max_tool_calls_per_message", 10), + tool_call_timeout=payload.get("tool_call_timeout", 300), + session_id=payload.get("session_id"), + ) + + prev_text = "" + _usage = None + _timings = None + + while True: + if await request.is_disconnected(): + cancel_event.set() + return + event = await asyncio.to_thread(next, gen, _sentinel) + if event is _sentinel: + break + if event["type"] == "status": + yield f"data: {json.dumps({'type': 'tool_status', 'content': event['text']})}\n\n" + continue + if event["type"] in ("tool_start", "tool_end"): + yield f"data: {json.dumps(event)}\n\n" + continue + if event["type"] == "metadata": + _usage = event.get("usage") + _timings = event.get("timings") + continue + cumulative = event.get("text", "") + new_text = cumulative[len(prev_text):] + prev_text = cumulative + if not new_text: + continue + chunk = {"id": completion_id, "object": "chat.completion.chunk", + "created": created, "model": model_name, + "choices": [{"index": 0, "delta": {"content": new_text}, "finish_reason": None}]} + yield f"data: {json.dumps(chunk, separators=_SEP)}\n\n" + + final = {"id": completion_id, "object": "chat.completion.chunk", + "created": created, "model": model_name, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]} + yield f"data: {json.dumps(final, separators=_SEP)}\n\n" + + if _usage or _timings: + uc = {"id": completion_id, "object": "chat.completion.chunk", + "created": created, "model": model_name, "choices": [], + "usage": {"prompt_tokens": (_usage or {}).get("prompt_tokens", 0), + "completion_tokens": (_usage or {}).get("completion_tokens", 0), + "total_tokens": (_usage or {}).get("total_tokens", 0)}} + if _timings: + uc["timings"] = _timings + yield f"data: {json.dumps(uc, separators=_SEP)}\n\n" + + yield "data: [DONE]\n\n" + except asyncio.CancelledError: + cancel_event.set() + raise + except Exception as e: + yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" + + return StreamingResponse(tool_sse(), media_type="text/event-stream", headers=_SSE_HEADERS) + + +# ── Path C: Non-streaming ───────────────────────────────────── + + +async def _handle_non_streaming(payload, llama_backend, gguf_messages, image_b64, + completion_id, created, model_name): + """Handle a non-streaming request. Returns a JSON response.""" + cancel_event = threading.Event() + + def _run_sync(): + gen = llama_backend.generate_chat_completion( + messages=gguf_messages, + image_b64=image_b64, + temperature=payload.get("temperature", 0.6), + top_p=payload.get("top_p", 0.95), + top_k=payload.get("top_k", 20), + min_p=payload.get("min_p", 0.01), + max_tokens=payload.get("max_tokens"), + repetition_penalty=payload.get("repetition_penalty", 1.0), + presence_penalty=payload.get("presence_penalty", 0.0), + stop=payload.get("stop"), + cancel_event=cancel_event, + enable_thinking=payload.get("enable_thinking"), + ) + text = "" + usage = None + timings = None + for item in gen: + if isinstance(item, dict) and item.get("type") == "metadata": + usage = item.get("usage") + timings = item.get("timings") + elif isinstance(item, str): + text = item + return text, usage, timings + + text, usage, timings = await asyncio.to_thread(_run_sync) + + result = { + "id": completion_id, + "object": "chat.completion", + "created": created, + "model": model_name, + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": text}, + "finish_reason": "stop", + }], + "usage": { + "prompt_tokens": (usage or {}).get("prompt_tokens", 0), + "completion_tokens": (usage or {}).get("completion_tokens", 0), + "total_tokens": (usage or {}).get("total_tokens", 0), + }, + } + if timings: + result["timings"] = timings + + return JSONResponse( + content=result, + headers={"Access-Control-Allow-Origin": "*"}, + ) + + +# ── Main endpoint ───────────────────────────────────────────── + + +@stream_app.post("/stream") +async def stream_endpoint(request: Request): + """ + Stream chat completions with minimal overhead. + + Three paths: + - Path A (hot): async httpx streaming direct to llama-server + - Path B: tool calling via asyncio.to_thread (tool exec is the bottleneck) + - Path C: non-streaming one-shot JSON response + """ + payload, llama_backend, gguf_messages, image_b64, completion_id, created, model_name = \ + await _validate_request(request) + + # Path C: Non-streaming + stream = payload.get("stream", True) + if not stream: + return await _handle_non_streaming( + payload, llama_backend, gguf_messages, image_b64, + completion_id, created, model_name, + ) + + # Path B: Tool calling + use_tools = payload.get("use_tools", False) + if use_tools and llama_backend.supports_tools: + return await _handle_tool_stream( + request, payload, llama_backend, gguf_messages, image_b64, + completion_id, created, model_name, + ) + + # Path A: Direct async streaming (hot path) + return await _handle_async_stream( + request, payload, llama_backend, gguf_messages, image_b64, + completion_id, created, model_name, + ) + + +# ── Server lifecycle ────────────────────────────────────────── + + +def start_streaming_server(port: int) -> None: + """Start the streaming server in the current thread (blocking). Use in a daemon thread.""" + import uvicorn + + uvicorn.run( + stream_app, + host="127.0.0.1", + port=port, + log_level="warning", + access_log=False, + ) + + +def find_free_port() -> int: + """Find a free TCP port.""" + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 57cbcccf66..662003ddef 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -177,12 +177,46 @@ export async function* streamChatCompletions( payload: OpenAIChatCompletionsRequest, signal: AbortSignal, ): AsyncGenerator { - const response = await authFetch("/v1/chat/completions", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(payload), - signal, - }); + // Try to acquire a fast-path token for the dedicated streaming server. + // Falls back silently to the baseline endpoint on any failure. + let useFastPath = false; + let fastUrl = ""; + let fastToken = ""; + + try { + const streamUrlResp = await authFetch("/api/inference/stream-url", { signal }); + if (streamUrlResp.ok) { + const info = (await streamUrlResp.json()) as { + supported?: boolean; + stream_url?: string; + token?: string; + }; + if (info.supported && info.stream_url && info.token) { + fastUrl = info.stream_url; + fastToken = info.token; + useFastPath = true; + } + } + } catch { + /* fall back silently to baseline */ + } + + const response = useFastPath + ? await fetch(fastUrl, { + method: "POST", + headers: { + "Content-Type": "application/json", + "X-Stream-Token": fastToken, + }, + body: JSON.stringify(payload), + signal, + }) + : await authFetch("/v1/chat/completions", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + signal, + }); if (!response.ok) { const body = await response.json().catch(() => null); From 0f0e02603bdac0f7cf718dfadf48d813e57608b4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 27 Mar 2026 02:04:08 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 43 ++- studio/backend/main.py | 11 +- studio/backend/routes/inference.py | 26 +- studio/backend/streaming_server.py | 335 ++++++++++++++------- 4 files changed, 277 insertions(+), 138 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 5c68981096..c4e1bef3fe 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1486,7 +1486,10 @@ class LlamaCppBackend: pool = 10, ) with client.stream( - "POST", url, json = payload, timeout = prefill_timeout, + "POST", + url, + json = payload, + timeout = prefill_timeout, headers = headers, ) as response: _response_ref[0] = response @@ -1561,10 +1564,16 @@ class LlamaCppBackend: # can finish. Cancel during streaming is handled by the # watcher thread (closes the response on cancel_event). stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) with httpx.Client(timeout = stream_timeout) as client: with self._stream_with_retry( - client, url, payload, cancel_event, headers = _auth_headers, + client, + url, + payload, + cancel_event, + headers = _auth_headers, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -1721,7 +1730,11 @@ class LlamaCppBackend: payload["stop"] = stop try: - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} + if self._api_key + else None + ) with httpx.Client(timeout = None) as client: resp = client.post(url, json = payload, headers = _auth_headers) if resp.status_code != 200: @@ -1966,10 +1979,16 @@ class LlamaCppBackend: try: stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + ) with httpx.Client(timeout = stream_timeout) as client: with self._stream_with_retry( - client, url, stream_payload, cancel_event, headers = _auth_headers, + client, + url, + stream_payload, + cancel_event, + headers = _auth_headers, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -2095,7 +2114,9 @@ class LlamaCppBackend: if not self.is_loaded: return None try: - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + ) with httpx.Client(timeout = 10, headers = _auth_headers) as client: def _detok(tid: int) -> str: @@ -2214,8 +2235,12 @@ class LlamaCppBackend: if need_ids: payload["n_probs"] = 1 - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} - with httpx.Client(timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers) as client: + _auth_headers = ( + {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + ) + with httpx.Client( + timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers + ) as client: resp = client.post(f"{self.base_url}/completion", json = payload) if resp.status_code != 200: raise RuntimeError( diff --git a/studio/backend/main.py b/studio/backend/main.py index 83373b1def..b6c509d727 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -115,7 +115,10 @@ async def lifespan(app: FastAPI): # Start dedicated streaming server in a daemon thread (Option B). # Accepts both UNSLOTH_FAST_SSE=1 and UNSLOTH_STREAM_SERVER=1. - if os.getenv("UNSLOTH_FAST_SSE", "0") == "1" or os.getenv("UNSLOTH_STREAM_SERVER", "0") == "1": + if ( + os.getenv("UNSLOTH_FAST_SSE", "0") == "1" + or os.getenv("UNSLOTH_STREAM_SERVER", "0") == "1" + ): import threading as _threading from streaming_server import start_streaming_server, find_free_port @@ -123,9 +126,9 @@ async def lifespan(app: FastAPI): stream_port = find_free_port() app.state.stream_port = stream_port _stream_thread = _threading.Thread( - target=start_streaming_server, - args=(stream_port,), - daemon=True, + target = start_streaming_server, + args = (stream_port,), + daemon = True, ) _stream_thread.start() print(f"[streaming_server] Started on 127.0.0.1:{stream_port}") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index f7cc0db78f..284c32a97d 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1633,25 +1633,27 @@ async def consume_stream_token_endpoint(request: Request): body = await request.json() token = body.get("token") if not token: - return JSONResponse({"valid": False}, status_code=400) + return JSONResponse({"valid": False}, status_code = 400) from stream_token_store import consume_stream_token username = consume_stream_token(token) if not username: - return JSONResponse({"valid": False}, status_code=401) + return JSONResponse({"valid": False}, status_code = 401) llama = get_llama_cpp_backend() - return JSONResponse({ - "valid": True, - "username": username, - "llama_port": llama._port if llama.is_loaded else None, - "llama_api_key": llama._api_key, - "model_name": llama.model_identifier or "unknown", - "supports_reasoning": llama.supports_reasoning, - "supports_tools": llama.supports_tools, - "is_vision": llama.is_vision, - }) + return JSONResponse( + { + "valid": True, + "username": username, + "llama_port": llama._port if llama.is_loaded else None, + "llama_api_key": llama._api_key, + "model_name": llama.model_identifier or "unknown", + "supports_reasoning": llama.supports_reasoning, + "supports_tools": llama.supports_tools, + "is_vision": llama.is_vision, + } + ) # ===================================================================== diff --git a/studio/backend/streaming_server.py b/studio/backend/streaming_server.py index 9b5f667eff..1a2b40e597 100644 --- a/studio/backend/streaming_server.py +++ b/studio/backend/streaming_server.py @@ -35,7 +35,7 @@ from fastapi.responses import JSONResponse, StreamingResponse from stream_token_store import consume_stream_token -stream_app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None) +stream_app = FastAPI(docs_url = None, redoc_url = None, openapi_url = None) # ── Shared helpers ──────────────────────────────────────────── @@ -92,17 +92,17 @@ def _process_image(image_b64, llama_backend): if img.mode == "RGBA": img = img.convert("RGB") buf = _BytesIO() - img.save(buf, format="PNG") + img.save(buf, format = "PNG") return _b64.b64encode(buf.getvalue()).decode("ascii") elif image_b64 and not llama_backend.is_vision: raise HTTPException( - status_code=400, - detail="Image provided but current GGUF model does not support vision.", + status_code = 400, + detail = "Image provided but current GGUF model does not support vision.", ) return image_b64 -def _build_llama_payload(llama_backend, openai_messages, payload, stream=True): +def _build_llama_payload(llama_backend, openai_messages, payload, stream = True): """Build the payload for llama-server /v1/chat/completions. Only sends repeat_penalty when the client explicitly provides it. @@ -125,7 +125,9 @@ def _build_llama_payload(llama_backend, openai_messages, payload, stream=True): if stream: llama_payload["stream_options"] = {"include_usage": True} if llama_backend.supports_reasoning and payload.get("enable_thinking") is not None: - llama_payload["chat_template_kwargs"] = {"enable_thinking": payload["enable_thinking"]} + llama_payload["chat_template_kwargs"] = { + "enable_thinking": payload["enable_thinking"] + } if payload.get("max_tokens") is not None: llama_payload["max_tokens"] = payload["max_tokens"] if payload.get("stop"): @@ -139,12 +141,15 @@ def _build_llama_payload(llama_backend, openai_messages, payload, stream=True): @stream_app.options("/stream") async def stream_preflight(): """Handle CORS preflight for the /stream endpoint.""" - return JSONResponse(content={}, headers={ - "Access-Control-Allow-Origin": "*", - "Access-Control-Allow-Methods": "POST, OPTIONS", - "Access-Control-Allow-Headers": "Content-Type, X-Stream-Token", - "Access-Control-Max-Age": "86400", - }) + return JSONResponse( + content = {}, + headers = { + "Access-Control-Allow-Origin": "*", + "Access-Control-Allow-Methods": "POST, OPTIONS", + "Access-Control-Allow-Headers": "Content-Type, X-Stream-Token", + "Access-Control-Max-Age": "86400", + }, + ) # ── Request validation ──────────────────────────────────────── @@ -154,21 +159,22 @@ async def _validate_request(request: Request): """Validate token, parse body, get backend. Returns all needed context.""" token = request.headers.get("X-Stream-Token") if not token: - raise HTTPException(status_code=401, detail="Missing X-Stream-Token header") + raise HTTPException(status_code = 401, detail = "Missing X-Stream-Token header") username = consume_stream_token(token) if username is None: - raise HTTPException(status_code=401, detail="Invalid or expired stream token") + raise HTTPException(status_code = 401, detail = "Invalid or expired stream token") body_bytes = await request.body() try: payload = json.loads(body_bytes) except (json.JSONDecodeError, UnicodeDecodeError): - raise HTTPException(status_code=400, detail="Invalid JSON body") + raise HTTPException(status_code = 400, detail = "Invalid JSON body") from routes.inference import get_llama_cpp_backend + llama_backend = get_llama_cpp_backend() if not llama_backend.is_loaded: - raise HTTPException(status_code=400, detail="No GGUF model loaded") + raise HTTPException(status_code = 400, detail = "No GGUF model loaded") messages = payload.get("messages", []) gguf_messages, image_b64 = _extract_content_parts(messages) @@ -184,7 +190,15 @@ async def _validate_request(request: Request): created = int(time.time()) model_name = llama_backend.model_identifier or "unknown" - return payload, llama_backend, gguf_messages, image_b64, completion_id, created, model_name + return ( + payload, + llama_backend, + gguf_messages, + image_b64, + completion_id, + created, + model_name, + ) # ── Path A: Direct async streaming (HOT PATH) ──────────────── @@ -200,8 +214,16 @@ _SSE_HEADERS = { _SEP = (",", ":") -async def _handle_async_stream(request, payload, llama_backend, gguf_messages, image_b64, - completion_id, created, model_name): +async def _handle_async_stream( + request, + payload, + llama_backend, + gguf_messages, + image_b64, + completion_id, + created, + model_name, +): """ Stream directly from llama-server using httpx.AsyncClient. @@ -210,24 +232,34 @@ async def _handle_async_stream(request, payload, llama_backend, gguf_messages, i with delta tokens natively, so no cumulative-to-delta conversion needed. """ openai_messages = llama_backend._build_openai_messages(gguf_messages, image_b64) - llama_payload = _build_llama_payload(llama_backend, openai_messages, payload, stream=True) + llama_payload = _build_llama_payload( + llama_backend, openai_messages, payload, stream = True + ) port = llama_backend._port api_key = llama_backend._api_key headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} url = f"http://127.0.0.1:{port}/v1/chat/completions" - timeout = httpx.Timeout(connect=30, read=120.0, write=10, pool=10) + timeout = httpx.Timeout(connect = 30, read = 120.0, write = 10, pool = 10) async def sse_generator(): try: # Role chunk - role = {"id": completion_id, "object": "chat.completion.chunk", - "created": created, "model": model_name, - "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]} - yield f"data: {json.dumps(role, separators=_SEP)}\n\n" + role = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": model_name, + "choices": [ + {"index": 0, "delta": {"role": "assistant"}, "finish_reason": None} + ], + } + yield f"data: {json.dumps(role, separators = _SEP)}\n\n" - async with httpx.AsyncClient(timeout=timeout) as client: - async with client.stream("POST", url, json=llama_payload, headers=headers) as resp: + async with httpx.AsyncClient(timeout = timeout) as client: + async with client.stream( + "POST", url, json = llama_payload, headers = headers + ) as resp: if resp.status_code != 200: error_body = await resp.aread() raise RuntimeError( @@ -253,9 +285,9 @@ async def _handle_async_stream(request, payload, llama_backend, gguf_messages, i if line == "data: [DONE]": if in_thinking: if has_content_tokens: - yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': ''}, 'finish_reason': None}]}, separators=_SEP)}\n\n" + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': ''}, 'finish_reason': None}]}, separators = _SEP)}\n\n" else: - yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning_text}, 'finish_reason': None}]}, separators=_SEP)}\n\n" + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning_text}, 'finish_reason': None}]}, separators = _SEP)}\n\n" stream_done = True break if not line.startswith("data: "): @@ -284,8 +316,8 @@ async def _handle_async_stream(request, payload, llama_backend, gguf_messages, i reasoning_text += reasoning if not in_thinking: in_thinking = True - yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': ''}, 'finish_reason': None}]}, separators=_SEP)}\n\n" - yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning}, 'finish_reason': None}]}, separators=_SEP)}\n\n" + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': ''}, 'finish_reason': None}]}, separators = _SEP)}\n\n" + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning}, 'finish_reason': None}]}, separators = _SEP)}\n\n" # Handle content tokens token = delta.get("content", "") @@ -293,32 +325,41 @@ async def _handle_async_stream(request, payload, llama_backend, gguf_messages, i has_content_tokens = True if in_thinking: in_thinking = False - yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': ''}, 'finish_reason': None}]}, separators=_SEP)}\n\n" - yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': token}, 'finish_reason': None}]}, separators=_SEP)}\n\n" + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': ''}, 'finish_reason': None}]}, separators = _SEP)}\n\n" + yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': token}, 'finish_reason': None}]}, separators = _SEP)}\n\n" if stream_done: break # Final stop chunk - final = {"id": completion_id, "object": "chat.completion.chunk", - "created": created, "model": model_name, - "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]} - yield f"data: {json.dumps(final, separators=_SEP)}\n\n" + final = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": model_name, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + } + yield f"data: {json.dumps(final, separators = _SEP)}\n\n" # Usage chunk if stream_usage or stream_timings: usage_chunk = { - "id": completion_id, "object": "chat.completion.chunk", - "created": created, "model": model_name, "choices": [], + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": model_name, + "choices": [], "usage": { "prompt_tokens": (stream_usage or {}).get("prompt_tokens", 0), - "completion_tokens": (stream_usage or {}).get("completion_tokens", 0), + "completion_tokens": (stream_usage or {}).get( + "completion_tokens", 0 + ), "total_tokens": (stream_usage or {}).get("total_tokens", 0), }, } if stream_timings: usage_chunk["timings"] = stream_timings - yield f"data: {json.dumps(usage_chunk, separators=_SEP)}\n\n" + yield f"data: {json.dumps(usage_chunk, separators = _SEP)}\n\n" yield "data: [DONE]\n\n" @@ -327,14 +368,24 @@ async def _handle_async_stream(request, payload, llama_backend, gguf_messages, i except Exception as e: yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" - return StreamingResponse(sse_generator(), media_type="text/event-stream", headers=_SSE_HEADERS) + return StreamingResponse( + sse_generator(), media_type = "text/event-stream", headers = _SSE_HEADERS + ) # ── Path B: Tool calling (asyncio.to_thread) ───────────────── -async def _handle_tool_stream(request, payload, llama_backend, gguf_messages, image_b64, - completion_id, created, model_name): +async def _handle_tool_stream( + request, + payload, + llama_backend, + gguf_messages, + image_b64, + completion_id, + created, + model_name, +): """Handle a tool-calling streaming request via asyncio.to_thread. Tool execution is the bottleneck, not streaming, so the thread overhead @@ -346,7 +397,9 @@ async def _handle_tool_stream(request, payload, llama_backend, gguf_messages, im p_enabled_tools = payload.get("enabled_tools") if p_enabled_tools is not None: - tools_to_use = [t for t in ALL_TOOLS if t["function"]["name"] in p_enabled_tools] + tools_to_use = [ + t for t in ALL_TOOLS if t["function"]["name"] in p_enabled_tools + ] else: tools_to_use = ALL_TOOLS @@ -354,27 +407,33 @@ async def _handle_tool_stream(request, payload, llama_backend, gguf_messages, im async def tool_sse(): try: - first = {"id": completion_id, "object": "chat.completion.chunk", - "created": created, "model": model_name, - "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]} - yield f"data: {json.dumps(first, separators=_SEP)}\n\n" + first = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": model_name, + "choices": [ + {"index": 0, "delta": {"role": "assistant"}, "finish_reason": None} + ], + } + yield f"data: {json.dumps(first, separators = _SEP)}\n\n" gen = llama_backend.generate_chat_completion_with_tools( - messages=gguf_messages, - tools=tools_to_use, - temperature=payload.get("temperature", 0.6), - top_p=payload.get("top_p", 0.95), - top_k=payload.get("top_k", 20), - min_p=payload.get("min_p", 0.01), - max_tokens=payload.get("max_tokens"), - repetition_penalty=payload.get("repetition_penalty", 1.1), - presence_penalty=payload.get("presence_penalty", 0.0), - cancel_event=cancel_event, - enable_thinking=payload.get("enable_thinking"), - auto_heal_tool_calls=payload.get("auto_heal_tool_calls", True), - max_tool_iterations=payload.get("max_tool_calls_per_message", 10), - tool_call_timeout=payload.get("tool_call_timeout", 300), - session_id=payload.get("session_id"), + messages = gguf_messages, + tools = tools_to_use, + temperature = payload.get("temperature", 0.6), + top_p = payload.get("top_p", 0.95), + top_k = payload.get("top_k", 20), + min_p = payload.get("min_p", 0.01), + max_tokens = payload.get("max_tokens"), + repetition_penalty = payload.get("repetition_penalty", 1.1), + presence_penalty = payload.get("presence_penalty", 0.0), + cancel_event = cancel_event, + enable_thinking = payload.get("enable_thinking"), + auto_heal_tool_calls = payload.get("auto_heal_tool_calls", True), + max_tool_iterations = payload.get("max_tool_calls_per_message", 10), + tool_call_timeout = payload.get("tool_call_timeout", 300), + session_id = payload.get("session_id"), ) prev_text = "" @@ -399,29 +458,50 @@ async def _handle_tool_stream(request, payload, llama_backend, gguf_messages, im _timings = event.get("timings") continue cumulative = event.get("text", "") - new_text = cumulative[len(prev_text):] + new_text = cumulative[len(prev_text) :] prev_text = cumulative if not new_text: continue - chunk = {"id": completion_id, "object": "chat.completion.chunk", - "created": created, "model": model_name, - "choices": [{"index": 0, "delta": {"content": new_text}, "finish_reason": None}]} - yield f"data: {json.dumps(chunk, separators=_SEP)}\n\n" + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": model_name, + "choices": [ + { + "index": 0, + "delta": {"content": new_text}, + "finish_reason": None, + } + ], + } + yield f"data: {json.dumps(chunk, separators = _SEP)}\n\n" - final = {"id": completion_id, "object": "chat.completion.chunk", - "created": created, "model": model_name, - "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]} - yield f"data: {json.dumps(final, separators=_SEP)}\n\n" + final = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": model_name, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + } + yield f"data: {json.dumps(final, separators = _SEP)}\n\n" if _usage or _timings: - uc = {"id": completion_id, "object": "chat.completion.chunk", - "created": created, "model": model_name, "choices": [], - "usage": {"prompt_tokens": (_usage or {}).get("prompt_tokens", 0), - "completion_tokens": (_usage or {}).get("completion_tokens", 0), - "total_tokens": (_usage or {}).get("total_tokens", 0)}} + uc = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": model_name, + "choices": [], + "usage": { + "prompt_tokens": (_usage or {}).get("prompt_tokens", 0), + "completion_tokens": (_usage or {}).get("completion_tokens", 0), + "total_tokens": (_usage or {}).get("total_tokens", 0), + }, + } if _timings: uc["timings"] = _timings - yield f"data: {json.dumps(uc, separators=_SEP)}\n\n" + yield f"data: {json.dumps(uc, separators = _SEP)}\n\n" yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -430,31 +510,34 @@ async def _handle_tool_stream(request, payload, llama_backend, gguf_messages, im except Exception as e: yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" - return StreamingResponse(tool_sse(), media_type="text/event-stream", headers=_SSE_HEADERS) + return StreamingResponse( + tool_sse(), media_type = "text/event-stream", headers = _SSE_HEADERS + ) # ── Path C: Non-streaming ───────────────────────────────────── -async def _handle_non_streaming(payload, llama_backend, gguf_messages, image_b64, - completion_id, created, model_name): +async def _handle_non_streaming( + payload, llama_backend, gguf_messages, image_b64, completion_id, created, model_name +): """Handle a non-streaming request. Returns a JSON response.""" cancel_event = threading.Event() def _run_sync(): gen = llama_backend.generate_chat_completion( - messages=gguf_messages, - image_b64=image_b64, - temperature=payload.get("temperature", 0.6), - top_p=payload.get("top_p", 0.95), - top_k=payload.get("top_k", 20), - min_p=payload.get("min_p", 0.01), - max_tokens=payload.get("max_tokens"), - repetition_penalty=payload.get("repetition_penalty", 1.0), - presence_penalty=payload.get("presence_penalty", 0.0), - stop=payload.get("stop"), - cancel_event=cancel_event, - enable_thinking=payload.get("enable_thinking"), + messages = gguf_messages, + image_b64 = image_b64, + temperature = payload.get("temperature", 0.6), + top_p = payload.get("top_p", 0.95), + top_k = payload.get("top_k", 20), + min_p = payload.get("min_p", 0.01), + max_tokens = payload.get("max_tokens"), + repetition_penalty = payload.get("repetition_penalty", 1.0), + presence_penalty = payload.get("presence_penalty", 0.0), + stop = payload.get("stop"), + cancel_event = cancel_event, + enable_thinking = payload.get("enable_thinking"), ) text = "" usage = None @@ -474,11 +557,13 @@ async def _handle_non_streaming(payload, llama_backend, gguf_messages, image_b64 "object": "chat.completion", "created": created, "model": model_name, - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": text}, - "finish_reason": "stop", - }], + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": text}, + "finish_reason": "stop", + } + ], "usage": { "prompt_tokens": (usage or {}).get("prompt_tokens", 0), "completion_tokens": (usage or {}).get("completion_tokens", 0), @@ -489,8 +574,8 @@ async def _handle_non_streaming(payload, llama_backend, gguf_messages, image_b64 result["timings"] = timings return JSONResponse( - content=result, - headers={"Access-Control-Allow-Origin": "*"}, + content = result, + headers = {"Access-Control-Allow-Origin": "*"}, ) @@ -507,29 +592,53 @@ async def stream_endpoint(request: Request): - Path B: tool calling via asyncio.to_thread (tool exec is the bottleneck) - Path C: non-streaming one-shot JSON response """ - payload, llama_backend, gguf_messages, image_b64, completion_id, created, model_name = \ - await _validate_request(request) + ( + payload, + llama_backend, + gguf_messages, + image_b64, + completion_id, + created, + model_name, + ) = await _validate_request(request) # Path C: Non-streaming stream = payload.get("stream", True) if not stream: return await _handle_non_streaming( - payload, llama_backend, gguf_messages, image_b64, - completion_id, created, model_name, + payload, + llama_backend, + gguf_messages, + image_b64, + completion_id, + created, + model_name, ) # Path B: Tool calling use_tools = payload.get("use_tools", False) if use_tools and llama_backend.supports_tools: return await _handle_tool_stream( - request, payload, llama_backend, gguf_messages, image_b64, - completion_id, created, model_name, + request, + payload, + llama_backend, + gguf_messages, + image_b64, + completion_id, + created, + model_name, ) # Path A: Direct async streaming (hot path) return await _handle_async_stream( - request, payload, llama_backend, gguf_messages, image_b64, - completion_id, created, model_name, + request, + payload, + llama_backend, + gguf_messages, + image_b64, + completion_id, + created, + model_name, ) @@ -542,10 +651,10 @@ def start_streaming_server(port: int) -> None: uvicorn.run( stream_app, - host="127.0.0.1", - port=port, - log_level="warning", - access_log=False, + host = "127.0.0.1", + port = port, + log_level = "warning", + access_log = False, )