diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 7cf3141ad1..f559a41631 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -246,209 +246,3 @@ def token_counter(model_name: str | None = None): return max(1, len(text) // 4) return _count - - -# --- Late chunking (Jina technique) --- - -_LATE_WINDOW_OVERLAP_TOKENS = 512 - - -def late_chunk_encode( - doc_text: str, - char_spans: list[tuple[int, int]], - *, - model_name: str | None = None, - normalize: bool = True, -): - """Single forward pass over the doc, mean-pool token embeddings per chunk span.""" - import numpy as np - - if not char_spans: - return [] - model = get_embedder(model_name) - tokenizer = model.tokenizer - max_length = int(getattr(model, "max_seq_length", None) or 8192) - - encoded = tokenizer( - doc_text, - return_tensors = "pt", - return_offsets_mapping = True, - add_special_tokens = True, - truncation = False, - ) - offsets = encoded.pop("offset_mapping")[0].tolist() - n_tokens = int(encoded["input_ids"].shape[1]) - - if n_tokens <= max_length: - token_embeddings = _encode_tokens(model, encoded) - return _pool_spans( - token_embeddings, - offsets, - char_spans, - normalize = normalize, - np_module = np, - model = model, - doc_text = doc_text, - ) - - logger.info( - "Late chunking: doc has %d tokens > model max %d; using windowed pass", - n_tokens, - max_length, - ) - return _windowed_late_chunk_encode( - doc_text = doc_text, - char_spans = char_spans, - model = model, - max_length = max_length, - normalize = normalize, - np_module = np, - ) - - -def _encode_tokens(model, encoded): - import torch - - transformer = model[0].auto_model - device = next(transformer.parameters()).device - inputs_on_device = {k: v.to(device) for k, v in encoded.items()} - with torch.no_grad(): - outputs = transformer(**inputs_on_device) - return outputs.last_hidden_state[0].detach().cpu().numpy() - - -def _pool_spans( - token_embeddings, - offsets, - char_spans, - *, - normalize: bool, - np_module, - model, - doc_text: str, - token_index_offset: int = 0, -): - """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: - # Skip special tokens whose offsets are (0, 0). - indices = [ - i - token_index_offset - for i, (ts, te) in enumerate(offsets) - if te > ts and te > char_start and ts < char_end - ] - indices = [i for i in indices if 0 <= i < n_rows] - if not indices: - vec = model.encode( - doc_text[char_start:char_end], - normalize_embeddings = normalize, - convert_to_numpy = True, - show_progress_bar = False, - ) - vectors.append(vec) - continue - pooled = token_embeddings[indices].mean(axis = 0) - if normalize: - denom = float(np_module.linalg.norm(pooled)) - if denom > 0: - pooled = pooled / denom - vectors.append(pooled) - return vectors - - -def _windowed_late_chunk_encode( - *, - doc_text: str, - char_spans: list[tuple[int, int]], - model, - max_length: int, - normalize: bool, - np_module, -): - """Doc > ctx window: pool each chunk against the window containing most of its tokens.""" - import torch - - tokenizer = model.tokenizer - transformer = model[0].auto_model - device = next(transformer.parameters()).device - - full = tokenizer( - doc_text, - return_tensors = "pt", - return_offsets_mapping = True, - add_special_tokens = False, - truncation = False, - ) - all_input_ids = full["input_ids"][0] - all_offsets = full["offset_mapping"][0].tolist() - n_tokens = int(all_input_ids.shape[0]) - stride = max(1, max_length - _LATE_WINDOW_OVERLAP_TOKENS) - - windows: list[tuple[int, int]] = [] - pos = 0 - while pos < n_tokens: - end = min(pos + max_length, n_tokens) - windows.append((pos, end)) - if end >= n_tokens: - break - pos += stride - - window_embeddings: dict[int, "np_module.ndarray"] = {} - - def _window_embeddings(window_index: int): - if window_index in window_embeddings: - return window_embeddings[window_index] - ws, we = windows[window_index] - win_ids = all_input_ids[ws:we].unsqueeze(0).to(device) - win_attn = torch.ones_like(win_ids) - with torch.no_grad(): - outputs = transformer(input_ids = win_ids, attention_mask = win_attn) - emb = outputs.last_hidden_state[0].detach().cpu().numpy() - window_embeddings[window_index] = emb - return emb - - vectors = [] - for char_start, char_end in char_spans: - chunk_token_indices = [ - i - for i, (ts, te) in enumerate(all_offsets) - if te > ts and te > char_start and ts < char_end - ] - if not chunk_token_indices: - vec = model.encode( - doc_text[char_start:char_end], - normalize_embeddings = normalize, - convert_to_numpy = True, - show_progress_bar = False, - ) - vectors.append(vec) - continue - best_window = 0 - best_overlap = 0 - for wi, (ws, we) in enumerate(windows): - overlap = sum(1 for ti in chunk_token_indices if ws <= ti < we) - if overlap > best_overlap: - best_overlap = overlap - best_window = wi - ws, _we = windows[best_window] - emb = _window_embeddings(best_window) - local_indices = [ - ti - ws for ti in chunk_token_indices if ws <= ti < ws + emb.shape[0] - ] - if not local_indices: - vec = model.encode( - doc_text[char_start:char_end], - normalize_embeddings = normalize, - convert_to_numpy = True, - show_progress_bar = False, - ) - vectors.append(vec) - continue - pooled = emb[local_indices].mean(axis = 0) - if normalize: - denom = float(np_module.linalg.norm(pooled)) - if denom > 0: - pooled = pooled / denom - vectors.append(pooled) - return vectors diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index bd151855b2..4aa51bffdd 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -58,7 +58,6 @@ def _subprocess_worker( overlap: int, batch_size: int, out_queue: Any, - chunking_strategy: str = "standard", mode: str = "text", document_id: str = "", vlm_url: str | None = None, @@ -128,7 +127,6 @@ def _subprocess_worker( out_queue.put({"type": "progress", "stage": "load_model", "progress": 0.1}) from core.rag.embeddings import ( get_embedder, - late_chunk_encode, token_counter, ) @@ -137,18 +135,6 @@ def _subprocess_worker( dim = int(model.get_sentence_embedding_dimension()) out_queue.put({"type": "dim", "dim": dim}) - if chunking_strategy == "late": - _run_late_chunking( - pages = pages, - stored_path = Path(stored_path), - chunk_size = chunk_size, - overlap = overlap, - counter = counter, - model_name = model_name, - late_chunk_encode = late_chunk_encode, - out_queue = out_queue, - ) - return text_count = _run_standard_chunking( pages = pages, @@ -352,67 +338,6 @@ def _stream_image_chunks( return len(out_chunks) -def _run_late_chunking( - *, - pages, - stored_path, - chunk_size, - overlap, - counter, - model_name, - late_chunk_encode, - out_queue, -) -> None: - """Chunk once, embed in one pass, ship all chunks in one chunks_batch.""" - from core.rag.chunking import chunk_pages_with_spans - from core.rag.locators import pdf_regions_for_chunks - - out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2}) - full_doc, chunks, char_spans = chunk_pages_with_spans( - pages, - max_tokens = chunk_size, - overlap_tokens = overlap, - token_counter = counter, - ) - if not chunks: - out_queue.put({"type": "error", "error": "chunker produced no chunks"}) - return - - out_queue.put({"type": "progress", "stage": "embed", "progress": 0.4}) - vectors = late_chunk_encode( - full_doc, - char_spans, - model_name = model_name, - normalize = True, - ) - pdf_regions = pdf_regions_for_chunks(stored_path, pages, chunks) - - out_queue.put({"type": "progress", "stage": "embed", "progress": 0.9}) - out_queue.put( - { - "type": "chunks_batch", - "first_index": 0, - "chunks": [ - { - "text": c.text, - "token_count": c.token_count, - "page_number": c.page_number, - "source_page_index": c.source_page_index, - "page_char_start": c.page_char_start, - "page_char_end": c.page_char_end, - "line_start": c.line_start, - "line_end": c.line_end, - "pdf_regions": pdf_regions[index], - "kind": "text", - } - for index, c in enumerate(chunks) - ], - "vectors": [v.tolist() for v in vectors], - } - ) - out_queue.put({"type": "complete", "num_chunks": len(chunks)}) - - # --- Job manager (parent side) --- @@ -830,7 +755,6 @@ def enqueue_ingestion( kb_id: str | None = None, thread_id: str | None = None, embedding_model: str | None = None, - chunking_strategy: str = "standard", mode: str = "text", enable_captions: bool = True, ) -> str: @@ -840,7 +764,7 @@ def enqueue_ingestion( scope = _scope_for(kb_id, thread_id) model_name = ( embedding_model - or resolve_embedder(mode, chunking_strategy) + or resolve_embedder(mode) or RAG_EMBEDDING_MODEL ) # Probe the loaded chat backend so the subprocess can caption figures with the @@ -892,7 +816,6 @@ def enqueue_ingestion( RAG_CHUNK_OVERLAP, RAG_EMBED_BATCH_SIZE, out_queue, - chunking_strategy, mode, document_id, vlm_url, diff --git a/studio/backend/core/rag/scope.py b/studio/backend/core/rag/scope.py index da5c224cda..c618fd0a14 100644 --- a/studio/backend/core/rag/scope.py +++ b/studio/backend/core/rag/scope.py @@ -39,11 +39,6 @@ def resolve_scope_embedder(scope: str) -> str | None: if explicit: return explicit mode = per_thread.get("mode") or defaults.get("mode") or "text" - chunking_strategy = ( - per_thread.get("chunking_strategy") - or defaults.get("chunking_strategy") - or "standard" - ) - return resolve_embedder(mode, chunking_strategy) + return resolve_embedder(mode) return None diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 37c2d232a9..263e5782ba 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -63,7 +63,6 @@ logger = get_logger(__name__) # --- Pydantic schemas --- -ChunkingStrategy = Literal["standard", "late"] KBMode = Literal["text", "multimodal"] @@ -71,7 +70,6 @@ class CreateKBRequest(BaseModel): name: str = Field(min_length = 1, max_length = 200) description: str | None = None embedding_model: str | None = None - chunking_strategy: ChunkingStrategy = "standard" mode: KBMode = "text" @@ -80,7 +78,6 @@ class KBResponse(BaseModel): name: str description: str | None embedding_model: str - chunking_strategy: ChunkingStrategy mode: KBMode created_at: int @@ -173,34 +170,17 @@ from core.rag.scope import resolve_scope_embedder as _resolve_scope_embedder # def _row_to_kb(row: Any) -> KBResponse: keys = row.keys() if hasattr(row, "keys") else () - chunking_strategy = ( - row["chunking_strategy"] if "chunking_strategy" in keys else "standard" - ) mode = row["mode"] if "mode" in keys else "text" return KBResponse( id = row["id"], name = row["name"], description = row["description"], embedding_model = row["embedding_model"], - chunking_strategy = chunking_strategy, mode = mode, created_at = row["created_at"], ) -def _validate_mode_combo(mode: KBMode, chunking_strategy: ChunkingStrategy) -> None: - """Reject (multimodal, late) — no embedder supports both at once.""" - if mode == "multimodal" and chunking_strategy == "late": - raise HTTPException( - status_code = 400, - detail = ( - "Late chunking is not supported in multimodal mode — " - "the multimodal embedder does not expose per-token " - "embeddings. Pick 'standard' chunking or 'text' mode." - ), - ) - - def _row_to_document(row: Any) -> DocumentResponse: return DocumentResponse( id = row["id"], @@ -300,7 +280,6 @@ def _start_ingestion( kb_id: str | None, thread_id: str | None, embedding_model: str, - chunking_strategy: str = "standard", mode: str = "text", caption_images: bool = True, content_hash: str | None = None, @@ -360,7 +339,6 @@ def _start_ingestion( kb_id = kb_id, thread_id = thread_id, embedding_model = embedding_model, - chunking_strategy = chunking_strategy, mode = mode, enable_captions = caption_images, ) @@ -387,13 +365,9 @@ def create_knowledge_base( ) -> KBResponse: from utils.rag.config import resolve_embedder - _validate_mode_combo(payload.mode, payload.chunking_strategy) - kb_id = str(uuid4()) - # No override: resolve from (mode, strategy) matrix. - embedding_model = payload.embedding_model or resolve_embedder( - payload.mode, payload.chunking_strategy - ) + # No override: resolve the embedder from the KB mode. + embedding_model = payload.embedding_model or resolve_embedder(payload.mode) created_at = _now_ms() with closing_connection() as conn: try: @@ -401,8 +375,8 @@ def create_knowledge_base( """ INSERT INTO rag_knowledge_bases (id, name, description, owner_user_id, embedding_model, - chunking_strategy, mode, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) + mode, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?) """, ( kb_id, @@ -410,7 +384,6 @@ def create_knowledge_base( payload.description, current_subject, embedding_model, - payload.chunking_strategy, payload.mode, created_at, ), @@ -426,7 +399,6 @@ def create_knowledge_base( name = payload.name, description = payload.description, embedding_model = embedding_model, - chunking_strategy = payload.chunking_strategy, mode = payload.mode, created_at = created_at, ) @@ -444,7 +416,6 @@ def list_knowledge_bases( class RagDefaults(BaseModel): - chunking_strategy: ChunkingStrategy = "standard" mode: KBMode = "text" embedding_model: str | None = None @@ -452,7 +423,6 @@ class RagDefaults(BaseModel): class UpdateRagDefaultsRequest(BaseModel): """Patch shape — only fields present overwrite stored values.""" - chunking_strategy: ChunkingStrategy | None = None mode: KBMode | None = None embedding_model: str | None = None @@ -466,7 +436,6 @@ def _load_rag_defaults() -> RagDefaults: if not isinstance(raw, dict): raw = {} return RagDefaults( - chunking_strategy = raw.get("chunking_strategy") or "standard", mode = raw.get("mode") or "text", embedding_model = raw.get("embedding_model"), ) @@ -492,10 +461,7 @@ def warmup_rag_embedder( from utils.rag.config import resolve_embedder defaults = _load_rag_defaults() - model_name = defaults.embedding_model or resolve_embedder( - defaults.mode, - defaults.chunking_strategy, - ) + model_name = defaults.embedding_model or resolve_embedder(defaults.mode) try: embeddings.get_embedder(model_name) except Exception as exc: # noqa: BLE001 @@ -511,7 +477,6 @@ def set_rag_defaults( current_subject: str = Depends(get_current_subject), ) -> RagDefaults: current = _load_rag_defaults() - new_strategy = payload.chunking_strategy or current.chunking_strategy new_mode = payload.mode or current.mode # PATCH-style: empty string clears, null/missing keeps current. if payload.embedding_model is None: @@ -520,32 +485,27 @@ def set_rag_defaults( new_embedder = None else: new_embedder = payload.embedding_model.strip() - _validate_mode_combo(new_mode, new_strategy) upsert_chat_settings_merge( { _DEFAULTS_KEY: { - "chunking_strategy": new_strategy, "mode": new_mode, "embedding_model": new_embedder, } } ) return RagDefaults( - chunking_strategy = new_strategy, mode = new_mode, embedding_model = new_embedder, ) class ThreadRagSettings(BaseModel): - chunking_strategy: ChunkingStrategy = "standard" mode: KBMode = "text" embedding_model: str | None = None class UpdateThreadRagSettingsRequest(BaseModel): - chunking_strategy: ChunkingStrategy | None = None mode: KBMode | None = None embedding_model: str | None = None # Reingest-only (not persisted); omit or None keeps captioning on. @@ -564,7 +524,6 @@ def _load_thread_settings(thread_id: str) -> ThreadRagSettings: raw = {} fallback = _load_rag_defaults() return ThreadRagSettings( - chunking_strategy = (raw.get("chunking_strategy") or fallback.chunking_strategy), mode = raw.get("mode") or fallback.mode, embedding_model = raw.get("embedding_model") or fallback.embedding_model, ) @@ -591,7 +550,6 @@ def set_thread_rag_settings( current_subject: str = Depends(get_current_subject), ) -> ThreadRagSettings: current = _load_thread_settings(thread_id) - new_strategy = payload.chunking_strategy or current.chunking_strategy new_mode = payload.mode or current.mode if payload.embedding_model is None: new_embedder = current.embedding_model @@ -599,19 +557,16 @@ def set_thread_rag_settings( new_embedder = None else: new_embedder = payload.embedding_model.strip() - _validate_mode_combo(new_mode, new_strategy) upsert_chat_settings_merge( { _thread_settings_key(thread_id): { - "chunking_strategy": new_strategy, "mode": new_mode, "embedding_model": new_embedder, } } ) return ThreadRagSettings( - chunking_strategy = new_strategy, mode = new_mode, embedding_model = new_embedder, ) @@ -620,7 +575,6 @@ def set_thread_rag_settings( class ReingestKBRequest(BaseModel): """All fields optional — omitting one keeps the KB's current value.""" - chunking_strategy: ChunkingStrategy | None = None mode: KBMode | None = None embedding_model: str | None = None # Not persisted on the KB; omit or None keeps captioning on for the rebuild. @@ -636,7 +590,6 @@ def _reingest_scope( *, kb_id: str | None, thread_id: str | None, - chunking_strategy: str, mode: str, embedding_model: str, caption_images: bool = True, @@ -685,7 +638,6 @@ def _reingest_scope( kb_id = kb_id, thread_id = thread_id, embedding_model = embedding_model, - chunking_strategy = chunking_strategy, mode = mode, caption_images = caption_images, ) @@ -707,37 +659,31 @@ def reingest_knowledge_base( kb_row = _kb_or_404(kb_id) keys = kb_row.keys() if hasattr(kb_row, "keys") else () - current_strategy = ( - kb_row["chunking_strategy"] if "chunking_strategy" in keys else "standard" - ) current_mode = kb_row["mode"] if "mode" in keys else "text" current_embedder = kb_row["embedding_model"] - new_strategy = payload.chunking_strategy or current_strategy new_mode = payload.mode or current_mode - _validate_mode_combo(new_mode, new_strategy) new_embedder = payload.embedding_model or ( current_embedder - if (new_strategy == current_strategy and new_mode == current_mode) - else resolve_embedder(new_mode, new_strategy) + if new_mode == current_mode + else resolve_embedder(new_mode) ) with closing_connection() as conn: conn.execute( """ UPDATE rag_knowledge_bases - SET chunking_strategy = ?, mode = ?, embedding_model = ? + SET mode = ?, embedding_model = ? WHERE id = ? """, - (new_strategy, new_mode, new_embedder, kb_id), + (new_mode, new_embedder, kb_id), ) conn.commit() return _reingest_scope( kb_id = kb_id, thread_id = None, - chunking_strategy = new_strategy, mode = new_mode, embedding_model = new_embedder, caption_images = payload.caption_images is not False, @@ -758,11 +704,7 @@ def reingest_thread_documents( if payload is None: payload = UpdateThreadRagSettingsRequest() - if ( - payload.chunking_strategy is not None - or payload.mode is not None - or payload.embedding_model is not None - ): + if payload.mode is not None or payload.embedding_model is not None: settings = set_thread_rag_settings( thread_id, payload, @@ -771,14 +713,10 @@ def reingest_thread_documents( else: settings = _load_thread_settings(thread_id) - embedder = settings.embedding_model or resolve_embedder( - settings.mode, - settings.chunking_strategy, - ) + embedder = settings.embedding_model or resolve_embedder(settings.mode) return _reingest_scope( kb_id = None, thread_id = thread_id, - chunking_strategy = settings.chunking_strategy, mode = settings.mode, embedding_model = embedder, caption_images = payload.caption_images is not False, @@ -816,11 +754,8 @@ async def upload_kb_document( ) -> UploadResponse: kb_row = _kb_or_404(kb_id) stored_path, filename, byte_size, content_hash = await _save_upload(file) - # Tolerate pre-Phase-3 rows missing chunking_strategy/mode. + # Tolerate pre-Phase-3 rows missing 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" - ) mode = kb_row["mode"] if "mode" in kb_keys else "text" return _start_ingestion( filename = filename, @@ -830,7 +765,6 @@ async def upload_kb_document( kb_id = kb_id, thread_id = None, embedding_model = kb_row["embedding_model"], - chunking_strategy = chunking_strategy, mode = mode, caption_images = caption_images, content_hash = content_hash, @@ -849,10 +783,7 @@ async def upload_thread_document( # No chat_threads check — fresh threads aren't persisted until first run. stored_path, filename, byte_size, content_hash = await _save_upload(file) settings = _load_thread_settings(thread_id) - embedder = settings.embedding_model or resolve_embedder( - settings.mode, - settings.chunking_strategy, - ) + embedder = settings.embedding_model or resolve_embedder(settings.mode) return _start_ingestion( filename = filename, stored_path = stored_path, @@ -861,7 +792,6 @@ async def upload_thread_document( kb_id = None, thread_id = thread_id, embedding_model = embedder, - chunking_strategy = settings.chunking_strategy, mode = settings.mode, caption_images = caption_images, content_hash = content_hash, diff --git a/studio/backend/utils/rag/config.py b/studio/backend/utils/rag/config.py index 278a866613..b627bf3e59 100644 --- a/studio/backend/utils/rag/config.py +++ b/studio/backend/utils/rag/config.py @@ -31,8 +31,7 @@ RAG_EMBEDDING_MODEL: str = ( or "BAAI/bge-small-en-v1.5" ) -# Default embedder per (mode, chunking). (multimodal, late) is rejected at -# KB-create time in routes/rag.py. +# Default embedder per mode. # # Text mode is the default: PDF figures are captioned at ingest (chat VLM or # helper gemma-3n fallback) and spliced into the page markdown before chunking, @@ -43,19 +42,15 @@ RAG_EMBEDDING_MODEL: str = ( # Alternative multimodal embedders kept for manual override: # - "BAAI/BGE-VL-large" — smaller (~400 M / 768-d) but CLIP-family with a # 77-token text cap; routed via `_BGEVLAdapter` in core/rag/embeddings.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", - ("multimodal", "standard"): "Qwen/Qwen3-VL-Embedding-2B", +RAG_EMBEDDER_MATRIX: dict[str, str] = { + "text": "BAAI/bge-small-en-v1.5", + "multimodal": "Qwen/Qwen3-VL-Embedding-2B", } -def resolve_embedder(mode: str, chunking_strategy: str) -> str: - """Embedder for (mode, chunking); unknown combos fall back to RAG_EMBEDDING_MODEL.""" - return RAG_EMBEDDER_MATRIX.get( - (mode, chunking_strategy), - RAG_EMBEDDING_MODEL, - ) +def resolve_embedder(mode: str) -> str: + """Embedder for the given mode; unknown modes fall back to RAG_EMBEDDING_MODEL.""" + return RAG_EMBEDDER_MATRIX.get(mode, RAG_EMBEDDING_MODEL) RAG_CHUNK_SIZE: int = _env_int("UNSLOTH_RAG_CHUNK_SIZE", 512) diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 8f28119150..749a34cc74 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -45,10 +45,7 @@ import { TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { - type KBMode, - type ChunkingStrategy as RagChunkingStrategy, -} from "@/features/rag/api/rag-api"; +import { type KBMode } from "@/features/rag/api/rag-api"; import { DocumentRow } from "@/features/rag/components/document-row"; import { KBCreateDialog } from "@/features/rag/components/kb-create-dialog"; import { PreviewPanel } from "@/features/rag/components/preview-panel"; @@ -562,10 +559,6 @@ export function ChatSettingsPanel({ } }, [ragSource.kind, activeThreadId, loadThreadSettings]); - const effectiveThreadChunking: RagChunkingStrategy = - threadSettings?.chunking_strategy ?? - ragDefaults?.chunking_strategy ?? - "standard"; const effectiveThreadMode: KBMode = threadSettings?.mode ?? ragDefaults?.mode ?? "text"; @@ -592,7 +585,6 @@ export function ChatSettingsPanel({ }; const applyThreadSettingChange = (patch: { - chunking_strategy?: RagChunkingStrategy; mode?: KBMode; }) => { void (async () => { @@ -1422,7 +1414,6 @@ export function ChatSettingsPanel({ ) : null} {knowledgeBases.map((kb) => { const isActive = kb.id === activeKbId; - const isLate = kb.chunking_strategy === "late"; const isMultimodal = kb.mode === "multimodal"; return ( {kb.name} - {isLate ? ( - - ⚡ Late - - ) : null} {isMultimodal ? ( {ragSource.kind === "thread" ? ( <> -
-
- - -
-
- - -
+
+ +

- Changing either setting will re-index this thread's - existing documents. + Changing the mode will re-index this thread's existing + documents.

-
- - -

- Late chunking embeds the whole document in one pass, so each - chunk vector carries full-document context. Slower to ingest - (one forward pass per doc) but improves retrieval on long, - cross-referenced text. Cannot be combined with multimodal - mode. + (~1.5 GB VRAM).

diff --git a/studio/frontend/src/features/rag/components/kb-detail-panel.tsx b/studio/frontend/src/features/rag/components/kb-detail-panel.tsx index b64ab87d63..f0cb689d65 100644 --- a/studio/frontend/src/features/rag/components/kb-detail-panel.tsx +++ b/studio/frontend/src/features/rag/components/kb-detail-panel.tsx @@ -134,7 +134,6 @@ export function KBDetailPanel({ wrapping next to the action buttons. */}

{kb.mode === "multimodal" ? "🖼️ Multimodal · " : ""} - {kb.chunking_strategy === "late" ? "⚡ Late · " : ""} Embedder: {kb.embedding_model}

diff --git a/studio/frontend/src/features/rag/components/kb-list.tsx b/studio/frontend/src/features/rag/components/kb-list.tsx index 404746efde..223e06d229 100644 --- a/studio/frontend/src/features/rag/components/kb-list.tsx +++ b/studio/frontend/src/features/rag/components/kb-list.tsx @@ -50,14 +50,6 @@ export function KBList({
{kb.name} - {kb.chunking_strategy === "late" ? ( - - ⚡ Late - - ) : null} {kb.mode === "multimodal" ? ( s.reingestKB); - const [chunkingStrategy, setChunkingStrategy] = useState( - kb.chunking_strategy, - ); const [mode, setMode] = useState(kb.mode); const [embeddingModel, setEmbeddingModel] = useState(""); const [submitting, setSubmitting] = useState(false); @@ -51,28 +44,17 @@ export function KBReconfigureDialog({ // Re-sync when the dialog opens against a different KB. useEffect(() => { if (open) { - setChunkingStrategy(kb.chunking_strategy); setMode(kb.mode); setEmbeddingModel(""); setError(null); setSubmitting(false); } - }, [open, kb.id, kb.chunking_strategy, kb.mode]); + }, [open, kb.id, kb.mode]); - const lateDisabled = mode === "multimodal"; - const multimodalDisabled = chunkingStrategy === "late"; - - const placeholderEmbedder = - mode === "multimodal" - ? `Current: ${kb.embedding_model}` - : chunkingStrategy === "late" - ? `Current: ${kb.embedding_model}` - : `Current: ${kb.embedding_model}`; + const placeholderEmbedder = `Current: ${kb.embedding_model}`; const changedSettings = - chunkingStrategy !== kb.chunking_strategy || - mode !== kb.mode || - embeddingModel.trim() !== ""; + mode !== kb.mode || embeddingModel.trim() !== ""; const handleSubmit = async (e: React.FormEvent) => { e.preventDefault(); @@ -91,7 +73,6 @@ export function KBReconfigureDialog({ setError(null); try { await reingestKB(kb.id, { - chunking_strategy: chunkingStrategy, mode, embedding_model: embeddingModel.trim() || undefined, caption_images: useChatRuntimeStore.getState().ragCaptionImages, @@ -110,7 +91,7 @@ export function KBReconfigureDialog({ Reconfigure “{kb.name}” - Change the chunking strategy, mode, or embedder for this KB. + Change the mode or embedder for this KB. All {documentCount} document{documentCount === 1 ? "" : "s"}{" "} will be re-ingested from the originals on disk. @@ -127,47 +108,12 @@ export function KBReconfigureDialog({ Text only - + Multimodal — text + images
-
- - -

- Leave blank to keep the current model (or pick the matrix - default when mode/strategy changes). + Leave blank to keep the current model (or pick the default + when the mode changes).

{error ? ( diff --git a/studio/frontend/src/features/rag/components/rag-defaults-section.tsx b/studio/frontend/src/features/rag/components/rag-defaults-section.tsx index 78f1b801cd..10be81174d 100644 --- a/studio/frontend/src/features/rag/components/rag-defaults-section.tsx +++ b/studio/frontend/src/features/rag/components/rag-defaults-section.tsx @@ -10,17 +10,15 @@ import { SelectValue, } from "@/components/ui/select"; import { useEffect, useState } from "react"; -import type { ChunkingStrategy, KBMode } from "../api/rag-api"; +import type { KBMode } from "../api/rag-api"; import { useRagStore } from "../stores/rag-store"; -/** Defaults pre-fill the KB create dialog. Same (multimodal, late) rejection as create. */ +/** Defaults pre-fill the KB create dialog. */ export function RagDefaultsSection() { const defaults = useRagStore((s) => s.defaults); const loadDefaults = useRagStore((s) => s.loadDefaults); const updateDefaults = useRagStore((s) => s.updateDefaults); - const [chunkingStrategy, setChunkingStrategy] = - useState("standard"); const [mode, setMode] = useState("text"); const [error, setError] = useState(null); @@ -30,16 +28,11 @@ export function RagDefaultsSection() { useEffect(() => { if (defaults) { - setChunkingStrategy(defaults.chunking_strategy); setMode(defaults.mode); } }, [defaults]); - const lateDisabled = mode === "multimodal"; - const multimodalDisabled = chunkingStrategy === "late"; - const persist = (patch: { - chunking_strategy?: ChunkingStrategy; mode?: KBMode; embedding_model?: string | null; }) => { @@ -74,46 +67,7 @@ export function RagDefaultsSection() { Text only - - Multimodal - - - -
-
- -
diff --git a/tests/python/test_rag_late_chunking.py b/tests/python/test_rag_late_chunking.py deleted file mode 100644 index d9747d1d13..0000000000 --- a/tests/python/test_rag_late_chunking.py +++ /dev/null @@ -1,113 +0,0 @@ -"""Late chunking tests (Phase 3B-late). - -Pure-python coverage of `chunk_pages_with_spans` runs always. The -encoder test loads a small SentenceTransformer and is gated behind the -existing `server` marker so default `pytest` runs skip it. -""" - -import sys -from pathlib import Path - -import pytest - -REPO_ROOT = Path(__file__).resolve().parents[2] -STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" -if str(STUDIO_BACKEND) not in sys.path: - sys.path.insert(0, str(STUDIO_BACKEND)) - -from core.rag.chunking import chunk_pages_with_spans -from core.rag.parsers import ParsedPage - - -def _wc_counter(text: str) -> int: - return max(1, len(text.split())) - - -def test_spans_index_back_to_full_doc_text(): - pages = [ - ParsedPage(text = "# Section A\n\n" + ("alpha " * 20), page_number = 1), - ParsedPage(text = "# Section B\n\n" + ("beta " * 20), page_number = 2), - ] - full_doc, chunks, char_spans = chunk_pages_with_spans( - pages, - max_tokens = 12, - overlap_tokens = 0, - token_counter = _wc_counter, - ) - assert chunks - assert len(chunks) == len(char_spans) - for chunk, (start, end) in zip(chunks, char_spans): - # Chunk text must equal the full_doc slice it claims. - assert full_doc[start:end] == chunk.text - - -def test_chunks_inherit_page_number_by_overlap(): - pages = [ - ParsedPage(text = "page-one text here", page_number = 1), - ParsedPage(text = "page-two text here", page_number = 2), - ] - _full_doc, chunks, _spans = chunk_pages_with_spans( - pages, - max_tokens = 4, - overlap_tokens = 0, - token_counter = _wc_counter, - ) - pages_seen = {c.page_number for c in chunks} - assert pages_seen <= {1, 2} - # Each page contributes at least one chunk. - assert 1 in pages_seen - assert 2 in pages_seen - - -def test_full_doc_joins_pages_with_blank_line_separator(): - pages = [ - ParsedPage(text = "first", page_number = 1), - ParsedPage(text = "second", page_number = 2), - ] - full_doc, _chunks, _spans = chunk_pages_with_spans( - pages, - max_tokens = 5, - overlap_tokens = 0, - token_counter = _wc_counter, - ) - assert "first" in full_doc - assert "second" in full_doc - # Pages separated by exactly one blank line. - assert "first\n\nsecond" in full_doc - - -@pytest.mark.server -def test_late_chunk_encode_returns_one_vector_per_span(): - pytest.importorskip("sentence_transformers") - pytest.importorskip("torch") - # all-MiniLM-L6-v2: ~80MB, 384 dims. - import os - - os.environ.setdefault( - "UNSLOTH_RAG_EMBEDDING_MODEL", - "sentence-transformers/all-MiniLM-L6-v2", - ) - from core.rag import embeddings as embeddings_module - - embeddings_module._model = None # force re-load - embeddings_module._model_name = None - - doc_text = ( - "# Intro\n\n" - "The quick brown fox jumps over the lazy dog.\n\n" - "# Methods\n\n" - "We trained the model on a corpus of 100M tokens.\n\n" - "# Results\n\n" - "Accuracy improved by 12% over the baseline." - ) - # char_spans: one per section, picked manually. - char_spans = [ - (doc_text.index("The quick"), doc_text.index("\n\n# Methods")), - (doc_text.index("We trained"), doc_text.index("\n\n# Results")), - (doc_text.index("Accuracy"), len(doc_text)), - ] - vectors = embeddings_module.late_chunk_encode(doc_text, char_spans) - assert len(vectors) == len(char_spans) - dim = vectors[0].shape[0] - for v in vectors: - assert v.shape == (dim,) diff --git a/tests/python/test_rag_multimodal.py b/tests/python/test_rag_multimodal.py index 823c3609e6..17007bd013 100644 --- a/tests/python/test_rag_multimodal.py +++ b/tests/python/test_rag_multimodal.py @@ -7,7 +7,6 @@ returns images when asked, route accepts the mode field, constraint validator rejects illegal combos) run in every test invocation. """ -import importlib.util import sys from pathlib import Path @@ -19,27 +18,6 @@ if str(STUDIO_BACKEND) not in sys.path: sys.path.insert(0, str(STUDIO_BACKEND)) -def _rag_route(): - """Load ``routes/rag.py`` directly, bypassing the ``routes`` package. - - ``from routes.rag import X`` first runs ``routes/__init__.py``, which eagerly - imports every router — including the datasets router, whose chain does - ``from datasets import IterableDataset`` at import time. On a GPU-less CI - runner the unsloth bootstrap can leave ``datasets`` half-initialized, so that - eager import raises. These tests only need pure helpers from rag.py, so load - the file on its own (it has no intra-``routes`` imports). - """ - mod = sys.modules.get("_rag_route_under_test") - if mod is None: - spec = importlib.util.spec_from_file_location( - "_rag_route_under_test", STUDIO_BACKEND / "routes" / "rag.py" - ) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - sys.modules["_rag_route_under_test"] = mod - return mod - - def test_html_parser_returns_images_when_requested(tmp_path): pytest.importorskip("bs4") pytest.importorskip("lxml") @@ -74,32 +52,14 @@ def test_html_parser_returns_images_when_requested(tmp_path): assert img.nearest_caption == "A tiny figure" -def test_multimodal_late_combo_validator(): - from fastapi import HTTPException - - _validate_mode_combo = _rag_route()._validate_mode_combo - - # Allowed combos → None. - assert _validate_mode_combo("text", "standard") is None - assert _validate_mode_combo("text", "late") is None - assert _validate_mode_combo("multimodal", "standard") is None - - # Forbidden combo → 400. - with pytest.raises(HTTPException) as excinfo: - _validate_mode_combo("multimodal", "late") - assert excinfo.value.status_code == 400 - - -def test_rag_embedder_matrix_excludes_multimodal_late(): +def test_rag_embedder_matrix_is_keyed_by_mode(): from utils.rag.config import RAG_EMBEDDER_MATRIX, resolve_embedder - assert ("multimodal", "late") not in RAG_EMBEDDER_MATRIX - assert ("text", "standard") in RAG_EMBEDDER_MATRIX - assert ("text", "late") in RAG_EMBEDDER_MATRIX - assert ("multimodal", "standard") in RAG_EMBEDDER_MATRIX + assert "text" in RAG_EMBEDDER_MATRIX + assert "multimodal" in RAG_EMBEDDER_MATRIX - # Unknown combos fall back to the legacy default, not KeyError. - fallback = resolve_embedder("multimodal", "late") + # Unknown modes fall back to the default, not KeyError. + fallback = resolve_embedder("unknown-mode") assert isinstance(fallback, str) and fallback diff --git a/tests/python/test_rag_multimodal_integration.py b/tests/python/test_rag_multimodal_integration.py index 62eff4e683..2995fb7a26 100644 --- a/tests/python/test_rag_multimodal_integration.py +++ b/tests/python/test_rag_multimodal_integration.py @@ -93,7 +93,6 @@ def test_multimodal_subprocess_emits_image_and_caption_chunks( overlap = 20, batch_size = 4, out_queue = out_queue, - chunking_strategy = "standard", mode = "multimodal", document_id = "test-doc-1", ) diff --git a/tests/python/test_rag_reingest.py b/tests/python/test_rag_reingest.py index 37d2e66767..0891904f2a 100644 --- a/tests/python/test_rag_reingest.py +++ b/tests/python/test_rag_reingest.py @@ -2,8 +2,7 @@ Full end-to-end reingest needs a running studio + a real embedder; that's covered manually via the curl smoke flow in the plan. Here we cover the -parts that are testable without external models: payload validation and -the (multimodal, late) constraint propagation. +parts that are testable without external models: payload validation. """ import importlib.util @@ -43,22 +42,12 @@ def test_reingest_request_accepts_all_optional_fields(): ReingestKBRequest = _rag_route().ReingestKBRequest empty = ReingestKBRequest() - assert empty.chunking_strategy is None assert empty.mode is None assert empty.embedding_model is None - partial = ReingestKBRequest(chunking_strategy = "late") - assert partial.chunking_strategy == "late" - assert partial.mode is None - - -def test_reingest_request_rejects_unknown_strategy(): - from pydantic import ValidationError - - ReingestKBRequest = _rag_route().ReingestKBRequest - - with pytest.raises(ValidationError): - ReingestKBRequest(chunking_strategy = "telekinetic") + partial = ReingestKBRequest(mode = "multimodal") + assert partial.mode == "multimodal" + assert partial.embedding_model is None def test_reingest_request_rejects_unknown_mode(): @@ -68,14 +57,3 @@ def test_reingest_request_rejects_unknown_mode(): with pytest.raises(ValidationError): ReingestKBRequest(mode = "augmented") - - -def test_constraint_still_enforced_for_reingest_combos(): - """The combination guard is shared with create — verify it still bites.""" - from fastapi import HTTPException - - _validate_mode_combo = _rag_route()._validate_mode_combo - - with pytest.raises(HTTPException) as excinfo: - _validate_mode_combo("multimodal", "late") - assert excinfo.value.status_code == 400