diff --git a/studio/backend/core/rag/locators.py b/studio/backend/core/rag/locators.py index 0ca02699da..2e2786b8e7 100644 --- a/studio/backend/core/rag/locators.py +++ b/studio/backend/core/rag/locators.py @@ -31,19 +31,6 @@ class LocatorMatch: line_end: int -@dataclass(frozen = True) -class BackfillResult: - document_id: str - total_chunks: int - matched: int - already_located: int - ambiguous: int - missing: int - skipped: int - regions_matched: int - pages_refreshed: int - - def _line_bounds(text: str, start: int, end: int) -> tuple[int, int]: line_start = text.count("\n", 0, start) + 1 line_end = text.count("\n", 0, max(start, end - 1)) + 1 @@ -334,167 +321,3 @@ def pdf_regions_for_chunks( return regions finally: doc.close() - - -def _scope_for_document(kb_id: str | None, thread_id: str | None) -> str | None: - if kb_id: - return kb_scope(kb_id) - if thread_id: - return thread_scope(thread_id) - return None - - -def _update_vector_payloads( - scope: str | None, updates: dict[str, dict[str, Any]] -) -> None: - if not scope or not updates: - return - try: - vector_store.update_chunk_payload_fields(scope, updates) - except Exception as exc: - logger.warning( - "RAG locator backfill: vector payload update failed", - error = str(exc), - ) - - -def backfill_document_locators(document_id: str, stored_path: Path) -> BackfillResult: - parsed = parse(stored_path, want_images = False) - pages = parsed.pages - _replace_document_pages(document_id, pages) - - with closing_connection() as conn: - doc_row = conn.execute( - "SELECT kb_id, thread_id FROM rag_documents WHERE id = ?", - (document_id,), - ).fetchone() - if doc_row is None: - return BackfillResult(document_id, 0, 0, 0, 0, 0, 0, 0, len(pages)) - - rows = conn.execute( - """ - SELECT id, text, kind, page_number, source_page_index, - page_char_start, page_char_end, line_start, line_end, - pdf_regions_json - FROM rag_chunks - WHERE document_id = ? - ORDER BY chunk_index ASC - """, - (document_id,), - ).fetchall() - - scope = _scope_for_document(doc_row["kb_id"], doc_row["thread_id"]) - total = len(rows) - matched = 0 - already_located = 0 - ambiguous = 0 - missing = 0 - skipped = 0 - regions_matched = 0 - sql_updates: list[tuple[Any, ...]] = [] - vector_updates: dict[str, dict[str, Any]] = {} - - for row in rows: - kind = row["kind"] or "text" - text = row["text"] or "" - if kind not in ("text", "caption") or not text.strip(): - skipped += 1 - continue - - existing_complete = ( - row["source_page_index"] is not None - and row["page_char_start"] is not None - and row["page_char_end"] is not None - and row["line_start"] is not None - and row["line_end"] is not None - ) - - match: LocatorMatch | None - status: str - if existing_complete: - already_located += 1 - page_index = int(row["source_page_index"]) - if 0 <= page_index < len(pages): - match = LocatorMatch( - page_index = page_index, - page_number = row["page_number"], - start = int(row["page_char_start"]), - end = int(row["page_char_end"]), - line_start = int(row["line_start"]), - line_end = int(row["line_end"]), - ) - else: - match = None - status = "already_located" - else: - match, status = _locate_unique(text, pages) - if status == "matched" and match is not None: - matched += 1 - elif status == "ambiguous": - ambiguous += 1 - continue - else: - missing += 1 - continue - - if match is None: - continue - - regions = pdf_regions_for_match(stored_path, pages, match) - regions_json = json.dumps(regions, separators = (",", ":")) if regions else None - if regions: - regions_matched += 1 - - if status == "matched" or (regions and not row["pdf_regions_json"]): - sql_updates.append( - ( - match.page_number, - match.page_index, - match.start, - match.end, - match.line_start, - match.line_end, - regions_json, - row["id"], - ) - ) - vector_updates[row["id"]] = { - "page_number": match.page_number, - "source_page_index": match.page_index, - "page_char_start": match.start, - "page_char_end": match.end, - "line_start": match.line_start, - "line_end": match.line_end, - "pdf_regions": regions, - } - - if sql_updates: - with closing_connection() as conn: - conn.executemany( - """ - UPDATE rag_chunks - SET page_number = COALESCE(page_number, ?), - source_page_index = ?, - page_char_start = ?, - page_char_end = ?, - line_start = ?, - line_end = ?, - pdf_regions_json = COALESCE(?, pdf_regions_json) - WHERE id = ? - """, - sql_updates, - ) - conn.commit() - _update_vector_payloads(scope, vector_updates) - - return BackfillResult( - document_id = document_id, - total_chunks = total, - matched = matched, - already_located = already_located, - ambiguous = ambiguous, - missing = missing, - skipped = skipped, - regions_matched = regions_matched, - pages_refreshed = len(pages), - ) diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index d4b9858a1c..ebbc327346 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -43,7 +43,6 @@ async def _sse_auth( from core.rag import embeddings, ingestion, retrieval, vector_store from core.rag.authorization import document_for_subject_or_404 -from core.rag.locators import backfill_document_locators from core.rag.vector_store import kb_scope, thread_scope from loggers import get_logger from storage.studio_db import ( @@ -976,18 +975,6 @@ class PreviewFileUrlResponse(BaseModel): expiresAt: int -class LocatorBackfillResponse(BaseModel): - documentId: str - totalChunks: int - matched: int - alreadyLocated: int - ambiguous: int - missing: int - skipped: int - regionsMatched: int - pagesRefreshed: int - - # Extension allowlist for inline rendering / disposition. Unlisted ext collapses # to ("application/octet-stream", attachment, "unknown"). .html / .htm serve as # text/plain attachment (decisions Q7 + Risk #3) so uploaded HTML can't execute in the app origin. @@ -1332,47 +1319,6 @@ def get_document_preview_target( ) -@router.post( - "/documents/{document_id}/locators/backfill", - response_model = LocatorBackfillResponse, -) -def backfill_document_locators_route( - document_id: str, - current_subject: str = Depends(get_current_subject), -) -> LocatorBackfillResponse: - """In-place locator backfill for existing citations. - - This preserves ``document_id`` and ``chunk_id``. Chunks are updated only - when their text has one unambiguous match in the parsed document text; - duplicate or missing matches remain null. - """ - doc_row = document_for_subject_or_404(document_id, current_subject) - resolved = _resolve_document_file_or_404(doc_row, document_id) - try: - result = backfill_document_locators(document_id, resolved) - except Exception as exc: - logger.warning( - "RAG locator backfill failed", - document_id = document_id, - error = str(exc), - ) - raise HTTPException( - status_code = 400, - detail = "Document locators could not be backfilled", - ) from exc - return LocatorBackfillResponse( - documentId = result.document_id, - totalChunks = result.total_chunks, - matched = result.matched, - alreadyLocated = result.already_located, - ambiguous = result.ambiguous, - missing = result.missing, - skipped = result.skipped, - regionsMatched = result.regions_matched, - pagesRefreshed = result.pages_refreshed, - ) - - @router.get( "/documents/{document_id}/file-url", response_model = PreviewFileUrlResponse, diff --git a/studio/backend/tests/test_rag_locator_backfill.py b/studio/backend/tests/test_rag_locator_backfill.py deleted file mode 100644 index 2503ab1d53..0000000000 --- a/studio/backend/tests/test_rag_locator_backfill.py +++ /dev/null @@ -1,140 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -from __future__ import annotations - -import uuid -from pathlib import Path - -import pytest -from fastapi.testclient import TestClient - -import storage.studio_db as studio_db -from auth.authentication import get_current_subject - - -@pytest.fixture(scope = "module") -def app(): - import sys - - backend_dir = str(Path(__file__).resolve().parent.parent) - if backend_dir not in sys.path: - sys.path.insert(0, backend_dir) - from main import app as _app - - return _app - - -@pytest.fixture -def db_env(tmp_path, monkeypatch): - monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) - monkeypatch.setattr(studio_db, "_schema_ready", False) - return tmp_path - - -def _uid() -> str: - return str(uuid.uuid4()) - - -def _make_client(app, subject: str = "alice"): - app.dependency_overrides[get_current_subject] = lambda: subject - return TestClient(app, raise_server_exceptions = True) - - -def _clear_overrides(app): - app.dependency_overrides.clear() - - -def _insert_kb(conn, kb_id: str, owner: str = "alice") -> None: - conn.execute( - "INSERT INTO rag_knowledge_bases " - "(id, name, embedding_model, owner_user_id, created_at) " - "VALUES (?, ?, ?, ?, ?)", - (kb_id, f"KB-{kb_id[:6]}", "bge-small", owner, 1_700_000_000), - ) - - -def _insert_doc(conn, doc_id: str, kb_id: str, stored_path: str, filename: str) -> None: - conn.execute( - "INSERT INTO rag_documents " - "(id, kb_id, thread_id, filename, content_type, stored_path, status, " - "num_chunks, byte_size, created_at) " - "VALUES (?, ?, NULL, ?, 'text/plain', ?, 'completed', 1, 64, ?)", - (doc_id, kb_id, filename, stored_path, 1_700_000_000), - ) - - -def _insert_chunk(conn, chunk_id: str, doc_id: str, text: str) -> None: - conn.execute( - "INSERT INTO rag_chunks " - "(id, document_id, chunk_index, text, token_count, page_number) " - "VALUES (?, ?, 0, ?, 5, NULL)", - (chunk_id, doc_id, text), - ) - - -def test_backfill_preserves_ids_and_updates_unique_locator(app, db_env, monkeypatch): - doc_id, kb_id, chunk_id = _uid(), _uid(), _uid() - stored = db_env / "rag" / "uploads" / "paper.txt" - stored.parent.mkdir(parents = True, exist_ok = True) - stored.write_text("Intro line\nUnique quote here.\nEnd.", encoding = "utf-8") - monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env)) - - with studio_db.get_connection() as conn: - _insert_kb(conn, kb_id) - _insert_doc(conn, doc_id, kb_id, str(stored), "paper.txt") - _insert_chunk(conn, chunk_id, doc_id, "Unique quote here.") - - client = _make_client(app, "alice") - try: - resp = client.post(f"/api/rag/documents/{doc_id}/locators/backfill") - target_resp = client.get( - f"/api/rag/documents/{doc_id}/preview-target?chunk_id={chunk_id}" - ) - finally: - _clear_overrides(app) - - assert resp.status_code == 200 - body = resp.json() - assert body["documentId"] == doc_id - assert body["matched"] == 1 - assert body["ambiguous"] == 0 - - target = target_resp.json() - assert target["documentId"] == doc_id - assert target["chunkId"] == chunk_id - assert target["sourcePageIndex"] == 0 - assert target["lineStart"] == 2 - assert target["pageCharStart"] in (len("Intro line\n"), len("Intro line\r\n")) - - -def test_backfill_leaves_ambiguous_matches_null(app, db_env, monkeypatch): - doc_id, kb_id, chunk_id = _uid(), _uid(), _uid() - stored = db_env / "rag" / "uploads" / "paper.txt" - stored.parent.mkdir(parents = True, exist_ok = True) - stored.write_text("Repeat me.\nOther text.\nRepeat me.", encoding = "utf-8") - monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env)) - - with studio_db.get_connection() as conn: - _insert_kb(conn, kb_id) - _insert_doc(conn, doc_id, kb_id, str(stored), "paper.txt") - _insert_chunk(conn, chunk_id, doc_id, "Repeat me.") - - client = _make_client(app, "alice") - try: - resp = client.post(f"/api/rag/documents/{doc_id}/locators/backfill") - target_resp = client.get( - f"/api/rag/documents/{doc_id}/preview-target?chunk_id={chunk_id}" - ) - finally: - _clear_overrides(app) - - assert resp.status_code == 200 - body = resp.json() - assert body["matched"] == 0 - assert body["ambiguous"] == 1 - - target = target_resp.json() - assert target["sourcePageIndex"] is None - assert target["pageCharStart"] is None - assert target["lineStart"] is None diff --git a/studio/backend/tests/test_rag_locator_migration.py b/studio/backend/tests/test_rag_locator_migration.py deleted file mode 100644 index 627af23d05..0000000000 --- a/studio/backend/tests/test_rag_locator_migration.py +++ /dev/null @@ -1,86 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -from __future__ import annotations - -import uuid - -import storage.studio_db as studio_db - - -def _uid() -> str: - return str(uuid.uuid4()) - - -def test_locator_schema_is_additive_and_nullable(tmp_path, monkeypatch): - monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) - monkeypatch.setattr(studio_db, "_schema_ready", False) - - with studio_db.get_connection() as conn: - chunk_cols = { - row["name"] for row in conn.execute("PRAGMA table_info(rag_chunks)") - } - assert { - "source_page_index", - "page_char_start", - "page_char_end", - "line_start", - "line_end", - }.issubset(chunk_cols) - - page_cols = { - row["name"] for row in conn.execute("PRAGMA table_info(rag_document_pages)") - } - assert { - "document_id", - "page_index", - "page_number", - "text", - "char_count", - "line_count", - }.issubset(page_cols) - - kb_id = _uid() - doc_id = _uid() - chunk_id = _uid() - conn.execute( - """ - INSERT INTO rag_knowledge_bases - (id, name, embedding_model, owner_user_id, created_at) - VALUES (?, ?, ?, ?, ?) - """, - (kb_id, "KB", "embedder", "alice", 1_700_000_000), - ) - conn.execute( - """ - INSERT INTO rag_documents - (id, kb_id, thread_id, filename, content_type, stored_path, status, - num_chunks, byte_size, created_at) - VALUES (?, ?, NULL, ?, ?, ?, 'completed', 1, 10, ?) - """, - (doc_id, kb_id, "old.pdf", "application/pdf", "old.pdf", 1_700_000_001), - ) - conn.execute( - """ - INSERT INTO rag_chunks - (id, document_id, chunk_index, text, token_count, page_number) - VALUES (?, ?, 0, ?, 3, 1) - """, - (chunk_id, doc_id, "legacy chunk"), - ) - - row = conn.execute( - """ - SELECT source_page_index, page_char_start, page_char_end, - line_start, line_end - FROM rag_chunks WHERE id = ? - """, - (chunk_id,), - ).fetchone() - assert dict(row) == { - "source_page_index": None, - "page_char_start": None, - "page_char_end": None, - "line_start": None, - "line_end": None, - } diff --git a/studio/frontend/src/__tests__/rag-api.test.ts b/studio/frontend/src/__tests__/rag-api.test.ts index 25c0cd270b..8335147829 100644 --- a/studio/frontend/src/__tests__/rag-api.test.ts +++ b/studio/frontend/src/__tests__/rag-api.test.ts @@ -45,7 +45,6 @@ vi.mock("event-source-polyfill", () => ({ })); import { - backfillDocumentLocators, fetchPreviewFileUrl, fetchPreviewTarget, subscribeToJobEvents, @@ -122,34 +121,6 @@ describe("RAG API preview target", () => { expect(result.url).not.toContain("Authorization"); }); - it("posts the explicit locator backfill action", async () => { - mockAuthFetch.mockResolvedValue( - new Response( - JSON.stringify({ - documentId: "doc-abc", - totalChunks: 1, - matched: 1, - alreadyLocated: 0, - ambiguous: 0, - missing: 0, - skipped: 0, - regionsMatched: 0, - pagesRefreshed: 1, - }), - { - status: 200, - headers: { "Content-Type": "application/json" }, - }, - ), - ); - - await backfillDocumentLocators("doc id/with?slash"); - - expect(mockAuthFetch).toHaveBeenCalledWith( - "/api/rag/documents/doc%20id%2Fwith%3Fslash/locators/backfill", - { method: "POST" }, - ); - }); }); describe("RAG API job events", () => { diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts index f5bcdb4d11..2b4b37211d 100644 --- a/studio/frontend/src/features/rag/api/rag-api.ts +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -102,18 +102,6 @@ export interface PreviewFileUrl { expiresAt: number; } -export interface LocatorBackfillResult { - documentId: string; - totalChunks: number; - matched: number; - alreadyLocated: number; - ambiguous: number; - missing: number; - skipped: number; - regionsMatched: number; - pagesRefreshed: number; -} - export interface SearchRequest { query: string; kb_id?: string; @@ -490,16 +478,6 @@ export async function fetchPreviewFileUrl( }; } -export async function backfillDocumentLocators( - documentId: string, -): Promise { - const response = await authFetch( - `/api/rag/documents/${encodeURIComponent(documentId)}/locators/backfill`, - { method: "POST" }, - ); - return parseJsonOrThrow(response); -} - /** Download the original uploaded file as a Blob via `authFetch` (bearer * token in the Authorization header, never a query string). The caller * (preview-store) creates and revokes the object URL so blob lifecycle