Studio: trim verbose comments/docstrings across RAG code
This commit is contained in:
Roland Tannous
parent
3b477a816c
commit
86b52503dd
34 changed files with 154 additions and 833 deletions
|
|
@ -503,19 +503,13 @@ TERMINAL_TOOL = {
|
|||
}
|
||||
|
||||
|
||||
# Lazy import — keeps studio.db init lazy so tools.py doesn't pull in
|
||||
# the whole rag stack on inference paths that never see RAG.
|
||||
# Lazy import: don't pull rag stack on inference paths that never see RAG.
|
||||
def _get_rag_tool_spec():
|
||||
from core.rag.tool import SEARCH_KNOWLEDGE_BASE_TOOL
|
||||
|
||||
return SEARCH_KNOWLEDGE_BASE_TOOL
|
||||
|
||||
|
||||
# RAG_SEARCH_TOOL is included in ALL_TOOLS; routes/inference.py filters
|
||||
# the list against payload.enabled_tools so each request only sees the
|
||||
# tools the frontend explicitly enabled. When the RAG button is off
|
||||
# the tool name won't be in enabled_tools and the LLM will never see
|
||||
# the spec.
|
||||
ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL, _get_rag_tool_spec()]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,17 +1,9 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Per-scope BM25 lexical index using the ``bm25s`` library.
|
||||
"""Per-scope BM25 index (rebuild on change; bm25s has no cheap incremental insert).
|
||||
|
||||
bm25s does not support incremental insertion cheaply, so we rebuild the
|
||||
full per-scope index whenever its document set changes. At studio
|
||||
scale (a few hundred to a few tens of thousands of chunks per KB) this
|
||||
is fast enough; the upside is that deletes are trivial.
|
||||
|
||||
A scope is ``kb_<uuid>`` or ``thread_<uuid>``. Each scope stores:
|
||||
- ``<scope>/`` directory holding the bm25s index files
|
||||
- ``<scope>/ids.json`` mapping the index's positional ids back to
|
||||
chunk-id strings (bm25s returns row indices, not our ids).
|
||||
Each scope dir holds the bm25s files + ids.json mapping row index → chunk_id.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -48,11 +40,7 @@ def _evict(scope: str) -> None:
|
|||
|
||||
|
||||
def rebuild_index(scope: str, chunks: list[dict]) -> None:
|
||||
"""Rebuild the BM25 index for ``scope`` from the full chunk list.
|
||||
|
||||
Each chunk dict must contain ``id`` and ``text``. Passing an empty
|
||||
list deletes the scope's index files.
|
||||
"""
|
||||
"""Rebuild scope's BM25 from full chunk list. Empty list deletes the index."""
|
||||
import bm25s
|
||||
|
||||
base = _scope_dir(scope)
|
||||
|
|
@ -64,9 +52,7 @@ def rebuild_index(scope: str, chunks: list[dict]) -> None:
|
|||
tokens = bm25s.tokenize(texts, show_progress = False)
|
||||
retriever = bm25s.BM25()
|
||||
retriever.index(tokens, show_progress = False)
|
||||
# Drop stale files from any previous build/library version before
|
||||
# writing the new index so we never mix old + new artifacts in the
|
||||
# scope dir (bm25s.BM25.save does not unlink files it does not write).
|
||||
# bm25s.BM25.save does not unlink stale files; clear the dir first.
|
||||
delete_scope(scope)
|
||||
ensure_dir(base)
|
||||
retriever.save(str(base))
|
||||
|
|
@ -87,10 +73,7 @@ def _load(scope: str) -> tuple[Any, list[str]] | None:
|
|||
retriever = bm25s.BM25.load(str(_scope_dir(scope)), load_corpus = False)
|
||||
ids = json.loads(_ids_path(scope).read_text())
|
||||
except (FileNotFoundError, OSError, json.JSONDecodeError, ValueError) as exc:
|
||||
# Corrupt or partially-written index files: treat as a
|
||||
# missing index so search returns empty and a future
|
||||
# re-ingest can rebuild cleanly. Log so the failure
|
||||
# is visible without crashing the request.
|
||||
# Corrupt/partial index: treat as missing so re-ingest rebuilds cleanly.
|
||||
logger.warning(
|
||||
"bm25 index unreadable for scope %s (%s: %s); treating as missing",
|
||||
scope,
|
||||
|
|
|
|||
|
|
@ -20,7 +20,6 @@ TokenCounter = Callable[[str], int]
|
|||
|
||||
|
||||
def _char_token_estimate(text: str) -> int:
|
||||
# Rough 4 chars / token heuristic — only used if no real tokenizer is provided.
|
||||
return max(1, (len(text) + 3) // 4)
|
||||
|
||||
|
||||
|
|
@ -60,7 +59,6 @@ def _atomic_split(
|
|||
tail = separators[separators.index(sep) + 1 :]
|
||||
out.extend(_atomic_split(piece, tail, max_tokens, count))
|
||||
return out
|
||||
# No separator made progress — hard-slice by characters.
|
||||
approx_chars = max(1, max_tokens * 4)
|
||||
return [text[i : i + approx_chars] for i in range(0, len(text), approx_chars)]
|
||||
|
||||
|
|
@ -71,7 +69,7 @@ def _merge(
|
|||
overlap_tokens: int,
|
||||
count: TokenCounter,
|
||||
) -> list[str]:
|
||||
"""Greedy-merge atomic pieces into chunks <= max_tokens with overlap between adjacent chunks."""
|
||||
"""Greedy-merge into <= max_tokens chunks with overlap."""
|
||||
chunks: list[str] = []
|
||||
buffer: list[str] = []
|
||||
buffer_tokens = 0
|
||||
|
|
@ -100,12 +98,8 @@ def _merge(
|
|||
return [c.strip() for c in chunks if c.strip()]
|
||||
|
||||
|
||||
# Markdown headings first so layout-aware parser output splits at sections.
|
||||
DEFAULT_SEPARATORS: tuple[str, ...] = (
|
||||
# Markdown heading boundaries first — when the parser emits
|
||||
# layout-aware Markdown (PDF via pymupdf4llm, DOCX via mammoth,
|
||||
# HTML via markdownify) chunks split at section breaks rather
|
||||
# than mid-paragraph. Falls back to the original separators on
|
||||
# plain text input where headings are absent.
|
||||
"\n# ",
|
||||
"\n## ",
|
||||
"\n### ",
|
||||
|
|
@ -126,11 +120,7 @@ def chunk_pages(
|
|||
token_counter: TokenCounter | None = None,
|
||||
separators: tuple[str, ...] = DEFAULT_SEPARATORS,
|
||||
) -> list[Chunk]:
|
||||
"""Split parsed pages into overlapping chunks.
|
||||
|
||||
Each page is split independently so page_number stays meaningful for
|
||||
PDFs — cross-page chunks would lose source attribution.
|
||||
"""
|
||||
"""Split pages independently so page_number stays attached to chunks."""
|
||||
count = token_counter or _char_token_estimate
|
||||
out: list[Chunk] = []
|
||||
for page in pages:
|
||||
|
|
@ -158,23 +148,11 @@ def chunk_pages_with_spans(
|
|||
token_counter: TokenCounter | None = None,
|
||||
separators: tuple[str, ...] = DEFAULT_SEPARATORS,
|
||||
) -> tuple[str, list[Chunk], list[tuple[int, int]]]:
|
||||
"""Late-chunking-friendly variant of :func:`chunk_pages`.
|
||||
"""Late-chunking variant: joins pages so the embedder sees the whole doc.
|
||||
|
||||
Joins all pages into a single document so the embedder sees the
|
||||
whole text in one pass (that's the point of late chunking — chunk
|
||||
vectors that carry full-document context via the model's
|
||||
bidirectional attention).
|
||||
|
||||
Returns ``(full_doc, chunks, char_spans)`` where
|
||||
``char_spans[i] = (start, end)`` are byte-character offsets of
|
||||
``chunks[i].text`` inside ``full_doc``. The embedder layer maps
|
||||
char spans → token spans via the tokenizer's offsets_mapping and
|
||||
mean-pools per chunk.
|
||||
|
||||
Page-number metadata on each :class:`Chunk` is recovered from the
|
||||
chunk's char span — the first page whose range overlaps the chunk
|
||||
wins. PDFs keep useful citations even though chunking ignores page
|
||||
boundaries here.
|
||||
Returns ``(full_doc, chunks, char_spans)``; ``char_spans[i]`` is the
|
||||
(start, end) char offset of ``chunks[i].text`` inside ``full_doc``.
|
||||
Page numbers are recovered by overlap with the original page ranges.
|
||||
"""
|
||||
count = token_counter or _char_token_estimate
|
||||
|
||||
|
|
@ -203,13 +181,9 @@ def chunk_pages_with_spans(
|
|||
continue
|
||||
idx = full_doc.find(text, search_cursor)
|
||||
if idx < 0:
|
||||
# Overlap can push the search cursor past a chunk's true
|
||||
# start — restart from the document head as a fallback.
|
||||
# Overlap can push past a chunk's true start; restart from head.
|
||||
idx = full_doc.find(text)
|
||||
if idx < 0:
|
||||
# Chunker output diverged from the source (rare — happens
|
||||
# if a separator-splice mangled the text). Skip the chunk
|
||||
# rather than corrupt the vector store with a wrong span.
|
||||
continue
|
||||
end_idx = idx + len(text)
|
||||
page_number = _page_for_span(idx, end_idx, page_ranges)
|
||||
|
|
@ -221,8 +195,7 @@ def chunk_pages_with_spans(
|
|||
)
|
||||
)
|
||||
char_spans.append((idx, end_idx))
|
||||
# Advance past the *start* of this chunk so an overlapping
|
||||
# next chunk can still be found.
|
||||
# Advance past start (not end) so overlapping next chunk is findable.
|
||||
search_cursor = idx + 1
|
||||
|
||||
return full_doc, chunks, char_spans
|
||||
|
|
|
|||
|
|
@ -1,13 +1,7 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""sqlite-vec backed connection helper for RAG vectors.
|
||||
|
||||
Single process-wide connection, opened lazily on first use. The
|
||||
extension-load step runs once at open time. studio.db (chat history,
|
||||
RAG metadata) stays untouched, so the extension-load surface is scoped
|
||||
to the RAG code path only — chat code keeps its plain sqlite handle.
|
||||
"""
|
||||
"""Lazy process-wide rag.db connection with sqlite-vec loaded."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -29,14 +23,6 @@ def rag_db_path() -> Path:
|
|||
|
||||
|
||||
def _load_sqlite_vec(conn: sqlite3.Connection) -> None:
|
||||
"""Enable extension loading and pull in sqlite-vec.
|
||||
|
||||
install.sh creates the studio venv via `uv venv --python <ver>`,
|
||||
which uses uv's managed python-build-standalone build. That CPython
|
||||
is compiled with --enable-loadable-sqlite-extensions, so this path
|
||||
succeeds on standard installs. The actionable error message is
|
||||
here for the rare custom-interpreter case.
|
||||
"""
|
||||
try:
|
||||
conn.enable_load_extension(True)
|
||||
except AttributeError as exc:
|
||||
|
|
@ -77,12 +63,6 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
|
|||
|
||||
|
||||
def get_rag_connection() -> sqlite3.Connection:
|
||||
"""Lazy process-wide sqlite connection to rag.db with sqlite-vec loaded.
|
||||
|
||||
Returns the cached connection on subsequent calls. FastAPI's thread
|
||||
pool plus check_same_thread=False + WAL mode handles concurrent
|
||||
reads; writes are serialized by SQLite itself.
|
||||
"""
|
||||
global _conn
|
||||
with _conn_lock:
|
||||
if _conn is None:
|
||||
|
|
@ -101,7 +81,6 @@ def get_rag_connection() -> sqlite3.Connection:
|
|||
|
||||
|
||||
def _reset_for_tests() -> None:
|
||||
"""Drop the cached connection. Test-only — production never calls."""
|
||||
global _conn
|
||||
with _conn_lock:
|
||||
if _conn is not None:
|
||||
|
|
|
|||
|
|
@ -1,15 +1,7 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Embedding model singleton for RAG.
|
||||
|
||||
Loads the configured embedder via Unsloth's ``FastSentenceTransformer``
|
||||
wrapper with ``for_inference=True`` (which returns a plain
|
||||
``sentence_transformers.SentenceTransformer`` instance with proper dtype
|
||||
and device handling). Lifecycle is fully independent of the chat
|
||||
``InferenceBackend`` so loading an embedder cannot evict the active
|
||||
chat model.
|
||||
"""
|
||||
"""RAG embedder singleton. Independent of the chat InferenceBackend."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -30,19 +22,13 @@ _embedding_dim: int | None = None
|
|||
def _load(model_name: str) -> Any:
|
||||
logger.info("Loading RAG embedder: %s", model_name)
|
||||
|
||||
# BGE-VL ships a sentence-transformers shim that's tightly coupled
|
||||
# to a specific ST internal API and breaks across ST version bumps.
|
||||
# Bypass ST entirely and load via the canonical transformers
|
||||
# AutoModel path, wrapped to match the SentenceTransformer API
|
||||
# slice the RAG ingester uses.
|
||||
# BGE-VL's ST shim breaks across ST versions; load via AutoModel.
|
||||
if model_name.startswith("BAAI/BGE-VL"):
|
||||
return _BGEVLAdapter(model_name)
|
||||
|
||||
from unsloth import FastSentenceTransformer
|
||||
|
||||
# trust_remote_code is required for nomic-embed-text-v1.5 (custom
|
||||
# modeling for 8K context). Safe to enable because the embedder
|
||||
# matrix is config-pinned — users don't supply arbitrary names.
|
||||
# trust_remote_code: nomic-embed-text-v1.5 needs custom modeling for 8K ctx.
|
||||
return FastSentenceTransformer.from_pretrained(
|
||||
model_name,
|
||||
for_inference = True,
|
||||
|
|
@ -51,18 +37,7 @@ def _load(model_name: str) -> Any:
|
|||
|
||||
|
||||
class _BGEVLAdapter:
|
||||
"""Adapter exposing the slice of SentenceTransformer API the RAG
|
||||
ingester depends on, backed by BGE-VL's transformers AutoModel.
|
||||
|
||||
Supports:
|
||||
- ``encode(list_of_strings, ...)`` → text embeddings
|
||||
- ``encode(list_of_PIL_images, ...)`` → image embeddings
|
||||
- ``get_sentence_embedding_dimension()``
|
||||
- ``tokenize([text])`` for token-aware chunking (best-effort)
|
||||
|
||||
Auto-detects image vs text inputs from the first element. Returns
|
||||
L2-normalized numpy arrays when ``normalize_embeddings=True``.
|
||||
"""
|
||||
"""SentenceTransformer-shaped adapter over BGE-VL's AutoModel."""
|
||||
|
||||
def __init__(self, hf_model_name: str):
|
||||
from transformers import AutoModel
|
||||
|
|
@ -72,9 +47,7 @@ class _BGEVLAdapter:
|
|||
hf_model_name,
|
||||
trust_remote_code = True,
|
||||
)
|
||||
# BGE-VL's encode() requires set_processor to install the
|
||||
# tokenizer / image processor on the model. Without it, the
|
||||
# first encode() raises with a missing-processor error.
|
||||
# Required: BGE-VL's encode() raises without an installed processor.
|
||||
self._model.set_processor(hf_model_name)
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
self._model.to(device).eval()
|
||||
|
|
@ -86,10 +59,7 @@ class _BGEVLAdapter:
|
|||
|
||||
return F.normalize(tensor, p = 2.0, dim = -1)
|
||||
|
||||
# CLIP-family text encoder context cap. BGE-VL inherits CLIP's
|
||||
# 77-token text positional embedding table — exceeding it triggers
|
||||
# a shape-mismatch in the embedding layer. Pre-truncate any text
|
||||
# chunk to this length before calling the model.
|
||||
# CLIP positional embedding cap; longer text triggers shape mismatch.
|
||||
_CLIP_TEXT_MAX_TOKENS = 77
|
||||
|
||||
def encode(
|
||||
|
|
@ -140,13 +110,7 @@ class _BGEVLAdapter:
|
|||
return out.numpy() if convert_to_numpy else out
|
||||
|
||||
def _encode_text_truncated(self, texts: list[str]):
|
||||
"""Tokenize with explicit truncation to CLIP's 77-token limit, then
|
||||
call get_text_features directly. BGE-VL's high-level encode() does
|
||||
not truncate, so longer chunks overflow the position-embedding
|
||||
table and crash inside the text model. Text chunks beyond the cap
|
||||
are silently truncated — the multimodal mode is primarily about
|
||||
the image side; large text chunks should land in text-mode RAG.
|
||||
"""
|
||||
"""Truncate to CLIP's 77-token limit; long text in multimodal mode is lossy."""
|
||||
import torch
|
||||
|
||||
tokenizer = self._get_text_tokenizer()
|
||||
|
|
@ -158,14 +122,7 @@ class _BGEVLAdapter:
|
|||
max_length = self._CLIP_TEXT_MAX_TOKENS,
|
||||
)
|
||||
inputs = {k: v.to(self._device) for k, v in inputs.items()}
|
||||
# Some chunks were almost certainly truncated; flag it once per
|
||||
# batch so users running multimodal on long-form text know the
|
||||
# text channel is lossy by design.
|
||||
overflowed = any(
|
||||
len(t.split()) > 30 # ~rough proxy; tokens vary by lang
|
||||
for t in texts
|
||||
)
|
||||
if overflowed:
|
||||
if any(len(t.split()) > 30 for t in texts):
|
||||
logger.info(
|
||||
"BGE-VL text encode: truncating chunks to %d tokens (CLIP cap)",
|
||||
self._CLIP_TEXT_MAX_TOKENS,
|
||||
|
|
@ -174,7 +131,6 @@ class _BGEVLAdapter:
|
|||
return self._model.get_text_features(**inputs)
|
||||
|
||||
def _get_text_tokenizer(self):
|
||||
"""Locate the tokenizer set up by ``set_processor`` for text input."""
|
||||
processor = getattr(self._model, "processor", None)
|
||||
if processor is not None:
|
||||
tok = getattr(processor, "tokenizer", None)
|
||||
|
|
@ -192,11 +148,6 @@ class _BGEVLAdapter:
|
|||
return self._dim
|
||||
|
||||
def tokenize(self, texts):
|
||||
"""Best-effort tokenize for the token_counter chunking path.
|
||||
|
||||
Falls back gracefully — the caller already handles exceptions
|
||||
by approximating tokens as ``len(text) // 4`` when this raises.
|
||||
"""
|
||||
return self._get_text_tokenizer()(
|
||||
texts,
|
||||
return_tensors = "pt",
|
||||
|
|
@ -205,7 +156,6 @@ class _BGEVLAdapter:
|
|||
|
||||
|
||||
def get_embedder(model_name: str | None = None) -> Any:
|
||||
"""Return the cached SentenceTransformer, loading it on first use."""
|
||||
global _model, _model_name, _embedding_dim
|
||||
target = model_name or RAG_EMBEDDING_MODEL
|
||||
with _lock:
|
||||
|
|
@ -255,14 +205,7 @@ def encode_images(
|
|||
batch_size: int | None = None,
|
||||
normalize: bool = True,
|
||||
):
|
||||
"""Embed raw image bytes via a multimodal SentenceTransformer.
|
||||
|
||||
Works with CLIP-family models (BGE-VL, openai/clip-*) whose
|
||||
`encode` accepts PIL.Image objects in the same call as text. The
|
||||
returned vectors live in the same 512-d (or model-specific) space
|
||||
as text vectors from this model, so a single scope's vector rows
|
||||
hold both kinds.
|
||||
"""
|
||||
"""Embed image bytes via a CLIP-family multimodal encoder."""
|
||||
from io import BytesIO
|
||||
|
||||
from PIL import Image
|
||||
|
|
@ -281,11 +224,7 @@ def encode_images(
|
|||
|
||||
|
||||
def token_counter(model_name: str | None = None):
|
||||
"""Return a ``len(tokenize(text))`` callable using the embedder's tokenizer.
|
||||
|
||||
Avoid loading the model just for chunking by reaching through the
|
||||
SentenceTransformer's ``tokenize`` API.
|
||||
"""
|
||||
"""Return a token-count callable backed by the embedder's tokenizer."""
|
||||
model = get_embedder(model_name)
|
||||
|
||||
def _count(text: str) -> int:
|
||||
|
|
@ -301,9 +240,7 @@ def token_counter(model_name: str | None = None):
|
|||
return _count
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Late chunking (Phase 3B-late)
|
||||
# ------------------------------------------------------------------
|
||||
# --- Late chunking (Jina technique) ---
|
||||
|
||||
_LATE_WINDOW_OVERLAP_TOKENS = 512
|
||||
|
||||
|
|
@ -315,18 +252,7 @@ def late_chunk_encode(
|
|||
model_name: str | None = None,
|
||||
normalize: bool = True,
|
||||
):
|
||||
"""Embed each chunk via late-chunking pooling.
|
||||
|
||||
Single forward pass over the full document, then mean-pool the
|
||||
token embeddings whose offset ranges fall inside each chunk's
|
||||
char span. Chunks therefore carry full-document context via the
|
||||
encoder's bidirectional attention — Jina's published technique,
|
||||
works with any encoder that exposes per-token outputs.
|
||||
|
||||
When the doc exceeds the embedder's context, falls back to
|
||||
windowed late chunking with a 512-token overlap between windows
|
||||
so cross-window context is partially preserved.
|
||||
"""
|
||||
"""Single forward pass over the doc, mean-pool token embeddings per chunk span."""
|
||||
import numpy as np
|
||||
|
||||
if not char_spans:
|
||||
|
|
@ -373,7 +299,6 @@ def late_chunk_encode(
|
|||
|
||||
|
||||
def _encode_tokens(model, encoded):
|
||||
"""Run the embedder's underlying transformer to get per-token last_hidden_state."""
|
||||
import torch
|
||||
|
||||
transformer = model[0].auto_model
|
||||
|
|
@ -395,15 +320,11 @@ def _pool_spans(
|
|||
doc_text: str,
|
||||
token_index_offset: int = 0,
|
||||
):
|
||||
"""Mean-pool token embeddings per (char_start, char_end) span.
|
||||
|
||||
`token_index_offset` shifts char_span-derived token indices into
|
||||
a sub-window's local frame (used by the windowed code path).
|
||||
"""
|
||||
"""Mean-pool token embeddings per (char_start, char_end) span."""
|
||||
vectors = []
|
||||
n_rows = token_embeddings.shape[0]
|
||||
for char_start, char_end in char_spans:
|
||||
# Special tokens (CLS / SEP) report offsets (0, 0) — exclude them.
|
||||
# Skip special tokens whose offsets are (0, 0).
|
||||
indices = [
|
||||
i - token_index_offset
|
||||
for i, (ts, te) in enumerate(offsets)
|
||||
|
|
@ -411,9 +332,6 @@ def _pool_spans(
|
|||
]
|
||||
indices = [i for i in indices if 0 <= i < n_rows]
|
||||
if not indices:
|
||||
# Fall back to a standalone encode of the chunk text — rare
|
||||
# (would mean tokenizer produced zero non-special tokens for
|
||||
# the span), but keeps the pipeline alive.
|
||||
vec = model.encode(
|
||||
doc_text[char_start:char_end],
|
||||
normalize_embeddings = normalize,
|
||||
|
|
@ -440,12 +358,7 @@ def _windowed_late_chunk_encode(
|
|||
normalize: bool,
|
||||
np_module,
|
||||
):
|
||||
"""Doc exceeds context window — slice into overlapping windows.
|
||||
|
||||
Each chunk is pooled against the window that contains the most of
|
||||
its tokens. The 512-token window overlap means chunks near a
|
||||
boundary still see context from both sides.
|
||||
"""
|
||||
"""Doc > ctx window: pool each chunk against the window containing most of its tokens."""
|
||||
import torch
|
||||
|
||||
tokenizer = model.tokenizer
|
||||
|
|
@ -464,7 +377,6 @@ def _windowed_late_chunk_encode(
|
|||
n_tokens = int(all_input_ids.shape[0])
|
||||
stride = max(1, max_length - _LATE_WINDOW_OVERLAP_TOKENS)
|
||||
|
||||
# Build (start_token, end_token) windows.
|
||||
windows: list[tuple[int, int]] = []
|
||||
pos = 0
|
||||
while pos < n_tokens:
|
||||
|
|
@ -474,7 +386,6 @@ def _windowed_late_chunk_encode(
|
|||
break
|
||||
pos += stride
|
||||
|
||||
# Cache window → token embeddings (only encode when needed).
|
||||
window_embeddings: dict[int, "np_module.ndarray"] = {}
|
||||
|
||||
def _window_embeddings(window_index: int):
|
||||
|
|
@ -491,7 +402,6 @@ def _windowed_late_chunk_encode(
|
|||
|
||||
vectors = []
|
||||
for char_start, char_end in char_spans:
|
||||
# Collect global token indices in the chunk.
|
||||
chunk_token_indices = [
|
||||
i
|
||||
for i, (ts, te) in enumerate(all_offsets)
|
||||
|
|
@ -506,7 +416,6 @@ def _windowed_late_chunk_encode(
|
|||
)
|
||||
vectors.append(vec)
|
||||
continue
|
||||
# Pick the window covering the most of this chunk's tokens.
|
||||
best_window = 0
|
||||
best_overlap = 0
|
||||
for wi, (ws, we) in enumerate(windows):
|
||||
|
|
|
|||
|
|
@ -1,19 +1,10 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Document ingestion pipeline.
|
||||
"""RAG ingestion pipeline.
|
||||
|
||||
Follows the studio's existing job pattern (`core/data_recipe/jobs/manager.py`):
|
||||
spawn a fresh subprocess per job with ``mp.get_context("spawn")`` and stream
|
||||
progress events back over a queue. The subprocess does the heavy work
|
||||
(parse → chunk → load embedder → embed in batches) and ships
|
||||
``(chunks, vectors)`` batches back. The parent persists everything:
|
||||
sqlite rows, vector_store rows in rag.db, and (at job completion) a
|
||||
rebuilt BM25 index.
|
||||
|
||||
Only the parent process owns the rag.db connection (sqlite-vec loaded
|
||||
there); subprocesses never open it directly. This keeps search
|
||||
available throughout the lifetime of an ingestion job.
|
||||
Spawn-subprocess per job (parse/chunk/embed); parent persists chunks,
|
||||
vectors, and rebuilds BM25 on completion. Only the parent opens rag.db.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -44,9 +35,7 @@ _CTX = mp.get_context("spawn")
|
|||
_QUEUE_TIMEOUT_SECONDS = 300
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Subprocess worker
|
||||
# ------------------------------------------------------------------
|
||||
# --- Subprocess worker ---
|
||||
|
||||
_MIME_TO_EXT = {
|
||||
"image/png": ".png",
|
||||
|
|
@ -108,7 +97,6 @@ def _subprocess_worker(
|
|||
)
|
||||
return
|
||||
|
||||
# Standard chunking path (text + optional images for multimodal mode).
|
||||
text_count = _run_standard_chunking(
|
||||
pages = pages,
|
||||
chunk_size = chunk_size,
|
||||
|
|
@ -147,12 +135,7 @@ def _run_standard_chunking(
|
|||
out_queue,
|
||||
send_complete: bool = True,
|
||||
) -> int:
|
||||
"""Stream text chunks back to the parent. Returns the number streamed.
|
||||
|
||||
When called as part of the multimodal pipeline, `send_complete` is
|
||||
False because image chunks still need to be streamed before the
|
||||
document is marked complete.
|
||||
"""
|
||||
"""Stream text chunks; returns count. send_complete=False when images follow."""
|
||||
out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2})
|
||||
chunks = chunk_pages(
|
||||
pages,
|
||||
|
|
@ -207,17 +190,7 @@ def _stream_image_chunks(
|
|||
out_queue,
|
||||
first_index: int,
|
||||
) -> int:
|
||||
"""Save extracted images to disk and stream image+caption chunks.
|
||||
|
||||
Each image becomes an `image`-kind chunk; if the parser found an
|
||||
adjacent caption, a paired `caption`-kind chunk is also emitted.
|
||||
Pairs share a ``pair_group`` so the parent can link them via
|
||||
rag_chunks.linked_chunk_id.
|
||||
|
||||
Images are written to ``rag_uploads_root() / 'images' /
|
||||
<document_id> / img-NNN.<ext>`` so the parent can serve them via
|
||||
the static-image route without holding bytes in memory.
|
||||
"""
|
||||
"""Persist images, emit image+caption chunks; pairs share pair_group."""
|
||||
from core.rag.embeddings import encode, encode_images
|
||||
from utils.paths.storage_roots import ensure_dir, rag_uploads_root
|
||||
|
||||
|
|
@ -228,7 +201,6 @@ def _stream_image_chunks(
|
|||
|
||||
img_dir = ensure_dir(rag_uploads_root() / "images" / document_id)
|
||||
|
||||
# Persist bytes, build parallel lists for encoding.
|
||||
paths: list[str] = []
|
||||
bytes_for_encoding: list[bytes] = []
|
||||
captions: list[str] = []
|
||||
|
|
@ -251,7 +223,6 @@ def _stream_image_chunks(
|
|||
|
||||
image_vectors = encode_images(bytes_for_encoding, model_name = model_name)
|
||||
|
||||
# Embed only the non-empty captions; track which images they map to.
|
||||
caption_to_image: list[int] = [i for i, cap in enumerate(captions) if cap.strip()]
|
||||
if caption_to_image:
|
||||
caption_vectors_arr = encode(
|
||||
|
|
@ -280,8 +251,6 @@ def _stream_image_chunks(
|
|||
}
|
||||
)
|
||||
out_vectors.append(image_vectors[idx].tolist())
|
||||
# Emit the caption chunk right after its image so the parent
|
||||
# sees them adjacent (simplifies pair linking).
|
||||
if next_cap is not None and next_cap[0] == idx:
|
||||
_cap_index, cap_vec = next_cap
|
||||
out_chunks.append(
|
||||
|
|
@ -319,13 +288,7 @@ def _run_late_chunking(
|
|||
late_chunk_encode,
|
||||
out_queue,
|
||||
) -> None:
|
||||
"""Late chunking: chunk once over the whole doc, embed in a single pass.
|
||||
|
||||
There's no per-batch streaming here — the whole doc is encoded in
|
||||
one forward pass (or one per window for long docs). We ship all
|
||||
chunks back to the parent in one message; the parent's pump still
|
||||
handles them via the same chunks_batch handler.
|
||||
"""
|
||||
"""Chunk once, embed in one pass, ship all chunks in one chunks_batch."""
|
||||
from core.rag.chunking import chunk_pages_with_spans
|
||||
|
||||
out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2})
|
||||
|
|
@ -366,9 +329,7 @@ def _run_late_chunking(
|
|||
out_queue.put({"type": "complete", "num_chunks": len(chunks)})
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Job manager (parent side)
|
||||
# ------------------------------------------------------------------
|
||||
# --- Job manager (parent side) ---
|
||||
|
||||
|
||||
class _JobState:
|
||||
|
|
@ -450,13 +411,7 @@ def _insert_chunks_and_collect_for_bm25(
|
|||
chunks_meta: list[dict],
|
||||
vectors: list[list[float]],
|
||||
) -> list[dict]:
|
||||
"""Insert chunks into sqlite + vector_store; return [{id, text}] for BM25.
|
||||
|
||||
Image-kind chunks ship a stable image_path and skip BM25 (no text
|
||||
body to tokenise). Paired image/caption chunks share a pair_group
|
||||
field — the second pass links them via rag_chunks.linked_chunk_id
|
||||
so retrieval can dereference an image hit to its caption.
|
||||
"""
|
||||
"""Insert chunks into sqlite + vector_store; return [{id, text}] for BM25."""
|
||||
rows: list[tuple] = []
|
||||
points: list[dict] = []
|
||||
bm25_rows: list[dict] = []
|
||||
|
|
@ -496,9 +451,6 @@ def _insert_chunks_and_collect_for_bm25(
|
|||
},
|
||||
}
|
||||
)
|
||||
# BM25 indexes text + caption chunks. Image chunks have no
|
||||
# tokenisable body — their caption (if any) is in a separate
|
||||
# caption-kind chunk that BM25 will index.
|
||||
if kind in ("text", "caption") and meta["text"]:
|
||||
bm25_rows.append({"id": chunk_id, "text": meta["text"]})
|
||||
with get_connection() as conn:
|
||||
|
|
@ -511,9 +463,7 @@ def _insert_chunks_and_collect_for_bm25(
|
|||
""",
|
||||
rows,
|
||||
)
|
||||
# Link image ↔ caption pairs. We only set linked_chunk_id when
|
||||
# a pair_group has exactly two members; lone images stay
|
||||
# unlinked (no caption was paired).
|
||||
# Link only when exactly two members in a pair_group.
|
||||
for ids in pair_groups.values():
|
||||
if len(ids) != 2:
|
||||
continue
|
||||
|
|
@ -560,7 +510,7 @@ def _pump(
|
|||
proc: Any,
|
||||
out_queue: Any,
|
||||
) -> None:
|
||||
"""Drain queue messages until the subprocess signals complete/error or dies."""
|
||||
"""Drain queue until subprocess completes/errors/dies."""
|
||||
bm25_buffer: list[dict] = []
|
||||
embedding_dim: int | None = None
|
||||
final_status = "failed"
|
||||
|
|
@ -597,7 +547,6 @@ def _pump(
|
|||
vector_store.ensure_collection(state.scope, embedding_dim)
|
||||
elif mtype == "chunks_batch":
|
||||
if embedding_dim is None:
|
||||
# defensive: subprocess should always emit "dim" first
|
||||
embedding_dim = len(msg["vectors"][0]) if msg["vectors"] else None
|
||||
if embedding_dim is not None:
|
||||
vector_store.ensure_collection(state.scope, embedding_dim)
|
||||
|
|
@ -675,15 +624,7 @@ def enqueue_ingestion(
|
|||
chunking_strategy: str = "standard",
|
||||
mode: str = "text",
|
||||
) -> str:
|
||||
"""Create the job row, spawn the subprocess, and start the pump thread.
|
||||
|
||||
Returns the job_id. The caller can poll via ``GET /api/rag/jobs/{job_id}/events``
|
||||
or read the ``rag_ingestion_jobs`` table directly.
|
||||
|
||||
chunking_strategy / mode default to today's behaviour. KB-scoped
|
||||
uploads should pass the KB's stored values; per-thread uploads
|
||||
default unless an override is set in chat_settings.
|
||||
"""
|
||||
"""Create the job row, spawn the subprocess, start the pump; return job_id."""
|
||||
from utils.rag.config import resolve_embedder
|
||||
|
||||
scope = _scope_for(kb_id, thread_id)
|
||||
|
|
@ -736,11 +677,7 @@ def enqueue_ingestion(
|
|||
|
||||
|
||||
def delete_document_artifacts(document_id: str, scope: str) -> None:
|
||||
"""Remove a document's vectors, then rebuild BM25 for the scope.
|
||||
|
||||
The caller is responsible for the sqlite cascade (deleting the
|
||||
rag_documents row triggers ON DELETE CASCADE on rag_chunks).
|
||||
"""
|
||||
"""Drop the doc's vectors, rebuild BM25. Caller deletes the rag_documents row."""
|
||||
vector_store.delete_document(scope, document_id)
|
||||
remaining = _all_scope_chunks(scope)
|
||||
if remaining:
|
||||
|
|
@ -755,11 +692,7 @@ def delete_scope_artifacts(scope: str) -> None:
|
|||
|
||||
|
||||
def purge_thread_documents(thread_ids: list[str]) -> None:
|
||||
"""Remove all RAG artifacts owned by the given chat thread ids.
|
||||
|
||||
Used by the chat-thread DELETE handlers because rag_documents has
|
||||
no FK cascade to chat_threads (see schema comment).
|
||||
"""
|
||||
"""Drop RAG artifacts for the given thread ids (no FK cascade to chat_threads)."""
|
||||
if not thread_ids:
|
||||
return
|
||||
import os
|
||||
|
|
@ -791,7 +724,7 @@ def purge_thread_documents(thread_ids: list[str]) -> None:
|
|||
|
||||
|
||||
def purge_all_thread_documents() -> None:
|
||||
"""Drop every per-thread RAG artifact. Used by clear-all-history."""
|
||||
"""Drop every per-thread RAG artifact."""
|
||||
with get_connection() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT DISTINCT thread_id FROM rag_documents WHERE thread_id IS NOT NULL"
|
||||
|
|
|
|||
|
|
@ -9,15 +9,7 @@ from pathlib import Path
|
|||
|
||||
@dataclass(frozen = True)
|
||||
class ParsedPage:
|
||||
"""One page (or page-equivalent) of Markdown-rendered text from a source document.
|
||||
|
||||
For PDFs `page_number` is the 1-indexed physical page. For DOCX / HTML /
|
||||
TXT / MD the whole document is one ParsedPage with `page_number = None`.
|
||||
|
||||
Text is expected to be Markdown — heading markers (`#`, `##`, …),
|
||||
pipe-tables, and list bullets survive extraction so the chunker can
|
||||
split on them. Parsers MUST emit Markdown, not bare plain text.
|
||||
"""
|
||||
"""Markdown text from one page (PDF) or whole doc (others)."""
|
||||
|
||||
text: str
|
||||
page_number: int | None = None
|
||||
|
|
@ -25,14 +17,7 @@ class ParsedPage:
|
|||
|
||||
@dataclass(frozen = True)
|
||||
class ParsedImage:
|
||||
"""One image extracted from a source document.
|
||||
|
||||
Captured only when `parse(..., want_images=True)` is set — the
|
||||
multimodal ingestion path in Phase 3B-multimodal consumes these.
|
||||
`nearest_caption` is best-effort paragraph-adjacency; can be empty
|
||||
when no caption could be paired (the image still ingests, just
|
||||
without the paired-caption chunk).
|
||||
"""
|
||||
"""Image extracted with want_images=True; nearest_caption may be empty."""
|
||||
|
||||
image_bytes: bytes
|
||||
mime_type: str
|
||||
|
|
@ -42,13 +27,6 @@ class ParsedImage:
|
|||
|
||||
@dataclass(frozen = True)
|
||||
class ParseResult:
|
||||
"""Result of parsing a single source document.
|
||||
|
||||
`pages` is always populated; `images` is empty unless the caller
|
||||
passed `want_images=True`. Iteration aliases for `pages` so legacy
|
||||
code that did `for page in parse(path)` keeps working.
|
||||
"""
|
||||
|
||||
pages: list[ParsedPage] = field(default_factory = list)
|
||||
images: list[ParsedImage] = field(default_factory = list)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,17 +1,7 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""DOCX parsing via mammoth.
|
||||
|
||||
Mammoth converts Word documents to Markdown while preserving Heading
|
||||
styles (`# `, `## `, ...), bullet/numbered lists, tables, and basic
|
||||
emphasis. This replaces the previous python-docx paragraph-concat
|
||||
approach that lost all heading metadata.
|
||||
|
||||
Images are captured via the `convert_image` handler when
|
||||
`want_images=True`, falling back to python-docx for inline image bytes
|
||||
if mammoth's docx adapter can't reach them.
|
||||
"""
|
||||
"""DOCX → Markdown via mammoth (preserves headings, lists, tables)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -24,9 +14,7 @@ from . import ParsedImage, ParsedPage, ParseResult
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Mammoth uses some default DOCX-style-name → Markdown mappings, but a
|
||||
# few common variants ship with non-default names. Map them explicitly
|
||||
# so we don't lose headings.
|
||||
# Force common DOCX heading style names to Markdown headings.
|
||||
_STYLE_MAP = """
|
||||
p[style-name='Title'] => h1.title:fresh
|
||||
p[style-name='Subtitle'] => h2.subtitle:fresh
|
||||
|
|
@ -40,11 +28,9 @@ p[style-name='Heading 6'] => h6:fresh
|
|||
|
||||
|
||||
def _html_to_markdown(html: str) -> str:
|
||||
"""Convert mammoth's HTML output to Markdown via markdownify."""
|
||||
from markdownify import markdownify
|
||||
|
||||
md = markdownify(html, heading_style = "ATX", strip = ["script", "style"])
|
||||
# markdownify can emit excessive blank lines on tables; tighten up.
|
||||
md = re.sub(r"\n{3,}", "\n\n", md)
|
||||
return md.strip()
|
||||
|
||||
|
|
@ -55,12 +41,7 @@ def extract(path: Path, *, want_images: bool = False) -> ParseResult:
|
|||
images: list[ParsedImage] = []
|
||||
|
||||
if want_images:
|
||||
# mammoth's image converter is called for every inline image.
|
||||
# We capture the bytes here and substitute a stable placeholder
|
||||
# in the rendered Markdown so the chunker doesn't trip over
|
||||
# base64 blobs. Caption-pairing is approximate — we use the
|
||||
# full document text as the caption pool (better than nothing
|
||||
# for DOCX where heading→figure adjacency isn't reliable).
|
||||
# Capture bytes; suppress src so base64 doesn't land in Markdown.
|
||||
def _convert(image):
|
||||
with image.open() as image_bytes:
|
||||
blob = image_bytes.read()
|
||||
|
|
@ -77,8 +58,6 @@ def extract(path: Path, *, want_images: bool = False) -> ParseResult:
|
|||
|
||||
convert_image = mammoth.images.img_element(_convert)
|
||||
else:
|
||||
# Drop image elements entirely — cheaper and avoids embedding
|
||||
# base64 in Markdown when the caller doesn't want images.
|
||||
convert_image = mammoth.images.img_element(lambda _image: {"src": ""})
|
||||
|
||||
with open(path, "rb") as fp:
|
||||
|
|
@ -92,9 +71,7 @@ def extract(path: Path, *, want_images: bool = False) -> ParseResult:
|
|||
|
||||
markdown = _html_to_markdown(result.value)
|
||||
if want_images and images:
|
||||
# Best-effort: every image inherits the whole doc text as a
|
||||
# caption pool. Phase 3B-multimodal will improve this once
|
||||
# multimodal embedders consume captions directly.
|
||||
# Approximate caption: first 1500 chars of doc.
|
||||
caption_pool = markdown[:1500]
|
||||
images = [
|
||||
ParsedImage(
|
||||
|
|
|
|||
|
|
@ -1,18 +1,7 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""HTML parsing via markdownify.
|
||||
|
||||
Converts HTML to Markdown so headings (`<h1>`…`<h6>`), tables, and lists
|
||||
arrive at the chunker as Markdown structure. The previous
|
||||
`BeautifulSoup.get_text()` approach stripped all tags and made every
|
||||
heading indistinguishable from body text.
|
||||
|
||||
Image extraction (when `want_images=True`) only handles local file
|
||||
references — remote URLs are skipped to avoid network calls during
|
||||
ingestion. Phase 3B-multimodal can revisit this if HTML inputs with
|
||||
remote images become a common pattern.
|
||||
"""
|
||||
"""HTML → Markdown via markdownify. Images only resolve local file refs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -35,13 +24,12 @@ def _collect_local_images(soup, html_path: Path) -> list[ParsedImage]:
|
|||
src = tag.get("src") or ""
|
||||
parsed = urlparse(src)
|
||||
if parsed.scheme and parsed.scheme not in ("file", ""):
|
||||
# Remote / data URLs — skip; we don't fetch over network.
|
||||
continue
|
||||
local_path = (base_dir / unquote(parsed.path or src)).resolve()
|
||||
try:
|
||||
local_path.relative_to(base_dir.resolve())
|
||||
except ValueError:
|
||||
# Refuse to read outside the source's own directory.
|
||||
# Path traversal: refuse to read outside the source's directory.
|
||||
continue
|
||||
if not local_path.is_file():
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -1,16 +1,7 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Layout-aware PDF parsing via pymupdf + pymupdf4llm.
|
||||
|
||||
Produces Markdown per page (headings, pipe-tables, lists survive) so the
|
||||
recursive chunker can split on heading boundaries. Falls back to pypdf
|
||||
text-extraction only when pymupdf fails to open the file — keeps the
|
||||
pipeline alive for malformed PDFs.
|
||||
|
||||
Image extraction is gated behind `want_images=True` so text-only KBs
|
||||
pay zero cost for images they don't index.
|
||||
"""
|
||||
"""PDF → Markdown via pymupdf4llm; pypdf fallback for malformed files."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -39,8 +30,7 @@ def _extract_with_pymupdf(path: Path, want_images: bool) -> ParseResult:
|
|||
show_progress = False,
|
||||
)
|
||||
except Exception:
|
||||
# pymupdf4llm can choke on individual pages (rare). Fall
|
||||
# back to plain text extraction for just that page.
|
||||
# pymupdf4llm can choke on a single page; fall back to plain text.
|
||||
md = doc[page_index].get_text("text") or ""
|
||||
md = md.strip()
|
||||
if md:
|
||||
|
|
@ -55,7 +45,6 @@ def _extract_with_pymupdf(path: Path, want_images: bool) -> ParseResult:
|
|||
|
||||
|
||||
def _extract_images_pymupdf(doc, pages: list[ParsedPage]) -> list[ParsedImage]:
|
||||
"""Pull embedded images and pair each with the nearest text on the same page."""
|
||||
captions_by_page: dict[int, str] = {
|
||||
p.page_number: p.text for p in pages if p.page_number
|
||||
}
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ from . import ParsedPage, ParseResult
|
|||
|
||||
|
||||
def extract(path: Path, *, want_images: bool = False) -> ParseResult:
|
||||
# want_images is ignored — plain text / Markdown have no embedded images.
|
||||
raw = path.read_bytes()
|
||||
try:
|
||||
text = raw.decode("utf-8")
|
||||
|
|
|
|||
|
|
@ -1,18 +1,7 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Optional cross-encoder reranking stage.
|
||||
|
||||
Off by default. Callers opt in per request via ``enable_rerank`` on
|
||||
``SearchRequest``. The reranker model is lazy-loaded on first opt-in
|
||||
query and competes with the active chat model for GPU memory — keeping
|
||||
it opt-in protects chat latency on smaller GPUs.
|
||||
|
||||
``sentence_transformers.CrossEncoder`` has no unsloth wrapper today;
|
||||
this module loads it directly. A future ``FastCrossEncoder`` addition
|
||||
to ``unsloth/models/sentence_transformer.py`` would slot in here by
|
||||
replacing the import in ``_load``.
|
||||
"""
|
||||
"""Opt-in CrossEncoder reranker (off by default; shares GPU with chat model)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -51,12 +40,7 @@ def get_reranker(model_name: str | None = None) -> Any:
|
|||
|
||||
|
||||
def unload() -> None:
|
||||
"""Drop the reranker reference and trigger a GC pass.
|
||||
|
||||
Useful when memory pressure is high — callers can free the
|
||||
reranker without restarting the studio process. Next ``rerank``
|
||||
call lazy-loads it again.
|
||||
"""
|
||||
"""Drop the reranker; next call lazy-loads again."""
|
||||
global _model, _model_name
|
||||
with _lock:
|
||||
if _model is not None:
|
||||
|
|
@ -79,16 +63,9 @@ def rerank(
|
|||
model_name: str | None = None,
|
||||
top_k: int | None = None,
|
||||
) -> list[Hit]:
|
||||
"""Re-order ``pairs`` by CrossEncoder relevance to ``query``.
|
||||
|
||||
Each pair is ``(Hit, chunk_text)``. Returns Hits with the new
|
||||
cross-encoder scores attached. If ``top_k`` is given, truncates.
|
||||
"""
|
||||
"""Re-order (Hit, text) pairs by CrossEncoder score; image hits are appended last."""
|
||||
if not pairs:
|
||||
return []
|
||||
# CrossEncoder is text-only; image-kind hits get appended at the end
|
||||
# in their original relative order so they're never dropped, just
|
||||
# never reranked. Caption-kind hits are eligible (they carry text).
|
||||
text_pairs = [(h, t) for h, t in pairs if h.kind != "image"]
|
||||
image_hits = [h for h, _t in pairs if h.kind == "image"]
|
||||
|
||||
|
|
|
|||
|
|
@ -1,17 +1,7 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""High-level retrieval surface for RAG: BM25, dense, and RRF hybrid.
|
||||
|
||||
Reciprocal Rank Fusion is parameter-light: each candidate's fused score
|
||||
is the sum of ``1 / (rrf_k + rank)`` across rankers. It avoids the
|
||||
need to calibrate score scales between BM25 (raw, unbounded) and cosine
|
||||
similarity (-1..1).
|
||||
|
||||
Hits also carry the raw dense cosine score when available so callers
|
||||
can apply a meaningful similarity threshold (e.g., "drop hits below
|
||||
0.3 cosine") — the fused RRF score isn't on a comparable scale.
|
||||
"""
|
||||
"""RAG retrieval: BM25, dense, and RRF hybrid. Hits carry dense_score for thresholding."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -34,10 +24,7 @@ class Hit:
|
|||
document_id: str | None = None
|
||||
chunk_index: int | None = None
|
||||
kind: str = "text"
|
||||
# Raw cosine similarity from the dense retriever (0..1 for
|
||||
# normalized embeddings). None when this chunk wasn't returned by
|
||||
# the dense pass (BM25-only hit) — callers applying a similarity
|
||||
# floor should treat None as "no signal" and exclude it.
|
||||
# Raw cosine; None for BM25-only hits.
|
||||
dense_score: float | None = None
|
||||
|
||||
|
||||
|
|
@ -54,12 +41,7 @@ def retrieve_dense(
|
|||
document_ids: list[str] | None = None,
|
||||
embedder_model: str | None = None,
|
||||
) -> list[Hit]:
|
||||
"""Dense retrieval. `embedder_model` MUST match the model that
|
||||
populated this scope's vectors — using a different one yields a
|
||||
dim mismatch (shape (N, scope_dim) vs (query_dim,)) at distance
|
||||
compute time. Callers should resolve from the KB / thread settings
|
||||
before passing.
|
||||
"""
|
||||
"""Dense retrieval. embedder_model MUST match the model that populated this scope."""
|
||||
limit = k or RAG_TOP_K_DENSE
|
||||
vector = embeddings.encode(
|
||||
[query],
|
||||
|
|
@ -96,8 +78,7 @@ def _rrf_fuse(
|
|||
) -> list[Hit]:
|
||||
fused: dict[str, float] = {}
|
||||
seen: dict[str, Hit] = {}
|
||||
# Track the dense cosine score per chunk so it survives fusion —
|
||||
# callers downstream filter on this, not the RRF score.
|
||||
# Preserve dense_score through fusion for downstream thresholding.
|
||||
dense_scores: dict[str, float] = {}
|
||||
for ranking in rankings:
|
||||
for rank, hit in enumerate(ranking):
|
||||
|
|
@ -148,12 +129,7 @@ def retrieve_hybrid(
|
|||
|
||||
|
||||
def filter_by_min_score(hits: list[Hit], min_score: float) -> list[Hit]:
|
||||
"""Drop hits whose dense cosine score is below ``min_score``.
|
||||
|
||||
Hits without a dense score (BM25-only) are dropped too — there's
|
||||
no comparable signal to evaluate them against the similarity floor.
|
||||
Use ``min_score = 0.0`` (or negative) to disable the filter.
|
||||
"""
|
||||
"""Drop hits whose dense_score < min_score; BM25-only hits dropped too."""
|
||||
if min_score <= 0.0:
|
||||
return hits
|
||||
return [h for h in hits if h.dense_score is not None and h.dense_score >= min_score]
|
||||
|
|
|
|||
|
|
@ -1,25 +1,13 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Scope identifiers + per-scope embedder resolution.
|
||||
|
||||
A "scope" is the namespace key carried by vector rows and BM25
|
||||
indexes — ``kb_<uuid>`` for stand-alone Knowledge Bases or
|
||||
``thread_<uuid>`` for per-thread document sets.
|
||||
|
||||
`resolve_scope_embedder` looks up which embedder populated a scope's
|
||||
vectors so the query side can re-use the same model. Lives in
|
||||
``core/rag`` rather than ``routes`` so the inference-side tool
|
||||
handler can call it without a route-→-core import cycle.
|
||||
"""
|
||||
"""Scope identifiers (kb_<id> / thread_<id>) + per-scope embedder resolver."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from storage.studio_db import get_connection, list_chat_settings
|
||||
from utils.rag.config import resolve_embedder
|
||||
|
||||
# Persisted chat-settings keys. Defined here so the resolver (core)
|
||||
# and the route handlers (routes/rag.py) share one source of truth.
|
||||
RAG_DEFAULTS_KEY = "rag.defaults"
|
||||
|
||||
|
||||
|
|
@ -28,16 +16,7 @@ def thread_settings_key(thread_id: str) -> str:
|
|||
|
||||
|
||||
def resolve_scope_embedder(scope: str) -> str | None:
|
||||
"""Return the embedder used to populate ``scope``'s vector rows.
|
||||
|
||||
Resolution order:
|
||||
- ``kb_<id>`` → ``rag_knowledge_bases.embedding_model`` column.
|
||||
- ``thread_<id>`` → per-thread override → app-level defaults
|
||||
override → ``RAG_EMBEDDER_MATRIX[(mode, chunking)]``.
|
||||
|
||||
Returns ``None`` for unrecognised scope strings; callers treat
|
||||
``None`` as "fall back to the configured default embedder".
|
||||
"""
|
||||
"""KB → kb.embedding_model; thread → per-thread/defaults/matrix. None = use default."""
|
||||
if scope.startswith("kb_"):
|
||||
kb_id = scope[len("kb_") :]
|
||||
with get_connection() as conn:
|
||||
|
|
|
|||
|
|
@ -3,15 +3,8 @@
|
|||
|
||||
"""`search_knowledge_base` tool — RAG retrieval surfaced to the LLM.
|
||||
|
||||
Invoked from `core/inference/tools.execute_tool` when the local model
|
||||
emits a `search_knowledge_base` call. The handler runs the existing
|
||||
hybrid retrieval, hydrates chunk text + filename + page number from
|
||||
sqlite, and returns a Markdown-with-numbered-citations string that
|
||||
the LLM consumes as the tool-result message.
|
||||
|
||||
Scope (`kb_id` / `thread_id`) is not exposed as a tool argument — it
|
||||
comes from the chat-completions request body (`rag_scope`) so the
|
||||
LLM doesn't need to know about KB UUIDs.
|
||||
Scope comes from the request body (`rag_scope`), not from the tool args,
|
||||
so the model never sees KB UUIDs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -61,12 +54,7 @@ SEARCH_KNOWLEDGE_BASE_TOOL = {
|
|||
|
||||
|
||||
def _format_hits_for_llm(hits: list[Any]) -> str:
|
||||
"""Render hits as numbered Markdown citations for the LLM.
|
||||
|
||||
Empty results produce a one-line message rather than an empty
|
||||
string — the model needs to know the search ran but found nothing
|
||||
so it can fall back to its own knowledge or ask the user.
|
||||
"""
|
||||
"""Render hits as numbered Markdown citations; empty results return a message, not ''."""
|
||||
if not hits:
|
||||
return (
|
||||
"No matching chunks were found in the attached documents. "
|
||||
|
|
@ -95,13 +83,7 @@ def search_knowledge_base(
|
|||
min_score: float = 0.0,
|
||||
mode: Literal["bm25", "dense", "hybrid"] = "hybrid",
|
||||
) -> str:
|
||||
"""Execute the RAG search and return a tool-result string.
|
||||
|
||||
`kb_id` takes precedence over `thread_id` when both are set —
|
||||
matches the create/upload contract that a document belongs to one
|
||||
or the other, never both. ``min_score`` is a cosine-similarity
|
||||
floor on dense hits; chunks below it are dropped.
|
||||
"""
|
||||
"""Run RAG and return a tool-result string. kb_id takes precedence over thread_id."""
|
||||
if not query or not query.strip():
|
||||
return "Error: empty query."
|
||||
|
||||
|
|
@ -215,8 +197,7 @@ def search_knowledge_base(
|
|||
else:
|
||||
hits = hits[:k]
|
||||
|
||||
# Image-kind hits don't carry LLM-friendly text — skip them. The
|
||||
# paired caption (linked_chunk_id) usually surfaces separately.
|
||||
# Skip image-kind hits; the paired caption surfaces separately.
|
||||
formatted = [
|
||||
lookup[hit.chunk_id]
|
||||
for hit in hits
|
||||
|
|
|
|||
|
|
@ -1,21 +1,8 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""SQLite-backed vector store for RAG, distance math via sqlite-vec.
|
||||
|
||||
Replaces the previous Qdrant local-mode store. Vectors are stored as
|
||||
BLOBs in a single `rag_vectors` table inside rag.db, keyed by
|
||||
chunk_id; cosine distance is computed by sqlite-vec's
|
||||
`vec_distance_cosine(blob, blob)` scalar function.
|
||||
|
||||
A "scope" is `kb_<uuid>` for standalone knowledge bases or
|
||||
`thread_<uuid>` for per-thread document sets. The scope column is
|
||||
indexed; queries filter by scope before computing distances so
|
||||
different scopes can hold vectors of different dimensions without
|
||||
breaking the cosine math. The per-scope embedder resolver
|
||||
(core/rag/scope.py:resolve_scope_embedder) guarantees one embedder
|
||||
per scope, so dims within a scope are always consistent.
|
||||
"""
|
||||
"""sqlite-vec vector store. Scope filter (kb_<id>/thread_<id>) keeps mixed-dim
|
||||
scopes safe — per-scope embedder resolver guarantees one dim per scope."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -36,11 +23,6 @@ def thread_scope(thread_id: str) -> str:
|
|||
|
||||
|
||||
def collection_exists(scope: str) -> bool:
|
||||
"""Whether the scope has any indexed vectors.
|
||||
|
||||
Used by callers (notably retrieval.retrieve_dense) to short-circuit
|
||||
when the scope was never populated. Cheap — a covering index hit.
|
||||
"""
|
||||
from core.rag.db import get_rag_connection
|
||||
|
||||
conn = get_rag_connection()
|
||||
|
|
@ -52,21 +34,12 @@ def collection_exists(scope: str) -> bool:
|
|||
|
||||
|
||||
def ensure_collection(scope: str, dim: int) -> None:
|
||||
"""No-op for sqlite-vec — vectors get inserted directly into the
|
||||
shared table. Kept for API parity with the previous Qdrant store
|
||||
so ingestion callers don't need conditional logic.
|
||||
"""
|
||||
_ = scope, dim # unused; signature preserved
|
||||
"""No-op; kept for API parity. Vectors go straight into the shared table."""
|
||||
_ = scope, dim
|
||||
|
||||
|
||||
def upsert_chunks(scope: str, points: Iterable[dict]) -> None:
|
||||
"""Insert/update vectors. Each point: {id, vector, payload}.
|
||||
|
||||
Conflict resolution is per chunk_id (the primary key): re-ingesting
|
||||
overwrites in place. Payload is round-tripped as JSON so the
|
||||
Qdrant-shaped {filename, page_number, kind, ...} dicts callers
|
||||
already build can be reused unchanged.
|
||||
"""
|
||||
"""Insert/update vectors. Each point: {id, vector, payload}."""
|
||||
import sqlite_vec
|
||||
|
||||
from core.rag.db import get_rag_connection
|
||||
|
|
@ -116,14 +89,7 @@ def search(
|
|||
top_k: int,
|
||||
document_ids: list[str] | None = None,
|
||||
) -> list[dict]:
|
||||
"""Cosine-distance search filtered to a single scope.
|
||||
|
||||
Returns rows in the same shape as the old Qdrant path —
|
||||
{chunk_id, score, payload} — where ``score`` is cosine similarity
|
||||
in [0, 1] (vec_distance_cosine returns 1 - similarity, so we
|
||||
invert). Filtered-by-document_ids variant for the per-thread
|
||||
"search only these uploads" case.
|
||||
"""
|
||||
"""Cosine search; returns {chunk_id, score, payload} with score = 1 - distance."""
|
||||
import sqlite_vec
|
||||
|
||||
from core.rag.db import get_rag_connection
|
||||
|
|
|
|||
|
|
@ -1,21 +1,7 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""RAG API routes.
|
||||
|
||||
Surface:
|
||||
- Knowledge-base CRUD
|
||||
- Document upload (KB-scoped and per-thread)
|
||||
- Document list/delete
|
||||
- Ingestion-job SSE stream
|
||||
- Search (BM25 / dense / hybrid)
|
||||
|
||||
Per-thread document uploads are scoped to a single chat thread and
|
||||
share the same chunk/embed/index pipeline as KB documents — they only
|
||||
differ in the scope key (``thread_<id>`` vs ``kb_<id>``) and lifecycle
|
||||
(per-thread docs are dropped when the thread is deleted, via the
|
||||
ON DELETE CASCADE on rag_documents.thread_id).
|
||||
"""
|
||||
"""RAG API: KB CRUD, document upload (KB + per-thread), ingestion SSE, search."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -69,9 +55,7 @@ router = APIRouter()
|
|||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Pydantic schemas
|
||||
# ------------------------------------------------------------------
|
||||
# --- Pydantic schemas ---
|
||||
|
||||
ChunkingStrategy = Literal["standard", "late"]
|
||||
KBMode = Literal["text", "multimodal"]
|
||||
|
|
@ -81,9 +65,6 @@ class CreateKBRequest(BaseModel):
|
|||
name: str = Field(min_length = 1, max_length = 200)
|
||||
description: str | None = None
|
||||
embedding_model: str | None = None
|
||||
# Phase 3 introduces two per-KB knobs. Both default to today's
|
||||
# behaviour so existing API clients are unaffected. The (multimodal,
|
||||
# late) combination is rejected at create time — see _validate_mode_combo.
|
||||
chunking_strategy: ChunkingStrategy = "standard"
|
||||
mode: KBMode = "text"
|
||||
|
||||
|
|
@ -145,9 +126,6 @@ class SearchRequest(BaseModel):
|
|||
document_ids: list[str] | None = None
|
||||
enable_rerank: bool = False
|
||||
reranker_model: str | None = None
|
||||
# Cosine-similarity floor on the dense retrieval score. Hits whose
|
||||
# dense_score is below this (or absent — BM25-only hits) are
|
||||
# dropped before the response is sent. 0.0 disables the filter.
|
||||
min_score: float = Field(default = 0.0, ge = 0.0, le = 1.0)
|
||||
|
||||
|
||||
|
|
@ -167,9 +145,7 @@ class SearchResponse(BaseModel):
|
|||
hits: list[SearchHit]
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
# --- Helpers ---
|
||||
|
||||
|
||||
def _sanitize_filename(filename: str) -> str:
|
||||
|
|
@ -181,15 +157,10 @@ def _now_ms() -> int:
|
|||
return int(time.time())
|
||||
|
||||
|
||||
# Per-scope embedder lookup lives in core/rag/scope.py so the
|
||||
# inference-side tool handler can use it without a routes-→-core
|
||||
# import cycle.
|
||||
from core.rag.scope import resolve_scope_embedder as _resolve_scope_embedder # noqa: E402
|
||||
|
||||
|
||||
def _row_to_kb(row: Any) -> KBResponse:
|
||||
# chunking_strategy / mode may be absent on rows fetched through a
|
||||
# pre-Phase-3 connection in tests; fall back to the schema defaults.
|
||||
keys = row.keys() if hasattr(row, "keys") else ()
|
||||
chunking_strategy = (
|
||||
row["chunking_strategy"] if "chunking_strategy" in keys else "standard"
|
||||
|
|
@ -207,12 +178,7 @@ def _row_to_kb(row: Any) -> KBResponse:
|
|||
|
||||
|
||||
def _validate_mode_combo(mode: KBMode, chunking_strategy: ChunkingStrategy) -> None:
|
||||
"""Reject the one illegal (mode, strategy) combination.
|
||||
|
||||
No public open-weight embedder supports both late-chunking pooling
|
||||
and shared text/image embedding. Surface the constraint as a 400
|
||||
rather than failing silently during ingestion.
|
||||
"""
|
||||
"""Reject (multimodal, late) — no embedder supports both at once."""
|
||||
if mode == "multimodal" and chunking_strategy == "late":
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
|
|
@ -287,11 +253,9 @@ async def _save_upload(file: UploadFile) -> tuple[Path, str, int]:
|
|||
stored_path = upload_dir / stored_name
|
||||
max_bytes = RAG_MAX_UPLOAD_MB * 1024 * 1024
|
||||
written = 0
|
||||
# anyio.open_file routes each write through a worker thread so a
|
||||
# multi-MB upload doesn't stall concurrent requests on the event
|
||||
# loop. The async-with handles close on both happy and error
|
||||
# paths; the outer try/except cleans up the partial file after
|
||||
# the file handle is closed (Windows refuses unlink on an open fd).
|
||||
# Route writes through anyio worker thread so the event loop stays free.
|
||||
# Outer try/except cleans up partial files after async-with closes the fd
|
||||
# (Windows refuses unlink on an open fd).
|
||||
try:
|
||||
async with await anyio.open_file(stored_path, "wb") as f:
|
||||
while True:
|
||||
|
|
@ -369,9 +333,7 @@ def _unlink_if_under_uploads(path: Path) -> None:
|
|||
real.unlink(missing_ok = True)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Knowledge bases
|
||||
# ------------------------------------------------------------------
|
||||
# --- Knowledge bases ---
|
||||
|
||||
|
||||
@router.post("/knowledge-bases", response_model = KBResponse)
|
||||
|
|
@ -384,9 +346,7 @@ def create_knowledge_base(
|
|||
_validate_mode_combo(payload.mode, payload.chunking_strategy)
|
||||
|
||||
kb_id = str(uuid4())
|
||||
# If the caller didn't override embedding_model, resolve from the
|
||||
# Phase-3 matrix using their (mode, strategy) selection. Unknown
|
||||
# combos fall back to the legacy default — see resolve_embedder.
|
||||
# No override: resolve from (mode, strategy) matrix.
|
||||
embedding_model = payload.embedding_model or resolve_embedder(
|
||||
payload.mode, payload.chunking_strategy
|
||||
)
|
||||
|
|
@ -483,8 +443,7 @@ def set_rag_defaults(
|
|||
current = _load_rag_defaults()
|
||||
new_strategy = payload.chunking_strategy or current.chunking_strategy
|
||||
new_mode = payload.mode or current.mode
|
||||
# PATCH-style — passing an empty string clears the override; a
|
||||
# null/missing field keeps the current value.
|
||||
# PATCH-style: empty string clears, null/missing keeps current.
|
||||
if payload.embedding_model is None:
|
||||
new_embedder = current.embedding_model
|
||||
elif payload.embedding_model.strip() == "":
|
||||
|
|
@ -526,11 +485,7 @@ def _thread_settings_key(thread_id: str) -> str:
|
|||
|
||||
|
||||
def _load_thread_settings(thread_id: str) -> ThreadRagSettings:
|
||||
"""Per-thread RAG settings, falling back to app-level defaults.
|
||||
|
||||
Stored in chat_settings under "thread:<id>:rag" as a nested JSON
|
||||
dict — same shape as RagDefaults.
|
||||
"""
|
||||
"""Per-thread RAG settings (chat_settings['thread:<id>:rag']) with defaults fallback."""
|
||||
settings = list_chat_settings()
|
||||
raw = settings.get(_thread_settings_key(thread_id)) or {}
|
||||
if not isinstance(raw, dict):
|
||||
|
|
@ -611,11 +566,7 @@ def _reingest_scope(
|
|||
mode: str,
|
||||
embedding_model: str,
|
||||
) -> ReingestResponse:
|
||||
"""Wipe scope artifacts and re-enqueue every stored document.
|
||||
|
||||
The chat_settings / per-thread defaults aren't touched — caller is
|
||||
responsible for updating any associated metadata before calling.
|
||||
"""
|
||||
"""Wipe scope artifacts and re-enqueue every document; metadata untouched."""
|
||||
scope = kb_scope(kb_id) if kb_id else thread_scope(thread_id) # type: ignore[arg-type]
|
||||
with get_connection() as conn:
|
||||
if kb_id:
|
||||
|
|
@ -628,9 +579,7 @@ def _reingest_scope(
|
|||
"SELECT id, stored_path FROM rag_documents WHERE thread_id = ?",
|
||||
(thread_id,),
|
||||
).fetchall()
|
||||
# Delete the rag_documents rows (cascade drops chunks); the
|
||||
# uploaded file on disk is preserved so we can re-ingest from
|
||||
# it. We re-INSERT a fresh row per stored_path below.
|
||||
# Drop rag_documents (chunks cascade); files on disk are reused below.
|
||||
doc_ids = [r["id"] for r in rows]
|
||||
if doc_ids:
|
||||
placeholders = ",".join("?" for _ in doc_ids)
|
||||
|
|
@ -648,8 +597,7 @@ def _reingest_scope(
|
|||
if not stored_path.is_file():
|
||||
continue
|
||||
filename = stored_path.name
|
||||
# Strip the UUID prefix we attached at upload time so the
|
||||
# re-inserted document carries the original name.
|
||||
# Strip the upload-time UUID prefix; keep the original filename.
|
||||
if "_" in filename:
|
||||
_uuid_prefix, _, original = filename.partition("_")
|
||||
if original:
|
||||
|
|
@ -728,14 +676,7 @@ def reingest_thread_documents(
|
|||
payload: UpdateThreadRagSettingsRequest | None = None,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> ReingestResponse:
|
||||
"""Rebuild a thread's RAG index.
|
||||
|
||||
Optional body lets the caller change the thread's chunking
|
||||
strategy / mode / embedder at the same time — persisted into
|
||||
chat_settings before re-ingestion so subsequent uploads pick up
|
||||
the new values too. With an empty body, current settings are
|
||||
reused.
|
||||
"""
|
||||
"""Rebuild a thread's RAG index; optional body updates settings before reingest."""
|
||||
from utils.rag.config import resolve_embedder
|
||||
|
||||
if payload is None:
|
||||
|
|
@ -745,7 +686,6 @@ def reingest_thread_documents(
|
|||
or payload.mode is not None
|
||||
or payload.embedding_model is not None
|
||||
):
|
||||
# set_thread_rag_settings handles validation + persistence.
|
||||
settings = set_thread_rag_settings(
|
||||
thread_id,
|
||||
payload,
|
||||
|
|
@ -786,9 +726,7 @@ def delete_knowledge_base(
|
|||
return {"ok": True}
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Document upload (KB and per-thread)
|
||||
# ------------------------------------------------------------------
|
||||
# --- Document upload (KB and per-thread) ---
|
||||
|
||||
|
||||
@router.post("/knowledge-bases/{kb_id}/documents", response_model = UploadResponse)
|
||||
|
|
@ -799,9 +737,7 @@ async def upload_kb_document(
|
|||
) -> UploadResponse:
|
||||
kb_row = _kb_or_404(kb_id)
|
||||
stored_path, filename, byte_size = await _save_upload(file)
|
||||
# Defensive .get() — rows fetched through a connection that pre-dates
|
||||
# the Phase 3 schema (e.g. in tests) lack chunking_strategy/mode;
|
||||
# fall back to the same defaults as the column.
|
||||
# Tolerate pre-Phase-3 rows missing chunking_strategy/mode.
|
||||
kb_keys = kb_row.keys() if hasattr(kb_row, "keys") else ()
|
||||
chunking_strategy = (
|
||||
kb_row["chunking_strategy"] if "chunking_strategy" in kb_keys else "standard"
|
||||
|
|
@ -828,12 +764,8 @@ async def upload_thread_document(
|
|||
) -> UploadResponse:
|
||||
from utils.rag.config import resolve_embedder
|
||||
|
||||
# Don't validate against chat_threads — a brand-new chat won't be
|
||||
# persisted there until after the first runStart/runEnd. Users who
|
||||
# attach a document on a fresh thread would otherwise hit a 404.
|
||||
# No chat_threads check — fresh threads aren't persisted until first run.
|
||||
stored_path, filename, byte_size = await _save_upload(file)
|
||||
# Per-thread settings fall back to app-level defaults inside the
|
||||
# helper, so first-time-uploaded threads inherit user preferences.
|
||||
settings = _load_thread_settings(thread_id)
|
||||
embedder = settings.embedding_model or resolve_embedder(
|
||||
settings.mode,
|
||||
|
|
@ -852,9 +784,7 @@ async def upload_thread_document(
|
|||
)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Document list / delete
|
||||
# ------------------------------------------------------------------
|
||||
# --- Document list / delete ---
|
||||
|
||||
|
||||
@router.get("/knowledge-bases/{kb_id}/documents", response_model = DocumentListResponse)
|
||||
|
|
@ -890,13 +820,7 @@ def get_rag_image(
|
|||
filename: str,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> FileResponse:
|
||||
"""Serve an image extracted during multimodal ingestion.
|
||||
|
||||
Files live under ``rag_uploads_root() / 'images' / <document_id>``.
|
||||
We realpath-check the resolved file against that root to refuse
|
||||
path-traversal attempts (``..`` segments, symlinks pointing
|
||||
elsewhere). Filenames are constrained to a single path component.
|
||||
"""
|
||||
"""Serve an extracted image; realpath-check against the uploads root."""
|
||||
if "/" in filename or "\\" in filename or filename.startswith("."):
|
||||
raise HTTPException(status_code = 400, detail = "Invalid filename")
|
||||
root = Path(os.path.realpath(rag_uploads_root() / "images"))
|
||||
|
|
@ -930,12 +854,7 @@ def delete_document(
|
|||
def list_thread_indexes(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> ThreadIndexListResponse:
|
||||
"""List every chat thread that has at least one RAG document.
|
||||
|
||||
LEFT JOIN to chat_threads so threads that were never persisted
|
||||
(user attached a file but never sent a message) still show up —
|
||||
just with a null title.
|
||||
"""
|
||||
"""List threads with >=1 RAG doc. LEFT JOIN keeps unpersisted threads (null title)."""
|
||||
with get_connection() as conn:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
|
|
@ -969,19 +888,12 @@ def clear_thread_documents(
|
|||
thread_id: str,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> dict:
|
||||
"""Purge every RAG document attached to ``thread_id``.
|
||||
|
||||
Removes the per-thread vector rows, the bm25 index, the
|
||||
rag_documents/rag_chunks rows, and the uploaded files. The chat
|
||||
thread itself is untouched.
|
||||
"""
|
||||
"""Drop all RAG artifacts for thread_id; chat thread itself untouched."""
|
||||
ingestion.purge_thread_documents([thread_id])
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Ingestion job SSE
|
||||
# ------------------------------------------------------------------
|
||||
# --- Ingestion job SSE ---
|
||||
|
||||
|
||||
@router.get("/jobs/{job_id}/events")
|
||||
|
|
@ -1050,9 +962,7 @@ async def _replay_terminal_state(row: Any):
|
|||
yield f"data: {json.dumps(payload)}\n\n"
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Search
|
||||
# ------------------------------------------------------------------
|
||||
# --- Search ---
|
||||
|
||||
|
||||
@router.post("/search", response_model = SearchResponse)
|
||||
|
|
@ -1071,9 +981,7 @@ def search(
|
|||
else:
|
||||
scope = thread_scope(payload.thread_id)
|
||||
|
||||
# Embed the query with the same model that populated this scope's
|
||||
# vectors. Mixing spaces (e.g. Qwen3-VL 2048-d docs vs bge-small
|
||||
# 384-d query) crashes inside the cosine-distance compute.
|
||||
# Query must use the same embedder as the scope (dim must match).
|
||||
scope_embedder = _resolve_scope_embedder(scope)
|
||||
logger.info(
|
||||
"RAG search: scope=%s embedder=%s mode=%s top_k=%d min_score=%.3f rerank=%s query=%r",
|
||||
|
|
@ -1086,8 +994,7 @@ def search(
|
|||
payload.query[:120],
|
||||
)
|
||||
|
||||
# When reranking is opt-in, pull a wider candidate pool so the
|
||||
# CrossEncoder has more to choose from before truncating to top_k.
|
||||
# Reranker needs a wider candidate pool than top_k.
|
||||
candidate_k = (
|
||||
max(payload.top_k, RAG_RERANK_CANDIDATE_K)
|
||||
if payload.enable_rerank
|
||||
|
|
@ -1168,8 +1075,7 @@ def search(
|
|||
image_url: str | None = None
|
||||
if kind == "image" and meta.get("image_path"):
|
||||
image_url = (
|
||||
f"/api/rag/images/{meta['document_id']}/"
|
||||
f"{Path(meta['image_path']).name}"
|
||||
f"/api/rag/images/{meta['document_id']}/{Path(meta['image_path']).name}"
|
||||
)
|
||||
out.append(
|
||||
SearchHit(
|
||||
|
|
|
|||
|
|
@ -205,13 +205,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
|
|||
) WITHOUT ROWID
|
||||
"""
|
||||
)
|
||||
# RAG: knowledge bases, documents, chunks, ingestion jobs.
|
||||
# rag_documents enforces XOR on (kb_id, thread_id): a document belongs
|
||||
# to either a standalone KB or a single chat thread, never both.
|
||||
# chunking_strategy and mode are set at KB-creation time and are
|
||||
# immutable thereafter — changing either invalidates existing chunks
|
||||
# because they were ingested through a specific pipeline (different
|
||||
# chunker, different embedder). See Phase 3 in the plan.
|
||||
# RAG schema. rag_documents enforces XOR on (kb_id, thread_id).
|
||||
# chunking_strategy/mode are immutable post-create (invalidates chunks).
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS rag_knowledge_bases (
|
||||
|
|
@ -226,9 +221,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
|
|||
)
|
||||
"""
|
||||
)
|
||||
# Idempotent ALTER for existing installs that pre-date the columns.
|
||||
# Mirrors the chat_threads display_name / *_code_exec_container_id
|
||||
# pattern earlier in this file.
|
||||
# Idempotent ALTER for pre-existing installs.
|
||||
kb_cols = {
|
||||
row[1]
|
||||
for row in conn.execute("PRAGMA table_info(rag_knowledge_bases)").fetchall()
|
||||
|
|
@ -243,11 +236,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
|
|||
"ALTER TABLE rag_knowledge_bases "
|
||||
"ADD COLUMN mode TEXT NOT NULL DEFAULT 'text'"
|
||||
)
|
||||
# thread_id has no FK to chat_threads. A user can attach a document
|
||||
# to a thread that hasn't yet been persisted (saveThread only runs
|
||||
# after the first model exchange — see runtime-provider.tsx). The
|
||||
# chat-thread DELETE handlers in routes/chat_history.py purge
|
||||
# matching rag_documents explicitly so lifecycle stays clean.
|
||||
# thread_id has no FK: docs can attach before the thread is persisted.
|
||||
# chat_history DELETE handlers purge matching rag_documents explicitly.
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS rag_documents (
|
||||
|
|
@ -272,11 +262,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
|
|||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_rag_documents_thread_id ON rag_documents(thread_id)"
|
||||
)
|
||||
# kind = 'text' | 'image' | 'caption'. image_path is set when kind = 'image'
|
||||
# (path to the extracted figure under rag_uploads_root() / 'images/').
|
||||
# linked_chunk_id pairs an 'image' chunk with its 'caption' chunk (and vice
|
||||
# versa) so retrieval can hydrate the matching half. Both default to NULL
|
||||
# so today's text-only ingestion is unaffected. See Phase 3B-multimodal.
|
||||
# kind: text|image|caption. linked_chunk_id pairs image↔caption (both null for text).
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS rag_chunks (
|
||||
|
|
|
|||
|
|
@ -31,32 +31,17 @@ RAG_EMBEDDING_MODEL: str = (
|
|||
or "BAAI/bge-small-en-v1.5"
|
||||
)
|
||||
|
||||
# Phase 3: default embedders per (mode, chunking_strategy). Ingestion in
|
||||
# Phase 3B-late and Phase 3B-multimodal looks the embedder up here at job
|
||||
# start, falling back to RAG_EMBEDDING_MODEL (above) for legacy KBs that
|
||||
# pre-date the columns. The (multimodal, late) combo is intentionally
|
||||
# absent — no public open-weight embedder supports both at once, and
|
||||
# routes/rag.py rejects the combo with a 400 at KB create time.
|
||||
# Default embedder per (mode, chunking). (multimodal, late) is unsupported
|
||||
# and rejected at KB-create time in routes/rag.py.
|
||||
RAG_EMBEDDER_MATRIX: dict[tuple[str, str], str] = {
|
||||
("text", "standard"): "BAAI/bge-small-en-v1.5",
|
||||
("text", "late"): "nomic-ai/nomic-embed-text-v1.5",
|
||||
# Qwen3-VL-Embedding-2B: 2B-param multimodal embedder built on
|
||||
# Qwen3 (not CLIP), so no 77-token text cap — long chunks embed
|
||||
# losslessly. Loads through vanilla SentenceTransformer with
|
||||
# trust_remote_code, no shim. 2048-d shared text/image space.
|
||||
# Trade-off: ~4 GB download / VRAM vs BGE-VL-base's ~600 MB.
|
||||
("multimodal", "standard"): "Qwen/Qwen3-VL-Embedding-2B",
|
||||
}
|
||||
|
||||
|
||||
def resolve_embedder(mode: str, chunking_strategy: str) -> str:
|
||||
"""Look up the default embedder for a (mode, chunking_strategy) pair.
|
||||
|
||||
Unknown combos fall back to the legacy single default so old KBs
|
||||
keep working. Callers that explicitly require the new matrix
|
||||
behaviour (Phase 3B paths) should validate the inputs before
|
||||
calling.
|
||||
"""
|
||||
"""Embedder for (mode, chunking); unknown combos fall back to RAG_EMBEDDING_MODEL."""
|
||||
return RAG_EMBEDDER_MATRIX.get(
|
||||
(mode, chunking_strategy),
|
||||
RAG_EMBEDDING_MODEL,
|
||||
|
|
@ -76,9 +61,6 @@ RAG_MAX_UPLOAD_MB: int = _env_int("UNSLOTH_RAG_MAX_UPLOAD_MB", 50)
|
|||
|
||||
RAG_EMBED_BATCH_SIZE: int = _env_int("UNSLOTH_RAG_EMBED_BATCH_SIZE", 32)
|
||||
|
||||
# Reranking is off by default. The CrossEncoder runs on GPU and competes
|
||||
# with the active chat model — callers opt in per-request via
|
||||
# `enable_rerank` on SearchRequest.
|
||||
RAG_RERANKER_MODEL: str = (
|
||||
os.environ.get("UNSLOTH_RAG_RERANKER_MODEL", "").strip() or "BAAI/bge-reranker-base"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -337,8 +337,7 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => {
|
|||
event.preventDefault();
|
||||
return;
|
||||
}
|
||||
// Send is going through — drop the chips. The docs themselves
|
||||
// remain in the backend thread KB and stay searchable.
|
||||
// Drop chips on send; docs stay searchable in the backend.
|
||||
clearDocs();
|
||||
},
|
||||
[
|
||||
|
|
@ -941,9 +940,7 @@ const ImagesToggle: FC = () => {
|
|||
);
|
||||
};
|
||||
|
||||
// Mirror of shared-composer's RAG pill (the master switch for
|
||||
// retrieval). Visible on every model; the sidebar Retrieval section
|
||||
// configures source / top-K / reranker once this is on.
|
||||
// Master RAG switch (mirrors shared-composer); sidebar configures the rest.
|
||||
const RagToggle: FC = () => {
|
||||
const modelLoaded = useChatRuntimeStore(
|
||||
(s) => !!s.params.checkpoint && !s.modelLoading,
|
||||
|
|
@ -1037,11 +1034,7 @@ const ToolStatusDisplay: FC = () => {
|
|||
);
|
||||
};
|
||||
|
||||
// Custom + button used when RAG is on: opens a picker accepting
|
||||
// doc formats the ingester handles and routes selected files to
|
||||
// the RAG thread-document pipeline. Replaces the stock
|
||||
// ComposerAddAttachment (which base64-attaches files inline) so
|
||||
// docs become indexed chunks instead of one-shot model context.
|
||||
// RAG-aware + button: picks doc formats and routes to ingest pipeline.
|
||||
const RagDocAttachment: FC<{ onSelect: (file: File) => void }> = ({
|
||||
onSelect,
|
||||
}) => {
|
||||
|
|
@ -1061,7 +1054,6 @@ const RagDocAttachment: FC<{ onSelect: (file: File) => void }> = ({
|
|||
const f = files[i];
|
||||
if (f && isDocumentFile(f)) onSelect(f);
|
||||
}
|
||||
// Reset so re-selecting the same file fires onChange.
|
||||
e.target.value = "";
|
||||
}}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -992,19 +992,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
});
|
||||
}
|
||||
|
||||
// RAG: optionally retrieve context for the last user turn and
|
||||
// prepend it as a system-role block. Failures are logged but
|
||||
// don't break the chat — better to answer without context than
|
||||
// to drop a message the user just sent.
|
||||
//
|
||||
// Pre-fetch RAG context unconditionally when the RAG button is
|
||||
// on and a source is selected. This runs for every provider —
|
||||
// local-tool, local-no-tool, and external — so users don't have
|
||||
// to phrase their query as "the document I attached" for
|
||||
// retrieval to fire. On local tool-capable models the
|
||||
// `search_knowledge_base` tool is *also* registered below as an
|
||||
// optional refinement path (the LLM can run a second, narrower
|
||||
// query if the pre-fetched chunks weren't enough).
|
||||
// Pre-fetch RAG context for the last user turn; failures don't block chat.
|
||||
// Runs for all providers; local tool-capable models also get the tool below
|
||||
// for a narrower follow-up query if needed.
|
||||
const ragSource = runtime.ragSource;
|
||||
const ragToolEnabled = runtime.ragToolEnabled;
|
||||
const ragToolPathTaken =
|
||||
|
|
@ -1647,11 +1637,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
...(codeToolsEnabled ? ["python", "terminal"] : []),
|
||||
...(ragToolPathTaken ? ["search_knowledge_base"] : []),
|
||||
],
|
||||
// Phase 4: per-request RAG context the backend's
|
||||
// `search_knowledge_base` handler reads when the LLM
|
||||
// invokes the tool. Only sent when the tool path
|
||||
// is taken — external providers fall through to the
|
||||
// pre-fetch block above.
|
||||
// Per-request scope for the LLM-invoked tool; tool path only.
|
||||
...(ragToolPathTaken
|
||||
? {
|
||||
rag_scope: {
|
||||
|
|
|
|||
|
|
@ -464,10 +464,8 @@ export function ChatSettingsPanel({
|
|||
const updateThreadSettings = useRagStore((s) => s.updateThreadSettings);
|
||||
const ragDefaults = useRagStore((s) => s.defaults);
|
||||
|
||||
// Load this thread's RAG settings once per threadId. Ref-guarded so
|
||||
// `threadSettings` isn't a dep — if it were, the post-load
|
||||
// store-mutation re-triggers the effect and any failure mode where
|
||||
// the selector flickers undefined produces an update loop.
|
||||
// Load thread RAG settings once per threadId. Ref-guarded; keep
|
||||
// `threadSettings` out of deps to avoid post-load update loops.
|
||||
const threadSettingsLoadedRef = useRef<string | null>(null);
|
||||
useEffect(() => {
|
||||
if (
|
||||
|
|
@ -492,7 +490,6 @@ export function ChatSettingsPanel({
|
|||
) => {
|
||||
if (!activeThreadId) return;
|
||||
if (threadDocs.length === 0) {
|
||||
// No existing chunks to invalidate — just persist.
|
||||
void updateThreadSettings(activeThreadId, patch);
|
||||
return;
|
||||
}
|
||||
|
|
@ -503,8 +500,7 @@ export function ChatSettingsPanel({
|
|||
if (ok) {
|
||||
void reingestThread(activeThreadId, patch);
|
||||
} else {
|
||||
// User declined — refresh the store so the select snaps back
|
||||
// to the unchanged settings.
|
||||
// User declined: refresh so the select snaps back.
|
||||
void loadThreadSettings(activeThreadId);
|
||||
}
|
||||
};
|
||||
|
|
@ -645,8 +641,6 @@ export function ChatSettingsPanel({
|
|||
activeExternalProvider.baseUrl,
|
||||
) &&
|
||||
activeExternalProvider.providerType === "openai";
|
||||
// (activeThreadId is declared earlier in this component — see the
|
||||
// RAG retrieval-section block above.)
|
||||
const openAiApiKeyForSection = activeExternalProvider
|
||||
? getExternalProviderApiKey(activeExternalProvider.id) || null
|
||||
: null;
|
||||
|
|
|
|||
|
|
@ -12,9 +12,6 @@ interface PendingDocChipsProps {
|
|||
onRemove: (id: string) => void;
|
||||
}
|
||||
|
||||
// Renders the upload/index/ready/error chips for in-flight RAG
|
||||
// document uploads above the composer textarea. Mirrors the
|
||||
// inline chip markup SharedComposer uses on the empty state.
|
||||
export const PendingDocChips: FC<PendingDocChipsProps> = ({ docs, onRemove }) => {
|
||||
if (docs.length === 0) return null;
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -27,8 +27,6 @@ const DOCUMENT_EXTENSIONS = new Set([
|
|||
".htm",
|
||||
]);
|
||||
|
||||
// File-input accept attribute matching DOCUMENT_EXTENSIONS so the
|
||||
// browser picker filters to formats the RAG ingester actually handles.
|
||||
export const DOCUMENT_ACCEPT =
|
||||
".pdf,.txt,.md,.markdown,.docx,.html,.htm,application/pdf,text/plain,text/markdown,application/vnd.openxmlformats-officedocument.wordprocessingml.document,text/html";
|
||||
|
||||
|
|
@ -47,12 +45,7 @@ export interface UseThreadDocUploadsResult {
|
|||
isIndexing: boolean;
|
||||
}
|
||||
|
||||
// Encapsulates the per-thread RAG document upload lifecycle:
|
||||
// pick file → POST /api/rag/threads/{id}/documents → subscribe to
|
||||
// ingestion SSE → flip chip status → on send, clear chips (docs
|
||||
// live in the backend KB and don't need to be re-attached).
|
||||
// Used by both the empty-state SharedComposer and the in-thread
|
||||
// assistant-ui Composer so the upload UX is identical in both.
|
||||
/** Per-thread RAG upload: file → POST → SSE → chip status. */
|
||||
export function useThreadDocUploads(): UseThreadDocUploadsResult {
|
||||
const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
|
||||
const [pendingDocs, setPendingDocs] = useState<PendingDoc[]>([]);
|
||||
|
|
@ -83,9 +76,7 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
|
|||
d.id === id ? { ...d, status: "ready" } : d,
|
||||
),
|
||||
);
|
||||
// First successful ingest in an off-source thread should
|
||||
// flip the source to 'thread' so the tool / pre-fetch
|
||||
// path has somewhere to search.
|
||||
// First ingest on an off-source thread → flip to 'thread'.
|
||||
if (useChatRuntimeStore.getState().ragSource.kind === "off") {
|
||||
useChatRuntimeStore
|
||||
.getState()
|
||||
|
|
@ -126,9 +117,7 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult {
|
|||
void useRagStore
|
||||
.getState()
|
||||
.deleteDocument(doc.documentId, `thread:${activeThreadId ?? ""}`)
|
||||
.catch(() => {
|
||||
// Best effort — the chip is going away regardless.
|
||||
});
|
||||
.catch(() => {});
|
||||
}
|
||||
return prev.filter((d) => d.id !== id);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -580,12 +580,10 @@ export function SharedComposer({
|
|||
});
|
||||
continue;
|
||||
}
|
||||
// Handle RAG document files (route to per-thread ingest pipeline).
|
||||
if (isDocumentFile(file)) {
|
||||
addDoc(file);
|
||||
continue;
|
||||
}
|
||||
// Handle image files
|
||||
if (!file.type.match(/^image\/(jpeg|png|webp|gif)$/i)) continue;
|
||||
if (file.size > MAX_IMAGE_SIZE) continue;
|
||||
if (attachUnavailableReason) {
|
||||
|
|
@ -614,9 +612,7 @@ export function SharedComposer({
|
|||
doc.documentId,
|
||||
`thread:${activeThreadId ?? ""}`,
|
||||
)
|
||||
.catch(() => {
|
||||
// Best effort — the chip is going away regardless.
|
||||
});
|
||||
.catch(() => {});
|
||||
}
|
||||
return prev.filter((d) => d.id !== id);
|
||||
});
|
||||
|
|
@ -709,9 +705,7 @@ export function SharedComposer({
|
|||
setPendingImages([]);
|
||||
setPendingAudio(null);
|
||||
clearPendingAudioStore();
|
||||
// Docs remain in the backend (already uploaded); just drop the
|
||||
// composer-side chips. The settings panel still shows them in the
|
||||
// thread's document list.
|
||||
// Docs stay in backend; drop chips only.
|
||||
setPendingDocs([]);
|
||||
textareaRef.current?.focus();
|
||||
|
||||
|
|
@ -1277,12 +1271,7 @@ export function SharedComposer({
|
|||
<span>Images</span>
|
||||
</button>
|
||||
)}
|
||||
{/* RAG: master switch for retrieval. On local models with
|
||||
tool-use support, registers `search_knowledge_base` as a
|
||||
tool the LLM can call. On external providers, falls back
|
||||
to the pre-fetch path. The sidebar Retrieval section
|
||||
configures the source / top-K / reranker; the button is
|
||||
the only on/off control. */}
|
||||
{/* Master RAG toggle; sidebar Retrieval section configures the rest. */}
|
||||
<button
|
||||
type="button"
|
||||
disabled={!modelLoaded}
|
||||
|
|
|
|||
|
|
@ -301,10 +301,7 @@ type ChatRuntimeStore = {
|
|||
ragMode: RagMode;
|
||||
enableRerank: boolean;
|
||||
ragTopK: number;
|
||||
// Cosine-similarity floor for RAG hits. Off (0) by default — set
|
||||
// > 0 to drop chunks below the threshold so unrelated docs don't
|
||||
// get injected into the prompt (e.g., asking about the weather
|
||||
// when the indexed docs are about economics).
|
||||
// Cosine floor; 0 disables. Set > 0 to drop off-topic hits.
|
||||
ragMinScore: number;
|
||||
hydratePersistedSettings: () => Promise<void>;
|
||||
setModelLoading: (loading: boolean) => void;
|
||||
|
|
@ -593,9 +590,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
supportsBuiltinCodeExecution: false,
|
||||
supportsBuiltinImageGeneration: false,
|
||||
toolsEnabled: loadBool(CHAT_TOOLS_ENABLED_KEY, false),
|
||||
// Phase 4: RAG button defaults off. Migration nudge happens after
|
||||
// settings hydration, when the persisted ragSource becomes visible —
|
||||
// see hydratePersistedSettings.
|
||||
// Defaults off; hydratePersistedSettings nudges it on for existing users.
|
||||
ragToolEnabled: loadBool(CHAT_RAG_TOOL_ENABLED_KEY, false),
|
||||
codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false),
|
||||
imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false),
|
||||
|
|
@ -651,11 +646,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
),
|
||||
...getHydratedSettingsState(settings, state, hydrationVersions),
|
||||
};
|
||||
// Phase 4 migration: pre-existing users with ragSource set
|
||||
// before the RAG button shipped should keep getting RAG —
|
||||
// auto-flip ragToolEnabled so the button starts ON for them.
|
||||
// The CHAT_RAG_TOOL_ENABLED_KEY localStorage write makes the
|
||||
// migration stick across reloads.
|
||||
// Migration: pre-existing ragSource → auto-enable ragToolEnabled.
|
||||
const hydratedRagSource =
|
||||
(nextState.ragSource as RagSource | undefined) ?? state.ragSource;
|
||||
if (
|
||||
|
|
|
|||
|
|
@ -45,9 +45,6 @@ export interface SearchHit {
|
|||
score: number;
|
||||
page_number: number | null;
|
||||
filename: string | null;
|
||||
// Phase 3B-multimodal: "image" hits have an image_url (served by
|
||||
// GET /api/rag/images/{document_id}/{filename}); "caption" hits
|
||||
// carry the paired image's caption text and are LLM-friendly.
|
||||
kind?: "text" | "image" | "caption";
|
||||
image_url?: string | null;
|
||||
}
|
||||
|
|
@ -94,9 +91,7 @@ async function throwOnError(response: Response): Promise<void> {
|
|||
throw new Error(parseErrorText(response.status, body));
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
// Knowledge bases
|
||||
// ------------------------------------------------------------------
|
||||
// --- Knowledge bases ---
|
||||
|
||||
export async function listKnowledgeBases(): Promise<KnowledgeBase[]> {
|
||||
const response = await authFetch("/api/rag/knowledge-bases");
|
||||
|
|
@ -108,8 +103,6 @@ export interface CreateKnowledgeBaseRequest {
|
|||
name: string;
|
||||
description?: string;
|
||||
embedding_model?: string;
|
||||
// Phase 3: both default server-side to "standard" / "text" — clients
|
||||
// that pre-date the field send the same payloads as before.
|
||||
chunking_strategy?: ChunkingStrategy;
|
||||
mode?: KBMode;
|
||||
}
|
||||
|
|
@ -133,9 +126,7 @@ export async function deleteKnowledgeBase(kbId: string): Promise<void> {
|
|||
await throwOnError(response);
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
// Documents
|
||||
// ------------------------------------------------------------------
|
||||
// --- Documents ---
|
||||
|
||||
export async function listKBDocuments(kbId: string): Promise<RagDocument[]> {
|
||||
const response = await authFetch(
|
||||
|
|
@ -313,9 +304,7 @@ export async function setRagDefaults(
|
|||
return parseJsonOrThrow<RagDefaults>(response);
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
// Search
|
||||
// ------------------------------------------------------------------
|
||||
// --- Search ---
|
||||
|
||||
export async function search(req: SearchRequest): Promise<SearchHit[]> {
|
||||
const response = await authFetch("/api/rag/search", {
|
||||
|
|
@ -327,19 +316,10 @@ export async function search(req: SearchRequest): Promise<SearchHit[]> {
|
|||
return body.hits;
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
// Ingestion SSE
|
||||
// ------------------------------------------------------------------
|
||||
// --- Ingestion SSE ---
|
||||
|
||||
/**
|
||||
* Subscribe to an ingestion job's SSE event stream.
|
||||
*
|
||||
* EventSource cannot send custom headers, so we pass the bearer token
|
||||
* as a `?token=…` query param. The backend's /jobs/{id}/events route
|
||||
* accepts either the query param or the Authorization header.
|
||||
*
|
||||
* Returns an unsubscribe function that closes the EventSource.
|
||||
*/
|
||||
/** Subscribe to a job's SSE stream; returns an unsubscribe fn.
|
||||
* Token goes via ?token= since EventSource cannot send headers. */
|
||||
export function subscribeToJobEvents(
|
||||
jobId: string,
|
||||
handlers: {
|
||||
|
|
@ -367,8 +347,7 @@ export function subscribeToJobEvents(
|
|||
};
|
||||
|
||||
source.onerror = () => {
|
||||
// Browser auto-reconnects unless we close explicitly. We close on
|
||||
// any error so the consumer can decide whether to resubscribe.
|
||||
// Browser auto-reconnects unless we close; let the consumer decide instead.
|
||||
source.close();
|
||||
handlers.onError?.(new Error("SSE connection lost"));
|
||||
handlers.onClose?.();
|
||||
|
|
|
|||
|
|
@ -9,35 +9,19 @@ import { useEffect, useRef, useState } from "react";
|
|||
import { useRagStore } from "../stores/rag-store";
|
||||
import { IngestionProgress } from "./ingestion-progress";
|
||||
|
||||
/**
|
||||
* Floating progress stack for in-flight RAG ingestion jobs.
|
||||
*
|
||||
* Watches the rag-store `jobs` map (populated by `subscribeJob`) and
|
||||
* renders one card per job that's neither completed nor errored. On
|
||||
* completion the card stays for a few seconds with a success state
|
||||
* before fading out, so users notice the run finished even if they
|
||||
* weren't watching the per-doc progress chips.
|
||||
*
|
||||
* Mounted once at the app root (`__root.tsx`) so it follows users
|
||||
* across pages.
|
||||
*/
|
||||
/** Floating per-job progress stack; mounted at the app root. */
|
||||
|
||||
const DISMISS_DELAY_MS = 4000;
|
||||
|
||||
export function IngestionToastStack() {
|
||||
const jobs = useRagStore((s) => s.jobs);
|
||||
const reduced = useReducedMotion();
|
||||
// Per-job timeout handles so terminal toasts auto-dismiss.
|
||||
const [dismissedJobs, setDismissedJobs] = useState<Set<string>>(
|
||||
() => new Set(),
|
||||
);
|
||||
|
||||
// Schedule auto-dismiss for jobs that have reached a terminal state.
|
||||
// Ref-tracked so `dismissedJobs` isn't a useEffect dep — the setter
|
||||
// fires *inside* the effect, and depending on its output here is a
|
||||
// recipe for update-depth loops if the scheduler ever runs faster
|
||||
// than the cleanup. We snapshot the latest dismissed set into a ref
|
||||
// and read from it inside the scheduling loop instead.
|
||||
// Ref-tracked dismissed set so it stays out of the effect deps
|
||||
// (avoids update-depth loops when the scheduler races cleanup).
|
||||
const scheduledJobsRef = useRef<Set<string>>(new Set());
|
||||
const dismissedJobsRef = useRef<Set<string>>(dismissedJobs);
|
||||
dismissedJobsRef.current = dismissedJobs;
|
||||
|
|
|
|||
|
|
@ -41,8 +41,6 @@ export function KBCreateDialog({
|
|||
const defaults = useRagStore((s) => s.defaults);
|
||||
const loadDefaults = useRagStore((s) => s.loadDefaults);
|
||||
|
||||
// Cache the initial values so reset() puts us back to the latest
|
||||
// saved defaults rather than the hard-coded ones.
|
||||
const initialStrategy: ChunkingStrategy =
|
||||
defaults?.chunking_strategy ?? "standard";
|
||||
const initialMode: KBMode = defaults?.mode ?? "text";
|
||||
|
|
@ -57,7 +55,6 @@ export function KBCreateDialog({
|
|||
const [submitting, setSubmitting] = useState(false);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
|
||||
// Fetch defaults on first open and re-sync the form when they arrive.
|
||||
useEffect(() => {
|
||||
if (open && !defaults) {
|
||||
void loadDefaults();
|
||||
|
|
@ -70,7 +67,7 @@ export function KBCreateDialog({
|
|||
setMode(defaults.mode);
|
||||
setEmbeddingModel(defaults.embedding_model ?? "");
|
||||
}
|
||||
// Intentional: only when `open` flips, not on every defaults change.
|
||||
// Only on open-flip, not on every defaults change.
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [open]);
|
||||
|
||||
|
|
@ -84,8 +81,7 @@ export function KBCreateDialog({
|
|||
setSubmitting(false);
|
||||
};
|
||||
|
||||
// Forbidden combo: multimodal + late chunking. We disable each side
|
||||
// when the other side is picking it, with a tooltip explaining why.
|
||||
// Forbid (multimodal, late); disable the other side when one is picked.
|
||||
const lateDisabled = mode === "multimodal";
|
||||
const multimodalDisabled = chunkingStrategy === "late";
|
||||
|
||||
|
|
|
|||
|
|
@ -14,15 +14,7 @@ import { useEffect, useState } from "react";
|
|||
import type { ChunkingStrategy, KBMode } from "../api/rag-api";
|
||||
import { useRagStore } from "../stores/rag-store";
|
||||
|
||||
/**
|
||||
* Defaults for newly-created KBs, surfaced in the Settings → Knowledge
|
||||
* Bases tab. Values pre-fill the KB create dialog so power users can
|
||||
* pick their preferred strategy/mode once instead of toggling per KB.
|
||||
*
|
||||
* Backed by chat_settings via PUT /api/rag/defaults. The selects
|
||||
* enforce the same (multimodal, late) constraint as the create
|
||||
* dialog.
|
||||
*/
|
||||
/** Defaults pre-fill the KB create dialog. Same (multimodal, late) rejection as create. */
|
||||
export function RagDefaultsSection() {
|
||||
const defaults = useRagStore((s) => s.defaults);
|
||||
const loadDefaults = useRagStore((s) => s.loadDefaults);
|
||||
|
|
@ -34,7 +26,6 @@ export function RagDefaultsSection() {
|
|||
const [embeddingModel, setEmbeddingModel] = useState("");
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
|
||||
// Load defaults on mount; reflect them in local state.
|
||||
useEffect(() => {
|
||||
void loadDefaults();
|
||||
}, [loadDefaults]);
|
||||
|
|
|
|||
|
|
@ -4,10 +4,7 @@
|
|||
import { useEffect } from "react";
|
||||
import { useRagStore } from "../stores/rag-store";
|
||||
|
||||
/**
|
||||
* Subscribe to an ingestion job's SSE events. Returns the latest event
|
||||
* for that job from the global jobs map. Pass `null` to skip.
|
||||
*/
|
||||
/** Subscribe to a job's SSE; returns latest event, null skips. */
|
||||
export function useIngestionEvents(jobId: string | null) {
|
||||
const event = useRagStore((s) =>
|
||||
jobId ? (s.jobs[jobId] ?? null) : null,
|
||||
|
|
|
|||
|
|
@ -5,11 +5,8 @@ import { useEffect } from "react";
|
|||
import type { RagDocument } from "../api/rag-api";
|
||||
import { kbScopeKey, threadScopeKey, useRagStore } from "../stores/rag-store";
|
||||
|
||||
// Module-scope sentinel so the selector returns a stable reference
|
||||
// when the scope key isn't populated yet. A `[]` literal in the
|
||||
// selector returns a new array on every call → Zustand's Object.is
|
||||
// snapshot check flags it as changed → re-render → selector reruns
|
||||
// → new `[]` → infinite loop → React error #185.
|
||||
// Stable sentinel: inline `[]` in the selector causes React error #185
|
||||
// (new ref each call → Zustand re-renders → infinite loop).
|
||||
const EMPTY_DOCS: RagDocument[] = [];
|
||||
|
||||
export function useKBDocuments(kbId: string | null) {
|
||||
|
|
|
|||
|
|
@ -194,14 +194,11 @@ export const useRagStore = create<RagStoreState>((set, get) => ({
|
|||
scope.kind === "kb"
|
||||
? kbScopeKey(scope.kbId)
|
||||
: threadScopeKey(scope.threadId);
|
||||
// Refresh the doc list so the new pending row appears.
|
||||
if (scope.kind === "kb") {
|
||||
void get().loadKBDocuments(scope.kbId);
|
||||
} else {
|
||||
void get().loadThreadDocuments(scope.threadId);
|
||||
}
|
||||
// Subscribe to the job; refresh docs on completion to pick up
|
||||
// the final num_chunks / status.
|
||||
get().subscribeJob(result.job_id, () => {
|
||||
if (scope.kind === "kb") {
|
||||
void get().loadKBDocuments(scope.kbId);
|
||||
|
|
@ -254,13 +251,8 @@ export const useRagStore = create<RagStoreState>((set, get) => ({
|
|||
|
||||
async reingestKB(kbId, opts) {
|
||||
const response = await apiReingestKB(kbId, opts ?? {});
|
||||
// Refresh the KB list so updated chunking_strategy / mode / embedder
|
||||
// values flow back into the UI, and the doc list so old chunk
|
||||
// counts reset to 0 until each job completes.
|
||||
void get().loadKnowledgeBases();
|
||||
void get().loadKBDocuments(kbId);
|
||||
// Subscribe to every new job so progress chips render and the doc
|
||||
// list refreshes on completion (mirrors uploadDocument's pattern).
|
||||
for (const jobId of response.job_ids) {
|
||||
get().subscribeJob(jobId, () => {
|
||||
void get().loadKBDocuments(kbId);
|
||||
|
|
@ -274,9 +266,7 @@ export const useRagStore = create<RagStoreState>((set, get) => ({
|
|||
const defaults = await apiGetRagDefaults();
|
||||
set({ defaults });
|
||||
} catch {
|
||||
// Defaults endpoint is best-effort; a 401 / network blip just
|
||||
// leaves defaults null and dialogs fall back to hard-coded
|
||||
// ('standard', 'text', no embedder override).
|
||||
// Best-effort; null defaults fall back to hard-coded UI defaults.
|
||||
}
|
||||
},
|
||||
|
||||
|
|
@ -290,8 +280,6 @@ export const useRagStore = create<RagStoreState>((set, get) => ({
|
|||
void get().loadThreadDocuments(threadId);
|
||||
void get().loadThreadIndexes();
|
||||
if (opts) {
|
||||
// The reingest endpoint persists the new settings as a side
|
||||
// effect; refresh the local copy so the UI reflects them.
|
||||
void get().loadThreadSettings(threadId);
|
||||
}
|
||||
for (const jobId of response.job_ids) {
|
||||
|
|
@ -309,7 +297,7 @@ export const useRagStore = create<RagStoreState>((set, get) => ({
|
|||
threadSettings: { ...state.threadSettings, [threadId]: settings },
|
||||
}));
|
||||
} catch {
|
||||
// Best-effort — falls back to defaults UI-side when missing.
|
||||
// Best-effort; UI falls back to defaults when missing.
|
||||
}
|
||||
},
|
||||
|
||||
|
|
@ -322,9 +310,7 @@ export const useRagStore = create<RagStoreState>((set, get) => ({
|
|||
},
|
||||
|
||||
subscribeJob(jobId, onComplete) {
|
||||
// Object indexing in TS returns V (not V|undefined) without
|
||||
// noUncheckedIndexedAccess, so we test membership explicitly to
|
||||
// avoid the "always-truthy function reference" lint.
|
||||
// `in` check avoids always-truthy lint without noUncheckedIndexedAccess.
|
||||
if (jobId in get().jobUnsubscribers) return;
|
||||
const unsubscribe = subscribeToJobEvents(jobId, {
|
||||
onEvent: (event) => {
|
||||
|
|
|
|||
|
|
@ -1205,9 +1205,7 @@ def install_python_stack() -> int:
|
|||
req = REQ_ROOT / "studio.txt",
|
||||
)
|
||||
|
||||
# 9. RAG dependencies (vector store + lexical index + layout-aware
|
||||
# parsers). Skipped in no-torch mode because RAG embeddings go
|
||||
# through sentence-transformers, which requires torch.
|
||||
# 9. RAG deps; skipped in no-torch mode (sentence-transformers requires torch).
|
||||
if not NO_TORCH:
|
||||
_progress("rag deps")
|
||||
pip_install(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue