ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/models/inference.py
Daniel Han e159b93b97
 studio: improve GGUF tool calling accuracy and reliability (#4700) 
* studio: improve GGUF tool calling accuracy and reliability

- Add URL fetching to web_search tool so models can read full page
  content instead of only getting search snippets. Uses html2text for
  clean markdown conversion with regex fallback.
- Inject current date and behavioral guidance (URL fetch workflow,
  no repeated queries, use code for data processing) into the
  tool-use system prompt.
- Append error recovery nudge to tool results that indicate failure,
  helping small models avoid looping on the same broken call.
- Strip leaked <tool_call> XML from assistant messages in conversation
  history and from the outgoing SSE stream.
- Raise default max tool iterations from 10 to 25 across backend,
  model schema, and frontend defaults.
- Increase _MAX_PAGE_CHARS from 4k to 16k so fetched pages contain
  enough content for the model to extract useful information.
- Add "IMPORTANT: These are only short snippets" hint to search
  results so models know to fetch full pages when needed.

Tested with Qwen3.5-4B-GGUF (UD-Q4_K_XL), 10 runs before/after:
- XML leaks in responses: 10/10 -> 0/10
- URL fetch usage: 0 -> 4/10 runs
- Runs producing actual correct answers: 0/10 -> 2/10
- Average tool calls per query: 5.5 -> 3.8 (more efficient)
- Average response time: 12.3s -> 9.8s

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

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

* Add tool calling benchmark results across model sizes and quants

Tested 16 configurations (4 models x 2 quants x 2 KV cache types)
with 10 runs each on NVIDIA B200.

Best config: 27B UD-Q4_K_XL + bf16 KV -- 6/10 runs found all 4
correct songs, 0 XML leaks, 131s average response time.

* Add duplicate tool-call detection and final-answer synthesis

When the model repeats the exact same tool call (same name + arguments)
twice in a row, skip execution and return a redirect message telling it
to try a different approach. This prevents the 8x-repeated-query loops
observed on 27B and 35B models.

When the tool iteration cap (25) is reached, inject a "provide your
final answer now" message before the final streaming pass. This lets
the model synthesize a useful answer from everything it gathered
instead of being silently cut off.

Tested on Qwen3.5-27B UD-Q4_K_XL (10 runs):
- Repeated query runs: 4/10 -> 2/10
- Cap hits: 1/10 -> 0/10
- All 4/4 accuracy: 5/10 -> 7/10

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

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

* Fix CodeQL alert: handle whitespace in script/style closing tags

The regex fallback for HTML stripping did not match closing tags
with whitespace before the angle bracket (e.g. </script >).
Use \s* before > in both script and style patterns.

* Address reviewer findings: SSRF, timeout crash, XML regex, dedup

- SSRF: resolve hostname via getaddrinfo and reject private, loopback,
  link-local, multicast, and reserved addresses before fetching
- Timeout: handle timeout=None (unlimited mode) in URL fetch path
  by defaulting to 60s instead of crashing on min(None, 60)
- Download cap: read at most max_chars*4+1 bytes instead of the
  full response body before truncating
- XML regex: match both <tool_call> and <function=...> markup in
  the history/stream cleanup (inference.py)
- CodeQL: use [^>]* in closing script/style tags to handle any
  whitespace or attributes before >
- Dedup: track whether each tool call failed so retries after
  transient errors are allowed; only block consecutive identical
  calls that both succeeded
- Final-answer synthesis: guard on max_tool_iterations > 0 so
  callers who disable tools do not get a false "used all calls" turn

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

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

* Fix redirect SSRF, SSE streaming regression, dedup off-by-one

- SSRF redirect bypass: disable auto-redirect in urllib, manually
  follow up to 5 hops with host validation at each step. Prevents
  public URLs from redirecting to loopback/private targets.
- SSE streaming: track prev_text on the raw cumulative and strip
  XML from the delta only, so completed tool_call tags do not cause
  the cumulative to shrink and drop trailing real text.
- Dedup off-by-one: check the immediately previous call (window=1)
  instead of requiring 2 matching history entries, so the second
  identical successful call is blocked rather than the third.

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

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

* Fix redirect HTTPError handling and tighten error prefixes

- Redirect fix: urllib raises HTTPError (not a normal response) when
  the redirect handler returns None. Catch HTTPError for 3xx codes
  and extract the Location header from the exception object.
- Error prefixes: remove overly broad "No " prefix that matched
  "No results found." (a valid empty-search outcome, not an error).
  Replace with specific prefixes like "Blocked:", "No query provided",
  "Failed to resolve". This ensures empty search results are correctly
  classified as non-errors for duplicate-call tracking.

* Fix SSE cross-chunk XML leaks, cleanup review findings

- SSE streaming: sanitize the full cumulative text before diffing
  against the previous sanitized snapshot, so XML tags that span
  chunk boundaries are stripped correctly. The previous delta-based
  approach leaked split tags.
- DRAINING fallback: use _strip_tool_markup() helper instead of a
  manual regex that only handled <tool_call> but not <function=...>.
- Move hashlib import, _TOOL_XML_RE compile, and datetime import to
  module level per style guide.
- Remove unused _hit_tool_cap variable.

* Fix DNS rebinding, charset detection, HTTPError handling, dedup double-record

- DNS rebinding: resolve hostname once via getaddrinfo, pin the
  returned IP, rewrite the URL to connect to the pinned IP with
  a Host header. Each redirect hop re-resolves and re-validates.
  Closes the TOCTOU window between validation and connection.
- Charset: use resp.headers.get_content_charset() instead of
  hardcoding utf-8, so pages with other encodings decode correctly.
- HTTPError: return descriptive "HTTP {code} {reason}" instead of
  re-raising into a generic "Search failed" message.
- Dedup: remove redundant _record_tool_call in the duplicate branch;
  the single call at the end of the loop handles all cases.

* [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>
2026-03-31 03:06:44 -07:00

426 lines
15 KiB
Python
Raw Permalink Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Pydantic schemas for Inference API
"""
from __future__ import annotations
import time
import uuid
from typing import Annotated, Any, Dict, Literal, Optional, List, Union
from pydantic import BaseModel, Discriminator, Field, Tag
class LoadRequest(BaseModel):
"""Request to load a model for inference"""
model_path: str = Field(..., description = "Model identifier or local path")
hf_token: Optional[str] = Field(
None, description = "HuggingFace token for gated models"
)
max_seq_length: int = Field(
0,
ge = 0,
le = 1048576,
description = "Maximum sequence length (0 = model default for GGUF)",
)
load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization")
is_lora: bool = Field(False, description = "Whether this is a LoRA adapter")
gguf_variant: Optional[str] = Field(
None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
)
trust_remote_code: bool = Field(
False,
description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
)
chat_template_override: Optional[str] = Field(
None,
description = "Custom Jinja2 chat template to use instead of the model's default",
)
cache_type_kv: Optional[str] = Field(
None,
description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')",
)
gpu_ids: Optional[List[int]] = Field(
None,
description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.",
)
class UnloadRequest(BaseModel):
"""Request to unload a model"""
model_path: str = Field(..., description = "Model identifier to unload")
class ValidateModelRequest(BaseModel):
"""
Lightweight validation request to check whether a model identifier
*can be resolved* into a ModelConfig.
This does NOT actually load weights into GPU memory.
"""
model_path: str = Field(..., description = "Model identifier or local path")
hf_token: Optional[str] = Field(
None, description = "HuggingFace token for gated models"
)
gguf_variant: Optional[str] = Field(
None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
)
class ValidateModelResponse(BaseModel):
"""
Result of model validation.
valid == True means ModelConfig.from_identifier() succeeded and basic
introspection (GGUF / LoRA / vision flags) is available.
"""
valid: bool = Field(..., description = "Whether the model identifier looks valid")
message: str = Field(..., description = "Human-readable validation message")
identifier: Optional[str] = Field(None, description = "Resolved model identifier")
display_name: Optional[str] = Field(
None, description = "Display name derived from identifier"
)
is_gguf: bool = Field(False, description = "Whether this is a GGUF model (llama.cpp)")
is_lora: bool = Field(False, description = "Whether this is a LoRA adapter")
is_vision: bool = Field(False, description = "Whether this is a vision-capable model")
class GenerateRequest(BaseModel):
"""Request for text generation (legacy /generate/stream endpoint)"""
messages: List[dict] = Field(..., description = "Chat messages in OpenAI format")
system_prompt: str = Field("", description = "System prompt")
temperature: float = Field(0.6, ge = 0.0, le = 2.0, description = "Sampling temperature")
top_p: float = Field(0.95, ge = 0.0, le = 1.0, description = "Top-p sampling")
top_k: int = Field(20, ge = -1, le = 100, description = "Top-k sampling")
max_new_tokens: int = Field(
2048, ge = 1, le = 4096, description = "Maximum tokens to generate"
)
repetition_penalty: float = Field(
1.0, ge = 1.0, le = 2.0, description = "Repetition penalty"
)
presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty")
image_base64: Optional[str] = Field(
None, description = "Base64 encoded image for vision models"
)
class LoadResponse(BaseModel):
"""Response after loading a model"""
status: str = Field(..., description = "Load status")
model: str = Field(..., description = "Model identifier")
display_name: str = Field(..., description = "Display name of the model")
is_vision: bool = Field(False, description = "Whether model is a vision model")
is_lora: bool = Field(False, description = "Whether model is a LoRA adapter")
is_gguf: bool = Field(
False, description = "Whether model is a GGUF model (llama.cpp)"
)
is_audio: bool = Field(False, description = "Whether model is a TTS audio model")
audio_type: Optional[str] = Field(
None, description = "Audio codec type: snac, csm, bicodec, dac"
)
has_audio_input: bool = Field(
False, description = "Whether model accepts audio input (ASR)"
)
inference: dict = Field(
..., description = "Inference parameters (temperature, top_p, top_k, min_p)"
)
context_length: Optional[int] = Field(
None, description = "Model's native context length (from GGUF metadata)"
)
max_context_length: Optional[int] = Field(
None, description = "Maximum context length currently available on this hardware"
)
supports_reasoning: bool = Field(
False,
description = "Whether model supports thinking/reasoning mode (enable_thinking)",
)
reasoning_always_on: bool = Field(
False,
description = "Whether reasoning is always on (hardcoded <think> tags, not toggleable)",
)
supports_tools: bool = Field(
False,
description = "Whether model supports tool calling (web search, etc.)",
)
cache_type_kv: Optional[str] = Field(
None,
description = "KV cache data type for K and V (e.g. 'f16', 'bf16', 'q8_0')",
)
chat_template: Optional[str] = Field(
None,
description = "Jinja2 chat template string (from GGUF metadata or tokenizer)",
)
class UnloadResponse(BaseModel):
"""Response after unloading a model"""
status: str = Field(..., description = "Unload status")
model: str = Field(..., description = "Model identifier that was unloaded")
class InferenceStatusResponse(BaseModel):
"""Current inference backend status"""
active_model: Optional[str] = Field(
None, description = "Currently active model identifier"
)
is_vision: bool = Field(
False, description = "Whether the active model is a vision model"
)
is_gguf: bool = Field(
False, description = "Whether the active model is a GGUF model (llama.cpp)"
)
gguf_variant: Optional[str] = Field(
None, description = "GGUF quantization variant (e.g. Q4_K_M)"
)
is_audio: bool = Field(
False, description = "Whether the active model is a TTS audio model"
)
audio_type: Optional[str] = Field(
None, description = "Audio codec type: snac, csm, bicodec, dac"
)
has_audio_input: bool = Field(
False, description = "Whether model accepts audio input (ASR)"
)
loading: List[str] = Field(
default_factory = list, description = "Models currently being loaded"
)
loaded: List[str] = Field(
default_factory = list, description = "Models currently loaded"
)
inference: Optional[Dict[str, Any]] = Field(
None, description = "Recommended inference parameters for the active model"
)
supports_reasoning: bool = Field(
False, description = "Whether the active model supports reasoning/thinking mode"
)
reasoning_always_on: bool = Field(
False, description = "Whether reasoning is always on (not toggleable)"
)
supports_tools: bool = Field(
False, description = "Whether the active model supports tool calling"
)
context_length: Optional[int] = Field(
None, description = "Context length of the active model"
)
max_context_length: Optional[int] = Field(
None,
description = "Maximum context length currently available for the active model",
)
# =====================================================================
# OpenAI-Compatible Chat Completions Models
# =====================================================================
# ── Multimodal content parts (OpenAI vision format) ──────────────
class TextContentPart(BaseModel):
"""Text content part in a multimodal message."""
type: Literal["text"]
text: str
class ImageUrl(BaseModel):
"""Image URL object — supports data URIs and remote URLs."""
url: str = Field(..., description = "data:image/png;base64,... or https://...")
detail: Optional[Literal["auto", "low", "high"]] = "auto"
class ImageContentPart(BaseModel):
"""Image content part in a multimodal message."""
type: Literal["image_url"]
image_url: ImageUrl
def _content_part_discriminator(v):
if isinstance(v, dict):
return v.get("type")
return getattr(v, "type", None)
ContentPart = Annotated[
Union[
Annotated[TextContentPart, Tag("text")],
Annotated[ImageContentPart, Tag("image_url")],
],
Discriminator(_content_part_discriminator),
]
"""Union type for multimodal content parts, discriminated by the 'type' field."""
# ── Messages ─────────────────────────────────────────────────────
class ChatMessage(BaseModel):
"""
A single message in the conversation.
``content`` may be a plain string (text-only) or a list of
content parts for multimodal messages (OpenAI vision format).
"""
role: Literal["system", "user", "assistant"] = Field(
..., description = "Message role"
)
content: Union[str, list[ContentPart]] = Field(
..., description = "Message content (string or multimodal parts)"
)
class ChatCompletionRequest(BaseModel):
"""
OpenAI-compatible chat completion request.
Extensions (non-OpenAI fields) are marked with 'x-unsloth'.
"""
model: str = Field(
"default",
description = "Model identifier (informational; the active model is used)",
)
messages: list[ChatMessage] = Field(..., description = "Conversation messages")
stream: bool = Field(True, description = "Whether to stream the response via SSE")
temperature: float = Field(0.6, ge = 0.0, le = 2.0)
top_p: float = Field(0.95, ge = 0.0, le = 1.0)
max_tokens: Optional[int] = Field(
None, ge = 1, description = "Maximum tokens to generate (None = until EOS)"
)
presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty")
# ── Unsloth extensions (ignored by standard OpenAI clients) ──
top_k: int = Field(20, ge = -1, le = 100, description = "[x-unsloth] Top-k sampling")
min_p: float = Field(
0.01, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold"
)
repetition_penalty: float = Field(
1.0, ge = 1.0, le = 2.0, description = "[x-unsloth] Repetition penalty"
)
image_base64: Optional[str] = Field(
None, description = "[x-unsloth] Base64-encoded image for vision models"
)
audio_base64: Optional[str] = Field(
None, description = "[x-unsloth] Base64-encoded WAV for audio-input models (ASR)"
)
use_adapter: Optional[Union[bool, str]] = Field(
None,
description = (
"[x-unsloth] Adapter control for compare mode. "
"null = no change (default), "
"false = disable adapters (base model), "
"true = enable the current adapter, "
"string = enable a specific adapter by name."
),
)
enable_thinking: Optional[bool] = Field(
None,
description = "[x-unsloth] Enable/disable thinking/reasoning mode for supported models",
)
enable_tools: Optional[bool] = Field(
None,
description = "[x-unsloth] Enable tool calling for supported models",
)
enabled_tools: Optional[list[str]] = Field(
None,
description = "[x-unsloth] List of enabled tool names (e.g. ['web_search', 'python', 'terminal']). If None, all tools are enabled.",
)
auto_heal_tool_calls: Optional[bool] = Field(
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
)
max_tool_calls_per_message: Optional[int] = Field(
25,
ge = 0,
description = "[x-unsloth] Maximum number of tool call iterations per message (0 = disabled, 9999 = unlimited).",
)
tool_call_timeout: Optional[int] = Field(
300,
ge = 1,
description = "[x-unsloth] Timeout in seconds for each tool call execution (9999 = no limit).",
)
session_id: Optional[str] = Field(
None,
description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.",
)
# ── Streaming response chunks ────────────────────────────────────
class ChoiceDelta(BaseModel):
"""Delta content for a streaming chunk."""
role: Optional[str] = None
content: Optional[str] = None
class ChunkChoice(BaseModel):
"""A single choice in a streaming chunk."""
index: int = 0
delta: ChoiceDelta
finish_reason: Optional[Literal["stop", "length"]] = None
class ChatCompletionChunk(BaseModel):
"""A single SSE chunk in OpenAI streaming format."""
id: str = Field(default_factory = lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
created: int = Field(default_factory = lambda: int(time.time()))
model: str = "default"
choices: list[ChunkChoice]
usage: Optional[CompletionUsage] = None
timings: Optional[dict] = None
# ── Non-streaming response ───────────────────────────────────────
class CompletionMessage(BaseModel):
"""The assistant's complete response message."""
role: Literal["assistant"] = "assistant"
content: str
class CompletionChoice(BaseModel):
"""A single choice in a non-streaming response."""
index: int = 0
message: CompletionMessage
finish_reason: Literal["stop", "length"] = "stop"
class CompletionUsage(BaseModel):
"""Token usage statistics (approximate)."""
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
class ChatCompletion(BaseModel):
"""Non-streaming chat completion response."""
id: str = Field(default_factory = lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}")
object: Literal["chat.completion"] = "chat.completion"
created: int = Field(default_factory = lambda: int(time.time()))
model: str = "default"
choices: list[CompletionChoice]
usage: CompletionUsage = Field(default_factory = CompletionUsage)
Reference in a new issue View git blame Copy permalink