unsloth/studio/backend/tests/test_studio_api.py
oobabooga 72e67ae5a6
Studio: Add Tensor-Parallel llama.cpp support (#6040)
* Studio: Add Tensor-Parallel llama.cpp support

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: harden Tensor-Parallel fallback and GPU selection

* Studio: reconcile split-mode extras and harden tensor-split planning

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: reconcile split-mode extras in backend duplicate-load guard

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: preserve inherited non-tensor split modes on reload

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: honor cancellation in tensor fallback, preserve tensor mode on rollback, and don't raise an explicit small context

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: reconcile split-mode in reload check and strip it on tensor downgrade

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Strip --tensor-split alongside --split-mode so inherited ratios don't override the tensor planner

An inherited or stale --tensor-split in llama_extra_args was appended after
Studio's computed --tensor-split and won last in llama.cpp, re-introducing the
asymmetric-GPU OOM tensor mode is meant to prevent. Group -ts/--tensor-split
into the split-mode shadow set so it is stripped on inherit and on the layer
fallback; parse_split_mode_override still keys on the mode value only.

* Drop quantized KV for the tensor attempt and report native max context

Tensor mode aborts on a quantized KV cache, so a user with q8_0/q4_1 etc. who
enabled Tensor Parallelism silently fell back to layer split. Clear the cache
type (and strip inherited/explicit --cache-type) for the tensor attempt only;
the layer fallback re-runs with tensor off and keeps the user's choice.

Also report max_available_ctx from the native context, not an explicit small
-c, so the context slider no longer warns too early in tensor mode.

* Reconcile inherited split-mode extras in the already-loaded check

When a same-model load omitted llama_extra_args, the tensor comparison resolved
the raw (None) request and treated an inherited --split-mode tensor server as a
mismatch, forcing a needless reload. Compare using the stored extras stripped
the same way the reload strips them.

* Pass tensor_parallel through compare-mode loads

The generalized compare path loaded each GGUF without tensor_parallel, so
compare ran layer split even with the toggle on and left the settings sheet
stale. Send the toggle and hydrate the loaded state from the response, matching
the main chat and recipe load paths.

* Add --tensor-parallel flag to unsloth studio run

The headless one-liner could only reach tensor mode by passing --split-mode
tensor as a raw llama.cpp extra. Add a first-class --tensor-parallel/
--no-tensor-parallel option that sets the tensor_parallel field on the
/api/inference/load payload, forwarded through the studio-venv re-exec like the
other polarity flags. Matches the web UI toggle and the API field.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
2026-06-12 04:00:52 -07:00

944 lines
33 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
End-to-end tests for Unsloth Studio's HTTP API surface.
Covers the OpenAI- 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 -- Studio server-side tools (enable_tools=true)
5. curl -- Standard OpenAI function calling (non-streaming)
6. curl -- Standard OpenAI function calling (streaming)
7. curl -- Standard OpenAI function calling (multi-turn tool loop)
8. OpenAI Python SDK -- Standard function calling
9. Anthropic Messages API -- basic non-streaming
10. Anthropic Messages API -- streaming SSE
11. Anthropic Python SDK -- non-streaming
12. Anthropic Messages API -- streaming with tools
13. Anthropic Messages API -- tool_choice={"type":"any"} honored
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
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",
"--tensor-parallel",
]:
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.
When ``enable_tools`` is set the server always returns SSE streaming
regardless of the ``stream`` flag, so we parse SSE chunks. The model 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")
# Standard OpenAI function-calling pass-through tests.
#
# Regression coverage for unslothai/unsloth#4999: /v1/chat/completions used
# to strip standard OpenAI `tools`/`tool_choice`, so clients never got
# structured tool_calls back. These exercise the pass-through that forwards
# those fields to llama-server verbatim. Require a tool-capable GGUF
# (supports_tools=True); the default unsloth/Qwen3-1.7B-GGUF qualifies.
_WEATHER_TOOL = {
"type": "function",
"function": {
"name": "get_weather",
"description": "Look up the current weather for a given city.",
"parameters": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "The name of the city, e.g. 'Paris'.",
},
},
"required": ["city"],
},
},
}
def _collect_streamed_tool_calls(chunks: list[dict]) -> list[dict]:
"""Reassemble OpenAI streaming delta.tool_calls into full tool calls.
OpenAI streams partial tool calls across chunks — the first chunk for a
given index carries ``id`` + ``function.name``, and later chunks append
fragments to ``function.arguments``.
"""
by_index: dict[int, dict] = {}
for c in chunks:
choices = c.get("choices") or []
if not choices:
continue
delta = choices[0].get("delta") or {}
tool_calls = delta.get("tool_calls") or []
for tc in tool_calls:
idx = tc.get("index", 0)
slot = by_index.setdefault(
idx,
{
"id": None,
"type": "function",
"function": {"name": None, "arguments": ""},
},
)
if tc.get("id"):
slot["id"] = tc["id"]
fn = tc.get("function") or {}
if fn.get("name"):
slot["function"]["name"] = fn["name"]
if fn.get("arguments"):
slot["function"]["arguments"] += fn["arguments"]
return [by_index[i] for i in sorted(by_index)]
def _final_finish_reason(chunks: list[dict]) -> str | None:
for c in reversed(chunks):
choices = c.get("choices") or []
if not choices:
continue
fr = choices[0].get("finish_reason")
if fr is not None:
return fr
return None
def test_openai_tools_nonstream(base_url: str, api_key: str):
"""Standard OpenAI function calling, non-streaming, tool_choice='required'.
Regression: before the fix, Studio stripped `tools` and the model
returned plain text with finish_reason='stop'. After the fix,
llama-server's response is forwarded verbatim so the client sees
finish_reason='tool_calls' with a structured tool_calls array and
non-zero usage.prompt_tokens.
"""
status, text = _http(
"POST",
f"{base_url}/v1/chat/completions",
body = {
"messages": [{"role": "user", "content": "What is the weather in Paris?"}],
"tools": [_WEATHER_TOOL],
"tool_choice": "required",
"stream": False,
},
headers = {"Authorization": f"Bearer {api_key}"},
timeout = 120,
)
assert status == 200, f"Expected 200, got {status}: {text[:500]}"
data = json.loads(text)
assert "choices" in data, f"Missing 'choices': {text[:300]}"
choice = data["choices"][0]
assert (
choice["finish_reason"] == "tool_calls"
), f"Expected finish_reason='tool_calls', got {choice['finish_reason']!r}"
msg = choice["message"]
tool_calls = msg.get("tool_calls") or []
assert len(tool_calls) >= 1, f"No tool_calls in response: {msg}"
first = tool_calls[0]
assert first["type"] == "function"
assert (
first["function"]["name"] == "get_weather"
), f"Wrong tool name: {first['function']['name']!r}"
# arguments must be valid JSON
parsed = json.loads(first["function"]["arguments"])
assert "city" in parsed, f"Tool call missing required 'city' arg: {parsed}"
# Usage must be non-zero (was 0 before the fix)
usage = data.get("usage") or {}
assert usage.get("prompt_tokens", 0) > 0, f"Expected non-zero prompt_tokens; got {usage}"
assert data.get("id"), "Missing response id"
print(
f" PASS openai tools non-stream: "
f"tool={first['function']['name']}, args={parsed}, "
f"prompt_tokens={usage['prompt_tokens']}"
)
def test_openai_tools_stream(base_url: str, api_key: str):
"""Standard OpenAI function calling, streaming, tool_choice='required'."""
status, chunks = _stream_http(
f"{base_url}/v1/chat/completions",
body = {
"messages": [{"role": "user", "content": "What is the weather in Tokyo?"}],
"tools": [_WEATHER_TOOL],
"tool_choice": "required",
"stream": True,
},
headers = {"Authorization": f"Bearer {api_key}"},
timeout = 120,
)
assert status == 200, f"Expected 200, got {status}"
assert len(chunks) > 0, "No SSE chunks received"
assert _final_finish_reason(chunks) == "tool_calls", (
f"Expected final finish_reason='tool_calls', got " f"{_final_finish_reason(chunks)!r}"
)
assembled = _collect_streamed_tool_calls(chunks)
assert len(assembled) >= 1, "No tool_calls reassembled from stream"
first = assembled[0]
assert first["function"]["name"] == "get_weather"
parsed = json.loads(first["function"]["arguments"])
assert "city" in parsed
print(
f" PASS openai tools stream: {len(chunks)} chunks, "
f"tool={first['function']['name']}, args={parsed}"
)
def test_openai_tools_multiturn(base_url: str, api_key: str):
"""Multi-turn client-side tool loop: validates that role='tool' result
messages and assistant messages carrying tool_calls are accepted.
Regression: before the fix, ChatMessage.role was restricted to
{system,user,assistant} and rejected role='tool' at Pydantic
validation. This test sends a full round trip so the model receives the
simulated tool result and responds with final text.
"""
status, text = _http(
"POST",
f"{base_url}/v1/chat/completions",
body = {
"messages": [
{"role": "user", "content": "What is the weather in Paris?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_test_1",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"city": "Paris"}',
},
}
],
},
{
"role": "tool",
"tool_call_id": "call_test_1",
"content": '{"temperature_c": 14, "condition": "cloudy"}',
},
],
"tools": [_WEATHER_TOOL],
"stream": False,
},
headers = {"Authorization": f"Bearer {api_key}"},
timeout = 120,
)
assert status == 200, f"Expected 200, got {status}: {text[:500]}"
data = json.loads(text)
msg = data["choices"][0]["message"]
# The model should respond with text now it has the tool result
content = msg.get("content") or ""
assert len(content) > 0 or msg.get(
"tool_calls"
), f"Expected text or follow-up tool call, got empty message: {msg}"
print(f" PASS openai tools multiturn: {content[:80]!r}")
def test_openai_sdk_tool_calling(base_url: str, api_key: str):
"""OpenAI Python SDK round trip — the real client shape opencode et al. use."""
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)
resp = client.chat.completions.create(
model = "current",
messages = [{"role": "user", "content": "What's the weather in Berlin?"}],
tools = [_WEATHER_TOOL],
tool_choice = "required",
stream = False,
)
assert resp.choices[0].finish_reason == "tool_calls", (
f"Expected finish_reason='tool_calls', got " f"{resp.choices[0].finish_reason!r}"
)
tool_calls = resp.choices[0].message.tool_calls
assert tool_calls and len(tool_calls) >= 1, "No tool_calls from SDK"
tc = tool_calls[0]
assert tc.function.name == "get_weather"
parsed = json.loads(tc.function.arguments)
assert "city" in parsed
print(f" PASS openai SDK tool calling: " f"tool={tc.function.name}, args={parsed}")
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")
def test_anthropic_tool_choice_any(base_url: str, api_key: str):
"""Anthropic Messages API: ``tool_choice: {"type": "any"}`` must be
honored (forwarded as OpenAI ``tool_choice: "required"`` to
llama-server). Regression for the secondary fix bundled with #4999 —
previously this field was accepted on the request model but dropped with
a warning log, so the model could answer from memory instead of using
the tool.
"""
status, events = _stream_anthropic_http(
f"{base_url}/v1/messages",
body = {
"model": "default",
"max_tokens": 256,
"messages": [
# A question the model could answer from memory if
# tool_choice were not enforced.
{
"role": "user",
"content": "What is the weather in London right now?",
}
],
"tools": [
{
"name": "get_weather",
"description": "Look up current weather for a city.",
"input_schema": {
"type": "object",
"properties": {
"city": {"type": "string"},
},
"required": ["city"],
},
}
],
"tool_choice": {"type": "any"},
"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"
# With tool_choice=any, stop_reason must be tool_use, not end_turn
stop_reason = None
for etype, data in events:
if etype == "message_delta":
stop_reason = data.get("delta", {}).get("stop_reason") or stop_reason
assert stop_reason == "tool_use", (
f"Expected stop_reason='tool_use' with tool_choice=any, got "
f"{stop_reason!r} — tool_choice may not be forwarded to llama-server."
)
# And at least one tool_use content block must be emitted
tool_use_starts = [
e
for e in events
if e[0] == "content_block_start" and e[1].get("content_block", {}).get("type") == "tool_use"
]
assert len(tool_use_starts) >= 1, "No tool_use content block emitted"
print(
f" PASS anthropic tool_choice=any honored: "
f"{len(tool_use_starts)} tool_use blocks, stop_reason={stop_reason}"
)
# 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. --help (no server needed)
print("\n[1/16] Testing --help output")
run_test(test_help_output)
# 2-16. 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/16] Testing curl basic (non-streaming)")
run_test(test_curl_basic, base_url, api_key)
print("[3/16] Testing curl streaming")
run_test(test_curl_streaming, base_url, api_key)
print("[4/16] Testing OpenAI Python SDK (streaming)")
run_test(test_openai_sdk, base_url, api_key)
print("[5/16] Testing curl with tools (server-side enable_tools)")
run_test(test_curl_with_tools, base_url, api_key)
print("[6/16] Testing OpenAI standard tools (non-streaming)")
run_test(test_openai_tools_nonstream, base_url, api_key)
print("[7/16] Testing OpenAI standard tools (streaming)")
run_test(test_openai_tools_stream, base_url, api_key)
print("[8/16] Testing OpenAI standard tools (multi-turn)")
run_test(test_openai_tools_multiturn, base_url, api_key)
print("[9/16] Testing OpenAI SDK tool calling")
run_test(test_openai_sdk_tool_calling, base_url, api_key)
print("[10/16] Testing invalid API key rejection")
run_test(test_invalid_key_rejected, base_url)
print("[11/16] Testing no API key rejection")
run_test(test_no_key_rejected, base_url)
print("[12/16] Testing Anthropic basic (non-streaming)")
run_test(test_anthropic_basic, base_url, api_key)
print("[13/16] Testing Anthropic streaming")
run_test(test_anthropic_streaming, base_url, api_key)
print("[14/16] Testing Anthropic Python SDK")
run_test(test_anthropic_sdk, base_url, api_key)
print("[15/16] Testing Anthropic with tools")
run_test(test_anthropic_with_tools, base_url, api_key)
print("[16/16] Testing Anthropic tool_choice=any honored")
run_test(test_anthropic_tool_choice_any, base_url, api_key)
except RuntimeError as exc:
print(f"\nFATAL: Server failed to start: {exc}")
failed += 16 # remaining tests count 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()