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
|
import jwt
|
||||||
|
|
||||||
from .storage import (
|
from .storage import (
|
||||||
|
API_KEY_PREFIX,
|
||||||
get_jwt_secret,
|
get_jwt_secret,
|
||||||
get_user_and_secret,
|
get_user_and_secret,
|
||||||
load_jwt_secret,
|
load_jwt_secret,
|
||||||
save_refresh_token,
|
save_refresh_token,
|
||||||
|
validate_api_key,
|
||||||
verify_refresh_token,
|
verify_refresh_token,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -137,6 +139,18 @@ async def _get_current_subject(
|
||||||
...
|
...
|
||||||
"""
|
"""
|
||||||
token = credentials.credentials
|
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)
|
subject = _decode_subject_without_verification(token)
|
||||||
if subject is None:
|
if subject is None:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
|
|
|
||||||
|
|
@ -72,7 +72,22 @@ def clear_bootstrap_password() -> None:
|
||||||
|
|
||||||
|
|
||||||
def _hash_token(token: str) -> str:
|
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()
|
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)")}
|
columns = {row["name"] for row in conn.execute("PRAGMA table_info(auth_user)")}
|
||||||
if "must_change_password" not in columns:
|
if "must_change_password" not in columns:
|
||||||
conn.execute(
|
conn.execute(
|
||||||
|
|
@ -112,6 +150,89 @@ def get_connection() -> sqlite3.Connection:
|
||||||
return conn
|
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:
|
def is_initialized() -> bool:
|
||||||
"""Check if auth is ready for login (at least one user exists in DB)."""
|
"""Check if auth is ready for login (at least one user exists in DB)."""
|
||||||
conn = get_connection()
|
conn = get_connection()
|
||||||
|
|
@ -357,3 +478,105 @@ def revoke_user_refresh_tokens(username: str) -> None:
|
||||||
conn.commit()
|
conn.commit()
|
||||||
finally:
|
finally:
|
||||||
conn.close()
|
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,
|
speculative_type: Optional[str] = None,
|
||||||
n_threads: Optional[int] = None,
|
n_threads: Optional[int] = None,
|
||||||
n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused
|
n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused
|
||||||
|
n_parallel: int = 1,
|
||||||
) -> bool:
|
) -> bool:
|
||||||
"""
|
"""
|
||||||
Start llama-server with a GGUF model.
|
Start llama-server with a GGUF model.
|
||||||
|
|
@ -1283,7 +1284,7 @@ class LlamaCppBackend:
|
||||||
"-c",
|
"-c",
|
||||||
str(effective_ctx) if effective_ctx > 0 else "0",
|
str(effective_ctx) if effective_ctx > 0 else "0",
|
||||||
"--parallel",
|
"--parallel",
|
||||||
"1", # Single-user studio, saves VRAM
|
str(n_parallel),
|
||||||
"--flash-attn",
|
"--flash-attn",
|
||||||
"on", # Force flash attention for speed
|
"on", # Force flash attention for speed
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -5,6 +5,8 @@
|
||||||
Pydantic schemas for Authentication API
|
Pydantic schemas for Authentication API
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -45,3 +47,44 @@ class ChangePasswordRequest(BaseModel):
|
||||||
new_password: str = Field(
|
new_password: str = Field(
|
||||||
..., min_length = 8, description = "Replacement password (minimum 8 characters)"
|
..., 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"
|
model: str = "default"
|
||||||
choices: list[CompletionChoice]
|
choices: list[CompletionChoice]
|
||||||
usage: CompletionUsage = Field(default_factory = CompletionUsage)
|
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 fastapi import APIRouter, Depends, HTTPException, status
|
||||||
|
|
||||||
|
from datetime import datetime, timedelta, timezone
|
||||||
|
|
||||||
from models.auth import (
|
from models.auth import (
|
||||||
|
ApiKeyListResponse,
|
||||||
|
ApiKeyResponse,
|
||||||
AuthLoginRequest,
|
AuthLoginRequest,
|
||||||
RefreshTokenRequest,
|
|
||||||
AuthStatusResponse,
|
AuthStatusResponse,
|
||||||
ChangePasswordRequest,
|
ChangePasswordRequest,
|
||||||
|
CreateApiKeyRequest,
|
||||||
|
CreateApiKeyResponse,
|
||||||
|
RefreshTokenRequest,
|
||||||
)
|
)
|
||||||
from models.users import Token
|
from models.users import Token
|
||||||
from auth import storage, hashing
|
from auth import storage, hashing
|
||||||
|
|
@ -131,3 +137,68 @@ async def change_password(
|
||||||
token_type = "bearer",
|
token_type = "bearer",
|
||||||
must_change_password = False,
|
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,
|
port: int = 8888,
|
||||||
frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist",
|
frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist",
|
||||||
silent: bool = False,
|
silent: bool = False,
|
||||||
|
llama_parallel_slots: int = 1,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Start the FastAPI server.
|
Start the FastAPI server.
|
||||||
|
|
@ -257,6 +258,7 @@ def run_server(
|
||||||
port: Port to bind to (auto-increments if in use)
|
port: Port to bind to (auto-increments if in use)
|
||||||
frontend_path: Path to frontend build directory (optional)
|
frontend_path: Path to frontend build directory (optional)
|
||||||
silent: Suppress startup messages
|
silent: Suppress startup messages
|
||||||
|
llama_parallel_slots: Number of parallel slots for llama-server
|
||||||
|
|
||||||
Note:
|
Note:
|
||||||
Signal handlers are NOT registered here so that embedders
|
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
|
# binds (port==0) leave it unset and let request handlers fall back
|
||||||
# to the ASGI request scope or request.base_url.
|
# to the ASGI request scope or request.base_url.
|
||||||
app.state.server_port = port if port and port > 0 else None
|
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
|
# Run server in a daemon thread
|
||||||
def _run():
|
def _run():
|
||||||
|
|
|
||||||
|
|
@ -3,14 +3,136 @@
|
||||||
|
|
||||||
"""
|
"""
|
||||||
Shared pytest configuration for the backend test suite.
|
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
|
import sys
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
# Add backend root to sys.path (mirrors how the app itself is launched)
|
# Add backend root to sys.path (mirrors how the app itself is launched)
|
||||||
_backend_root = Path(__file__).resolve().parent.parent
|
_backend_root = Path(__file__).resolve().parent.parent
|
||||||
if str(_backend_root) not in sys.path:
|
if str(_backend_root) not in sys.path:
|
||||||
sys.path.insert(0, str(_backend_root))
|
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 onboardingRoute } from "./routes/onboarding";
|
||||||
import { Route as changePasswordRoute } from "./routes/change-password";
|
import { Route as changePasswordRoute } from "./routes/change-password";
|
||||||
import { Route as studioRoute } from "./routes/studio";
|
import { Route as studioRoute } from "./routes/studio";
|
||||||
|
import { Route as apiKeysRoute } from "./routes/api-keys";
|
||||||
|
|
||||||
const routeTree = rootRoute.addChildren([
|
const routeTree = rootRoute.addChildren([
|
||||||
indexRoute,
|
indexRoute,
|
||||||
|
|
@ -25,6 +26,7 @@ const routeTree = rootRoute.addChildren([
|
||||||
exportRoute,
|
exportRoute,
|
||||||
dataRecipesRoute,
|
dataRecipesRoute,
|
||||||
dataRecipeRoute,
|
dataRecipeRoute,
|
||||||
|
apiKeysRoute,
|
||||||
]);
|
]);
|
||||||
|
|
||||||
export const router = createRouter({ routeTree });
|
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,
|
ChefHatIcon,
|
||||||
Copy01Icon,
|
Copy01Icon,
|
||||||
CursorInfo02Icon,
|
CursorInfo02Icon,
|
||||||
|
Key01Icon,
|
||||||
PackageIcon,
|
PackageIcon,
|
||||||
Tick02Icon,
|
Tick02Icon,
|
||||||
ZapIcon,
|
ZapIcon,
|
||||||
|
|
@ -416,6 +417,20 @@ export function Navbar() {
|
||||||
</HoverCardContent>
|
</HoverCardContent>
|
||||||
</HoverCard>
|
</HoverCard>
|
||||||
</div>
|
</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 ? (
|
{tourId ? (
|
||||||
<div className="flex shrink-0 items-center">
|
<div className="flex shrink-0 items-center">
|
||||||
<button
|
<button
|
||||||
|
|
@ -529,11 +544,24 @@ export function Navbar() {
|
||||||
</Link>
|
</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
|
<a
|
||||||
href="https://unsloth.ai/docs"
|
href="https://unsloth.ai/docs"
|
||||||
target="_blank"
|
target="_blank"
|
||||||
rel="noopener noreferrer"
|
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)}
|
onClick={() => setMobileOpen(false)}
|
||||||
>
|
>
|
||||||
<HugeiconsIcon icon={Book03Icon} className="size-4" />
|
<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
|
// SPDX-License-Identifier: AGPL-3.0-only
|
||||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
// 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 { LoginPage } from "./login-page";
|
||||||
export { ChangePasswordPage } from "./change-password-page";
|
export { ChangePasswordPage } from "./change-password-page";
|
||||||
export { authFetch, refreshSession } from "./api";
|
export { authFetch, refreshSession } from "./api";
|
||||||
|
|
|
||||||
|
|
@ -69,6 +69,80 @@ def _find_setup_script() -> Optional[Path]:
|
||||||
return None
|
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) ──────────────────────────────────────────
|
# ── unsloth studio (server) ──────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -166,6 +240,174 @@ def studio_default(
|
||||||
typer.echo("\nShutting down...")
|
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 ───────────────────────────────────────────────
|
# ── unsloth studio stop ───────────────────────────────────────────────
|
||||||
|
|
||||||
_PID_FILE = STUDIO_HOME / "studio.pid"
|
_PID_FILE = STUDIO_HOME / "studio.pid"
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue