diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index f559a41631..5e099ef655 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -22,10 +22,6 @@ _embedding_dim: int | None = None def _load(model_name: str) -> Any: logger.info("Loading RAG embedder: %s", model_name) - # 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: nomic-embed-text-v1.5 needs custom modeling for 8K ctx. @@ -36,133 +32,6 @@ def _load(model_name: str) -> Any: ) -class _BGEVLAdapter: - """SentenceTransformer-shaped adapter over BGE-VL's AutoModel.""" - - def __init__(self, hf_model_name: str): - from transformers import AutoModel - import torch - - self._model = AutoModel.from_pretrained( - hf_model_name, - trust_remote_code = True, - ) - # 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() - self._device = device - self._dim: int | None = None - - def _normalize(self, tensor): - import torch.nn.functional as F - - return F.normalize(tensor, p = 2.0, dim = -1) - - # CLIP positional embedding cap; longer text triggers shape mismatch. - _CLIP_TEXT_MAX_TOKENS = 77 - - def encode( - self, - inputs, - *, - batch_size: int = 32, - normalize_embeddings: bool = True, - convert_to_numpy: bool = True, - show_progress_bar: bool = False, - **_ignored, - ): - import io - - import numpy as np - import torch - from PIL import Image - - if inputs is None or len(inputs) == 0: - return np.zeros( - (0, self.get_sentence_embedding_dimension()), dtype = np.float32 - ) - - sample = inputs[0] - is_image = isinstance(sample, Image.Image) or isinstance( - sample, (bytes, bytearray) - ) - - chunks_out = [] - for start in range(0, len(inputs), batch_size): - batch = list(inputs[start : start + batch_size]) - if is_image: - # BGE-VL's data_process re-opens each item via Image.open(...), - # which needs a file-like (.read()) or path — NOT a pre-opened PIL - # Image. Pass BytesIO; PIL Images get rebuffered via an in-memory PNG. - file_likes: list[Any] = [] - for b in batch: - if isinstance(b, (bytes, bytearray)): - file_likes.append(io.BytesIO(b)) - elif isinstance(b, Image.Image): - buf = io.BytesIO() - b.save(buf, format = "PNG") - buf.seek(0) - file_likes.append(buf) - else: - file_likes.append(b) - with torch.no_grad(): - vecs = self._model.encode(images = file_likes) - else: - vecs = self._encode_text_truncated([str(t) for t in batch]) - if normalize_embeddings: - vecs = self._normalize(vecs) - chunks_out.append(vecs.detach().cpu()) - - out = torch.cat(chunks_out, dim = 0) - return out.numpy() if convert_to_numpy else out - - def _encode_text_truncated(self, texts: list[str]): - """Truncate to CLIP's 77-token limit; long text in multimodal mode is lossy.""" - import torch - - tokenizer = self._get_text_tokenizer() - inputs = tokenizer( - texts, - return_tensors = "pt", - padding = True, - truncation = True, - max_length = self._CLIP_TEXT_MAX_TOKENS, - ) - inputs = {k: v.to(self._device) for k, v in inputs.items()} - 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, - ) - with torch.no_grad(): - return self._model.get_text_features(**inputs) - - def _get_text_tokenizer(self): - processor = getattr(self._model, "processor", None) - if processor is not None: - tok = getattr(processor, "tokenizer", None) - if tok is not None: - return tok - tok = getattr(self._model, "tokenizer", None) - if tok is not None: - return tok - raise AttributeError("BGE-VL adapter could not locate a text tokenizer") - - def get_sentence_embedding_dimension(self) -> int: - if self._dim is None: - v = self.encode(["dim-probe"], batch_size = 1) - self._dim = int(v.shape[-1]) - return self._dim - - def tokenize(self, texts): - return self._get_text_tokenizer()( - texts, - return_tensors = "pt", - padding = True, - ) - - def get_embedder(model_name: str | None = None) -> Any: global _model, _model_name, _embedding_dim target = model_name or RAG_EMBEDDING_MODEL @@ -206,31 +75,6 @@ def encode( ) -def encode_images( - image_bytes_list: list[bytes], - *, - model_name: str | None = None, - batch_size: int | None = None, - normalize: bool = True, -): - """Embed image bytes via a CLIP-family multimodal encoder.""" - from io import BytesIO - - from PIL import Image - - if not image_bytes_list: - return [] - model = get_embedder(model_name) - images = [Image.open(BytesIO(b)).convert("RGB") for b in image_bytes_list] - return model.encode( - images, - batch_size = batch_size or RAG_EMBED_BATCH_SIZE, - normalize_embeddings = normalize, - convert_to_numpy = True, - show_progress_bar = False, - ) - - def token_counter(model_name: str | None = None): """Return a token-count callable backed by the embedder's tokenizer.""" model = get_embedder(model_name) diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 75b7a38e77..37178e84c0 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -58,8 +58,6 @@ def _subprocess_worker( overlap: int, batch_size: int, out_queue: Any, - mode: str = "text", - document_id: str = "", vlm_url: str | None = None, vlm_model: str | None = None, enable_captions: bool = True, @@ -83,8 +81,8 @@ def _subprocess_worker( from core.rag.parsers import inline_image_captions, parse out_queue.put({"type": "progress", "stage": "parse", "progress": 0.05}) - # Always extract images to caption + splice for both modes. Text mode uses - # captions inline in markdown; multimodal also embeds raw images as image-kind chunks. + # Extract images so figures can be captioned and spliced into the page + # markdown, where the chunker indexes them like any other text. parsed = parse(Path(stored_path), want_images = True) pages = parsed.pages if not pages and not parsed.images: @@ -95,7 +93,6 @@ def _subprocess_worker( # Caption figures once (chat VLM if available, else helper VLM), then splice # captions into the page markdown so the chunker indexes them like any text. - # Multimodal reuses these captions in _stream_image_chunks below — no duplicate VLM calls. captions: list[str] = [] if parsed.images and enable_captions: out_queue.put( @@ -147,17 +144,7 @@ def _subprocess_worker( out_queue = out_queue, send_complete = False, ) - image_count = 0 - if mode == "multimodal" and parsed.images and document_id: - image_count = _stream_image_chunks( - images = parsed.images, - document_id = document_id, - model_name = model_name, - out_queue = out_queue, - first_index = text_count, - precomputed_captions = captions, - ) - out_queue.put({"type": "complete", "num_chunks": text_count + image_count}) + out_queue.put({"type": "complete", "num_chunks": text_count}) except Exception as exc: # noqa: BLE001 logger.exception("ingestion subprocess failed") out_queue.put({"type": "error", "error": f"{type(exc).__name__}: {exc}"}) @@ -232,111 +219,6 @@ def _run_standard_chunking( return total -def _stream_image_chunks( - *, - images, - document_id: str, - model_name: str, - out_queue, - first_index: int, - precomputed_captions: list[str] | None = None, -) -> int: - """Persist images, emit image+caption chunks; pairs share pair_group. - - ``precomputed_captions`` come from the parent's earlier - caption_images call (used so we don't VLM-caption the same images - twice — once for markdown splicing, once for the caption-kind chunk). - If absent we fall back to each image's nearest_caption (page text). - """ - from core.rag.embeddings import encode, encode_images - from utils.paths.storage_roots import ensure_dir, rag_uploads_root - - if not images: - return 0 - - out_queue.put({"type": "progress", "stage": "extract_images", "progress": 0.85}) - - img_dir = ensure_dir(rag_uploads_root() / "images" / document_id) - - paths: list[str] = [] - bytes_for_encoding: list[bytes] = [] - captions: list[str] = [] - pages: list[int | None] = [] - pre_caps = precomputed_captions or [] - for idx, img in enumerate(images): - ext = _MIME_TO_EXT.get(img.mime_type, ".bin") - path = img_dir / f"img-{idx:04d}{ext}" - try: - path.write_bytes(img.image_bytes) - except OSError: - logger.warning("failed to save image; skipping", path = str(path)) - continue - paths.append(str(path)) - bytes_for_encoding.append(img.image_bytes) - vlm_cap = pre_caps[idx].strip() if idx < len(pre_caps) and pre_caps[idx] else "" - captions.append(vlm_cap or (img.nearest_caption or "")) - pages.append(img.page_number) - - if not paths: - return 0 - - image_vectors = encode_images(bytes_for_encoding, model_name = model_name) - - caption_to_image: list[int] = [i for i, cap in enumerate(captions) if cap.strip()] - if caption_to_image: - caption_vectors_arr = encode( - [captions[i] for i in caption_to_image], - model_name = model_name, - ) - caption_vectors = caption_vectors_arr.tolist() - else: - caption_vectors = [] - - out_chunks: list[dict] = [] - out_vectors: list[list[float]] = [] - cap_iter = iter(zip(caption_to_image, caption_vectors)) - next_cap = next(cap_iter, None) - - for idx, (path, page, caption) in enumerate(zip(paths, pages, captions)): - group_id = f"img-{idx:04d}" - out_chunks.append( - { - "text": caption[:1000] if caption else "", - "token_count": 0, - "page_number": page, - "kind": "image", - "image_path": path, - "pair_group": group_id, - } - ) - out_vectors.append(image_vectors[idx].tolist()) - if next_cap is not None and next_cap[0] == idx: - _cap_index, cap_vec = next_cap - out_chunks.append( - { - "text": caption, - "token_count": max(1, len(caption.split())), - "page_number": page, - "kind": "caption", - "image_path": None, - "pair_group": group_id, - } - ) - out_vectors.append(cap_vec) - next_cap = next(cap_iter, None) - - out_queue.put( - { - "type": "chunks_batch", - "first_index": first_index, - "chunks": out_chunks, - "vectors": out_vectors, - } - ) - out_queue.put({"type": "progress", "stage": "extract_images", "progress": 0.95}) - return len(out_chunks) - - # --- Job manager (parent side) --- @@ -754,17 +636,15 @@ def enqueue_ingestion( kb_id: str | None = None, thread_id: str | None = None, embedding_model: str | None = None, - mode: str = "text", enable_captions: bool = True, ) -> str: """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) - model_name = embedding_model or resolve_embedder(mode) or RAG_EMBEDDING_MODEL + model_name = embedding_model or resolve_embedder() or RAG_EMBEDDING_MODEL # Probe the loaded chat backend so the subprocess can caption figures with the - # user's own vision model (no extra VRAM). Runs for both modes — text splices - # captions into markdown, multimodal also feeds them to the image-vector encoder. + # user's own vision model (no extra VRAM). Text splices captions into markdown. # No vision chat model loaded → falls back to the helper VLM (pre-cached at startup). # Skipped when captioning is disabled for this upload. vlm_url: str | None = None @@ -811,8 +691,6 @@ def enqueue_ingestion( RAG_CHUNK_OVERLAP, RAG_EMBED_BATCH_SIZE, out_queue, - mode, - document_id, vlm_url, vlm_model, enable_captions, diff --git a/studio/backend/core/rag/scope.py b/studio/backend/core/rag/scope.py index c618fd0a14..76c1edd224 100644 --- a/studio/backend/core/rag/scope.py +++ b/studio/backend/core/rag/scope.py @@ -38,7 +38,6 @@ def resolve_scope_embedder(scope: str) -> str | None: explicit = per_thread.get("embedding_model") or defaults.get("embedding_model") if explicit: return explicit - mode = per_thread.get("mode") or defaults.get("mode") or "text" - return resolve_embedder(mode) + return resolve_embedder() return None diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index 3a3ebb3afa..3b5e796d17 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -120,15 +120,6 @@ def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str: tokens = hit.get("token_count") if tokens: attrs.append(f'tokens="{tokens}"') - kind = hit.get("kind") - if kind and kind != "text": - attrs.append(f'kind="{_xml_attr(kind)}"') - image_path = hit.get("image_path") - if kind == "image" and image_path and document_id: - # Mirror routes/rag.py search-response shape so the frontend tool card - # can render the image inline via the same route. - image_url = f"/api/rag/images/{document_id}/{Path(image_path).name}" - attrs.append(f'image_url="{_xml_attr(image_url)}"') text = (hit.get("text") or "").strip() blocks.append(f"\n{text}\n") return "\n\n".join(blocks) diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index d139a4af76..d4b9858a1c 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -63,14 +63,11 @@ logger = get_logger(__name__) # --- Pydantic schemas --- -KBMode = Literal["text", "multimodal"] - class CreateKBRequest(BaseModel): name: str = Field(min_length = 1, max_length = 200) description: str | None = None embedding_model: str | None = None - mode: KBMode = "text" class KBResponse(BaseModel): @@ -78,7 +75,6 @@ class KBResponse(BaseModel): name: str description: str | None embedding_model: str - mode: KBMode created_at: int @@ -169,14 +165,11 @@ 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 () - 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"], - mode = mode, created_at = row["created_at"], ) @@ -280,7 +273,6 @@ def _start_ingestion( kb_id: str | None, thread_id: str | None, embedding_model: str, - mode: str = "text", caption_images: bool = True, content_hash: str | None = None, ) -> UploadResponse: @@ -339,7 +331,6 @@ def _start_ingestion( kb_id = kb_id, thread_id = thread_id, embedding_model = embedding_model, - mode = mode, enable_captions = caption_images, ) return UploadResponse(document_id = document_id, job_id = job_id, filename = filename) @@ -366,8 +357,7 @@ def create_knowledge_base( from utils.rag.config import resolve_embedder kb_id = str(uuid4()) - # No override: resolve the embedder from the KB mode. - embedding_model = payload.embedding_model or resolve_embedder(payload.mode) + embedding_model = payload.embedding_model or resolve_embedder() created_at = _now_ms() with closing_connection() as conn: try: @@ -375,8 +365,8 @@ def create_knowledge_base( """ INSERT INTO rag_knowledge_bases (id, name, description, owner_user_id, embedding_model, - mode, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?) + created_at) + VALUES (?, ?, ?, ?, ?, ?) """, ( kb_id, @@ -384,7 +374,6 @@ def create_knowledge_base( payload.description, current_subject, embedding_model, - payload.mode, created_at, ), ) @@ -399,7 +388,6 @@ def create_knowledge_base( name = payload.name, description = payload.description, embedding_model = embedding_model, - mode = payload.mode, created_at = created_at, ) @@ -416,14 +404,12 @@ def list_knowledge_bases( class RagDefaults(BaseModel): - mode: KBMode = "text" embedding_model: str | None = None class UpdateRagDefaultsRequest(BaseModel): """Patch shape — only fields present overwrite stored values.""" - mode: KBMode | None = None embedding_model: str | None = None @@ -436,7 +422,6 @@ def _load_rag_defaults() -> RagDefaults: if not isinstance(raw, dict): raw = {} return RagDefaults( - mode = raw.get("mode") or "text", embedding_model = raw.get("embedding_model"), ) @@ -461,7 +446,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) + model_name = defaults.embedding_model or resolve_embedder() try: embeddings.get_embedder(model_name) except Exception as exc: # noqa: BLE001 @@ -477,7 +462,6 @@ def set_rag_defaults( current_subject: str = Depends(get_current_subject), ) -> RagDefaults: current = _load_rag_defaults() - new_mode = payload.mode or current.mode # PATCH-style: empty string clears, null/missing keeps current. if payload.embedding_model is None: new_embedder = current.embedding_model @@ -489,24 +473,20 @@ def set_rag_defaults( upsert_chat_settings_merge( { _DEFAULTS_KEY: { - "mode": new_mode, "embedding_model": new_embedder, } } ) return RagDefaults( - mode = new_mode, embedding_model = new_embedder, ) class ThreadRagSettings(BaseModel): - mode: KBMode = "text" embedding_model: str | None = None class UpdateThreadRagSettingsRequest(BaseModel): - mode: KBMode | None = None embedding_model: str | None = None # Reingest-only (not persisted); omit or None keeps captioning on. caption_images: bool | None = None @@ -524,7 +504,6 @@ def _load_thread_settings(thread_id: str) -> ThreadRagSettings: raw = {} fallback = _load_rag_defaults() return ThreadRagSettings( - mode = raw.get("mode") or fallback.mode, embedding_model = raw.get("embedding_model") or fallback.embedding_model, ) @@ -550,7 +529,6 @@ def set_thread_rag_settings( current_subject: str = Depends(get_current_subject), ) -> ThreadRagSettings: current = _load_thread_settings(thread_id) - new_mode = payload.mode or current.mode if payload.embedding_model is None: new_embedder = current.embedding_model elif payload.embedding_model.strip() == "": @@ -561,13 +539,11 @@ def set_thread_rag_settings( upsert_chat_settings_merge( { _thread_settings_key(thread_id): { - "mode": new_mode, "embedding_model": new_embedder, } } ) return ThreadRagSettings( - mode = new_mode, embedding_model = new_embedder, ) @@ -575,7 +551,6 @@ def set_thread_rag_settings( class ReingestKBRequest(BaseModel): """All fields optional — omitting one keeps the KB's current value.""" - mode: KBMode | None = None embedding_model: str | None = None # Not persisted on the KB; omit or None keeps captioning on for the rebuild. caption_images: bool | None = None @@ -590,7 +565,6 @@ def _reingest_scope( *, kb_id: str | None, thread_id: str | None, - mode: str, embedding_model: str, caption_images: bool = True, ) -> ReingestResponse: @@ -638,7 +612,6 @@ def _reingest_scope( kb_id = kb_id, thread_id = thread_id, embedding_model = embedding_model, - mode = mode, caption_images = caption_images, ) job_ids.append(upload.job_id) @@ -655,34 +628,24 @@ def reingest_knowledge_base( payload: ReingestKBRequest, current_subject: str = Depends(get_current_subject), ) -> ReingestResponse: - from utils.rag.config import resolve_embedder - kb_row = _kb_or_404(kb_id) - keys = kb_row.keys() if hasattr(kb_row, "keys") else () - current_mode = kb_row["mode"] if "mode" in keys else "text" current_embedder = kb_row["embedding_model"] - - new_mode = payload.mode or current_mode - - new_embedder = payload.embedding_model or ( - current_embedder if new_mode == current_mode else resolve_embedder(new_mode) - ) + new_embedder = payload.embedding_model or current_embedder with closing_connection() as conn: conn.execute( """ UPDATE rag_knowledge_bases - SET mode = ?, embedding_model = ? + SET embedding_model = ? WHERE id = ? """, - (new_mode, new_embedder, kb_id), + (new_embedder, kb_id), ) conn.commit() return _reingest_scope( kb_id = kb_id, thread_id = None, - mode = new_mode, embedding_model = new_embedder, caption_images = payload.caption_images is not False, ) @@ -702,7 +665,7 @@ def reingest_thread_documents( if payload is None: payload = UpdateThreadRagSettingsRequest() - if payload.mode is not None or payload.embedding_model is not None: + if payload.embedding_model is not None: settings = set_thread_rag_settings( thread_id, payload, @@ -711,11 +674,10 @@ def reingest_thread_documents( else: settings = _load_thread_settings(thread_id) - embedder = settings.embedding_model or resolve_embedder(settings.mode) + embedder = settings.embedding_model or resolve_embedder() return _reingest_scope( kb_id = None, thread_id = thread_id, - mode = settings.mode, embedding_model = embedder, caption_images = payload.caption_images is not False, ) @@ -752,9 +714,6 @@ 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 mode. - kb_keys = kb_row.keys() if hasattr(kb_row, "keys") else () - mode = kb_row["mode"] if "mode" in kb_keys else "text" return _start_ingestion( filename = filename, stored_path = stored_path, @@ -763,7 +722,6 @@ async def upload_kb_document( kb_id = kb_id, thread_id = None, embedding_model = kb_row["embedding_model"], - mode = mode, caption_images = caption_images, content_hash = content_hash, ) @@ -781,7 +739,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) + embedder = settings.embedding_model or resolve_embedder() return _start_ingestion( filename = filename, stored_path = stored_path, @@ -790,7 +748,6 @@ async def upload_thread_document( kb_id = None, thread_id = thread_id, embedding_model = embedder, - mode = settings.mode, caption_images = caption_images, content_hash = content_hash, ) @@ -826,28 +783,6 @@ def list_thread_documents( return DocumentListResponse(documents = [_row_to_document(r) for r in rows]) -@router.get("/images/{document_id}/{filename}") -def get_rag_image( - document_id: str, - filename: str, - current_subject: str = Depends(get_current_subject), -) -> FileResponse: - """Serve an extracted image; realpath-check against the uploads root.""" - document_for_subject_or_404(document_id, current_subject) - 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")) - candidate = rag_uploads_root() / "images" / document_id / filename - try: - real = Path(os.path.realpath(candidate)) - real.relative_to(root) - except (OSError, ValueError) as exc: - raise HTTPException(status_code = 404, detail = "Image not found") from exc - if not real.is_file(): - raise HTTPException(status_code = 404, detail = "Image not found") - return FileResponse(str(real)) - - @router.delete("/documents/{document_id}") def delete_document( document_id: str, @@ -1375,10 +1310,6 @@ def get_document_preview_target( chunk_kind: PreviewChunkKind = chunk_row["kind"] or "text" # type: ignore[assignment] image_url: str | None = None - if chunk_kind == "image" and chunk_row["image_path"]: - image_url = ( - f"/api/rag/images/{doc_row['id']}/" f"{Path(chunk_row['image_path']).name}" - ) return PreviewTargetResponse( **base, @@ -1600,10 +1531,6 @@ def search( continue kind = meta.get("kind", "text") or "text" image_url: str | None = None - if kind == "image" and meta.get("image_path"): - image_url = ( - f"/api/rag/images/{meta['document_id']}/{Path(meta['image_path']).name}" - ) out.append( SearchHit( chunk_id = hit.chunk_id, diff --git a/studio/backend/tests/test_rag_preview_routes.py b/studio/backend/tests/test_rag_preview_routes.py index 0774328296..514e3df150 100644 --- a/studio/backend/tests/test_rag_preview_routes.py +++ b/studio/backend/tests/test_rag_preview_routes.py @@ -570,57 +570,3 @@ class TestFileRoute: assert resp.headers.get("x-content-type-options") == "nosniff" cc = resp.headers.get("cache-control", "") assert "private" in cc - - -# ── /images tests ───────────────────────────────────────────────────── - - -class TestImageRoute: - def test_image_route_wrong_subject_returns_404(self, app, db_env, monkeypatch): - """Extracted images require the same document authorization as /file.""" - doc_id, kb_id = _uid(), _uid() - uploads = db_env / "rag" / "uploads" - images = uploads / "images" / doc_id - images.mkdir(parents = True, exist_ok = True) - image = images / "figure.png" - image.write_bytes(b"\x89PNG\r\n\x1a\n") - stored = uploads / "report.pdf" - stored.write_bytes(b"%PDF-1.4") - monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env)) - - with studio_db.get_connection() as conn: - _insert_kb(conn, kb_id, owner = "alice") - _insert_doc(conn, doc_id, kb_id, str(stored), "report.pdf") - - client = _make_client(app, "mallory") - try: - resp = client.get(f"/api/rag/images/{doc_id}/figure.png") - finally: - _clear_overrides(app) - - assert resp.status_code == 404 - - def test_image_route_authorized_subject_gets_image(self, app, db_env, monkeypatch): - """Authorized subject can still fetch an extracted image.""" - doc_id, kb_id = _uid(), _uid() - uploads = db_env / "rag" / "uploads" - images = uploads / "images" / doc_id - images.mkdir(parents = True, exist_ok = True) - image = images / "figure.png" - image.write_bytes(b"\x89PNG\r\n\x1a\n") - stored = uploads / "report.pdf" - stored.write_bytes(b"%PDF-1.4") - monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env)) - - with studio_db.get_connection() as conn: - _insert_kb(conn, kb_id, owner = "alice") - _insert_doc(conn, doc_id, kb_id, str(stored), "report.pdf") - - client = _make_client(app, "alice") - try: - resp = client.get(f"/api/rag/images/{doc_id}/figure.png") - finally: - _clear_overrides(app) - - assert resp.status_code == 200 - assert resp.content.startswith(b"\x89PNG") diff --git a/studio/backend/utils/rag/config.py b/studio/backend/utils/rag/config.py index b627bf3e59..448e5a5f74 100644 --- a/studio/backend/utils/rag/config.py +++ b/studio/backend/utils/rag/config.py @@ -31,26 +31,12 @@ RAG_EMBEDDING_MODEL: str = ( or "BAAI/bge-small-en-v1.5" ) -# 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, -# so a single 384-d text embedder handles retrieval. Multimodal adds image-vector -# rows on top via Qwen3-VL-Embedding-2B (2 B, 2048-d, no CLIP text cap — full -# 512-token chunks embed losslessly). -# -# 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[str, str] = { - "text": "BAAI/bge-small-en-v1.5", - "multimodal": "Qwen/Qwen3-VL-Embedding-2B", -} - - -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) +# A single text embedder handles retrieval. PDF figures are captioned at ingest +# (chat VLM or helper gemma-3n fallback) and spliced into the page markdown +# before chunking, so the 384-d text embedder covers figure content too. +def resolve_embedder() -> str: + """The configured RAG embedder.""" + return RAG_EMBEDDING_MODEL RAG_CHUNK_SIZE: int = _env_int("UNSLOTH_RAG_CHUNK_SIZE", 512) diff --git a/studio/frontend/src/__tests__/knowledge-bases-tab.test.tsx b/studio/frontend/src/__tests__/knowledge-bases-tab.test.tsx index 8a468bfd52..21575e3018 100644 --- a/studio/frontend/src/__tests__/knowledge-bases-tab.test.tsx +++ b/studio/frontend/src/__tests__/knowledge-bases-tab.test.tsx @@ -48,11 +48,6 @@ vi.mock("@/features/rag/components/thread-index-list", () => ({ React.createElement("div", { "data-testid": "thread-index-list" }), })); -vi.mock("@/features/rag/components/rag-defaults-section", () => ({ - RagDefaultsSection: () => - React.createElement("div", { "data-testid": "rag-defaults-section" }), -})); - beforeEach(() => { mockUsePreviewStore.__state.target = null; mockUsePreviewStore.__state.status = "idle"; diff --git a/studio/frontend/src/__tests__/preview-a11y.test.tsx b/studio/frontend/src/__tests__/preview-a11y.test.tsx index 1fa87afe9d..d72817238b 100644 --- a/studio/frontend/src/__tests__/preview-a11y.test.tsx +++ b/studio/frontend/src/__tests__/preview-a11y.test.tsx @@ -47,7 +47,6 @@ function target(overrides: Partial = {}): PreviewTarget { targetPage: 1, snippet: "safe extracted text", kind: "text", - imageUrl: null, sourcePageIndex: null, pageCharStart: null, pageCharEnd: null, diff --git a/studio/frontend/src/__tests__/preview-panel.test.tsx b/studio/frontend/src/__tests__/preview-panel.test.tsx index abc88a3b84..67d8a52575 100644 --- a/studio/frontend/src/__tests__/preview-panel.test.tsx +++ b/studio/frontend/src/__tests__/preview-panel.test.tsx @@ -163,7 +163,6 @@ function makeTarget(overrides: Partial = {}): PreviewTarget { targetPage: null, snippet: null, kind: null, - imageUrl: null, sourcePageIndex: null, pageCharStart: null, pageCharEnd: null, diff --git a/studio/frontend/src/__tests__/preview-pdf-smoke.test.tsx b/studio/frontend/src/__tests__/preview-pdf-smoke.test.tsx index 2382cd3804..2762ad8ceb 100644 --- a/studio/frontend/src/__tests__/preview-pdf-smoke.test.tsx +++ b/studio/frontend/src/__tests__/preview-pdf-smoke.test.tsx @@ -86,7 +86,6 @@ function target(overrides: Partial = {}): PreviewTarget { targetPage: 1, snippet: "target phrase appears here", kind: "text", - imageUrl: null, sourcePageIndex: 0, pageCharStart: 0, pageCharEnd: 13, diff --git a/studio/frontend/src/__tests__/preview-store.test.ts b/studio/frontend/src/__tests__/preview-store.test.ts index 346d807998..e4f83e36f4 100644 --- a/studio/frontend/src/__tests__/preview-store.test.ts +++ b/studio/frontend/src/__tests__/preview-store.test.ts @@ -100,7 +100,6 @@ function makeTarget(overrides: Partial = {}): PreviewTarget { targetPage: null, snippet: null, kind: null, - imageUrl: null, sourcePageIndex: null, pageCharStart: null, pageCharEnd: null, diff --git a/studio/frontend/src/__tests__/preview-target-locator.test.tsx b/studio/frontend/src/__tests__/preview-target-locator.test.tsx index 9713f86420..1aec595c50 100644 --- a/studio/frontend/src/__tests__/preview-target-locator.test.tsx +++ b/studio/frontend/src/__tests__/preview-target-locator.test.tsx @@ -23,7 +23,6 @@ function target(overrides: Partial = {}): PreviewTarget { targetPage: 2, snippet: "alpha\nhighlighted line\nomega", kind: "text", - imageUrl: null, sourcePageIndex: null, pageCharStart: null, pageCharEnd: null, diff --git a/studio/frontend/src/__tests__/rag-api.test.ts b/studio/frontend/src/__tests__/rag-api.test.ts index 8e3a4d0b73..25c0cd270b 100644 --- a/studio/frontend/src/__tests__/rag-api.test.ts +++ b/studio/frontend/src/__tests__/rag-api.test.ts @@ -66,7 +66,6 @@ function target(): PreviewTarget { targetPage: 1, snippet: "excerpt", kind: "text", - imageUrl: null, sourcePageIndex: 0, pageCharStart: 0, pageCharEnd: 7, diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-search-knowledge-base.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-search-knowledge-base.tsx index d6eac4ed14..8a731da82e 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-search-knowledge-base.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-search-knowledge-base.tsx @@ -3,14 +3,13 @@ "use client"; -import { authFetch } from "@/features/auth"; import { usePreviewStore } from "@/features/rag/stores/preview-store"; import { cn } from "@/lib/utils"; import { type ToolCallMessagePartComponent, useAuiState, } from "@assistant-ui/react"; -import { FileTextIcon, ImageIcon, LoaderIcon } from "lucide-react"; +import { FileTextIcon, LoaderIcon } from "lucide-react"; import { memo, useCallback, useEffect, useState } from "react"; import { ToolFallbackContent, @@ -31,7 +30,6 @@ export interface ParsedChunk { lineStart?: string; lineEnd?: string; kind?: string; - imageUrl?: string; text: string; /** Durable `rag_documents.id` from tool XML `document_id=`. Absent on * legacy tool output. */ @@ -79,7 +77,6 @@ export function parseChunks(raw: string): ParsedChunk[] { lineStart: attrs.line_start, lineEnd: attrs.line_end, kind: attrs.kind, - imageUrl: attrs.image_url, text, // Durable backend ids (legacy XML omits both → preview gated off). ...(attrs.document_id ? { documentId: attrs.document_id } : {}), @@ -93,59 +90,6 @@ export function parseChunks(raw: string): ParsedChunk[] { return out; } -/** Fetch a backend image via bearer-authed `authFetch`, expose it as a - * blob URL for ``. Revokes the object URL on unmount. */ -function useAuthedImageUrl(path: string | undefined): string | undefined { - const [url, setUrl] = useState(undefined); - useEffect(() => { - if (!path) { - setUrl(undefined); - return; - } - let cancelled = false; - let objectUrl: string | undefined; - authFetch(path) - .then((response) => { - if (!response.ok) { - throw new Error(`image fetch ${response.status}`); - } - return response.blob(); - }) - .then((blob) => { - if (cancelled) return; - objectUrl = URL.createObjectURL(blob); - setUrl(objectUrl); - }) - .catch(() => { - if (!cancelled) setUrl(undefined); - }); - return () => { - cancelled = true; - if (objectUrl) URL.revokeObjectURL(objectUrl); - }; - }, [path]); - return url; -} - -function ChunkImage({ url, alt }: { url: string; alt: string }) { - const blobUrl = useAuthedImageUrl(url); - if (!blobUrl) { - return ( -
- - Loading image… -
- ); - } - return ( - {alt} - ); -} - function ChunkCard({ chunk }: { chunk: ParsedChunk }) { const openPreview = usePreviewStore((s) => s.open); const meta: string[] = []; @@ -174,11 +118,7 @@ function ChunkCard({ chunk }: { chunk: ParsedChunk }) { [{chunk.id}] - {chunk.kind === "image" ? ( - - ) : ( - - )} + {chunk.source} @@ -213,9 +153,6 @@ function ChunkCard({ chunk }: { chunk: ParsedChunk }) { ) : null} - {chunk.kind === "image" && chunk.imageUrl ? ( - - ) : null} {chunk.text ? ( 
           {chunk.text}
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index 749a34cc74..3a24c5ee99 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -45,7 +45,6 @@ import {
   TooltipContent,
   TooltipTrigger,
 } from "@/components/ui/tooltip";
-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";
@@ -559,9 +558,6 @@ export function ChatSettingsPanel({
     }
   }, [ragSource.kind, activeThreadId, loadThreadSettings]);
 
-  const effectiveThreadMode: KBMode =
-    threadSettings?.mode ?? ragDefaults?.mode ?? "text";
-
   const aui = useAui();
   // Brand-new chat has no backend thread yet — initialize the local
   // assistant-ui thread to mint a remoteId so per-thread RAG settings
@@ -584,31 +580,6 @@ export function ChatSettingsPanel({
     }
   };
 
-  const applyThreadSettingChange = (patch: {
-    mode?: KBMode;
-  }) => {
-    void (async () => {
-      const threadId = await ensureThreadId();
-      if (!threadId) return;
-      if (threadDocs.length === 0) {
-        void updateThreadSettings(threadId, patch);
-        return;
-      }
-      const ok = window.confirm(
-        `Re-index ${threadDocs.length} document${threadDocs.length === 1 ? "" : "s"} ` +
-          `with the new settings? Existing chunks will be deleted and rebuilt.`,
-      );
-      if (ok) {
-        void reingestThread(threadId, {
-          ...patch,
-          caption_images: ragCaptionImages,
-        });
-      } else {
-        // User declined: refresh so the select snaps back.
-        void loadThreadSettings(threadId);
-      }
-    })();
-  };
   const [kbCreateOpen, setKbCreateOpen] = useState(false);
   const ragEnabled = ragSource.kind !== "off";
   const activeKbId = ragSource.kind === "kb" ? ragSource.kbId : null;
@@ -1414,7 +1385,6 @@ export function ChatSettingsPanel({
                       ) : null}
                       {knowledgeBases.map((kb) => {
                         const isActive = kb.id === activeKbId;
-                        const isMultimodal = kb.mode === "multimodal";
                         return (
                           
                             
                               {kb.name}
-                              {isMultimodal ? (
-                                
-                                  🖼️ MM
-                                
-                              ) : null}