unsloth/studio/backend/models/inference.py
Daniel Han 0533efe3f8
Harden model fetching (#6391)
* Harden model fetching: consent gate for trust_remote_code

Add a load-path consent gate that scans a model's auto_map repository code
before it executes and blocks CRITICAL/HIGH findings unless the user pins
approval of that exact code version. Capability detection stays code-free,
reading raw config.json instead of AutoConfig.

- Scan config.json and tokenizer_config.json auto_map, nested local helpers,
  and external owner/name--module repos; fail closed on partial downloads.
- Gate inference, training, and export workers, including the MLX path and a
  LoRA's base model, and report requires_trust_remote_code from the raw config
  so chat and auto-load surface the dialog.
- Verify trusted-org auto-enable against the Hub with the request token and key
  the verdict cache by token; reject local-path and spoofed names.
- Add a consent dialog showing the flagged file, line, and surrounding code.
- Thread hf_token through the scan and load paths for gated repos.

* Address review: token handling, tokenizer/LoRA scan coverage, rollback

- Send the HF token for remote-code scans in the POST body, not the URL, so it
  never lands in a log or browser history.
- Collect tokenizer_config.json auto_map files directly instead of relying only
  on the repo file listing.
- Resolve a LoRA's base model for the validate flag and the scan endpoint so the
  dialog scans the code the workers actually gate.
- Pass the request token to the training YAML trusted-org auto-enable.
- Resend a previously approved fingerprint when rolling back to a custom-code
  model after a failed switch.

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

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

* Consent UX: drop legacy chat toggle, fix decline copy, purge declined downloads

The per-model consent dialog is now the single approval path for custom
(auto_map) code in chat, so three leftovers from before it existed are removed:

- Remove the "Enable custom code" switch from Chat Settings and stop persisting
  trust_remote_code, so a previously saved blanket-on cannot linger and load a
  model without going through per-version review. The flag stays as an internal
  YAML/preset default (e.g. first-party auto-enable); the load path still gates
  every custom-code load on a fingerprint only the dialog produces.
- Reword the decline message and the auto-load toast to describe approving the
  model's code from the dialog, not a missing settings toggle.
- On decline, purge the repo the scan downloaded so untrusted code is not left
  on disk. A new /api/models/discard-remote-code endpoint deletes only a
  metadata-only cache entry the scan created; it refuses local paths, loaded
  models, and any repo with weight files cached, so a model the user already had
  or pre-downloaded is always left untouched. The frontend only calls it when
  the scan reported created_by_scan.

Adds discard-endpoint tests (delete metadata-only, refuse on weights/gguf,
refuse local, no-op when not cached) and a created_by_scan payload assertion.

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

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

* Export: remove the user-facing trust remote code toggle

The Export page kept a "Trust remote code" switch (default on) next to the HF
token field. Like chat, custom (auto_map) code should be approved per model
through the load-time review dialog, not a persistent blanket switch, so the
toggle is removed. The export load path already routes through the same consent
dialog: an HF source now starts with trust_remote_code off and only enables it
when the user approves the scanned code in the dialog (a local checkpoint the
user exported stays trusted by default). With the dialog unreachable and no
approval, an HF source loads with trust_remote_code off, which fails closed
rather than running unreviewed code.

* Block loads of repos with unsafe files using Hugging Face's security scan

The trust_remote_code consent gate covers one load-time RCE vector (a repo's
auto_map Python). It does not cover the other: a malicious pickle inside a weight
file (pytorch_model.bin, *.pkl, *.dat) deserializes during from_pretrained even
with trust_remote_code False, so a repo with a normal config plus a poisoned
pickle slips past the existing gate.

Add a metadata-only malware gate that uses Hugging Face's own scan (picklescan +
ClamAV), read via model_info(securityStatus=True).security_repo_status. It never
downloads, opens, or unpickles the flagged files; it only reads the Hub's verdict
and surfaces the flagged file names. New evaluate_file_security runs
unconditionally (independent of trust_remote_code) in every load path (inference,
training SFT/MLX, export), blocking the load when a file is flagged
unsafe/suspicious/malicious. The /remote-code-scan preflight and the validate
endpoint also report the result so the consent dialog opens as a hard block (no
override) listing the flagged files, even for a repo with no custom code.

Policy: hard block with no user override; fail open when the scan is unavailable
(offline/unscanned) so legitimate loads are not broken; no first-party exemption
(a poisoned pickle in a compromised trusted repo still blocks); local paths and
GGUF are skipped (no Hub scan, non-pickle format). Blocking does not gate on
scansDone, since that is often false for clean repos and a file already flagged
unsafe is unsafe regardless.

Adds test_file_security.py covering the block/allow/fail-open/skip matrix.

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

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

* Address review: scan list-form tokenizer auto_map, gate unsafe files on all load paths

Fixes from a 10-reviewer pass on the model-fetching hardening:

- The remote-code scanner skipped tokenizer auto_map encoded as a [slow, fast]
  list (transformers' standard tokenizer shape, e.g.
  {"AutoTokenizer": ["owner/repo--tokenization_x.Slow", null]}). External
  tokenizer code in that form was never fetched, scanned, or fingerprinted, so an
  AutoTokenizer(trust_remote_code=True) load could run it. _auto_map_refs now
  flattens string, list, and nested values. Adds a regression test.

- Compare-mode chat loads and background auto-load only gated on
  requires_trust_remote_code, so a repo flagged unsafe by the Hub scan but with no
  custom code skipped the hard-block dialog. Both now also gate on
  requires_security_review, matching the main chat path.

- The /remote-code-scan and /validate routes collapsed a LoRA adapter to its base
  before the malware scan, so unsafe files in the adapter repo itself were missed
  in the pre-load review (the workers already scan both). Both routes now run the
  file-security scan over the adapter and the base.

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

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

* Require approval for all HIGH remote code, fail closed when unscannable

Tighten the load-time security gates based on review:

Consent gate
- HIGH-severity auto_map code now requires explicit, per-version approval for
  every repo, including first-party unsloth/nvidia. The org is no longer a
  blanket bypass: a compromised first-party repo with HIGH code still warrants
  review. CRITICAL stays a hard block; clean code still loads after the consent
  prompt.
- Fail closed when auto_map code is present but cannot be fully fetched or
  listed to scan (gated, offline, transient, or a repo-listing failure that
  could hide an imported helper). We cannot fingerprint code we cannot see, so
  this is a non-approvable block, retryable once the repo is reachable.
- Scan auto_map from every config that can carry one (model, tokenizer, image
  and feature processor, processor, video processor), not just config.json and
  tokenizer_config.json, so a custom-processor model is not missed. The file
  list is the single source of truth in remote_code_scan and is pinned to the
  transformers filename constants by a guard test.
- Distinguish a genuine 404 (config truly absent) from a transient error: only
  the latter forces a scan, so a repo with no config is correctly a no-op.

Malware gate
- Scan a remote repo even when its name ends in .gguf; only local paths skip the
  Hub scan, so a repo cannot dodge the scan by naming itself "*.gguf".
- Correct the docstring: a file already flagged unsafe blocks regardless of
  scansDone; the only fail-open path is an unavailable scan.

Coverage
- Resolve a remote LoRA adapter's base model (not just local directories) so the
  base, where the code and weights actually execute, is scanned in validate,
  the scan route, and the training and export workers.
- Gate the embedding training path (FastSentenceTransformer) with the malware
  and consent checks, matching the other load paths.

Tests updated and added for each change.

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

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

* Scope malware gate to the load-path vector; stop false-blocking first-party models

Follow-up hardening from a second review pass + a broad live model matrix
(unsloth/* , nvidia/* , third-party, and the eicar malware repo).

Malware / unsafe-file gate
- Scope the block to the actual RCE vector: a root-level file in a code-executing
  format. from_pretrained deserializes weight files at the repo ROOT, so a flag is
  only a load-path pickle vector there. Two exclusions, because neither is loaded:
  inert formats (safetensors is tensor-only, gguf is non-pickle, configs/text/
  images) and files in subdirectories. This keeps eicar blocked (its *.pkl/*.dat/
  eicar_test_file sit at the repo root) while no longer false-blocking legitimate
  first-party repos: nvidia/Nemotron-H-8B-Base-8K ships root safetensors plus NeMo
  pickle checkpoints under nemo/ that the loader never touches, and the Hub flags
  both; the gate previously hard-blocked it.
- Unknown / future non-"safe" levels now fail closed (block) instead of being
  silently allowed, so Hub schema drift cannot introduce a bypass; in-progress
  ("pending"/"scanning"/"error") levels stay non-blocking to avoid false blocks.

Consent gate
- Ignore a STALE own-repo auto_map target that is absent from the repo listing (an
  older config pointing at a file the repo no longer ships) instead of failing the
  whole repo closed as unscannable. The present .py are still fully scanned, which
  is the stronger coverage, and a file that is not there cannot execute. This
  unblocks first-party models like unsloth/PaddleOCR-VL (its tokenizer_config.json
  names processing_ppocrvl.py while the repo ships processing_paddleocr_vl.py). A
  referenced .py that IS present but cannot be fetched, and a repo-listing failure,
  still fail closed.

Remote LoRA base resolution
- Distinguish a genuine 404 (not a LoRA / repo absent -> None) from a transient
  error: the transient case is retried once, then logged as a WARNING (a missed
  base is scanned by neither gate) rather than silently skipped.

Discard endpoint
- Treat .onnx and .ckpt as weights so a repo whose only heavy artifact is one of
  those is never eligible for the declined-download purge.

Tests added for each: load-path scoping (safetensors/subdir/Nemotron-H shapes,
unknown-level fail-closed, pending non-block), stale own-repo auto_map ref, remote
LoRA transient retry, and the empty-config-list (all-404 -> []) semantics.

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

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

* Make LoRA-base transient-warning test robust to logging backend

Assert on the logger object directly instead of capsys, so the test does not
depend on whether the real structlog logger or the module-stub logger is active
(which varies with test collection order).

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

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

* Allow a repo with auto_map but no executable code (e.g. GGUF) instead of blocking

A config can declare an auto_map yet the repo ship NO executable .py -- most
commonly a GGUF repo whose config.json carries an auto_map copied from the original
model (e.g. unsloth/Llama-3_1-Nemotron-Ultra-253B-v1-GGUF references
modeling_decilm.py, which the GGUF-only repo does not contain). A GGUF model loads
through llama.cpp, which never executes auto_map, and transformers cannot run a file
that is not present, so there is nothing to scan and trust_remote_code is a no-op.

The fail-closed change treated this empty result the same as "code is present but we
could not fetch it" and hard-blocked the load. Distinguish the two: repo_remote_code_files
now RAISES RemoteCodeUnscannable when code is present but cannot be fully fetched or
listed (offline / gated / transient / a present .py that 404s / a listing failure),
and returns an empty dict only when the listing succeeded and the repo genuinely ships
no executable .py. The consent gate blocks on the exception (fail closed) and allows the
empty case as a no-op. Real unscannable code still hard-blocks; eicar and CRITICAL/HIGH
custom code are unaffected.

Verified against all 37 unsloth/*Nemotron* models (two GGUF repos were false-blocked,
now load) and the existing matrix (eicar still blocks; DeepSeek-OCR / NVLM-D-72B still
prompt approvable consent). Tests updated to expect the raise for unscannable cases and
added for the no-executable-code no-op.

* Ignore vestigial auto_map in GGUF repos (llama.cpp never runs it)

A GGUF repo's config.json is often copied verbatim from the original
transformers model, auto_map and all, but a GGUF load goes through
llama.cpp which never executes auto_map, so the config is inert. Treat
a direct .gguf reference, and a repo that ships .gguf weights with no
.safetensors, as having no remote code so the consent flow is never
triggered. A mixed repo with both .gguf and .safetensors is still gated,
since the safetensors variant would load through transformers where
auto_map does run. The check sits behind the existing auto_map-present
gate so normal models pay no extra repo listing.

* Add scanner-result copy to the remote-code consent dialog

Make the consent dialog state the scan outcome in plain language for
every model. When the static scan finds nothing, reassure the user with
'Our automatic scanner did not flag any worrying files, but please
double check.' (shown only for the clean, approvable case). When the
scan flags custom code or unsafe files, label the list with 'Our
automatic scanner flagged issues including:'. The Hugging Face
attribution for unsafe files stays in the dialog description.

* Close GGUF-suffix consent bypass for repo ids ending in .gguf

The .gguf short-circuit in _config_has_auto_map skipped the scan for any
model name ending in .gguf, including a bare two-segment repo id like
'evil/model.gguf'. Such a repo can still ship safetensors plus auto_map
Python that transformers would execute, so skipping the scan was an
asymmetric bypass (file_security already scans those repos). Restrict the
short-circuit to genuine direct GGUF file references via
_is_direct_gguf_file_ref: a local .gguf path, or a remote repo_id plus
filename (three or more segments). A two-segment repo id named *.gguf now
falls through to the config scan and _is_gguf_repo file inspection, so it
only skips consent when it actually ships .gguf weights and no safetensors.

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

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

* Align consent dialog body with the title and fix narrow-width overflow

The scan results (the 'Our automatic scanner...' label, finding/unsafe
cards, and the clean-scan reassurance) sat at the dialog's left padding
while the title and description were indented past the status icon, so
the body did not line up under the description. Move the title,
description and results into one column to the right of the icon so they
share a left edge, and let that column fill its width so the description
no longer wraps early.

Also stop a wide code snippet from pushing the dialog off-screen on
narrow viewports: AlertDialogHeader is a grid with place-items-center,
which sized the content row to its content; give the row w-full so it
fills the track, and add min-w-0 down the results chain so the snippet
scrolls inside its card instead of widening the dialog. Verified aligned
and contained from mobile portrait through ultrawide.

* Treat a repo as GGUF-only only when it ships no transformers weights

_is_gguf_repo excluded only .safetensors, so a repo with a .gguf and a
pytorch_model.bin (or .pt/.pth/.h5/.msgpack/.onnx/.ckpt) and no
safetensors was treated as GGUF-only and skipped the consent scan, even
though transformers can load that weight set and execute the repo's
auto_map code. Require the absence of ANY transformers-loadable weight
before treating the repo as a llama.cpp-only GGUF load. A genuine
GGUF-only repo (only .gguf) is still inert; a mixed repo with any pickle
or safetensors weight is gated. Adds a regression test across all the
non-safetensors weight formats.

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

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

* Block flagged subdir weight shards referenced by a root index

The malware gate treated every subdirectory file as non-loadable, but
from_pretrained deserializes a subdir shard a root index references
(pytorch_model.bin.index.json -> shards/...-00001-of-00002.bin). Read the
root weight indexes and block a flagged subdir pickle the weight_map
points at; a flagged subdir pickle no index lists (NeMo nemo/*.distcp)
stays non-blocking, and an inconclusive index lookup fails closed.

* Pass hf_token to the export checkpoint load

ExportBackend.load_checkpoint scanned with hf_token in the worker but
loaded the weights unauthenticated, so a gated/private checkpoint passed
preflight then 401'd at from_pretrained. Add hf_token to load_checkpoint
and forward token to every from_pretrained branch; the worker passes the
command's hf_token.

* Scope created_by_scan to every HF cache the discard searches

created_by_scan used get_cache_path (active HF_HUB_CACHE only) while
/discard-remote-code deletes across active, legacy, and default caches. A
repo the user already had in a legacy/default cache was marked
scan-created and deleted on decline. Check all three caches for the repo
dir before declaring the scan created it.

* Scan the full .py closure of external auto_map repos

An auto_map cross-repo ref (owner/name--module.Class) only had its entry
file downloaded, but transformers also fetches that file's relative
imports from the same repo, so a dangerous helper.py was left outside the
scanned fingerprint. List each external repo's .py and scan the whole set
(plus the referenced entry files); fail closed if the repo cannot be
listed or fetched.

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

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

* Fail closed when a weight index cannot be fully read

_indexed_shard_paths treated a partial result as definitive: if one weight
index read cleanly but another failed transiently, it returned the shard
paths it did see. A flagged subdirectory pickle listed only by the index we
could not read would then be classed as "not a load input" and skipped,
re-opening the very fail-open this guard was added to close.

Return None whenever any index read is inconclusive, even if another read
cleanly, so the caller blocks the already-flagged subdir pickle. A repo that
ships no index files raises EntryNotFoundError for each (never inconclusive)
and still returns an empty set.

* Match cached repos case-insensitively in the created_by_scan guard

_repo_in_any_hf_cache resolved casing only against the active cache and then
probed every cache with an exact directory name. A case-variant already
present in a legacy or default cache (models--Unsloth--Foo for a scan of
unsloth/foo) was missed, so the repo was marked created_by_scan and deleted
on decline -- but discard_remote_code_download deletes case-insensitively,
so that delete would hit the user's pre-existing cache entry. Detect
case-insensitively too, mirroring the deletion path.

* Skip remote-code and security review for selected GGUF variants

validate_model ran the trust_remote_code and Hugging Face security-scan
preflight against the repo even when the selected artifact is a .gguf. A
GGUF loads through llama.cpp, which never executes the repo's auto_map
Python and never deserializes root pickle weights, so repo-level Transformers
artifacts (a config.json with auto_map, or an unsafe pytorch_model.bin next
to the .gguf in a mixed repo) are inert for that load. Gating the GGUF on
them is a false positive. Run both preflights only for non-GGUF loads.

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

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

* Scope the malware gate to actual load roots and serialized files

Two fixes to evaluate_file_security so it neither misses a load-path pickle nor
false-blocks an inert file:

- Honor subdirectory load roots. Spark-TTS / BiCodec call from_pretrained on the
  snapshot's LLM subdirectory, so a flagged pickle directly under it is a
  root-level load artifact there. A new load_subdirs parameter (set from the
  model's audio type via security_load_subdirs) reclassifies those files relative
  to the load root and looks for weight indexes under it, so a flagged shard in
  that subdir is no longer skipped as "not root-level".
- Exempt source files. A root .py is never deserialized by from_pretrained;
  executable repo code runs only through auto_map, which the remote-code consent
  gate scans. Flagging a Python helper here would false-block a repo that merely
  ships a build or train script.

* Scan a LoRA adapter and base as one consent unit, and gate MEDIUM code

A LoRA load runs both the adapter's and the base's repo code. The consent gate
scanned them separately and pinned one fingerprint per repo, so an adapter that
shipped its own auto_map code was either never shown in the dialog (which only
saw the base) or impossible to approve with the base's fingerprint.

evaluate_remote_code_consent_for_targets now scans all of a load's repos as a
single combined unit and pins ONE fingerprint over the union of their code, so
approving the load approves every repo's code together. evaluate_remote_code_consent
becomes a thin single-target wrapper, and an unscannable target fails the whole
load closed.

Also gate MEDIUM findings: like HIGH they now block pending pinned approval, so a
direct API caller cannot run flagged code by setting trust_remote_code=True
without consenting. Only a clean scan loads without a fingerprint.

* Preflight a LoRA load's adapter and base as one combined consent scan

scan_model_remote_code rewrote a LoRA adapter to its base and scanned only the
base for remote code, so the dialog never surfaced an adapter's own auto_map
code. Scan the adapter and base together through
preflight_remote_code_consent_for_targets, which pins one combined fingerprint
the worker gate accepts. The malware preflight is also scoped to each target's
load subdirectories.

* Apply combined consent and subdir-aware malware scan in load workers

Each load worker (inference, export, training) evaluated remote-code consent
once per target with a single shared fingerprint, so a LoRA adapter that ships
its own auto_map code could not be approved by the base's fingerprint. They now
scan the adapter and base together via evaluate_remote_code_consent_for_targets,
which pins one combined fingerprint over the union of their code. The malware
scan in each worker is also scoped to the model's load subdirectories so a
flagged pickle under a from_pretrained load subdir is not missed.

* Report a consistent trust_remote_code requirement after a model loads

validate_model reports requires_trust_remote_code from the YAML default OR the
raw auto_map, but the load, already-loaded, and status responses reported only
the YAML default. A custom-code model approved and loaded via auto_map was then
reported as not requiring trust_remote_code, so the frontend stored false and a
later retry or rollback sent trust_remote_code=false and failed.

A shared resolver reports the same requirement for a loaded model (a value
stored at load time, else the trust_remote_code the load used, else the YAML
default, else the raw auto_map check), and the load response persists it so the
status and already-loaded paths stay consistent. The selected-GGUF security
review is also scoped to the model's load subdirectories.

* Run the consent gate on training resume and for YAML-only trust_remote_code

Three frontend gaps left a model loading without the trust_remote_code it needs:

- The shared consent helper returned early when the scan found no auto_map and no
  unsafe files, dropping a requirement that comes from a model's Studio YAML
  default (e.g. GLM-4.7-Flash). It now grants the caller's requirement with an
  empty pin instead of sending trust_remote_code=false.
- Resume-from-history called startTraining directly with no consent gate, so a
  resumed run whose model needs custom code (or an old run with no approved
  fingerprint) hit the worker block with no dialog. It now runs the same gate as
  a fresh start.
- HF export passed requiresTrustRemoteCode=false for every HF source, so a
  YAML-only model could not flip the flag before export. It now signals the
  requirement for HF sources.

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

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

* Cover both LoRA repos in validate, report GGUF as inert, purge all declined repos

Three follow-on gaps from the combined adapter+base consent work:

- validate_model resolved requires_trust_remote_code from the base alone, so a
  LoRA adapter that ships its OWN auto_map code (with a plain base) was reported
  as not needing trust_remote_code and the consent dialog never opened. It now
  checks the [adapter, base] target set, matching the scan route and the workers
  (which already gate both) and the security review already running over both.

- The already-loaded, loaded, and status responses for a selected GGUF reported
  requires_trust_remote_code from the model's YAML default. A GGUF loads through
  llama.cpp, which never executes the repo's auto_map Python, so the requirement
  is inert for that load. They now report False, matching validate_model (which
  already skips both gates for GGUF) so a status refresh cannot flip the flag
  back on.

- The remote-code scan downloads both the adapter's and the base's config, but
  created_by_scan tracked only the primary, so a base the scan was first to pull
  into the cache was left on disk when the user declined. The scan now reports
  scan_created_repos (every repo it newly cached) and the decline cleanup purges
  each; created_by_scan stays for older clients. The frontend falls back to the
  primary flag when the list is absent.

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

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

* Scan the repo the load fetches, purge external code on decline, harden consent pins

Six follow-on hardening fixes from a fresh review pass over the gate:

- The malware gate scanned the literal "Spark-TTS-0.5B/LLM" alias, but the trainer
  downloads it as unsloth/Spark-TTS-0.5B and loads LLM/, so the alias 404'd and
  failed open, missing a flagged LLM/ pickle. evaluate_file_security now resolves
  the alias to the repo the loader fetches and scans LLM/ as a load root.

- security_load_subdirs relied only on tokenizer detection, which fails on an
  unresolved alias or offline; it now also honors the Studio YAML audio_type
  default, so a BiCodec LLM/ load root is not missed.

- The remote-code scan downloads external auto_map repos (owner/name--module.Class),
  but the decline cleanup tracked only the model/adapter/base, leaving the external
  untrusted code cached. The scan now enumerates external auto_map repos and reports
  the ones it created in scan_created_repos, so a decline purges them too.

- External auto_map refs failed the whole load closed on a stale or mis-derived
  dotted ref (sub.mod.py vs the real sub/mod.py) even though the actual file was
  present and scanned. They now drop such refs when the repo listing is real, exactly
  like the own-repo path; an empty/incomplete listing still fetches and fails closed.

- The combined consent fingerprint keyed code by the raw target string, so the scan
  endpoint's canonicalized casing and a worker's raw user input produced different
  pins for identical code, rejecting a valid approval. Hub repo ids are now folded to
  lowercase in the key (local paths stay case-sensitive), so the pin tracks the code.

- Export threaded hf_token into the weight load but not into detect_audio_type /
  is_vision_model, so a gated multimodal base 404'd in detection and fell through to
  the text loader. Both probes now use the same token.

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

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

* Thread the token through check-vision and guard the gate's parallel sites

The /check-vision endpoint classified a model without the hf_token, so a gated or
private vision model 404'd in the probe and was reported as a plain text model --
the same dropped-token shape as the export probes, at a sibling site. It now passes
the token like the neighboring /check-embedding endpoint.

Add deterministic consistency guards (tests/test_security_gate_consistency.py) that
enumerate the gate's parallel sites mechanically instead of relying on a review to
spot a missed sibling: every is_vision_model / is_embedding_model / detect_audio_type
caller under routes/ and core/ must thread the token, every GGUF response must report
trust_remote_code via the resolver or False (never the raw YAML default), and every
load worker that runs the malware or consent gate must resolve the LoRA base. A new
site that drops the token or mis-reports the requirement now fails CI directly.

* Narrow the LLM alias rewrite and make audio detection token-aware

Three fixes from the confirmatory review, one a regression from the previous round:

- _load_scan_target rewrote EVERY remote repo ending in "/LLM" to unsloth/<parent>,
  so a real third-party repo named "<owner>/LLM" was scanned as unsloth/<owner>
  while the loader still fetched the real repo -- a fail-open hole introduced when
  the Spark-TTS alias handling was added. It now rewrites only a registry-known
  bicodec alias; every other "/LLM" repo is scanned as itself.

- detect_audio_type cached results under the bare model name, so an unauthenticated
  probe of a gated/private repo cached None and poisoned a later authenticated call
  with the token. The cache is now keyed by (normalized_name, token_fingerprint),
  matching the vision cache.

- The training fallback /check-vision call dropped the hf_token, misclassifying a
  gated/private VLM when the config endpoint failed. It now passes the token, like
  the getModelConfig call it falls back from; checkEmbeddingModel takes the token too.

Extend the consistency guards: every capability cache must be keyed by a tuple
including the token, so a cache re-declared as Dict[str, ...] fails CI.

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

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

* Document the broad .py scan as deliberate and enforce it with a test

The remote-code scanner scans every .py in a repo once an auto_map exists, not
just the auto_map entry's static import closure. This is intentional: the entry
module can reach a sibling via an absolute import, importlib, or exec, none of
which a static relative-import closure follows, so closure-only scanning would be
a real bypass of a load-time RCE gate. The broad scan never under-scans; the cost
is that an unrelated benign script can over-block, which is the safe failure
direction (HIGH stays approvable; only CRITICAL hard-blocks).

Spell this out at both the local and remote scan sites so the choice reads as
deliberate, and add a test asserting an unrelated, never-imported .py is still
scanned -- so a future narrowing to the static closure fails CI.

* Purge a declined remote LoRA adapter the scan downloaded

scan_model_remote_code probed the created-by-scan state AFTER resolving the base,
but get_base_model_from_lora_identifier downloads a remote adapter's own
adapter_config.json, so the adapter looked already-cached and was dropped from
scan_created_repos. On decline the adapter -- including the auto_map .py the
preflight fetched -- was left on disk, defeating the "untrusted code is not left
on disk" guarantee for the adapter itself.

Snapshot the primary's cache state BEFORE base resolution and use it when marking
the adapter scan-created; on any probe error treat it as pre-existing so a decline
never deletes it. The base and external repos are unaffected (their configs are not
downloaded before their own probe). Add a test that models the mid-scan download
side effect, which the prior static-stub tests did not.

* Clear remote-code approval when the training model changes

Switching the training model from an approved custom-code model to a clean one
kept the previous model's trust_remote_code=true and approved fingerprint in the
store: setSelectedModel reset visionImageSize on a true switch but not the
remote-code approval. The clean model then trained with trust_remote_code=true,
which bypasses the compiler and disables fused cross-entropy.

Reset trustRemoteCode and approvedRemoteCodeFingerprint on a true model switch.
The new model's own YAML default is re-applied by loadAndApplyModelDefaults, and a
custom-code model still re-opens the consent dialog before training starts, so the
only change is that a clean model no longer inherits a stale approval.

* Trim verbose comments across the model-fetching hardening changes

Condense the explanatory comments and docstrings introduced across the
trust_remote_code consent gate, the malware/unsafe-file gate, the remote-code
scanner, the load workers, the model routes, and the security frontend into
fewer, tighter lines while preserving every security rationale (fail-open vs
fail-closed direction, the deliberate broad-scan anti-bypass note, the
empty-vs-unscannable distinction, stale-ref handling, and the alias-rewrite
spoof guard).

Comments and docstrings only. No code, logic, identifiers, or test behaviour
changed; verified comment-only via the AST/TypeScript checker (40/40), with the
backend test suite and frontend tsc green.

* Do not cache transient audio-detection failures

detect_audio_type cached _detect_audio_from_tokenizer's result
unconditionally, so a transient read failure (network error or 5xx,
returned as None) poisoned the cache and the later successful probe never
ran. Mirror the vision cache: _detect_audio_from_tokenizer now returns
(audio_type, definitive) and the caller caches only definitive results.

A read that succeeds with no audio tokens, or clean 404s for every
tokenizer path, stays a cacheable None; only a genuine transient failure
(connection error, timeout, 5xx, malformed body) skips the cache so the
next call retries.

---------

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

1657 lines
65 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
"""Pydantic schemas for the 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,
field_validator,
model_validator,
)
class LoadRequest(BaseModel):
"""Request to load a model for inference"""
model_path: str = Field(..., description = "Model identifier or local path")
native_path_lease: Optional[str] = Field(
None, description = "Frontend-visible signed native path grant"
)
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.",
)
approved_remote_code_fingerprint: Optional[str] = Field(
None,
description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.",
)
chat_template_override: Optional[str] = Field(
None,
description = "Custom Jinja2 chat template to use instead of the model's default",
)
@field_validator("chat_template_override")
@classmethod
def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optional[str]:
if value is not None and value.strip() == "":
return None
return value
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.",
)
speculative_type: Optional[str] = Field(
None,
description = (
"Speculative decoding mode for GGUF models. Canonical values: "
"'auto' (platform-aware: MTP on MTP GGUFs, ngram-mod fallback "
"for sub-3B), 'mtp' (force draft-mtp only on both GPU and CPU), "
"'ngram' (force ngram-mod only), 'mtp+ngram' (force "
"ngram-mod+draft-mtp chain on both platforms), 'off' (disabled). "
"Legacy values 'default' (-> auto), 'draft-mtp' (-> mtp), "
"'ngram-mod' (-> ngram), and 'ngram-simple' (kept as-is) are "
"still accepted. Ignored for non-GGUF models."
),
)
spec_draft_n_max: Optional[int] = Field(
None,
ge = 1,
le = 16,
description = (
"Max draft tokens per step for MTP speculative decoding "
"(--spec-draft-n-max). Defaults to 2 on GPU and 3 on CPU/Mac "
"when unset (upstream-bench sweet spot for dense Qwen3.6 MTP "
"quants). Only applied when speculative_type resolves to "
"'mtp' or 'mtp+ngram'."
),
)
tensor_parallel: bool = Field(
False,
description = (
"Split the model across GPUs by tensor (--split-mode tensor) "
"instead of by layer for GGUF models. Only affects multi-GPU "
"setups, where it can make generation significantly faster. "
"No effect on a single GPU. Ignored for non-GGUF models."
),
)
llama_extra_args: Optional[List[str]] = Field(
None,
description = (
"Extra arguments forwarded verbatim to llama-server for GGUF models. "
"One token per list entry, e.g. ['--top-k', '20', '--seed', '42']. "
"Studio-managed flags (model identity, port, context length, GPU placement, "
"auth, --flash-attn, --no-context-shift, --jinja) are rejected. Ignored for "
"non-GGUF models."
),
)
class UnloadRequest(BaseModel):
"""Request to unload a model"""
model_path: str = Field(..., description = "Model identifier to unload")
class ValidateModelRequest(BaseModel):
"""Check whether an identifier resolves to a ModelConfig; does NOT load weights."""
model_path: str = Field(..., description = "Model identifier or local path")
native_path_lease: Optional[str] = Field(
None, description = "Frontend-visible signed native path grant"
)
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')"
)
include_context_length: bool = Field(
False,
description = "Also read the native context length from the local GGUF header. "
"Opt-in so the normal load preflight doesn't pay for a cache scan it doesn't need.",
)
class ValidateModelResponse(BaseModel):
"""Result of model validation.
valid == True means from_identifier() succeeded and GGUF/LoRA/vision flags are 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")
requires_trust_remote_code: bool = Field(
False,
description = "Whether the model defaults require trust_remote_code to be enabled for loading.",
)
requires_security_review: bool = Field(
False,
description = "Whether Hugging Face's security scan flagged unsafe files (e.g. a "
"malicious pickle), so the load is hard-blocked pending review.",
)
context_length: Optional[int] = Field(
None,
description = "Native training context length, read from the GGUF header when the file "
"is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.",
)
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_diffusion: bool = Field(
False, description = "Whether model is a block-diffusion model (DiffusionGemma)"
)
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)"
)
requires_trust_remote_code: bool = Field(
False,
description = "Whether the model defaults require trust_remote_code to be enabled for loading.",
)
context_length: Optional[int] = Field(
None, description = "Runtime context length in tokens for the loaded model"
)
max_context_length: Optional[int] = Field(
None, description = "Maximum context length currently available on this hardware"
)
native_context_length: Optional[int] = Field(
None,
description = "Model's native context length from GGUF metadata (not capped by VRAM)",
)
supports_reasoning: bool = Field(
False,
description = "Whether model supports thinking/reasoning mode (enable_thinking or reasoning_effort)",
)
reasoning_style: Literal["enable_thinking", "reasoning_effort"] = Field(
"enable_thinking",
description = "Reasoning control style: 'enable_thinking' (boolean) or 'reasoning_effort' (low|medium|high)",
)
reasoning_always_on: bool = Field(
False,
description = "Whether reasoning is always on (hardcoded <think> tags, not toggleable)",
)
supports_preserve_thinking: bool = Field(
False,
description = "Whether the template understands the optional preserve_thinking kwarg (Qwen3.6-style)",
)
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)",
)
speculative_type: Optional[str] = Field(
None,
description = (
"Canonical UI-facing requested speculative decoding mode "
"('auto' / 'mtp' / 'ngram' / 'mtp+ngram' / 'off' / "
"'ngram-simple'), round-tripped from the original LoadRequest "
"via _canonicalize_spec_mode. None when no model is loaded."
),
)
spec_draft_n_max: Optional[int] = Field(
None,
description = (
"Active --spec-draft-n-max for MTP speculative decoding, or "
"None when the platform default is in effect."
),
)
tensor_parallel: bool = Field(
False,
description = "Whether tensor-parallel split (--split-mode tensor) is active.",
)
class UnloadResponse(BaseModel):
"""Response after unloading a model"""
status: str = Field(..., description = "Unload status")
model: str = Field(..., description = "Model identifier that was unloaded")
class LoadProgressResponse(BaseModel):
"""Progress of the active GGUF load, sampled on demand.
Drives a real progress bar during the post-download warmup (mmap + CUDA upload)
instead of a spinner that freezes for minutes on large MoE models.
"""
phase: Optional[str] = Field(
None,
description = (
"Load phase: 'mmap' (weights paging into RAM via mmap), "
"'ready' (llama-server reported healthy), or null when no "
"load is in flight."
),
)
bytes_loaded: int = Field(
0,
description = (
"Bytes of the model already resident in the llama-server process (VmRSS on Linux)."
),
)
bytes_total: int = Field(
0,
description = "Total bytes across all GGUF shards for the active model.",
)
fraction: float = Field(0.0, description = "bytes_loaded / bytes_total, clamped to 0..1.")
class InferenceStatusResponse(BaseModel):
"""Current inference backend status"""
active_model: Optional[str] = Field(
None, description = "Currently active model display identifier"
)
model_identifier: Optional[str] = Field(
None,
description = "Loadable identifier for the active model.",
)
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)")
is_diffusion: bool = Field(
False, description = "Whether the active model is a block-diffusion model (DiffusionGemma)"
)
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"
)
requires_trust_remote_code: bool = Field(
False,
description = "Whether the active model requires trust_remote_code to be enabled for loading.",
)
supports_reasoning: bool = Field(
False, description = "Whether the active model supports reasoning/thinking mode"
)
reasoning_style: Literal["enable_thinking", "reasoning_effort"] = Field(
"enable_thinking",
description = "Reasoning control style: 'enable_thinking' (boolean) or 'reasoning_effort' (low|medium|high)",
)
reasoning_always_on: bool = Field(
False, description = "Whether reasoning is always on (not toggleable)"
)
supports_preserve_thinking: bool = Field(
False,
description = "Whether the active model's template understands the optional preserve_thinking kwarg",
)
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",
)
native_context_length: Optional[int] = Field(
None,
description = "Model's native context length from GGUF metadata (not capped by VRAM)",
)
cache_type_kv: Optional[str] = Field(
None,
description = "KV cache quantization dtype (e.g. 'q8_0'), or None for default",
)
chat_template: Optional[str] = Field(
None, description = "Model's default chat template (Jinja2 source), if any"
)
chat_template_override: Optional[str] = Field(
None,
description = "Active chat template override applied at load time, or None if model is using its default",
)
speculative_type: Optional[str] = Field(
None,
description = (
"Canonical UI-facing requested speculative decoding mode "
"('auto' / 'mtp' / 'ngram' / 'mtp+ngram' / 'off' / "
"'ngram-simple'), round-tripped from the original LoadRequest. "
"None when no model is loaded."
),
)
spec_draft_n_max: Optional[int] = Field(
None,
description = (
"Active --spec-draft-n-max for MTP speculative decoding, or "
"None when the platform default is in effect."
),
)
tensor_parallel: bool = Field(
False,
description = "Whether tensor-parallel split (--split-mode tensor) is active.",
)
llama_cpp_supports_mtp: bool = Field(
True,
description = (
"Whether llama.cpp supports MTP (--spec-type mtp/draft-mtp). "
"False -> recommend `unsloth studio update`."
),
)
spec_fallback_reason: Optional[str] = Field(
None,
description = (
"Why MTP was disabled on the loaded model despite being requested "
"(auto on an MTP model, or forced mtp / mtp+ngram). "
"'binary_no_mtp' / 'binary_outdated' -> a newer prebuilt would "
"re-enable it (show the update affordance); 'runtime_error' -> the "
"current build could not run it. None when MTP engaged or was not "
"requested."
),
)
llama_cpp_prebuilt_stale: bool = Field(
False,
description = (
"Installed llama.cpp prebuilt is >=3 days behind the latest "
"release. True -> show `unsloth studio update` banner."
),
)
llama_cpp_installed_tag: Optional[str] = Field(
None,
description = "Installed llama.cpp tag, or None if unknown.",
)
llama_cpp_latest_tag: Optional[str] = Field(
None,
description = "Latest published llama.cpp tag, or None if GitHub unreachable.",
)
# =====================================================================
# 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", "original"]] = "auto"
class ImageContentPart(BaseModel):
"""Image content part in a multimodal message."""
type: Literal["image_url"]
image_url: ImageUrl
class InputDocumentContentPart(BaseModel):
"""Document (PDF / file) content part in a multimodal message.
Studio-normalised shape (file_data or file_url, plus optional filename/media_type).
Mapped onto Anthropic ``document`` / OpenAI ``input_file`` for vision providers;
dropped for non-vision providers.
"""
type: Literal["input_document"]
file_data: Optional[str] = Field(
None,
description = "data:<media_type>;base64,<DATA> URI for inline payloads. Either file_data or file_url must be set; otherwise the part is dropped.",
)
file_url: Optional[str] = Field(
None,
description = "Remote URL pointing to the document (https://...).",
)
filename: Optional[str] = Field(
None,
description = "Display filename, forwarded to providers as `title`/`filename`.",
)
media_type: Optional[str] = Field(
None,
description = 'Override the media type sniffed from the data URI (e.g. "application/pdf").',
)
class OpenAIReasoningContentPart(BaseModel):
"""OpenAI Responses reasoning item paired with a tool output.
Reasoning models may require this replayed before an ``image_generation_call``
id. OpenAI-only; routes strip it for other providers before proxying.
"""
type: Literal["reasoning"]
id: str = Field(..., description = "OpenAI reasoning output item id.")
summary: list[dict[str, Any]] = Field(default_factory = list)
status: Optional[Literal["in_progress", "completed", "incomplete"]] = None
class ImageGenerationCallContentPart(BaseModel):
"""OpenAI Responses image_generation call reference.
Prior ``image_generation_call`` items let follow-up prompts edit a generated
image without resending the payload. The frontend forwards it as a synthetic
assistant part; ``external_provider`` maps it back to a top-level input item.
"""
type: Literal["image_generation_call"]
id: str = Field(..., description = "OpenAI image_generation_call output item id.")
response_id: Optional[str] = Field(
None,
description = "OpenAI Responses response id to use as previous_response_id for follow-up edits.",
)
class CompactionContentPart(BaseModel):
"""Anthropic server-side compaction state, round-tripped on the next turn.
Anthropic returns a ``compaction`` block on the assistant message; the next
request must forward it back so Anthropic reuses the compaction state instead
of re-summarising. See ``external_provider._stream_anthropic`` and
https://platform.claude.com/docs/en/build-with-claude/compaction
"""
type: Literal["compaction"]
content: str = Field(
...,
description = "Anthropic-produced summary of the compacted-away conversation prefix.",
)
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")],
Annotated[InputDocumentContentPart, Tag("input_document")],
Annotated[OpenAIReasoningContentPart, Tag("reasoning")],
Annotated[ImageGenerationCallContentPart, Tag("image_generation_call")],
Annotated[CompactionContentPart, Tag("compaction")],
],
Discriminator(_content_part_discriminator),
]
"""Union type for multimodal content parts, discriminated by the 'type' field."""
# ── Messages ─────────────────────────────────────────────────────
class ChatMessage(BaseModel):
"""Single message in a chat conversation.
``content`` is a string or list of multimodal parts. Assistant messages with
only ``tool_calls`` may set ``content=None``. Missing ``tool_call_id`` on
``role="tool"`` is resolved at the ``ChatCompletionRequest`` layer.
"""
role: Literal["system", "user", "assistant", "tool", "developer"] = Field(
..., description = "Message role"
)
content: Optional[Union[str, list[ContentPart]]] = Field(
None, description = "Message content (string or multimodal parts)"
)
tool_call_id: Optional[str] = Field(
None,
description = "OpenAI tool-result messages: id of the tool call this result belongs to.",
)
tool_calls: Optional[list[dict]] = Field(
None,
description = "OpenAI assistant messages: structured tool calls the model decided to make.",
)
name: Optional[str] = Field(
None,
description = "OpenAI tool-result messages: name of the tool whose result this is.",
)
extra_content: Optional[dict] = Field(
None,
description = (
"Provider-specific extra fields the translator may read. "
"Gemini reads `extra_content.google.thought_signature` "
"from assistant messages to replay text-part signatures."
),
)
@model_validator(mode = "after")
def _validate_role_shape(self) -> "ChatMessage":
if self.tool_calls is not None and self.role != "assistant":
raise ValueError('"tool_calls" is only valid on role="assistant" messages.')
if self.tool_call_id is not None and self.role != "tool":
raise ValueError('"tool_call_id" is only valid on role="tool" messages.')
if self.name is not None and self.role != "tool":
raise ValueError('"name" is only valid on role="tool" messages.')
if self.role == "tool":
# tool_call_id resolution happens at ChatCompletionRequest scope.
# OpenAI accepts empty tool results (commands with no output);
# normalize to "" instead of a 400 agentic clients treat as fatal.
if self.content is None or self.content == []:
self.content = ""
elif self.role == "assistant":
# Post-Stop sentinel: collapse content="" / [] to None.
if (self.content == "" or self.content == []) and not self.tool_calls:
self.content = None
else: # "user" | "system"
if self.content is None or self.content == []:
raise ValueError(f'role="{self.role}" messages require "content".')
return self
class ThinkingConfig(BaseModel):
"""Anthropic-compatible thinking/reasoning configuration.
Use type='disabled' to turn off thinking, or type='enabled' to turn it on.
Only type is read; extra fields (e.g. budget_tokens) are ignored, since
Studio sets provider thinking budgets itself.
"""
type: Literal["disabled", "enabled"] = "disabled"
class ChatCompletionRequest(BaseModel):
"""OpenAI-compatible chat completion request.
Non-OpenAI extension fields are marked with 'x-unsloth'.
"""
# Accept unknown fields so future OpenAI fields aren't dropped before route
# code runs. Mirrors AnthropicMessagesRequest and ResponsesRequest.
model_config = {"extra": "allow"}
model: str = Field(
"default",
description = "Model identifier (informational; the active model is used)",
)
messages: list[ChatMessage] = Field(..., description = "Conversation messages")
stream: bool = Field(
False,
description = (
"Whether to stream the response via SSE. Default matches OpenAI's "
"spec (`false`); opt into streaming by sending `stream: true`."
),
)
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")
stop: Optional[Union[str, list[str]]] = Field(
None,
description = "OpenAI stop sequences: a single string or list of strings at which generation halts.",
)
tools: Optional[list[dict]] = Field(
None,
description = (
"OpenAI function-tool definitions. When provided without `enable_tools=true`, "
"Studio forwards the tools to the backend so the model returns structured "
"tool_calls for the client to execute (standard OpenAI function calling)."
),
)
tool_choice: Optional[Union[str, dict]] = Field(
None,
description = (
"OpenAI tool choice: 'auto' | 'required' | 'none' | "
"{'type': 'function', 'function': {'name': ...}}"
),
)
max_completion_tokens: Optional[int] = Field(
None,
ge = 1,
description = "OpenAI upper bound on generated tokens (supersedes the deprecated max_tokens).",
)
n: Optional[int] = Field(
None,
ge = 1,
le = 128,
description = "Number of chat completion choices to generate.",
)
logprobs: Optional[bool] = Field(
None, description = "Whether to return log probabilities of the output tokens."
)
top_logprobs: Optional[int] = Field(
None,
ge = 0,
le = 20,
description = "Number of most likely tokens (0-20) to return per position; requires logprobs=true.",
)
parallel_tool_calls: Optional[bool] = Field(
None, description = "Whether to enable parallel function calling during tool use."
)
seed: Optional[int] = Field(None, description = "Best-effort deterministic sampling seed.")
stream_options: Optional[dict] = Field(
None,
description = 'Streaming options, e.g. {"include_usage": true} to emit a final usage chunk.',
)
# ── 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 audio (wav/mp3/ogg/flac/m4a) for audio-input models",
)
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",
)
reasoning_effort: Optional[
Literal["none", "minimal", "low", "medium", "high", "max", "xhigh"]
] = Field(
None,
description = "[x-unsloth] Reasoning effort level ('none'|'minimal'|'low'|'medium'|'high'|'max'|'xhigh'). OpenAI `/v1/responses` accepts model-dependent subsets; Anthropic adaptive thinking uses `max` as the top tier on Claude 4.6 Opus/Sonnet (inbound `xhigh` is mapped to `max`) and `xhigh` on Claude 4.7 Opus; local Harmony/gpt-oss templates support low|medium|high.",
)
preserve_thinking: Optional[bool] = Field(
None,
description = "[x-unsloth] When true, keep historical <think> blocks from past assistant turns in the prompt (Qwen3.6 templates). Independent of enable_thinking / reasoning_effort.",
)
thinking: Optional[ThinkingConfig] = Field(
None,
description = "[Anthropic-compatible] Thinking configuration. "
"Use {type: 'disabled'} to disable thinking, {type: 'enabled'} to enable.",
)
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. Local GGUF/safetensors models "
"accept ['web_search', 'python', 'terminal', 'render_html']. External "
"providers accept ['web_search', 'web_fetch', 'code_execution'] for "
"Anthropic and ['web_search', 'code_execution', 'image_generation'] for "
"OpenAI Responses. If None, all local tools are enabled and no "
"server-side tools are forwarded."
),
)
mcp_enabled: Optional[bool] = Field(
None,
description = "[x-unsloth] When true, append tools from every enabled MCP server to this request's tool list.",
)
confirm_tool_calls: Optional[bool] = Field(
None,
description = "[x-unsloth] When true, pause before each tool call and wait for the user to allow/deny it via POST /api/inference/tool-confirm.",
)
bypass_permissions: Optional[bool] = Field(
False,
description = "[x-unsloth] Bypass Permissions: when true, skip the tool-call confirmation gate AND disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits). Secret env vars are still stripped. Takes precedence over confirm_tool_calls.",
)
auto_heal_tool_calls: Optional[bool] = Field(
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
)
context_overflow: Optional[Literal["error", "truncate_middle"]] = Field(
None,
description = (
"[x-unsloth] Passthrough behavior when the prompt exceeds the real "
"context window. 'error' (default) returns a 400 with "
"code=context_length_exceeded. 'truncate_middle' drops middle "
"turn-groups (system prompt, first turn, and recent turns kept; "
"tool calls stay paired with their results) and retries."
),
)
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.",
)
rag_scope: Optional[dict] = Field(
None,
description = (
"[x-unsloth] Hidden RAG retrieval scope for the search_knowledge_base "
"tool: {kb_id?, thread_id?, default_top_k?, mode?, autoinject?, "
"autoinject_min_score?}. Candidate pools and the RRF constant come from "
"server config. The model never sees this; the server resolves which "
"documents to search."
),
)
cancel_id: Optional[str] = Field(
None,
description = "[x-unsloth] Per-request cancellation token. Frontend sends a fresh UUID per run so /inference/cancel matches one specific generation.",
)
# ── External provider routing (x-unsloth extensions) ──────────
provider_id: Optional[str] = Field(
None,
description = "[x-unsloth] Saved provider config ID. If set with encrypted_api_key, routes to external LLM.",
)
provider_type: Optional[str] = Field(
None,
description = "[x-unsloth] Provider type (e.g. 'openai', 'mistral'). Used if provider_id is not set.",
)
external_model: Optional[str] = Field(
None,
description = "[x-unsloth] Model ID at the external provider.",
)
encrypted_api_key: Optional[str] = Field(
None,
description = "[x-unsloth] RSA-encrypted, base64-encoded API key for the external provider.",
)
provider_base_url: Optional[str] = Field(
None,
description = "[x-unsloth] Override base URL for the external provider.",
)
enable_prompt_caching: Optional[Union[bool, str]] = Field(
None,
description = (
"[x-unsloth] Opt in to provider-side prompt caching. On Anthropic, "
"boolean true attaches cache_control={type:ephemeral} to the system "
"block so the static prefix is reused across turns. On OpenAI cloud, "
"caching is automatic for prompts >=1024 tokens and the boolean is "
"informational. On Gemini, pass a string cache resource name such "
"as `cachedContents/abc123` to attach `cachedContent` on the native "
"request (boolean true is a no-op on Gemini because creating the "
"cache requires a separate POST /cachedContents call). Ignored for "
"every other provider. Treated as enabled when omitted."
),
)
@field_validator("enable_prompt_caching", mode = "before")
@classmethod
def _coerce_enable_prompt_caching(cls, value: Any) -> Any:
"""Coerce JSON bool strings back to bool. Widening to Union[bool, str] for
Gemini cache names would let `"false"` read as truthy, so canonical bool
literals are coerced to keep explicit opt-outs working."""
if isinstance(value, str):
lowered = value.strip().lower()
# Match Pydantic v1's bool coercion table; anything else stays a
# string for Gemini's cachedContent resource path.
if lowered in ("true", "t", "1", "yes", "y", "on"):
return True
if lowered in ("false", "f", "0", "no", "n", "off"):
return False
return value
prompt_cache_ttl: Optional[str] = Field(
None,
description = (
"[x-unsloth] Anthropic cache_control TTL. Defaults to the 5-minute "
"ephemeral pool when omitted. Pass `1h` to write into the 1-hour "
"pool instead -- 1h writes are billed at 2x base input vs 1.25x "
"for 5m, but reads stay at 0.1x for both, so 1h pays off the "
"moment a single extra read lands more than 5 minutes after the "
"write. Only `5m` and `1h` are forwarded; any other value is "
"silently ignored downstream so a stale frontend can't make the "
"API 422 on the request. No-op on every non-Anthropic provider."
),
)
compaction_threshold: Optional[int] = Field(
None,
ge = 1,
le = 2_000_000,
description = (
"[x-unsloth] Server-side context compaction trigger, in tokens. "
"Per-provider routing:\n"
" - Anthropic (Opus 4.6+, Sonnet 4.6, Mythos preview): attaches "
"the `compact_20260112` edit and the `compact-2026-01-12` beta "
"header. The upstream floor is 50k; `_stream_anthropic` clamps "
"lower values up.\n"
" - OpenAI cloud (api.openai.com) and Azure OpenAI Foundry "
"(*.openai.azure.com): attaches "
"`context_management:[{type:'compaction', compact_threshold:N}]` "
"to /v1/responses. Effective floor is around 200k (OpenAI's "
"canonical example); values below it surface "
"`compact_threshold is not enabled` 400s upstream.\n"
"Schema floor stays at ge=1 (any positive int) so the field is a "
"silent no-op on non-cloud OpenAI-compatible bases (ollama / "
"llama.cpp / vLLM) and every non-compaction-capable provider "
"rather than returning 422 at request validation time. Per-"
"provider floors are enforced in the corresponding stream helpers."
),
)
openai_code_exec_container_id: Optional[str] = Field(
None,
description = (
"[x-unsloth] OpenAI shell-tool container id from the prior response "
"in the same chat thread. When set and `code_execution` is in "
"`enabled_tools`, the next /v1/responses call uses "
"environment.type='container_reference' so filesystem state "
"persists across turns. Unset → environment.type='container_auto' "
"and OpenAI creates a fresh container. Only meaningful for the "
"OpenAI cloud + gpt-5.5 family path; ignored otherwise."
),
)
anthropic_code_exec_container_id: Optional[str] = Field(
None,
description = (
"[x-unsloth] Anthropic code_execution container id from the prior "
"response in the same chat thread. When set and `code_execution` "
"is in `enabled_tools`, the next /v1/messages call carries a "
"top-level `container` field so the model sees filesystem state "
"from earlier turns. Unset → Anthropic auto-creates a fresh "
"container. Stale ids surface a 4xx with a `container_expired` / "
"`container_not_found` hint; the backend emits a synthetic "
"`container_invalidated` _toolEvent so the next turn falls back "
"to auto-create."
),
)
fast_mode: Optional[bool] = Field(
None,
description = (
"[x-unsloth] Anthropic fast-mode toggle. On Claude Opus 4.6 / "
"4.7 adds the `fast-mode-2026-02-01` beta header and sends "
"`speed: 'fast'` for higher OTPS at premium pricing. Silently "
"ignored on every other model + provider. See "
"https://platform.claude.com/docs/en/build-with-claude/fast-mode"
),
)
@model_validator(mode = "after")
def _resolve_missing_tool_call_ids(self) -> "ChatCompletionRequest":
"""Fill missing tool_call_id by walking back to the preceding assistant.
OpenAI / Anthropic passthrough require the result id to match the
assistant's tool_calls[].id. Prefer function.name match, else first
unconsumed tool_call; synth a random id only if none exists. A user
turn breaks the lookup.
"""
# Pre-mark explicit ids so a missing-id sibling can't steal a claimed one.
consumed: set[tuple[int, int]] = set()
def _mark_consumed(start_idx: int, tool_call_id: str) -> None:
for asst_idx in range(start_idx - 1, -1, -1):
prev = self.messages[asst_idx]
if prev.role == "user":
break
if prev.role != "assistant" or not prev.tool_calls:
continue
for tc_idx, tc in enumerate(prev.tool_calls):
if isinstance(tc, dict) and tc.get("id") == tool_call_id:
consumed.add((asst_idx, tc_idx))
return
for tool_idx, msg in enumerate(self.messages):
if msg.role == "tool" and msg.tool_call_id:
_mark_consumed(tool_idx, msg.tool_call_id)
for tool_idx, msg in enumerate(self.messages):
if msg.role != "tool" or msg.tool_call_id:
continue
picked: str | None = None
for asst_idx in range(tool_idx - 1, -1, -1):
prev = self.messages[asst_idx]
if prev.role != "assistant" or not prev.tool_calls:
if prev.role == "user":
break
continue
name_match = None
fallback = None
for tc_idx, tc in enumerate(prev.tool_calls):
if (asst_idx, tc_idx) in consumed:
continue
if not isinstance(tc, dict):
continue
tc_id = tc.get("id")
if not tc_id:
continue
function = tc.get("function")
function_name = function.get("name") if isinstance(function, dict) else None
if msg.name and function_name == msg.name:
name_match = (tc_id, asst_idx, tc_idx)
break
if fallback is None:
fallback = (tc_id, asst_idx, tc_idx)
chosen = name_match or fallback
if chosen is not None:
picked, a, t = chosen
consumed.add((a, t))
break
if picked is None:
import secrets as _secrets
picked = f"call_{_secrets.token_hex(8)}"
msg.tool_call_id = picked
return self
@model_validator(mode = "after")
def _map_thinking_to_enable_thinking(self) -> "ChatCompletionRequest":
"""Map Anthropic-style ``thinking`` parameter to internal ``enable_thinking``.
``thinking: {type: 'enabled'}`` sets ``enable_thinking = True`` and
``thinking: {type: 'disabled'}`` sets ``enable_thinking = False``.
``enable_thinking`` takes precedence when both are provided so that
callers who already use the internal field are unaffected. Invalid
``thinking`` shapes are rejected at validation time (422).
"""
if self.thinking is not None and self.enable_thinking is None:
self.enable_thinking = self.thinking.type == "enabled"
return self
class ToolConfirmRequest(BaseModel):
session_id: Optional[str] = None
approval_id: Optional[str] = None
decision: Literal["allow", "deny"] = "deny"
# ── OpenAI shell-tool container management ─────────────────────
class OpenAIContainerRequest(BaseModel):
"""Shared body for the OpenAI container endpoints (list / create / delete).
Carries the encrypted API key + base URL so the route can decrypt and proxy
to the user's account, keeping the key off backend persistent storage.
"""
encrypted_api_key: str = Field(
...,
description = "[x-unsloth] RSA-encrypted, base64-encoded OpenAI API key.",
)
provider_base_url: Optional[str] = Field(
None,
description = "[x-unsloth] OpenAI base URL. Only api.openai.com is supported; non-cloud bases are rejected with 400.",
)
class CreateOpenAIContainerBody(OpenAIContainerRequest):
name: str = Field(
...,
min_length = 1,
max_length = 256,
description = "Human-readable container name. Surfaces in the picker UI.",
)
ttl_minutes: int = Field(
20,
ge = 1,
le = 20,
description = (
"Idle-timeout TTL the new container will inherit (anchor="
"last_active_at). OpenAI hard-caps this at 20 minutes and "
"rejects larger values with integer_above_max_value."
),
)
class DeleteOpenAIContainerBody(OpenAIContainerRequest):
container_id: str = Field(
...,
description = "OpenAI container id (cntr_...) to delete.",
)
class OpenAIContainerSummary(BaseModel):
"""One row from GET /v1/containers, reshaped for the UI."""
id: str
name: Optional[str] = None
created_at: Optional[int] = None
last_active_at: Optional[int] = None
expires_after_minutes: Optional[int] = None
status: Optional[str] = None
class ListOpenAIContainersResponse(BaseModel):
containers: list[OpenAIContainerSummary]
# ── Streaming response chunks ────────────────────────────────────
class ChoiceDelta(BaseModel):
"""Delta content for a streaming chunk."""
role: Optional[str] = None
content: Optional[str] = None
OpenAIFinishReason = Literal["stop", "length", "tool_calls", "content_filter", "function_call"]
class ChunkChoice(BaseModel):
"""A single choice in a streaming chunk."""
index: int = 0
delta: ChoiceDelta
finish_reason: Optional[OpenAIFinishReason] = None
logprobs: Optional[dict] = 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
refusal: Optional[str] = None
class CompletionChoice(BaseModel):
"""A single choice in a non-streaming response."""
index: int = 0
message: CompletionMessage
finish_reason: OpenAIFinishReason = "stop"
logprobs: Optional[dict] = None
class CompletionUsage(BaseModel):
"""Token usage statistics (approximate)."""
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
prompt_tokens_details: Optional[dict] = Field(
default_factory = lambda: {"cached_tokens": 0, "audio_tokens": 0}
)
completion_tokens_details: Optional[dict] = Field(
default_factory = lambda: {
"reasoning_tokens": 0,
"audio_tokens": 0,
"accepted_prediction_tokens": 0,
"rejected_prediction_tokens": 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)
system_fingerprint: Optional[str] = None
# =====================================================================
# 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", "original"]] = "auto"
class ResponsesOutputTextPart(BaseModel):
"""Assistant ``output_text`` content part replayed on subsequent turns.
Clients looping on a stateless Responses endpoint round-trip prior assistant
messages as ``output_text`` parts; we keep the text and ignore the
annotations/logprobs when flattening into Chat Completions.
"""
type: Literal["output_text"]
text: str
annotations: Optional[list] = None
logprobs: Optional[list] = None
model_config = {"extra": "allow"}
class ResponsesUnknownContentPart(BaseModel):
"""Catch-all for unmodelled content-part types.
Keeps validation green for newer part types (e.g. ``input_audio``); skipped
during normalisation rather than rejected with a 422.
"""
type: str
model_config = {"extra": "allow"}
ResponsesContentPart = Union[
ResponsesInputTextPart,
ResponsesInputImagePart,
ResponsesOutputTextPart,
ResponsesUnknownContentPart,
]
class ResponsesInputMessage(BaseModel):
"""A single message in the Responses API input array."""
type: Optional[Literal["message"]] = None
role: Literal["system", "user", "assistant", "developer"]
content: Union[str, list[ResponsesContentPart]]
# Codex attaches a `phase` field to assistant messages and requires clients
# to preserve it across turns; we round-trip it, llama-server ignores it.
model_config = {"extra": "allow"}
class ResponsesFunctionCallInputItem(BaseModel):
"""A prior assistant function_call replayed in a multi-turn Responses input.
Tool calls are top-level input items (not nested), correlated by ``call_id``.
"""
type: Literal["function_call"]
id: Optional[str] = Field(None, description = "Item id assigned by the server (e.g. fc_...)")
call_id: str = Field(
...,
description = "Correlation id matching a function_call_output on the next turn.",
)
name: str
arguments: str = Field(..., description = "JSON string of the arguments the model produced.")
status: Optional[Literal["in_progress", "completed", "incomplete"]] = None
class ResponsesFunctionCallOutputInputItem(BaseModel):
"""A tool result supplied by the client for a prior function_call.
Replaces Chat Completions' ``role="tool"`` message. Correlated to its
originating call by ``call_id``.
"""
type: Literal["function_call_output"]
id: Optional[str] = None
call_id: str
output: Union[str, list] = Field(
..., description = "String or content-array result of the tool call."
)
status: Optional[Literal["in_progress", "completed", "incomplete"]] = None
class ResponsesUnknownInputItem(BaseModel):
"""Catch-all for unmodelled Responses input item types.
Covers ``reasoning`` items and future types. Dropped during normalisation
(GGUFs can't consume them), but kept in the union so unrelated turns don't 422.
"""
type: str
model_config = {"extra": "allow"}
def _responses_input_item_discriminator(v: Any) -> str:
"""Route a Responses input item to the correct tagged variant.
Pydantic's smart-union matching misreports errors when a strict-``Literal``
variant doesn't match; an explicit discriminator makes routing deterministic
and falls through to the catch-all.
"""
if isinstance(v, dict):
t = v.get("type")
r = v.get("role")
else:
t = getattr(v, "type", None)
r = getattr(v, "role", None)
if t == "function_call":
return "function_call"
if t == "function_call_output":
return "function_call_output"
if r is not None or t == "message":
return "message"
return "unknown"
ResponsesInputItem = Annotated[
Union[
Annotated[ResponsesInputMessage, Tag("message")],
Annotated[ResponsesFunctionCallInputItem, Tag("function_call")],
Annotated[ResponsesFunctionCallOutputInputItem, Tag("function_call_output")],
Annotated[ResponsesUnknownInputItem, Tag("unknown")],
],
Discriminator(_responses_input_item_discriminator),
]
class ResponsesFunctionTool(BaseModel):
"""Flat function-tool definition for the Responses API request.
Unlike Chat Completions (nested under a ``"function"`` key), this uses a flat
shape with ``type``/``name``/``description``/``parameters``/``strict`` at top level.
"""
type: Literal["function"]
name: str
description: Optional[str] = None
parameters: Optional[dict] = None
strict: Optional[bool] = None
class ResponsesRequest(BaseModel):
"""OpenAI Responses API request."""
model: str = Field("default", description = "Model identifier")
input: Union[str, list[ResponsesInputItem]] = Field(
default = [],
description = "Input text or list of messages / function_call / function_call_output items",
)
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")
# OpenAI function-calling fields, forwarded via the Chat Completions
# pass-through. Plain list so built-in tool shapes round-trip without
# validation errors; the translator forwards only ``type=="function"`` entries.
tools: Optional[list[dict]] = Field(
None,
description = (
"Responses-shape function tool definitions. Entries with "
'`type="function"` are translated to the Chat Completions nested '
"shape before being forwarded to llama-server; other tool types "
"(built-in web_search, file_search, mcp, ...) are accepted for SDK "
"compatibility but ignored on the llama-server passthrough."
),
)
tool_choice: Optional[Any] = Field(
None,
description = (
"'auto' | 'required' | 'none' | {'type': 'function', 'name': ...} — "
"the Responses-shape forcing object is translated to the Chat "
"Completions nested shape internally."
),
)
parallel_tool_calls: Optional[bool] = 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 ResponsesOutputReasoningContent(BaseModel):
"""A reasoning text content block inside a reasoning output item."""
type: Literal["reasoning_text"] = "reasoning_text"
text: str
class ResponsesOutputReasoning(BaseModel):
"""A top-level reasoning output item in the Responses API response."""
type: Literal["reasoning"] = "reasoning"
id: str = Field(default_factory = lambda: f"rs_{uuid.uuid4().hex[:12]}")
status: Literal["completed", "in_progress", "incomplete"] = "completed"
summary: list = Field(default_factory = list)
content: Optional[list[ResponsesOutputReasoningContent]] = None
class ResponsesOutputFunctionCall(BaseModel):
"""A function-call output item in the Responses API response.
Each tool call is its own top-level ``output`` item, correlated via ``call_id``.
"""
type: Literal["function_call"] = "function_call"
id: str = Field(default_factory = lambda: f"fc_{uuid.uuid4().hex[:12]}")
call_id: str
name: str
arguments: str = Field(..., description = "JSON string of the arguments the model produced.")
status: Literal["completed", "in_progress", "incomplete"] = "completed"
ResponsesOutputItem = Union[
ResponsesOutputMessage,
ResponsesOutputReasoning,
ResponsesOutputFunctionCall,
]
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[ResponsesOutputItem] = 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,
]
def _anthropic_content_to_system_text(content: Any) -> str:
"""Convert misplaced system message content into Anthropic system text."""
if content is None: # null content must not become the literal "None"
return ""
if isinstance(content, str):
return content
if isinstance(content, list):
parts: list[str] = []
for block in content:
if isinstance(block, dict) and block.get("type") == "text":
text = block.get("text")
if isinstance(text, str):
parts.append(text)
continue
if block is not None:
parts.append(str(block))
return "\n\n".join(part for part in parts if part)
return str(content)
def _merge_anthropic_system(system: Any, additions: list[str]) -> Any:
if not additions:
return system
addition_blocks = [{"type": "text", "text": text} for text in additions if text.strip()]
if not addition_blocks:
return system
if system is None:
return addition_blocks[0]["text"] if len(addition_blocks) == 1 else addition_blocks
if isinstance(system, str):
return "\n\n".join([system, *[block["text"] for block in addition_blocks]])
if isinstance(system, list):
return [*system, *addition_blocks]
return system
class AnthropicMessage(BaseModel):
role: Literal["user", "assistant"]
content: Union[str, list[AnthropicContentBlock]]
class AnthropicTool(BaseModel):
# Client tools have input_schema; server tools may only have type/name.
type: Optional[str] = None
name: Optional[str] = None
description: Optional[str] = None
input_schema: Optional[dict] = None
model_config = {"extra": "allow"}
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 mirroring 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
cancel_id: Optional[str] = None
bypass_permissions: Optional[bool] = Field(
False,
description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.",
)
model_config = {"extra": "allow"}
@model_validator(mode = "before")
@classmethod
def normalize_system_messages(cls, data: Any) -> Any:
if not isinstance(data, dict):
return data
messages = data.get("messages")
if not isinstance(messages, list):
return data
normalized_messages: list[Any] = []
system_additions: list[str] = []
changed = False
for message in messages:
if isinstance(message, dict) and message.get("role") == "system":
system_additions.append(
_anthropic_content_to_system_text(message.get("content", ""))
)
changed = True
continue
normalized_messages.append(message)
if not changed:
return data
normalized = dict(data)
normalized["messages"] = normalized_messages
normalized["system"] = _merge_anthropic_system(normalized.get("system"), system_additions)
return normalized
# ── Response models ────────────────────────────────────────────
class AnthropicUsage(BaseModel):
input_tokens: int = 0
cache_creation_input_tokens: int = 0
cache_read_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)