diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index b39f915764..da59ba9a1a 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -10,10 +10,12 @@ from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer import jwt from .storage import ( + API_KEY_PREFIX, get_jwt_secret, get_user_and_secret, load_jwt_secret, save_refresh_token, + validate_api_key, verify_refresh_token, ) @@ -137,6 +139,18 @@ async def _get_current_subject( ... """ token = credentials.credentials + + # --- API key path (sk-unsloth-...) --- + if token.startswith(API_KEY_PREFIX): + username = validate_api_key(token) + if username is None: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Invalid or expired API key", + ) + return username + + # --- JWT path --- subject = _decode_subject_without_verification(token) if subject is None: raise HTTPException( diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 1395574cce..7d55a2dc59 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -72,7 +72,22 @@ def clear_bootstrap_password() -> None: def _hash_token(token: str) -> str: - """SHA-256 hash helper used for refresh token storage.""" + """SHA-256 hash helper used for refresh token storage. + + Plain SHA-256 is intentional here: refresh tokens are high-entropy + random strings from ``secrets.token_urlsafe(48)`` (384 bits of + entropy), so a slow KDF (Argon2 / bcrypt / PBKDF2) provides zero + additional security — no attacker can brute-force 2^384 regardless + of hash speed — while adding tens of ms of CPU to every refresh. + See the OWASP Password Storage Cheat Sheet on fast-vs-slow hashing + of high-entropy inputs. + + API keys use the separate ``_pbkdf2_api_key`` helper below, which + runs PBKDF2-HMAC-SHA256 with a persistent server-side salt — not + for cryptographic reasons (128-bit random tokens don't need slow + hashing), but because CodeQL's ``py/weak-sensitive-data-hashing`` + query mislabels API keys as passwords and demands a KDF. + """ return hashlib.sha256(token.encode("utf-8")).hexdigest() @@ -103,6 +118,29 @@ def get_connection() -> sqlite3.Connection: ); """ ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS api_keys ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + username TEXT NOT NULL, + key_prefix TEXT NOT NULL, + key_hash TEXT NOT NULL UNIQUE, + name TEXT NOT NULL DEFAULT '', + created_at TEXT NOT NULL, + last_used_at TEXT, + expires_at TEXT, + is_active INTEGER NOT NULL DEFAULT 1 + ); + """ + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS app_secrets ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL + ); + """ + ) columns = {row["name"] for row in conn.execute("PRAGMA table_info(auth_user)")} if "must_change_password" not in columns: conn.execute( @@ -112,6 +150,89 @@ def get_connection() -> sqlite3.Connection: return conn +# ── API-key PBKDF2 salt ──────────────────────────────────────────────── +# +# Module-level cache for the persistent API-key PBKDF2 salt. Populated +# lazily on first use via ``_get_or_create_api_key_pbkdf2_salt``. Not +# protected by a lock because (a) the ``INSERT OR IGNORE`` provides +# atomicity at the SQLite layer and (b) concurrent populations converge +# on the same value, so the worst case is a harmless duplicate read on +# startup. +_api_key_pbkdf2_salt_cache: Optional[bytes] = None + + +def _get_or_create_api_key_pbkdf2_salt() -> bytes: + """Return the persistent API-key PBKDF2 salt, generating it once if missing. + + Stored as a hex-encoded 32-byte random value in the ``app_secrets`` + table under key ``"api_key_pbkdf2_salt"``. Regenerated only if the row + is missing (i.e. fresh install, or operator manually deleted the row + and accepts invalidating existing API keys). + """ + global _api_key_pbkdf2_salt_cache + if _api_key_pbkdf2_salt_cache is not None: + return _api_key_pbkdf2_salt_cache + + conn = get_connection() + try: + cur = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + ("api_key_pbkdf2_salt",), + ) + row = cur.fetchone() + if row is None: + new_value = secrets.token_hex(32) # 32 bytes -> 64 hex chars + conn.execute( + "INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)", + ("api_key_pbkdf2_salt", new_value), + ) + conn.commit() + cur = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + ("api_key_pbkdf2_salt",), + ) + row = cur.fetchone() + salt = bytes.fromhex(row["value"]) + finally: + conn.close() + + _api_key_pbkdf2_salt_cache = salt + return salt + + +_API_KEY_PBKDF2_ITERATIONS = 100_000 + + +def _pbkdf2_api_key(raw_key: str) -> str: + """PBKDF2-HMAC-SHA256 an API key with a persistent server-side salt. + + Used for API-key storage ONLY, not refresh tokens. Matches the + PBKDF2 algorithm + iteration count used by the password hasher in + ``auth/hashing.py`` so the codebase is consistent on which KDF it + uses for credential storage. + + Notes on why a slow KDF here is *only* a CodeQL appeasement and + *not* a cryptographic requirement: API keys are cryptographically + random 128-bit tokens (via ``secrets.token_hex``), so brute force + against 2^128 is infeasible regardless of hash speed. CodeQL's + ``py/weak-sensitive-data-hashing`` query mislabels these tokens as + "password" sensitive data and then demands a KDF from its + allowlist (Argon2 / scrypt / bcrypt / PBKDF2). Per the query's + own recommendation page we use PBKDF2. The persistent salt is + still loaded from ``app_secrets`` so an attacker dumping the + ``api_keys`` table alone cannot derive hashes for candidate + tokens without also obtaining the salt row. + """ + salt = _get_or_create_api_key_pbkdf2_salt() + dk = hashlib.pbkdf2_hmac( + "sha256", + raw_key.encode("utf-8"), + salt, + _API_KEY_PBKDF2_ITERATIONS, + ) + return dk.hex() + + def is_initialized() -> bool: """Check if auth is ready for login (at least one user exists in DB).""" conn = get_connection() @@ -357,3 +478,105 @@ def revoke_user_refresh_tokens(username: str) -> None: conn.commit() finally: conn.close() + + +# --------------------------------------------------------------------------- +# API key management +# --------------------------------------------------------------------------- + +API_KEY_PREFIX = "sk-unsloth-" + + +def create_api_key( + username: str, + name: str, + expires_at: Optional[str] = None, +) -> Tuple[str, dict]: + """Create a new API key for *username*. + + Returns ``(raw_key, row_dict)`` where *raw_key* is shown to the user + exactly once. The database only stores the SHA-256 hash. + """ + raw_key = API_KEY_PREFIX + secrets.token_hex(16) + key_hash = _pbkdf2_api_key(raw_key) + key_prefix = raw_key[len(API_KEY_PREFIX) : len(API_KEY_PREFIX) + 8] + now = datetime.now(timezone.utc).isoformat() + + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at) + VALUES (?, ?, ?, ?, ?, ?) + """, + (username, key_prefix, key_hash, name, now, expires_at), + ) + conn.commit() + cur = conn.execute("SELECT * FROM api_keys WHERE key_hash = ?", (key_hash,)) + row = cur.fetchone() + return raw_key, dict(row) + finally: + conn.close() + + +def list_api_keys(username: str) -> list: + """Return all API keys for *username* (never exposes ``key_hash``).""" + conn = get_connection() + try: + cur = conn.execute( + """ + SELECT id, username, key_prefix, name, created_at, last_used_at, expires_at, is_active + FROM api_keys + WHERE username = ? + ORDER BY created_at DESC + """, + (username,), + ) + return [dict(row) for row in cur.fetchall()] + finally: + conn.close() + + +def revoke_api_key(username: str, key_id: int) -> bool: + """Soft-delete an API key. Returns True if a matching row was found.""" + conn = get_connection() + try: + cursor = conn.execute( + "UPDATE api_keys SET is_active = 0 WHERE id = ? AND username = ?", + (key_id, username), + ) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + +def validate_api_key(raw_key: str) -> Optional[str]: + """Validate *raw_key* and return the owning username, or ``None``. + + Also updates ``last_used_at`` on success. + """ + key_hash = _pbkdf2_api_key(raw_key) + conn = get_connection() + try: + cur = conn.execute( + "SELECT id, username, is_active, expires_at FROM api_keys WHERE key_hash = ?", + (key_hash,), + ) + row = cur.fetchone() + if row is None: + return None + if not row["is_active"]: + return None + if row["expires_at"] is not None: + expires = datetime.fromisoformat(row["expires_at"]) + if datetime.now(timezone.utc) > expires: + return None + conn.execute( + "UPDATE api_keys SET last_used_at = ? WHERE id = ?", + (datetime.now(timezone.utc).isoformat(), row["id"]), + ) + conn.commit() + return row["username"] + finally: + conn.close() diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py new file mode 100644 index 0000000000..e7b40a60ce --- /dev/null +++ b/studio/backend/core/inference/anthropic_compat.py @@ -0,0 +1,488 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +""" +Anthropic Messages API ↔ OpenAI format translation utilities. + +Pure functions and a stateful stream emitter — no FastAPI, no I/O. +""" + +from __future__ import annotations + +import json +from typing import Any, Optional, Union + + +def anthropic_messages_to_openai( + messages: list[dict], + system: Optional[Union[str, list]] = None, +) -> list[dict]: + """Convert Anthropic messages + system to OpenAI-format message dicts.""" + result: list[dict] = [] + + # System prompt + if system: + if isinstance(system, str): + result.append({"role": "system", "content": system}) + elif isinstance(system, list): + parts = [] + for block in system: + if isinstance(block, dict) and block.get("type") == "text": + parts.append(block["text"]) + elif isinstance(block, str): + parts.append(block) + if parts: + result.append({"role": "system", "content": "\n".join(parts)}) + + for msg in messages: + role = msg["role"] if isinstance(msg, dict) else msg.role + content = msg["content"] if isinstance(msg, dict) else msg.content + + if isinstance(content, str): + result.append({"role": role, "content": content}) + continue + + # Content is a list of blocks + text_parts: list[str] = [] + tool_calls: list[dict] = [] + tool_results: list[dict] = [] + + for block in content: + b = block if isinstance(block, dict) else block.model_dump() + btype = b.get("type", "") + + if btype == "text": + text_parts.append(b["text"]) + elif btype == "tool_use": + tool_calls.append( + { + "id": b["id"], + "type": "function", + "function": { + "name": b["name"], + "arguments": json.dumps(b["input"]), + }, + } + ) + elif btype == "tool_result": + tc = b.get("content", "") + if isinstance(tc, list): + tc = " ".join( + p["text"] + for p in tc + if isinstance(p, dict) and p.get("type") == "text" + ) + tool_results.append( + { + "role": "tool", + "tool_call_id": b["tool_use_id"], + "content": str(tc), + } + ) + + if role == "assistant": + msg_dict: dict[str, Any] = {"role": "assistant"} + if text_parts: + msg_dict["content"] = "\n".join(text_parts) + if tool_calls: + msg_dict["tool_calls"] = tool_calls + result.append(msg_dict) + elif role == "user": + if text_parts: + result.append({"role": "user", "content": "\n".join(text_parts)}) + for tr in tool_results: + result.append(tr) + + return result + + +def anthropic_tools_to_openai(tools: list) -> list[dict]: + """Convert Anthropic tool definitions to OpenAI function-tool format.""" + result = [] + for t in tools: + td = t if isinstance(t, dict) else t.model_dump() + result.append( + { + "type": "function", + "function": { + "name": td["name"], + "description": td.get("description", ""), + "parameters": td.get("input_schema", {}), + }, + } + ) + return result + + +def build_anthropic_sse_event(event_type: str, data: dict) -> str: + """Format a single Anthropic SSE event.""" + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n" + + +class AnthropicStreamEmitter: + """Converts generator events from generate_chat_completion_with_tools() + into Anthropic Messages SSE strings.""" + + def __init__(self) -> None: + self.block_index: int = 0 + self._text_block_open: bool = False + self._prev_text: str = "" + self._usage: dict = {} + + def start(self, message_id: str, model: str) -> list[str]: + """Emit message_start and open the first text content block.""" + events = [] + events.append( + build_anthropic_sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "content": [], + "model": model, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + }, + }, + ) + ) + events.extend(self._open_text_block()) + return events + + def feed(self, event: dict) -> list[str]: + """Process one generator event, return SSE strings.""" + etype = event.get("type", "") + if etype == "content": + return self._handle_content(event) + elif etype == "tool_start": + return self._handle_tool_start(event) + elif etype == "tool_end": + return self._handle_tool_end(event) + elif etype == "metadata": + self._usage = event.get("usage", {}) + return [] + # status events — no Anthropic equivalent + return [] + + def finish(self, stop_reason: str = "end_turn") -> list[str]: + """Close any open block and emit message_delta + message_stop.""" + events = [] + if self._text_block_open: + events.append(self._close_block()) + events.append( + build_anthropic_sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "usage": { + "output_tokens": self._usage.get("completion_tokens", 0), + }, + }, + ) + ) + events.append( + build_anthropic_sse_event( + "message_stop", + { + "type": "message_stop", + }, + ) + ) + return events + + def _handle_content(self, event: dict) -> list[str]: + cumulative = event.get("text", "") + new_text = cumulative[len(self._prev_text) :] + self._prev_text = cumulative + if not new_text: + return [] + if not self._text_block_open: + events = self._open_text_block() + else: + events = [] + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": new_text}, + }, + ) + ) + return events + + def _handle_tool_start(self, event: dict) -> list[str]: + events = [] + # Close current text block if open + if self._text_block_open: + events.append(self._close_block()) + # Open a tool_use block + self.block_index += 1 + events.append( + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": { + "type": "tool_use", + "id": event.get("tool_call_id", ""), + "name": event.get("tool_name", ""), + "input": {}, + }, + }, + ) + ) + # Emit the arguments as input_json_delta + args = event.get("arguments", {}) + if args: + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": { + "type": "input_json_delta", + "partial_json": json.dumps(args), + }, + }, + ) + ) + return events + + def _handle_tool_end(self, event: dict) -> list[str]: + events = [] + # Close the tool_use block + events.append(self._close_block()) + # Emit custom tool_result event (non-standard, ignored by SDKs) + events.append( + build_anthropic_sse_event( + "tool_result", + { + "type": "tool_result", + "tool_use_id": event.get("tool_call_id", ""), + "content": event.get("result", ""), + }, + ) + ) + # Open a new text block for the model's next response + self.block_index += 1 + events.extend(self._open_text_block()) + # Reset text tracking for the next synthesis turn + self._prev_text = "" + return events + + def _open_text_block(self) -> list[str]: + self._text_block_open = True + return [ + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": {"type": "text", "text": ""}, + }, + ) + ] + + def _close_block(self) -> str: + self._text_block_open = False + return build_anthropic_sse_event( + "content_block_stop", + { + "type": "content_block_stop", + "index": self.block_index, + }, + ) + + +class AnthropicPassthroughEmitter: + """Converts llama-server's OpenAI-format streaming chunks into Anthropic SSE. + + Used for the client-side tool-use pass-through path: the client (e.g. Claude + Code) sends its own tool definitions in the ``tools`` field and expects to + execute them itself. We forward them to llama-server and translate the + streaming response back to Anthropic format without executing anything. + """ + + def __init__(self) -> None: + self.block_index: int = -1 + self._current_block_type: Optional[str] = None # "text" | "tool_use" | None + self._tool_call_states: dict = {} # delta index -> {block_index, id, name} + self._usage: dict = {} + self._stop_reason: str = "end_turn" + + def start(self, message_id: str, model: str) -> list[str]: + return [ + build_anthropic_sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "content": [], + "model": model, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + }, + }, + ) + ] + + def feed_chunk(self, chunk: dict) -> list[str]: + """Process one OpenAI streaming chat.completion.chunk.""" + events: list[str] = [] + + # usage-only chunks carry token totals + usage = chunk.get("usage") + if usage: + self._usage = usage + + choices = chunk.get("choices") or [] + if not choices: + return events + + choice = choices[0] + delta = choice.get("delta") or {} + finish_reason = choice.get("finish_reason") + + # ── Text content ── + content = delta.get("content") + if content: + if self._current_block_type != "text": + if self._current_block_type is not None: + events.append(self._close_current_block()) + events.extend(self._open_text_block()) + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": content}, + }, + ) + ) + + # ── Tool calls (streaming deltas) ── + tool_calls = delta.get("tool_calls") or [] + for tc in tool_calls: + tc_idx = tc.get("index", 0) + fn = tc.get("function") or {} + if tc_idx not in self._tool_call_states: + # New tool call — close prior block, open tool_use block + if self._current_block_type is not None: + events.append(self._close_current_block()) + tc_id = tc.get("id", "") + tc_name = fn.get("name", "") + self.block_index += 1 + self._current_block_type = "tool_use" + self._tool_call_states[tc_idx] = { + "block_index": self.block_index, + "id": tc_id, + "name": tc_name, + } + events.append( + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": { + "type": "tool_use", + "id": tc_id, + "name": tc_name, + "input": {}, + }, + }, + ) + ) + + args_delta = fn.get("arguments", "") + if args_delta: + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self._tool_call_states[tc_idx]["block_index"], + "delta": { + "type": "input_json_delta", + "partial_json": args_delta, + }, + }, + ) + ) + + # ── Finish reason ── + if finish_reason: + if finish_reason == "tool_calls": + self._stop_reason = "tool_use" + elif finish_reason == "length": + self._stop_reason = "max_tokens" + else: + self._stop_reason = "end_turn" + + return events + + def finish(self) -> list[str]: + events: list[str] = [] + if self._current_block_type is not None: + events.append(self._close_current_block()) + events.append( + build_anthropic_sse_event( + "message_delta", + { + "type": "message_delta", + "delta": { + "stop_reason": self._stop_reason, + "stop_sequence": None, + }, + "usage": { + "output_tokens": self._usage.get("completion_tokens", 0), + }, + }, + ) + ) + events.append( + build_anthropic_sse_event( + "message_stop", + {"type": "message_stop"}, + ) + ) + return events + + def _open_text_block(self) -> list[str]: + self.block_index += 1 + self._current_block_type = "text" + return [ + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": {"type": "text", "text": ""}, + }, + ) + ] + + def _close_current_block(self) -> str: + idx = self.block_index + self._current_block_type = None + return build_anthropic_sse_event( + "content_block_stop", + { + "type": "content_block_stop", + "index": idx, + }, + ) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c84ac640df..d4a33387dc 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1063,6 +1063,7 @@ class LlamaCppBackend: speculative_type: Optional[str] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused + n_parallel: int = 1, ) -> bool: """ Start llama-server with a GGUF model. @@ -1283,7 +1284,7 @@ class LlamaCppBackend: "-c", str(effective_ctx) if effective_ctx > 0 else "0", "--parallel", - "1", # Single-user studio, saves VRAM + str(n_parallel), "--flash-attn", "on", # Force flash attention for speed ] diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py index 73d21130ae..c55e646508 100644 --- a/studio/backend/models/auth.py +++ b/studio/backend/models/auth.py @@ -5,6 +5,8 @@ Pydantic schemas for Authentication API """ +from typing import Optional + from pydantic import BaseModel, Field @@ -45,3 +47,44 @@ class ChangePasswordRequest(BaseModel): new_password: str = Field( ..., min_length = 8, description = "Replacement password (minimum 8 characters)" ) + + +# --------------------------------------------------------------------------- +# API key schemas +# --------------------------------------------------------------------------- + + +class CreateApiKeyRequest(BaseModel): + """Request body to create a new API key.""" + + name: str = Field(..., description = "Human-readable label for this key") + expires_in_days: Optional[int] = Field( + None, description = "Number of days until the key expires (None = never)" + ) + + +class ApiKeyResponse(BaseModel): + """Public representation of an API key (never contains the raw key).""" + + id: int + name: str + key_prefix: str = Field( + ..., description = "First 8 characters after sk-unsloth- for display" + ) + created_at: str + last_used_at: Optional[str] = None + expires_at: Optional[str] = None + is_active: bool + + +class CreateApiKeyResponse(BaseModel): + """Returned once when a key is created -- ``key`` is never shown again.""" + + key: str = Field(..., description = "Full API key (shown once)") + api_key: ApiKeyResponse + + +class ApiKeyListResponse(BaseModel): + """List of API keys for the authenticated user.""" + + api_keys: list[ApiKeyResponse] diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 8a27f4f630..8d9dc9830d 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -456,3 +456,241 @@ class ChatCompletion(BaseModel): model: str = "default" choices: list[CompletionChoice] usage: CompletionUsage = Field(default_factory = CompletionUsage) + + +# ===================================================================== +# OpenAI Responses API Models (/v1/responses) +# ===================================================================== + + +# ── Request models ────────────────────────────────────────────── + + +class ResponsesInputTextPart(BaseModel): + """Text content part in a Responses API message (type=input_text).""" + + type: Literal["input_text"] + text: str + + +class ResponsesInputImagePart(BaseModel): + """Image content part in a Responses API message (type=input_image).""" + + type: Literal["input_image"] + image_url: str = Field(..., description = "data:image/png;base64,... or https://...") + detail: Optional[Literal["auto", "low", "high"]] = "auto" + + +ResponsesContentPart = Union[ResponsesInputTextPart, ResponsesInputImagePart] + + +class ResponsesInputMessage(BaseModel): + """A single message in the Responses API input array.""" + + role: Literal["system", "user", "assistant", "developer"] + content: Union[str, list[ResponsesContentPart]] + + +class ResponsesRequest(BaseModel): + """OpenAI Responses API request.""" + + model: str = Field("default", description = "Model identifier") + input: Union[str, list[ResponsesInputMessage]] = Field( + default = [], + description = "Input text or message list", + ) + instructions: Optional[str] = Field( + None, description = "System / developer instructions" + ) + temperature: Optional[float] = Field(None, ge = 0.0, le = 2.0) + top_p: Optional[float] = Field(None, ge = 0.0, le = 1.0) + max_output_tokens: Optional[int] = Field(None, ge = 1) + stream: bool = Field(False, description = "Whether to stream the response via SSE") + + # Accepted but ignored -- keeps SDK clients from failing on unsupported fields + tools: Optional[list] = None + tool_choice: Optional[Any] = None + previous_response_id: Optional[str] = None + store: Optional[bool] = None + metadata: Optional[dict] = None + truncation: Optional[Any] = None + user: Optional[str] = None + text: Optional[Any] = None + reasoning: Optional[Any] = None + + model_config = {"extra": "allow"} + + +# ── Response models ───────────────────────────────────────────── + + +class ResponsesOutputTextContent(BaseModel): + """A text content block inside an output message.""" + + type: Literal["output_text"] = "output_text" + text: str + annotations: list = Field(default_factory = list) + + +class ResponsesOutputMessage(BaseModel): + """An output message in the Responses API response.""" + + type: Literal["message"] = "message" + id: str = Field(default_factory = lambda: f"msg_{uuid.uuid4().hex[:12]}") + status: Literal["completed", "in_progress"] = "completed" + role: Literal["assistant"] = "assistant" + content: list[ResponsesOutputTextContent] = Field(default_factory = list) + + +class ResponsesUsage(BaseModel): + """Token usage for a Responses API response (input_tokens, not prompt_tokens).""" + + input_tokens: int = 0 + output_tokens: int = 0 + total_tokens: int = 0 + + +class ResponsesResponse(BaseModel): + """Top-level Responses API response object.""" + + id: str = Field(default_factory = lambda: f"resp_{uuid.uuid4().hex[:12]}") + object: Literal["response"] = "response" + created_at: int = Field(default_factory = lambda: int(time.time())) + status: Literal["completed", "in_progress", "failed"] = "completed" + model: str = "default" + output: list[ResponsesOutputMessage] = Field(default_factory = list) + usage: ResponsesUsage = Field(default_factory = ResponsesUsage) + error: Optional[Any] = None + incomplete_details: Optional[Any] = None + instructions: Optional[str] = None + metadata: dict = Field(default_factory = dict) + temperature: Optional[float] = None + top_p: Optional[float] = None + max_output_tokens: Optional[int] = None + previous_response_id: Optional[str] = None + text: Optional[Any] = None + tool_choice: Optional[Any] = None + tools: list = Field(default_factory = list) + truncation: Optional[Any] = None + + +# ===================================================================== +# Anthropic Messages API Models (/v1/messages) +# ===================================================================== + + +# ── Request models ───────────────────────────────────────────── + + +class AnthropicTextBlock(BaseModel): + type: Literal["text"] + text: str + + +class AnthropicImageSource(BaseModel): + type: Literal["base64", "url"] + media_type: Optional[str] = None + data: Optional[str] = None + url: Optional[str] = None + + +class AnthropicImageBlock(BaseModel): + type: Literal["image"] + source: AnthropicImageSource + + +class AnthropicToolUseBlock(BaseModel): + type: Literal["tool_use"] + id: str + name: str + input: dict + + +class AnthropicToolResultBlock(BaseModel): + type: Literal["tool_result"] + tool_use_id: str + content: Union[str, list] = "" + + +AnthropicContentBlock = Union[ + AnthropicTextBlock, + AnthropicImageBlock, + AnthropicToolUseBlock, + AnthropicToolResultBlock, +] + + +class AnthropicMessage(BaseModel): + role: Literal["user", "assistant"] + content: Union[str, list[AnthropicContentBlock]] + + +class AnthropicTool(BaseModel): + name: str + description: Optional[str] = None + input_schema: dict + + +class AnthropicMessagesRequest(BaseModel): + model: str = "default" + max_tokens: Optional[int] = None + messages: list[AnthropicMessage] + system: Optional[Union[str, list]] = None + tools: Optional[list[AnthropicTool]] = None + tool_choice: Optional[Any] = None + stream: bool = False + temperature: Optional[float] = None + top_p: Optional[float] = None + top_k: Optional[int] = None + stop_sequences: Optional[list[str]] = None + metadata: Optional[dict] = None + # [x-unsloth] extensions — mirror the OpenAI endpoint convenience fields + min_p: Optional[float] = Field( + None, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold" + ) + repetition_penalty: Optional[float] = Field( + None, ge = 1.0, le = 2.0, description = "[x-unsloth] Repetition penalty" + ) + presence_penalty: Optional[float] = Field( + None, ge = 0.0, le = 2.0, description = "[x-unsloth] Presence penalty" + ) + enable_tools: Optional[bool] = None + enabled_tools: Optional[list[str]] = None + session_id: Optional[str] = None + model_config = {"extra": "allow"} + + +# ── Response models ──────────────────────────────────────────── + + +class AnthropicUsage(BaseModel): + input_tokens: int = 0 + output_tokens: int = 0 + + +class AnthropicResponseTextBlock(BaseModel): + type: Literal["text"] = "text" + text: str + + +class AnthropicResponseToolUseBlock(BaseModel): + type: Literal["tool_use"] = "tool_use" + id: str + name: str + input: dict + + +AnthropicResponseBlock = Union[ + AnthropicResponseTextBlock, AnthropicResponseToolUseBlock +] + + +class AnthropicMessagesResponse(BaseModel): + id: str = Field(default_factory = lambda: f"msg_{uuid.uuid4().hex[:24]}") + type: Literal["message"] = "message" + role: Literal["assistant"] = "assistant" + content: list[AnthropicResponseBlock] = Field(default_factory = list) + model: str = "default" + stop_reason: Optional[str] = None + stop_sequence: Optional[str] = None + usage: AnthropicUsage = Field(default_factory = AnthropicUsage) diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index db37ed837d..5cd23bd450 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -7,11 +7,17 @@ Authentication API routes from fastapi import APIRouter, Depends, HTTPException, status +from datetime import datetime, timedelta, timezone + from models.auth import ( + ApiKeyListResponse, + ApiKeyResponse, AuthLoginRequest, - RefreshTokenRequest, AuthStatusResponse, ChangePasswordRequest, + CreateApiKeyRequest, + CreateApiKeyResponse, + RefreshTokenRequest, ) from models.users import Token from auth import storage, hashing @@ -131,3 +137,68 @@ async def change_password( token_type = "bearer", must_change_password = False, ) + + +# --------------------------------------------------------------------------- +# API key management +# --------------------------------------------------------------------------- + + +def _row_to_api_key_response(row: dict) -> ApiKeyResponse: + return ApiKeyResponse( + id = row["id"], + name = row["name"], + key_prefix = row["key_prefix"], + created_at = row["created_at"], + last_used_at = row.get("last_used_at"), + expires_at = row.get("expires_at"), + is_active = bool(row["is_active"]), + ) + + +@router.post("/api-keys", response_model = CreateApiKeyResponse) +async def create_api_key( + payload: CreateApiKeyRequest, + current_subject: str = Depends(get_current_subject), +) -> CreateApiKeyResponse: + """Create a new API key. The raw key is returned once and cannot be retrieved later.""" + expires_at = None + if payload.expires_in_days is not None: + expires_at = ( + datetime.now(timezone.utc) + timedelta(days = payload.expires_in_days) + ).isoformat() + + raw_key, row = storage.create_api_key( + username = current_subject, + name = payload.name, + expires_at = expires_at, + ) + return CreateApiKeyResponse( + key = raw_key, + api_key = _row_to_api_key_response(row), + ) + + +@router.get("/api-keys", response_model = ApiKeyListResponse) +async def list_api_keys( + current_subject: str = Depends(get_current_subject), +) -> ApiKeyListResponse: + """List all API keys for the authenticated user (raw keys are never exposed).""" + rows = storage.list_api_keys(current_subject) + return ApiKeyListResponse( + api_keys = [_row_to_api_key_response(r) for r in rows], + ) + + +@router.delete("/api-keys/{key_id}") +async def revoke_api_key( + key_id: int, + current_subject: str = Depends(get_current_subject), +) -> dict: + """Revoke (soft-delete) an API key.""" + if not storage.revoke_api_key(current_subject, key_id): + raise HTTPException( + status_code = status.HTTP_404_NOT_FOUND, + detail = "API key not found", + ) + return {"detail": "API key revoked"} diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 7952d3a30b..0717b3bc93 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11,9 +11,10 @@ import time import uuid from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status -from fastapi.responses import StreamingResponse, JSONResponse +from fastapi.responses import StreamingResponse, JSONResponse, Response from typing import Optional import json +import httpx import structlog from loggers import get_logger import asyncio @@ -76,6 +77,7 @@ from models.inference import ( ChatCompletionRequest, ChatCompletionChunk, ChatCompletion, + ChatMessage, ChunkChoice, ChoiceDelta, CompletionChoice, @@ -83,6 +85,28 @@ from models.inference import ( CompletionUsage, ValidateModelRequest, ValidateModelResponse, + TextContentPart, + ImageContentPart, + ImageUrl, + ResponsesRequest, + ResponsesInputMessage, + ResponsesInputTextPart, + ResponsesInputImagePart, + ResponsesOutputTextContent, + ResponsesOutputMessage, + ResponsesUsage, + ResponsesResponse, + AnthropicMessagesRequest, + AnthropicMessagesResponse, + AnthropicResponseTextBlock, + AnthropicResponseToolUseBlock, + AnthropicUsage, +) +from core.inference.anthropic_compat import ( + anthropic_messages_to_openai, + anthropic_tools_to_openai, + AnthropicStreamEmitter, + AnthropicPassthroughEmitter, ) from auth.authentication import get_current_subject @@ -121,6 +145,7 @@ def get_llama_cpp_backend() -> LlamaCppBackend: @router.post("/load", response_model = LoadResponse) async def load_model( request: LoadRequest, + fastapi_request: Request, current_subject: str = Depends(get_current_subject), ): """ @@ -258,6 +283,8 @@ async def load_model( # Run in a thread so the event loop stays free for progress # polling and other requests during the (potentially long) # GGUF download + llama-server startup. + _n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1) + if config.gguf_hf_repo: # HF mode: download via huggingface_hub then start llama-server success = await asyncio.to_thread( @@ -271,6 +298,7 @@ async def load_model( chat_template_override = request.chat_template_override, cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, + n_parallel = _n_parallel, ) else: # Local mode: llama-server loads via -m @@ -284,6 +312,7 @@ async def load_model( chat_template_override = request.chat_template_override, cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, + n_parallel = _n_parallel, ) if not success: @@ -1843,3 +1872,1071 @@ async def openai_list_models( ) return {"object": "list", "data": models} + + +# ===================================================================== +# OpenAI-Compatible Completions Proxy (/completions → /v1/completions) +# ===================================================================== + + +@router.post("/completions") +async def openai_completions( + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + OpenAI-compatible text completions endpoint (non-chat). + + Transparently proxies to the running llama-server's ``/v1/completions``. + Only available when a GGUF model is loaded. + """ + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + body = await request.json() + target_url = f"{llama_backend.base_url}/v1/completions" + is_stream = body.get("stream", False) + + if is_stream: + + async def _stream(): + # Manual httpx client/response lifecycle — see + # _anthropic_passthrough_stream for the full rationale. Briefly: + # `async with` inside an async generator causes + # "Attempted to exit cancel scope in a different task" / + # "async generator ignored GeneratorExit" on Python 3.13 + + # httpcore 1.0.x when the generator is orphaned and finalized + # by GC. Closing via a finally block that catches Exception + # (but not BaseException) suppresses the anyio cleanup noise + # while letting GeneratorExit propagate cleanly. + client = httpx.AsyncClient(timeout = 600) + resp = None + try: + req = client.build_request("POST", target_url, json = body) + resp = await client.send(req, stream = True) + async for chunk in resp.aiter_bytes(): + yield chunk + except Exception as e: + logger.error("openai_completions stream error: %s", e) + finally: + if resp is not None: + try: + await resp.aclose() + except Exception: + pass + try: + await client.aclose() + except Exception: + pass + + return StreamingResponse(_stream(), media_type = "text/event-stream") + else: + async with httpx.AsyncClient() as client: + resp = await client.post(target_url, json = body, timeout = 600) + return Response( + content = resp.content, + status_code = resp.status_code, + media_type = "application/json", + ) + + +# ===================================================================== +# OpenAI-Compatible Embeddings Proxy (/embeddings → /v1/embeddings) +# ===================================================================== + + +@router.post("/embeddings") +async def openai_embeddings( + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + OpenAI-compatible embeddings endpoint. + + Transparently proxies to the running llama-server's ``/v1/embeddings``. + Only available when a GGUF model is loaded. + Note: the loaded model must support pooling; otherwise llama-server + will return an error (expected). + """ + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + body = await request.json() + target_url = f"{llama_backend.base_url}/v1/embeddings" + + async with httpx.AsyncClient() as client: + resp = await client.post(target_url, json = body, timeout = 600) + return Response( + content = resp.content, + status_code = resp.status_code, + media_type = "application/json", + ) + + +# ===================================================================== +# OpenAI Responses API (/responses → /v1/responses) +# ===================================================================== + + +def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]: + """Convert a ResponsesRequest into a list of ChatMessage for the completions backend.""" + messages: list[ChatMessage] = [] + + # System / developer instructions + if payload.instructions: + messages.append(ChatMessage(role = "system", content = payload.instructions)) + + # Simple string input + if isinstance(payload.input, str): + if payload.input: + messages.append(ChatMessage(role = "user", content = payload.input)) + return messages + + # List of ResponsesInputMessage + for msg in payload.input: + role = "system" if msg.role == "developer" else msg.role + + if isinstance(msg.content, str): + messages.append(ChatMessage(role = role, content = msg.content)) + else: + # Convert Responses content parts -> Chat content parts + parts = [] + for part in msg.content: + if isinstance(part, ResponsesInputTextPart): + parts.append(TextContentPart(type = "text", text = part.text)) + elif isinstance(part, ResponsesInputImagePart): + parts.append( + ImageContentPart( + type = "image_url", + image_url = ImageUrl(url = part.image_url, detail = part.detail), + ) + ) + messages.append(ChatMessage(role = role, content = parts if parts else "")) + + return messages + + +def _build_chat_request( + payload: ResponsesRequest, messages: list[ChatMessage], stream: bool +) -> ChatCompletionRequest: + """Build a ChatCompletionRequest from a ResponsesRequest.""" + chat_kwargs = dict( + model = payload.model, + messages = messages, + stream = stream, + ) + if payload.temperature is not None: + chat_kwargs["temperature"] = payload.temperature + if payload.top_p is not None: + chat_kwargs["top_p"] = payload.top_p + if payload.max_output_tokens is not None: + chat_kwargs["max_tokens"] = payload.max_output_tokens + return ChatCompletionRequest(**chat_kwargs) + + +async def _responses_non_streaming( + payload: ResponsesRequest, + messages: list[ChatMessage], + request: Request, +) -> JSONResponse: + """Handle a non-streaming Responses API call.""" + chat_req = _build_chat_request(payload, messages, stream = False) + result = await openai_chat_completions(chat_req, request) + + # openai_chat_completions returns a JSONResponse for non-streaming + if isinstance(result, JSONResponse): + body = json.loads(result.body.decode()) + elif isinstance(result, Response): + body = json.loads(result.body.decode()) + else: + body = result + + # Extract content and usage from the Chat Completions response + choices = body.get("choices", []) + text = "" + if choices: + msg = choices[0].get("message", {}) + text = msg.get("content", "") or "" + + usage_data = body.get("usage", {}) + input_tokens = usage_data.get("prompt_tokens", 0) + output_tokens = usage_data.get("completion_tokens", 0) + + resp_id = f"resp_{uuid.uuid4().hex[:12]}" + msg_id = f"msg_{uuid.uuid4().hex[:12]}" + + response = ResponsesResponse( + id = resp_id, + created_at = int(time.time()), + status = "completed", + model = body.get("model", payload.model), + output = [ + ResponsesOutputMessage( + id = msg_id, + status = "completed", + role = "assistant", + content = [ + ResponsesOutputTextContent(text = text), + ], + ), + ], + usage = ResponsesUsage( + input_tokens = input_tokens, + output_tokens = output_tokens, + total_tokens = input_tokens + output_tokens, + ), + temperature = payload.temperature, + top_p = payload.top_p, + max_output_tokens = payload.max_output_tokens, + instructions = payload.instructions, + ) + return JSONResponse(content = response.model_dump()) + + +async def _responses_stream( + payload: ResponsesRequest, + messages: list[ChatMessage], + request: Request, +): + """Handle a streaming Responses API call, emitting named SSE events.""" + resp_id = f"resp_{uuid.uuid4().hex[:12]}" + msg_id = f"msg_{uuid.uuid4().hex[:12]}" + item_id = f"item_{uuid.uuid4().hex[:12]}" + created_at = int(time.time()) + + chat_req = _build_chat_request(payload, messages, stream = True) + result = await openai_chat_completions(chat_req, request) + + async def event_generator(): + full_text = "" + input_tokens = 0 + output_tokens = 0 + + # ── Preamble events ── + yield f"event: response.created\ndata: {json.dumps({'type': 'response.created', 'response': {'id': resp_id, 'object': 'response', 'created_at': created_at, 'status': 'in_progress', 'model': payload.model, 'output': [], 'usage': {'input_tokens': 0, 'output_tokens': 0, 'total_tokens': 0}}})}\n\n" + + # output_item.added + output_item = { + "type": "message", + "id": msg_id, + "status": "in_progress", + "role": "assistant", + "content": [], + } + yield f"event: response.output_item.added\ndata: {json.dumps({'type': 'response.output_item.added', 'output_index': 0, 'item': output_item})}\n\n" + + # content_part.added + content_part = {"type": "output_text", "text": "", "annotations": []} + yield f"event: response.content_part.added\ndata: {json.dumps({'type': 'response.content_part.added', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'part': content_part})}\n\n" + + # ── Stream delta events from the inner chat completions stream ── + if isinstance(result, StreamingResponse): + async for raw_chunk in result.body_iterator: + if isinstance(raw_chunk, bytes): + raw_chunk = raw_chunk.decode("utf-8", errors = "replace") + + for line in raw_chunk.split("\n"): + line = line.strip() + if not line.startswith("data: "): + continue + data_str = line[6:] + if data_str == "[DONE]": + continue + try: + chunk_data = json.loads(data_str) + except json.JSONDecodeError: + continue + + choices = chunk_data.get("choices", []) + if not choices: + # Check for usage in final chunk + usage = chunk_data.get("usage") + if usage: + input_tokens = usage.get("prompt_tokens", input_tokens) + output_tokens = usage.get( + "completion_tokens", output_tokens + ) + continue + + delta = choices[0].get("delta", {}) + content = delta.get("content") + if content: + full_text += content + delta_event = { + "type": "response.output_text.delta", + "item_id": msg_id, + "output_index": 0, + "content_index": 0, + "delta": content, + } + yield f"event: response.output_text.delta\ndata: {json.dumps(delta_event)}\n\n" + + # Check for usage in chunk + usage = chunk_data.get("usage") + if usage: + input_tokens = usage.get("prompt_tokens", input_tokens) + output_tokens = usage.get("completion_tokens", output_tokens) + + # ── Closing events ── + # output_text.done + yield f"event: response.output_text.done\ndata: {json.dumps({'type': 'response.output_text.done', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'text': full_text})}\n\n" + + # content_part.done + yield f"event: response.content_part.done\ndata: {json.dumps({'type': 'response.content_part.done', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'part': {'type': 'output_text', 'text': full_text, 'annotations': []}})}\n\n" + + # output_item.done + yield f"event: response.output_item.done\ndata: {json.dumps({'type': 'response.output_item.done', 'output_index': 0, 'item': {'type': 'message', 'id': msg_id, 'status': 'completed', 'role': 'assistant', 'content': [{'type': 'output_text', 'text': full_text, 'annotations': []}]}})}\n\n" + + # response.completed + total_tokens = input_tokens + output_tokens + completed_response = { + "type": "response.completed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "completed", + "model": payload.model, + "output": [ + { + "type": "message", + "id": msg_id, + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": full_text, + "annotations": [], + } + ], + } + ], + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": total_tokens, + }, + }, + } + yield f"event: response.completed\ndata: {json.dumps(completed_response)}\n\n" + + return StreamingResponse( + event_generator(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + +@router.post("/responses") +async def openai_responses( + payload: ResponsesRequest, + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + OpenAI Responses API endpoint. + + Accepts the Responses-format request, converts it to a + ChatCompletionRequest internally, and returns a response + matching the OpenAI Responses API schema (output array, + input_tokens/output_tokens, named SSE events for streaming). + """ + messages = _normalise_responses_input(payload) + if not messages: + raise HTTPException(status_code = 400, detail = "No input provided.") + + if payload.stream: + return await _responses_stream(payload, messages, request) + return await _responses_non_streaming(payload, messages, request) + + +# ===================================================================== +# Anthropic-Compatible Messages API (/messages → /v1/messages) +# ===================================================================== + + +@router.post("/messages") +async def anthropic_messages( + payload: AnthropicMessagesRequest, + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + Anthropic-compatible Messages API endpoint. + + Translates Anthropic message format to internal OpenAI format, runs + through the existing agentic tool loop when tools are provided, and + returns responses in Anthropic Messages API format (streaming SSE or + non-streaming JSON). + """ + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + model_name = getattr(llama_backend, "model_identifier", None) or payload.model + message_id = f"msg_{uuid.uuid4().hex[:24]}" + + # ── Translate Anthropic → OpenAI ────────────────────────── + openai_messages = anthropic_messages_to_openai( + [m.model_dump() for m in payload.messages], + payload.system, + ) + + temperature = payload.temperature if payload.temperature is not None else 0.6 + top_p = payload.top_p if payload.top_p is not None else 0.95 + top_k = payload.top_k if payload.top_k is not None else 20 + min_p = payload.min_p if payload.min_p is not None else 0.01 + repetition_penalty = ( + payload.repetition_penalty if payload.repetition_penalty is not None else 1.0 + ) + presence_penalty = ( + payload.presence_penalty if payload.presence_penalty is not None else 0.0 + ) + stop = payload.stop_sequences or None + + # tool_choice is declared on AnthropicMessagesRequest for Anthropic SDK + # compatibility (the SDK often sets it by default), but it is not + # currently honored by Unsloth's backend. Warn once per request so the + # silent drop is visible to operators instead of looking like a model + # quality issue to clients. + if payload.tool_choice is not None: + logger.warning( + "anthropic_messages.tool_choice_ignored", + tool_choice = payload.tool_choice, + note = ( + "tool_choice is accepted for Anthropic SDK compatibility but not " + "honored by Unsloth. Use enable_tools / enabled_tools (server-side " + "built-in tools) or restrict the `tools` array (client-side) to " + "control which tools the model sees." + ), + ) + + cancel_event = threading.Event() + + # ── Tool routing ────────────────────────────────────────── + # Three paths: + # 1. enable_tools=true → server-side execution of built-in tools (Unsloth shorthand) + # 2. tools=[...] only → client-side pass-through (standard Anthropic behavior) + # 3. neither → plain chat + server_tools = payload.enable_tools and llama_backend.supports_tools + client_tools = ( + not server_tools + and payload.tools + and len(payload.tools) > 0 + and llama_backend.supports_tools + ) + + # ── Client-side pass-through path ───────────────────────── + if client_tools: + openai_tools = anthropic_tools_to_openai(payload.tools) + + if payload.stream: + return await _anthropic_passthrough_stream( + request, + cancel_event, + llama_backend, + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + payload.max_tokens, + message_id, + model_name, + stop = stop, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ) + return await _anthropic_passthrough_non_streaming( + llama_backend, + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + payload.max_tokens, + message_id, + model_name, + stop = stop, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ) + + if server_tools: + from core.inference.tools import ALL_TOOLS + + if payload.enabled_tools is not None: + openai_tools = [ + t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools + ] + else: + openai_tools = ALL_TOOLS + + # Build tool-use system prompt nudge (same logic as /chat/completions) + _tool_names = {t["function"]["name"] for t in openai_tools} + _has_web = "web_search" in _tool_names + _has_code = "python" in _tool_names or "terminal" in _tool_names + + _date_line = f"The current date is {_date.today().isoformat()}." + _model_size_b = _extract_model_size_b(model_name) + _is_small_model = _model_size_b is not None and _model_size_b < 9 + + if _is_small_model: + _web_tips = "Do not repeat the same search query." + else: + _web_tips = ( + "When you search and find a relevant URL in the results, " + "fetch its full content by calling web_search with the url parameter. " + "Do not repeat the same search query. If a search returns " + "no useful results, try rephrasing or fetching a result URL directly." + ) + _code_tips = ( + "Use code execution for math, calculations, data processing, " + "or to parse and analyze information from tool results." + ) + + if _has_web and _has_code: + _nudge = ( + _date_line + " " + "You have access to tools. When appropriate, prefer using " + "tools rather than answering from memory. " + + _web_tips + + " " + + _code_tips + ) + elif _has_code: + _nudge = ( + _date_line + " " + "You have access to tools. When appropriate, prefer using " + "code execution rather than answering from memory. " + _code_tips + ) + elif _has_web: + _nudge = ( + _date_line + " " + "You have access to tools. When appropriate, prefer using " + "web search for up-to-date or uncertain factual " + "information rather than answering from memory. " + _web_tips + ) + else: + _nudge = "" + + if _nudge: + _nudge += _TOOL_ACTION_NUDGE + # Inject into system prompt + if openai_messages and openai_messages[0].get("role") == "system": + openai_messages[0]["content"] = ( + openai_messages[0]["content"].rstrip() + "\n\n" + _nudge + ) + else: + openai_messages.insert(0, {"role": "system", "content": _nudge}) + + # Strip stale tool-call XML from conversation + for _msg in openai_messages: + if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): + _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip() + + def _run_tool_gen(): + return llama_backend.generate_chat_completion_with_tools( + messages = openai_messages, + tools = openai_tools, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + max_tokens = payload.max_tokens, + stop = stop, + cancel_event = cancel_event, + max_tool_iterations = 25, + auto_heal_tool_calls = True, + tool_call_timeout = 300, + session_id = payload.session_id, + ) + + if payload.stream: + return await _anthropic_tool_stream( + request, + cancel_event, + _run_tool_gen, + message_id, + model_name, + ) + return await _anthropic_tool_non_streaming( + _run_tool_gen, + message_id, + model_name, + ) + + # ── No-tool path ────────────────────────────────────────── + def _run_plain_gen(): + return llama_backend.generate_chat_completion( + messages = openai_messages, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + max_tokens = payload.max_tokens, + stop = stop, + cancel_event = cancel_event, + ) + + if payload.stream: + return await _anthropic_plain_stream( + request, + cancel_event, + _run_plain_gen, + message_id, + model_name, + ) + return await _anthropic_plain_non_streaming( + _run_plain_gen, + message_id, + model_name, + ) + + +async def _anthropic_tool_stream( + request, + cancel_event, + run_gen, + message_id, + model_name, +): + """Streaming response for the tool-calling path.""" + _sentinel = object() + + async def _stream(): + emitter = AnthropicStreamEmitter() + for line in emitter.start(message_id, model_name): + yield line + + gen = run_gen() + try: + while True: + if await request.is_disconnected(): + cancel_event.set() + return + event = await asyncio.to_thread(next, gen, _sentinel) + if event is _sentinel: + break + # Strip leaked tool-call XML from content events + if event.get("type") == "content": + event = dict(event) + event["text"] = _TOOL_XML_RE.sub("", event["text"]) + for line in emitter.feed(event): + yield line + except Exception as e: + logger.error("anthropic_messages stream error: %s", e) + + for line in emitter.finish("end_turn"): + yield line + + return StreamingResponse( + _stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + +async def _anthropic_plain_stream( + request, + cancel_event, + run_gen, + message_id, + model_name, +): + """Streaming response for the no-tool path.""" + _sentinel = object() + + async def _stream(): + emitter = AnthropicStreamEmitter() + for line in emitter.start(message_id, model_name): + yield line + + gen = run_gen() + try: + while True: + if await request.is_disconnected(): + cancel_event.set() + return + cumulative = await asyncio.to_thread(next, gen, _sentinel) + if cumulative is _sentinel: + break + if isinstance(cumulative, dict): + if cumulative.get("type") == "metadata": + for line in emitter.feed(cumulative): + yield line + continue + # Plain generator yields cumulative text strings + for line in emitter.feed({"type": "content", "text": cumulative}): + yield line + except Exception as e: + logger.error("anthropic_messages stream error: %s", e) + + for line in emitter.finish("end_turn"): + yield line + + return StreamingResponse( + _stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + +async def _anthropic_tool_non_streaming(run_gen, message_id, model_name): + """Non-streaming response for the tool-calling path. + + Builds ``content_blocks`` in generation order (text → tool_use → text → + tool_use → ...), mirroring the streaming emitter's behavior. Deltas + within a single synthesis turn are merged into the trailing text block; + tool_use blocks interrupt the text sequence and open a new text block on + the next content event. + + ``prev_text`` is reset on ``tool_end`` because + ``generate_chat_completion_with_tools`` yields cumulative content *per + turn* — the first content event of turn N+1 must diff against an empty + baseline, not against turn N's final length. + """ + content_blocks: list = [] + usage = {} + prev_text = "" + + for event in run_gen(): + etype = event.get("type", "") + if etype == "content": + # Strip leaked tool-call XML + clean = _TOOL_XML_RE.sub("", event["text"]) + new = clean[len(prev_text) :] + prev_text = clean + if new: + if content_blocks and isinstance( + content_blocks[-1], AnthropicResponseTextBlock + ): + content_blocks[-1].text += new + else: + content_blocks.append(AnthropicResponseTextBlock(text = new)) + elif etype == "tool_start": + content_blocks.append( + AnthropicResponseToolUseBlock( + id = event["tool_call_id"], + name = event["tool_name"], + input = event.get("arguments", {}), + ) + ) + elif etype == "tool_end": + prev_text = "" + elif etype == "metadata": + usage = event.get("usage", {}) + + resp = AnthropicMessagesResponse( + id = message_id, + model = model_name, + content = content_blocks, + stop_reason = "end_turn", + usage = AnthropicUsage( + input_tokens = usage.get("prompt_tokens", 0), + output_tokens = usage.get("completion_tokens", 0), + ), + ) + return JSONResponse(content = resp.model_dump()) + + +async def _anthropic_plain_non_streaming(run_gen, message_id, model_name): + """Non-streaming response for the no-tool path.""" + text_parts = [] + usage = {} + prev_text = "" + + for cumulative in run_gen(): + if isinstance(cumulative, dict): + if cumulative.get("type") == "metadata": + usage = cumulative.get("usage", {}) + continue + new = cumulative[len(prev_text) :] + prev_text = cumulative + if new: + text_parts.append(new) + + full_text = "".join(text_parts) + content_blocks = [] + if full_text: + content_blocks.append(AnthropicResponseTextBlock(text = full_text)) + + resp = AnthropicMessagesResponse( + id = message_id, + model = model_name, + content = content_blocks, + stop_reason = "end_turn", + usage = AnthropicUsage( + input_tokens = usage.get("prompt_tokens", 0), + output_tokens = usage.get("completion_tokens", 0), + ), + ) + return JSONResponse(content = resp.model_dump()) + + +# ===================================================================== +# Client-side tool pass-through (Anthropic-native tools field) +# ===================================================================== + + +def _build_passthrough_payload( + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + max_tokens, + stream, + stop = None, + min_p = None, + repetition_penalty = None, + presence_penalty = None, +): + body = { + "messages": openai_messages, + "tools": openai_tools, + "tool_choice": "auto", + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "stream": stream, + } + if stream: + body["stream_options"] = {"include_usage": True} + if max_tokens is not None: + body["max_tokens"] = max_tokens + if stop: + body["stop"] = stop + if min_p is not None: + body["min_p"] = min_p + if repetition_penalty is not None: + # llama-server's field is "repeat_penalty", not "repetition_penalty" + body["repeat_penalty"] = repetition_penalty + if presence_penalty is not None: + body["presence_penalty"] = presence_penalty + return body + + +async def _anthropic_passthrough_stream( + request, + cancel_event, + llama_backend, + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + max_tokens, + message_id, + model_name, + stop = None, + min_p = None, + repetition_penalty = None, + presence_penalty = None, +): + """Streaming client-side pass-through: forward tools to llama-server and + translate its streaming response to Anthropic SSE without executing anything.""" + target_url = f"{llama_backend.base_url}/v1/chat/completions" + body = _build_passthrough_payload( + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + max_tokens, + True, + stop = stop, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ) + + async def _stream(): + emitter = AnthropicPassthroughEmitter() + for line in emitter.start(message_id, model_name): + yield line + + # Manage the httpx client and response MANUALLY — no `async with`. + # + # On Python 3.13 + httpcore 1.0.x, an orphaned async generator (e.g. + # when the client disconnects mid-stream and Starlette drops the + # StreamingResponse iterator without explicitly calling aclose()) + # is finalized by Python's asyncgen GC hook in a DIFFERENT asyncio + # task than the one that originally entered the httpx context + # managers. When `async with` exits run in the wrong task, httpcore's + # internal `HTTP11ConnectionByteStream.aclose()` hits + # `anyio.CancelScope.__exit__` with a mismatched task and raises + # RuntimeError("Attempted to exit cancel scope in a different task"), + # which escapes as "Exception ignored in:" because it happens during + # GC finalization outside any user-owned try/except. + # + # The fix: do not use `async with` for the client/response. Close + # them in a finally block wrapped in `try: ... except Exception: pass`. + # This narrowly suppresses RuntimeError / other Exception subclasses + # from the anyio cleanup noise while letting GeneratorExit (a + # BaseException, not Exception) propagate through cleanly so the + # generator terminates as Python expects. + client = httpx.AsyncClient(timeout = 600) + resp = None + try: + req = client.build_request("POST", target_url, json = body) + resp = await client.send(req, stream = True) + + async for raw_line in resp.aiter_lines(): + if await request.is_disconnected(): + cancel_event.set() + break + if not raw_line or not raw_line.startswith("data: "): + continue + data_str = raw_line[6:] + if data_str.strip() == "[DONE]": + break + try: + chunk = json.loads(data_str) + except json.JSONDecodeError: + continue + for line in emitter.feed_chunk(chunk): + yield line + except Exception as e: + logger.error("anthropic_messages passthrough stream error: %s", e) + finally: + if resp is not None: + try: + await resp.aclose() + except Exception: + pass + try: + await client.aclose() + except Exception: + pass + + for line in emitter.finish(): + yield line + + return StreamingResponse( + _stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + +async def _anthropic_passthrough_non_streaming( + llama_backend, + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + max_tokens, + message_id, + model_name, + stop = None, + min_p = None, + repetition_penalty = None, + presence_penalty = None, +): + """Non-streaming client-side pass-through.""" + target_url = f"{llama_backend.base_url}/v1/chat/completions" + body = _build_passthrough_payload( + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + max_tokens, + False, + stop = stop, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ) + + async with httpx.AsyncClient() as client: + resp = await client.post(target_url, json = body, timeout = 600) + + if resp.status_code != 200: + raise HTTPException( + status_code = resp.status_code, + detail = f"llama-server error: {resp.text[:500]}", + ) + + data = resp.json() + choice = (data.get("choices") or [{}])[0] + message = choice.get("message") or {} + finish_reason = choice.get("finish_reason") + + content_blocks = [] + text = message.get("content") or "" + if text: + text = _TOOL_XML_RE.sub("", text).strip() + if text: + content_blocks.append(AnthropicResponseTextBlock(text = text)) + + tool_calls = message.get("tool_calls") or [] + for tc in tool_calls: + fn = tc.get("function") or {} + try: + args = json.loads(fn.get("arguments", "{}")) + except json.JSONDecodeError: + args = {} + content_blocks.append( + AnthropicResponseToolUseBlock( + id = tc.get("id", ""), + name = fn.get("name", ""), + input = args, + ) + ) + + if tool_calls: + stop_reason = "tool_use" + elif finish_reason == "length": + stop_reason = "max_tokens" + else: + stop_reason = "end_turn" + + usage = data.get("usage") or {} + resp_obj = AnthropicMessagesResponse( + id = message_id, + model = model_name, + content = content_blocks, + stop_reason = stop_reason, + usage = AnthropicUsage( + input_tokens = usage.get("prompt_tokens", 0), + output_tokens = usage.get("completion_tokens", 0), + ), + ) + return JSONResponse(content = resp_obj.model_dump()) diff --git a/studio/backend/run.py b/studio/backend/run.py index 86c1194661..9675b9ea4c 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -248,6 +248,7 @@ def run_server( port: int = 8888, frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist", silent: bool = False, + llama_parallel_slots: int = 1, ): """ Start the FastAPI server. @@ -257,6 +258,7 @@ def run_server( port: Port to bind to (auto-increments if in use) frontend_path: Path to frontend build directory (optional) silent: Suppress startup messages + llama_parallel_slots: Number of parallel slots for llama-server Note: Signal handlers are NOT registered here so that embedders @@ -331,6 +333,7 @@ def run_server( # binds (port==0) leave it unset and let request handlers fall back # to the ASGI request scope or request.base_url. app.state.server_port = port if port and port > 0 else None + app.state.llama_parallel_slots = llama_parallel_slots # Run server in a daemon thread def _run(): diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py index 053e9b85d9..6aa6d314c1 100644 --- a/studio/backend/tests/conftest.py +++ b/studio/backend/tests/conftest.py @@ -3,14 +3,136 @@ """ Shared pytest configuration for the backend test suite. -Ensures that the backend root is on sys.path so that -`import utils.utils` (and similar flat imports) resolve correctly. + +Responsibilities: +1. Put the backend root on sys.path so `from models.inference import ...` + (and similar flat imports) resolve in test modules — mirrors how the + app itself is launched. +2. Provide a hybrid ``studio_server`` session fixture for end-to-end tests + (see ``test_studio_api.py``). The fixture supports two invocation modes: + + a. **External server.** If ``UNSLOTH_E2E_BASE_URL`` is set, tests point + at an already-running Studio instance. ``UNSLOTH_E2E_API_KEY`` must + also be set. This is the fast-iteration mode: start the server once + with ``unsloth studio run ...``, then run pytest against it many + times with no per-run GGUF load cost. + + b. **Fixture-managed server.** Otherwise, the fixture launches a fresh + server via ``_start_server`` and tears it down at session end. This + is the one-shot mode for CI or a clean-slate verification run. + + The model / variant for mode (b) come from ``--unsloth-model`` / + ``--unsloth-gguf-variant`` pytest options, then ``UNSLOTH_E2E_MODEL`` / + ``UNSLOTH_E2E_VARIANT`` env vars, then the defaults in + ``test_studio_api.py``. """ +import os import sys from pathlib import Path +import pytest + # Add backend root to sys.path (mirrors how the app itself is launched) _backend_root = Path(__file__).resolve().parent.parent if str(_backend_root) not in sys.path: sys.path.insert(0, str(_backend_root)) + + +# ── Pytest CLI options ─────────────────────────────────────────────── + + +def pytest_addoption(parser): + group = parser.getgroup( + "unsloth-e2e", + "Unsloth Studio end-to-end test options", + ) + group.addoption( + "--unsloth-model", + action = "store", + default = None, + help = ( + "GGUF model id used when starting a server for e2e tests. " + "Ignored if UNSLOTH_E2E_BASE_URL is set. Overrides " + "UNSLOTH_E2E_MODEL env var. Defaults to test_studio_api.py's " + "DEFAULT_MODEL." + ), + ) + group.addoption( + "--unsloth-gguf-variant", + action = "store", + default = None, + help = ( + "GGUF variant used when starting a server for e2e tests. " + "Ignored if UNSLOTH_E2E_BASE_URL is set. Overrides " + "UNSLOTH_E2E_VARIANT env var. Defaults to test_studio_api.py's " + "DEFAULT_VARIANT." + ), + ) + + +# ── E2E server fixtures ────────────────────────────────────────────── + + +@pytest.fixture(scope = "session") +def studio_server(request): + """Yield ``(base_url, api_key)`` for e2e tests. + + Resolution order: + + 1. If ``UNSLOTH_E2E_BASE_URL`` is set → point at that server, + require ``UNSLOTH_E2E_API_KEY`` alongside (skip if missing). + 2. Otherwise → start a fresh ``unsloth studio run`` subprocess via + the existing ``_start_server`` helper in ``test_studio_api.py`` + and tear it down on session teardown. + + Session-scoped so the expensive GGUF load happens at most once per + pytest invocation. Lazily instantiated — tests that don't request + the fixture (e.g. the unit tests in ``test_anthropic_messages.py`` + or ``test_help_output``) do not trigger server startup. + """ + external_url = os.environ.get("UNSLOTH_E2E_BASE_URL") + if external_url: + api_key = os.environ.get("UNSLOTH_E2E_API_KEY") + if not api_key: + pytest.skip( + "UNSLOTH_E2E_BASE_URL is set but UNSLOTH_E2E_API_KEY is " + "missing — tests that require auth cannot run against an " + "external server without it.", + ) + yield external_url, api_key + return + + # Lazy import: pytest has already loaded test_studio_api into + # sys.modules by the time any test requests this fixture, so this + # is a cache hit, not a re-execution. + import test_studio_api as _e2e + + model = ( + request.config.getoption("--unsloth-model") + or os.environ.get("UNSLOTH_E2E_MODEL") + or _e2e.DEFAULT_MODEL + ) + variant = ( + request.config.getoption("--unsloth-gguf-variant") + or os.environ.get("UNSLOTH_E2E_VARIANT") + or _e2e.DEFAULT_VARIANT + ) + + proc, api_key = _e2e._start_server(model, variant) + try: + yield f"http://{_e2e.HOST}:{_e2e.PORT}", api_key + finally: + _e2e._kill_server(proc) + + +@pytest.fixture +def base_url(studio_server): + """Base URL for the e2e Studio server (from ``studio_server``).""" + return studio_server[0] + + +@pytest.fixture +def api_key(studio_server): + """API key for the e2e Studio server (from ``studio_server``).""" + return studio_server[1] diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py new file mode 100644 index 0000000000..ec432df9e6 --- /dev/null +++ b/studio/backend/tests/test_anthropic_messages.py @@ -0,0 +1,774 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +""" +Tests for the Anthropic Messages API schemas and translation layer. +No running server or GPU required. +""" + +import sys +import os +import json + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from models.inference import ( + AnthropicMessagesRequest, + AnthropicMessagesResponse, + AnthropicMessage, + AnthropicTextBlock, + AnthropicToolUseBlock, + AnthropicToolResultBlock, + AnthropicTool, + AnthropicUsage, + AnthropicResponseTextBlock, + AnthropicResponseToolUseBlock, +) +from core.inference.anthropic_compat import ( + anthropic_messages_to_openai, + anthropic_tools_to_openai, + build_anthropic_sse_event, + AnthropicStreamEmitter, + AnthropicPassthroughEmitter, +) + + +# ===================================================================== +# Pydantic model tests +# ===================================================================== + + +class TestAnthropicModels: + def test_minimal_request(self): + req = AnthropicMessagesRequest( + messages = [{"role": "user", "content": "Hi"}], + ) + assert req.max_tokens is None + assert req.model == "default" + assert req.stream is False + + def test_max_tokens_optional(self): + req = AnthropicMessagesRequest( + max_tokens = 100, + messages = [{"role": "user", "content": "Hi"}], + ) + assert req.max_tokens == 100 + + def test_system_as_string(self): + req = AnthropicMessagesRequest( + max_tokens = 50, + messages = [{"role": "user", "content": "Hi"}], + system = "You are helpful.", + ) + assert req.system == "You are helpful." + + def test_tools_field_parses(self): + req = AnthropicMessagesRequest( + max_tokens = 100, + messages = [{"role": "user", "content": "Hi"}], + tools = [{"name": "web_search", "input_schema": {"type": "object"}}], + ) + assert len(req.tools) == 1 + assert req.tools[0].name == "web_search" + + def test_extra_fields_accepted(self): + req = AnthropicMessagesRequest( + max_tokens = 100, + messages = [{"role": "user", "content": "Hi"}], + some_future_field = "hello", + ) + assert req.max_tokens == 100 + + def test_stream_defaults_false(self): + req = AnthropicMessagesRequest( + max_tokens = 100, + messages = [{"role": "user", "content": "Hi"}], + ) + assert req.stream is False + + def test_enable_tools_shorthand(self): + req = AnthropicMessagesRequest( + messages = [{"role": "user", "content": "Hi"}], + enable_tools = True, + enabled_tools = ["web_search", "python"], + session_id = "my-session", + ) + assert req.enable_tools is True + assert req.enabled_tools == ["web_search", "python"] + assert req.session_id == "my-session" + + def test_extension_fields_default_none(self): + req = AnthropicMessagesRequest( + messages = [{"role": "user", "content": "Hi"}], + ) + assert req.enable_tools is None + assert req.enabled_tools is None + assert req.session_id is None + + def test_response_model_defaults(self): + resp = AnthropicMessagesResponse() + assert resp.type == "message" + assert resp.role == "assistant" + assert resp.id.startswith("msg_") + assert resp.content == [] + assert resp.usage.input_tokens == 0 + + +# ===================================================================== +# Message translation tests +# ===================================================================== + + +class TestAnthropicMessagesToOpenAI: + def test_simple_user_message(self): + msgs = [{"role": "user", "content": "Hello"}] + result = anthropic_messages_to_openai(msgs) + assert result == [{"role": "user", "content": "Hello"}] + + def test_system_string_prepended(self): + msgs = [{"role": "user", "content": "Hello"}] + result = anthropic_messages_to_openai(msgs, system = "Be brief.") + assert result[0] == {"role": "system", "content": "Be brief."} + assert result[1] == {"role": "user", "content": "Hello"} + + def test_system_as_block_list(self): + system = [ + {"type": "text", "text": "Be brief."}, + {"type": "text", "text": "Be accurate."}, + ] + msgs = [{"role": "user", "content": "Hello"}] + result = anthropic_messages_to_openai(msgs, system = system) + assert result[0]["role"] == "system" + assert "Be brief." in result[0]["content"] + assert "Be accurate." in result[0]["content"] + + def test_multi_turn_conversation(self): + msgs = [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello!"}, + {"role": "user", "content": "How are you?"}, + ] + result = anthropic_messages_to_openai(msgs) + assert len(result) == 3 + assert result[0]["role"] == "user" + assert result[1]["role"] == "assistant" + assert result[2]["role"] == "user" + + def test_assistant_tool_use_maps_to_tool_calls(self): + msgs = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me search."}, + { + "type": "tool_use", + "id": "tu_1", + "name": "web_search", + "input": {"query": "test"}, + }, + ], + } + ] + result = anthropic_messages_to_openai(msgs) + assert len(result) == 1 + m = result[0] + assert m["role"] == "assistant" + assert m["content"] == "Let me search." + assert len(m["tool_calls"]) == 1 + tc = m["tool_calls"][0] + assert tc["id"] == "tu_1" + assert tc["function"]["name"] == "web_search" + assert json.loads(tc["function"]["arguments"]) == {"query": "test"} + + def test_tool_result_maps_to_tool_role(self): + msgs = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": "Result text", + }, + ], + } + ] + result = anthropic_messages_to_openai(msgs) + assert len(result) == 1 + assert result[0]["role"] == "tool" + assert result[0]["tool_call_id"] == "tu_1" + assert result[0]["content"] == "Result text" + + def test_mixed_text_and_tool_use_blocks(self): + msgs = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Thinking..."}, + { + "type": "tool_use", + "id": "tu_1", + "name": "python", + "input": {"code": "1+1"}, + }, + { + "type": "tool_use", + "id": "tu_2", + "name": "terminal", + "input": {"command": "ls"}, + }, + ], + } + ] + result = anthropic_messages_to_openai(msgs) + assert len(result) == 1 + m = result[0] + assert m["content"] == "Thinking..." + assert len(m["tool_calls"]) == 2 + + def test_tool_result_with_list_content(self): + msgs = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": [ + {"type": "text", "text": "Line 1"}, + {"type": "text", "text": "Line 2"}, + ], + }, + ], + } + ] + result = anthropic_messages_to_openai(msgs) + assert result[0]["content"] == "Line 1 Line 2" + + +# ===================================================================== +# Tool translation tests +# ===================================================================== + + +class TestAnthropicToolsToOpenAI: + def test_single_tool(self): + tools = [ + { + "name": "web_search", + "description": "Search", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + } + ] + result = anthropic_tools_to_openai(tools) + assert len(result) == 1 + assert result[0]["type"] == "function" + assert result[0]["function"]["name"] == "web_search" + assert result[0]["function"]["parameters"]["type"] == "object" + + def test_multiple_tools(self): + tools = [ + {"name": "a", "description": "Tool A", "input_schema": {}}, + {"name": "b", "description": "Tool B", "input_schema": {}}, + ] + result = anthropic_tools_to_openai(tools) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_empty_list(self): + assert anthropic_tools_to_openai([]) == [] + + def test_pydantic_model_input(self): + tool = AnthropicTool( + name = "test", description = "desc", input_schema = {"type": "object"} + ) + result = anthropic_tools_to_openai([tool]) + assert result[0]["function"]["name"] == "test" + + +# ===================================================================== +# SSE event helper tests +# ===================================================================== + + +class TestBuildAnthropicSSEEvent: + def test_basic_event(self): + result = build_anthropic_sse_event("message_start", {"type": "message_start"}) + assert result.startswith("event: message_start\n") + assert "data: " in result + assert result.endswith("\n\n") + + def test_data_is_valid_json(self): + result = build_anthropic_sse_event("test", {"key": "value"}) + data_line = result.split("\n")[1] + payload = json.loads(data_line.removeprefix("data: ")) + assert payload == {"key": "value"} + + +# ===================================================================== +# Stream emitter tests +# ===================================================================== + + +class TestAnthropicStreamEmitter: + def test_start_emits_message_start_and_content_block_start(self): + e = AnthropicStreamEmitter() + events = e.start("msg_123", "test-model") + assert len(events) == 2 + assert "message_start" in events[0] + assert "content_block_start" in events[1] + assert '"type": "text"' in events[1] + + def test_content_delta_emits_text_delta(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + events = e.feed({"type": "content", "text": "Hello"}) + assert len(events) == 1 + parsed = json.loads(events[0].split("data: ")[1]) + assert parsed["delta"]["type"] == "text_delta" + assert parsed["delta"]["text"] == "Hello" + + def test_cumulative_content_diffs_correctly(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed({"type": "content", "text": "Hel"}) + events = e.feed({"type": "content", "text": "Hello"}) + parsed = json.loads(events[0].split("data: ")[1]) + assert parsed["delta"]["text"] == "lo" + + def test_empty_content_diff_no_event(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed({"type": "content", "text": "Hi"}) + events = e.feed({"type": "content", "text": "Hi"}) + assert events == [] + + def test_tool_start_closes_text_opens_tool_block(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed({"type": "content", "text": "Thinking"}) + events = e.feed( + { + "type": "tool_start", + "tool_name": "web_search", + "tool_call_id": "tc_1", + "arguments": {"query": "test"}, + } + ) + # content_block_stop + content_block_start(tool_use) + content_block_delta(input_json) + assert len(events) == 3 + assert "content_block_stop" in events[0] + assert "tool_use" in events[1] + assert "input_json_delta" in events[2] + + def test_tool_end_closes_tool_opens_new_text_block(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed( + { + "type": "tool_start", + "tool_name": "t", + "tool_call_id": "tc_1", + "arguments": {}, + } + ) + events = e.feed( + { + "type": "tool_end", + "tool_name": "t", + "tool_call_id": "tc_1", + "result": "done", + } + ) + # content_block_stop (tool) + tool_result + content_block_start (new text) + assert len(events) == 3 + assert "content_block_stop" in events[0] + assert "tool_result" in events[1] + parsed = json.loads(events[1].split("data: ")[1]) + assert parsed["content"] == "done" + assert parsed["tool_use_id"] == "tc_1" + assert "content_block_start" in events[2] + assert '"type": "text"' in events[2] + + def test_finish_emits_stop_events(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + events = e.finish("end_turn") + # content_block_stop + message_delta + message_stop + assert len(events) == 3 + assert "content_block_stop" in events[0] + assert "message_delta" in events[1] + assert "end_turn" in events[1] + assert "message_stop" in events[2] + + def test_metadata_captured_in_finish_usage(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed( + { + "type": "metadata", + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + } + ) + events = e.finish("end_turn") + delta_event = [ev for ev in events if "message_delta" in ev][0] + parsed = json.loads(delta_event.split("data: ")[1]) + assert parsed["usage"]["output_tokens"] == 20 + + def test_status_events_ignored(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + events = e.feed({"type": "status", "text": "Searching..."}) + assert events == [] + + def test_no_tool_calls_simple_text_flow(self): + e = AnthropicStreamEmitter() + start_events = e.start("msg_1", "m") + content_events = e.feed({"type": "content", "text": "Hello world"}) + meta_events = e.feed( + {"type": "metadata", "usage": {"prompt_tokens": 5, "completion_tokens": 2}} + ) + end_events = e.finish("end_turn") + + assert len(start_events) == 2 + assert len(content_events) == 1 + assert meta_events == [] + assert len(end_events) == 3 + + def test_block_index_increments(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + assert e.block_index == 0 + e.feed( + { + "type": "tool_start", + "tool_name": "t", + "tool_call_id": "tc_1", + "arguments": {}, + } + ) + assert e.block_index == 1 + e.feed( + { + "type": "tool_end", + "tool_name": "t", + "tool_call_id": "tc_1", + "result": "ok", + } + ) + assert e.block_index == 2 + + def test_text_after_tool_resets_prev_text(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed({"type": "content", "text": "Before tool"}) + e.feed( + { + "type": "tool_start", + "tool_name": "t", + "tool_call_id": "tc_1", + "arguments": {}, + } + ) + e.feed( + { + "type": "tool_end", + "tool_name": "t", + "tool_call_id": "tc_1", + "result": "ok", + } + ) + # After tool_end, prev_text should be reset + events = e.feed({"type": "content", "text": "After tool"}) + parsed = json.loads(events[0].split("data: ")[1]) + assert parsed["delta"]["text"] == "After tool" + + +# ===================================================================== +# Pass-through emitter tests (client-side tool execution path) +# ===================================================================== + + +class TestAnthropicPassthroughEmitter: + def _parse(self, event_str): + return json.loads(event_str.split("data: ")[1]) + + def test_start_emits_message_start_only(self): + e = AnthropicPassthroughEmitter() + events = e.start("msg_1", "test-model") + assert len(events) == 1 + assert "message_start" in events[0] + parsed = self._parse(events[0]) + assert parsed["message"]["id"] == "msg_1" + assert parsed["message"]["model"] == "test-model" + + def test_text_chunk_opens_text_block_and_emits_delta(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + chunk = {"choices": [{"delta": {"content": "Hello"}}]} + events = e.feed_chunk(chunk) + # content_block_start + content_block_delta + assert len(events) == 2 + assert "content_block_start" in events[0] + assert '"type": "text"' in events[0] + delta = self._parse(events[1]) + assert delta["delta"]["type"] == "text_delta" + assert delta["delta"]["text"] == "Hello" + + def test_sequential_text_chunks_single_block(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + events1 = e.feed_chunk({"choices": [{"delta": {"content": "Hello"}}]}) + events2 = e.feed_chunk({"choices": [{"delta": {"content": " world"}}]}) + # First chunk opens the block, second only emits delta + assert len(events1) == 2 + assert len(events2) == 1 + assert self._parse(events2[0])["delta"]["text"] == " world" + + def test_tool_call_opens_tool_use_block(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + chunk = { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "Bash", "arguments": ""}, + } + ] + } + } + ] + } + events = e.feed_chunk(chunk) + assert len(events) == 1 + parsed = self._parse(events[0]) + assert parsed["type"] == "content_block_start" + assert parsed["content_block"]["type"] == "tool_use" + assert parsed["content_block"]["id"] == "call_1" + assert parsed["content_block"]["name"] == "Bash" + + def test_tool_call_arguments_streamed_as_input_json_delta(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + # Open the tool call + e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "Bash", "arguments": ""}, + } + ] + } + } + ] + } + ) + # Stream argument fragments + events1 = e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + {"index": 0, "function": {"arguments": '{"cmd'}} + ] + } + } + ] + } + ) + events2 = e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + {"index": 0, "function": {"arguments": '": "ls"}'}} + ] + } + } + ] + } + ) + parsed1 = self._parse(events1[0]) + parsed2 = self._parse(events2[0]) + assert parsed1["delta"]["type"] == "input_json_delta" + assert parsed1["delta"]["partial_json"] == '{"cmd' + assert parsed2["delta"]["partial_json"] == '": "ls"}' + + def test_text_then_tool_closes_text_block(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk({"choices": [{"delta": {"content": "Let me check."}}]}) + events = e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "Bash", "arguments": ""}, + } + ] + } + } + ] + } + ) + # Should close text block and open tool_use block + assert "content_block_stop" in events[0] + assert "content_block_start" in events[1] + assert '"type": "tool_use"' in events[1] + + def test_finish_reason_tool_calls_sets_tool_use_stop(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "Bash", "arguments": "{}"}, + } + ] + } + } + ] + } + ) + e.feed_chunk({"choices": [{"delta": {}, "finish_reason": "tool_calls"}]}) + events = e.finish() + delta_event = [ev for ev in events if "message_delta" in ev][0] + parsed = self._parse(delta_event) + assert parsed["delta"]["stop_reason"] == "tool_use" + + def test_finish_reason_stop_sets_end_turn(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk({"choices": [{"delta": {"content": "Hi"}}]}) + e.feed_chunk({"choices": [{"delta": {}, "finish_reason": "stop"}]}) + events = e.finish() + delta_event = [ev for ev in events if "message_delta" in ev][0] + parsed = self._parse(delta_event) + assert parsed["delta"]["stop_reason"] == "end_turn" + + def test_finish_reason_length_sets_max_tokens(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk({"choices": [{"delta": {"content": "Hi"}}]}) + e.feed_chunk({"choices": [{"delta": {}, "finish_reason": "length"}]}) + events = e.finish() + delta_event = [ev for ev in events if "message_delta" in ev][0] + parsed = self._parse(delta_event) + assert parsed["delta"]["stop_reason"] == "max_tokens" + + def test_finish_closes_current_block(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk({"choices": [{"delta": {"content": "Hi"}}]}) + events = e.finish() + assert "content_block_stop" in events[0] + assert "message_delta" in events[1] + assert "message_stop" in events[2] + + def test_usage_chunk_captured(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk({"choices": [{"delta": {"content": "Hi"}}]}) + e.feed_chunk( + { + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 5}, + } + ) + events = e.finish() + delta_event = [ev for ev in events if "message_delta" in ev][0] + parsed = self._parse(delta_event) + assert parsed["usage"]["output_tokens"] == 5 + + def test_empty_chunk_returns_no_events(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + events = e.feed_chunk({"choices": []}) + assert events == [] + + def test_no_blocks_at_all_still_produces_valid_finish(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + events = e.finish() + # No content_block_stop because no block was opened + assert not any("content_block_stop" in ev for ev in events) + assert any("message_delta" in ev for ev in events) + assert any("message_stop" in ev for ev in events) + + def test_multiple_tool_calls_distinct_blocks(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + # First tool call + e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "Bash", "arguments": "{}"}, + } + ] + } + } + ] + } + ) + # Second tool call (different index) + events = e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 1, + "id": "c2", + "type": "function", + "function": {"name": "Read", "arguments": "{}"}, + } + ] + } + } + ] + } + ) + # Should close block 0, open block 1 + assert "content_block_stop" in events[0] + assert "content_block_start" in events[1] + parsed = self._parse(events[1]) + assert parsed["content_block"]["name"] == "Read" + assert parsed["content_block"]["id"] == "c2" diff --git a/studio/backend/tests/test_responses_api.py b/studio/backend/tests/test_responses_api.py new file mode 100644 index 0000000000..5b55f87259 --- /dev/null +++ b/studio/backend/tests/test_responses_api.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +""" +Tests for the OpenAI Responses API schemas and input normalisation. +These tests do NOT require a running server or GPU -- they validate +the Pydantic models and the _normalise_responses_input helper. +""" + +import sys +import os +import json +import re + +# Ensure backend is on path +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from models.inference import ( + ResponsesRequest, + ResponsesInputMessage, + ResponsesInputTextPart, + ResponsesInputImagePart, + ResponsesOutputTextContent, + ResponsesOutputMessage, + ResponsesUsage, + ResponsesResponse, + ChatMessage, + TextContentPart, + ImageContentPart, + ImageUrl, + ChatCompletionRequest, +) + + +# ── _normalise_responses_input: copied from routes/inference.py ── +# We cannot import routes.inference directly because routes/__init__.py +# pulls in heavy dependencies (structlog/twisted/torch). This is a +# direct copy of the function for testing purposes. + + +def _normalise_responses_input(payload: ResponsesRequest) -> list: + """Convert a ResponsesRequest into a list of ChatMessage for the completions backend.""" + messages = [] + + # System / developer instructions + if payload.instructions: + messages.append(ChatMessage(role = "system", content = payload.instructions)) + + # Simple string input + if isinstance(payload.input, str): + if payload.input: + messages.append(ChatMessage(role = "user", content = payload.input)) + return messages + + # List of ResponsesInputMessage + for msg in payload.input: + role = "system" if msg.role == "developer" else msg.role + + if isinstance(msg.content, str): + messages.append(ChatMessage(role = role, content = msg.content)) + else: + # Convert Responses content parts -> Chat content parts + parts = [] + for part in msg.content: + if isinstance(part, ResponsesInputTextPart): + parts.append(TextContentPart(type = "text", text = part.text)) + elif isinstance(part, ResponsesInputImagePart): + parts.append( + ImageContentPart( + type = "image_url", + image_url = ImageUrl(url = part.image_url, detail = part.detail), + ) + ) + messages.append(ChatMessage(role = role, content = parts if parts else "")) + + return messages + + +# ===================================================================== +# Schema validation tests +# ===================================================================== + + +class TestResponsesRequest: + """Validate ResponsesRequest accepts the shapes the OpenAI SDK sends.""" + + def test_minimal_string_input(self): + req = ResponsesRequest(input = "Hello") + assert req.input == "Hello" + assert req.stream is False + assert req.model == "default" + + def test_message_list_input(self): + req = ResponsesRequest( + input = [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello!"}, + ], + ) + assert len(req.input) == 2 + assert req.input[0].role == "user" + assert req.input[0].content == "Hi" + + def test_multimodal_input(self): + req = ResponsesRequest( + input = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "What is in this image?"}, + { + "type": "input_image", + "image_url": "https://example.com/img.png", + }, + ], + }, + ], + ) + parts = req.input[0].content + assert len(parts) == 2 + assert isinstance(parts[0], ResponsesInputTextPart) + assert isinstance(parts[1], ResponsesInputImagePart) + + def test_instructions_field(self): + req = ResponsesRequest( + input = "test", + instructions = "You are a helpful assistant.", + ) + assert req.instructions == "You are a helpful assistant." + + def test_extra_fields_accepted(self): + """OpenAI SDK may send fields we don't model -- extra='allow' should pass.""" + req = ResponsesRequest( + input = "test", + tools = [{"type": "web_search_preview"}], + store = True, + metadata = {"key": "value"}, + previous_response_id = "resp_abc123", + ) + assert req.tools == [{"type": "web_search_preview"}] + assert req.store is True + + def test_stream_flag(self): + req = ResponsesRequest(input = "test", stream = True) + assert req.stream is True + + def test_temperature_and_top_p(self): + req = ResponsesRequest(input = "test", temperature = 0.8, top_p = 0.9) + assert req.temperature == 0.8 + assert req.top_p == 0.9 + + def test_max_output_tokens(self): + req = ResponsesRequest(input = "test", max_output_tokens = 512) + assert req.max_output_tokens == 512 + + def test_developer_role(self): + req = ResponsesRequest( + input = [{"role": "developer", "content": "System instructions"}], + ) + assert req.input[0].role == "developer" + + +# ===================================================================== +# Response model tests +# ===================================================================== + + +class TestResponsesResponse: + """Validate response models serialise correctly.""" + + def test_basic_response(self): + resp = ResponsesResponse( + model = "test-model", + output = [ + ResponsesOutputMessage( + content = [ResponsesOutputTextContent(text = "Hello!")] + ), + ], + usage = ResponsesUsage(input_tokens = 10, output_tokens = 5, total_tokens = 15), + ) + d = resp.model_dump() + assert d["object"] == "response" + assert d["status"] == "completed" + assert d["output"][0]["type"] == "message" + assert d["output"][0]["content"][0]["type"] == "output_text" + assert d["output"][0]["content"][0]["text"] == "Hello!" + assert d["usage"]["input_tokens"] == 10 + assert d["usage"]["output_tokens"] == 5 + assert d["usage"]["total_tokens"] == 15 + # Must NOT have prompt_tokens / completion_tokens + assert "prompt_tokens" not in d["usage"] + assert "completion_tokens" not in d["usage"] + + def test_id_format(self): + resp = ResponsesResponse() + assert resp.id.startswith("resp_") + + def test_output_message_id_format(self): + msg = ResponsesOutputMessage() + assert msg.id.startswith("msg_") + + def test_annotations_default_empty(self): + part = ResponsesOutputTextContent(text = "hi") + assert part.annotations == [] + + def test_response_json_roundtrip(self): + resp = ResponsesResponse( + model = "gpt-4", + output = [ + ResponsesOutputMessage( + content = [ResponsesOutputTextContent(text = "ok")], + ), + ], + usage = ResponsesUsage(input_tokens = 1, output_tokens = 1, total_tokens = 2), + ) + j = json.loads(resp.model_dump_json()) + assert j["object"] == "response" + assert j["output"][0]["role"] == "assistant" + assert j["output"][0]["status"] == "completed" + + +# ===================================================================== +# Input normalisation tests +# ===================================================================== + + +class TestNormaliseResponsesInput: + """Test _normalise_responses_input converts Responses input to ChatMessages.""" + + def test_string_input(self): + payload = ResponsesRequest(input = "Hello world") + msgs = _normalise_responses_input(payload) + assert len(msgs) == 1 + assert msgs[0].role == "user" + assert msgs[0].content == "Hello world" + + def test_instructions_become_system_message(self): + payload = ResponsesRequest( + input = "Hi", + instructions = "Be concise.", + ) + msgs = _normalise_responses_input(payload) + assert len(msgs) == 2 + assert msgs[0].role == "system" + assert msgs[0].content == "Be concise." + assert msgs[1].role == "user" + assert msgs[1].content == "Hi" + + def test_message_list(self): + payload = ResponsesRequest( + input = [ + {"role": "user", "content": "First"}, + {"role": "assistant", "content": "Response"}, + {"role": "user", "content": "Second"}, + ], + ) + msgs = _normalise_responses_input(payload) + assert len(msgs) == 3 + assert msgs[0].role == "user" + assert msgs[1].role == "assistant" + assert msgs[2].role == "user" + + def test_developer_role_maps_to_system(self): + payload = ResponsesRequest( + input = [{"role": "developer", "content": "Instructions"}], + ) + msgs = _normalise_responses_input(payload) + assert msgs[0].role == "system" + assert msgs[0].content == "Instructions" + + def test_multimodal_parts(self): + payload = ResponsesRequest( + input = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Describe this:"}, + { + "type": "input_image", + "image_url": "data:image/png;base64,abc", + }, + ], + }, + ], + ) + msgs = _normalise_responses_input(payload) + assert len(msgs) == 1 + content = msgs[0].content + assert isinstance(content, list) + assert len(content) == 2 + assert isinstance(content[0], TextContentPart) + assert content[0].text == "Describe this:" + assert isinstance(content[1], ImageContentPart) + assert content[1].image_url.url == "data:image/png;base64,abc" + + def test_empty_string_input(self): + payload = ResponsesRequest(input = "") + msgs = _normalise_responses_input(payload) + assert len(msgs) == 0 + + def test_empty_list_input(self): + payload = ResponsesRequest(input = []) + msgs = _normalise_responses_input(payload) + assert len(msgs) == 0 + + def test_instructions_only(self): + payload = ResponsesRequest(input = "", instructions = "System msg") + msgs = _normalise_responses_input(payload) + assert len(msgs) == 1 + assert msgs[0].role == "system" + + def test_instructions_plus_message_list(self): + payload = ResponsesRequest( + input = [{"role": "user", "content": "Hello"}], + instructions = "Be brief.", + ) + msgs = _normalise_responses_input(payload) + assert len(msgs) == 2 + assert msgs[0].role == "system" + assert msgs[0].content == "Be brief." + assert msgs[1].role == "user" + + +if __name__ == "__main__": + import pytest + + pytest.main([__file__, "-v"]) diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py new file mode 100644 index 0000000000..9cc17c89fb --- /dev/null +++ b/studio/backend/tests/test_studio_api.py @@ -0,0 +1,643 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +End-to-end tests for Unsloth Studio's HTTP API surface. + +Covers the OpenAI-compatible and Anthropic-compatible endpoints exposed +by the server that ``unsloth studio run`` boots, plus API key +authentication and the CLI's ``--help`` output: + + 1. curl -- basic chat completions (non-streaming) + 2. curl -- streaming chat completions + 3. Python OpenAI SDK -- streaming completions + 4. curl -- with tools (web_search + python) + 5. Anthropic Messages API -- basic non-streaming + 6. Anthropic Messages API -- streaming SSE + 7. Anthropic Python SDK -- non-streaming + 8. Anthropic Messages API -- streaming with tools + +Training, export, fine-tuning, and chat-UI concerns are out of scope — +see the unit suites elsewhere under ``studio/backend/tests/`` for those. + +Usage: + + # Script mode — launches its own server via ``unsloth studio run``. + python tests/test_studio_api.py + python tests/test_studio_api.py --model unsloth/... --gguf-variant ... + + # Pytest mode, external server — start a Studio server yourself, + # then point pytest at it. Fastest iteration loop. + unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL & + export UNSLOTH_E2E_BASE_URL=http://127.0.0.1:8080 + export UNSLOTH_E2E_API_KEY=sk-unsloth-... # from the server banner + pytest tests/test_studio_api.py -v + + # Pytest mode, fixture-managed server — pytest launches and tears + # down the server itself. One-shot verification, CI-friendly. + pytest tests/test_studio_api.py -v \\ + --unsloth-model unsloth/Qwen3-1.7B-GGUF \\ + --unsloth-gguf-variant UD-Q4_K_XL + +The ``base_url`` / ``api_key`` parameters on the test functions resolve +via the ``studio_server`` session fixture in ``conftest.py``. + +Requires a GPU and ~2 GB of disk for the GGUF download. +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import signal +import subprocess +import sys +import time +import urllib.error +import urllib.request +from pathlib import Path + + +# ── Configuration ──────────────────────────────────────────────────── + +DEFAULT_MODEL = "unsloth/Qwen3-1.7B-GGUF" +DEFAULT_VARIANT = "UD-Q4_K_XL" +PORT = 18222 # high port unlikely to collide +HOST = "127.0.0.1" +STARTUP_TIMEOUT = 120 # seconds to wait for banner +LOG_FILE = ( + Path(__file__).resolve().parent.parent.parent.parent + / "temp" + / "test_studio_api.log" +) + + +# ── Helpers ────────────────────────────────────────────────────────── + + +def _http( + method: str, + url: str, + *, + body: dict | None = None, + headers: dict | None = None, + timeout: int = 60, +) -> tuple[int, str]: + """Minimal stdlib HTTP helper. Returns (status_code, body_text).""" + data = json.dumps(body).encode() if body else None + req = urllib.request.Request(url, data = data, headers = headers or {}, method = method) + if body: + req.add_header("Content-Type", "application/json") + try: + with urllib.request.urlopen(req, timeout = timeout) as resp: + return resp.status, resp.read().decode() + except urllib.error.HTTPError as exc: + return exc.code, exc.read().decode(errors = "replace") + + +def _stream_http( + url: str, + *, + body: dict, + headers: dict, + timeout: int = 60, +) -> tuple[int, list[dict]]: + """POST a streaming request and collect SSE chunks.""" + data = json.dumps(body).encode() + req = urllib.request.Request(url, data = data, headers = headers, method = "POST") + req.add_header("Content-Type", "application/json") + chunks: list[dict] = [] + try: + with urllib.request.urlopen(req, timeout = timeout) as resp: + status = resp.status + for raw_line in resp: + line = raw_line.decode().strip() + if line.startswith("data: ") and line != "data: [DONE]": + try: + chunks.append(json.loads(line[6:])) + except json.JSONDecodeError: + pass + return status, chunks + except urllib.error.HTTPError as exc: + return exc.code, [] + + +# ── Test functions ─────────────────────────────────────────────────── + + +def test_help_output(): + """``unsloth studio run --help`` should show all documented options.""" + result = subprocess.run( + ["unsloth", "studio", "run", "--help"], + capture_output = True, + text = True, + timeout = 15, + ) + out = result.stdout + assert result.returncode == 0, f"--help exited with {result.returncode}" + + for flag in [ + "--model", + "--gguf-variant", + "--max-seq-length", + "--load-in-4bit", + "--api-key-name", + "--port", + "--host", + "--frontend", + "--silent", + ]: + assert flag in out, f"Missing flag {flag!r} in --help output" + print(" PASS --help shows all flags") + + +def test_curl_basic(base_url: str, api_key: str): + """Example 1: basic non-streaming chat completion via HTTP.""" + status, text = _http( + "POST", + f"{base_url}/v1/chat/completions", + body = { + "messages": [{"role": "user", "content": "Say just the word hello"}], + "stream": False, + }, + headers = {"Authorization": f"Bearer {api_key}"}, + ) + assert status == 200, f"Expected 200, got {status}: {text[:300]}" + data = json.loads(text) + assert "choices" in data, f"Missing 'choices' in response: {text[:300]}" + content = data["choices"][0]["message"]["content"] + assert len(content) > 0, "Empty assistant content" + print(f" PASS curl basic: {content[:80]!r}") + + +def _collect_streamed_content(chunks: list[dict]) -> str: + """Extract text from SSE chunks, skipping role-only and usage chunks.""" + parts = [] + for c in chunks: + choices = c.get("choices", []) + if not choices: + continue + delta = choices[0].get("delta", {}) + part = delta.get("content") + if part: + parts.append(part) + return "".join(parts) + + +def test_curl_streaming(base_url: str, api_key: str): + """Example 2: streaming chat completion via HTTP SSE.""" + status, chunks = _stream_http( + f"{base_url}/v1/chat/completions", + body = { + "messages": [{"role": "user", "content": "Count from 1 to 3"}], + "stream": True, + }, + headers = {"Authorization": f"Bearer {api_key}"}, + ) + assert status == 200, f"Expected 200, got {status}" + assert len(chunks) > 0, "No SSE chunks received" + full = _collect_streamed_content(chunks) + assert len(full) > 0, "Streamed content is empty" + print(f" PASS curl streaming: got {len(chunks)} chunks, {len(full)} chars") + + +def test_openai_sdk(base_url: str, api_key: str): + """Example 3: OpenAI Python SDK streaming completion.""" + try: + from openai import OpenAI + except ImportError: + print(" SKIP openai SDK not installed") + return + + client = OpenAI(base_url = f"{base_url}/v1", api_key = api_key) + response = client.chat.completions.create( + model = "current", + messages = [ + {"role": "user", "content": "What is 2+2? Answer with just the number."} + ], + stream = True, + ) + content_parts = [] + for chunk in response: + if not chunk.choices: + continue + delta_content = chunk.choices[0].delta.content + if delta_content: + content_parts.append(delta_content) + full = "".join(content_parts) + assert len(full) > 0, "OpenAI SDK returned empty content" + print(f" PASS OpenAI SDK streaming: {full.strip()[:80]!r}") + + +def test_curl_with_tools(base_url: str, api_key: str): + """Example 4: chat completion with tool calling enabled. + + Note: when ``enable_tools`` is set the server always returns SSE + streaming regardless of the ``stream`` flag, so we parse SSE chunks. + The model may or may not produce visible content -- tool orchestration + can intercept the response -- so we only assert the endpoint succeeds. + """ + status, chunks = _stream_http( + f"{base_url}/v1/chat/completions", + body = { + "messages": [ + { + "role": "user", + "content": "What is 123 * 456? Use code to compute it.", + } + ], + "stream": True, + "enable_tools": True, + "enabled_tools": ["python"], + "session_id": "test-session", + }, + headers = {"Authorization": f"Bearer {api_key}"}, + timeout = 120, + ) + assert status == 200, f"Expected 200, got {status}" + assert len(chunks) > 0, "No SSE chunks received for tools request" + + # Check that at least one chunk has the expected shape + has_valid_chunk = any("choices" in c or "type" in c for c in chunks) + assert has_valid_chunk, "No valid chunks in tools response" + full = _collect_streamed_content(chunks) + print(f" PASS curl with tools: {len(chunks)} chunks, {len(full)} chars content") + + +def test_invalid_key_rejected(base_url: str): + """Requests with a bad API key should be rejected.""" + status, _text = _http( + "POST", + f"{base_url}/v1/chat/completions", + body = { + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + }, + headers = {"Authorization": "Bearer sk-unsloth-boguskey123"}, + ) + assert status == 401, f"Expected 401 for invalid key, got {status}" + print(" PASS invalid API key rejected (401)") + + +def test_no_key_rejected(base_url: str): + """Requests without any auth header should be rejected.""" + status, _text = _http( + "POST", + f"{base_url}/v1/chat/completions", + body = { + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + }, + ) + assert status == 401 or status == 403, f"Expected 401/403 for no key, got {status}" + print(f" PASS no API key rejected ({status})") + + +# ── Anthropic SSE helper ───────────────────────────────────────────── + + +def _stream_anthropic_http( + url: str, + *, + body: dict, + headers: dict, + timeout: int = 60, +) -> tuple[int, list[tuple[str, dict]]]: + """POST a streaming request and collect Anthropic SSE events. + + Returns (status, [(event_type, data_dict), ...]). + """ + data = json.dumps(body).encode() + req = urllib.request.Request(url, data = data, headers = headers, method = "POST") + req.add_header("Content-Type", "application/json") + events: list[tuple[str, dict]] = [] + try: + with urllib.request.urlopen(req, timeout = timeout) as resp: + status = resp.status + current_event = None + for raw_line in resp: + line = raw_line.decode().strip() + if line.startswith("event: "): + current_event = line[7:] + elif line.startswith("data: ") and current_event: + try: + events.append((current_event, json.loads(line[6:]))) + except json.JSONDecodeError: + pass + current_event = None + return status, events + except urllib.error.HTTPError as exc: + return exc.code, [] + + +def _collect_anthropic_text(events: list[tuple[str, dict]]) -> str: + """Extract text content from Anthropic SSE events.""" + parts = [] + for etype, data in events: + if etype == "content_block_delta": + delta = data.get("delta", {}) + if delta.get("type") == "text_delta": + parts.append(delta.get("text", "")) + return "".join(parts) + + +# ── Anthropic /v1/messages test functions ──────────────────────────── + + +def test_anthropic_basic(base_url: str, api_key: str): + """Anthropic Messages API: non-streaming.""" + status, text = _http( + "POST", + f"{base_url}/v1/messages", + body = { + "model": "default", + "max_tokens": 100, + "messages": [{"role": "user", "content": "Say just the word hello"}], + }, + headers = {"Authorization": f"Bearer {api_key}"}, + ) + assert status == 200, f"Expected 200, got {status}: {text[:300]}" + data = json.loads(text) + assert data.get("type") == "message", f"Expected type 'message': {text[:300]}" + assert data.get("role") == "assistant" + content = data.get("content", []) + assert len(content) > 0, "Empty content array" + text_block = content[-1] + assert text_block.get("type") == "text", f"Expected text block: {text_block}" + assert len(text_block.get("text", "")) > 0, "Empty text in response" + print(f" PASS anthropic basic: {text_block['text'][:80]!r}") + + +def test_anthropic_streaming(base_url: str, api_key: str): + """Anthropic Messages API: streaming SSE.""" + status, events = _stream_anthropic_http( + f"{base_url}/v1/messages", + body = { + "model": "default", + "max_tokens": 100, + "messages": [{"role": "user", "content": "Count from 1 to 3"}], + "stream": True, + }, + headers = {"Authorization": f"Bearer {api_key}"}, + ) + assert status == 200, f"Expected 200, got {status}" + assert len(events) > 0, "No SSE events received" + + event_types = [e[0] for e in events] + assert "message_start" in event_types, "Missing message_start event" + assert "message_stop" in event_types, "Missing message_stop event" + + full = _collect_anthropic_text(events) + assert len(full) > 0, "Streamed text content is empty" + print(f" PASS anthropic streaming: {len(events)} events, {len(full)} chars") + + +def test_anthropic_sdk(base_url: str, api_key: str): + """Anthropic Python SDK: non-streaming.""" + try: + from anthropic import Anthropic + except ImportError: + print(" SKIP anthropic SDK not installed") + return + + client = Anthropic(base_url = f"{base_url}/v1", api_key = api_key) + message = client.messages.create( + model = "default", + max_tokens = 100, + messages = [ + {"role": "user", "content": "What is 2+2? Answer with just the number."} + ], + ) + assert message.role == "assistant" + assert len(message.content) > 0, "Empty content" + text = message.content[0].text + assert len(text) > 0, "Empty text" + print(f" PASS Anthropic SDK: {text.strip()[:80]!r}") + + +def test_anthropic_with_tools(base_url: str, api_key: str): + """Anthropic Messages API: streaming with tools.""" + status, events = _stream_anthropic_http( + f"{base_url}/v1/messages", + body = { + "model": "default", + "max_tokens": 1024, + "messages": [ + { + "role": "user", + "content": "What is 123 * 456? Use code to compute it.", + } + ], + "tools": [ + { + "name": "python", + "description": "Execute Python code in a sandbox and return stdout/stderr.", + "input_schema": { + "type": "object", + "properties": { + "code": { + "type": "string", + "description": "The Python code to run", + }, + }, + "required": ["code"], + }, + } + ], + "stream": True, + }, + headers = {"Authorization": f"Bearer {api_key}"}, + timeout = 120, + ) + assert status == 200, f"Expected 200, got {status}" + assert len(events) > 0, "No SSE events received for tools request" + + event_types = [e[0] for e in events] + assert "message_start" in event_types, "Missing message_start" + assert "message_stop" in event_types, "Missing message_stop" + + full = _collect_anthropic_text(events) + print( + f" PASS anthropic with tools: {len(events)} events, {len(full)} chars content" + ) + + +# ── Server lifecycle ───────────────────────────────────────────────── + + +def _start_server(model: str, variant: str | None) -> tuple[subprocess.Popen, str]: + """Launch ``unsloth studio run`` and parse the API key from its banner. + + Returns (process, api_key). + """ + cmd = [ + "unsloth", + "studio", + "run", + "--model", + model, + "--port", + str(PORT), + "--host", + HOST, + "--api-key-name", + "test", + ] + if variant: + cmd.extend(["--gguf-variant", variant]) + + LOG_FILE.parent.mkdir(parents = True, exist_ok = True) + log_fh = open(LOG_FILE, "w") + proc = subprocess.Popen( + cmd, + stdout = log_fh, + stderr = subprocess.STDOUT, + preexec_fn = os.setsid, + ) + + # Wait for the banner containing the API key + api_key = None + deadline = time.monotonic() + STARTUP_TIMEOUT + while time.monotonic() < deadline: + time.sleep(2) + if proc.poll() is not None: + log_fh.flush() + log_text = LOG_FILE.read_text() + raise RuntimeError( + f"Server exited early (code {proc.returncode}):\n{log_text[-2000:]}" + ) + log_text = LOG_FILE.read_text() + m = re.search(r"API Key:\s+(sk-unsloth-[a-f0-9]+)", log_text) + if m: + api_key = m.group(1) + break + + if not api_key: + log_text = LOG_FILE.read_text() + _kill_server(proc) + raise RuntimeError( + f"Timed out waiting for API key in server output:\n{log_text[-2000:]}" + ) + + # Wait a moment for the model to be fully loaded + time.sleep(2) + return proc, api_key + + +def _kill_server(proc: subprocess.Popen): + """Send SIGTERM to the process group and wait for cleanup.""" + try: + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) + except (ProcessLookupError, PermissionError): + pass + try: + proc.wait(timeout = 10) + except subprocess.TimeoutExpired: + try: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + except (ProcessLookupError, PermissionError): + pass + proc.wait(timeout = 5) + + +# ── Main ───────────────────────────────────────────────────────────── + + +def main(): + parser = argparse.ArgumentParser( + description = "End-to-end tests for unsloth studio run" + ) + parser.add_argument( + "--model", + default = DEFAULT_MODEL, + help = f"Model to test with (default: {DEFAULT_MODEL})", + ) + parser.add_argument( + "--gguf-variant", + default = DEFAULT_VARIANT, + help = f"GGUF variant (default: {DEFAULT_VARIANT})", + ) + args = parser.parse_args() + + passed = 0 + failed = 0 + skipped = 0 + + def run_test(fn, *a, **kw): + nonlocal passed, failed, skipped + try: + fn(*a, **kw) + passed += 1 + except AssertionError as exc: + failed += 1 + print(f" FAIL {fn.__name__}: {exc}") + except Exception as exc: + failed += 1 + print(f" ERROR {fn.__name__}: {type(exc).__name__}: {exc}") + + # ── 1. Test --help (no server needed) ──────────────────────────── + print("\n[1/11] Testing --help output") + run_test(test_help_output) + + # ── 2-11. Start server and run API tests ───────────────────────── + print( + f"\nStarting server: {args.model} (variant={args.gguf_variant}) on port {PORT}..." + ) + proc = None + try: + proc, api_key = _start_server(args.model, args.gguf_variant) + base_url = f"http://{HOST}:{PORT}" + print(f"Server ready. API Key: {api_key[:20]}...\n") + + print("[2/11] Testing curl basic (non-streaming)") + run_test(test_curl_basic, base_url, api_key) + + print("[3/11] Testing curl streaming") + run_test(test_curl_streaming, base_url, api_key) + + print("[4/11] Testing OpenAI Python SDK (streaming)") + run_test(test_openai_sdk, base_url, api_key) + + print("[5/11] Testing curl with tools") + run_test(test_curl_with_tools, base_url, api_key) + + print("[6/11] Testing invalid API key rejection") + run_test(test_invalid_key_rejected, base_url) + + print("[7/11] Testing no API key rejection") + run_test(test_no_key_rejected, base_url) + + print("[8/11] Testing Anthropic basic (non-streaming)") + run_test(test_anthropic_basic, base_url, api_key) + + print("[9/11] Testing Anthropic streaming") + run_test(test_anthropic_streaming, base_url, api_key) + + print("[10/11] Testing Anthropic Python SDK") + run_test(test_anthropic_sdk, base_url, api_key) + + print("[11/11] Testing Anthropic with tools") + run_test(test_anthropic_with_tools, base_url, api_key) + + except RuntimeError as exc: + print(f"\nFATAL: Server failed to start: {exc}") + failed += 11 # count remaining tests as failed + finally: + if proc: + print("\nStopping server...") + _kill_server(proc) + print("Server stopped.") + + # ── Summary ────────────────────────────────────────────────────── + total = passed + failed + print(f"\n{'=' * 40}") + print(f"Results: {passed}/{total} passed, {failed} failed") + print(f"Log: {LOG_FILE}") + print(f"{'=' * 40}") + sys.exit(1 if failed else 0) + + +if __name__ == "__main__": + main() diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index 13ff8a5cbe..d507929758 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -13,6 +13,7 @@ import { Route as loginRoute } from "./routes/login"; import { Route as onboardingRoute } from "./routes/onboarding"; import { Route as changePasswordRoute } from "./routes/change-password"; import { Route as studioRoute } from "./routes/studio"; +import { Route as apiKeysRoute } from "./routes/api-keys"; const routeTree = rootRoute.addChildren([ indexRoute, @@ -25,6 +26,7 @@ const routeTree = rootRoute.addChildren([ exportRoute, dataRecipesRoute, dataRecipeRoute, + apiKeysRoute, ]); export const router = createRouter({ routeTree }); diff --git a/studio/frontend/src/app/routes/api-keys.tsx b/studio/frontend/src/app/routes/api-keys.tsx new file mode 100644 index 0000000000..5846690d7b --- /dev/null +++ b/studio/frontend/src/app/routes/api-keys.tsx @@ -0,0 +1,18 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { createRoute } from "@tanstack/react-router"; +import { lazy } from "react"; +import { requireAuth } from "../auth-guards"; +import { Route as rootRoute } from "./__root"; + +const ApiKeysPage = lazy(() => + import("@/features/auth/api-keys-page").then((m) => ({ default: m.ApiKeysPage })), +); + +export const Route = createRoute({ + getParentRoute: () => rootRoute, + path: "/api-keys", + beforeLoad: () => requireAuth(), + component: ApiKeysPage, +}); diff --git a/studio/frontend/src/components/navbar.tsx b/studio/frontend/src/components/navbar.tsx index 121c559db8..2d310a81c3 100644 --- a/studio/frontend/src/components/navbar.tsx +++ b/studio/frontend/src/components/navbar.tsx @@ -29,6 +29,7 @@ import { ChefHatIcon, Copy01Icon, CursorInfo02Icon, + Key01Icon, PackageIcon, Tick02Icon, ZapIcon, @@ -416,6 +417,20 @@ export function Navbar() { +
+ + + API Keys + +
{tourId ? (
+ ); +} + +function RevealKeyDialog({ + open, + rawKey, + onClose, +}: { + open: boolean; + rawKey: string; + onClose: () => void; +}) { + return ( + !o && onClose()}> + + + API Key Created + + Copy this key now. It will not be shown again. + + +
+ + {rawKey} + + +
+
+ +

+ Store this key securely. You will not be able to see it again after closing this dialog. +

+
+ + + + + + ); +} + +function CreateKeyForm({ onCreated }: { onCreated: (rawKey: string) => void }) { + const [name, setName] = useState(""); + const [expiresInDays, setExpiresInDays] = useState(""); + const [loading, setLoading] = useState(false); + + const handleSubmit = async (e: React.FormEvent) => { + e.preventDefault(); + if (!name.trim()) return; + setLoading(true); + try { + const days = expiresInDays ? parseInt(expiresInDays, 10) : null; + const result = await createApiKey(name.trim(), days); + onCreated(result.key); + setName(""); + setExpiresInDays(""); + } finally { + setLoading(false); + } + }; + + return ( +
+
+ + setName(e.target.value)} + required + /> +
+
+ + setExpiresInDays(e.target.value)} + /> +
+ +
+ ); +} + +function KeysTable({ + keys, + onRevoke, +}: { + keys: ApiKey[]; + onRevoke: (id: number) => void; +}) { + if (keys.length === 0) { + return ( +

+ No API keys yet. Create one above. +

+ ); + } + + return ( +
+ + + + + + + + + + + + {keys.map((k) => ( + + + + + + + + + ))} + +
NameKeyCreatedLast usedExpires +
{k.name} + + sk-unsloth-{k.key_prefix}... + + {formatDate(k.created_at)}{formatDate(k.last_used_at)}{formatDate(k.expires_at)} + {k.is_active ? ( + + ) : ( + Revoked + )} +
+
+ ); +} + +function UsageExamples() { + const base = window.location.origin; + + const curlExample = `curl ${base}/v1/chat/completions \\ + -H "Authorization: Bearer sk-unsloth-YOUR_KEY" \\ + -H "Content-Type: application/json" \\ + -d '{ + "messages": [{"role": "user", "content": "Hello"}], + "stream": true + }'`; + + const pythonExample = `from openai import OpenAI + +client = OpenAI( + base_url="${base}/v1", + api_key="sk-unsloth-YOUR_KEY", +) + +response = client.chat.completions.create( + model="current", + messages=[{"role": "user", "content": "Hello"}], + stream=True, +) +for chunk in response: + print(chunk.choices[0].delta.content or "", end="")`; + + const toolsExample = `curl ${base}/v1/chat/completions \\ + -H "Authorization: Bearer sk-unsloth-YOUR_KEY" \\ + -H "Content-Type: application/json" \\ + -d '{ + "messages": [{"role": "user", "content": "Search for Python 3.13 features"}], + "stream": true, + "enable_tools": true, + "enabled_tools": ["web_search", "python"], + "session_id": "my-session" + }'`; + + return ( +
+

Usage examples

+
+
+

curl

+ 
+            {curlExample}
+
+
+
+

Python (OpenAI SDK)

+ 
+            {pythonExample}
+
+
+
+

With tools (web search + code execution)

+ 
+            {toolsExample}
+
+
+
+
+ ); +} + +// --------------------------------------------------------------------------- +// Page +// --------------------------------------------------------------------------- + +export function ApiKeysPage() { + const [keys, setKeys] = useState([]); + const [revealedKey, setRevealedKey] = useState(null); + const [error, setError] = useState(null); + + const loadKeys = useCallback(async () => { + try { + setError(null); + const loaded = await fetchApiKeys(); + setKeys(loaded); + } catch (e) { + setError(e instanceof Error ? e.message : "Failed to load API keys"); + } + }, []); + + useEffect(() => { + void loadKeys(); + }, [loadKeys]); + + const handleCreated = (rawKey: string) => { + setRevealedKey(rawKey); + void loadKeys(); + }; + + const handleRevoke = async (keyId: number) => { + try { + await revokeApiKey(keyId); + void loadKeys(); + } catch { + setError("Failed to revoke key"); + } + }; + + return ( + +
+
+
+ +
+
+

API Keys

+

+ Create keys to access Unsloth Studio programmatically via the OpenAI-compatible API. +

+
+
+ + {error && ( +
+ + {error} +
+ )} + + + + +
+ + setRevealedKey(null)} + /> + + ); +} diff --git a/studio/frontend/src/features/auth/index.ts b/studio/frontend/src/features/auth/index.ts index 75db92432c..9962f6d431 100644 --- a/studio/frontend/src/features/auth/index.ts +++ b/studio/frontend/src/features/auth/index.ts @@ -1,6 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +export { ApiKeysPage } from "./api-keys-page"; export { LoginPage } from "./login-page"; export { ChangePasswordPage } from "./change-password-page"; export { authFetch, refreshSession } from "./api"; diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 2fecb9d6b1..a3c0840be1 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -69,6 +69,80 @@ def _find_setup_script() -> Optional[Path]: return None +# ── helpers for `unsloth studio run` ──────────────────────────────── + + +def _wait_for_server(port: int, timeout: int = 30) -> bool: + """Poll ``GET /api/health`` until the server responds 200 or *timeout* expires.""" + import urllib.request + import urllib.error + + url = f"http://127.0.0.1:{port}/api/health" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + with urllib.request.urlopen(url, timeout = 2) as resp: + if resp.status == 200: + return True + except (urllib.error.URLError, OSError, ConnectionError): + pass + time.sleep(0.5) + return False + + +def _create_api_key_inprocess(name: str) -> str: + """Create an API key via direct storage call (no HTTP needed). + + Bypasses the ``must_change_password`` gate that blocks HTTP + ``POST /api/auth/api-keys`` on fresh installs. Safe because the + CLI already has filesystem access to ``~/.unsloth/studio``. + """ + from auth.storage import create_api_key, DEFAULT_ADMIN_USERNAME + + raw_key, _row = create_api_key(username = DEFAULT_ADMIN_USERNAME, name = name) + return raw_key + + +def _load_model_via_http( + port: int, + api_key: str, + model: str, + gguf_variant: Optional[str], + max_seq_length: int, + load_in_4bit: bool, + timeout: int = 600, +) -> dict: + """POST to ``/api/inference/load`` using the API key for auth.""" + import json + import urllib.request + import urllib.error + + payload: dict = { + "model_path": model, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + } + if gguf_variant: + payload["gguf_variant"] = gguf_variant + + data = json.dumps(payload).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{port}/api/inference/load", + data = data, + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + }, + method = "POST", + ) + try: + with urllib.request.urlopen(req, timeout = timeout) as resp: + return json.loads(resp.read()) + except urllib.error.HTTPError as exc: + body = exc.read().decode(errors = "replace") + raise RuntimeError(f"Model load failed (HTTP {exc.code}): {body}") from exc + + # ── unsloth studio (server) ────────────────────────────────────────── @@ -166,6 +240,174 @@ def studio_default( typer.echo("\nShutting down...") +# ── unsloth studio run ─────────────────────────────────────────────── + + +@studio_app.command() +def run( + model: str = typer.Option(..., "--model", "-m", help = "Model path or HF repo"), + gguf_variant: Optional[str] = typer.Option( + None, "--gguf-variant", help = "GGUF quant variant (e.g. UD-Q4_K_XL)" + ), + max_seq_length: int = typer.Option( + 0, "--max-seq-length", help = "Max sequence length (0 = model default)" + ), + load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"), + api_key_name: str = typer.Option( + "cli", "--api-key-name", help = "Label for the auto-generated API key" + ), + port: int = typer.Option(8888, "--port", "-p"), + host: str = typer.Option("0.0.0.0", "--host", "-H"), + frontend: Optional[Path] = typer.Option(None, "--frontend", "-f"), + silent: bool = typer.Option(False, "--silent", "-q"), +): + """Start Studio, load a model, and print an API key -- one-liner server. + + Example: + unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL + """ + # ── 1. Venv re-exec (same pattern as studio_default) ────────────── + studio_venv_dir = STUDIO_HOME / "unsloth_studio" + in_studio_venv = sys.prefix.startswith(str(studio_venv_dir)) + + if not in_studio_venv: + studio_python = _studio_venv_python() + if not studio_python: + typer.echo("Studio not set up. Run install.sh first.") + raise typer.Exit(1) + # Re-exec into the studio venv via its `unsloth` entry point + studio_bin = studio_python.parent / "unsloth" + if not studio_bin.is_file(): + typer.echo( + "Studio venv missing 'unsloth' entry point. Re-run: unsloth studio setup" + ) + raise typer.Exit(1) + args = [ + str(studio_bin), + "studio", + "run", + "--model", + model, + "--max-seq-length", + str(max_seq_length), + "--api-key-name", + api_key_name, + "--port", + str(port), + "--host", + host, + ] + if gguf_variant: + args.extend(["--gguf-variant", gguf_variant]) + if not load_in_4bit: + args.append("--no-load-in-4bit") + if frontend: + args.extend(["--frontend", str(frontend)]) + if silent: + args.append("--silent") + + if sys.platform == "win32": + proc = subprocess.Popen(args) + try: + rc = proc.wait() + except KeyboardInterrupt: + rc = proc.wait() + raise typer.Exit(rc) + else: + os.execvp(str(studio_bin), args) + + # ── 2. Start server (always suppress built-in banner) ───────────── + from studio.backend.run import run_server, _resolve_external_ip + + run_kwargs = dict(host = host, port = port, silent = True, llama_parallel_slots = 4) + if frontend is not None: + run_kwargs["frontend_path"] = frontend + app = run_server(**run_kwargs) + actual_port = getattr(app.state, "server_port", port) or port + + # ── 3. Wait for server health ───────────────────────────────────── + if not silent: + typer.echo("Starting Unsloth Studio...") + if not _wait_for_server(actual_port): + typer.echo("Error: server did not become healthy within 30 seconds.", err = True) + raise typer.Exit(1) + + # ── 4. Create API key in-process ────────────────────────────────── + api_key = _create_api_key_inprocess(api_key_name) + + # ── 5. Load model via HTTP ──────────────────────────────────────── + if not silent: + typer.echo(f"Loading model: {model}...") + try: + result = _load_model_via_http( + port = actual_port, + api_key = api_key, + model = model, + gguf_variant = gguf_variant, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + ) + except RuntimeError as exc: + typer.echo(f"Error: {exc}", err = True) + raise typer.Exit(1) + + loaded_model = result.get("model", model) + display_variant = f" ({gguf_variant})" if gguf_variant else "" + + # ── 6. Print banner ─────────────────────────────────────────────── + display_host = _resolve_external_ip() if host == "0.0.0.0" else host + base_url = f"http://{display_host}:{actual_port}" + sdk_base_url = f"{base_url}/v1" + + if not silent: + typer.echo("") + typer.echo("=" * 56) + typer.echo(f" Unsloth Studio running at {base_url}") + typer.echo(f" Model loaded: {loaded_model}{display_variant}") + typer.echo(f" API Key: {api_key}") + typer.echo("") + typer.echo(" OpenAI / Anthropic SDK base URL:") + typer.echo(f" {sdk_base_url}") + typer.echo("=" * 56) + typer.echo("") + typer.echo("OpenAI Chat Completions:") + typer.echo(f" curl {sdk_base_url}/chat/completions \\") + typer.echo(f' -H "Authorization: Bearer {api_key}" \\') + typer.echo(' -H "Content-Type: application/json" \\') + typer.echo( + """ -d '{"messages": [{"role": "user", "content": "Hello"}], "stream": true}'""" + ) + typer.echo("") + typer.echo("Anthropic Messages:") + typer.echo(f" curl {sdk_base_url}/messages \\") + typer.echo(f' -H "Authorization: Bearer {api_key}" \\') + typer.echo(' -H "Content-Type: application/json" \\') + typer.echo( + """ -d '{"max_tokens": 256, "messages": [{"role": "user", "content": "Hello"}], "stream": true}'""" + ) + typer.echo("") + typer.echo("OpenAI Responses:") + typer.echo(f" curl {sdk_base_url}/responses \\") + typer.echo(f' -H "Authorization: Bearer {api_key}" \\') + typer.echo(' -H "Content-Type: application/json" \\') + typer.echo(""" -d '{"input": "Hello", "stream": true}'""") + typer.echo("") + + # ── 7. Wait for Ctrl+C ──────────────────────────────────────────── + from studio.backend.run import _shutdown_event, _graceful_shutdown, _server + + try: + if _shutdown_event is not None: + while not _shutdown_event.is_set(): + _shutdown_event.wait(timeout = 1) + else: + while True: + time.sleep(1) + except KeyboardInterrupt: + _graceful_shutdown(_server) + typer.echo("\nShutting down...") + + # ── unsloth studio stop ─────────────────────────────────────────────── _PID_FILE = STUDIO_HOME / "studio.pid"