Compare commits
21 commits
main
...
studio/api
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
69d80ca926 |
||
|
|
035519069f |
||
|
|
42a08e579c | ||
|
|
2010901666 | ||
|
|
d388807ccb | ||
|
|
5cedd9a5b2 | ||
|
|
a59d74d7d1 |
||
|
|
587a84d7d3 |
||
|
|
2edc74fe62 |
||
|
|
8302af06a9 |
||
|
|
ac8d5aee02 | ||
|
|
c45fa95c20 | ||
|
|
315f3d7f81 | ||
|
|
e7d2f2ee5a | ||
|
|
c40ed3766a | ||
|
|
242ec57305 | ||
|
|
924f1620a7 | ||
|
|
d69e294ec6 | ||
|
|
7838e8ad39 | ||
|
|
ae4ff49924 | ||
|
|
d2a917b9a4 |
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