unsloth/studio/backend/core/rag
Daniel Han 414503745e
Run the malware gate on the RAG embedding model before it loads (#6887)
* Run the malware gate on the RAG embedding model before it loads

Setting the RAG embedding model through PUT /api/settings/embedding-model
persisted an arbitrary repo and later handed it straight to
SentenceTransformer, which deserializes pickle weights. Unlike the normal
model-load paths, this route never ran evaluate_file_security, and force
skipped verification entirely, so a repo Hugging Face flags as unsafe (or
any repo under force) could be downloaded and loaded in the backend
process without a scan.

Run the malware/pickle scan at both ends: the settings endpoint now scans
before persisting and returns 409 on a flagged repo even under force
(force still only skips the is-embedding-model type check for offline or
local repos), and the embedder scans again at the load sink so a name that
arrives via env or default is covered too. Local paths and unreachable
scans fail open inside evaluate_file_security, and the sink never bricks
the embedder on a gate error.

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

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

* Thread the load token into the embedding scan and hard-fail on a block

The load-sink scan ran without a token, so evaluate_file_security (which
passes token=False when none is given) could not reach a gated or private
repo and failed open for exactly the model SentenceTransformer would still
load. Resolve the loader's own token (HF_TOKEN env or the cached login)
and pass it to the sink scan, and fall back to it in the settings endpoint
when the request omits one.

The sink previously raised a plain RuntimeError, which the llama-server
fallback in encode() and _build_st_backend_or_fallback() swallowed as a
routine ST failure, silently switching backends instead of blocking. Raise
a distinct UnsafeEmbeddingModelError that both fallback paths re-raise, so
a flagged model hard-fails.

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

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

* Scan sentence-transformers module dirs and scope the embedding pickle gate to the ST backend

Extend the RAG embedding malware gate so a poisoned pickle under a SentenceTransformer
module dir (for example 0_Transformer/pytorch_model.bin) blocks. Those dirs are read
from the repo's modules.json and passed as load roots to evaluate_file_security at both
the settings endpoint and the load sink, so such a pickle is treated as root-level there
instead of an unreferenced nested shard that was previously allowed.

Scope the ST pickle scan to the sentence-transformers backend. On the llama-server
backend the embedder loads GGUF files (inert) from the -GGUF companion repo, never the
ST repo's pickle, so a custom ST repo with a flagged pickle and a clean GGUF companion
is no longer rejected. The existing GGUF availability checks already cover that path.

Return 403 for the hard security block instead of 409. The settings UI routes every 409
into the forceable save-anyway flow, but this block cannot be bypassed by force, so it
now uses a distinct status the client treats as non-forceable.

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

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

* Base the embedding pickle scan on the actual backend, not just the resolver

_llama_backend_active only consulted the auto resolver, so on a GPU box
where auto resolves to sentence-transformers but the process already fell
back to the llama-server backend at runtime (a torch or CUDA load/encode
failure), it returned False and the settings endpoint hard-blocked a save
whose ST pickle is flagged even though the process loads only inert GGUF.

Add active_backend_is_llama, which reflects the actual built backend (True
when the cached backend is a LlamaServerBackend, including a runtime
fallback) and otherwise defers to the resolver as a fresh process would,
and delegate _llama_backend_active to it.

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

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

* Report the cached embedding backend verbatim, not the resolver

active_backend_is_llama() fell through to the config resolver whenever a
backend was already built but was not llama-server, so a live
sentence-transformers backend could report llama=True once the resolver
picked llama (GPU heuristic or a runtime config change) and wrongly skip
its pickle scan. Once a backend exists, return isinstance(backend,
LlamaServerBackend) directly; only defer to the resolver before any
backend is built.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-07 04:30:21 -07:00
..
__init__.py Studio: clean-room compact RAG (knowledge bases, hybrid search, fast indexing) (#5910) 2026-06-09 21:17:04 -07:00
captioner.py Studio RAG: disable trust_env on loopback llama-server httpx clients (#6775) 2026-07-01 03:58:08 -07:00
chunking.py Studio: clean-room compact RAG (knowledge bases, hybrid search, fast indexing) (#5910) 2026-06-09 21:17:04 -07:00
config.py Studio: customizable RAG embedding model with HF search, settings tab reorganization (#6800) 2026-07-02 05:26:33 -07:00
embed_llama_server.py Studio: customizable RAG embedding model with HF search, settings tab reorganization (#6800) 2026-07-02 05:26:33 -07:00
embeddings.py Run the malware gate on the RAG embedding model before it loads (#6887) 2026-07-07 04:30:21 -07:00
ingestion.py Studio: customizable RAG embedding model with HF search, settings tab reorganization (#6800) 2026-07-02 05:26:33 -07:00
locators.py Whole-document context for RAG chat attachments (#6693) 2026-06-30 15:55:23 +02:00
parsers.py Studio RAG: fix RTL/Indic PDF corruption and dropped DOCX tables (#6780) 2026-07-02 11:38:48 +01:00
retrieval.py Studio: customizable RAG embedding model with HF search, settings tab reorganization (#6800) 2026-07-02 05:26:33 -07:00
store.py Studio: customizable RAG embedding model with HF search, settings tab reorganization (#6800) 2026-07-02 05:26:33 -07:00
tool.py Whole-document context for RAG chat attachments (#6693) 2026-06-30 15:55:23 +02:00