Commit graph

5 commits

Author SHA1 Message Date
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
Michael Han
4f24b12cc9
Studio: customizable RAG embedding model with HF search, settings tab reorganization (#6800)
* Add customizable RAG embedding model setting and reorganize settings tabs

Chat with files, project sources, and knowledge bases previously always
embedded with unsloth/bge-small-en-v1.5. This adds a Settings option to
pick any Hugging Face embedding model (or local path), with HF search
autocomplete, server-side verification that the repo is actually an
embedding model, and a save anyway escape hatch for offline or local
models. The setting persists in app_settings and applies at runtime to
both the sentence-transformers and llama-server GGUF embedder backends
without a restart.

Also reorganizes the General settings tab: Documents & RAG sits above
Uploads, Helper LLM moved above the danger zone, and Model auto-switch
(OpenAI API) moved to the bottom of the API tab.

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

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

* Support local model paths on the GGUF embedder and normalize default saves

Found by simulation testing of the embedding model setting:

Local paths saved as the embedding model now work on the llama-server
GGUF backend (the default backend on macOS and CPU). A path to a .gguf
file is used directly and a directory is scanned for a variant-matching
non-mmproj .gguf, with a clear error when none exists. Previously a
local path was sent to the HF hub API and failed with a repo lookup
error.

Saving the default model explicitly no longer stores an override, so
is_custom stays false and the UI does not show a reset button for the
default value.

* Address review: stale-vector handling, GGUF derivation, save-time guards

Review follow-ups, each verified by new tests:

Re-uploading a document after an embedding model change now re-indexes
instead of deduping by content hash. Documents record the embedder that
produced their vectors (lazy embedding_model column, NULL legacy rows
keep deduping) and a mismatch replaces the old document.

A vector width change no longer bricks the dense index. ensure_vec
drops and recreates chunks_vec when the dim changes (old vectors are in
a foreign space and only block inserts) and search_dense returns empty
on a width mismatch instead of surfacing a vec0 error, so lexical
search keeps working until documents are re-uploaded.

Saving a local sentence-transformers folder with no .gguf now returns
409 with a clear message when the install embeds via llama-server,
instead of failing at first index. force still saves.

A custom RAG_EMBEDDING_MODEL env without RAG_EMBED_GGUF_REPO now
derives the -GGUF companion repo instead of silently keeping the bge
GGUF on CPU and macOS installs.

The resolved GGUF path is tagged with the repo captured at entry, so a
setting change during a download cannot mark the old model as current.

GGUF repo detection matches gguf as a whole name segment rather than a
substring, hf_token is trimmed before verification, and the settings
combobox drops a redundant state mirror of its controlled value.

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

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

* Shrink embedding model font to 11px in the input and dropdown

The combobox wrapper applies className to the outer input group, so the
size utility must target the inner input element; the previous text-xs
never reached it and the field rendered at the browser default.

* Show curated unsloth embedding models when the search field is empty

The empty-query listing was the global top-downloads page, which holds
no unsloth mirrors for the unsloth-first float to reorder, so the
dropdown opened on third-party models. Match the model picker: curated
unsloth listing when empty, whole-Hub search once a query is typed.

* Address review: settings resilience and index consistency

Keep the last known embedding model on settings store errors, remove the
re-entrant dim lock in the llama-server backend, accept local GGUF saves
and verify GGUF availability for HF repos on that backend, match local
path embedders exactly in model list filters, drop same-width stale
vectors from dense search, pin the embedder per ingestion job, and only
replace completed documents after the re-index succeeds.

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

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

* Consolidate the GGUF repo derivation tests

* Trim to a single core embedding-model test

* Address review: GGUF repo saves and cache race

Accept a GGUF-named HF repo on the llama-server backend by verifying GGUF
availability instead of the sentence-transformers metadata gate, and guard
the settings cache with a generation counter so a read overlapping a save
cannot repopulate it with the pre-save value.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-02 05:26:33 -07:00
Daniel Han
935f6c50ef
studio: tighten torchao Windows-ROCm comments and test docstrings (#6610) 2026-06-23 05:49:25 -07:00
Daniel Han
55c392ff7c
studio: fix sentence-transformers RAG embedder on Windows ROCm (torchao) (#6608)
torchao has no working Windows ROCm build. transformers.quantizers imports it,
and it loads torch's c10d distributed backend at module level, which the AMD
Windows wheels omit (no RCCL). The import aborts, transformers can no longer
expose PreTrainedModel, and the sentence-transformers embedder silently falls
back to the llama-server GGUF embedder. Linux ROCm and NVIDIA are unaffected
(the c10d ops are present / torchao is real there).

The training and export workers already install the shared torchao stub before
importing transformers, but the RAG embedder runs in the main backend process,
which never did. Two fixes, both no-ops off Windows ROCm:

- embeddings.py: install_torchao_windows_rocm_stub() before the first
  sentence-transformers import, so an already-installed torchao is neutralized
  (fixes existing venvs).
- install_python_stack.py: stop installing torchao on Windows ROCm; it can only
  crash on import there, so new venvs never ship it.

Add tests covering the embedder stub call and the install skip.
2026-06-23 05:39:02 -07:00
Daniel Han
8848a310df
Studio: clean-room compact RAG (knowledge bases, hybrid search, fast indexing) (#5910)
Adds a self-contained RAG stack to Studio: knowledge bases with chunked indexing, hybrid (dense + lexical) retrieval, and an automatic first-pass context inject into chat. Embeddings run through a local llama-server GGUF backend (default unsloth/bge-small-en-v1.5-GGUF) with a sentence-transformers fallback. The chat tool loop gains a search_knowledge_base tool, a per-turn re-search cap, and source citation, layered on top of the shared ToolLoopController.
2026-06-09 21:17:04 -07:00