e945f43652
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 -> <think> 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.
65 lines
2 KiB
Python
65 lines
2 KiB
Python
# 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)
|