ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio RAG: remove locator backfill (retroactive re-locator path)

Browse source 
Drops backfill_document_locators + the /documents/{id}/locators/backfill
route and its response model, the BackfillResult dataclass and the
backfill-only helpers (_scope_for_document, _update_vector_payloads), the
frontend backfillDocumentLocators client, and the backfill/migration
tests. Live preview-highlight locators (pdf_regions_for_chunks, computed
at ingest) are untouched.
This commit is contained in:
Roland Tannous 2026-06-02 18:53:15 +04:00
commit 45207c2bf1
6 changed files with 0 additions and 508 deletions
Download patch file Download diff file Expand all files Collapse all files

177
studio/backend/core/rag/locators.py
View file

@ -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),
)

54
studio/backend/routes/rag.py
View file

@ -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,

140
studio/backend/tests/test_rag_locator_backfill.py
View file

@ -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

86
studio/backend/tests/test_rag_locator_migration.py
View file

@ -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,
}

29
studio/frontend/src/__tests__/rag-api.test.ts
View file

@ -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", () => {

22
studio/frontend/src/features/rag/api/rag-api.ts
View file

@ -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<LocatorBackfillResult> {
const response = await authFetch(
`/api/rag/documents/${encodeURIComponent(documentId)}/locators/backfill`,
{ method: "POST" },
);
return parseJsonOrThrow<LocatorBackfillResult>(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