Studio: Expose openai and anthropic compatible external API end points (#4956)
* Studio: add API key authentication for programmatic access External users want to hit the Studio API (chat completions with tool calling, training, export, etc.) without going through the browser login flow. This adds sk-unsloth- prefixed API keys that work as a drop-in replacement for JWTs in the Authorization: Bearer header. Backend: - New api_keys table in SQLite (storage.py) - create/list/revoke/validate functions with SHA-256 hashed storage - API key detection in _get_current_subject before the JWT path - POST/GET/DELETE /api/auth/api-keys endpoints on the auth router Frontend: - /api-keys page with create form, one-time key reveal, keys table - API Keys link in desktop and mobile navbar - Route registered with requireAuth guard Zero changes to any existing route handler -- every endpoint that uses Depends(get_current_subject) automatically works with API keys. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use actual origin in API key usage examples The examples on /api-keys were hardcoded to localhost:8888 which is wrong for remote users. Use window.location.origin so the examples show the correct URL regardless of where the user is connecting from. * Add `unsloth studio run` CLI command for one-liner model serving Adds a `run` subcommand that starts Studio, loads a model, creates an API key, and prints a ready-to-use curl command -- similar to `ollama run` or `vllm serve`. Usage: unsloth studio run -m unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add end-to-end tests for `unsloth studio run` and API key usage Tests the 4 usage examples from the API Keys page: 1. curl basic (non-streaming) chat completions 2. curl streaming (SSE) chat completions 3. OpenAI Python SDK streaming completions 4. curl with tools (web_search + python) Also tests --help output, invalid key rejection, and no-key rejection. All 7 tests pass against Qwen3-1.7B-GGUF. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add /v1/completions, /v1/embeddings, /v1/responses endpoints and --parallel support - llama_cpp.py: accept n_parallel param, pass to llama-server --parallel - run.py: plumb llama_parallel_slots through to app.state - inference.py: add /completions and /embeddings as transparent proxies to llama-server, add /responses as application-level endpoint that converts to ChatCompletionRequest; thread n_parallel through load_model - studio.py: set llama_parallel_slots=4 for `unsloth studio run` path * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make /v1/responses endpoint match OpenAI Responses API format The existing /v1/responses shim returned Chat Completions format, which broke OpenAI SDK clients using openai.responses.create(). This commit replaces the endpoint with a proper implementation that: - Returns `output` array with `output_text` content parts instead of `choices` with `message` - Uses `input_tokens`/`output_tokens` instead of `prompt_tokens`/ `completion_tokens` in usage - Sets `object: "response"` and `id: "resp_..."` - Emits named SSE events for streaming (response.created, response.output_text.delta, response.completed, etc.) - Accepts all OpenAI Responses API fields (tools, store, metadata, previous_response_id) without erroring -- silently ignored - Maps `developer` role to `system` and `input_text`/`input_image` content parts to the internal Chat format Adds Pydantic schemas for request/response models and 23 unit tests covering schema validation, input normalisation, and response format. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: add Anthropic-compatible /v1/messages endpoint (#4981) * Add Anthropic-compatible /v1/messages endpoint with tool support Translate Anthropic Messages API format to/from internal OpenAI format and reuse the existing server-side agentic tool loop. Supports streaming SSE (message_start, content_block_delta, etc.) and non-streaming JSON. Includes offline unit tests and e2e tests in test_studio_run.py. * Add enable_tools, enabled_tools, session_id to /v1/messages endpoint Support the same shorthand as /v1/chat/completions: enable_tools=true with an optional enabled_tools list uses built-in server tools without requiring full Anthropic tool definitions. session_id is passed through for sandbox isolation. max_tokens is now optional. * Strip leaked tool-call XML from Anthropic endpoint content Apply _TOOL_XML_RE to content events in both streaming and non-streaming tool paths, matching the OpenAI endpoint behavior. * Emit custom tool_result SSE event in Anthropic stream Adds a non-standard tool_result event between the tool_use block close and the next text block, so clients can see server-side tool execution results. Anthropic SDKs ignore unknown event types. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Split /v1/messages into server-side and client-side tool paths enable_tools=true runs the existing server-side agentic loop with built-in tools (web_search/python/terminal). A bare tools=[...] field now triggers a client-side pass-through: client-provided tools are forwarded to llama-server and any tool_use output is returned to the caller with stop_reason=tool_use for client execution. This fixes Claude Code (and any Anthropic SDK client) which sends tools=[...] expecting client-side execution but was previously routed through execute_tool() and failing with 'Unknown tool'. Adds AnthropicPassthroughEmitter to convert llama-server OpenAI SSE chunks into Anthropic SSE events, plus unit tests covering text blocks, tool_use blocks, mixed, stop reasons, and usage. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix httpcore GeneratorExit in /v1/messages passthrough stream Explicitly aclose aiter_lines() before the surrounding async with blocks unwind, mirroring the prior fix in external_provider.py (a41160d3) and cc757b78's RuntimeError suppression. * Wire stop_sequences through /v1/messages; warn on tool_choice Plumb payload.stop_sequences to all three code paths (server-side tool loop, no-tool plain, client-side passthrough) so Anthropic SDK clients setting stop_sequences get the behavior they expect. The llama_cpp backend already accepted `stop` on both generate_chat_ completion and generate_chat_completion_with_tools; the Anthropic handler simply wasn't passing it. tool_choice remains declared on the request model for Anthropic SDK compatibility (the SDK often sets it by default) but is not yet honored. Log a structured warning on each request carrying a non- null tool_choice so the silent drop is visible to operators. * Wire min_p / repetition_penalty / presence_penalty through /v1/messages Align the Anthropic endpoint's sampling surface with /v1/chat/completions. Adds the three fields as x-unsloth extensions on AnthropicMessagesRequest and threads them through all three code paths: server-side tool loop, no-tool plain, and client-side passthrough. The passthrough builder emits "repeat_penalty" (not "repetition_penalty") because that is llama-server's field name; the backend methods already apply the same rename internally. * Fix block ordering and prev_text reset in non-streaming tool path _anthropic_tool_non_streaming was building the response by appending all tool_use blocks first, then a single concatenated text block at the end — losing generation order and merging pre-tool and post-tool text into one block. It also never reset prev_text between synthesis turns, so the first N characters of each post-tool turn were dropped (where N = length of the prior turn's final cumulative text). Rewrite to build content_blocks incrementally in generation order, matching the streaming emitter's behavior: deltas within a turn are merged into the trailing text block, tool_use blocks interrupt the text sequence, and prev_text is reset on tool_end so turn N+1 diffs against an empty baseline. Caught by gemini-code-assist[bot] review on #4981. * Make test_studio_run.py e2e tests pytest-compatible Add a hybrid session-scoped studio_server fixture in conftest.py that feeds base_url / api_key into the existing e2e test functions. Three invocation modes are now supported: 1. Script mode (unchanged) — python tests/test_studio_run.py 2. Pytest + external server — point at a running instance via UNSLOTH_E2E_BASE_URL / UNSLOTH_E2E_API_KEY env vars, no per-run GGUF load cost 3. Pytest + fixture-managed server — pytest drives _start_server / _kill_server itself via --unsloth-model / --unsloth-gguf-variant, CI-friendly The existing _start_server / _kill_server helpers and main() stay untouched so the script entry point keeps working exactly as before. Test function signatures are unchanged — the (base_url, api_key) parameters now resolve via the new fixtures when running under pytest. * Rename test_studio_run.py -> test_studio_api.py The file is entirely about HTTP API endpoint testing (OpenAI-compatible /v1/chat/completions, Anthropic-compatible /v1/messages, API key auth, plus a CLI --help sanity check on the command that runs the API). None of its tests cover training, export, chat-UI, or internal-Python-API concerns. The old name misleadingly suggested "tests for the unsloth studio run CLI subcommand" — the new name reflects the actual scope. Updates: - git mv the file (rename tracked, history preserved) - Rewrite opening docstring to state the API surface focus and call out what is explicitly out of scope - Update all 4 Usage-block path references to the new filename - LOG_FILE renamed to test_studio_api.log - conftest.py fixture import rewritten from test_studio_run to test_studio_api, plus 7 docstring/comment references updated No functional changes to test logic, signatures, or main(). --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * Fix httpcore asyncgen cleanup in /v1/messages and /v1/completions The earlier fix in985e92a9was incomplete: it closed aiter_lines() explicitly but still used `async with httpx.AsyncClient()` / `async with client.stream()` inside the generator. When the generator is orphaned (e.g. client disconnects mid-stream and Starlette drops the StreamingResponse iterator without explicitly calling aclose()), Python's asyncgen finalizer runs the cleanup in a DIFFERENT task than the one that originally entered the httpx context managers. The `async with` exits then trigger httpcore's HTTP11ConnectionByteStream .aclose(), which enters anyio.CancelScope.__exit__ with a mismatched task and raises RuntimeError("Attempted to exit cancel scope in a different task"). That error escapes any user-owned try/except because it happens during GC finalization. Replace `async with` with manual client/response lifecycle in both /v1/messages passthrough and /v1/completions proxy. Close the response and client in a finally block wrapped in `try: ... except Exception: pass`. This suppresses RuntimeError (and other Exception subclasses) from the anyio cleanup noise while letting GeneratorExit (a BaseException, not Exception) propagate cleanly so the generator terminates as Python expects. Traceback observed in user report: File ".../httpcore/_async/connection_pool.py", line 404, in __aiter__ yield part RuntimeError: async generator ignored GeneratorExit ... File ".../anyio/_backends/_asyncio.py", line 455, in __exit__ raise RuntimeError( RuntimeError: Attempted to exit cancel scope in a different task * Expand unsloth studio run banner with SDK base URL and more curl examples Add an explicit "OpenAI / Anthropic SDK base URL" line inside the info box so SDK users don't accidentally copy the bare server URL (without /v1) into their OpenAI/Anthropic SDK constructors and hit 404s. Replace the single /v1/chat/completions curl example with three labeled blocks: chat/completions, Anthropic /messages, and OpenAI Responses. The Anthropic example includes max_tokens (Anthropic SDKs require it even though Studio accepts None). All examples derived from a computed sdk_base_url so the /v1 prefix stays in sync if the public path ever changes. * Hash API keys with HMAC-SHA256 + persistent server secret Stores the HMAC secret in a new app_secrets singleton table. Fixes CodeQL py/weak-sensitive-data-hashing alert on storage.py:74-76, 394-395. Refresh tokens stay on plain SHA-256 (unchanged _hash_token) so existing user sessions survive upgrade — API keys are new on this branch so there is no migration. * Use PBKDF2 for API key hashing per CodeQL recommendation HMAC-SHA256 was still flagged by py/weak-sensitive-data-hashing. Switch to hashlib.pbkdf2_hmac, which is in CodeQL's recommended allowlist (Argon2/scrypt/bcrypt/PBKDF2). Persistent server-side salt stays in app_secrets for defense-in-depth. 100k iterations to match auth/hashing.py's password hasher. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
This commit is contained in:
Daniel Han
GitHub
parent
3bb72a557f
commit
9a261aec5f
19 changed files with 4760 additions and 7 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
488
studio/backend/core/inference/anthropic_compat.py
Normal file
488
studio/backend/core/inference/anthropic_compat.py
Normal file
|
|
@ -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,
|
||||
},
|
||||
)
|
||||
|
|
@ -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
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
774
studio/backend/tests/test_anthropic_messages.py
Normal file
774
studio/backend/tests/test_anthropic_messages.py
Normal file
|
|
@ -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"
|
||||
328
studio/backend/tests/test_responses_api.py
Normal file
328
studio/backend/tests/test_responses_api.py
Normal file
|
|
@ -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"])
|
||||
643
studio/backend/tests/test_studio_api.py
Normal file
643
studio/backend/tests/test_studio_api.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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 });
|
||||
|
|
|
|||
18
studio/frontend/src/app/routes/api-keys.tsx
Normal file
18
studio/frontend/src/app/routes/api-keys.tsx
Normal file
|
|
@ -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,
|
||||
});
|
||||
|
|
@ -29,6 +29,7 @@ import {
|
|||
ChefHatIcon,
|
||||
Copy01Icon,
|
||||
CursorInfo02Icon,
|
||||
Key01Icon,
|
||||
PackageIcon,
|
||||
Tick02Icon,
|
||||
ZapIcon,
|
||||
|
|
@ -416,6 +417,20 @@ export function Navbar() {
|
|||
</HoverCardContent>
|
||||
</HoverCard>
|
||||
</div>
|
||||
<div className="flex shrink-0 items-center">
|
||||
<Link
|
||||
to="/api-keys"
|
||||
className={cn(
|
||||
"flex h-9 items-center gap-1.5 rounded-md px-3 text-sm font-medium transition-colors hover:bg-accent",
|
||||
pathname === "/api-keys"
|
||||
? "text-foreground"
|
||||
: "text-muted-foreground hover:text-foreground",
|
||||
)}
|
||||
>
|
||||
<HugeiconsIcon icon={Key01Icon} className="size-4" />
|
||||
API Keys
|
||||
</Link>
|
||||
</div>
|
||||
{tourId ? (
|
||||
<div className="flex shrink-0 items-center">
|
||||
<button
|
||||
|
|
@ -529,11 +544,24 @@ export function Navbar() {
|
|||
</Link>
|
||||
);
|
||||
})}
|
||||
<Link
|
||||
to="/api-keys"
|
||||
onClick={() => setMobileOpen(false)}
|
||||
className={cn(
|
||||
"mt-3 flex items-center gap-2 rounded-md border px-3 py-2 text-sm font-medium",
|
||||
pathname === "/api-keys"
|
||||
? "border-foreground bg-foreground text-background"
|
||||
: "border-border text-foreground hover:bg-accent",
|
||||
)}
|
||||
>
|
||||
<HugeiconsIcon icon={Key01Icon} className="size-4" />
|
||||
API Keys
|
||||
</Link>
|
||||
<a
|
||||
href="https://unsloth.ai/docs"
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="mt-3 flex items-center gap-2 rounded-md border border-border px-3 py-2 text-sm font-medium text-foreground hover:bg-accent"
|
||||
className="flex items-center gap-2 rounded-md border border-border px-3 py-2 text-sm font-medium text-foreground hover:bg-accent"
|
||||
onClick={() => setMobileOpen(false)}
|
||||
>
|
||||
<HugeiconsIcon icon={Book03Icon} className="size-4" />
|
||||
|
|
|
|||
417
studio/frontend/src/features/auth/api-keys-page.tsx
Normal file
417
studio/frontend/src/features/auth/api-keys-page.tsx
Normal file
|
|
@ -0,0 +1,417 @@
|
|||
// 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 { DashboardLayout } from "@/components/layout/dashboard-layout";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
DialogDescription,
|
||||
DialogFooter,
|
||||
DialogHeader,
|
||||
DialogTitle,
|
||||
} from "@/components/ui/dialog";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { Label } from "@/components/ui/label";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import {
|
||||
AlertCircleIcon,
|
||||
Copy01Icon,
|
||||
Delete02Icon,
|
||||
Key01Icon,
|
||||
Tick02Icon,
|
||||
} from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { useCallback, useEffect, useRef, useState } from "react";
|
||||
import { authFetch } from "./api";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Types
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
interface ApiKey {
|
||||
id: number;
|
||||
name: string;
|
||||
key_prefix: string;
|
||||
created_at: string;
|
||||
last_used_at: string | null;
|
||||
expires_at: string | null;
|
||||
is_active: boolean;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// API helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async function fetchApiKeys(): Promise<ApiKey[]> {
|
||||
const res = await authFetch("/api/auth/api-keys");
|
||||
if (!res.ok) throw new Error("Failed to load API keys");
|
||||
const data = (await res.json()) as { api_keys: ApiKey[] };
|
||||
return data.api_keys;
|
||||
}
|
||||
|
||||
async function createApiKey(
|
||||
name: string,
|
||||
expiresInDays: number | null,
|
||||
): Promise<{ key: string; api_key: ApiKey }> {
|
||||
const res = await authFetch("/api/auth/api-keys", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
name,
|
||||
expires_in_days: expiresInDays,
|
||||
}),
|
||||
});
|
||||
if (!res.ok) throw new Error("Failed to create API key");
|
||||
return res.json();
|
||||
}
|
||||
|
||||
async function revokeApiKey(keyId: number): Promise<void> {
|
||||
const res = await authFetch(`/api/auth/api-keys/${keyId}`, {
|
||||
method: "DELETE",
|
||||
});
|
||||
if (!res.ok) throw new Error("Failed to revoke API key");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function formatDate(iso: string | null): string {
|
||||
if (!iso) return "--";
|
||||
const d = new Date(iso);
|
||||
return d.toLocaleDateString(undefined, {
|
||||
year: "numeric",
|
||||
month: "short",
|
||||
day: "numeric",
|
||||
hour: "2-digit",
|
||||
minute: "2-digit",
|
||||
});
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Components
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function CopyButton({ text }: { text: string }) {
|
||||
const [copied, setCopied] = useState(false);
|
||||
const timerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
if (timerRef.current) clearTimeout(timerRef.current);
|
||||
};
|
||||
}, []);
|
||||
|
||||
const handleCopy = () => {
|
||||
if (!copyToClipboard(text)) return;
|
||||
setCopied(true);
|
||||
if (timerRef.current) clearTimeout(timerRef.current);
|
||||
timerRef.current = setTimeout(() => setCopied(false), 2000);
|
||||
};
|
||||
|
||||
return (
|
||||
<Button variant="outline" size="sm" onClick={handleCopy} className="shrink-0">
|
||||
<HugeiconsIcon
|
||||
icon={copied ? Tick02Icon : Copy01Icon}
|
||||
className={cn("size-3.5 mr-1.5", copied && "text-emerald-600")}
|
||||
/>
|
||||
{copied ? "Copied" : "Copy"}
|
||||
</Button>
|
||||
);
|
||||
}
|
||||
|
||||
function RevealKeyDialog({
|
||||
open,
|
||||
rawKey,
|
||||
onClose,
|
||||
}: {
|
||||
open: boolean;
|
||||
rawKey: string;
|
||||
onClose: () => void;
|
||||
}) {
|
||||
return (
|
||||
<Dialog open={open} onOpenChange={(o) => !o && onClose()}>
|
||||
<DialogContent className="max-w-lg">
|
||||
<DialogHeader>
|
||||
<DialogTitle>API Key Created</DialogTitle>
|
||||
<DialogDescription>
|
||||
Copy this key now. It will not be shown again.
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
<div className="flex items-center gap-2 rounded-md border border-border bg-muted/40 p-3">
|
||||
<code className="min-w-0 flex-1 break-all font-mono text-sm">
|
||||
{rawKey}
|
||||
</code>
|
||||
<CopyButton text={rawKey} />
|
||||
</div>
|
||||
<div className="flex items-start gap-2 rounded-md border border-amber-500/20 bg-amber-50 p-3 text-amber-800 dark:border-amber-400/20 dark:bg-amber-950/30 dark:text-amber-300">
|
||||
<HugeiconsIcon icon={AlertCircleIcon} className="mt-0.5 size-4 shrink-0" />
|
||||
<p className="text-xs leading-relaxed">
|
||||
Store this key securely. You will not be able to see it again after closing this dialog.
|
||||
</p>
|
||||
</div>
|
||||
<DialogFooter>
|
||||
<Button onClick={onClose}>Done</Button>
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</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 (
|
||||
<form onSubmit={handleSubmit} className="flex flex-col gap-4 rounded-lg border border-border p-4">
|
||||
<div className="flex flex-col gap-1.5">
|
||||
<Label htmlFor="key-name">Key name</Label>
|
||||
<Input
|
||||
id="key-name"
|
||||
placeholder="e.g. My application"
|
||||
value={name}
|
||||
onChange={(e) => setName(e.target.value)}
|
||||
required
|
||||
/>
|
||||
</div>
|
||||
<div className="flex flex-col gap-1.5">
|
||||
<Label htmlFor="key-expiry">Expires in (days)</Label>
|
||||
<Input
|
||||
id="key-expiry"
|
||||
type="number"
|
||||
min={1}
|
||||
placeholder="Leave blank for no expiry"
|
||||
value={expiresInDays}
|
||||
onChange={(e) => setExpiresInDays(e.target.value)}
|
||||
/>
|
||||
</div>
|
||||
<Button type="submit" disabled={loading || !name.trim()} className="self-start">
|
||||
{loading ? "Creating..." : "Create API key"}
|
||||
</Button>
|
||||
</form>
|
||||
);
|
||||
}
|
||||
|
||||
function KeysTable({
|
||||
keys,
|
||||
onRevoke,
|
||||
}: {
|
||||
keys: ApiKey[];
|
||||
onRevoke: (id: number) => void;
|
||||
}) {
|
||||
if (keys.length === 0) {
|
||||
return (
|
||||
<p className="py-8 text-center text-sm text-muted-foreground">
|
||||
No API keys yet. Create one above.
|
||||
</p>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="overflow-x-auto rounded-lg border border-border">
|
||||
<table className="w-full text-sm">
|
||||
<thead>
|
||||
<tr className="border-b border-border bg-muted/40">
|
||||
<th className="px-4 py-2.5 text-left font-medium text-muted-foreground">Name</th>
|
||||
<th className="px-4 py-2.5 text-left font-medium text-muted-foreground">Key</th>
|
||||
<th className="px-4 py-2.5 text-left font-medium text-muted-foreground">Created</th>
|
||||
<th className="px-4 py-2.5 text-left font-medium text-muted-foreground">Last used</th>
|
||||
<th className="px-4 py-2.5 text-left font-medium text-muted-foreground">Expires</th>
|
||||
<th className="px-4 py-2.5 text-right font-medium text-muted-foreground" />
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
{keys.map((k) => (
|
||||
<tr
|
||||
key={k.id}
|
||||
className={cn(
|
||||
"border-b border-border last:border-b-0",
|
||||
!k.is_active && "opacity-50",
|
||||
)}
|
||||
>
|
||||
<td className="px-4 py-2.5 font-medium">{k.name}</td>
|
||||
<td className="px-4 py-2.5">
|
||||
<code className="rounded bg-muted px-1.5 py-0.5 font-mono text-xs">
|
||||
sk-unsloth-{k.key_prefix}...
|
||||
</code>
|
||||
</td>
|
||||
<td className="px-4 py-2.5 text-muted-foreground">{formatDate(k.created_at)}</td>
|
||||
<td className="px-4 py-2.5 text-muted-foreground">{formatDate(k.last_used_at)}</td>
|
||||
<td className="px-4 py-2.5 text-muted-foreground">{formatDate(k.expires_at)}</td>
|
||||
<td className="px-4 py-2.5 text-right">
|
||||
{k.is_active ? (
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
onClick={() => onRevoke(k.id)}
|
||||
className="text-destructive hover:text-destructive"
|
||||
>
|
||||
<HugeiconsIcon icon={Delete02Icon} className="size-3.5 mr-1" />
|
||||
Revoke
|
||||
</Button>
|
||||
) : (
|
||||
<span className="text-xs text-muted-foreground">Revoked</span>
|
||||
)}
|
||||
</td>
|
||||
</tr>
|
||||
))}
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
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 (
|
||||
<div className="flex flex-col gap-4">
|
||||
<h3 className="text-sm font-semibold">Usage examples</h3>
|
||||
<div className="flex flex-col gap-3">
|
||||
<div>
|
||||
<p className="mb-1.5 text-xs font-medium text-muted-foreground">curl</p>
|
||||
<pre className="overflow-x-auto rounded-md border border-border bg-muted/40 p-3 font-mono text-xs leading-relaxed">
|
||||
{curlExample}
|
||||
</pre>
|
||||
</div>
|
||||
<div>
|
||||
<p className="mb-1.5 text-xs font-medium text-muted-foreground">Python (OpenAI SDK)</p>
|
||||
<pre className="overflow-x-auto rounded-md border border-border bg-muted/40 p-3 font-mono text-xs leading-relaxed">
|
||||
{pythonExample}
|
||||
</pre>
|
||||
</div>
|
||||
<div>
|
||||
<p className="mb-1.5 text-xs font-medium text-muted-foreground">With tools (web search + code execution)</p>
|
||||
<pre className="overflow-x-auto rounded-md border border-border bg-muted/40 p-3 font-mono text-xs leading-relaxed">
|
||||
{toolsExample}
|
||||
</pre>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Page
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export function ApiKeysPage() {
|
||||
const [keys, setKeys] = useState<ApiKey[]>([]);
|
||||
const [revealedKey, setRevealedKey] = useState<string | null>(null);
|
||||
const [error, setError] = useState<string | null>(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 (
|
||||
<DashboardLayout>
|
||||
<div className="flex flex-col gap-8">
|
||||
<div className="flex items-center gap-3">
|
||||
<div className="flex h-10 w-10 items-center justify-center rounded-lg border border-border bg-muted/40">
|
||||
<HugeiconsIcon icon={Key01Icon} className="size-5" />
|
||||
</div>
|
||||
<div>
|
||||
<h1 className="text-xl font-bold font-heading">API Keys</h1>
|
||||
<p className="text-sm text-muted-foreground">
|
||||
Create keys to access Unsloth Studio programmatically via the OpenAI-compatible API.
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{error && (
|
||||
<div className="flex items-center gap-2 rounded-md border border-destructive/20 bg-destructive/5 p-3 text-sm text-destructive">
|
||||
<HugeiconsIcon icon={AlertCircleIcon} className="size-4 shrink-0" />
|
||||
{error}
|
||||
</div>
|
||||
)}
|
||||
|
||||
<CreateKeyForm onCreated={handleCreated} />
|
||||
<KeysTable keys={keys} onRevoke={handleRevoke} />
|
||||
<UsageExamples />
|
||||
</div>
|
||||
|
||||
<RevealKeyDialog
|
||||
open={revealedKey !== null}
|
||||
rawKey={revealedKey ?? ""}
|
||||
onClose={() => setRevealedKey(null)}
|
||||
/>
|
||||
</DashboardLayout>
|
||||
);
|
||||
}
|
||||
|
|
@ -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";
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue