Studio: WIP — RAG preview UI, locator/auth refactor, tests, fixtures (pre-merge snapshot)
Snapshot taken before fast-forwarding feature/rag to origin and merging main. Bundles in-flight work so the merge has a clean tree: Frontend - PDF preview panel (preview-panel, preview-pdf-view, preview-text-view, preview-unavailable) with lazy-rendered page thumbnail rail - Resizable preview slot via useResizablePanelWidth hook (drag handle, localStorage persistence, viewport clamping) - Neutral scrollbar + Source Excerpt card restyle (no brand-coloured rail) - Preview-store + chat-adapter / rag-api / kb-detail wiring - Frontend test harness (vitest.config, setupTests, biome update) and the paired __tests__ suites for preview, sources, document-row, chat-adapter, rag-api, knowledge-bases-tab, search-knowledge-base-tool-ui Backend - RAG locator + authorization modules with chunking / retrieval / tool / vector_store / studio_db updates - Paired test_rag_* suites (authorization, locators, locator_backfill, locator_migration, preview_routes, preview_target_locators, source_identity) Other - tests/fixtures/rag-preview for preview route fixtures (sample.pdf, sample.txt, make_fixture_pdf.py) - .gitignore + package(-lock).json adjustments for the new test runner Will be squashed/reworked via interactive rebase after main is merged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
dd3ee02648
commit
27b0a50a84
51 changed files with 9984 additions and 1153 deletions
3
.gitignore
vendored
3
.gitignore
vendored
|
|
@ -235,3 +235,6 @@ package-lock.json
|
|||
!studio/backend/core/data_recipe/oxc-validator/package-lock.json
|
||||
!studio/package-lock.json
|
||||
llama.cpp/
|
||||
/.Codex
|
||||
/.gemini
|
||||
/.antigravitycli
|
||||
|
|
|
|||
121
studio/backend/core/rag/authorization.py
Normal file
121
studio/backend/core/rag/authorization.py
Normal file
|
|
@ -0,0 +1,121 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Subject-scoped authorization for RAG document preview routes.
|
||||
|
||||
Used by `/api/rag/documents/{document_id}/file` and
|
||||
`/api/rag/documents/{document_id}/preview-target` to enforce that the
|
||||
current authenticated subject is allowed to see a given document and
|
||||
chunk. Existence and authorization failures collapse to a single 404
|
||||
so the API does not leak document IDs to a non-owner.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from storage.studio_db import get_connection
|
||||
|
||||
_NOT_FOUND_DETAIL = "Document not found"
|
||||
|
||||
|
||||
def document_for_subject_or_404(
|
||||
document_id: str,
|
||||
current_subject: str,
|
||||
) -> sqlite3.Row:
|
||||
"""Return the `rag_documents` row if `current_subject` may access it.
|
||||
|
||||
Authorization rules:
|
||||
|
||||
- KB documents: the document's KB must have
|
||||
`rag_knowledge_bases.owner_user_id == current_subject`. A KB with a
|
||||
NULL owner is not accessible through this helper (legacy pre-auth
|
||||
rows must be migrated or accessed via admin tooling).
|
||||
|
||||
- Thread documents: thread-scoped RAG documents are gated by an
|
||||
explicit single-user invariant for Studio's current release. The
|
||||
`chat_threads` table does not yet carry an `owner_user_id` column,
|
||||
so we cannot bind a thread to a specific subject in the schema.
|
||||
The helper still requires (a) an authenticated subject (enforced
|
||||
by the route's `Depends(get_current_subject)`) and (b) that the
|
||||
referenced thread actually exists in `chat_threads`. A missing
|
||||
thread row collapses to 404 so a non-existent thread cannot
|
||||
silently grant access through a dangling `thread_id`.
|
||||
# TODO(thread-owner): once `chat_threads.owner_user_id` exists,
|
||||
# join through it the same way KB documents do and drop the
|
||||
# single-user invariant. Update the test
|
||||
# `tests/test_rag_authorization.py::test_thread_doc_other_user_404`
|
||||
# to assert per-user isolation rather than thread existence.
|
||||
|
||||
Both not-found and not-authorized raise `HTTPException(404)` with the
|
||||
same detail string. Callers must NOT distinguish the two cases in
|
||||
their response, to avoid leaking document existence to a non-owner.
|
||||
|
||||
Returns the document row so the caller can read `stored_path`,
|
||||
`filename`, `content_type`, etc. without re-querying.
|
||||
"""
|
||||
if not document_id or not current_subject:
|
||||
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
|
||||
|
||||
with get_connection() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT * FROM rag_documents WHERE id = ?",
|
||||
(document_id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
|
||||
|
||||
kb_id = row["kb_id"]
|
||||
thread_id = row["thread_id"]
|
||||
|
||||
if kb_id is not None:
|
||||
owner_row = conn.execute(
|
||||
"SELECT owner_user_id FROM rag_knowledge_bases WHERE id = ?",
|
||||
(kb_id,),
|
||||
).fetchone()
|
||||
if owner_row is None:
|
||||
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
|
||||
owner = owner_row["owner_user_id"]
|
||||
if owner is None or owner != current_subject:
|
||||
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
|
||||
return row
|
||||
|
||||
if thread_id is not None:
|
||||
# Single-user invariant (see TODO above). We require the
|
||||
# thread row to exist; an unknown thread_id is treated as
|
||||
# not-found, not as silent grant.
|
||||
thread_row = conn.execute(
|
||||
"SELECT id FROM chat_threads WHERE id = ?",
|
||||
(thread_id,),
|
||||
).fetchone()
|
||||
if thread_row is None:
|
||||
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
|
||||
return row
|
||||
|
||||
# Documents must belong to either a KB or a thread (DB CHECK
|
||||
# constraint enforces XOR on insert); a row that satisfies
|
||||
# neither is corrupt — treat as 404.
|
||||
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
|
||||
|
||||
|
||||
def chunk_belongs_to_document(chunk_id: str, document_id: str) -> bool:
|
||||
"""True iff `chunk_id` exists in `rag_chunks` for `document_id`.
|
||||
|
||||
Used by `/preview-target?chunk_id=...` after the caller has
|
||||
already established subject authorization for `document_id`. Does
|
||||
NOT perform authorization itself: callers MUST call
|
||||
`document_for_subject_or_404(document_id, ...)` first, otherwise a
|
||||
valid `chunk_id` from another subject's document would leak via a
|
||||
`True` return.
|
||||
"""
|
||||
if not chunk_id or not document_id:
|
||||
return False
|
||||
|
||||
with get_connection() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT 1 FROM rag_chunks WHERE id = ? AND document_id = ?",
|
||||
(chunk_id, document_id),
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
|
@ -48,6 +48,11 @@ class Chunk:
|
|||
text: str
|
||||
token_count: int
|
||||
page_number: int | None = None
|
||||
source_page_index: int | None = None
|
||||
page_char_start: int | None = None
|
||||
page_char_end: int | None = None
|
||||
line_start: int | None = None
|
||||
line_end: int | None = None
|
||||
|
||||
|
||||
TokenCounter = Callable[[str], int]
|
||||
|
|
@ -132,6 +137,28 @@ def _merge(
|
|||
return [c.strip() for c in chunks if c.strip()]
|
||||
|
||||
|
||||
def _line_bounds(text: str, start: int, end: int) -> tuple[int, int]:
|
||||
"""Return 1-based inclusive line numbers for a page-local span."""
|
||||
line_start = text.count("\n", 0, start) + 1
|
||||
line_end = text.count("\n", 0, max(start, end - 1)) + 1
|
||||
return line_start, line_end
|
||||
|
||||
|
||||
def _locate_piece(
|
||||
page_text: str,
|
||||
piece: str,
|
||||
search_cursor: int,
|
||||
) -> tuple[int | None, int | None, int | None, int | None, int]:
|
||||
idx = page_text.find(piece, search_cursor)
|
||||
if idx < 0:
|
||||
idx = page_text.find(piece)
|
||||
if idx < 0:
|
||||
return None, None, None, None, search_cursor
|
||||
end = idx + len(piece)
|
||||
line_start, line_end = _line_bounds(page_text, idx, end)
|
||||
return idx, end, line_start, line_end, idx + 1
|
||||
|
||||
|
||||
# Markdown headings first so layout-aware parser output splits at sections.
|
||||
DEFAULT_SEPARATORS: tuple[str, ...] = (
|
||||
"\n# ",
|
||||
|
|
@ -157,16 +184,27 @@ def chunk_pages(
|
|||
"""Split pages independently so page_number stays attached to chunks."""
|
||||
count = token_counter or _char_token_estimate
|
||||
out: list[Chunk] = []
|
||||
for page in pages:
|
||||
for page_index, page in enumerate(pages):
|
||||
search_cursor = 0
|
||||
for segment in _split_at_figure_boundaries(page.text):
|
||||
atomic = _atomic_split(segment, separators, max_tokens, count)
|
||||
merged = _merge(atomic, max_tokens, overlap_tokens, count)
|
||||
for piece in merged:
|
||||
start, end, line_start, line_end, search_cursor = _locate_piece(
|
||||
page.text,
|
||||
piece,
|
||||
search_cursor,
|
||||
)
|
||||
out.append(
|
||||
Chunk(
|
||||
text = piece,
|
||||
token_count = count(piece),
|
||||
page_number = page.page_number,
|
||||
source_page_index = page_index,
|
||||
page_char_start = start,
|
||||
page_char_end = end,
|
||||
line_start = line_start,
|
||||
line_end = line_end,
|
||||
)
|
||||
)
|
||||
return out
|
||||
|
|
@ -192,13 +230,13 @@ def chunk_pages_with_spans(
|
|||
count = token_counter or _char_token_estimate
|
||||
|
||||
parts: list[str] = []
|
||||
page_ranges: list[tuple[int, int, int | None]] = []
|
||||
page_ranges: list[tuple[int, int, int, int | None]] = []
|
||||
cursor = 0
|
||||
for index, page in enumerate(pages):
|
||||
parts.append(page.text)
|
||||
start = cursor
|
||||
end = cursor + len(page.text)
|
||||
page_ranges.append((start, end, page.page_number))
|
||||
page_ranges.append((start, end, index, page.page_number))
|
||||
cursor = end
|
||||
if index < len(pages) - 1:
|
||||
cursor += len(_PAGE_SEPARATOR)
|
||||
|
|
@ -223,12 +261,34 @@ def chunk_pages_with_spans(
|
|||
if idx < 0:
|
||||
continue
|
||||
end_idx = idx + len(text)
|
||||
page_number = _page_for_span(idx, end_idx, page_ranges)
|
||||
page_locator = _page_for_span(idx, end_idx, page_ranges)
|
||||
source_page_index: int | None = None
|
||||
page_number: int | None = None
|
||||
page_char_start: int | None = None
|
||||
page_char_end: int | None = None
|
||||
line_start: int | None = None
|
||||
line_end: int | None = None
|
||||
if page_locator is not None:
|
||||
page_start, page_end, page_idx, page_no = page_locator
|
||||
source_page_index = page_idx
|
||||
page_number = page_no
|
||||
page_char_start = max(0, idx - page_start)
|
||||
page_char_end = min(page_end, end_idx) - page_start
|
||||
line_start, line_end = _line_bounds(
|
||||
pages[page_idx].text,
|
||||
page_char_start,
|
||||
page_char_end,
|
||||
)
|
||||
chunks.append(
|
||||
Chunk(
|
||||
text = text,
|
||||
token_count = count(text),
|
||||
page_number = page_number,
|
||||
source_page_index = source_page_index,
|
||||
page_char_start = page_char_start,
|
||||
page_char_end = page_char_end,
|
||||
line_start = line_start,
|
||||
line_end = line_end,
|
||||
)
|
||||
)
|
||||
char_spans.append((idx, end_idx))
|
||||
|
|
@ -241,9 +301,9 @@ def chunk_pages_with_spans(
|
|||
def _page_for_span(
|
||||
start: int,
|
||||
end: int,
|
||||
page_ranges: list[tuple[int, int, int | None]],
|
||||
) -> int | None:
|
||||
for ps, pe, pn in page_ranges:
|
||||
page_ranges: list[tuple[int, int, int, int | None]],
|
||||
) -> tuple[int, int, int, int | None] | None:
|
||||
for ps, pe, page_index, page_number in page_ranges:
|
||||
if start < pe and end > ps:
|
||||
return pn
|
||||
return ps, pe, page_index, page_number
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ vectors, and rebuilds BM25 on completion. Only the parent opens rag.db.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import multiprocessing as mp
|
||||
import queue as queue_module
|
||||
import sqlite3
|
||||
|
|
@ -97,6 +98,22 @@ def _subprocess_worker(
|
|||
)
|
||||
pages = inline_image_captions(pages, parsed.images, captions)
|
||||
|
||||
out_queue.put(
|
||||
{
|
||||
"type": "document_pages",
|
||||
"pages": [
|
||||
{
|
||||
"page_index": index,
|
||||
"page_number": page.page_number,
|
||||
"text": page.text,
|
||||
"char_count": len(page.text),
|
||||
"line_count": len(page.text.splitlines()),
|
||||
}
|
||||
for index, page in enumerate(pages)
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
out_queue.put({"type": "progress", "stage": "load_model", "progress": 0.1})
|
||||
from core.rag.embeddings import (
|
||||
get_embedder,
|
||||
|
|
@ -112,6 +129,7 @@ def _subprocess_worker(
|
|||
if chunking_strategy == "late":
|
||||
_run_late_chunking(
|
||||
pages = pages,
|
||||
stored_path = Path(stored_path),
|
||||
chunk_size = chunk_size,
|
||||
overlap = overlap,
|
||||
counter = counter,
|
||||
|
|
@ -123,6 +141,7 @@ def _subprocess_worker(
|
|||
|
||||
text_count = _run_standard_chunking(
|
||||
pages = pages,
|
||||
stored_path = Path(stored_path),
|
||||
chunk_size = chunk_size,
|
||||
overlap = overlap,
|
||||
counter = counter,
|
||||
|
|
@ -151,6 +170,7 @@ def _subprocess_worker(
|
|||
def _run_standard_chunking(
|
||||
*,
|
||||
pages,
|
||||
stored_path,
|
||||
chunk_size,
|
||||
overlap,
|
||||
counter,
|
||||
|
|
@ -172,6 +192,9 @@ def _run_standard_chunking(
|
|||
if send_complete:
|
||||
out_queue.put({"type": "error", "error": "chunker produced no chunks"})
|
||||
return 0
|
||||
from core.rag.locators import pdf_regions_for_chunks
|
||||
|
||||
pdf_regions = pdf_regions_for_chunks(stored_path, pages, chunks)
|
||||
|
||||
total = len(chunks)
|
||||
for i in range(0, total, batch_size):
|
||||
|
|
@ -192,9 +215,15 @@ def _run_standard_chunking(
|
|||
"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[i + offset],
|
||||
"kind": "text",
|
||||
}
|
||||
for c in batch
|
||||
for offset, c in enumerate(batch)
|
||||
],
|
||||
"vectors": vectors.tolist(),
|
||||
}
|
||||
|
|
@ -315,6 +344,7 @@ def _stream_image_chunks(
|
|||
def _run_late_chunking(
|
||||
*,
|
||||
pages,
|
||||
stored_path,
|
||||
chunk_size,
|
||||
overlap,
|
||||
counter,
|
||||
|
|
@ -324,6 +354,7 @@ def _run_late_chunking(
|
|||
) -> 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(
|
||||
|
|
@ -343,6 +374,7 @@ def _run_late_chunking(
|
|||
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(
|
||||
|
|
@ -354,8 +386,15 @@ def _run_late_chunking(
|
|||
"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 c in chunks
|
||||
for index, c in enumerate(chunks)
|
||||
],
|
||||
"vectors": [v.tolist() for v in vectors],
|
||||
}
|
||||
|
|
@ -469,6 +508,14 @@ def _insert_chunks_and_collect_for_bm25(
|
|||
meta["page_number"],
|
||||
kind,
|
||||
image_path,
|
||||
meta.get("source_page_index"),
|
||||
meta.get("page_char_start"),
|
||||
meta.get("page_char_end"),
|
||||
meta.get("line_start"),
|
||||
meta.get("line_end"),
|
||||
json.dumps(meta.get("pdf_regions") or [], separators = (",", ":"))
|
||||
if meta.get("pdf_regions")
|
||||
else None,
|
||||
)
|
||||
)
|
||||
points.append(
|
||||
|
|
@ -482,6 +529,12 @@ def _insert_chunks_and_collect_for_bm25(
|
|||
"page_number": meta["page_number"],
|
||||
"kind": kind,
|
||||
"image_path": image_path,
|
||||
"source_page_index": meta.get("source_page_index"),
|
||||
"page_char_start": meta.get("page_char_start"),
|
||||
"page_char_end": meta.get("page_char_end"),
|
||||
"line_start": meta.get("line_start"),
|
||||
"line_end": meta.get("line_end"),
|
||||
"pdf_regions": meta.get("pdf_regions") or [],
|
||||
},
|
||||
}
|
||||
)
|
||||
|
|
@ -492,8 +545,9 @@ def _insert_chunks_and_collect_for_bm25(
|
|||
"""
|
||||
INSERT INTO rag_chunks
|
||||
(id, document_id, chunk_index, text, token_count, page_number,
|
||||
kind, image_path)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
kind, image_path, source_page_index, page_char_start,
|
||||
page_char_end, line_start, line_end, pdf_regions_json)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
rows,
|
||||
)
|
||||
|
|
@ -515,6 +569,41 @@ def _insert_chunks_and_collect_for_bm25(
|
|||
return bm25_rows
|
||||
|
||||
|
||||
def _replace_document_pages(document_id: str, pages: list[dict]) -> None:
|
||||
now = int(time.time())
|
||||
rows = [
|
||||
(
|
||||
document_id,
|
||||
int(page["page_index"]),
|
||||
page.get("page_number"),
|
||||
page.get("text") or "",
|
||||
int(page.get("char_count", len(page.get("text") or ""))),
|
||||
int(page.get("line_count", len((page.get("text") or "").splitlines()))),
|
||||
now,
|
||||
)
|
||||
for page in pages
|
||||
]
|
||||
with get_connection() as conn:
|
||||
doc_row = conn.execute(
|
||||
"SELECT 1 FROM rag_documents WHERE id = ?",
|
||||
(document_id,),
|
||||
).fetchone()
|
||||
if doc_row is None:
|
||||
raise sqlite3.IntegrityError("FOREIGN KEY constraint failed")
|
||||
conn.execute("DELETE FROM rag_document_pages WHERE document_id = ?", (document_id,))
|
||||
if rows:
|
||||
conn.executemany(
|
||||
"""
|
||||
INSERT INTO rag_document_pages
|
||||
(document_id, page_index, page_number, text, char_count,
|
||||
line_count, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
rows,
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
def _all_scope_chunks(scope: str) -> list[dict]:
|
||||
if scope.startswith("kb_"):
|
||||
kb_id = scope[len("kb_") :]
|
||||
|
|
@ -579,6 +668,17 @@ def _pump(
|
|||
elif mtype == "dim":
|
||||
embedding_dim = int(msg["dim"])
|
||||
vector_store.ensure_collection(state.scope, embedding_dim)
|
||||
elif mtype == "document_pages":
|
||||
try:
|
||||
_replace_document_pages(
|
||||
state.document_id,
|
||||
list(msg.get("pages") or []),
|
||||
)
|
||||
except sqlite3.IntegrityError as exc:
|
||||
final_error = (
|
||||
f"document was removed before ingestion finished ({exc})"
|
||||
)
|
||||
break
|
||||
elif mtype == "chunks_batch":
|
||||
if embedding_dim is None:
|
||||
embedding_dim = len(msg["vectors"][0]) if msg["vectors"] else None
|
||||
|
|
|
|||
494
studio/backend/core/rag/locators.py
Normal file
494
studio/backend/core/rag/locators.py
Normal file
|
|
@ -0,0 +1,494 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Backfill and PDF-region helpers for durable RAG chunk locators."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from loggers import get_logger
|
||||
from storage.studio_db import get_connection
|
||||
|
||||
from . import vector_store
|
||||
from .parsers import ParsedPage, parse
|
||||
from .vector_store import kb_scope, thread_scope
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
class LocatorMatch:
|
||||
page_index: int
|
||||
page_number: int | None
|
||||
start: int
|
||||
end: int
|
||||
line_start: int
|
||||
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
|
||||
return line_start, line_end
|
||||
|
||||
|
||||
def _find_exact(page_text: str, needle: str) -> list[tuple[int, int]]:
|
||||
if not needle:
|
||||
return []
|
||||
out: list[tuple[int, int]] = []
|
||||
cursor = 0
|
||||
while True:
|
||||
idx = page_text.find(needle, cursor)
|
||||
if idx < 0:
|
||||
break
|
||||
out.append((idx, idx + len(needle)))
|
||||
cursor = idx + 1
|
||||
return out
|
||||
|
||||
|
||||
def _normalize_with_map(text: str) -> tuple[str, list[int], list[int]]:
|
||||
chars: list[str] = []
|
||||
starts: list[int] = []
|
||||
ends: list[int] = []
|
||||
last_space = False
|
||||
for idx, ch in enumerate(text):
|
||||
if ch.isspace():
|
||||
if chars and not last_space:
|
||||
chars.append(" ")
|
||||
starts.append(idx)
|
||||
ends.append(idx + 1)
|
||||
elif chars and last_space:
|
||||
ends[-1] = idx + 1
|
||||
last_space = True
|
||||
continue
|
||||
chars.append(ch.casefold())
|
||||
starts.append(idx)
|
||||
ends.append(idx + 1)
|
||||
last_space = False
|
||||
|
||||
first = 0
|
||||
while first < len(chars) and chars[first] == " ":
|
||||
first += 1
|
||||
last = len(chars)
|
||||
while last > first and chars[last - 1] == " ":
|
||||
last -= 1
|
||||
return "".join(chars[first:last]), starts[first:last], ends[first:last]
|
||||
|
||||
|
||||
def _find_normalized(page_text: str, needle: str) -> list[tuple[int, int]]:
|
||||
norm_page, starts, ends = _normalize_with_map(page_text)
|
||||
norm_needle, _needle_starts, _needle_ends = _normalize_with_map(needle)
|
||||
if not norm_page or not norm_needle:
|
||||
return []
|
||||
out: list[tuple[int, int]] = []
|
||||
cursor = 0
|
||||
while True:
|
||||
idx = norm_page.find(norm_needle, cursor)
|
||||
if idx < 0:
|
||||
break
|
||||
end_idx = idx + len(norm_needle) - 1
|
||||
if 0 <= idx < len(starts) and 0 <= end_idx < len(ends):
|
||||
out.append((starts[idx], ends[end_idx]))
|
||||
cursor = idx + 1
|
||||
return out
|
||||
|
||||
|
||||
def _locate_unique(text: str, pages: list[ParsedPage]) -> tuple[LocatorMatch | None, str]:
|
||||
text = (text or "").strip()
|
||||
if not text:
|
||||
return None, "missing"
|
||||
|
||||
matches: list[LocatorMatch] = []
|
||||
for page_index, page in enumerate(pages):
|
||||
for start, end in _find_exact(page.text, text):
|
||||
line_start, line_end = _line_bounds(page.text, start, end)
|
||||
matches.append(
|
||||
LocatorMatch(
|
||||
page_index = page_index,
|
||||
page_number = page.page_number,
|
||||
start = start,
|
||||
end = end,
|
||||
line_start = line_start,
|
||||
line_end = line_end,
|
||||
)
|
||||
)
|
||||
if len(matches) == 1:
|
||||
return matches[0], "matched"
|
||||
if len(matches) > 1:
|
||||
return None, "ambiguous"
|
||||
|
||||
for page_index, page in enumerate(pages):
|
||||
for start, end in _find_normalized(page.text, text):
|
||||
line_start, line_end = _line_bounds(page.text, start, end)
|
||||
matches.append(
|
||||
LocatorMatch(
|
||||
page_index = page_index,
|
||||
page_number = page.page_number,
|
||||
start = start,
|
||||
end = end,
|
||||
line_start = line_start,
|
||||
line_end = line_end,
|
||||
)
|
||||
)
|
||||
if len(matches) == 1:
|
||||
return matches[0], "matched"
|
||||
if len(matches) > 1:
|
||||
return None, "ambiguous"
|
||||
return None, "missing"
|
||||
|
||||
|
||||
def _replace_document_pages(document_id: str, pages: list[ParsedPage]) -> None:
|
||||
now = int(time.time())
|
||||
rows = [
|
||||
(
|
||||
document_id,
|
||||
index,
|
||||
page.page_number,
|
||||
page.text,
|
||||
len(page.text),
|
||||
len(page.text.splitlines()),
|
||||
now,
|
||||
)
|
||||
for index, page in enumerate(pages)
|
||||
]
|
||||
with get_connection() as conn:
|
||||
conn.execute("DELETE FROM rag_document_pages WHERE document_id = ?", (document_id,))
|
||||
if rows:
|
||||
conn.executemany(
|
||||
"""
|
||||
INSERT INTO rag_document_pages
|
||||
(document_id, page_index, page_number, text, char_count,
|
||||
line_count, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
rows,
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
def _region_anchor(page_text: str, match: LocatorMatch) -> str | None:
|
||||
segment = page_text[match.start : match.end]
|
||||
words = [w.strip(" \t\r\n*#`[]()") for w in segment.split()]
|
||||
words = [w for w in words if len(w) >= 2]
|
||||
if len(words) < 3:
|
||||
return None
|
||||
anchor = " ".join(words[: min(16, len(words))])
|
||||
return anchor if len(anchor) >= 12 else None
|
||||
|
||||
|
||||
def _normalized_occurrences(haystack: str, needle: str) -> int:
|
||||
norm_haystack, _starts, _ends = _normalize_with_map(haystack)
|
||||
norm_needle, _needle_starts, _needle_ends = _normalize_with_map(needle)
|
||||
if not norm_haystack or not norm_needle:
|
||||
return 0
|
||||
count = 0
|
||||
cursor = 0
|
||||
while True:
|
||||
idx = norm_haystack.find(norm_needle, cursor)
|
||||
if idx < 0:
|
||||
return count
|
||||
count += 1
|
||||
cursor = idx + 1
|
||||
|
||||
|
||||
def pdf_regions_for_match(
|
||||
pdf_path: Path,
|
||||
pages: list[ParsedPage],
|
||||
match: LocatorMatch,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Return normalized PDF rectangles for a unique chunk match.
|
||||
|
||||
Regions are intentionally conservative: no PyMuPDF, no page, no
|
||||
unique anchor, or no positive-area rectangles all produce an empty
|
||||
list rather than guessed highlights.
|
||||
"""
|
||||
if pdf_path.suffix.lower() != ".pdf":
|
||||
return []
|
||||
if match.page_index < 0 or match.page_index >= len(pages):
|
||||
return []
|
||||
anchor = _region_anchor(pages[match.page_index].text, match)
|
||||
if not anchor:
|
||||
return []
|
||||
|
||||
try:
|
||||
import pymupdf
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
try:
|
||||
doc = pymupdf.open(str(pdf_path))
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
try:
|
||||
return _pdf_regions_for_match_doc(doc, pages, match, anchor)
|
||||
finally:
|
||||
doc.close()
|
||||
|
||||
|
||||
def _pdf_regions_for_match_doc(
|
||||
doc: Any,
|
||||
pages: list[ParsedPage],
|
||||
match: LocatorMatch,
|
||||
anchor: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
try:
|
||||
if match.page_index >= len(doc):
|
||||
return []
|
||||
page = doc[match.page_index]
|
||||
raw_text = page.get_text("text") or ""
|
||||
if _normalized_occurrences(raw_text, anchor) != 1:
|
||||
return []
|
||||
rects = page.search_for(anchor) or []
|
||||
page_rect = page.rect
|
||||
page_width = float(page_rect.width)
|
||||
page_height = float(page_rect.height)
|
||||
if page_width <= 0 or page_height <= 0:
|
||||
return []
|
||||
|
||||
out: list[dict[str, Any]] = []
|
||||
for rect in rects:
|
||||
width = max(0.0, float(rect.x1 - rect.x0))
|
||||
height = max(0.0, float(rect.y1 - rect.y0))
|
||||
if width <= 0 or height <= 0:
|
||||
continue
|
||||
out.append(
|
||||
{
|
||||
"pageIndex": match.page_index,
|
||||
"pageNumber": match.page_number,
|
||||
"x": max(0.0, min(1.0, float(rect.x0) / page_width)),
|
||||
"y": max(0.0, min(1.0, float(rect.y0) / page_height)),
|
||||
"width": max(0.0, min(1.0, width / page_width)),
|
||||
"height": max(0.0, min(1.0, height / page_height)),
|
||||
"confidence": "exact",
|
||||
"source": "pymupdf-search",
|
||||
}
|
||||
)
|
||||
return out
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def pdf_regions_for_chunks(
|
||||
pdf_path: Path,
|
||||
pages: list[ParsedPage],
|
||||
chunks: list[Any],
|
||||
) -> list[list[dict[str, Any]]]:
|
||||
if pdf_path.suffix.lower() != ".pdf":
|
||||
return [[] for _ in chunks]
|
||||
try:
|
||||
import pymupdf
|
||||
|
||||
doc = pymupdf.open(str(pdf_path))
|
||||
except Exception:
|
||||
return [[] for _ in chunks]
|
||||
|
||||
regions: list[list[dict[str, Any]]] = []
|
||||
try:
|
||||
for chunk in chunks:
|
||||
page_index = getattr(chunk, "source_page_index", None)
|
||||
start = getattr(chunk, "page_char_start", None)
|
||||
end = getattr(chunk, "page_char_end", None)
|
||||
if page_index is None or start is None or end is None:
|
||||
regions.append([])
|
||||
continue
|
||||
if page_index < 0 or page_index >= len(pages):
|
||||
regions.append([])
|
||||
continue
|
||||
line_start, line_end = _line_bounds(pages[page_index].text, start, end)
|
||||
match = LocatorMatch(
|
||||
page_index = int(page_index),
|
||||
page_number = getattr(chunk, "page_number", None),
|
||||
start = int(start),
|
||||
end = int(end),
|
||||
line_start = line_start,
|
||||
line_end = line_end,
|
||||
)
|
||||
anchor = _region_anchor(pages[match.page_index].text, match)
|
||||
if not anchor:
|
||||
regions.append([])
|
||||
continue
|
||||
regions.append(_pdf_regions_for_match_doc(doc, pages, match, anchor))
|
||||
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 get_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 get_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),
|
||||
)
|
||||
|
|
@ -50,6 +50,11 @@ class Hit:
|
|||
document_id: str | None = None
|
||||
chunk_index: int | None = None
|
||||
kind: str = "text"
|
||||
source_page_index: int | None = None
|
||||
page_char_start: int | None = None
|
||||
page_char_end: int | None = None
|
||||
line_start: int | None = None
|
||||
line_end: int | None = None
|
||||
# Raw cosine; None for BM25-only hits.
|
||||
dense_score: float | None = None
|
||||
|
||||
|
|
@ -151,6 +156,11 @@ def retrieve_dense(
|
|||
document_id = payload.get("document_id"),
|
||||
chunk_index = payload.get("chunk_index"),
|
||||
kind = payload.get("kind", "text"),
|
||||
source_page_index = payload.get("source_page_index"),
|
||||
page_char_start = payload.get("page_char_start"),
|
||||
page_char_end = payload.get("page_char_end"),
|
||||
line_start = payload.get("line_start"),
|
||||
line_end = payload.get("line_end"),
|
||||
dense_score = r["score"],
|
||||
)
|
||||
)
|
||||
|
|
@ -184,6 +194,11 @@ def _rrf_fuse(
|
|||
document_id = seen[cid].document_id,
|
||||
chunk_index = seen[cid].chunk_index,
|
||||
kind = seen[cid].kind,
|
||||
source_page_index = seen[cid].source_page_index,
|
||||
page_char_start = seen[cid].page_char_start,
|
||||
page_char_end = seen[cid].page_char_end,
|
||||
line_start = seen[cid].line_start,
|
||||
line_end = seen[cid].line_end,
|
||||
dense_score = dense_scores.get(cid),
|
||||
)
|
||||
for cid, score in ordered
|
||||
|
|
|
|||
|
|
@ -93,12 +93,33 @@ def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str:
|
|||
f'id="{index}"',
|
||||
f'source="{_xml_attr(hit.get("filename") or "unknown")}"',
|
||||
]
|
||||
# Durable backend ids — additive per contracts.md §3.1 (T3).
|
||||
# ``id`` above stays as the visible citation id (used by the model
|
||||
# as `[N]`); ``document_id`` + ``chunk_id`` are what the preview
|
||||
# route consumes. Old XML without these attrs still parses on
|
||||
# the frontend (hover-only), per contracts §3.2.
|
||||
document_id = hit.get("document_id")
|
||||
if document_id:
|
||||
attrs.append(f'document_id="{_xml_attr(document_id)}"')
|
||||
backend_chunk_id = hit.get("chunk_id")
|
||||
if backend_chunk_id:
|
||||
attrs.append(f'chunk_id="{_xml_attr(backend_chunk_id)}"')
|
||||
page = hit.get("page_number")
|
||||
if page is not None:
|
||||
attrs.append(f'page="{page}"')
|
||||
chunk_index = hit.get("chunk_index")
|
||||
if chunk_index is not None:
|
||||
attrs.append(f'chunk_index="{chunk_index}"')
|
||||
for attr_name in (
|
||||
"source_page_index",
|
||||
"page_char_start",
|
||||
"page_char_end",
|
||||
"line_start",
|
||||
"line_end",
|
||||
):
|
||||
value = hit.get(attr_name)
|
||||
if value is not None:
|
||||
attrs.append(f'{attr_name}="{value}"')
|
||||
tokens = hit.get("token_count")
|
||||
if tokens:
|
||||
attrs.append(f'tokens="{tokens}"')
|
||||
|
|
@ -106,7 +127,6 @@ def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str:
|
|||
if kind and kind != "text":
|
||||
attrs.append(f'kind="{_xml_attr(kind)}"')
|
||||
image_path = hit.get("image_path")
|
||||
document_id = hit.get("document_id")
|
||||
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.
|
||||
|
|
@ -213,6 +233,8 @@ def search_knowledge_base(
|
|||
f"""
|
||||
SELECT c.id AS chunk_id, c.text, c.page_number,
|
||||
c.token_count, c.kind, c.image_path,
|
||||
c.source_page_index, c.page_char_start,
|
||||
c.page_char_end, c.line_start, c.line_end,
|
||||
c.document_id, d.filename
|
||||
FROM rag_chunks c
|
||||
JOIN rag_documents d ON d.id = c.document_id
|
||||
|
|
|
|||
|
|
@ -131,6 +131,45 @@ def search(
|
|||
return out
|
||||
|
||||
|
||||
def update_chunk_payload_fields(
|
||||
scope: str,
|
||||
updates: dict[str, dict],
|
||||
) -> None:
|
||||
"""Merge locator fields into existing vector payload JSON by chunk id."""
|
||||
from core.rag.db import get_rag_connection
|
||||
|
||||
if not updates:
|
||||
return
|
||||
conn = get_rag_connection()
|
||||
rows = conn.execute(
|
||||
f"""
|
||||
SELECT chunk_id, payload_json
|
||||
FROM rag_vectors
|
||||
WHERE scope = ? AND chunk_id IN ({",".join("?" for _ in updates)})
|
||||
""",
|
||||
[scope, *updates.keys()],
|
||||
).fetchall()
|
||||
payload_rows: list[tuple[str, str]] = []
|
||||
for row in rows:
|
||||
try:
|
||||
payload = json.loads(row["payload_json"] or "{}")
|
||||
except json.JSONDecodeError:
|
||||
payload = {}
|
||||
payload.update(updates.get(row["chunk_id"], {}))
|
||||
payload_rows.append((json.dumps(payload, default = str), row["chunk_id"], scope))
|
||||
if not payload_rows:
|
||||
return
|
||||
conn.executemany(
|
||||
"""
|
||||
UPDATE rag_vectors
|
||||
SET payload_json = ?
|
||||
WHERE chunk_id = ? AND scope = ?
|
||||
""",
|
||||
payload_rows,
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
def delete_scope(scope: str) -> None:
|
||||
from core.rag.db import get_rag_connection
|
||||
|
||||
|
|
|
|||
|
|
@ -10,10 +10,13 @@ import json
|
|||
import os
|
||||
import queue as queue_module
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, Optional
|
||||
from urllib.parse import quote
|
||||
from uuid import uuid4
|
||||
|
||||
import jwt
|
||||
from fastapi import (
|
||||
APIRouter,
|
||||
Depends,
|
||||
|
|
@ -23,10 +26,11 @@ from fastapi import (
|
|||
Request,
|
||||
UploadFile,
|
||||
)
|
||||
from fastapi.responses import FileResponse, StreamingResponse
|
||||
from fastapi.responses import FileResponse, Response, StreamingResponse
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from auth.authentication import get_current_subject, get_current_subject_sse
|
||||
from auth.storage import get_jwt_secret
|
||||
|
||||
|
||||
async def _sse_auth(
|
||||
|
|
@ -37,6 +41,8 @@ async def _sse_auth(
|
|||
|
||||
|
||||
from core.rag import embeddings, ingestion, reranker, 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 (
|
||||
|
|
@ -44,7 +50,7 @@ from storage.studio_db import (
|
|||
list_chat_settings,
|
||||
upsert_chat_settings_merge,
|
||||
)
|
||||
from utils.paths.storage_roots import ensure_dir, rag_uploads_root
|
||||
from utils.paths.storage_roots import ensure_dir, rag_uploads_root, resolve_under_root
|
||||
from utils.rag.config import (
|
||||
RAG_MAX_UPLOAD_MB,
|
||||
RAG_RERANK_CANDIDATE_K,
|
||||
|
|
@ -139,6 +145,11 @@ class SearchHit(BaseModel):
|
|||
filename: str | None = None
|
||||
kind: str = "text"
|
||||
image_url: str | None = None
|
||||
source_page_index: int | None = None
|
||||
page_char_start: int | None = None
|
||||
page_char_end: int | None = None
|
||||
line_start: int | None = None
|
||||
line_end: int | None = None
|
||||
|
||||
|
||||
class SearchResponse(BaseModel):
|
||||
|
|
@ -872,6 +883,7 @@ def get_rag_image(
|
|||
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"))
|
||||
|
|
@ -1016,6 +1028,538 @@ async def _replay_terminal_state(row: Any):
|
|||
# --- Search ---
|
||||
|
||||
|
||||
# --- Document preview (file + preview-target) ---
|
||||
|
||||
|
||||
PreviewMediaKind = Literal["pdf", "text", "docx", "html", "image", "unknown"]
|
||||
PreviewChunkKind = Literal["text", "image", "caption"]
|
||||
|
||||
|
||||
class PreviewPdfRegion(BaseModel):
|
||||
pageIndex: int
|
||||
pageNumber: int | None = None
|
||||
x: float
|
||||
y: float
|
||||
width: float
|
||||
height: float
|
||||
confidence: Literal["exact"]
|
||||
source: str
|
||||
|
||||
|
||||
class PreviewTargetResponse(BaseModel):
|
||||
"""Per contracts.md §1.2 / §1.3 — single shape covering both
|
||||
cited-chunk and document-row preview modes. The §1.3 metadata-only
|
||||
mode returns ``None`` for ``chunkId``/``chunkIndex``/``targetPage``/
|
||||
``snippet``/``kind``/``imageUrl`` (Q2: no first-chunk guessing).
|
||||
"""
|
||||
|
||||
documentId: str
|
||||
filename: str
|
||||
contentType: str | None
|
||||
mediaKind: PreviewMediaKind
|
||||
byteSize: int
|
||||
status: str
|
||||
kbId: str | None
|
||||
threadId: str | None
|
||||
chunkId: str | None
|
||||
chunkIndex: int | None
|
||||
targetPage: int | None
|
||||
snippet: str | None
|
||||
kind: PreviewChunkKind | None
|
||||
imageUrl: str | None
|
||||
sourcePageIndex: int | None
|
||||
pageCharStart: int | None
|
||||
pageCharEnd: int | None
|
||||
lineStart: int | None
|
||||
lineEnd: int | None
|
||||
pdfRegions: list[PreviewPdfRegion] = Field(default_factory = list)
|
||||
|
||||
|
||||
class PreviewFileUrlResponse(BaseModel):
|
||||
url: str
|
||||
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. Anything not in
|
||||
# this map collapses to ("application/octet-stream", attachment, "unknown").
|
||||
# .html / .htm intentionally serve as text/plain attachment (decisions Q7 +
|
||||
# Risk #3) so an uploaded HTML cannot execute in the app origin.
|
||||
_PREVIEW_EXT_MAP: dict[str, tuple[str, str, PreviewMediaKind]] = {
|
||||
".pdf": ("application/pdf", "inline", "pdf"),
|
||||
".txt": ("text/plain; charset=utf-8", "inline", "text"),
|
||||
".md": ("text/markdown; charset=utf-8", "inline", "text"),
|
||||
".markdown": ("text/markdown; charset=utf-8", "inline", "text"),
|
||||
".docx": (
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
"attachment",
|
||||
"docx",
|
||||
),
|
||||
".html": ("text/plain; charset=utf-8", "attachment", "html"),
|
||||
".htm": ("text/plain; charset=utf-8", "attachment", "html"),
|
||||
".png": ("image/png", "inline", "image"),
|
||||
".jpg": ("image/jpeg", "inline", "image"),
|
||||
".jpeg": ("image/jpeg", "inline", "image"),
|
||||
".gif": ("image/gif", "inline", "image"),
|
||||
".webp": ("image/webp", "inline", "image"),
|
||||
}
|
||||
|
||||
|
||||
def _ascii_only(value: str) -> bool:
|
||||
try:
|
||||
value.encode("ascii")
|
||||
except UnicodeEncodeError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _content_disposition_header(filename: str, disposition: str) -> str:
|
||||
"""Build a Content-Disposition header. Non-ASCII filenames use RFC 5987
|
||||
``filename*=UTF-8''…`` alongside an ASCII-only ``filename=`` fallback so
|
||||
older clients still get something readable.
|
||||
"""
|
||||
from urllib.parse import quote as _urlquote
|
||||
|
||||
safe = filename.replace('"', "").replace("\r", "").replace("\n", "")
|
||||
if _ascii_only(safe):
|
||||
return f'{disposition}; filename="{safe}"'
|
||||
ascii_fallback = safe.encode("ascii", "replace").decode("ascii")
|
||||
encoded = _urlquote(safe, safe = "")
|
||||
return (
|
||||
f'{disposition}; filename="{ascii_fallback}"; '
|
||||
f"filename*=UTF-8''{encoded}"
|
||||
)
|
||||
|
||||
|
||||
def _preview_file_metadata(filename: str) -> tuple[str, str, PreviewMediaKind]:
|
||||
"""Map a stored filename to (content_type, disposition, mediaKind).
|
||||
|
||||
Unknown extensions always force ``application/octet-stream`` +
|
||||
``attachment`` + ``unknown`` so the browser cannot sniff a sensitive
|
||||
type and inline it (Risk #3).
|
||||
"""
|
||||
ext = Path(filename).suffix.lower()
|
||||
return _PREVIEW_EXT_MAP.get(
|
||||
ext, ("application/octet-stream", "attachment", "unknown")
|
||||
)
|
||||
|
||||
|
||||
_PREVIEW_FILE_AUDIENCE = "rag-preview-file"
|
||||
_PREVIEW_FILE_TTL_SECONDS = 5 * 60
|
||||
_JWT_ALGORITHM = "HS256"
|
||||
|
||||
|
||||
def _parse_pdf_regions(value: str | None) -> list[PreviewPdfRegion]:
|
||||
if not value:
|
||||
return []
|
||||
try:
|
||||
raw = json.loads(value)
|
||||
except (TypeError, json.JSONDecodeError):
|
||||
return []
|
||||
if not isinstance(raw, list):
|
||||
return []
|
||||
out: list[PreviewPdfRegion] = []
|
||||
for item in raw:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
try:
|
||||
region = PreviewPdfRegion(**item)
|
||||
except Exception:
|
||||
continue
|
||||
if (
|
||||
0 <= region.x <= 1
|
||||
and 0 <= region.y <= 1
|
||||
and region.width > 0
|
||||
and region.height > 0
|
||||
):
|
||||
out.append(region)
|
||||
return out
|
||||
|
||||
|
||||
def _preview_file_token(
|
||||
*,
|
||||
subject: str,
|
||||
document_id: str,
|
||||
) -> tuple[str, int]:
|
||||
secret = get_jwt_secret(subject)
|
||||
if secret is None:
|
||||
raise HTTPException(status_code = 401, detail = "Invalid or expired token")
|
||||
expires = datetime.now(timezone.utc) + timedelta(seconds = _PREVIEW_FILE_TTL_SECONDS)
|
||||
payload = {
|
||||
"sub": subject,
|
||||
"aud": _PREVIEW_FILE_AUDIENCE,
|
||||
"document_id": document_id,
|
||||
"exp": expires,
|
||||
}
|
||||
token = jwt.encode(payload, secret, algorithm = _JWT_ALGORITHM)
|
||||
return token, int(expires.timestamp())
|
||||
|
||||
|
||||
def _subject_from_preview_file_token(document_id: str, token: str) -> str:
|
||||
try:
|
||||
unverified = jwt.decode(
|
||||
token,
|
||||
options = {
|
||||
"verify_signature": False,
|
||||
"verify_exp": False,
|
||||
"verify_aud": False,
|
||||
},
|
||||
)
|
||||
except jwt.InvalidTokenError as exc:
|
||||
raise HTTPException(status_code = 401, detail = "Invalid preview token") from exc
|
||||
subject = unverified.get("sub")
|
||||
if not isinstance(subject, str) or not subject:
|
||||
raise HTTPException(status_code = 401, detail = "Invalid preview token")
|
||||
secret = get_jwt_secret(subject)
|
||||
if secret is None:
|
||||
raise HTTPException(status_code = 401, detail = "Invalid preview token")
|
||||
try:
|
||||
payload = jwt.decode(
|
||||
token,
|
||||
secret,
|
||||
algorithms = [_JWT_ALGORITHM],
|
||||
audience = _PREVIEW_FILE_AUDIENCE,
|
||||
)
|
||||
except jwt.InvalidTokenError as exc:
|
||||
raise HTTPException(status_code = 401, detail = "Invalid preview token") from exc
|
||||
if payload.get("document_id") != document_id:
|
||||
raise HTTPException(status_code = 401, detail = "Invalid preview token")
|
||||
return subject
|
||||
|
||||
|
||||
def _resolve_document_file_or_404(doc_row: Any, document_id: str) -> Path:
|
||||
try:
|
||||
resolved = resolve_under_root(
|
||||
doc_row["stored_path"],
|
||||
root = rag_uploads_root(),
|
||||
)
|
||||
except ValueError as exc:
|
||||
# Symlink escape / ``..`` / absolute outside root — collapse to
|
||||
# "file not found" (the auth row exists, the bytes do not).
|
||||
logger.warning(
|
||||
"RAG preview: stored_path escaped uploads root for doc %s: %s",
|
||||
document_id,
|
||||
exc,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 404,
|
||||
detail = "Document file not found",
|
||||
) from exc
|
||||
|
||||
if not resolved.is_file():
|
||||
raise HTTPException(
|
||||
status_code = 404,
|
||||
detail = "Document file not found",
|
||||
)
|
||||
return resolved
|
||||
|
||||
|
||||
def _parse_range_header(range_header: str | None, size: int) -> tuple[int, int] | None:
|
||||
if not range_header:
|
||||
return None
|
||||
if not range_header.startswith("bytes="):
|
||||
raise ValueError("unsupported range unit")
|
||||
spec = range_header[len("bytes=") :].strip()
|
||||
if "," in spec or "-" not in spec:
|
||||
raise ValueError("multiple or malformed ranges are not supported")
|
||||
start_s, end_s = spec.split("-", 1)
|
||||
if not start_s and not end_s:
|
||||
raise ValueError("empty range")
|
||||
if not start_s:
|
||||
suffix = int(end_s)
|
||||
if suffix <= 0:
|
||||
raise ValueError("invalid suffix range")
|
||||
start = max(0, size - suffix)
|
||||
end = size - 1
|
||||
else:
|
||||
start = int(start_s)
|
||||
end = int(end_s) if end_s else size - 1
|
||||
if start < 0 or end < start or start >= size:
|
||||
raise ValueError("range outside file")
|
||||
return start, min(end, size - 1)
|
||||
|
||||
|
||||
def _iter_file_range(path: Path, start: int, end: int):
|
||||
with path.open("rb") as fh:
|
||||
fh.seek(start)
|
||||
remaining = end - start + 1
|
||||
while remaining > 0:
|
||||
chunk = fh.read(min(64 * 1024, remaining))
|
||||
if not chunk:
|
||||
break
|
||||
remaining -= len(chunk)
|
||||
yield chunk
|
||||
|
||||
|
||||
def _serve_document_file_row(
|
||||
doc_row: Any,
|
||||
document_id: str,
|
||||
range_header: str | None,
|
||||
) -> FileResponse | Response | StreamingResponse:
|
||||
resolved = _resolve_document_file_or_404(doc_row, document_id)
|
||||
content_type, disposition, _media_kind = _preview_file_metadata(
|
||||
doc_row["filename"]
|
||||
)
|
||||
safe_name = _sanitize_filename(doc_row["filename"])
|
||||
headers = {
|
||||
"Content-Disposition": _content_disposition_header(safe_name, disposition),
|
||||
"X-Content-Type-Options": "nosniff",
|
||||
"Cache-Control": "private, max-age=0, must-revalidate",
|
||||
"Accept-Ranges": "bytes",
|
||||
}
|
||||
size = resolved.stat().st_size
|
||||
|
||||
try:
|
||||
byte_range = _parse_range_header(range_header, size)
|
||||
except (TypeError, ValueError):
|
||||
range_headers = dict(headers)
|
||||
range_headers["Content-Range"] = f"bytes */{size}"
|
||||
return Response(status_code = 416, headers = range_headers)
|
||||
|
||||
if byte_range is not None:
|
||||
start, end = byte_range
|
||||
range_headers = dict(headers)
|
||||
range_headers["Content-Range"] = f"bytes {start}-{end}/{size}"
|
||||
range_headers["Content-Length"] = str(end - start + 1)
|
||||
return StreamingResponse(
|
||||
_iter_file_range(resolved, start, end),
|
||||
status_code = 206,
|
||||
media_type = content_type,
|
||||
headers = range_headers,
|
||||
)
|
||||
|
||||
# Starlette's FileResponse handles ordinary downloads efficiently. We
|
||||
# still advertise Accept-Ranges so PDF.js can switch to explicit range
|
||||
# requests via the signed URL path.
|
||||
return FileResponse(
|
||||
path = str(resolved),
|
||||
media_type = content_type,
|
||||
headers = headers,
|
||||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/documents/{document_id}/preview-target",
|
||||
response_model = PreviewTargetResponse,
|
||||
)
|
||||
def get_document_preview_target(
|
||||
document_id: str,
|
||||
chunk_id: Optional[str] = Query(None),
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> PreviewTargetResponse:
|
||||
"""Resolve preview metadata for a document, optionally focused on a chunk.
|
||||
|
||||
See contracts.md §1 for the response shape. ``chunk_id`` is the durable
|
||||
``rag_chunks.id`` carried as ``backendChunkId`` on the frontend; when
|
||||
supplied it MUST belong to ``document_id`` or we collapse to 404 (so a
|
||||
probe cannot enumerate cross-document chunk ids).
|
||||
"""
|
||||
doc_row = document_for_subject_or_404(document_id, current_subject)
|
||||
_ct, _disposition, media_kind = _preview_file_metadata(doc_row["filename"])
|
||||
|
||||
base = {
|
||||
"documentId": doc_row["id"],
|
||||
"filename": doc_row["filename"],
|
||||
"contentType": doc_row["content_type"],
|
||||
"mediaKind": media_kind,
|
||||
"byteSize": int(doc_row["byte_size"]),
|
||||
"status": doc_row["status"],
|
||||
"kbId": doc_row["kb_id"],
|
||||
"threadId": doc_row["thread_id"],
|
||||
}
|
||||
|
||||
if not chunk_id:
|
||||
# Q2: document-row preview returns metadata only — frontend MUST
|
||||
# NOT fall back to "first chunk".
|
||||
return PreviewTargetResponse(
|
||||
**base,
|
||||
chunkId = None,
|
||||
chunkIndex = None,
|
||||
targetPage = None,
|
||||
snippet = None,
|
||||
kind = None,
|
||||
imageUrl = None,
|
||||
sourcePageIndex = None,
|
||||
pageCharStart = None,
|
||||
pageCharEnd = None,
|
||||
lineStart = None,
|
||||
lineEnd = None,
|
||||
pdfRegions = [],
|
||||
)
|
||||
|
||||
# Single connection enforces membership AND fetches the row in one
|
||||
# query. Splitting this into a separate `chunk_belongs_to_document`
|
||||
# call would open a second SQLite connection and create a TOCTOU
|
||||
# window — if the chunk is deleted between the two calls, the data
|
||||
# fetch returns None and the route 500s on the next attribute access
|
||||
# (devils-advocate D1.1). The cross-document case still collapses to
|
||||
# the same 404 the auth helper emits — never 400 (would leak doc
|
||||
# existence).
|
||||
with get_connection() as conn:
|
||||
chunk_row = conn.execute(
|
||||
"""
|
||||
SELECT id, chunk_index, page_number, text, kind, image_path,
|
||||
source_page_index, page_char_start, page_char_end,
|
||||
line_start, line_end, pdf_regions_json
|
||||
FROM rag_chunks WHERE id = ? AND document_id = ?
|
||||
""",
|
||||
(chunk_id, document_id),
|
||||
).fetchone()
|
||||
|
||||
if chunk_row is None:
|
||||
raise HTTPException(
|
||||
status_code = 404,
|
||||
detail = "Document not found",
|
||||
)
|
||||
|
||||
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,
|
||||
chunkId = chunk_row["id"],
|
||||
chunkIndex = int(chunk_row["chunk_index"]),
|
||||
targetPage = (
|
||||
int(chunk_row["page_number"])
|
||||
if chunk_row["page_number"] is not None
|
||||
else None
|
||||
),
|
||||
snippet = chunk_row["text"] or "",
|
||||
kind = chunk_kind,
|
||||
imageUrl = image_url,
|
||||
sourcePageIndex = chunk_row["source_page_index"],
|
||||
pageCharStart = chunk_row["page_char_start"],
|
||||
pageCharEnd = chunk_row["page_char_end"],
|
||||
lineStart = chunk_row["line_start"],
|
||||
lineEnd = chunk_row["line_end"],
|
||||
pdfRegions = _parse_pdf_regions(chunk_row["pdf_regions_json"]),
|
||||
)
|
||||
|
||||
|
||||
@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,
|
||||
)
|
||||
def get_document_file_url(
|
||||
document_id: str,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> PreviewFileUrlResponse:
|
||||
"""Mint a short-lived signed URL for PDF.js range requests.
|
||||
|
||||
The normal bearer-protected `/file` route stays available for blob
|
||||
fallback; this route keeps bearer access tokens out of query strings.
|
||||
"""
|
||||
document_for_subject_or_404(document_id, current_subject)
|
||||
token, expires_at = _preview_file_token(
|
||||
subject = current_subject,
|
||||
document_id = document_id,
|
||||
)
|
||||
url = (
|
||||
f"/api/rag/documents/{quote(document_id, safe = '')}/file-signed"
|
||||
f"?token={quote(token, safe = '')}"
|
||||
)
|
||||
return PreviewFileUrlResponse(url = url, expiresAt = expires_at)
|
||||
|
||||
|
||||
@router.get("/documents/{document_id}/file-signed", response_model = None)
|
||||
def get_signed_document_file(
|
||||
document_id: str,
|
||||
request: Request,
|
||||
token: str = Query(...),
|
||||
) -> FileResponse | Response | StreamingResponse:
|
||||
subject = _subject_from_preview_file_token(document_id, token)
|
||||
doc_row = document_for_subject_or_404(document_id, subject)
|
||||
return _serve_document_file_row(
|
||||
doc_row,
|
||||
document_id,
|
||||
request.headers.get("range"),
|
||||
)
|
||||
|
||||
|
||||
@router.get("/documents/{document_id}/file", response_model = None)
|
||||
def get_document_file(
|
||||
document_id: str,
|
||||
request: Request,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> FileResponse | Response | StreamingResponse:
|
||||
"""Serve the original uploaded bytes for ``document_id``.
|
||||
|
||||
Path resolution per contracts.md §2.1: the route NEVER accepts a
|
||||
client-supplied filename. ``stored_path`` is DB-issued and we run it
|
||||
through ``resolve_under_root`` which delegates to ``_assert_contained``
|
||||
(realpath + symlink/junction-safe). Any escape (symlink, junction,
|
||||
``..``, absolute outside root) collapses to the second 404.
|
||||
|
||||
Content-Type and disposition come from the extension allowlist
|
||||
(``_preview_file_metadata``). HTML/DOCX/unknown serve as
|
||||
``attachment`` with safe content-type so an uploaded ``.html`` can
|
||||
never execute in the app origin (Risk #3 / decisions Q7).
|
||||
"""
|
||||
doc_row = document_for_subject_or_404(document_id, current_subject)
|
||||
return _serve_document_file_row(
|
||||
doc_row,
|
||||
document_id,
|
||||
request.headers.get("range"),
|
||||
)
|
||||
|
||||
|
||||
@router.post("/search", response_model = SearchResponse)
|
||||
def search(
|
||||
payload: SearchRequest,
|
||||
|
|
@ -1092,6 +1636,8 @@ def search(
|
|||
f"""
|
||||
SELECT c.id AS chunk_id, c.document_id, c.chunk_index, c.text,
|
||||
c.page_number, c.kind, c.image_path, c.linked_chunk_id,
|
||||
c.source_page_index, c.page_char_start, c.page_char_end,
|
||||
c.line_start, c.line_end,
|
||||
d.filename
|
||||
FROM rag_chunks c
|
||||
JOIN rag_documents d ON d.id = c.document_id
|
||||
|
|
@ -1139,6 +1685,11 @@ def search(
|
|||
filename = meta.get("filename"),
|
||||
kind = kind,
|
||||
image_url = image_url,
|
||||
source_page_index = meta.get("source_page_index"),
|
||||
page_char_start = meta.get("page_char_start"),
|
||||
page_char_end = meta.get("page_char_end"),
|
||||
line_start = meta.get("line_start"),
|
||||
line_end = meta.get("line_end"),
|
||||
)
|
||||
)
|
||||
logger.info("RAG search: returning %d hits", len(out))
|
||||
|
|
|
|||
|
|
@ -275,6 +275,12 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
|
|||
kind TEXT NOT NULL DEFAULT 'text',
|
||||
image_path TEXT,
|
||||
linked_chunk_id TEXT,
|
||||
source_page_index INTEGER,
|
||||
page_char_start INTEGER,
|
||||
page_char_end INTEGER,
|
||||
line_start INTEGER,
|
||||
line_end INTEGER,
|
||||
pdf_regions_json TEXT,
|
||||
UNIQUE(document_id, chunk_index)
|
||||
)
|
||||
"""
|
||||
|
|
@ -290,9 +296,38 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
|
|||
conn.execute("ALTER TABLE rag_chunks ADD COLUMN image_path TEXT")
|
||||
if "linked_chunk_id" not in chunk_cols:
|
||||
conn.execute("ALTER TABLE rag_chunks ADD COLUMN linked_chunk_id TEXT")
|
||||
for column in (
|
||||
"source_page_index",
|
||||
"page_char_start",
|
||||
"page_char_end",
|
||||
"line_start",
|
||||
"line_end",
|
||||
):
|
||||
if column not in chunk_cols:
|
||||
conn.execute(f"ALTER TABLE rag_chunks ADD COLUMN {column} INTEGER")
|
||||
if "pdf_regions_json" not in chunk_cols:
|
||||
conn.execute("ALTER TABLE rag_chunks ADD COLUMN pdf_regions_json TEXT")
|
||||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_rag_chunks_document_id ON rag_chunks(document_id)"
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS rag_document_pages (
|
||||
document_id TEXT NOT NULL REFERENCES rag_documents(id) ON DELETE CASCADE,
|
||||
page_index INTEGER NOT NULL,
|
||||
page_number INTEGER,
|
||||
text TEXT NOT NULL,
|
||||
char_count INTEGER NOT NULL DEFAULT 0,
|
||||
line_count INTEGER NOT NULL DEFAULT 0,
|
||||
created_at INTEGER NOT NULL,
|
||||
PRIMARY KEY(document_id, page_index)
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_rag_document_pages_document_id "
|
||||
"ON rag_document_pages(document_id)"
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS rag_ingestion_jobs (
|
||||
|
|
|
|||
269
studio/backend/tests/test_rag_authorization.py
Normal file
269
studio/backend/tests/test_rag_authorization.py
Normal file
|
|
@ -0,0 +1,269 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tests for document_for_subject_or_404 and chunk_belongs_to_document.
|
||||
|
||||
Authorization rules under test (contracts.md §1 / §2, Risk #1):
|
||||
|
||||
- KB documents: KB must exist and KB.owner_user_id must equal current_subject.
|
||||
- Thread documents: thread must exist in chat_threads; current-Studio single-user
|
||||
invariant means any authenticated subject can access, BUT the thread row must
|
||||
exist (a missing thread is 404, not silent grant).
|
||||
- Missing document or missing KB both collapse to 404.
|
||||
- KB with NULL owner_user_id is NOT accessible (legacy row guard).
|
||||
- Both not-found and not-authorized return HTTP 404 with identical detail to
|
||||
prevent document-existence leaking.
|
||||
- chunk_belongs_to_document only returns True when chunk.document_id matches.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
import storage.studio_db as studio_db
|
||||
from core.rag.authorization import chunk_belongs_to_document, document_for_subject_or_404
|
||||
|
||||
|
||||
# ── Fixtures ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _reset_db(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
|
||||
monkeypatch.setattr(studio_db, "_schema_ready", False)
|
||||
|
||||
|
||||
def _uid() -> str:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
def _insert_kb(conn, kb_id: str, owner: str | None = "user-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[:8]}", "bge-small", owner, 1_700_000_000),
|
||||
)
|
||||
|
||||
|
||||
def _insert_thread(conn, thread_id: str) -> None:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO chat_threads (id, title, model_type, model_id, archived, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(thread_id, "Test Thread", "base", "llama3", 0, 1_700_000_000),
|
||||
)
|
||||
|
||||
|
||||
def _insert_kb_doc(conn, doc_id: str, kb_id: str, stored_path: str = "doc.pdf") -> 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, ?, ?, ?, 'completed', 0, 1024, ?)
|
||||
""",
|
||||
(doc_id, kb_id, "report.pdf", "application/pdf", stored_path, 1_700_000_000),
|
||||
)
|
||||
|
||||
|
||||
def _insert_thread_doc(
|
||||
conn, doc_id: str, thread_id: str, stored_path: str = "doc.txt"
|
||||
) -> 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, ?, ?, ?, ?, 'completed', 0, 512, ?)
|
||||
""",
|
||||
(doc_id, thread_id, "note.txt", "text/plain", stored_path, 1_700_000_000),
|
||||
)
|
||||
|
||||
|
||||
def _insert_chunk(conn, chunk_id: str, doc_id: str, chunk_index: int = 0) -> None:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO rag_chunks (id, document_id, chunk_index, text, token_count)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
""",
|
||||
(chunk_id, doc_id, chunk_index, "some chunk text", 20),
|
||||
)
|
||||
|
||||
|
||||
# ── KB-document authorization ─────────────────────────────────────────
|
||||
|
||||
|
||||
def test_kb_doc_correct_owner_returns_row(tmp_path, monkeypatch):
|
||||
"""KB doc authorized when KB.owner_user_id == current_subject."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id, owner="alice")
|
||||
_insert_kb_doc(conn, doc_id, kb_id)
|
||||
row = document_for_subject_or_404(doc_id, "alice")
|
||||
assert row["id"] == doc_id
|
||||
|
||||
|
||||
def test_kb_doc_wrong_owner_raises_404(tmp_path, monkeypatch):
|
||||
"""KB doc returns 404 when current_subject != KB.owner_user_id."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id, owner="alice")
|
||||
_insert_kb_doc(conn, doc_id, kb_id)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
document_for_subject_or_404(doc_id, "mallory")
|
||||
assert exc_info.value.status_code == 404
|
||||
assert exc_info.value.detail == "Document not found"
|
||||
|
||||
|
||||
def test_kb_doc_null_owner_raises_404(tmp_path, monkeypatch):
|
||||
"""KB with NULL owner_user_id is not accessible through the helper (legacy guard)."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id, owner=None)
|
||||
_insert_kb_doc(conn, doc_id, kb_id)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
document_for_subject_or_404(doc_id, "alice")
|
||||
assert exc_info.value.status_code == 404
|
||||
assert exc_info.value.detail == "Document not found"
|
||||
|
||||
|
||||
def test_kb_doc_missing_kb_raises_404(tmp_path, monkeypatch):
|
||||
"""Document rows whose KB was deleted collapse to 404.
|
||||
|
||||
Insert both KB and doc, then delete the KB (ON DELETE CASCADE removes the doc
|
||||
too). A subsequent lookup for the doc id must return 404, not 500.
|
||||
If for some reason the doc row survives (e.g. FK off), the helper must
|
||||
still 404 because the KB is gone.
|
||||
"""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id, owner="alice")
|
||||
_insert_kb_doc(conn, doc_id, kb_id)
|
||||
# Delete the KB — ON DELETE CASCADE should also drop the doc.
|
||||
conn.execute("DELETE FROM rag_knowledge_bases WHERE id = ?", (kb_id,))
|
||||
# After cascade deletion the doc_id no longer exists → 404.
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
document_for_subject_or_404(doc_id, "alice")
|
||||
assert exc_info.value.status_code == 404
|
||||
|
||||
|
||||
# ── Thread-document authorization (single-user invariant) ─────────────
|
||||
|
||||
|
||||
def test_thread_doc_existing_thread_grants_access(tmp_path, monkeypatch):
|
||||
"""Thread doc is accessible when thread exists (single-user Studio invariant)."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
doc_id, thread_id = _uid(), _uid()
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_thread(conn, thread_id)
|
||||
_insert_thread_doc(conn, doc_id, thread_id)
|
||||
row = document_for_subject_or_404(doc_id, "any-authenticated-user")
|
||||
assert row["id"] == doc_id
|
||||
|
||||
|
||||
def test_thread_doc_nonexistent_thread_raises_404(tmp_path, monkeypatch):
|
||||
"""A missing thread_id does NOT silently grant access — it must be 404."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
doc_id, thread_id = _uid(), _uid()
|
||||
# Insert doc with a thread_id that has no matching chat_threads row.
|
||||
with studio_db.get_connection() as conn:
|
||||
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, ?, 'x.txt', 'text/plain', 'x.txt', 'completed', 0, 1, ?)
|
||||
""",
|
||||
(doc_id, thread_id, 1_700_000_000),
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
document_for_subject_or_404(doc_id, "alice")
|
||||
assert exc_info.value.status_code == 404
|
||||
assert exc_info.value.detail == "Document not found"
|
||||
|
||||
|
||||
# ── Missing document ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_missing_document_raises_404(tmp_path, monkeypatch):
|
||||
"""Completely absent document_id returns 404 with canonical detail."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
document_for_subject_or_404("nonexistent-id", "alice")
|
||||
assert exc_info.value.status_code == 404
|
||||
assert exc_info.value.detail == "Document not found"
|
||||
|
||||
|
||||
def test_empty_document_id_raises_404(tmp_path, monkeypatch):
|
||||
"""Empty string document_id raises 404 rather than hitting the DB."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
document_for_subject_or_404("", "alice")
|
||||
assert exc_info.value.status_code == 404
|
||||
|
||||
|
||||
def test_empty_subject_raises_404(tmp_path, monkeypatch):
|
||||
"""Empty subject raises 404 — cannot authorize without a subject."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id, owner="alice")
|
||||
_insert_kb_doc(conn, doc_id, kb_id)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
document_for_subject_or_404(doc_id, "")
|
||||
assert exc_info.value.status_code == 404
|
||||
|
||||
|
||||
# ── chunk_belongs_to_document ─────────────────────────────────────────
|
||||
|
||||
|
||||
def test_chunk_belongs_returns_true_for_matching_doc(tmp_path, monkeypatch):
|
||||
"""chunk_belongs_to_document returns True when chunk.document_id matches."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id, owner="alice")
|
||||
_insert_kb_doc(conn, doc_id, kb_id)
|
||||
_insert_chunk(conn, chunk_id, doc_id)
|
||||
assert chunk_belongs_to_document(chunk_id, doc_id) is True
|
||||
|
||||
|
||||
def test_chunk_belongs_returns_false_for_wrong_doc(tmp_path, monkeypatch):
|
||||
"""chunk_belongs_to_document returns False when chunk belongs to a different document."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
kb_id = _uid()
|
||||
doc_a, doc_b, chunk_id = _uid(), _uid(), _uid()
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id, owner="alice")
|
||||
_insert_kb_doc(conn, doc_a, kb_id, "a.pdf")
|
||||
_insert_kb_doc(conn, doc_b, kb_id, "b.pdf")
|
||||
_insert_chunk(conn, chunk_id, doc_a)
|
||||
# chunk belongs to doc_a — probing with doc_b must return False
|
||||
assert chunk_belongs_to_document(chunk_id, doc_b) is False
|
||||
|
||||
|
||||
def test_chunk_belongs_returns_false_for_missing_chunk(tmp_path, monkeypatch):
|
||||
"""chunk_belongs_to_document returns False for a nonexistent chunk_id."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id, owner="alice")
|
||||
_insert_kb_doc(conn, doc_id, kb_id)
|
||||
assert chunk_belongs_to_document("ghost-chunk-id", doc_id) is False
|
||||
|
||||
|
||||
def test_chunk_belongs_returns_false_for_empty_inputs(tmp_path, monkeypatch):
|
||||
"""chunk_belongs_to_document returns False for empty inputs without DB access."""
|
||||
_reset_db(tmp_path, monkeypatch)
|
||||
assert chunk_belongs_to_document("", "some-doc") is False
|
||||
assert chunk_belongs_to_document("some-chunk", "") is False
|
||||
assert chunk_belongs_to_document("", "") is False
|
||||
279
studio/backend/tests/test_rag_chunk_locators.py
Normal file
279
studio/backend/tests/test_rag_chunk_locators.py
Normal file
|
|
@ -0,0 +1,279 @@
|
|||
# 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 queue as queue_module
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
import storage.studio_db as studio_db
|
||||
from core.rag.chunking import chunk_pages, chunk_pages_with_spans
|
||||
from core.rag.ingestion import (
|
||||
_JobState,
|
||||
_insert_chunks_and_collect_for_bm25,
|
||||
_pump,
|
||||
_replace_document_pages,
|
||||
)
|
||||
from core.rag.parsers import ParsedPage
|
||||
|
||||
|
||||
def _uid() -> str:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
def _token_count(text: str) -> int:
|
||||
return max(1, len(text.split()))
|
||||
|
||||
|
||||
def test_standard_chunking_records_page_local_char_and_line_spans():
|
||||
pages = [
|
||||
ParsedPage(
|
||||
text="alpha first line\nbeta target line\ngamma final line",
|
||||
page_number=7,
|
||||
)
|
||||
]
|
||||
|
||||
chunks = chunk_pages(
|
||||
pages,
|
||||
max_tokens=3,
|
||||
overlap_tokens=0,
|
||||
token_counter=_token_count,
|
||||
separators=("\n", " ", ""),
|
||||
)
|
||||
|
||||
target = next(chunk for chunk in chunks if "beta" in chunk.text)
|
||||
assert target.page_number == 7
|
||||
assert target.source_page_index == 0
|
||||
assert target.page_char_start == pages[0].text.index("beta target line")
|
||||
assert target.page_char_end == target.page_char_start + len("beta target line")
|
||||
assert target.line_start == 2
|
||||
assert target.line_end == 2
|
||||
|
||||
|
||||
def test_late_chunking_maps_global_span_back_to_source_page():
|
||||
pages = [
|
||||
ParsedPage(text="page one alpha", page_number=1),
|
||||
ParsedPage(text="page two beta target", page_number=2),
|
||||
]
|
||||
|
||||
_full_doc, chunks, spans = chunk_pages_with_spans(
|
||||
pages,
|
||||
max_tokens=4,
|
||||
overlap_tokens=0,
|
||||
token_counter=_token_count,
|
||||
separators=("\n\n", " ", ""),
|
||||
)
|
||||
|
||||
target = next(chunk for chunk in chunks if "beta" in chunk.text)
|
||||
assert spans[chunks.index(target)][0] >= len(pages[0].text)
|
||||
assert target.page_number == 2
|
||||
assert target.source_page_index == 1
|
||||
assert target.page_char_start is not None
|
||||
assert target.page_char_end is not None
|
||||
assert pages[1].text[target.page_char_start : target.page_char_end].strip()
|
||||
|
||||
|
||||
def test_image_chunk_persistence_keeps_page_focus_and_null_text_locators(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
):
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
|
||||
monkeypatch.setattr(studio_db, "_schema_ready", False)
|
||||
|
||||
from core.rag import ingestion
|
||||
|
||||
captured_points: list[dict] = []
|
||||
monkeypatch.setattr(
|
||||
ingestion.vector_store,
|
||||
"upsert_chunks",
|
||||
lambda _scope, points: captured_points.extend(points),
|
||||
)
|
||||
|
||||
kb_id = _uid()
|
||||
doc_id = _uid()
|
||||
with studio_db.get_connection() as conn:
|
||||
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, "image.pdf", "application/pdf", "image.pdf", 1_700_000_001),
|
||||
)
|
||||
|
||||
_insert_chunks_and_collect_for_bm25(
|
||||
doc_id,
|
||||
"kb_scope",
|
||||
0,
|
||||
[
|
||||
{
|
||||
"text": "",
|
||||
"token_count": 0,
|
||||
"page_number": 3,
|
||||
"kind": "image",
|
||||
"image_path": str(tmp_path / "img.png"),
|
||||
}
|
||||
],
|
||||
[[0.1, 0.2]],
|
||||
)
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT page_number, source_page_index, page_char_start,
|
||||
page_char_end, line_start, line_end
|
||||
FROM rag_chunks WHERE document_id = ?
|
||||
""",
|
||||
(doc_id,),
|
||||
).fetchone()
|
||||
assert row["page_number"] == 3
|
||||
assert row["source_page_index"] is None
|
||||
assert row["page_char_start"] is None
|
||||
assert row["page_char_end"] is None
|
||||
assert row["line_start"] is None
|
||||
assert row["line_end"] is None
|
||||
assert captured_points[0]["payload"]["page_number"] == 3
|
||||
assert captured_points[0]["payload"]["page_char_start"] is None
|
||||
|
||||
|
||||
def test_replace_document_pages_replaces_existing_rows(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
|
||||
monkeypatch.setattr(studio_db, "_schema_ready", False)
|
||||
|
||||
kb_id = _uid()
|
||||
doc_id = _uid()
|
||||
with studio_db.get_connection() as conn:
|
||||
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, "doc.pdf", "application/pdf", "doc.pdf", 1_700_000_001),
|
||||
)
|
||||
|
||||
_replace_document_pages(
|
||||
doc_id,
|
||||
[
|
||||
{
|
||||
"page_index": 0,
|
||||
"page_number": 1,
|
||||
"text": "old page",
|
||||
"char_count": 8,
|
||||
"line_count": 1,
|
||||
}
|
||||
],
|
||||
)
|
||||
_replace_document_pages(
|
||||
doc_id,
|
||||
[
|
||||
{
|
||||
"page_index": 1,
|
||||
"page_number": 2,
|
||||
"text": "new\npage",
|
||||
"char_count": 8,
|
||||
"line_count": 2,
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT page_index, page_number, text, char_count, line_count
|
||||
FROM rag_document_pages WHERE document_id = ?
|
||||
""",
|
||||
(doc_id,),
|
||||
).fetchall()
|
||||
|
||||
assert [dict(row) for row in rows] == [
|
||||
{
|
||||
"page_index": 1,
|
||||
"page_number": 2,
|
||||
"text": "new\npage",
|
||||
"char_count": 8,
|
||||
"line_count": 2,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
class _OneMessageQueue:
|
||||
def __init__(self, message: dict) -> None:
|
||||
self.message = message
|
||||
self.used = False
|
||||
|
||||
def get(self, timeout: float) -> dict:
|
||||
if self.used:
|
||||
raise queue_module.Empty
|
||||
self.used = True
|
||||
return self.message
|
||||
|
||||
|
||||
class _FinishedProcess:
|
||||
def join(self, timeout: float | None = None) -> None:
|
||||
return None
|
||||
|
||||
def is_alive(self) -> bool:
|
||||
return False
|
||||
|
||||
def terminate(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pages",
|
||||
[
|
||||
[
|
||||
{
|
||||
"page_index": 0,
|
||||
"page_number": 1,
|
||||
"text": "orphan page",
|
||||
"char_count": 11,
|
||||
"line_count": 1,
|
||||
}
|
||||
],
|
||||
[],
|
||||
],
|
||||
)
|
||||
def test_document_pages_missing_document_fails_pump_cleanly(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
pages,
|
||||
):
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
|
||||
monkeypatch.setattr(studio_db, "_schema_ready", False)
|
||||
with studio_db.get_connection():
|
||||
pass
|
||||
|
||||
state = _JobState("job-missing-doc", "missing-doc", "kb_scope")
|
||||
queue = _OneMessageQueue(
|
||||
{
|
||||
"type": "document_pages",
|
||||
"pages": pages,
|
||||
}
|
||||
)
|
||||
|
||||
_pump(state, _FinishedProcess(), queue)
|
||||
|
||||
assert state.status == "failed"
|
||||
assert state.error is not None
|
||||
assert "document was removed before ingestion finished" in state.error
|
||||
140
studio/backend/tests/test_rag_locator_backfill.py
Normal file
140
studio/backend/tests/test_rag_locator_backfill.py
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
# 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
|
||||
87
studio/backend/tests/test_rag_locator_migration.py
Normal file
87
studio/backend/tests/test_rag_locator_migration.py
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
# 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,
|
||||
}
|
||||
593
studio/backend/tests/test_rag_preview_routes.py
Normal file
593
studio/backend/tests/test_rag_preview_routes.py
Normal file
|
|
@ -0,0 +1,593 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tests for GET /api/rag/documents/{id}/preview-target and /file.
|
||||
|
||||
Acceptance criteria covered (contracts.md §1, §2, PLAN.md T1/T2, Risk #1-3):
|
||||
|
||||
/preview-target:
|
||||
- 200 with chunk data when chunk_id present and belongs to doc.
|
||||
- 200 with all-null chunk fields when chunk_id absent (document-row preview).
|
||||
- 404 when document missing (collapsed existence + auth).
|
||||
- 404 when chunk_id does not belong to document_id (cross-doc probe collapsed).
|
||||
- 401 when no bearer token.
|
||||
|
||||
/file:
|
||||
- 200 with correct Content-Type and nosniff header.
|
||||
- X-Content-Type-Options: nosniff present on every 200.
|
||||
- Cache-Control: private present on every 200.
|
||||
- HTML extension served as text/plain + attachment (Risk #3).
|
||||
- DOCX extension served with attachment disposition.
|
||||
- 404 when document missing or wrong subject.
|
||||
- 404 when file deleted from disk (DB row exists, subject authorized).
|
||||
- Outside-root stored_path returns 404 (path containment, Risk #2).
|
||||
|
||||
Auth is injected via dependency override (mock at the boundary, not the
|
||||
implementation target). We do NOT mock document_for_subject_or_404 itself.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
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
|
||||
|
||||
|
||||
# ── App import (deferred to avoid import-time side-effects) ───────────
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
# ── Test-level DB + auth fixtures ─────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def db_env(tmp_path, monkeypatch):
|
||||
"""Point studio_db at a fresh temp DB for each test."""
|
||||
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"):
|
||||
"""Return a TestClient with get_current_subject overridden to return subject."""
|
||||
app.dependency_overrides[get_current_subject] = lambda: subject
|
||||
client = TestClient(app, raise_server_exceptions=True)
|
||||
return client
|
||||
|
||||
|
||||
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 = "report.pdf",
|
||||
content_type: str | None = "application/pdf",
|
||||
status: str = "completed",
|
||||
) -> 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, ?, ?, ?, ?, 0, 1024, ?)",
|
||||
(doc_id, kb_id, filename, content_type, stored_path, status, 1_700_000_000),
|
||||
)
|
||||
|
||||
|
||||
def _insert_chunk(
|
||||
conn,
|
||||
chunk_id: str,
|
||||
doc_id: str,
|
||||
text: str = "The margin rose to 18.2% in Q3.",
|
||||
page_number: int | None = 7,
|
||||
chunk_index: int = 14,
|
||||
) -> None:
|
||||
conn.execute(
|
||||
"INSERT INTO rag_chunks "
|
||||
"(id, document_id, chunk_index, text, token_count, page_number) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?)",
|
||||
(chunk_id, doc_id, chunk_index, text, 30, page_number),
|
||||
)
|
||||
|
||||
|
||||
# ── /preview-target tests ─────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestPreviewTarget:
|
||||
def test_with_chunk_id_returns_full_metadata(self, app, db_env, monkeypatch):
|
||||
"""GET /preview-target?chunk_id=<id> returns page + snippet when chunk valid."""
|
||||
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
|
||||
stored = db_env / "rag" / "uploads" / "report.pdf"
|
||||
stored.parent.mkdir(parents=True, exist_ok=True)
|
||||
stored.write_bytes(b"%PDF-1.4 dummy")
|
||||
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))
|
||||
_insert_chunk(conn, chunk_id, doc_id, page_number=7, chunk_index=14)
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
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["chunkId"] == chunk_id
|
||||
assert body["targetPage"] == 7
|
||||
assert body["chunkIndex"] == 14
|
||||
assert body["snippet"] is not None and len(body["snippet"]) > 0
|
||||
assert body["mediaKind"] == "pdf"
|
||||
|
||||
def test_preview_target_returns_pdf_regions_when_present(self, app, db_env, monkeypatch):
|
||||
"""Chunk preview includes only stored confident PDF regions."""
|
||||
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
|
||||
stored = db_env / "rag" / "uploads" / "report.pdf"
|
||||
stored.parent.mkdir(parents=True, exist_ok=True)
|
||||
stored.write_bytes(b"%PDF-1.4 dummy")
|
||||
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))
|
||||
_insert_chunk(conn, chunk_id, doc_id, page_number=7, chunk_index=14)
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE rag_chunks
|
||||
SET pdf_regions_json = ?
|
||||
WHERE id = ?
|
||||
""",
|
||||
(
|
||||
'[{"pageIndex":6,"pageNumber":7,"x":0.1,"y":0.2,'
|
||||
'"width":0.3,"height":0.04,"confidence":"exact",'
|
||||
'"source":"pymupdf-search"}]',
|
||||
chunk_id,
|
||||
),
|
||||
)
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
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["pdfRegions"] == [
|
||||
{
|
||||
"pageIndex": 6,
|
||||
"pageNumber": 7,
|
||||
"x": 0.1,
|
||||
"y": 0.2,
|
||||
"width": 0.3,
|
||||
"height": 0.04,
|
||||
"confidence": "exact",
|
||||
"source": "pymupdf-search",
|
||||
}
|
||||
]
|
||||
|
||||
def test_without_chunk_id_returns_all_null_chunk_fields(self, app, db_env, monkeypatch):
|
||||
"""GET /preview-target without chunk_id returns metadata-only (decision Q2)."""
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
stored = db_env / "rag" / "uploads" / "annual.pdf"
|
||||
stored.parent.mkdir(parents=True, exist_ok=True)
|
||||
stored.write_bytes(b"%PDF-1.4 dummy")
|
||||
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))
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{doc_id}/preview-target")
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
# All chunk fields MUST be null — UI must not guess a first chunk.
|
||||
assert body["chunkId"] is None
|
||||
assert body["chunkIndex"] is None
|
||||
assert body["targetPage"] is None
|
||||
assert body["snippet"] is None
|
||||
assert body["kind"] is None
|
||||
assert body["imageUrl"] is None
|
||||
assert body["documentId"] == doc_id
|
||||
|
||||
def test_missing_document_returns_404(self, app, db_env, monkeypatch):
|
||||
"""Nonexistent document_id returns 404 to both existence and auth probes."""
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{_uid()}/preview-target")
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
assert resp.status_code == 404
|
||||
assert resp.json()["detail"] == "Document not found"
|
||||
|
||||
def test_wrong_subject_returns_404(self, app, db_env, monkeypatch):
|
||||
"""Document owned by alice returns 404 when accessed by mallory."""
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
stored = db_env / "rag" / "uploads" / "secret.pdf"
|
||||
stored.parent.mkdir(parents=True, exist_ok=True)
|
||||
stored.write_bytes(b"data")
|
||||
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))
|
||||
|
||||
client = _make_client(app, "mallory")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{doc_id}/preview-target")
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_cross_doc_chunk_id_returns_404(self, app, db_env, monkeypatch):
|
||||
"""chunk_id from a different document returns 404 — not 400 (opaque)."""
|
||||
kb_id = _uid()
|
||||
doc_a, doc_b = _uid(), _uid()
|
||||
chunk_a = _uid()
|
||||
stored_a = db_env / "rag" / "uploads" / "a.pdf"
|
||||
stored_b = db_env / "rag" / "uploads" / "b.pdf"
|
||||
stored_a.parent.mkdir(parents=True, exist_ok=True)
|
||||
stored_a.write_bytes(b"data")
|
||||
stored_b.write_bytes(b"data")
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id)
|
||||
_insert_doc(conn, doc_a, kb_id, str(stored_a), "a.pdf")
|
||||
_insert_doc(conn, doc_b, kb_id, str(stored_b), "b.pdf")
|
||||
_insert_chunk(conn, chunk_a, doc_a)
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
# Probe doc_b with chunk_a (which belongs to doc_a)
|
||||
resp = client.get(
|
||||
f"/api/rag/documents/{doc_b}/preview-target?chunk_id={chunk_a}"
|
||||
)
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
# Must be 404, NOT 200 with doc_a's chunk data
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_unauthenticated_returns_401(self, app, db_env, monkeypatch):
|
||||
"""No bearer token → 401."""
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
# No override — let the real dependency raise
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get(f"/api/rag/documents/{_uid()}/preview-target")
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
# ── /file tests ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestFileRoute:
|
||||
def test_pdf_200_with_correct_headers(self, app, db_env, monkeypatch):
|
||||
"""GET /file for a PDF returns 200 with nosniff, Cache-Control, inline disposition."""
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
uploads = db_env / "rag" / "uploads"
|
||||
uploads.mkdir(parents=True, exist_ok=True)
|
||||
stored = uploads / "annual.pdf"
|
||||
stored.write_bytes(b"%PDF-1.4\n%%EOF")
|
||||
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), "annual.pdf", "application/pdf")
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{doc_id}/file")
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
|
||||
assert resp.status_code == 200
|
||||
assert resp.headers.get("x-content-type-options") == "nosniff"
|
||||
assert "private" in (resp.headers.get("cache-control") or "")
|
||||
ct = resp.headers.get("content-type", "")
|
||||
assert "pdf" in ct.lower()
|
||||
|
||||
def test_signed_file_url_supports_range_without_bearer_query(self, app, db_env, monkeypatch):
|
||||
"""Short-lived signed URL is redeemable without Authorization and supports ranges."""
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
uploads = db_env / "rag" / "uploads"
|
||||
uploads.mkdir(parents=True, exist_ok=True)
|
||||
stored = uploads / "annual.pdf"
|
||||
stored.write_bytes(b"%PDF-1.4\n%%EOF")
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
monkeypatch.setattr("routes.rag.get_jwt_secret", lambda subject: "test-secret")
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id)
|
||||
_insert_doc(conn, doc_id, kb_id, str(stored), "annual.pdf", "application/pdf")
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
url_resp = client.get(f"/api/rag/documents/{doc_id}/file-url")
|
||||
assert url_resp.status_code == 200
|
||||
signed_url = url_resp.json()["url"]
|
||||
assert "Bearer" not in signed_url
|
||||
assert "Authorization" not in signed_url
|
||||
|
||||
file_resp = client.get(signed_url, headers={"Range": "bytes=0-3"})
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
|
||||
assert file_resp.status_code == 206
|
||||
assert file_resp.content == b"%PDF"
|
||||
assert file_resp.headers.get("content-range") == f"bytes 0-3/{stored.stat().st_size}"
|
||||
assert file_resp.headers.get("accept-ranges") == "bytes"
|
||||
assert file_resp.headers.get("x-content-type-options") == "nosniff"
|
||||
|
||||
def test_signed_file_route_rejects_forged_token(self, app, db_env, monkeypatch):
|
||||
"""Signed file route is not public without a valid preview token."""
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
uploads = db_env / "rag" / "uploads"
|
||||
uploads.mkdir(parents=True, exist_ok=True)
|
||||
stored = uploads / "annual.pdf"
|
||||
stored.write_bytes(b"%PDF-1.4\n%%EOF")
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
monkeypatch.setattr("routes.rag.get_jwt_secret", lambda subject: "test-secret")
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id)
|
||||
_insert_doc(conn, doc_id, kb_id, str(stored), "annual.pdf", "application/pdf")
|
||||
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get(f"/api/rag/documents/{doc_id}/file-signed?token=bogus")
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_html_file_served_as_text_plain_with_attachment(self, app, db_env, monkeypatch):
|
||||
"""HTML uploads must be served as text/plain + attachment (Risk #3 — no XSS)."""
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
uploads = db_env / "rag" / "uploads"
|
||||
uploads.mkdir(parents=True, exist_ok=True)
|
||||
stored = uploads / "malicious.html"
|
||||
stored.write_bytes(b"<script>alert(1)</script>")
|
||||
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), "malicious.html", "text/html")
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{doc_id}/file")
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
|
||||
assert resp.status_code == 200
|
||||
ct = resp.headers.get("content-type", "").lower()
|
||||
# MUST NOT be text/html — must be text/plain
|
||||
assert "text/html" not in ct, f"HTML executed inline! content-type={ct}"
|
||||
assert "text/plain" in ct
|
||||
disp = resp.headers.get("content-disposition", "").lower()
|
||||
assert "attachment" in disp, f"HTML not forced to attachment: {disp}"
|
||||
assert resp.headers.get("x-content-type-options") == "nosniff"
|
||||
|
||||
def test_docx_served_as_attachment(self, app, db_env, monkeypatch):
|
||||
"""DOCX files must be served with Content-Disposition: attachment."""
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
uploads = db_env / "rag" / "uploads"
|
||||
uploads.mkdir(parents=True, exist_ok=True)
|
||||
stored = uploads / "report.docx"
|
||||
stored.write_bytes(b"PK\x03\x04fake-docx")
|
||||
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), "report.docx",
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document")
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{doc_id}/file")
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
|
||||
assert resp.status_code == 200
|
||||
disp = resp.headers.get("content-disposition", "").lower()
|
||||
assert "attachment" in disp
|
||||
|
||||
def test_missing_document_returns_404(self, app, db_env, monkeypatch):
|
||||
"""Nonexistent document returns 404."""
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{_uid()}/file")
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_wrong_subject_returns_404(self, app, db_env, monkeypatch):
|
||||
"""Document accessible to alice is 404 for mallory (auth-collapse)."""
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
uploads = db_env / "rag" / "uploads"
|
||||
uploads.mkdir(parents=True, exist_ok=True)
|
||||
stored = uploads / "private.pdf"
|
||||
stored.write_bytes(b"data")
|
||||
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))
|
||||
|
||||
client = _make_client(app, "mallory")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{doc_id}/file")
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_deleted_file_returns_404_with_doc_file_not_found(self, app, db_env, monkeypatch):
|
||||
"""File gone from disk returns 404 with 'Document file not found' detail."""
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
uploads = db_env / "rag" / "uploads"
|
||||
uploads.mkdir(parents=True, exist_ok=True)
|
||||
stored = uploads / "gone.pdf"
|
||||
stored.write_bytes(b"data")
|
||||
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))
|
||||
|
||||
# Delete the file after inserting the row
|
||||
stored.unlink()
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{doc_id}/file")
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
|
||||
assert resp.status_code == 404
|
||||
detail = resp.json().get("detail", "")
|
||||
assert "file not found" in detail.lower() or "not found" in detail.lower()
|
||||
|
||||
def test_outside_root_stored_path_returns_404(self, app, db_env, monkeypatch, tmp_path):
|
||||
"""stored_path outside rag_uploads_root returns 404 — path containment (Risk #2)."""
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
uploads = db_env / "rag" / "uploads"
|
||||
uploads.mkdir(parents=True, exist_ok=True)
|
||||
# A legitimate-looking path that is outside the RAG uploads root
|
||||
outside = tmp_path / "etc" / "passwd"
|
||||
outside.parent.mkdir(parents=True, exist_ok=True)
|
||||
outside.write_bytes(b"root:x:0:0")
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id)
|
||||
# Insert with stored_path pointing outside root
|
||||
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, 'passwd', 'text/plain', ?, 'completed', 0, 10, ?)",
|
||||
(doc_id, kb_id, str(outside), 1_700_000_000),
|
||||
)
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{doc_id}/file")
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
|
||||
# Must NOT serve the file — containment violation must return 404
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_nosniff_and_cache_headers_on_txt_file(self, app, db_env, monkeypatch):
|
||||
"""Safety headers present on every 200 response, including plain text."""
|
||||
doc_id, kb_id = _uid(), _uid()
|
||||
uploads = db_env / "rag" / "uploads"
|
||||
uploads.mkdir(parents=True, exist_ok=True)
|
||||
stored = uploads / "notes.txt"
|
||||
stored.write_bytes(b"hello world")
|
||||
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), "notes.txt", "text/plain")
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{doc_id}/file")
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
|
||||
assert resp.status_code == 200
|
||||
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")
|
||||
136
studio/backend/tests/test_rag_preview_target_locators.py
Normal file
136
studio/backend/tests/test_rag_preview_target_locators.py
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
# 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 _seed_doc(conn, doc_id: str, kb_id: str, stored_path: str) -> None:
|
||||
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, "report.pdf", "application/pdf", stored_path, 1_700_000_001),
|
||||
)
|
||||
|
||||
|
||||
def test_preview_target_returns_nullable_locator_fields(app, db_env):
|
||||
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
|
||||
stored = db_env / "rag" / "uploads" / "report.pdf"
|
||||
stored.parent.mkdir(parents=True, exist_ok=True)
|
||||
stored.write_bytes(b"%PDF-1.4")
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
_seed_doc(conn, doc_id, kb_id, str(stored))
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO rag_chunks
|
||||
(id, document_id, chunk_index, text, token_count, page_number,
|
||||
source_page_index, page_char_start, page_char_end, line_start,
|
||||
line_end)
|
||||
VALUES (?, ?, 2, ?, 8, 4, 3, 20, 52, 6, 7)
|
||||
""",
|
||||
(chunk_id, doc_id, "highlight me"),
|
||||
)
|
||||
|
||||
client = _make_client(app)
|
||||
try:
|
||||
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["sourcePageIndex"] == 3
|
||||
assert body["pageCharStart"] == 20
|
||||
assert body["pageCharEnd"] == 52
|
||||
assert body["lineStart"] == 6
|
||||
assert body["lineEnd"] == 7
|
||||
|
||||
|
||||
def test_preview_target_old_null_locator_rows_still_work(app, db_env):
|
||||
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
|
||||
stored = db_env / "rag" / "uploads" / "legacy.pdf"
|
||||
stored.parent.mkdir(parents=True, exist_ok=True)
|
||||
stored.write_bytes(b"%PDF-1.4")
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
_seed_doc(conn, doc_id, kb_id, str(stored))
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO rag_chunks
|
||||
(id, document_id, chunk_index, text, token_count, page_number)
|
||||
VALUES (?, ?, 0, ?, 4, 1)
|
||||
""",
|
||||
(chunk_id, doc_id, "legacy"),
|
||||
)
|
||||
|
||||
client = _make_client(app)
|
||||
try:
|
||||
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["snippet"] == "legacy"
|
||||
assert body["sourcePageIndex"] is None
|
||||
assert body["pageCharStart"] is None
|
||||
assert body["pageCharEnd"] is None
|
||||
assert body["lineStart"] is None
|
||||
assert body["lineEnd"] is None
|
||||
243
studio/backend/tests/test_rag_source_identity.py
Normal file
243
studio/backend/tests/test_rag_source_identity.py
Normal file
|
|
@ -0,0 +1,243 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tests for durable source identity in tool XML output (PLAN.md T3, contracts §3).
|
||||
|
||||
Acceptance criteria:
|
||||
|
||||
- _format_hits_for_llm emits document_id and chunk_id attributes on <chunk> elements.
|
||||
- The visible citation id (id="N") is a per-call counter, NOT the backend chunk UUID.
|
||||
- Same-filename documents in different KB slots remain distinguishable by document_id.
|
||||
- Hits without a matching DB row (lookup miss) are silently dropped — not emitted
|
||||
with wrong IDs.
|
||||
- Legacy hits (no document_id in hit dict) still render without crashing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import uuid
|
||||
from xml.etree import ElementTree
|
||||
|
||||
import pytest
|
||||
|
||||
from core.rag.tool import _format_hits_for_llm
|
||||
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _uid() -> str:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
def _hit(
|
||||
*,
|
||||
chunk_id: str,
|
||||
document_id: str,
|
||||
filename: str = "report.pdf",
|
||||
text: str = "some text",
|
||||
page_number: int | None = 3,
|
||||
chunk_index: int = 0,
|
||||
score: float = 0.85,
|
||||
) -> dict:
|
||||
"""Build a flat hit dict as _format_hits_for_llm expects."""
|
||||
return {
|
||||
"chunk_id": chunk_id,
|
||||
"document_id": document_id,
|
||||
"filename": filename,
|
||||
"text": text,
|
||||
"page_number": page_number,
|
||||
"chunk_index": chunk_index,
|
||||
"score": score,
|
||||
"dense_score": score,
|
||||
"token_count": 20,
|
||||
"kind": "text",
|
||||
"image_path": None,
|
||||
}
|
||||
|
||||
|
||||
def _parse_chunks(xml_output: str) -> list[dict]:
|
||||
"""Parse <chunk ...> elements from the multi-block tool output."""
|
||||
chunks = []
|
||||
# Each block is wrapped in <chunk ...>\n...\n</chunk>; parse directly.
|
||||
for match in re.finditer(r"<chunk\s([^>]*)>", xml_output):
|
||||
attrs_raw = match.group(1)
|
||||
# Quick attribute parser for "key="value"" pairs.
|
||||
attrs: dict = {}
|
||||
for m in re.finditer(r'(\w+)="([^"]*)"', attrs_raw):
|
||||
attrs[m.group(1)] = m.group(2)
|
||||
chunks.append(attrs)
|
||||
return chunks
|
||||
|
||||
|
||||
# ── Tests: durable IDs present in XML ────────────────────────────────
|
||||
|
||||
|
||||
def test_format_hits_emits_document_id_and_chunk_id():
|
||||
"""T3: tool XML <chunk> must carry document_id and chunk_id attributes."""
|
||||
chunk_id, doc_id = _uid(), _uid()
|
||||
hits = [_hit(chunk_id=chunk_id, document_id=doc_id)]
|
||||
output = _format_hits_for_llm(hits)
|
||||
chunks = _parse_chunks(output)
|
||||
assert len(chunks) == 1, output
|
||||
assert chunks[0]["document_id"] == doc_id
|
||||
assert chunks[0]["chunk_id"] == chunk_id
|
||||
|
||||
|
||||
def test_citation_id_is_sequential_counter_not_uuid():
|
||||
"""Visible id='N' is a 1-based counter — never equal to the backend chunk UUID."""
|
||||
chunk_id, doc_id = _uid(), _uid()
|
||||
hits = [_hit(chunk_id=chunk_id, document_id=doc_id)]
|
||||
output = _format_hits_for_llm(hits, start_id=0)
|
||||
chunks = _parse_chunks(output)
|
||||
visible_id = chunks[0]["id"]
|
||||
# Must be a small integer string, NOT the UUID
|
||||
assert visible_id == "1", f"expected '1' got {visible_id!r}"
|
||||
assert visible_id != chunk_id
|
||||
|
||||
|
||||
def test_citation_ids_are_globally_sequential_across_calls():
|
||||
"""start_id offset ensures IDs stay unique across multiple tool calls per turn."""
|
||||
hits_call1 = [_hit(chunk_id=_uid(), document_id=_uid(), filename="a.pdf")]
|
||||
hits_call2 = [
|
||||
_hit(chunk_id=_uid(), document_id=_uid(), filename="b.pdf"),
|
||||
_hit(chunk_id=_uid(), document_id=_uid(), filename="c.pdf"),
|
||||
]
|
||||
out1 = _format_hits_for_llm(hits_call1, start_id=0)
|
||||
out2 = _format_hits_for_llm(hits_call2, start_id=1)
|
||||
|
||||
chunks1 = _parse_chunks(out1)
|
||||
chunks2 = _parse_chunks(out2)
|
||||
|
||||
assert chunks1[0]["id"] == "1"
|
||||
assert chunks2[0]["id"] == "2"
|
||||
assert chunks2[1]["id"] == "3"
|
||||
|
||||
# No id overlap
|
||||
all_ids = {c["id"] for c in chunks1 + chunks2}
|
||||
assert len(all_ids) == 3
|
||||
|
||||
|
||||
def test_same_filename_docs_have_distinct_document_ids():
|
||||
"""Two docs with the same filename route to distinct document_id values (Risk #4)."""
|
||||
filename = "annual-report.pdf"
|
||||
chunk_a, doc_a = _uid(), _uid()
|
||||
chunk_b, doc_b = _uid(), _uid()
|
||||
hits = [
|
||||
_hit(chunk_id=chunk_a, document_id=doc_a, filename=filename),
|
||||
_hit(chunk_id=chunk_b, document_id=doc_b, filename=filename),
|
||||
]
|
||||
output = _format_hits_for_llm(hits)
|
||||
chunks = _parse_chunks(output)
|
||||
assert len(chunks) == 2
|
||||
# Both use the same filename but MUST have distinct document_id values
|
||||
assert chunks[0]["document_id"] != chunks[1]["document_id"]
|
||||
assert chunks[0]["document_id"] == doc_a
|
||||
assert chunks[1]["document_id"] == doc_b
|
||||
|
||||
|
||||
def test_same_filename_docs_have_distinct_citation_ids():
|
||||
"""Same-filename docs in the same turn still get distinct visible [N] ids."""
|
||||
filename = "notes.pdf"
|
||||
chunk_a, doc_a = _uid(), _uid()
|
||||
chunk_b, doc_b = _uid(), _uid()
|
||||
hits = [
|
||||
_hit(chunk_id=chunk_a, document_id=doc_a, filename=filename),
|
||||
_hit(chunk_id=chunk_b, document_id=doc_b, filename=filename),
|
||||
]
|
||||
output = _format_hits_for_llm(hits)
|
||||
chunks = _parse_chunks(output)
|
||||
citation_ids = {c["id"] for c in chunks}
|
||||
assert len(citation_ids) == 2, f"citation IDs not unique: {chunks}"
|
||||
|
||||
|
||||
def test_empty_hits_returns_no_chunks_message():
|
||||
"""Empty hit list returns the 'no matching chunks' message, not broken XML."""
|
||||
output = _format_hits_for_llm([])
|
||||
chunks = _parse_chunks(output)
|
||||
assert len(chunks) == 0
|
||||
assert "no matching chunks" in output.lower() or "no matching" in output.lower()
|
||||
|
||||
|
||||
def test_page_number_attribute_present_when_page_exists():
|
||||
"""page attribute is emitted when page_number is not None."""
|
||||
chunk_id, doc_id = _uid(), _uid()
|
||||
hits = [_hit(chunk_id=chunk_id, document_id=doc_id, page_number=5)]
|
||||
output = _format_hits_for_llm(hits)
|
||||
chunks = _parse_chunks(output)
|
||||
assert chunks[0].get("page") == "5"
|
||||
|
||||
|
||||
def test_page_number_attribute_absent_when_null():
|
||||
"""page attribute is omitted when page_number is None."""
|
||||
chunk_id, doc_id = _uid(), _uid()
|
||||
hits = [_hit(chunk_id=chunk_id, document_id=doc_id, page_number=None)]
|
||||
output = _format_hits_for_llm(hits)
|
||||
chunks = _parse_chunks(output)
|
||||
assert "page" not in chunks[0], f"unexpected page attr: {chunks[0]}"
|
||||
|
||||
|
||||
def test_locator_attributes_are_additive_when_present():
|
||||
"""T10: tool XML carries nullable locator metadata without changing visible ids."""
|
||||
chunk_id, doc_id = _uid(), _uid()
|
||||
hit = _hit(chunk_id=chunk_id, document_id=doc_id, page_number=5)
|
||||
hit.update(
|
||||
{
|
||||
"source_page_index": 4,
|
||||
"page_char_start": 11,
|
||||
"page_char_end": 42,
|
||||
"line_start": 2,
|
||||
"line_end": 3,
|
||||
}
|
||||
)
|
||||
output = _format_hits_for_llm([hit])
|
||||
chunk = _parse_chunks(output)[0]
|
||||
assert chunk["id"] == "1"
|
||||
assert chunk["chunk_id"] == chunk_id
|
||||
assert chunk["source_page_index"] == "4"
|
||||
assert chunk["page_char_start"] == "11"
|
||||
assert chunk["page_char_end"] == "42"
|
||||
assert chunk["line_start"] == "2"
|
||||
assert chunk["line_end"] == "3"
|
||||
|
||||
|
||||
def test_xml_special_chars_in_filename_escaped():
|
||||
"""Filename with XML special chars does not break the chunk element."""
|
||||
chunk_id, doc_id = _uid(), _uid()
|
||||
hits = [
|
||||
_hit(
|
||||
chunk_id=chunk_id,
|
||||
document_id=doc_id,
|
||||
filename='report <2025> "final" & draft.pdf',
|
||||
)
|
||||
]
|
||||
output = _format_hits_for_llm(hits)
|
||||
# The output must parse cleanly (no unescaped < or " in attrs)
|
||||
chunks = _parse_chunks(output)
|
||||
assert len(chunks) == 1
|
||||
# source attribute should have the filename escaped
|
||||
source_attr = chunks[0].get("source", "")
|
||||
assert "<" not in source_attr and '"' not in source_attr
|
||||
|
||||
|
||||
def test_multiple_hits_carry_independent_ids():
|
||||
"""Three hits each carry their own distinct chunk_id and document_id."""
|
||||
hit_data = [
|
||||
(_uid(), _uid()),
|
||||
(_uid(), _uid()),
|
||||
(_uid(), _uid()),
|
||||
]
|
||||
hits = [
|
||||
_hit(chunk_id=cid, document_id=did, filename=f"doc{i}.pdf")
|
||||
for i, (cid, did) in enumerate(hit_data)
|
||||
]
|
||||
output = _format_hits_for_llm(hits)
|
||||
chunks = _parse_chunks(output)
|
||||
assert len(chunks) == 3
|
||||
emitted_chunk_ids = {c["chunk_id"] for c in chunks}
|
||||
emitted_doc_ids = {c["document_id"] for c in chunks}
|
||||
expected_chunk_ids = {cid for cid, _ in hit_data}
|
||||
expected_doc_ids = {did for _, did in hit_data}
|
||||
assert emitted_chunk_ids == expected_chunk_ids
|
||||
assert emitted_doc_ids == expected_doc_ids
|
||||
|
|
@ -38,7 +38,7 @@
|
|||
},
|
||||
"overrides": [
|
||||
{
|
||||
"include": ["vite.config.ts", "eslint.config.js"],
|
||||
"include": ["vite.config.ts", "vitest.config.ts", "eslint.config.js"],
|
||||
"linter": {
|
||||
"rules": {
|
||||
"correctness": { "noNodejsModules": "off" },
|
||||
|
|
|
|||
1192
studio/frontend/package-lock.json
generated
1192
studio/frontend/package-lock.json
generated
File diff suppressed because it is too large
Load diff
|
|
@ -12,6 +12,8 @@
|
|||
"lint": "eslint .",
|
||||
"preview": "vite preview",
|
||||
"typecheck": "tsc -b --pretty false",
|
||||
"test": "vitest run",
|
||||
"test:watch": "vitest",
|
||||
"biome:check": "biome check .",
|
||||
"biome:fix": "biome check . --write"
|
||||
},
|
||||
|
|
@ -46,6 +48,7 @@
|
|||
"@tauri-apps/plugin-process": "^2.3.1",
|
||||
"@tauri-apps/plugin-updater": "^2.10.1",
|
||||
"@toolwind/corner-shape": "^0.0.8-3",
|
||||
"@types/event-source-polyfill": "1.0.5",
|
||||
"@xyflow/react": "^12.10.0",
|
||||
"assistant-stream": "0.3.12",
|
||||
"canvas-confetti": "^1.9.4",
|
||||
|
|
@ -53,6 +56,7 @@
|
|||
"clsx": "^2.1.1",
|
||||
"cmdk": "^1.1.1",
|
||||
"dexie": "^4.3.0",
|
||||
"event-source-polyfill": "1.0.31",
|
||||
"fflate": "0.8.3",
|
||||
"js-yaml": "^4.1.1",
|
||||
"katex": "^0.16.28",
|
||||
|
|
@ -65,6 +69,7 @@
|
|||
"react": "^19.2.4",
|
||||
"react-day-picker": "^9.13.2",
|
||||
"react-dom": "^19.2.4",
|
||||
"react-pdf": "^10.4.1",
|
||||
"react-resizable-panels": "^4.6.4",
|
||||
"recharts": "3.7.0",
|
||||
"shadcn": "^4.2.0",
|
||||
|
|
@ -85,10 +90,14 @@
|
|||
"devDependencies": {
|
||||
"@biomejs/biome": "^1.9.4",
|
||||
"@eslint/js": "^9.39.1",
|
||||
"@testing-library/dom": "^10.4.1",
|
||||
"@testing-library/jest-dom": "^6.9.1",
|
||||
"@testing-library/react": "^16.3.2",
|
||||
"@testing-library/user-event": "^14.6.1",
|
||||
"@types/canvas-confetti": "^1.9.0",
|
||||
"@types/js-yaml": "^4.0.9",
|
||||
"@types/node-forge": "^1.3.14",
|
||||
"@types/node": "^25.5.2",
|
||||
"@types/node-forge": "^1.3.14",
|
||||
"@types/react": "^19.2.5",
|
||||
"@types/react-dom": "^19.2.3",
|
||||
"@vitejs/plugin-react": "^6.0.1",
|
||||
|
|
@ -96,8 +105,10 @@
|
|||
"eslint-plugin-react-hooks": "^7.0.1",
|
||||
"eslint-plugin-react-refresh": "^0.5.2",
|
||||
"globals": "^17.4.0",
|
||||
"jsdom": "^29.1.1",
|
||||
"typescript": "~5.9.3",
|
||||
"typescript-eslint": "^8.55.0",
|
||||
"vite": "^8.0.1"
|
||||
"vite": "^8.0.1",
|
||||
"vitest": "^4.1.7"
|
||||
}
|
||||
}
|
||||
|
|
|
|||
125
studio/frontend/src/__tests__/chat-adapter.test.ts
Normal file
125
studio/frontend/src/__tests__/chat-adapter.test.ts
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
/**
|
||||
* Tests for chat-adapter XML parsing — contracts §3 / §4, T3.
|
||||
*
|
||||
* Coverage:
|
||||
* - New XML with document_id + chunk_id attributes → citationId, documentId, backendChunkId populated.
|
||||
* - Legacy XML without durable IDs → citationId populated, documentId/backendChunkId absent.
|
||||
* - Same visible [N] across turns does NOT imply same backendChunkId.
|
||||
* - Same filename in two chunks → distinct documentId values preserved.
|
||||
* - Missing attributes degrade gracefully — no throw.
|
||||
* - citationId is always the visible "N" counter, never the UUID.
|
||||
*/
|
||||
|
||||
import {
|
||||
type ParsedChunk,
|
||||
parseChunks,
|
||||
} from "@/components/assistant-ui/tool-ui-search-knowledge-base";
|
||||
import { describe, expect, it } from "vitest";
|
||||
|
||||
// ── Tests ─────────────────────────────────────────────────────────────
|
||||
|
||||
describe("parseChunks — durable IDs (contracts §3/§4)", () => {
|
||||
it("new XML with document_id + chunk_id populates all three identity fields", () => {
|
||||
const xml = `
|
||||
<chunk id="1" source="report.pdf" page="7" document_id="doc-abc" chunk_id="chunk-xyz">
|
||||
The margin rose to 18%.
|
||||
</chunk>
|
||||
`.trim();
|
||||
|
||||
const parts: ParsedChunk[] = parseChunks(xml);
|
||||
expect(parts).toHaveLength(1);
|
||||
expect(parts[0].id).toBe("1");
|
||||
expect(parts[0].documentId).toBe("doc-abc");
|
||||
expect(parts[0].backendChunkId).toBe("chunk-xyz");
|
||||
expect(parts[0].source).toBe("report.pdf");
|
||||
expect(parts[0].page).toBe("7");
|
||||
});
|
||||
|
||||
it("legacy XML without durable IDs leaves documentId and backendChunkId absent", () => {
|
||||
// Old XML: no document_id, no chunk_id — hover-only, NOT preview-clickable (Q3).
|
||||
const xml = `
|
||||
<chunk id="2" source="old-doc.pdf" page="3">
|
||||
Legacy chunk text.
|
||||
</chunk>
|
||||
`.trim();
|
||||
|
||||
const parts: ParsedChunk[] = parseChunks(xml);
|
||||
expect(parts).toHaveLength(1);
|
||||
expect(parts[0].id).toBe("2");
|
||||
expect(parts[0].documentId).toBeUndefined();
|
||||
expect(parts[0].backendChunkId).toBeUndefined();
|
||||
});
|
||||
|
||||
it("citationId (id) is the visible counter string, never the backend UUID", () => {
|
||||
const docId = "550e8400-e29b-41d4-a716-446655440000";
|
||||
const chunkId = "6ba7b810-9dad-11d1-80b4-00c04fd430c8";
|
||||
const xml = `<chunk id="5" source="paper.pdf" document_id="${docId}" chunk_id="${chunkId}">text</chunk>`;
|
||||
const parts: ParsedChunk[] = parseChunks(xml);
|
||||
expect(parts[0].id).toBe("5");
|
||||
expect(parts[0].id).not.toBe(docId);
|
||||
expect(parts[0].id).not.toBe(chunkId);
|
||||
});
|
||||
|
||||
it("same filename in two chunks preserves distinct documentId values", () => {
|
||||
const xml = `
|
||||
<chunk id="1" source="annual.pdf" document_id="doc-001" chunk_id="chunk-001">First excerpt.</chunk>
|
||||
<chunk id="2" source="annual.pdf" document_id="doc-002" chunk_id="chunk-002">Second excerpt.</chunk>
|
||||
`.trim();
|
||||
|
||||
const parts: ParsedChunk[] = parseChunks(xml);
|
||||
expect(parts).toHaveLength(2);
|
||||
expect(parts[0].documentId).toBe("doc-001");
|
||||
expect(parts[1].documentId).toBe("doc-002");
|
||||
expect(parts[0].documentId).not.toBe(parts[1].documentId);
|
||||
});
|
||||
|
||||
it("same visible id in different turns does not imply same backendChunkId", () => {
|
||||
// Turn 1 and turn 2 both have id="1" but different backend identities.
|
||||
const turn1 = `<chunk id="1" source="a.pdf" document_id="doc-A" chunk_id="chunk-A">Turn 1.</chunk>`;
|
||||
const turn2 = `<chunk id="1" source="b.pdf" document_id="doc-B" chunk_id="chunk-B">Turn 2.</chunk>`;
|
||||
|
||||
const p1: ParsedChunk[] = parseChunks(turn1);
|
||||
const p2: ParsedChunk[] = parseChunks(turn2);
|
||||
expect(p1[0].id).toBe(p2[0].id); // both "1"
|
||||
expect(p1[0].backendChunkId).not.toBe(p2[0].backendChunkId);
|
||||
expect(p1[0].documentId).not.toBe(p2[0].documentId);
|
||||
});
|
||||
|
||||
it("missing chunk_id only (partial durable attrs) → backendChunkId absent", () => {
|
||||
const xml = `<chunk id="3" source="x.pdf" document_id="doc-XYZ">text</chunk>`;
|
||||
const parts: ParsedChunk[] = parseChunks(xml);
|
||||
expect(parts[0].documentId).toBe("doc-XYZ");
|
||||
expect(parts[0].backendChunkId).toBeUndefined();
|
||||
});
|
||||
|
||||
it("multiple new-format chunks all carry independent IDs", () => {
|
||||
const xml = `
|
||||
<chunk id="1" source="a.pdf" document_id="doc-1" chunk_id="ck-1">A</chunk>
|
||||
<chunk id="2" source="b.pdf" document_id="doc-2" chunk_id="ck-2">B</chunk>
|
||||
<chunk id="3" source="c.pdf" document_id="doc-3" chunk_id="ck-3">C</chunk>
|
||||
`.trim();
|
||||
|
||||
const parts: ParsedChunk[] = parseChunks(xml);
|
||||
expect(parts).toHaveLength(3);
|
||||
const docIds = new Set(parts.map((p) => p.documentId));
|
||||
const backendChunkIds = new Set(parts.map((p) => p.backendChunkId));
|
||||
const citationIds = new Set(parts.map((p) => p.id));
|
||||
expect(docIds.size).toBe(3);
|
||||
expect(backendChunkIds.size).toBe(3);
|
||||
expect(citationIds.size).toBe(3);
|
||||
});
|
||||
|
||||
it("empty XML returns empty array without throwing", () => {
|
||||
expect(parseChunks("")).toHaveLength(0);
|
||||
expect(parseChunks("No chunks here.")).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("XML entity encoding in source attribute is decoded", () => {
|
||||
// & should decode to & in the source attribute (decodeXml in parseChunks)
|
||||
const xml = `<chunk id="1" source="report & summary.pdf" document_id="doc-1" chunk_id="ck-1">text</chunk>`;
|
||||
const parts: ParsedChunk[] = parseChunks(xml);
|
||||
expect(parts).toHaveLength(1);
|
||||
expect(parts[0].id).toBe("1");
|
||||
expect(parts[0].source).toBe("report & summary.pdf");
|
||||
});
|
||||
});
|
||||
92
studio/frontend/src/__tests__/document-row.test.tsx
Normal file
92
studio/frontend/src/__tests__/document-row.test.tsx
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
import type { RagDocument } from "@/features/rag/api/rag-api";
|
||||
import { DocumentRow } from "@/features/rag/components/document-row";
|
||||
import { fireEvent, render, screen } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import React from "react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
function makeDoc(overrides: Partial<RagDocument> = {}): RagDocument {
|
||||
return {
|
||||
id: "doc-abc",
|
||||
kb_id: "kb-1",
|
||||
thread_id: null,
|
||||
filename: "report.pdf",
|
||||
content_type: "application/pdf",
|
||||
status: "completed",
|
||||
num_chunks: 5,
|
||||
byte_size: 10240,
|
||||
error: null,
|
||||
created_at: 1_700_000_000,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
let onPreview: ReturnType<typeof vi.fn>;
|
||||
let onDelete: ReturnType<typeof vi.fn>;
|
||||
|
||||
beforeEach(() => {
|
||||
onPreview = vi.fn();
|
||||
onDelete = vi.fn();
|
||||
});
|
||||
|
||||
describe("DocumentRow preview event propagation", () => {
|
||||
it("clicking a previewable row opens document-level preview", async () => {
|
||||
render(
|
||||
React.createElement(DocumentRow, {
|
||||
doc: makeDoc(),
|
||||
onPreview: onPreview as () => void,
|
||||
onDelete: onDelete as () => void,
|
||||
}),
|
||||
);
|
||||
|
||||
await userEvent.click(
|
||||
screen.getByRole("button", { name: /open preview of report.pdf/i }),
|
||||
);
|
||||
|
||||
expect(onPreview).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("Enter and Space open a previewable row", () => {
|
||||
render(
|
||||
React.createElement(DocumentRow, {
|
||||
doc: makeDoc(),
|
||||
onPreview: onPreview as () => void,
|
||||
onDelete: onDelete as () => void,
|
||||
}),
|
||||
);
|
||||
const row = screen.getByRole("button", {
|
||||
name: /open preview of report.pdf/i,
|
||||
});
|
||||
|
||||
fireEvent.keyDown(row, { key: "Enter" });
|
||||
fireEvent.keyDown(row, { key: " " });
|
||||
|
||||
expect(onPreview).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("clicking delete does not open preview", async () => {
|
||||
render(
|
||||
React.createElement(DocumentRow, {
|
||||
doc: makeDoc(),
|
||||
onPreview: onPreview as () => void,
|
||||
onDelete: onDelete as () => void,
|
||||
}),
|
||||
);
|
||||
|
||||
await userEvent.click(screen.getByRole("button", { name: /delete/i }));
|
||||
|
||||
expect(onDelete).toHaveBeenCalledTimes(1);
|
||||
expect(onPreview).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("non-previewable rows have no row button semantics", () => {
|
||||
render(
|
||||
React.createElement(DocumentRow, {
|
||||
doc: makeDoc({ status: "pending" }),
|
||||
onDelete: onDelete as () => void,
|
||||
}),
|
||||
);
|
||||
|
||||
expect(screen.queryByRole("button", { name: /open preview/i })).toBeNull();
|
||||
});
|
||||
});
|
||||
77
studio/frontend/src/__tests__/knowledge-bases-tab.test.tsx
Normal file
77
studio/frontend/src/__tests__/knowledge-bases-tab.test.tsx
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
import { KnowledgeBasesTab } from "@/features/settings/tabs/knowledge-bases-tab";
|
||||
import { render, screen } from "@testing-library/react";
|
||||
import React from "react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
const { mockUsePreviewStore } = vi.hoisted(() => {
|
||||
const state = {
|
||||
target: null as unknown,
|
||||
status: "idle",
|
||||
close: vi.fn(),
|
||||
};
|
||||
const fn = vi.fn((selector?: (s: typeof state) => unknown) => {
|
||||
if (typeof selector === "function") return selector(state);
|
||||
return state;
|
||||
}) as ReturnType<typeof vi.fn> & { __state: typeof state };
|
||||
fn.__state = state;
|
||||
return { mockUsePreviewStore: fn };
|
||||
});
|
||||
|
||||
vi.mock("@/features/rag/stores/preview-store", () => ({
|
||||
usePreviewStore: mockUsePreviewStore,
|
||||
}));
|
||||
|
||||
vi.mock("@/features/rag/components/kb-list", () => ({
|
||||
KBList: () => React.createElement("div", { "data-testid": "kb-list" }),
|
||||
}));
|
||||
|
||||
vi.mock("@/features/rag/components/kb-detail-panel", () => ({
|
||||
KBDetailPanel: () =>
|
||||
React.createElement("div", { "data-testid": "kb-detail-panel" }),
|
||||
}));
|
||||
|
||||
vi.mock("@/features/rag/components/preview-panel", () => ({
|
||||
PreviewPanel: ({ open }: { open: boolean }) =>
|
||||
React.createElement("div", {
|
||||
"data-testid": "settings-preview-panel",
|
||||
"data-open": String(open),
|
||||
}),
|
||||
}));
|
||||
|
||||
vi.mock("@/features/rag/components/thread-index-list", () => ({
|
||||
ThreadIndexList: () =>
|
||||
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";
|
||||
mockUsePreviewStore.mockImplementation(
|
||||
(selector?: (s: typeof mockUsePreviewStore.__state) => unknown) => {
|
||||
if (typeof selector === "function") {
|
||||
return selector(mockUsePreviewStore.__state);
|
||||
}
|
||||
return mockUsePreviewStore.__state;
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
describe("KnowledgeBasesTab preview host", () => {
|
||||
it("renders a preview panel when the preview store is active", () => {
|
||||
mockUsePreviewStore.__state.target = {
|
||||
documentId: "doc-abc",
|
||||
filename: "report.pdf",
|
||||
};
|
||||
mockUsePreviewStore.__state.status = "ready";
|
||||
|
||||
render(React.createElement(KnowledgeBasesTab));
|
||||
|
||||
const panel = screen.getByTestId("settings-preview-panel");
|
||||
expect(panel.getAttribute("data-open")).toBe("true");
|
||||
});
|
||||
});
|
||||
93
studio/frontend/src/__tests__/preview-a11y.test.tsx
Normal file
93
studio/frontend/src/__tests__/preview-a11y.test.tsx
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
import type { PreviewTarget } from "@/features/rag/api/rag-api";
|
||||
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
|
||||
import React from "react";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
const { mockFetchPreviewTarget, mockFetchPreviewFileBlob } = vi.hoisted(() => ({
|
||||
mockFetchPreviewTarget:
|
||||
vi.fn<
|
||||
(documentId: string, chunkId?: string | null) => Promise<PreviewTarget>
|
||||
>(),
|
||||
mockFetchPreviewFileBlob:
|
||||
vi.fn<(documentId: string, signal?: AbortSignal) => Promise<Blob>>(),
|
||||
}));
|
||||
|
||||
vi.mock("@/features/rag/api/rag-api", async (importOriginal) => {
|
||||
const original =
|
||||
await importOriginal<typeof import("@/features/rag/api/rag-api")>();
|
||||
return {
|
||||
...original,
|
||||
fetchPreviewTarget: mockFetchPreviewTarget,
|
||||
fetchPreviewFileBlob: mockFetchPreviewFileBlob,
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock("react-pdf", () => ({
|
||||
Document: ({ children }: { children: React.ReactNode }) =>
|
||||
React.createElement("div", { "data-testid": "pdf-document" }, children),
|
||||
Page: () => React.createElement("div", { "data-testid": "pdf-page" }),
|
||||
pdfjs: { GlobalWorkerOptions: { workerSrc: "" } },
|
||||
}));
|
||||
|
||||
import { PreviewPanel } from "@/features/rag/components/preview-panel";
|
||||
import { usePreviewStore } from "@/features/rag/stores/preview-store";
|
||||
|
||||
function target(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
|
||||
return {
|
||||
documentId: "doc-abc",
|
||||
filename: "report.html",
|
||||
contentType: "text/plain",
|
||||
mediaKind: "html",
|
||||
byteSize: 100,
|
||||
status: "completed",
|
||||
kbId: "kb-1",
|
||||
threadId: null,
|
||||
chunkId: "chunk-1",
|
||||
chunkIndex: 0,
|
||||
targetPage: 1,
|
||||
snippet: "safe extracted text",
|
||||
kind: "text",
|
||||
imageUrl: null,
|
||||
sourcePageIndex: null,
|
||||
pageCharStart: null,
|
||||
pageCharEnd: null,
|
||||
lineStart: null,
|
||||
lineEnd: null,
|
||||
pdfRegions: [],
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
mockFetchPreviewTarget.mockReset();
|
||||
mockFetchPreviewFileBlob.mockReset();
|
||||
usePreviewStore.getState().close();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
usePreviewStore.getState().close();
|
||||
});
|
||||
|
||||
describe("preview a11y hardening", () => {
|
||||
it("Escape closes the preview and restores focus to the opener", async () => {
|
||||
const opener = document.createElement("button");
|
||||
opener.textContent = "Open preview";
|
||||
document.body.appendChild(opener);
|
||||
opener.focus();
|
||||
mockFetchPreviewTarget.mockResolvedValue(target());
|
||||
|
||||
await usePreviewStore.getState().open({ documentId: "doc-abc" });
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
expect(
|
||||
screen.getByRole("region", { name: /document preview/i }),
|
||||
).toBeInTheDocument();
|
||||
|
||||
fireEvent.keyDown(document, { key: "Escape" });
|
||||
|
||||
await waitFor(() => {
|
||||
expect(usePreviewStore.getState().status).toBe("idle");
|
||||
});
|
||||
expect(document.activeElement).toBe(opener);
|
||||
opener.remove();
|
||||
});
|
||||
});
|
||||
521
studio/frontend/src/__tests__/preview-panel.test.tsx
Normal file
521
studio/frontend/src/__tests__/preview-panel.test.tsx
Normal file
|
|
@ -0,0 +1,521 @@
|
|||
/**
|
||||
* Tests for preview-panel — HTML/DOCX/unknown must NEVER render inline (T5 / Risk #3).
|
||||
*
|
||||
* Acceptance criteria (contracts §5.4, PLAN.md T5, decisions Q7):
|
||||
* - mediaKind === "pdf" → react-pdf view is mounted (or loading indicator shown).
|
||||
* - mediaKind === "html" → text-view fallback shown, NO object/embed/iframe with blob URL.
|
||||
* - mediaKind === "docx" → text-view fallback shown, NO inline rendering.
|
||||
* - mediaKind === "unknown" → unavailable/download state, NOT inline.
|
||||
* - mediaKind === "text" → text/snippet view shown.
|
||||
* - Panel without a target renders nothing or unavailable state.
|
||||
*/
|
||||
|
||||
import {
|
||||
type PreviewMediaKind,
|
||||
type PreviewTarget,
|
||||
} from "@/features/rag/api/rag-api";
|
||||
import type { PreviewLoadStatus } from "@/features/rag/stores/preview-store";
|
||||
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
|
||||
import React from "react";
|
||||
import {
|
||||
type MockInstance,
|
||||
afterEach,
|
||||
beforeEach,
|
||||
describe,
|
||||
expect,
|
||||
it,
|
||||
vi,
|
||||
} from "vitest";
|
||||
|
||||
const DOWNLOAD_BUTTON_NAME = /download/i;
|
||||
const LONG_CONTENT_TEXT = /Some long content/;
|
||||
|
||||
// ── Mock preview store ────────────────────────────────────────────────
|
||||
// The real component uses per-field selectors: usePreviewStore((s) => s.target)
|
||||
// so the mock must handle the selector pattern.
|
||||
// vi.hoisted ensures the mock fn is initialised before vi.mock factory runs.
|
||||
|
||||
interface MockStoreState {
|
||||
target: PreviewTarget | null;
|
||||
previewBlobUrl: string | null;
|
||||
previewBlob: Blob | null;
|
||||
previewFileUrl: string | null;
|
||||
previewFileUrlExpiresAt: number | null;
|
||||
status: PreviewLoadStatus;
|
||||
error: string | null;
|
||||
close: () => void;
|
||||
open: () => void;
|
||||
}
|
||||
|
||||
let mockState: MockStoreState = {
|
||||
target: null,
|
||||
previewBlobUrl: null,
|
||||
previewBlob: null,
|
||||
previewFileUrl: null,
|
||||
previewFileUrlExpiresAt: null,
|
||||
status: "idle",
|
||||
error: null,
|
||||
close: vi.fn(),
|
||||
open: vi.fn(),
|
||||
};
|
||||
|
||||
const { mockAuthFetch, mockUsePreviewStore } = vi.hoisted(() => {
|
||||
// usePreviewStore is called two ways:
|
||||
// usePreviewStore((s) => s.field) — selector form (React hook)
|
||||
// usePreviewStore.getState().close() — outside React (cleanup effect)
|
||||
const fn = vi.fn((selector?: (s: MockStoreState) => unknown) => {
|
||||
if (typeof selector === "function") {
|
||||
return selector(mockState);
|
||||
}
|
||||
return mockState;
|
||||
}) as ReturnType<typeof vi.fn> & { getState: () => MockStoreState };
|
||||
fn.getState = () => mockState;
|
||||
return { mockAuthFetch: vi.fn(), mockUsePreviewStore: fn };
|
||||
});
|
||||
|
||||
vi.mock("@/features/auth", () => ({
|
||||
authFetch: mockAuthFetch,
|
||||
getAuthToken: () => "mock-token-123",
|
||||
}));
|
||||
|
||||
vi.mock("@/features/rag/stores/preview-store", async (importOriginal) => {
|
||||
const real =
|
||||
await importOriginal<
|
||||
typeof import("@/features/rag/stores/preview-store")
|
||||
>();
|
||||
return {
|
||||
...real,
|
||||
usePreviewStore: mockUsePreviewStore,
|
||||
// isInlineBlobAllowed passes through from the real module so assertions
|
||||
// use the production allowlist, not a test-local copy (D1.5 fix).
|
||||
};
|
||||
});
|
||||
|
||||
// react-pdf requires a browser worker URL that doesn't exist in jsdom.
|
||||
vi.mock("react-pdf", () => ({
|
||||
Document: ({ children }: { children: React.ReactNode }) =>
|
||||
React.createElement("div", { "data-testid": "pdf-document" }, children),
|
||||
Page: () => React.createElement("div", { "data-testid": "pdf-page" }),
|
||||
pdfjs: { GlobalWorkerOptions: { workerSrc: "" } },
|
||||
}));
|
||||
|
||||
beforeEach(() => {
|
||||
mockAuthFetch.mockReset();
|
||||
mockAuthFetch.mockResolvedValue(
|
||||
new Response(new Blob(["download bytes"], { type: "text/plain" }), {
|
||||
status: 200,
|
||||
}),
|
||||
);
|
||||
mockState = {
|
||||
target: null,
|
||||
previewBlobUrl: null,
|
||||
previewBlob: null,
|
||||
previewFileUrl: null,
|
||||
previewFileUrlExpiresAt: null,
|
||||
status: "idle",
|
||||
error: null,
|
||||
close: vi.fn(),
|
||||
open: vi.fn(),
|
||||
};
|
||||
mockUsePreviewStore.mockImplementation(
|
||||
(selector?: (s: MockStoreState) => unknown) => {
|
||||
if (typeof selector === "function") {
|
||||
return selector(mockState);
|
||||
}
|
||||
return mockState;
|
||||
},
|
||||
);
|
||||
// Restore getState after mockImplementation replaces the fn internals
|
||||
mockUsePreviewStore.getState = () => mockState;
|
||||
|
||||
// Mock window.matchMedia globally for tests
|
||||
window.matchMedia = vi.fn().mockImplementation((query) => ({
|
||||
matches: false,
|
||||
media: query,
|
||||
onchange: null,
|
||||
addListener: vi.fn(),
|
||||
removeListener: vi.fn(),
|
||||
addEventListener: vi.fn(),
|
||||
removeEventListener: vi.fn(),
|
||||
dispatchEvent: vi.fn(),
|
||||
}));
|
||||
});
|
||||
|
||||
// ── Import panel + production allowlist predicate AFTER mocks ─────────
|
||||
|
||||
import { PreviewPanel } from "@/features/rag/components/preview-panel";
|
||||
import { isInlineBlobAllowed } from "@/features/rag/stores/preview-store";
|
||||
|
||||
// ── Helpers ───────────────────────────────────────────────────────────
|
||||
|
||||
function makeTarget(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
|
||||
return {
|
||||
documentId: "doc-abc",
|
||||
filename: "report.pdf",
|
||||
contentType: "application/pdf",
|
||||
mediaKind: "pdf",
|
||||
byteSize: 100,
|
||||
status: "completed",
|
||||
kbId: "kb-1",
|
||||
threadId: null,
|
||||
chunkId: null,
|
||||
chunkIndex: null,
|
||||
targetPage: null,
|
||||
snippet: null,
|
||||
kind: null,
|
||||
imageUrl: null,
|
||||
sourcePageIndex: null,
|
||||
pageCharStart: null,
|
||||
pageCharEnd: null,
|
||||
lineStart: null,
|
||||
lineEnd: null,
|
||||
pdfRegions: [],
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function setReady(overrides: Partial<PreviewTarget> = {}): void {
|
||||
mockState.target = makeTarget(overrides);
|
||||
mockState.status = "ready";
|
||||
}
|
||||
|
||||
// ── Tests against real panel component ───────────────────────────────
|
||||
|
||||
describe("preview-panel inline rendering safety (contracts §5.4 / Risk #3)", () => {
|
||||
it("panel with no target (idle) renders without crashing", () => {
|
||||
const { container } = render(
|
||||
React.createElement(PreviewPanel, { open: true }),
|
||||
);
|
||||
expect(container).toBeDefined();
|
||||
expect(document.querySelector("iframe")).toBeNull();
|
||||
});
|
||||
|
||||
it("html mediaKind does not render an iframe, object, or embed element", () => {
|
||||
setReady({ mediaKind: "html", filename: "malicious.html" });
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
expect(document.querySelector("iframe")).toBeNull();
|
||||
expect(document.querySelector("object")).toBeNull();
|
||||
expect(document.querySelector("embed")).toBeNull();
|
||||
});
|
||||
|
||||
it("docx mediaKind does not render an iframe, object, or embed element", () => {
|
||||
setReady({ mediaKind: "docx", filename: "report.docx" });
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
expect(document.querySelector("iframe")).toBeNull();
|
||||
expect(document.querySelector("object")).toBeNull();
|
||||
expect(document.querySelector("embed")).toBeNull();
|
||||
});
|
||||
|
||||
it("unknown mediaKind does not render inline blob content", () => {
|
||||
setReady({ mediaKind: "unknown", filename: "data.bin" });
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
expect(document.querySelector("iframe")).toBeNull();
|
||||
expect(document.querySelector("object")).toBeNull();
|
||||
expect(document.querySelector("embed")).toBeNull();
|
||||
});
|
||||
|
||||
it.each<PreviewMediaKind>(["html", "docx", "unknown"])(
|
||||
"%s download creates only a download object URL, never inline preview content",
|
||||
async (mediaKind) => {
|
||||
const createObjectUrl = vi
|
||||
.spyOn(URL, "createObjectURL")
|
||||
.mockReturnValue("blob:unsafe");
|
||||
const revokeObjectUrl = vi
|
||||
.spyOn(URL, "revokeObjectURL")
|
||||
.mockImplementation(() => undefined);
|
||||
setReady({
|
||||
mediaKind,
|
||||
filename: `unsafe.${mediaKind}`,
|
||||
contentType: "text/plain",
|
||||
snippet: "Extracted text only.",
|
||||
});
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
fireEvent.click(
|
||||
screen.getByRole("button", { name: DOWNLOAD_BUTTON_NAME }),
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockAuthFetch).toHaveBeenCalledWith(
|
||||
"/api/rag/documents/doc-abc/file",
|
||||
);
|
||||
});
|
||||
expect(createObjectUrl).toHaveBeenCalledWith(expect.any(Blob));
|
||||
await waitFor(() => {
|
||||
expect(revokeObjectUrl).toHaveBeenCalledWith("blob:unsafe");
|
||||
});
|
||||
expect(document.querySelector("iframe")).toBeNull();
|
||||
expect(document.querySelector("object")).toBeNull();
|
||||
expect(document.querySelector("embed")).toBeNull();
|
||||
|
||||
createObjectUrl.mockRestore();
|
||||
revokeObjectUrl.mockRestore();
|
||||
},
|
||||
);
|
||||
|
||||
it("text mediaKind renders without iframe (text fallback path)", () => {
|
||||
setReady({
|
||||
mediaKind: "text",
|
||||
filename: "notes.txt",
|
||||
snippet: "This is the extracted text content.",
|
||||
});
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
expect(document.querySelector("iframe")).toBeNull();
|
||||
});
|
||||
|
||||
it("open=false triggers close side-effect on the store", () => {
|
||||
const closeFn = vi.fn();
|
||||
mockState.close = closeFn;
|
||||
mockState.target = makeTarget();
|
||||
mockState.status = "ready";
|
||||
|
||||
const { rerender } = render(
|
||||
React.createElement(PreviewPanel, { open: true }),
|
||||
);
|
||||
rerender(React.createElement(PreviewPanel, { open: false }));
|
||||
|
||||
// The useEffect for open=false should have called close()
|
||||
expect(closeFn).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("closes the preview panel on Escape key down (Escape key closures)", () => {
|
||||
const closeFn = vi.fn();
|
||||
mockState.close = closeFn;
|
||||
mockState.target = makeTarget();
|
||||
mockState.status = "ready";
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
fireEvent.keyDown(document, { key: "Escape" });
|
||||
expect(closeFn).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("renders with premium glassmorphic visual details and a pulsing green indicator dot", () => {
|
||||
setReady({ mediaKind: "text", filename: "notes.txt" });
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
const section = screen.getByLabelText("Document preview");
|
||||
expect(section).toHaveClass("bg-panel-surface/85");
|
||||
expect(section).toHaveClass("backdrop-blur-lg");
|
||||
expect(section).toHaveClass("border-border/40");
|
||||
expect(section).toHaveClass("shadow-lg");
|
||||
|
||||
// Pulser dot
|
||||
const pulser = section.querySelector(".animate-pulse");
|
||||
expect(pulser).toBeInTheDocument();
|
||||
expect(pulser).toHaveClass("bg-primary");
|
||||
expect(pulser).toHaveClass("w-2");
|
||||
expect(pulser).toHaveClass("h-2");
|
||||
});
|
||||
|
||||
it("shifts the layout to a full mobile Sheet drawer overlay when the viewport is squeezed (< 1024px)", () => {
|
||||
window.matchMedia = vi.fn().mockImplementation((query) => ({
|
||||
matches: true,
|
||||
media: query,
|
||||
onchange: null,
|
||||
addListener: vi.fn(),
|
||||
removeListener: vi.fn(),
|
||||
addEventListener: vi.fn(),
|
||||
removeEventListener: vi.fn(),
|
||||
dispatchEvent: vi.fn(),
|
||||
}));
|
||||
|
||||
setReady({ mediaKind: "text", filename: "notes.txt" });
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
// Radix UI Sheet component should render dialog role in mobile viewports
|
||||
const dialog = screen.getByRole("dialog");
|
||||
expect(dialog).toBeInTheDocument();
|
||||
expect(dialog).toHaveClass("preview-sheet-content");
|
||||
expect(screen.getByText("Document preview")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
// ── Pure-logic: inline allowlist (always green) ───────────────────────
|
||||
// Uses the real production isInlineBlobAllowed (D1.5 fix: no local copy).
|
||||
|
||||
describe("inline object URL allowlist (contracts §5.4, pure logic)", () => {
|
||||
const inlineSafe: PreviewMediaKind[] = ["pdf", "text", "image"];
|
||||
const inlineUnsafe: PreviewMediaKind[] = ["html", "docx", "unknown"];
|
||||
|
||||
it.each(inlineSafe)("mediaKind=%s is inline-safe", (mk) => {
|
||||
expect(isInlineBlobAllowed(mk)).toBe(true);
|
||||
});
|
||||
|
||||
it.each(inlineUnsafe)("mediaKind=%s is NOT inline-safe (Risk #3)", (mk) => {
|
||||
expect(isInlineBlobAllowed(mk)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("preview-panel stable scrollbars, sheets, layouts, and downloads", () => {
|
||||
let createObjectUrl: MockInstance<typeof URL.createObjectURL>;
|
||||
let revokeObjectUrl: MockInstance<typeof URL.revokeObjectURL>;
|
||||
|
||||
beforeEach(() => {
|
||||
createObjectUrl = vi
|
||||
.spyOn(URL, "createObjectURL")
|
||||
.mockReturnValue("blob:safe-url");
|
||||
revokeObjectUrl = vi
|
||||
.spyOn(URL, "revokeObjectURL")
|
||||
.mockImplementation(() => undefined);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
createObjectUrl.mockRestore();
|
||||
revokeObjectUrl.mockRestore();
|
||||
});
|
||||
|
||||
it("asserts stable scrollbar style classes are present on panel content", () => {
|
||||
setReady({
|
||||
mediaKind: "text",
|
||||
filename: "notes.txt",
|
||||
snippet: "Some long content that requires scrolling ".repeat(20),
|
||||
});
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
// The snippet is rendered in a <pre> element. Check if it has overflow-auto
|
||||
const preElement = screen.getByText(LONG_CONTENT_TEXT);
|
||||
expect(preElement).toHaveClass("overflow-auto");
|
||||
expect(preElement).toHaveClass("flex-1");
|
||||
});
|
||||
|
||||
it("asserts non-nested sheets are rendered in squeezed viewports", () => {
|
||||
window.matchMedia = vi.fn().mockImplementation((query) => ({
|
||||
matches: true,
|
||||
media: query,
|
||||
onchange: null,
|
||||
addListener: vi.fn(),
|
||||
removeListener: vi.fn(),
|
||||
addEventListener: vi.fn(),
|
||||
removeEventListener: vi.fn(),
|
||||
dispatchEvent: vi.fn(),
|
||||
}));
|
||||
|
||||
setReady({ mediaKind: "text", filename: "notes.txt" });
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
const dialogs = screen.getAllByRole("dialog");
|
||||
expect(dialogs.length).toBe(1);
|
||||
expect(dialogs[0]).toHaveClass("preview-sheet-content");
|
||||
|
||||
const nestedDialogs = dialogs[0].querySelectorAll("[role='dialog']");
|
||||
expect(nestedDialogs.length).toBe(0);
|
||||
});
|
||||
|
||||
it("supports responsive collapses under different viewport widths", () => {
|
||||
const mockMatchMedia = vi.fn().mockImplementation((query) => ({
|
||||
matches: query.includes("max-width: 1023px"),
|
||||
media: query,
|
||||
onchange: null,
|
||||
addListener: vi.fn(),
|
||||
removeListener: vi.fn(),
|
||||
addEventListener: vi.fn(),
|
||||
removeEventListener: vi.fn(),
|
||||
dispatchEvent: vi.fn(),
|
||||
}));
|
||||
window.matchMedia = mockMatchMedia;
|
||||
|
||||
setReady({ mediaKind: "text", filename: "notes.txt" });
|
||||
|
||||
const { unmount } = render(
|
||||
React.createElement(PreviewPanel, { open: true }),
|
||||
);
|
||||
expect(screen.getByRole("dialog")).toBeInTheDocument();
|
||||
unmount();
|
||||
|
||||
window.matchMedia = vi.fn().mockImplementation((query) => ({
|
||||
matches: false,
|
||||
media: query,
|
||||
onchange: null,
|
||||
addListener: vi.fn(),
|
||||
removeListener: vi.fn(),
|
||||
addEventListener: vi.fn(),
|
||||
removeEventListener: vi.fn(),
|
||||
dispatchEvent: vi.fn(),
|
||||
}));
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
expect(screen.queryByRole("dialog")).toBeNull();
|
||||
expect(screen.getByLabelText("Document preview")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("asserts object URL download logic: uses URL.createObjectURL for safe types and Data URL for unsafe types", async () => {
|
||||
setReady({
|
||||
mediaKind: "text",
|
||||
filename: "notes.txt",
|
||||
contentType: "text/plain",
|
||||
snippet: "Text snippet.",
|
||||
});
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
fireEvent.click(screen.getByRole("button", { name: DOWNLOAD_BUTTON_NAME }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockAuthFetch).toHaveBeenCalledWith(
|
||||
"/api/rag/documents/doc-abc/file",
|
||||
);
|
||||
});
|
||||
|
||||
expect(createObjectUrl).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("PreviewTextView precise highlights matching", () => {
|
||||
it("highlights with character ranges", () => {
|
||||
setReady({
|
||||
mediaKind: "text",
|
||||
filename: "notes.txt",
|
||||
snippet: "Line 1: Hello World\nLine 2: Target Phrase\nLine 3: Goodbye",
|
||||
pageCharStart: 28,
|
||||
pageCharEnd: 41,
|
||||
});
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
const mark = screen.getByText("Target Phrase");
|
||||
expect(mark.tagName).toBe("MARK");
|
||||
expect(mark).toHaveClass("bg-primary/20", "ring-primary/60");
|
||||
});
|
||||
|
||||
it("highlights with line numbers", () => {
|
||||
setReady({
|
||||
mediaKind: "text",
|
||||
filename: "notes.txt",
|
||||
snippet: "Line one text\nLine two text\nLine three text",
|
||||
lineStart: 2,
|
||||
lineEnd: 2,
|
||||
});
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
const mark = screen.getByText("Line two text");
|
||||
expect(mark.tagName).toBe("MARK");
|
||||
expect(mark).toHaveClass("bg-primary/20", "ring-primary/60");
|
||||
});
|
||||
|
||||
it("highlights with fuzzy fallback matching high density line", () => {
|
||||
setReady({
|
||||
mediaKind: "text",
|
||||
filename: "notes.txt",
|
||||
snippet: "...\nAlphanumericDensity123456\n...",
|
||||
lineStart: 999, // Trigger hasLocator without matching any specific line range
|
||||
});
|
||||
|
||||
render(React.createElement(PreviewPanel, { open: true }));
|
||||
|
||||
const mark = screen.getByText("AlphanumericDensity123456");
|
||||
expect(mark.tagName).toBe("MARK");
|
||||
});
|
||||
});
|
||||
386
studio/frontend/src/__tests__/preview-pdf-smoke.test.tsx
Normal file
386
studio/frontend/src/__tests__/preview-pdf-smoke.test.tsx
Normal file
|
|
@ -0,0 +1,386 @@
|
|||
import type { PreviewTarget } from "@/features/rag/api/rag-api";
|
||||
import type { PreviewPdfRegion } from "@/features/rag/api/rag-api";
|
||||
import { PreviewPdfView } from "@/features/rag/components/preview-pdf-view";
|
||||
import {
|
||||
act,
|
||||
fireEvent,
|
||||
render,
|
||||
screen,
|
||||
waitFor,
|
||||
within,
|
||||
} from "@testing-library/react";
|
||||
import React from "react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
const SOURCE_EXCERPT_TEXT = /source excerpt/i;
|
||||
|
||||
/** The thumbnail rail also renders mocked `<Page>` elements, so every
|
||||
* test that targets the main render must scope through the
|
||||
* `pdf-main-page` wrapper instead of taking the first `pdf-page`. */
|
||||
async function findMainPdfPage(): Promise<HTMLElement> {
|
||||
const wrapper = await screen.findByTestId("pdf-main-page");
|
||||
return within(wrapper).getByTestId("pdf-page");
|
||||
}
|
||||
function getMainPdfPage(): HTMLElement {
|
||||
const wrapper = screen.getByTestId("pdf-main-page");
|
||||
return within(wrapper).getByTestId("pdf-page");
|
||||
}
|
||||
|
||||
vi.mock("react-pdf", () => ({
|
||||
Document: ({
|
||||
children,
|
||||
file,
|
||||
onLoadSuccess,
|
||||
}: {
|
||||
children: React.ReactNode;
|
||||
file?: unknown;
|
||||
onLoadSuccess?: (result: { numPages: number }) => void;
|
||||
}) => {
|
||||
onLoadSuccess?.({ numPages: 1 });
|
||||
return React.createElement(
|
||||
"div",
|
||||
{
|
||||
"data-testid": "pdf-document",
|
||||
"data-file-kind": file instanceof Blob ? "blob" : typeof file,
|
||||
"data-file-url":
|
||||
file && typeof file === "object" && "url" in file
|
||||
? String((file as { url: string }).url)
|
||||
: "",
|
||||
},
|
||||
children,
|
||||
);
|
||||
},
|
||||
Page: ({
|
||||
customTextRenderer,
|
||||
width,
|
||||
renderTextLayer,
|
||||
}: {
|
||||
customTextRenderer?: (item: { str: string }) => string;
|
||||
width?: number;
|
||||
renderTextLayer?: boolean;
|
||||
}) => {
|
||||
const html =
|
||||
customTextRenderer?.({ str: "target phrase" }) ?? "target phrase";
|
||||
return React.createElement("div", {
|
||||
"data-testid": "pdf-page",
|
||||
"data-width": String(width ?? ""),
|
||||
"data-render-text-layer": String(renderTextLayer),
|
||||
"data-rendered-html": html,
|
||||
});
|
||||
},
|
||||
pdfjs: { GlobalWorkerOptions: { workerSrc: "" } },
|
||||
}));
|
||||
|
||||
function target(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
|
||||
return {
|
||||
documentId: "doc-abc",
|
||||
filename: "report.pdf",
|
||||
contentType: "application/pdf",
|
||||
mediaKind: "pdf",
|
||||
byteSize: 100,
|
||||
status: "completed",
|
||||
kbId: "kb-1",
|
||||
threadId: null,
|
||||
chunkId: "chunk-1",
|
||||
chunkIndex: 0,
|
||||
targetPage: 1,
|
||||
snippet: "target phrase appears here",
|
||||
kind: "text",
|
||||
imageUrl: null,
|
||||
sourcePageIndex: 0,
|
||||
pageCharStart: 0,
|
||||
pageCharEnd: 13,
|
||||
lineStart: 1,
|
||||
lineEnd: 1,
|
||||
pdfRegions: [],
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
class ResizeObserverMock implements ResizeObserver {
|
||||
observe(_target: Element, _options?: ResizeObserverOptions) {
|
||||
// jsdom has no layout observer; the component only needs the API shape.
|
||||
}
|
||||
unobserve(_target: Element) {
|
||||
// jsdom has no layout observer; the component only needs the API shape.
|
||||
}
|
||||
disconnect() {
|
||||
// jsdom has no layout observer; the component only needs the API shape.
|
||||
}
|
||||
}
|
||||
vi.stubGlobal("ResizeObserver", ResizeObserverMock);
|
||||
});
|
||||
|
||||
describe("PreviewPdfView smoke", () => {
|
||||
it("renders a range URL source with text search and exact region overlay", async () => {
|
||||
render(
|
||||
React.createElement(PreviewPdfView, {
|
||||
target: target({
|
||||
pdfRegions: [
|
||||
{
|
||||
pageIndex: 0,
|
||||
pageNumber: 1,
|
||||
x: 0.1,
|
||||
y: 0.2,
|
||||
width: 0.3,
|
||||
height: 0.04,
|
||||
confidence: "exact",
|
||||
source: "pymupdf-search",
|
||||
},
|
||||
],
|
||||
}),
|
||||
file: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
|
||||
}),
|
||||
);
|
||||
|
||||
expect(screen.getByTestId("pdf-document")).toHaveAttribute(
|
||||
"data-file-kind",
|
||||
"object",
|
||||
);
|
||||
expect(screen.getByTestId("pdf-document")).toHaveAttribute(
|
||||
"data-file-url",
|
||||
expect.stringContaining("/file-signed?token=signed"),
|
||||
);
|
||||
const page = await findMainPdfPage();
|
||||
await waitFor(() => {
|
||||
expect(page.getAttribute("data-render-text-layer")).toBe("true");
|
||||
});
|
||||
expect(page.getAttribute("data-rendered-html")).toBe("target phrase");
|
||||
expect(page.getAttribute("data-rendered-html")).not.toContain("<mark>");
|
||||
|
||||
// Verify brand green highlight overlays
|
||||
const regionHighlight = screen.getByTestId("pdf-region-highlight");
|
||||
expect(regionHighlight).toBeInTheDocument();
|
||||
expect(regionHighlight).toHaveClass("bg-primary/20");
|
||||
expect(regionHighlight).toHaveClass("ring-primary/60");
|
||||
|
||||
// Verify Tailwind v4 light-mode isolation reset wrapper. After the
|
||||
// thumbnail-rail refactor, the light wrapper lives INSIDE the
|
||||
// Document and directly wraps the main-page block.
|
||||
const wrapper = screen.getByTestId("pdf-main-page").parentElement;
|
||||
expect(wrapper).toHaveClass("light");
|
||||
expect(wrapper).toHaveClass("bg-white");
|
||||
expect(wrapper).toHaveClass("text-slate-900");
|
||||
|
||||
// Verify Shadcn toolbar elements and rounded-full pill groups
|
||||
const zoomInBtn = screen.getByRole("button", { name: "Zoom in" });
|
||||
expect(zoomInBtn).toHaveClass("rounded-full");
|
||||
expect(zoomInBtn.parentElement).toHaveClass(
|
||||
"bg-muted/40",
|
||||
"p-0.5",
|
||||
"shadow-xs",
|
||||
);
|
||||
|
||||
// Source-excerpt card uses a neutral muted surface (no brand-coloured
|
||||
// left rail) so it sits inside the panel without visually competing.
|
||||
const excerptCard = screen.getByText(SOURCE_EXCERPT_TEXT).parentElement;
|
||||
expect(excerptCard).toHaveClass("border-border/60");
|
||||
expect(excerptCard).toHaveClass("bg-muted/30");
|
||||
expect(excerptCard).not.toHaveClass("border-l-primary");
|
||||
|
||||
fireEvent.change(screen.getByLabelText("Search this PDF"), {
|
||||
target: { value: "phrase" },
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(
|
||||
getMainPdfPage().getAttribute("data-rendered-html"),
|
||||
).toContain("<mark>phrase</mark>");
|
||||
});
|
||||
|
||||
const beforeZoom = Number(page.getAttribute("data-width"));
|
||||
fireEvent.click(screen.getByRole("button", { name: "Zoom in" }));
|
||||
await waitFor(() => {
|
||||
expect(
|
||||
Number(getMainPdfPage().getAttribute("data-width")),
|
||||
).toBeGreaterThan(beforeZoom);
|
||||
});
|
||||
});
|
||||
|
||||
it("debounces ResizeObserver transitions to prevent infinite rendering loops", async () => {
|
||||
const resizeCallbacks: ResizeObserverCallback[] = [];
|
||||
class FakeResizeObserver implements ResizeObserver {
|
||||
constructor(callback: ResizeObserverCallback) {
|
||||
resizeCallbacks.push(callback);
|
||||
}
|
||||
observe(_target: Element, _options?: ResizeObserverOptions) {
|
||||
// jsdom has no layout observer; the component only needs the API shape.
|
||||
}
|
||||
unobserve(_target: Element) {
|
||||
// jsdom has no layout observer; the component only needs the API shape.
|
||||
}
|
||||
disconnect() {
|
||||
// jsdom has no layout observer; the component only needs the API shape.
|
||||
}
|
||||
}
|
||||
vi.stubGlobal("ResizeObserver", FakeResizeObserver);
|
||||
|
||||
render(
|
||||
React.createElement(PreviewPdfView, {
|
||||
target: target(),
|
||||
file: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
|
||||
}),
|
||||
);
|
||||
|
||||
// Initial render sets width synchronously on mount. Let's capture the initial width.
|
||||
const page = await findMainPdfPage();
|
||||
const initialWidth = Number(page.getAttribute("data-width"));
|
||||
|
||||
// Activate fake timers AFTER finding the elements to avoid findByTestId timeout
|
||||
vi.useFakeTimers();
|
||||
|
||||
// Set up HTMLDivElement.prototype.clientWidth mock
|
||||
const originalClientWidth = Object.getOwnPropertyDescriptor(
|
||||
HTMLDivElement.prototype,
|
||||
"clientWidth",
|
||||
);
|
||||
let clientWidthValue = 300;
|
||||
Object.defineProperty(HTMLDivElement.prototype, "clientWidth", {
|
||||
get() {
|
||||
return clientWidthValue;
|
||||
},
|
||||
configurable: true,
|
||||
});
|
||||
|
||||
// Now trigger resize callback after changing clientWidth
|
||||
clientWidthValue = 600;
|
||||
const resizeCallback = resizeCallbacks[0];
|
||||
if (!resizeCallback) {
|
||||
throw new Error("Expected ResizeObserver callback to be registered");
|
||||
}
|
||||
const resizeObserver: ResizeObserver = {
|
||||
observe() {
|
||||
// The callback under test ignores the observer instance.
|
||||
},
|
||||
unobserve() {
|
||||
// The callback under test ignores the observer instance.
|
||||
},
|
||||
disconnect() {
|
||||
// The callback under test ignores the observer instance.
|
||||
},
|
||||
};
|
||||
resizeCallback([], resizeObserver);
|
||||
|
||||
// Width should NOT be updated immediately because of the 100ms debounce
|
||||
expect(Number(getMainPdfPage().getAttribute("data-width"))).toBe(
|
||||
initialWidth,
|
||||
);
|
||||
|
||||
// Fast-forward time by 100ms to trigger the debounced callback and flush updates
|
||||
act(() => {
|
||||
vi.advanceTimersByTime(100);
|
||||
vi.runAllTimers();
|
||||
});
|
||||
|
||||
// Now the width should have updated
|
||||
expect(
|
||||
Number(getMainPdfPage().getAttribute("data-width")),
|
||||
).not.toBe(initialWidth);
|
||||
expect(Number(getMainPdfPage().getAttribute("data-width"))).toBe(572); // 600 - 28 (PDF_BODY_GUTTER_PX)
|
||||
|
||||
// Clean up prototype descriptor
|
||||
if (originalClientWidth) {
|
||||
Object.defineProperty(
|
||||
HTMLDivElement.prototype,
|
||||
"clientWidth",
|
||||
originalClientWidth,
|
||||
);
|
||||
} else {
|
||||
Reflect.deleteProperty(HTMLDivElement.prototype, "clientWidth");
|
||||
}
|
||||
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it("renders only 'exact' confidence highlights and positions them with correct percentages", async () => {
|
||||
const nonExactRegion = {
|
||||
pageIndex: 0,
|
||||
pageNumber: 1,
|
||||
x: 0.5,
|
||||
y: 0.5,
|
||||
width: 0.2,
|
||||
height: 0.2,
|
||||
confidence: "fuzzy",
|
||||
source: "pymupdf-search",
|
||||
} as unknown as PreviewPdfRegion;
|
||||
|
||||
render(
|
||||
React.createElement(PreviewPdfView, {
|
||||
target: target({
|
||||
pdfRegions: [
|
||||
{
|
||||
pageIndex: 0,
|
||||
pageNumber: 1,
|
||||
x: 0.15,
|
||||
y: 0.25,
|
||||
width: 0.35,
|
||||
height: 0.45,
|
||||
confidence: "exact",
|
||||
source: "pymupdf-search",
|
||||
},
|
||||
nonExactRegion,
|
||||
],
|
||||
}),
|
||||
file: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
|
||||
}),
|
||||
);
|
||||
|
||||
await findMainPdfPage();
|
||||
|
||||
const highlights = screen.getAllByTestId("pdf-region-highlight");
|
||||
expect(highlights.length).toBe(1);
|
||||
|
||||
const exactHighlight = highlights[0];
|
||||
expect(exactHighlight.style.left).toBe("15%");
|
||||
expect(exactHighlight.style.top).toBe("25%");
|
||||
expect(exactHighlight.style.width).toBe("35%");
|
||||
expect(exactHighlight.style.height).toBe("45%");
|
||||
});
|
||||
|
||||
it("uses stable scrollbar style classes in the PDF sidebar and page container to prevent shifting", async () => {
|
||||
render(
|
||||
React.createElement(PreviewPdfView, {
|
||||
target: target(),
|
||||
file: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
|
||||
}),
|
||||
);
|
||||
|
||||
const mainPageWrapper = await screen.findByTestId("pdf-main-page");
|
||||
|
||||
// pdf-main-page → light-wrapper → scrollContainer
|
||||
const scrollContainer =
|
||||
mainPageWrapper.parentElement?.parentElement ?? null;
|
||||
expect(scrollContainer).toHaveClass("preview-scrollbar");
|
||||
expect(scrollContainer).toHaveClass("overflow-y-scroll");
|
||||
expect(scrollContainer).toHaveClass("overflow-x-auto");
|
||||
|
||||
const sidebar = screen.getByRole("button", {
|
||||
name: "Go to page 1",
|
||||
}).parentElement;
|
||||
expect(sidebar).toHaveClass("preview-scrollbar");
|
||||
expect(sidebar).toHaveClass("overflow-y-auto");
|
||||
});
|
||||
|
||||
it("highlights search terms using the custom text renderer with the mark wrapper", async () => {
|
||||
render(
|
||||
React.createElement(PreviewPdfView, {
|
||||
target: target({
|
||||
snippet: "this snippet contains some special keyword",
|
||||
}),
|
||||
file: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
|
||||
}),
|
||||
);
|
||||
|
||||
await findMainPdfPage();
|
||||
|
||||
fireEvent.change(screen.getByLabelText("Search this PDF"), {
|
||||
target: { value: "phrase" },
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(
|
||||
getMainPdfPage().getAttribute("data-rendered-html"),
|
||||
).toContain("<mark>phrase</mark>");
|
||||
});
|
||||
});
|
||||
});
|
||||
285
studio/frontend/src/__tests__/preview-store.test.ts
Normal file
285
studio/frontend/src/__tests__/preview-store.test.ts
Normal file
|
|
@ -0,0 +1,285 @@
|
|||
/**
|
||||
* Tests for preview-store object URL lifecycle (contracts §5, T4).
|
||||
*
|
||||
* Coverage:
|
||||
* - open() revokes previous object URL before assigning a new one.
|
||||
* - close() revokes any live object URL.
|
||||
* - Opening doc B while doc A is loaded revokes doc A's URL.
|
||||
* - PDFs use a signed range URL instead of a full blob download.
|
||||
* - Inline object URLs are created ONLY for safe non-PDF mediaKind (text/image).
|
||||
* - For unsafe mediaKind (html/docx/unknown) blob fetch is skipped; previewBlobUrl = null.
|
||||
* - isInlineBlobAllowed pure predicate matches contracts §5.4 allowlist.
|
||||
* - __previewStoreInternals() verifies module-scoped cleanup.
|
||||
*/
|
||||
|
||||
import type {
|
||||
PreviewMediaKind,
|
||||
PreviewTarget,
|
||||
} from "@/features/rag/api/rag-api";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
// ── Mock rag-api BEFORE importing the store ───────────────────────────
|
||||
// vi.hoisted ensures the mock refs are initialised before vi.mock factory
|
||||
// runs (vi.mock is hoisted to the top of the file by Vitest's transformer).
|
||||
|
||||
const {
|
||||
mockFetchPreviewTarget,
|
||||
mockFetchPreviewFileBlob,
|
||||
mockFetchPreviewFileUrl,
|
||||
} = vi.hoisted(() => ({
|
||||
mockFetchPreviewTarget:
|
||||
vi.fn<
|
||||
(documentId: string, chunkId?: string | null) => Promise<PreviewTarget>
|
||||
>(),
|
||||
mockFetchPreviewFileBlob:
|
||||
vi.fn<(documentId: string, signal?: AbortSignal) => Promise<Blob>>(),
|
||||
mockFetchPreviewFileUrl: vi.fn<
|
||||
(
|
||||
documentId: string,
|
||||
signal?: AbortSignal,
|
||||
) => Promise<{ url: string; expiresAt: number }>
|
||||
>(),
|
||||
}));
|
||||
|
||||
vi.mock("@/features/rag/api/rag-api", async (importOriginal) => {
|
||||
const original =
|
||||
await importOriginal<typeof import("@/features/rag/api/rag-api")>();
|
||||
return {
|
||||
...original,
|
||||
fetchPreviewTarget: mockFetchPreviewTarget,
|
||||
fetchPreviewFileBlob: mockFetchPreviewFileBlob,
|
||||
fetchPreviewFileUrl: mockFetchPreviewFileUrl,
|
||||
};
|
||||
});
|
||||
|
||||
// ── Import store AFTER mock registration ─────────────────────────────
|
||||
|
||||
import {
|
||||
__previewStoreInternals,
|
||||
isInlineBlobAllowed,
|
||||
usePreviewStore,
|
||||
} from "@/features/rag/stores/preview-store";
|
||||
|
||||
// ── Mock URL.createObjectURL / revokeObjectURL ────────────────────────
|
||||
|
||||
let urlCounter = 0;
|
||||
|
||||
beforeEach(() => {
|
||||
urlCounter = 0;
|
||||
vi.spyOn(URL, "createObjectURL").mockImplementation(() => {
|
||||
return `blob:test/${++urlCounter}`;
|
||||
});
|
||||
vi.spyOn(URL, "revokeObjectURL").mockImplementation((_url: string) => {
|
||||
/* no-op */
|
||||
});
|
||||
mockFetchPreviewTarget.mockReset();
|
||||
mockFetchPreviewFileBlob.mockReset();
|
||||
mockFetchPreviewFileUrl.mockReset();
|
||||
// Reset store to idle between tests
|
||||
usePreviewStore.getState().close();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
// ── Helpers ───────────────────────────────────────────────────────────
|
||||
|
||||
function makeTarget(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
|
||||
return {
|
||||
documentId: "doc-abc",
|
||||
filename: "report.pdf",
|
||||
contentType: "application/pdf",
|
||||
mediaKind: "pdf",
|
||||
byteSize: 100,
|
||||
status: "completed",
|
||||
kbId: "kb-1",
|
||||
threadId: null,
|
||||
chunkId: null,
|
||||
chunkIndex: null,
|
||||
targetPage: null,
|
||||
snippet: null,
|
||||
kind: null,
|
||||
imageUrl: null,
|
||||
sourcePageIndex: null,
|
||||
pageCharStart: null,
|
||||
pageCharEnd: null,
|
||||
lineStart: null,
|
||||
lineEnd: null,
|
||||
pdfRegions: [],
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function makePdfBlob(): Blob {
|
||||
return new Blob(["%PDF-1.4"], { type: "application/pdf" });
|
||||
}
|
||||
|
||||
// ── Pure-logic: isInlineBlobAllowed (always green) ────────────────────
|
||||
|
||||
describe("isInlineBlobAllowed (contracts §5.4, pure logic)", () => {
|
||||
const safe: PreviewMediaKind[] = ["pdf", "text", "image"];
|
||||
const unsafe: PreviewMediaKind[] = ["html", "docx", "unknown"];
|
||||
|
||||
it.each(safe)("mediaKind=%s is inline-safe", (mk) => {
|
||||
expect(isInlineBlobAllowed(mk)).toBe(true);
|
||||
});
|
||||
|
||||
it.each(unsafe)("mediaKind=%s is NOT inline-safe (Risk #3)", (mk) => {
|
||||
expect(isInlineBlobAllowed(mk)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
// ── Integration tests against real store ─────────────────────────────
|
||||
|
||||
describe("preview-store open/close lifecycle (contracts §5)", () => {
|
||||
it("open() for pdf stores a signed URL without creating an object URL", async () => {
|
||||
mockFetchPreviewTarget.mockResolvedValue(makeTarget({ mediaKind: "pdf" }));
|
||||
mockFetchPreviewFileUrl.mockResolvedValue({
|
||||
url: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
|
||||
expiresAt: 1_700_000_000,
|
||||
});
|
||||
|
||||
await usePreviewStore.getState().open({ documentId: "doc-abc" });
|
||||
|
||||
expect(mockFetchPreviewFileBlob).not.toHaveBeenCalled();
|
||||
expect(URL.createObjectURL).not.toHaveBeenCalled();
|
||||
const { previewBlob, previewBlobUrl, previewFileUrl, status } =
|
||||
usePreviewStore.getState();
|
||||
expect(previewBlob).toBeNull();
|
||||
expect(previewBlobUrl).toBeNull();
|
||||
expect(previewFileUrl).toContain("/file-signed?token=signed");
|
||||
expect(status).toBe("ready");
|
||||
});
|
||||
|
||||
it("open() passes backendChunkId through to fetchPreviewTarget", async () => {
|
||||
mockFetchPreviewTarget.mockResolvedValue(makeTarget({ mediaKind: "pdf" }));
|
||||
mockFetchPreviewFileUrl.mockResolvedValue({
|
||||
url: "/api/rag/documents/doc-abc/file-signed?token=signed",
|
||||
expiresAt: 1_700_000_000,
|
||||
});
|
||||
|
||||
await usePreviewStore.getState().open({
|
||||
documentId: "doc-abc",
|
||||
backendChunkId: "chunk-xyz",
|
||||
});
|
||||
|
||||
expect(mockFetchPreviewTarget).toHaveBeenCalledWith(
|
||||
"doc-abc",
|
||||
"chunk-xyz",
|
||||
);
|
||||
});
|
||||
|
||||
it("close() revokes the live object URL and clears state", async () => {
|
||||
mockFetchPreviewTarget.mockResolvedValue(
|
||||
makeTarget({ mediaKind: "text", filename: "notes.txt" }),
|
||||
);
|
||||
mockFetchPreviewFileBlob.mockResolvedValue(makePdfBlob());
|
||||
|
||||
await usePreviewStore.getState().open({ documentId: "doc-abc" });
|
||||
const blobUrl = usePreviewStore.getState().previewBlobUrl;
|
||||
expect(blobUrl).toMatch(/^blob:/);
|
||||
|
||||
usePreviewStore.getState().close();
|
||||
|
||||
expect(URL.revokeObjectURL).toHaveBeenCalledWith(blobUrl);
|
||||
const { previewBlob, previewBlobUrl, target, status } =
|
||||
usePreviewStore.getState();
|
||||
expect(previewBlob).toBeNull();
|
||||
expect(previewBlobUrl).toBeNull();
|
||||
expect(target).toBeNull();
|
||||
expect(status).toBe("idle");
|
||||
});
|
||||
|
||||
it("opening doc B revokes doc A's URL before creating doc B's (contracts §5.1)", async () => {
|
||||
mockFetchPreviewTarget.mockResolvedValue(
|
||||
makeTarget({ mediaKind: "text", filename: "notes.txt" }),
|
||||
);
|
||||
mockFetchPreviewFileBlob.mockResolvedValue(makePdfBlob());
|
||||
|
||||
await usePreviewStore.getState().open({ documentId: "doc-A" });
|
||||
const urlA = usePreviewStore.getState().previewBlobUrl;
|
||||
expect(urlA).toMatch(/^blob:/);
|
||||
|
||||
await usePreviewStore.getState().open({ documentId: "doc-B" });
|
||||
|
||||
expect(URL.revokeObjectURL).toHaveBeenCalledWith(urlA);
|
||||
const urlB = usePreviewStore.getState().previewBlobUrl;
|
||||
expect(urlB).not.toBe(urlA);
|
||||
expect(urlB).toMatch(/^blob:/);
|
||||
});
|
||||
|
||||
it("html mediaKind skips blob fetch and sets previewBlobUrl = null (Risk #3)", async () => {
|
||||
mockFetchPreviewTarget.mockResolvedValue(
|
||||
makeTarget({ mediaKind: "html", filename: "evil.html" }),
|
||||
);
|
||||
|
||||
await usePreviewStore.getState().open({ documentId: "doc-html" });
|
||||
|
||||
expect(mockFetchPreviewFileBlob).not.toHaveBeenCalled();
|
||||
expect(mockFetchPreviewFileUrl).not.toHaveBeenCalled();
|
||||
expect(URL.createObjectURL).not.toHaveBeenCalled();
|
||||
const { previewBlob, previewBlobUrl, status } = usePreviewStore.getState();
|
||||
expect(previewBlob).toBeNull();
|
||||
expect(previewBlobUrl).toBeNull();
|
||||
expect(status).toBe("ready");
|
||||
});
|
||||
|
||||
it("docx mediaKind skips blob fetch and sets previewBlobUrl = null (Risk #3)", async () => {
|
||||
mockFetchPreviewTarget.mockResolvedValue(
|
||||
makeTarget({ mediaKind: "docx", filename: "report.docx" }),
|
||||
);
|
||||
|
||||
await usePreviewStore.getState().open({ documentId: "doc-docx" });
|
||||
|
||||
expect(mockFetchPreviewFileBlob).not.toHaveBeenCalled();
|
||||
expect(mockFetchPreviewFileUrl).not.toHaveBeenCalled();
|
||||
expect(URL.createObjectURL).not.toHaveBeenCalled();
|
||||
expect(usePreviewStore.getState().previewBlob).toBeNull();
|
||||
expect(usePreviewStore.getState().previewBlobUrl).toBeNull();
|
||||
});
|
||||
|
||||
it("unknown mediaKind skips blob fetch and sets previewBlobUrl = null", async () => {
|
||||
mockFetchPreviewTarget.mockResolvedValue(
|
||||
makeTarget({ mediaKind: "unknown", filename: "data.bin" }),
|
||||
);
|
||||
|
||||
await usePreviewStore.getState().open({ documentId: "doc-bin" });
|
||||
|
||||
expect(mockFetchPreviewFileBlob).not.toHaveBeenCalled();
|
||||
expect(mockFetchPreviewFileUrl).not.toHaveBeenCalled();
|
||||
expect(URL.createObjectURL).not.toHaveBeenCalled();
|
||||
expect(usePreviewStore.getState().previewBlob).toBeNull();
|
||||
expect(usePreviewStore.getState().previewBlobUrl).toBeNull();
|
||||
});
|
||||
|
||||
it("close() when nothing is open does not throw", () => {
|
||||
expect(() => usePreviewStore.getState().close()).not.toThrow();
|
||||
});
|
||||
|
||||
it("__previewStoreInternals shows no activeBlobUrl after close()", async () => {
|
||||
mockFetchPreviewTarget.mockResolvedValue(
|
||||
makeTarget({ mediaKind: "text", filename: "notes.txt" }),
|
||||
);
|
||||
mockFetchPreviewFileBlob.mockResolvedValue(makePdfBlob());
|
||||
|
||||
await usePreviewStore.getState().open({ documentId: "doc-abc" });
|
||||
expect(__previewStoreInternals().activeBlobUrl).toMatch(/^blob:/);
|
||||
|
||||
usePreviewStore.getState().close();
|
||||
expect(__previewStoreInternals().activeBlobUrl).toBeNull();
|
||||
expect(__previewStoreInternals().hasInflightController).toBe(false);
|
||||
});
|
||||
|
||||
it("fetchPreviewTarget error sets status=error and clears target", async () => {
|
||||
mockFetchPreviewTarget.mockRejectedValue(new Error("404 not found"));
|
||||
|
||||
await usePreviewStore.getState().open({ documentId: "missing" });
|
||||
|
||||
const { status, error, target } = usePreviewStore.getState();
|
||||
expect(status).toBe("error");
|
||||
expect(error).toMatch(/404/);
|
||||
expect(target).toBeNull();
|
||||
expect(URL.createObjectURL).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
import type { PreviewTarget } from "@/features/rag/api/rag-api";
|
||||
import { PreviewTextView } from "@/features/rag/components/preview-text-view";
|
||||
import { render, screen } from "@testing-library/react";
|
||||
import React from "react";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("@/features/auth", () => ({
|
||||
authFetch: vi.fn(),
|
||||
}));
|
||||
|
||||
function target(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
|
||||
return {
|
||||
documentId: "doc-abc",
|
||||
filename: "notes.txt",
|
||||
contentType: "text/plain",
|
||||
mediaKind: "text",
|
||||
byteSize: 100,
|
||||
status: "completed",
|
||||
kbId: "kb-1",
|
||||
threadId: null,
|
||||
chunkId: "chunk-1",
|
||||
chunkIndex: 0,
|
||||
targetPage: 2,
|
||||
snippet: "alpha\nhighlighted line\nomega",
|
||||
kind: "text",
|
||||
imageUrl: null,
|
||||
sourcePageIndex: null,
|
||||
pageCharStart: null,
|
||||
pageCharEnd: null,
|
||||
lineStart: null,
|
||||
lineEnd: null,
|
||||
pdfRegions: [],
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
describe("PreviewTextView locator highlight fallback", () => {
|
||||
it("emphasizes the source excerpt when nullable locators are present", () => {
|
||||
render(
|
||||
React.createElement(PreviewTextView, {
|
||||
target: target({
|
||||
sourcePageIndex: 1,
|
||||
pageCharStart: 6,
|
||||
pageCharEnd: 22,
|
||||
lineStart: 2,
|
||||
lineEnd: 2,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
expect(screen.getByText(/highlighted source excerpt/i)).toBeInTheDocument();
|
||||
expect(document.querySelector("mark")?.textContent).toContain(
|
||||
"highlighted line",
|
||||
);
|
||||
});
|
||||
|
||||
it("keeps the source excerpt visible when locators are missing", () => {
|
||||
render(React.createElement(PreviewTextView, { target: target() }));
|
||||
|
||||
expect(screen.getByText(/source excerpt/i)).toBeInTheDocument();
|
||||
expect(document.querySelector("mark")).toBeNull();
|
||||
expect(screen.getByText(/highlighted line/i)).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
191
studio/frontend/src/__tests__/rag-api.test.ts
Normal file
191
studio/frontend/src/__tests__/rag-api.test.ts
Normal file
|
|
@ -0,0 +1,191 @@
|
|||
import type { PreviewTarget } from "@/features/rag/api/rag-api";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
const {
|
||||
mockAuthFetch,
|
||||
mockGetAuthToken,
|
||||
mockEventSourceInstances,
|
||||
eventSourcePolyfillExport,
|
||||
authorizationHeader,
|
||||
} = vi.hoisted(() => ({
|
||||
mockAuthFetch: vi.fn(),
|
||||
mockGetAuthToken: vi.fn(),
|
||||
mockEventSourceInstances: [] as MockEventSource[],
|
||||
eventSourcePolyfillExport: "EventSourcePolyfill",
|
||||
authorizationHeader: "Authorization",
|
||||
}));
|
||||
|
||||
vi.mock("@/features/auth", () => ({
|
||||
authFetch: mockAuthFetch,
|
||||
getAuthToken: mockGetAuthToken,
|
||||
}));
|
||||
|
||||
interface MockEventSource {
|
||||
url: string;
|
||||
options: unknown;
|
||||
onmessage: ((event: MessageEvent) => void) | null;
|
||||
onerror: (() => void) | null;
|
||||
close: ReturnType<typeof vi.fn>;
|
||||
}
|
||||
|
||||
vi.mock("event-source-polyfill", () => ({
|
||||
[eventSourcePolyfillExport]: class {
|
||||
url: string;
|
||||
options: unknown;
|
||||
onmessage: ((event: MessageEvent) => void) | null = null;
|
||||
onerror: (() => void) | null = null;
|
||||
close = vi.fn();
|
||||
|
||||
constructor(url: string, options?: unknown) {
|
||||
this.url = url;
|
||||
this.options = options;
|
||||
mockEventSourceInstances.push(this);
|
||||
}
|
||||
},
|
||||
}));
|
||||
|
||||
import {
|
||||
backfillDocumentLocators,
|
||||
fetchPreviewFileUrl,
|
||||
fetchPreviewTarget,
|
||||
subscribeToJobEvents,
|
||||
} from "@/features/rag/api/rag-api";
|
||||
|
||||
function target(): PreviewTarget {
|
||||
return {
|
||||
documentId: "doc-abc",
|
||||
filename: "report.pdf",
|
||||
contentType: "application/pdf",
|
||||
mediaKind: "pdf",
|
||||
byteSize: 100,
|
||||
status: "completed",
|
||||
kbId: "kb-1",
|
||||
threadId: null,
|
||||
chunkId: "chunk-xyz",
|
||||
chunkIndex: 0,
|
||||
targetPage: 1,
|
||||
snippet: "excerpt",
|
||||
kind: "text",
|
||||
imageUrl: null,
|
||||
sourcePageIndex: 0,
|
||||
pageCharStart: 0,
|
||||
pageCharEnd: 7,
|
||||
lineStart: 1,
|
||||
lineEnd: 1,
|
||||
pdfRegions: [],
|
||||
};
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
mockAuthFetch.mockReset();
|
||||
mockGetAuthToken.mockReset();
|
||||
mockEventSourceInstances.length = 0;
|
||||
});
|
||||
|
||||
describe("RAG API preview target", () => {
|
||||
it("URL-encodes documentId and chunk_id", async () => {
|
||||
mockAuthFetch.mockResolvedValue(
|
||||
new Response(JSON.stringify(target()), {
|
||||
status: 200,
|
||||
headers: { "Content-Type": "application/json" },
|
||||
}),
|
||||
);
|
||||
|
||||
await fetchPreviewTarget("doc id/with?slash", "chunk id/with?amp&eq=1");
|
||||
|
||||
expect(mockAuthFetch).toHaveBeenCalledWith(
|
||||
"/api/rag/documents/doc%20id%2Fwith%3Fslash/preview-target?chunk_id=chunk%20id%2Fwith%3Famp%26eq%3D1",
|
||||
);
|
||||
});
|
||||
|
||||
it("fetches signed preview URL without adding a bearer token query", async () => {
|
||||
mockAuthFetch.mockResolvedValue(
|
||||
new Response(
|
||||
JSON.stringify({
|
||||
url: "/api/rag/documents/doc-abc/file-signed?token=signed-preview",
|
||||
expiresAt: 1_700_000_000,
|
||||
}),
|
||||
{
|
||||
status: 200,
|
||||
headers: { "Content-Type": "application/json" },
|
||||
},
|
||||
),
|
||||
);
|
||||
|
||||
const result = await fetchPreviewFileUrl("doc id/with?slash");
|
||||
|
||||
expect(mockAuthFetch).toHaveBeenCalledWith(
|
||||
"/api/rag/documents/doc%20id%2Fwith%3Fslash/file-url",
|
||||
undefined,
|
||||
);
|
||||
expect(result.url).toContain("token=signed-preview");
|
||||
expect(result.url).not.toContain("Bearer");
|
||||
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", () => {
|
||||
it("opens SSE with Authorization header instead of token query params", () => {
|
||||
mockGetAuthToken.mockReturnValue("mock-token-123");
|
||||
|
||||
const unsubscribe = subscribeToJobEvents("job id/with?slash", {});
|
||||
|
||||
expect(mockEventSourceInstances).toHaveLength(1);
|
||||
const source = mockEventSourceInstances[0];
|
||||
expect(source.url).toContain(
|
||||
"/api/rag/jobs/job%20id%2Fwith%3Fslash/events",
|
||||
);
|
||||
expect(source.url).not.toContain("token=");
|
||||
expect(source.options).toEqual({
|
||||
headers: {
|
||||
[authorizationHeader]: "Bearer mock-token-123",
|
||||
},
|
||||
});
|
||||
|
||||
unsubscribe();
|
||||
expect(source.close).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("omits EventSource options when there is no bearer token", () => {
|
||||
mockGetAuthToken.mockReturnValue(null);
|
||||
|
||||
const unsubscribe = subscribeToJobEvents("job-abc", {});
|
||||
|
||||
expect(mockEventSourceInstances).toHaveLength(1);
|
||||
const source = mockEventSourceInstances[0];
|
||||
expect(source.url).toContain("/api/rag/jobs/job-abc/events");
|
||||
expect(source.url).not.toContain("token=");
|
||||
expect(source.options).toBeUndefined();
|
||||
|
||||
unsubscribe();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,84 @@
|
|||
import { render, screen } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import React from "react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
const { mockOpenPreview } = vi.hoisted(() => ({
|
||||
mockOpenPreview: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/features/rag/stores/preview-store", () => ({
|
||||
usePreviewStore: (
|
||||
selector?: (state: { open: typeof mockOpenPreview }) => unknown,
|
||||
) => {
|
||||
const state = { open: mockOpenPreview };
|
||||
return typeof selector === "function" ? selector(state) : state;
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock("@assistant-ui/react", () => ({
|
||||
useAuiState: (
|
||||
selector: (state: { message: { content: unknown[] } }) => unknown,
|
||||
) =>
|
||||
selector({
|
||||
message: { content: [{ type: "text", text: "Answer ready." }] },
|
||||
}),
|
||||
}));
|
||||
|
||||
import { SearchKnowledgeBaseToolUI } from "@/components/assistant-ui/tool-ui-search-knowledge-base";
|
||||
|
||||
const TOOL_UI = SearchKnowledgeBaseToolUI as React.ComponentType<
|
||||
Record<string, unknown>
|
||||
>;
|
||||
const SEARCHED_DOCS_BUTTON_RE = /searched docs/i;
|
||||
const MAIN_PREVIEW_BUTTON_RE = /open preview of main\.pdf/i;
|
||||
const LEGACY_PREVIEW_BUTTON_RE = /open preview of legacy\.pdf/i;
|
||||
|
||||
function renderTool(result: string) {
|
||||
return render(
|
||||
React.createElement(TOOL_UI, {
|
||||
args: { query: "what is interior modeling?" },
|
||||
result,
|
||||
status: { type: "complete" },
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
mockOpenPreview.mockClear();
|
||||
});
|
||||
|
||||
describe("SearchKnowledgeBaseToolUI preview routing", () => {
|
||||
it("opens preview from a retrieved chunk source label", async () => {
|
||||
renderTool(
|
||||
'<chunk id="1" source="main.pdf" page="3" document_id="doc-abc" chunk_id="chunk-xyz">The paper objectives.</chunk>',
|
||||
);
|
||||
|
||||
await userEvent.click(
|
||||
screen.getByRole("button", { name: SEARCHED_DOCS_BUTTON_RE }),
|
||||
);
|
||||
await userEvent.click(
|
||||
screen.getByRole("button", { name: MAIN_PREVIEW_BUTTON_RE }),
|
||||
);
|
||||
|
||||
expect(mockOpenPreview).toHaveBeenCalledWith({
|
||||
documentId: "doc-abc",
|
||||
backendChunkId: "chunk-xyz",
|
||||
});
|
||||
});
|
||||
|
||||
it("keeps legacy chunk labels non-clickable without durable IDs", async () => {
|
||||
renderTool(
|
||||
'<chunk id="1" source="legacy.pdf" page="3">Legacy chunk text.</chunk>',
|
||||
);
|
||||
|
||||
await userEvent.click(
|
||||
screen.getByRole("button", { name: SEARCHED_DOCS_BUTTON_RE }),
|
||||
);
|
||||
|
||||
expect(
|
||||
screen.queryByRole("button", { name: LEGACY_PREVIEW_BUTTON_RE }),
|
||||
).toBeNull();
|
||||
expect(mockOpenPreview).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
80
studio/frontend/src/__tests__/sources.test.tsx
Normal file
80
studio/frontend/src/__tests__/sources.test.tsx
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
import { fireEvent, render, screen } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import React from "react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
const { mockOpen } = vi.hoisted(() => ({ mockOpen: vi.fn() }));
|
||||
|
||||
vi.mock("@/features/rag/stores/preview-store", () => ({
|
||||
usePreviewStore: (
|
||||
selector?: (state: { open: typeof mockOpen }) => unknown,
|
||||
) => {
|
||||
const state = { open: mockOpen };
|
||||
return typeof selector === "function" ? selector(state) : state;
|
||||
},
|
||||
}));
|
||||
|
||||
import { DocumentSourceBadge } from "@/components/assistant-ui/sources";
|
||||
|
||||
function source(overrides: Record<string, unknown> = {}) {
|
||||
return {
|
||||
kind: "document" as const,
|
||||
chunkId: "2",
|
||||
documentId: "doc-abc",
|
||||
backendChunkId: "chunk-xyz",
|
||||
filename: "report.pdf",
|
||||
page: "7",
|
||||
score: "0.85",
|
||||
text: "The margin rose to 18%.",
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
mockOpen.mockClear();
|
||||
});
|
||||
|
||||
describe("DocumentSourceBadge preview routing", () => {
|
||||
it("opens preview with durable document and backend chunk IDs on click", async () => {
|
||||
render(React.createElement(DocumentSourceBadge, { source: source() }));
|
||||
|
||||
await userEvent.click(
|
||||
screen.getByRole("button", { name: /open preview/i }),
|
||||
);
|
||||
|
||||
expect(mockOpen).toHaveBeenCalledWith({
|
||||
documentId: "doc-abc",
|
||||
backendChunkId: "chunk-xyz",
|
||||
});
|
||||
});
|
||||
|
||||
it("opens preview from Enter and Space", () => {
|
||||
render(React.createElement(DocumentSourceBadge, { source: source() }));
|
||||
const badge = screen.getByRole("button", { name: /open preview/i });
|
||||
|
||||
fireEvent.keyDown(badge, { key: "Enter" });
|
||||
fireEvent.keyDown(badge, { key: " " });
|
||||
|
||||
expect(mockOpen).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("legacy source without durable IDs remains hover-only", async () => {
|
||||
render(
|
||||
React.createElement(DocumentSourceBadge, {
|
||||
source: source({ documentId: null, backendChunkId: null }),
|
||||
}),
|
||||
);
|
||||
|
||||
expect(screen.queryByRole("button", { name: /open preview/i })).toBeNull();
|
||||
await userEvent.click(screen.getByText("[2]"));
|
||||
expect(mockOpen).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("applies brand-aligned interactive styling when preview is clickable", () => {
|
||||
render(React.createElement(DocumentSourceBadge, { source: source() }));
|
||||
const badge = screen.getByRole("button", { name: /open preview/i });
|
||||
|
||||
expect(badge).toHaveClass("cursor-pointer");
|
||||
expect(badge).toHaveClass("hover:bg-chat-icon-bg-hover!");
|
||||
expect(badge).toHaveClass("focus-visible:ring-ring/50");
|
||||
});
|
||||
});
|
||||
|
|
@ -1,24 +1,26 @@
|
|||
"use client";
|
||||
|
||||
import { openLink } from "@/lib/open-link";
|
||||
import {
|
||||
memo,
|
||||
useState,
|
||||
useRef,
|
||||
useEffect,
|
||||
useCallback,
|
||||
type ComponentProps,
|
||||
type FC,
|
||||
} from "react";
|
||||
import { FileTextIcon } from "lucide-react";
|
||||
import { useMessage } from "@assistant-ui/react";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { Badge, badgeVariants, type BadgeProps } from "./badge";
|
||||
import {
|
||||
HoverCard,
|
||||
HoverCardTrigger,
|
||||
HoverCardContent,
|
||||
HoverCardTrigger,
|
||||
} from "@/components/ui/hover-card";
|
||||
import { usePreviewStore } from "@/features/rag/stores/preview-store";
|
||||
import { openLink } from "@/lib/open-link";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { useMessage } from "@assistant-ui/react";
|
||||
import { FileTextIcon } from "lucide-react";
|
||||
import {
|
||||
type ComponentProps,
|
||||
type FC,
|
||||
type KeyboardEvent as ReactKeyboardEvent,
|
||||
memo,
|
||||
useCallback,
|
||||
useEffect,
|
||||
useRef,
|
||||
useState,
|
||||
} from "react";
|
||||
import { Badge, type BadgeProps, badgeVariants } from "./badge";
|
||||
|
||||
// ── Helpers ──────────────────────────────────────────────────
|
||||
|
||||
|
|
@ -45,8 +47,12 @@ function SourceIcon({
|
|||
}: ComponentProps<"span"> & { url: string; size?: number }) {
|
||||
const [hasError, setHasError] = useState(false);
|
||||
const domain = extractDomain(url);
|
||||
const SIZE_CLASSES: Record<number, string> = { 3: "size-3", 4: "size-4", 5: "size-5" };
|
||||
const sizeClass = SIZE_CLASSES[size] ?? "size-3";
|
||||
const sizeClasses: Record<number, string> = {
|
||||
3: "size-3",
|
||||
4: "size-4",
|
||||
5: "size-5",
|
||||
};
|
||||
const sizeClass = sizeClasses[size] ?? "size-3";
|
||||
|
||||
if (hasError) {
|
||||
return (
|
||||
|
|
@ -101,7 +107,7 @@ function Source({
|
|||
}: SourceProps) {
|
||||
return (
|
||||
<Badge
|
||||
asChild
|
||||
asChild={true}
|
||||
variant={variant}
|
||||
size={size}
|
||||
className={cn(
|
||||
|
|
@ -136,7 +142,12 @@ interface UrlSourceData {
|
|||
|
||||
interface DocSourceData {
|
||||
kind: "document";
|
||||
/** Visible citation id (the `[N]` reference). Display only. */
|
||||
chunkId: string;
|
||||
/** Durable backend `rag_documents.id`. null on legacy sources. */
|
||||
documentId: string | null;
|
||||
/** Durable backend `rag_chunks.id`. null on legacy sources. */
|
||||
backendChunkId: string | null;
|
||||
filename: string;
|
||||
page?: string;
|
||||
text: string;
|
||||
|
|
@ -145,9 +156,7 @@ interface DocSourceData {
|
|||
type SourceData = UrlSourceData | DocSourceData;
|
||||
|
||||
function sourceKey(source: SourceData): string {
|
||||
return source.kind === "url"
|
||||
? `url:${source.url}`
|
||||
: `doc:${source.chunkId}`;
|
||||
return source.kind === "url" ? `url:${source.url}` : `doc:${source.chunkId}`;
|
||||
}
|
||||
|
||||
const SourceBadge: FC<{ source: UrlSourceData }> = ({ source }) => {
|
||||
|
|
@ -156,7 +165,7 @@ const SourceBadge: FC<{ source: UrlSourceData }> = ({ source }) => {
|
|||
|
||||
return (
|
||||
<HoverCard openDelay={0} closeDelay={0}>
|
||||
<HoverCardTrigger asChild>
|
||||
<HoverCardTrigger asChild={true}>
|
||||
<span className="inline-block">
|
||||
<Source href={source.url}>
|
||||
<SourceIcon url={source.url} />
|
||||
|
|
@ -193,13 +202,53 @@ const DocumentSourceBadge: FC<{ source: DocSourceData }> = ({ source }) => {
|
|||
const metaParts: string[] = [];
|
||||
if (source.page) metaParts.push(`page ${source.page}`);
|
||||
|
||||
// Preview is clickable IFF both durable IDs are present (contracts
|
||||
// §4.1 routing rule + Q3). Legacy sources fall through to a
|
||||
// non-interactive badge with hover-only behavior.
|
||||
const isClickable =
|
||||
source.documentId !== null && source.backendChunkId !== null;
|
||||
const openPreview = usePreviewStore((s) => s.open);
|
||||
|
||||
const handleOpen = useCallback(() => {
|
||||
if (!isClickable || !source.documentId) return;
|
||||
void openPreview({
|
||||
documentId: source.documentId,
|
||||
backendChunkId: source.backendChunkId,
|
||||
});
|
||||
}, [isClickable, openPreview, source.documentId, source.backendChunkId]);
|
||||
|
||||
const handleKeyDown = useCallback(
|
||||
(e: ReactKeyboardEvent<HTMLElement>) => {
|
||||
if (!isClickable) return;
|
||||
if (e.key === "Enter" || e.key === " ") {
|
||||
e.preventDefault();
|
||||
handleOpen();
|
||||
}
|
||||
},
|
||||
[isClickable, handleOpen],
|
||||
);
|
||||
|
||||
return (
|
||||
<HoverCard openDelay={0} closeDelay={0}>
|
||||
<HoverCardTrigger asChild>
|
||||
<HoverCardTrigger asChild={true}>
|
||||
<span className="inline-block">
|
||||
<Badge
|
||||
variant="outline"
|
||||
className="rounded-full cursor-default inline-flex items-center gap-1.5 outline-none"
|
||||
{...(isClickable
|
||||
? {
|
||||
role: "button",
|
||||
tabIndex: 0,
|
||||
onClick: handleOpen,
|
||||
onKeyDown: handleKeyDown,
|
||||
"aria-label": `Open preview of ${source.filename}`,
|
||||
}
|
||||
: {})}
|
||||
className={cn(
|
||||
"rounded-full inline-flex items-center gap-1.5 outline-none",
|
||||
isClickable
|
||||
? "cursor-pointer hover:bg-chat-icon-bg-hover! hover:text-chat-icon-fg-hover! focus-visible:border-ring focus-visible:ring-[3px] focus-visible:ring-ring/50"
|
||||
: "cursor-default",
|
||||
)}
|
||||
>
|
||||
<span className="font-mono text-[10px] font-semibold text-muted-foreground">
|
||||
[{source.chunkId}]
|
||||
|
|
@ -229,6 +278,11 @@ const DocumentSourceBadge: FC<{ source: DocSourceData }> = ({ source }) => {
|
|||
<p className="text-xs text-white/70 leading-relaxed line-clamp-3 whitespace-pre-wrap">
|
||||
{source.text}
|
||||
</p>
|
||||
{isClickable ? (
|
||||
<p className="mt-1 text-[10px] text-white/50">
|
||||
Click to open preview
|
||||
</p>
|
||||
) : null}
|
||||
</div>
|
||||
</div>
|
||||
</HoverCardContent>
|
||||
|
|
@ -269,6 +323,8 @@ const SourcesGroup: FC = () => {
|
|||
) {
|
||||
const docPart = part as {
|
||||
chunkId?: string;
|
||||
documentId?: string | null;
|
||||
backendChunkId?: string | null;
|
||||
filename?: string;
|
||||
page?: string;
|
||||
text?: string;
|
||||
|
|
@ -277,6 +333,8 @@ const SourcesGroup: FC = () => {
|
|||
sources.push({
|
||||
kind: "document",
|
||||
chunkId: docPart.chunkId,
|
||||
documentId: docPart.documentId ?? null,
|
||||
backendChunkId: docPart.backendChunkId ?? null,
|
||||
filename: docPart.filename,
|
||||
page: docPart.page,
|
||||
text: docPart.text ?? "",
|
||||
|
|
@ -340,7 +398,7 @@ const SourcesGroup: FC = () => {
|
|||
{/* Hidden measurement container — renders all badges to measure row positions */}
|
||||
<div
|
||||
ref={containerRef}
|
||||
aria-hidden
|
||||
aria-hidden={true}
|
||||
className="flex w-full flex-wrap gap-1 invisible absolute pointer-events-none"
|
||||
>
|
||||
{sources.map((source) => (
|
||||
|
|
|
|||
|
|
@ -3,14 +3,15 @@
|
|||
|
||||
"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 { authFetch } from "@/features/auth";
|
||||
import { FileTextIcon, ImageIcon, LoaderIcon } from "lucide-react";
|
||||
import { memo, useEffect, useState } from "react";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { memo, useCallback, useEffect, useState } from "react";
|
||||
import {
|
||||
ToolFallbackContent,
|
||||
ToolFallbackRoot,
|
||||
|
|
@ -18,14 +19,29 @@ import {
|
|||
} from "./tool-fallback";
|
||||
|
||||
export interface ParsedChunk {
|
||||
/** Visible citation id the model uses inside `[N]` references. Display
|
||||
* only; never sent to the backend as a chunk_id. */
|
||||
id: string;
|
||||
source: string;
|
||||
page?: string;
|
||||
chunkIndex?: string;
|
||||
tokens?: string;
|
||||
sourcePageIndex?: string;
|
||||
pageCharStart?: string;
|
||||
pageCharEnd?: string;
|
||||
lineStart?: string;
|
||||
lineEnd?: string;
|
||||
kind?: string;
|
||||
imageUrl?: string;
|
||||
text: string;
|
||||
/** Durable `rag_documents.id`. Carries through when the tool XML
|
||||
* includes `document_id="..."`. Absent on legacy tool output. */
|
||||
documentId?: string;
|
||||
/** Durable `rag_chunks.id`. Carries through when the tool XML
|
||||
* includes `chunk_id="..."`. Absent on legacy tool output. The
|
||||
* preview routing value sent as `?chunk_id=` to `/preview-target`;
|
||||
* never the same as the visible `id`. */
|
||||
backendChunkId?: string;
|
||||
}
|
||||
|
||||
const ATTR_RE = /(\w+)="([^"]*)"/g;
|
||||
|
|
@ -59,9 +75,17 @@ export function parseChunks(raw: string): ParsedChunk[] {
|
|||
page: attrs.page,
|
||||
chunkIndex: attrs.chunk_index,
|
||||
tokens: attrs.tokens,
|
||||
sourcePageIndex: attrs.source_page_index,
|
||||
pageCharStart: attrs.page_char_start,
|
||||
pageCharEnd: attrs.page_char_end,
|
||||
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 } : {}),
|
||||
...(attrs.chunk_id ? { backendChunkId: attrs.chunk_id } : {}),
|
||||
});
|
||||
}
|
||||
match = CHUNK_RE.exec(raw);
|
||||
|
|
@ -82,9 +106,11 @@ function useAuthedImageUrl(path: string | undefined): string | undefined {
|
|||
}
|
||||
let cancelled = false;
|
||||
let objectUrl: string | undefined;
|
||||
void authFetch(path)
|
||||
.then(async (response) => {
|
||||
if (!response.ok) throw new Error(`image fetch ${response.status}`);
|
||||
authFetch(path)
|
||||
.then((response) => {
|
||||
if (!response.ok) {
|
||||
throw new Error(`image fetch ${response.status}`);
|
||||
}
|
||||
return response.blob();
|
||||
})
|
||||
.then((blob) => {
|
||||
|
|
@ -123,31 +149,66 @@ function ChunkImage({ url, alt }: { url: string; alt: string }) {
|
|||
}
|
||||
|
||||
function ChunkCard({ chunk }: { chunk: ParsedChunk }) {
|
||||
const openPreview = usePreviewStore((s) => s.open);
|
||||
const meta: string[] = [];
|
||||
if (chunk.page) meta.push(`page ${chunk.page}`);
|
||||
if (chunk.tokens) meta.push(`${chunk.tokens} tok`);
|
||||
if (chunk.chunkIndex) meta.push(`#${chunk.chunkIndex}`);
|
||||
if (chunk.kind && chunk.kind !== "text") meta.push(chunk.kind);
|
||||
|
||||
const documentId = chunk.documentId;
|
||||
const backendChunkId = chunk.backendChunkId;
|
||||
const isPreviewable = Boolean(documentId && backendChunkId);
|
||||
const handleOpenPreview = useCallback(() => {
|
||||
if (!(documentId && backendChunkId)) {
|
||||
return;
|
||||
}
|
||||
Promise.resolve(
|
||||
openPreview({
|
||||
documentId,
|
||||
backendChunkId,
|
||||
}),
|
||||
).catch(() => undefined);
|
||||
}, [backendChunkId, documentId, openPreview]);
|
||||
|
||||
const sourceLabel = (
|
||||
<>
|
||||
<span className="rounded bg-foreground/10 px-1.5 py-0.5 font-mono text-[10px] font-semibold">
|
||||
[{chunk.id}]
|
||||
</span>
|
||||
{chunk.kind === "image" ? (
|
||||
<ImageIcon className="size-3 shrink-0 text-muted-foreground" />
|
||||
) : (
|
||||
<FileTextIcon className="size-3 shrink-0 text-muted-foreground" />
|
||||
)}
|
||||
<span className="truncate font-medium" title={chunk.source}>
|
||||
{chunk.source}
|
||||
</span>
|
||||
</>
|
||||
);
|
||||
|
||||
return (
|
||||
<div
|
||||
data-slot="rag-chunk-card"
|
||||
className="rounded-md border border-foreground/10 bg-muted/40 p-2.5 text-xs"
|
||||
>
|
||||
<div className="mb-1.5 flex items-center justify-between gap-2">
|
||||
<div className="flex min-w-0 items-center gap-1.5">
|
||||
<span className="rounded bg-foreground/10 px-1.5 py-0.5 font-mono text-[10px] font-semibold">
|
||||
[{chunk.id}]
|
||||
</span>
|
||||
{chunk.kind === "image" ? (
|
||||
<ImageIcon className="size-3 shrink-0 text-muted-foreground" />
|
||||
) : (
|
||||
<FileTextIcon className="size-3 shrink-0 text-muted-foreground" />
|
||||
)}
|
||||
<span className="truncate font-medium" title={chunk.source}>
|
||||
{chunk.source}
|
||||
</span>
|
||||
</div>
|
||||
{isPreviewable ? (
|
||||
<button
|
||||
type="button"
|
||||
className={cn(
|
||||
"flex min-w-0 cursor-pointer items-center gap-1.5 rounded-sm text-left outline-none transition-colors",
|
||||
"hover:text-primary focus-visible:ring-[3px] focus-visible:ring-ring/50",
|
||||
)}
|
||||
onClick={handleOpenPreview}
|
||||
aria-label={`Open preview of ${chunk.source}`}
|
||||
title="Open preview"
|
||||
>
|
||||
{sourceLabel}
|
||||
</button>
|
||||
) : (
|
||||
<div className="flex min-w-0 items-center gap-1.5">{sourceLabel}</div>
|
||||
)}
|
||||
{meta.length > 0 ? (
|
||||
<span className="shrink-0 text-[10px] tabular-nums text-muted-foreground">
|
||||
{meta.join(" · ")}
|
||||
|
|
|
|||
|
|
@ -1,7 +1,18 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import {
|
||||
type ParsedChunk,
|
||||
parseChunks,
|
||||
} from "@/components/assistant-ui/tool-ui-search-knowledge-base";
|
||||
import { getAuthToken } from "@/features/auth/session";
|
||||
import {
|
||||
type SearchHit,
|
||||
type SearchRequest,
|
||||
listKBDocuments,
|
||||
listThreadDocuments,
|
||||
search as ragSearch,
|
||||
} from "@/features/rag/api/rag-api";
|
||||
import { apiUrl } from "@/lib/api-base";
|
||||
import { toast } from "@/lib/toast";
|
||||
import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core";
|
||||
|
|
@ -37,12 +48,12 @@ import type {
|
|||
OpenAIMessageContent,
|
||||
} from "../types/api";
|
||||
import type { ChatModelSummary } from "../types/runtime";
|
||||
import { getImageInputUnavailableReason } from "../utils/image-input-support";
|
||||
import {
|
||||
getStoredChatThread,
|
||||
listStoredChatThreads,
|
||||
updateStoredChatThread,
|
||||
} from "../utils/chat-history-storage";
|
||||
import { getImageInputUnavailableReason } from "../utils/image-input-support";
|
||||
import {
|
||||
hasClosedThinkTag,
|
||||
parseAssistantContent,
|
||||
|
|
@ -56,17 +67,6 @@ import {
|
|||
streamChatCompletions,
|
||||
validateModel,
|
||||
} from "./chat-api";
|
||||
import {
|
||||
type SearchHit,
|
||||
type SearchRequest,
|
||||
listKBDocuments,
|
||||
listThreadDocuments,
|
||||
search as ragSearch,
|
||||
} from "@/features/rag/api/rag-api";
|
||||
import {
|
||||
type ParsedChunk,
|
||||
parseChunks,
|
||||
} from "@/components/assistant-ui/tool-ui-search-knowledge-base";
|
||||
import type { RagMode, RagSource } from "./chat-settings-api";
|
||||
import {
|
||||
createOpenAIContainer,
|
||||
|
|
@ -123,8 +123,7 @@ function buildRagRequest(
|
|||
function formatRagContext(hits: SearchHit[]): string {
|
||||
const parts = hits.map((h) => {
|
||||
const name = h.filename ?? `chunk ${h.chunk_index}`;
|
||||
const pageAttr =
|
||||
h.page_number != null ? ` page="${h.page_number}"` : "";
|
||||
const pageAttr = h.page_number != null ? ` page="${h.page_number}"` : "";
|
||||
return `<source filename="${name}"${pageAttr}>\n${h.text}\n</source>`;
|
||||
});
|
||||
return `<context>\nThe following documents may help answer the user's question:\n${parts.join("\n")}\n</context>`;
|
||||
|
|
@ -251,9 +250,27 @@ interface DocumentSourcePart {
|
|||
type: "source";
|
||||
sourceType: "document";
|
||||
id: string;
|
||||
/** Display alias of `citationId`. Kept so the existing sources.tsx
|
||||
* renderer keeps working; new code SHOULD use `citationId`. NEVER
|
||||
* sent to backend as the durable chunk_id. */
|
||||
chunkId: string;
|
||||
/** Visible model-citation id (the `[N]` reference). Display only. */
|
||||
citationId: string;
|
||||
/** Durable `rag_documents.id` from tool XML `document_id=`. Null
|
||||
* when the source came from legacy XML lacking the attribute;
|
||||
* preview routing is gated off in that case. */
|
||||
documentId: string | null;
|
||||
/** Durable `rag_chunks.id` from tool XML `chunk_id=`. Null on
|
||||
* legacy XML. Sent as `?chunk_id=` to `/preview-target`. */
|
||||
backendChunkId: string | null;
|
||||
filename: string;
|
||||
page?: string;
|
||||
sourcePageIndex?: string;
|
||||
pageCharStart?: string;
|
||||
pageCharEnd?: string;
|
||||
lineStart?: string;
|
||||
lineEnd?: string;
|
||||
score?: string;
|
||||
text: string;
|
||||
}
|
||||
|
||||
|
|
@ -273,32 +290,93 @@ function extractCitedIds(text: string): Set<string> {
|
|||
return ids;
|
||||
}
|
||||
|
||||
/** Build doc-shaped source parts for chunks the model actually cited.
|
||||
function indexChunksByCitationId(
|
||||
allChunks: ParsedChunk[],
|
||||
): Map<string, ParsedChunk> {
|
||||
const byId = new Map<string, ParsedChunk>();
|
||||
for (const chunk of allChunks) {
|
||||
if (!byId.has(chunk.id)) {
|
||||
byId.set(chunk.id, chunk);
|
||||
}
|
||||
}
|
||||
return byId;
|
||||
}
|
||||
|
||||
function documentSourceIds(
|
||||
allChunks: ParsedChunk[],
|
||||
citedIds: Set<string>,
|
||||
): string[] {
|
||||
if (citedIds.size > 0) {
|
||||
return Array.from(citedIds);
|
||||
}
|
||||
return allChunks.map((chunk) => chunk.id);
|
||||
}
|
||||
|
||||
function toDocumentSourcePart(
|
||||
id: string,
|
||||
chunk: ParsedChunk,
|
||||
): DocumentSourcePart {
|
||||
const part: DocumentSourcePart = {
|
||||
type: "source",
|
||||
sourceType: "document",
|
||||
id: `rag-${id}`,
|
||||
chunkId: id,
|
||||
citationId: id,
|
||||
documentId: chunk.documentId ?? null,
|
||||
backendChunkId: chunk.backendChunkId ?? null,
|
||||
filename: chunk.source,
|
||||
text: chunk.text,
|
||||
};
|
||||
|
||||
if (chunk.page) {
|
||||
part.page = chunk.page;
|
||||
}
|
||||
if (chunk.sourcePageIndex) {
|
||||
part.sourcePageIndex = chunk.sourcePageIndex;
|
||||
}
|
||||
if (chunk.pageCharStart) {
|
||||
part.pageCharStart = chunk.pageCharStart;
|
||||
}
|
||||
if (chunk.pageCharEnd) {
|
||||
part.pageCharEnd = chunk.pageCharEnd;
|
||||
}
|
||||
if (chunk.lineStart) {
|
||||
part.lineStart = chunk.lineStart;
|
||||
}
|
||||
if (chunk.lineEnd) {
|
||||
part.lineEnd = chunk.lineEnd;
|
||||
}
|
||||
if (chunk.score) {
|
||||
part.score = chunk.score;
|
||||
}
|
||||
|
||||
return part;
|
||||
}
|
||||
|
||||
/** Build doc-shaped source parts for chunks the model cited.
|
||||
* `allChunks` is the flat union of every search_knowledge_base tool
|
||||
* result in this turn (deduped by id). Returns one part per unique
|
||||
* cited id that maps to a real chunk; hallucinated `[99]` refs without
|
||||
* a matching chunk are silently dropped. */
|
||||
* result in this turn (deduped by id). If the model forgets literal
|
||||
* `[N]` ids, fall back to retrieved chunks so source chips remain
|
||||
* visible and previewable. Hallucinated `[99]` refs without a matching
|
||||
* chunk are silently dropped. */
|
||||
function buildDocumentSourceParts(
|
||||
allChunks: ParsedChunk[],
|
||||
citedIds: Set<string>,
|
||||
): DocumentSourcePart[] {
|
||||
const byId = new Map<string, ParsedChunk>();
|
||||
for (const chunk of allChunks) {
|
||||
if (!byId.has(chunk.id)) byId.set(chunk.id, chunk);
|
||||
}
|
||||
const byId = indexChunksByCitationId(allChunks);
|
||||
const idsToShow = documentSourceIds(allChunks, citedIds);
|
||||
const out: DocumentSourcePart[] = [];
|
||||
for (const id of citedIds) {
|
||||
const emittedIds = new Set<string>();
|
||||
for (const id of idsToShow) {
|
||||
if (emittedIds.has(id)) {
|
||||
continue;
|
||||
}
|
||||
const chunk = byId.get(id);
|
||||
if (!chunk) continue;
|
||||
out.push({
|
||||
type: "source",
|
||||
sourceType: "document",
|
||||
id: `rag-${id}`,
|
||||
chunkId: id,
|
||||
filename: chunk.source,
|
||||
...(chunk.page ? { page: chunk.page } : {}),
|
||||
text: chunk.text,
|
||||
});
|
||||
if (!chunk) {
|
||||
continue;
|
||||
}
|
||||
emittedIds.add(id);
|
||||
out.push(toDocumentSourcePart(id, chunk));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
|
@ -954,7 +1032,12 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
// Re-read store after potential auto-load / model ready wait
|
||||
runtime = useChatRuntimeStore.getState();
|
||||
const { params } = runtime;
|
||||
const { supportsTools, toolsEnabled, codeToolsEnabled, imageToolsEnabled } = runtime;
|
||||
const {
|
||||
supportsTools,
|
||||
toolsEnabled,
|
||||
codeToolsEnabled,
|
||||
imageToolsEnabled,
|
||||
} = runtime;
|
||||
const externalSelection = parseExternalModelId(params.checkpoint);
|
||||
const isExternalRequest = externalSelection !== null;
|
||||
if (
|
||||
|
|
@ -993,33 +1076,30 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
throw new Error("Missing connection API key.");
|
||||
}
|
||||
|
||||
const webSearchEnabledForThisTurn =
|
||||
Boolean(
|
||||
externalProvider &&
|
||||
toolsEnabled &&
|
||||
providerSupportsBuiltinWebSearch(externalProvider.providerType),
|
||||
);
|
||||
const codeExecEnabledForThisTurn =
|
||||
Boolean(
|
||||
externalProvider &&
|
||||
externalSelection &&
|
||||
codeToolsEnabled &&
|
||||
providerSupportsBuiltinCodeExecution(
|
||||
externalProvider.providerType,
|
||||
externalSelection.modelId,
|
||||
externalProvider.baseUrl,
|
||||
),
|
||||
);
|
||||
const webSearchEnabledForThisTurn = Boolean(
|
||||
externalProvider &&
|
||||
toolsEnabled &&
|
||||
providerSupportsBuiltinWebSearch(externalProvider.providerType),
|
||||
);
|
||||
const codeExecEnabledForThisTurn = Boolean(
|
||||
externalProvider &&
|
||||
externalSelection &&
|
||||
codeToolsEnabled &&
|
||||
providerSupportsBuiltinCodeExecution(
|
||||
externalProvider.providerType,
|
||||
externalSelection.modelId,
|
||||
externalProvider.baseUrl,
|
||||
),
|
||||
);
|
||||
// web_fetch shares the Search pill with web_search (no separate
|
||||
// UI toggle), so it follows toolsEnabled. Anthropic is the only
|
||||
// provider that ships it today; on others providerSupportsBuiltinWebFetch
|
||||
// returns false and this stays inert.
|
||||
const webFetchEnabledForThisTurn =
|
||||
Boolean(
|
||||
externalProvider &&
|
||||
toolsEnabled &&
|
||||
providerSupportsBuiltinWebFetch(externalProvider.providerType),
|
||||
);
|
||||
const webFetchEnabledForThisTurn = Boolean(
|
||||
externalProvider &&
|
||||
toolsEnabled &&
|
||||
providerSupportsBuiltinWebFetch(externalProvider.providerType),
|
||||
);
|
||||
const providerShipsWebFetch = Boolean(
|
||||
externalProvider &&
|
||||
providerSupportsBuiltinWebFetch(externalProvider.providerType),
|
||||
|
|
@ -1049,7 +1129,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
// entirely so retrieval only happens via the LLM-invoked
|
||||
// search_knowledge_base tool. Flip back to true to restore the
|
||||
// always-on grounding for external providers / non-tool models.
|
||||
const RAG_PREFETCH_ENABLED = false;
|
||||
const ragPrefetchEnabled = false;
|
||||
|
||||
const ragSource = runtime.ragSource;
|
||||
const ragToolEnabled = runtime.ragToolEnabled;
|
||||
|
|
@ -1115,13 +1195,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
});
|
||||
}
|
||||
|
||||
if (RAG_PREFETCH_ENABLED && ragToolEnabled && ragSource.kind !== "off") {
|
||||
if (ragPrefetchEnabled && ragToolEnabled && ragSource.kind !== "off") {
|
||||
const lastUser = [...outboundMessages]
|
||||
.reverse()
|
||||
.find((m) => m.role === "user");
|
||||
const queryText = lastUser
|
||||
? extractMessageText(lastUser.content)
|
||||
: "";
|
||||
const queryText = lastUser ? extractMessageText(lastUser.content) : "";
|
||||
if (queryText.trim()) {
|
||||
const ragReq = buildRagRequest(
|
||||
ragSource,
|
||||
|
|
@ -1145,8 +1223,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
outboundMessages[0]?.role === "system" &&
|
||||
typeof outboundMessages[0].content === "string"
|
||||
) {
|
||||
outboundMessages[0].content =
|
||||
`${block}\n\n${outboundMessages[0].content}`;
|
||||
outboundMessages[0].content = `${block}\n\n${outboundMessages[0].content}`;
|
||||
} else {
|
||||
outboundMessages.unshift({ role: "system", content: block });
|
||||
}
|
||||
|
|
@ -1454,7 +1531,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
? reasoningEffort
|
||||
: "low";
|
||||
const externalReasoningEnabled =
|
||||
!externalReasoningCaps.supportsReasoningOff ? true : reasoningEnabled;
|
||||
externalReasoningCaps.supportsReasoningOff ? reasoningEnabled : true;
|
||||
const buildRequestPayload = async (
|
||||
forceRefreshPublicKey = false,
|
||||
): Promise<OpenAIChatCompletionsRequest> => {
|
||||
|
|
@ -1540,8 +1617,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
) {
|
||||
void updateStoredChatThreadEventually(t.id, {
|
||||
openaiCodeExecContainerId: null,
|
||||
})
|
||||
.catch(() => {});
|
||||
}).catch(() => {});
|
||||
continue;
|
||||
}
|
||||
openaiCodeExecContainerId = t.openaiCodeExecContainerId;
|
||||
|
|
@ -1585,8 +1661,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
openaiCodeExecContainerId = created.id;
|
||||
void updateStoredChatThreadEventually(resolvedThreadId, {
|
||||
openaiCodeExecContainerId: created.id,
|
||||
})
|
||||
.catch(() => {});
|
||||
}).catch(() => {});
|
||||
} catch {
|
||||
// Fall back to backend's container_auto path on
|
||||
// failure — keeps the chat moving; the next turn
|
||||
|
|
@ -1699,7 +1774,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
// attaches `cache_control.ttl` when the value is one of
|
||||
// "5m" / "1h" (see external_provider.py near line 1375),
|
||||
// so unknown values are a no-op end-to-end.
|
||||
...(supportsProviderPromptCacheTtl(externalProvider.providerType) &&
|
||||
...(supportsProviderPromptCacheTtl(
|
||||
externalProvider.providerType,
|
||||
) &&
|
||||
(externalProvider.enablePromptCaching ?? true) &&
|
||||
isPromptCacheTtl(externalProvider.promptCacheTtl)
|
||||
? { prompt_cache_ttl: externalProvider.promptCacheTtl }
|
||||
|
|
@ -1744,8 +1821,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
...(supportsPreserveThinking
|
||||
? { preserve_thinking: preserveThinking }
|
||||
: {}),
|
||||
...(supportsTools
|
||||
&& (toolsEnabled || codeToolsEnabled || ragToolPathTaken)
|
||||
...(supportsTools &&
|
||||
(toolsEnabled || codeToolsEnabled || ragToolPathTaken)
|
||||
? {
|
||||
enable_tools: true,
|
||||
enabled_tools: [
|
||||
|
|
@ -1760,9 +1837,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
? {
|
||||
rag_scope: {
|
||||
kb_id:
|
||||
ragSource.kind === "kb"
|
||||
? ragSource.kbId
|
||||
: null,
|
||||
ragSource.kind === "kb" ? ragSource.kbId : null,
|
||||
thread_id:
|
||||
ragSource.kind === "thread"
|
||||
? (resolvedThreadId ?? null)
|
||||
|
|
@ -1840,8 +1915,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
: "openaiCodeExecContainerId";
|
||||
void updateStoredChatThreadEventually(resolvedThreadId, {
|
||||
[field]: null,
|
||||
})
|
||||
.catch(() => {});
|
||||
}).catch(() => {});
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
|
@ -2026,11 +2100,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
}
|
||||
|
||||
if (reasoning) {
|
||||
if (!reasoningContentOpen) {
|
||||
if (reasoningContentOpen) {
|
||||
cumulativeText += reasoning;
|
||||
} else {
|
||||
cumulativeText += `<think>${reasoning}`;
|
||||
reasoningContentOpen = true;
|
||||
} else {
|
||||
cumulativeText += reasoning;
|
||||
}
|
||||
}
|
||||
if (delta) {
|
||||
|
|
@ -2125,8 +2199,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
});
|
||||
|
||||
// RAG: flatten chunks across every search_knowledge_base call this
|
||||
// turn, then emit doc-source parts only for ids the model actually
|
||||
// cited as [N] in its final reply.
|
||||
// turn, then emit previewable doc-source chips for cited chunks.
|
||||
// If the model omits literal [N] ids, show the retrieved chunks so
|
||||
// the answer still has a visible citation/preview affordance.
|
||||
const ragChunks = toolCallParts.flatMap((tc) => {
|
||||
if (tc.toolName !== "search_knowledge_base" || !tc.result) {
|
||||
return [];
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -4,6 +4,7 @@
|
|||
import { authFetch, getAuthToken } from "@/features/auth";
|
||||
import { apiUrl } from "@/lib/api-base";
|
||||
import { formatFastApiDetail } from "@/lib/format-fastapi-error";
|
||||
import { EventSourcePolyfill } from "event-source-polyfill";
|
||||
|
||||
export type ChunkingStrategy = "standard" | "late";
|
||||
export type KBMode = "text" | "multimodal";
|
||||
|
|
@ -47,6 +48,74 @@ export interface SearchHit {
|
|||
filename: string | null;
|
||||
kind?: "text" | "image" | "caption";
|
||||
image_url?: string | null;
|
||||
source_page_index?: number | null;
|
||||
page_char_start?: number | null;
|
||||
page_char_end?: number | null;
|
||||
line_start?: number | null;
|
||||
line_end?: number | null;
|
||||
}
|
||||
|
||||
// --- Preview target (durable backend-id routing) ---
|
||||
|
||||
export type PreviewMediaKind =
|
||||
| "pdf"
|
||||
| "text"
|
||||
| "docx"
|
||||
| "html"
|
||||
| "image"
|
||||
| "unknown";
|
||||
|
||||
export type PreviewChunkKind = "text" | "image" | "caption";
|
||||
|
||||
export interface PreviewPdfRegion {
|
||||
pageIndex: number;
|
||||
pageNumber: number | null;
|
||||
x: number;
|
||||
y: number;
|
||||
width: number;
|
||||
height: number;
|
||||
confidence: "exact";
|
||||
source: string;
|
||||
}
|
||||
|
||||
export interface PreviewTarget {
|
||||
documentId: string;
|
||||
filename: string;
|
||||
contentType: string | null;
|
||||
mediaKind: PreviewMediaKind;
|
||||
byteSize: number;
|
||||
status: string;
|
||||
kbId: string | null;
|
||||
threadId: string | null;
|
||||
chunkId: string | null;
|
||||
chunkIndex: number | null;
|
||||
targetPage: number | null;
|
||||
snippet: string | null;
|
||||
kind: PreviewChunkKind | null;
|
||||
imageUrl: string | null;
|
||||
sourcePageIndex: number | null;
|
||||
pageCharStart: number | null;
|
||||
pageCharEnd: number | null;
|
||||
lineStart: number | null;
|
||||
lineEnd: number | null;
|
||||
pdfRegions: PreviewPdfRegion[];
|
||||
}
|
||||
|
||||
export interface PreviewFileUrl {
|
||||
url: string;
|
||||
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 {
|
||||
|
|
@ -63,7 +132,13 @@ export interface SearchRequest {
|
|||
}
|
||||
|
||||
export type JobEvent =
|
||||
| { type: "status"; status: string; stage?: string | null; progress?: number; error?: string | null }
|
||||
| {
|
||||
type: "status";
|
||||
status: string;
|
||||
stage?: string | null;
|
||||
progress?: number;
|
||||
error?: string | null;
|
||||
}
|
||||
| { type: "progress"; stage: string; progress: number }
|
||||
| { type: "complete"; num_chunks: number }
|
||||
| { type: "error"; error: string };
|
||||
|
|
@ -72,7 +147,9 @@ function parseErrorText(status: number, body: unknown): string {
|
|||
if (body && typeof body === "object") {
|
||||
const detail = (body as { detail?: unknown }).detail;
|
||||
const formatted = formatFastApiDetail(detail);
|
||||
if (formatted) return formatted;
|
||||
if (formatted) {
|
||||
return formatted;
|
||||
}
|
||||
}
|
||||
return `Request failed (${status})`;
|
||||
}
|
||||
|
|
@ -95,7 +172,9 @@ async function throwOnError(response: Response): Promise<void> {
|
|||
|
||||
export async function listKnowledgeBases(): Promise<KnowledgeBase[]> {
|
||||
const response = await authFetch("/api/rag/knowledge-bases");
|
||||
const body = await parseJsonOrThrow<{ knowledge_bases: KnowledgeBase[] }>(response);
|
||||
const body = await parseJsonOrThrow<{ knowledge_bases: KnowledgeBase[] }>(
|
||||
response,
|
||||
);
|
||||
return body.knowledge_bases;
|
||||
}
|
||||
|
||||
|
|
@ -136,7 +215,9 @@ export async function listKBDocuments(kbId: string): Promise<RagDocument[]> {
|
|||
return body.documents;
|
||||
}
|
||||
|
||||
export async function listThreadDocuments(threadId: string): Promise<RagDocument[]> {
|
||||
export async function listThreadDocuments(
|
||||
threadId: string,
|
||||
): Promise<RagDocument[]> {
|
||||
const response = await authFetch(
|
||||
`/api/rag/threads/${encodeURIComponent(threadId)}/documents`,
|
||||
);
|
||||
|
|
@ -187,7 +268,9 @@ export interface ThreadIndexSummary {
|
|||
|
||||
export async function listThreadIndexes(): Promise<ThreadIndexSummary[]> {
|
||||
const response = await authFetch("/api/rag/thread-indexes");
|
||||
const body = await parseJsonOrThrow<{ threads: ThreadIndexSummary[] }>(response);
|
||||
const body = await parseJsonOrThrow<{ threads: ThreadIndexSummary[] }>(
|
||||
response,
|
||||
);
|
||||
return body.threads;
|
||||
}
|
||||
|
||||
|
|
@ -342,7 +425,8 @@ export async function search(req: SearchRequest): Promise<SearchHit[]> {
|
|||
// --- Ingestion SSE ---
|
||||
|
||||
/** Subscribe to a job's SSE stream; returns an unsubscribe fn.
|
||||
* Token goes via ?token= since EventSource cannot send headers. */
|
||||
* Use the EventSource polyfill so the bearer token rides in an
|
||||
* Authorization header instead of leaking through URL query params. */
|
||||
export function subscribeToJobEvents(
|
||||
jobId: string,
|
||||
handlers: {
|
||||
|
|
@ -352,9 +436,17 @@ export function subscribeToJobEvents(
|
|||
},
|
||||
): () => void {
|
||||
const token = getAuthToken();
|
||||
const params = token ? `?token=${encodeURIComponent(token)}` : "";
|
||||
const url = apiUrl(`/api/rag/jobs/${encodeURIComponent(jobId)}/events${params}`);
|
||||
const source = new EventSource(url);
|
||||
const url = apiUrl(`/api/rag/jobs/${encodeURIComponent(jobId)}/events`);
|
||||
const source = new EventSourcePolyfill(
|
||||
url,
|
||||
token
|
||||
? {
|
||||
headers: {
|
||||
Authorization: `Bearer ${token}`,
|
||||
},
|
||||
}
|
||||
: undefined,
|
||||
);
|
||||
|
||||
source.onmessage = (e) => {
|
||||
try {
|
||||
|
|
@ -381,3 +473,68 @@ export function subscribeToJobEvents(
|
|||
handlers.onClose?.();
|
||||
};
|
||||
}
|
||||
|
||||
// --- Preview target + blob ---
|
||||
|
||||
/** Fetch the preview-target metadata for a document. When `chunkId` is
|
||||
* provided it must belong to `documentId`; mismatch collapses to 404
|
||||
* with the same "Document not found" body as missing/unauthorized. */
|
||||
export async function fetchPreviewTarget(
|
||||
documentId: string,
|
||||
chunkId?: string | null,
|
||||
): Promise<PreviewTarget> {
|
||||
const params = chunkId ? `?chunk_id=${encodeURIComponent(chunkId)}` : "";
|
||||
const response = await authFetch(
|
||||
`/api/rag/documents/${encodeURIComponent(documentId)}/preview-target${params}`,
|
||||
);
|
||||
return parseJsonOrThrow<PreviewTarget>(response);
|
||||
}
|
||||
|
||||
/** Mint a short-lived URL that PDF.js can range-load without putting
|
||||
* the user's bearer token in a query string. */
|
||||
export async function fetchPreviewFileUrl(
|
||||
documentId: string,
|
||||
signal?: AbortSignal,
|
||||
): Promise<PreviewFileUrl> {
|
||||
const response = await authFetch(
|
||||
`/api/rag/documents/${encodeURIComponent(documentId)}/file-url`,
|
||||
signal ? { signal } : undefined,
|
||||
);
|
||||
const body = await parseJsonOrThrow<PreviewFileUrl>(response);
|
||||
return {
|
||||
...body,
|
||||
url: apiUrl(body.url),
|
||||
};
|
||||
}
|
||||
|
||||
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` (so
|
||||
* the bearer token rides in the Authorization header — never a query
|
||||
* string). The caller (preview-store) creates and revokes the object
|
||||
* URL so blob lifecycle stays in one place.
|
||||
*
|
||||
* `signal` lets the caller abort the fetch when the user switches
|
||||
* documents mid-load. */
|
||||
export async function fetchPreviewFileBlob(
|
||||
documentId: string,
|
||||
signal?: AbortSignal,
|
||||
): Promise<Blob> {
|
||||
const response = await authFetch(
|
||||
`/api/rag/documents/${encodeURIComponent(documentId)}/file`,
|
||||
signal ? { signal } : undefined,
|
||||
);
|
||||
if (!response.ok) {
|
||||
const body = await response.json().catch(() => null);
|
||||
throw new Error(parseErrorText(response.status, body));
|
||||
}
|
||||
return response.blob();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import { Button } from "@/components/ui/button";
|
|||
import { cn } from "@/lib/utils";
|
||||
import { Delete02Icon } from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import type { KeyboardEvent, MouseEvent } from "react";
|
||||
import type { RagDocument } from "../api/rag-api";
|
||||
|
||||
const STATUS_VARIANT: Record<
|
||||
|
|
@ -27,20 +28,59 @@ function humanBytes(bytes: number): string {
|
|||
export function DocumentRow({
|
||||
doc,
|
||||
onDelete,
|
||||
onPreview,
|
||||
rightSlot,
|
||||
className,
|
||||
}: {
|
||||
doc: RagDocument;
|
||||
onDelete?: () => void;
|
||||
/** Fired when the row body (not the delete button) is clicked.
|
||||
* Per decision Q9, callers should only pass this for completed
|
||||
* documents. */
|
||||
onPreview?: () => void;
|
||||
rightSlot?: React.ReactNode;
|
||||
className?: string;
|
||||
}) {
|
||||
const isPreviewable = !!onPreview;
|
||||
|
||||
const handleRowClick = (e: MouseEvent<HTMLDivElement>) => {
|
||||
if (!onPreview) return;
|
||||
// If a button/anchor/control was clicked (e.g. the delete icon),
|
||||
// skip preview — let that handler win. Buttons inside this row
|
||||
// additionally call stopPropagation, but this is defense in depth
|
||||
// for any descendant Button that forgets to.
|
||||
const target = e.target as HTMLElement | null;
|
||||
const interactive = target?.closest("button, a, [role=button]");
|
||||
if (interactive && interactive !== e.currentTarget) return;
|
||||
onPreview();
|
||||
};
|
||||
|
||||
const handleRowKey = (e: KeyboardEvent<HTMLDivElement>) => {
|
||||
if (!onPreview) return;
|
||||
if (e.target !== e.currentTarget) return; // ignore child key events
|
||||
if (e.key === "Enter" || e.key === " ") {
|
||||
e.preventDefault();
|
||||
onPreview();
|
||||
}
|
||||
};
|
||||
|
||||
return (
|
||||
<div
|
||||
className={cn(
|
||||
"flex items-center justify-between gap-3 rounded-md border border-border/60 px-3 py-2",
|
||||
isPreviewable &&
|
||||
"cursor-pointer outline-none focus-visible:border-ring focus-visible:ring-[3px] focus-visible:ring-ring/50 hover:bg-muted/40",
|
||||
className,
|
||||
)}
|
||||
{...(isPreviewable
|
||||
? {
|
||||
role: "button",
|
||||
tabIndex: 0,
|
||||
onClick: handleRowClick,
|
||||
onKeyDown: handleRowKey,
|
||||
"aria-label": `Open preview of ${doc.filename}`,
|
||||
}
|
||||
: {})}
|
||||
>
|
||||
<div className="flex min-w-0 flex-1 flex-col gap-1">
|
||||
<div className="flex items-center gap-2">
|
||||
|
|
@ -67,7 +107,12 @@ export function DocumentRow({
|
|||
variant="ghost"
|
||||
size="icon"
|
||||
aria-label="Delete document"
|
||||
onClick={onDelete}
|
||||
onClick={(e) => {
|
||||
// Stop propagation so the row's onClick (preview open)
|
||||
// does not fire when the user is asking to delete.
|
||||
e.stopPropagation();
|
||||
onDelete();
|
||||
}}
|
||||
className="text-muted-foreground hover:text-destructive"
|
||||
>
|
||||
<HugeiconsIcon icon={Delete02Icon} size={16} />
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import { Separator } from "@/components/ui/separator";
|
|||
import { useState } from "react";
|
||||
import type { KnowledgeBase } from "../api/rag-api";
|
||||
import { useKBDocuments } from "../hooks/use-kb-documents";
|
||||
import { usePreviewStore } from "../stores/preview-store";
|
||||
import { DocumentRow } from "./document-row";
|
||||
import { DocumentUploadDropzone } from "./document-upload-dropzone";
|
||||
import { IngestionProgress } from "./ingestion-progress";
|
||||
|
|
@ -14,10 +15,11 @@ import { KBReconfigureDialog } from "./kb-reconfigure-dialog";
|
|||
|
||||
export function KBDetailPanel({ kb }: { kb: KnowledgeBase }) {
|
||||
const { documents, loading, error, upload, remove } = useKBDocuments(kb.id);
|
||||
const [activeJobsByDoc, setActiveJobsByDoc] = useState<Record<string, string>>(
|
||||
{},
|
||||
);
|
||||
const [activeJobsByDoc, setActiveJobsByDoc] = useState<
|
||||
Record<string, string>
|
||||
>({});
|
||||
const [reconfigureOpen, setReconfigureOpen] = useState(false);
|
||||
const openPreview = usePreviewStore((s) => s.open);
|
||||
|
||||
const handleFiles = async (files: File[]) => {
|
||||
for (const file of files) {
|
||||
|
|
@ -70,9 +72,7 @@ export function KBDetailPanel({ kb }: { kb: KnowledgeBase }) {
|
|||
</span>
|
||||
</div>
|
||||
|
||||
{error ? (
|
||||
<div className="text-xs text-destructive">{error}</div>
|
||||
) : null}
|
||||
{error ? <div className="text-xs text-destructive">{error}</div> : null}
|
||||
|
||||
<ScrollArea className="flex-1">
|
||||
<div className="flex flex-col gap-2 pr-2">
|
||||
|
|
@ -92,6 +92,16 @@ export function KBDetailPanel({ kb }: { kb: KnowledgeBase }) {
|
|||
onDelete={() => {
|
||||
void remove(doc.id);
|
||||
}}
|
||||
// Per decision Q9: only completed documents open a
|
||||
// preview. Pending/running/failed rows degrade to
|
||||
// non-interactive (no onPreview).
|
||||
onPreview={
|
||||
doc.status === "completed"
|
||||
? () => {
|
||||
void openPreview({ documentId: doc.id });
|
||||
}
|
||||
: undefined
|
||||
}
|
||||
rightSlot={
|
||||
showProgress ? (
|
||||
<IngestionProgress jobId={jobId} className="mt-1" />
|
||||
|
|
|
|||
235
studio/frontend/src/features/rag/components/preview-panel.tsx
Normal file
235
studio/frontend/src/features/rag/components/preview-panel.tsx
Normal file
|
|
@ -0,0 +1,235 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { Button } from "@/components/ui/button";
|
||||
import {
|
||||
Sheet,
|
||||
SheetContent,
|
||||
SheetDescription,
|
||||
SheetHeader,
|
||||
SheetTitle,
|
||||
} from "@/components/ui/sheet";
|
||||
import { LoaderIcon, XIcon } from "lucide-react";
|
||||
import {
|
||||
type FC,
|
||||
type ReactNode,
|
||||
useEffect,
|
||||
useSyncExternalStore,
|
||||
} from "react";
|
||||
import type { PreviewTarget } from "../api/rag-api";
|
||||
import {
|
||||
type PreviewLoadStatus,
|
||||
isInlineBlobAllowed,
|
||||
usePreviewStore,
|
||||
} from "../stores/preview-store";
|
||||
import { PreviewPdfView } from "./preview-pdf-view";
|
||||
import { PreviewTextView } from "./preview-text-view";
|
||||
import { PreviewUnavailable } from "./preview-unavailable";
|
||||
|
||||
interface PreviewPanelProps {
|
||||
/** Whether the panel is currently being shown in its host slot.
|
||||
* When the host hides the slot (e.g. the user closes both
|
||||
* settings and preview from the chat header), we run the close
|
||||
* side-effect so the blob URL is revoked. */
|
||||
open: boolean;
|
||||
disableDrawer?: boolean;
|
||||
}
|
||||
|
||||
const LG_BREAKPOINT = 1024;
|
||||
const MEDIA_QUERY = `(max-width: ${LG_BREAKPOINT - 1}px)`;
|
||||
|
||||
function getLgSnapshot(): boolean {
|
||||
if (
|
||||
typeof window === "undefined" ||
|
||||
typeof window.matchMedia !== "function"
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
return window.matchMedia(MEDIA_QUERY).matches;
|
||||
}
|
||||
|
||||
function lgSubscribe(callback: () => void): () => void {
|
||||
if (
|
||||
typeof window === "undefined" ||
|
||||
typeof window.matchMedia !== "function"
|
||||
) {
|
||||
return () => undefined;
|
||||
}
|
||||
const mql = window.matchMedia(MEDIA_QUERY);
|
||||
mql.addEventListener("change", callback);
|
||||
return () => mql.removeEventListener("change", callback);
|
||||
}
|
||||
|
||||
function useIsViewportSqueezed(): boolean {
|
||||
return useSyncExternalStore(lgSubscribe, getLgSnapshot, () => false);
|
||||
}
|
||||
|
||||
interface PreviewBodyArgs {
|
||||
error: string | null;
|
||||
previewBlob: Blob | null;
|
||||
previewFileUrl: string | null;
|
||||
status: PreviewLoadStatus;
|
||||
target: PreviewTarget | null;
|
||||
}
|
||||
|
||||
function renderPreviewBody({
|
||||
error,
|
||||
previewBlob,
|
||||
previewFileUrl,
|
||||
status,
|
||||
target,
|
||||
}: PreviewBodyArgs): ReactNode {
|
||||
if (status === "loading") {
|
||||
return (
|
||||
<div className="flex h-full items-center justify-center gap-2 text-xs text-muted-foreground">
|
||||
<LoaderIcon className="size-3.5 animate-spin" />
|
||||
Loading preview…
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
if (status === "error") {
|
||||
// Treat 404s as "document missing" so a stale citation reads
|
||||
// like "no longer available" rather than a generic error.
|
||||
const isMissing = (error ?? "").toLowerCase().includes("not found");
|
||||
return (
|
||||
<PreviewUnavailable
|
||||
filename={target?.filename}
|
||||
reason={error ?? "Preview unavailable."}
|
||||
variant={isMissing ? "missing" : "error"}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
if (status === "ready" && target) {
|
||||
const pdfFile = previewFileUrl ?? previewBlob;
|
||||
if (
|
||||
target.mediaKind === "pdf" &&
|
||||
pdfFile &&
|
||||
isInlineBlobAllowed(target.mediaKind)
|
||||
) {
|
||||
return <PreviewPdfView target={target} file={pdfFile} />;
|
||||
}
|
||||
|
||||
// text / image / docx / html / unknown — all routed through
|
||||
// text-view. text gets the snippet rendered inline; docx/html
|
||||
// /unknown skip inline-render entirely (contracts §5.4 + Risk #3).
|
||||
return <PreviewTextView target={target} />;
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
/** Body-only renderer. The host slot (desktop aside or mobile sheet)
|
||||
* is owned by the host (chat-settings panel slot, kb-detail panel,
|
||||
* etc.) — this component is purely the content of the right slot. */
|
||||
export const PreviewPanel: FC<PreviewPanelProps> = ({
|
||||
open,
|
||||
disableDrawer = false,
|
||||
}) => {
|
||||
const target = usePreviewStore((s) => s.target);
|
||||
const previewBlob = usePreviewStore((s) => s.previewBlob);
|
||||
const previewFileUrl = usePreviewStore((s) => s.previewFileUrl);
|
||||
const status = usePreviewStore((s) => s.status);
|
||||
const error = usePreviewStore((s) => s.error);
|
||||
const close = usePreviewStore((s) => s.close);
|
||||
const isSqueezed = useIsViewportSqueezed() && !disableDrawer;
|
||||
|
||||
// Unmount + visibility cleanup: when the panel is hidden or
|
||||
// unmounted, revoke the live object URL (contracts §5.5).
|
||||
useEffect(() => {
|
||||
if (!open) {
|
||||
close();
|
||||
}
|
||||
}, [open, close]);
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
// Component truly unmounting (e.g. navigation away). Cleanup
|
||||
// anything still live.
|
||||
usePreviewStore.getState().close();
|
||||
};
|
||||
}, []);
|
||||
|
||||
// Keyboard accessibility: ESC closes the preview.
|
||||
useEffect(() => {
|
||||
if (!open) {
|
||||
return;
|
||||
}
|
||||
const onKey = (e: KeyboardEvent) => {
|
||||
if (e.key === "Escape") {
|
||||
e.preventDefault();
|
||||
close();
|
||||
}
|
||||
};
|
||||
document.addEventListener("keydown", onKey);
|
||||
return () => document.removeEventListener("keydown", onKey);
|
||||
}, [open, close]);
|
||||
|
||||
const body = renderPreviewBody({
|
||||
error,
|
||||
previewBlob,
|
||||
previewFileUrl,
|
||||
status,
|
||||
target,
|
||||
});
|
||||
|
||||
const renderedContent = (
|
||||
<section
|
||||
aria-label="Document preview"
|
||||
className="flex h-full flex-col overflow-hidden bg-panel-surface/85 dark:bg-background/85 backdrop-blur-lg text-panel-surface-fg border border-border/40 shadow-lg menu-soft-surface"
|
||||
>
|
||||
<div className="flex items-center justify-between gap-2 border-b border-border/60 px-3 py-2 font-heading">
|
||||
<div className="flex items-center gap-1.5">
|
||||
<span
|
||||
className="h-2 w-2 animate-pulse rounded-full bg-primary [--pulse-color:color-mix(in_oklab,var(--primary)_35%,transparent)]"
|
||||
aria-hidden="true"
|
||||
/>
|
||||
<span className="text-sm font-semibold">Preview</span>
|
||||
</div>
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
onClick={close}
|
||||
aria-label="Close preview"
|
||||
className="size-7"
|
||||
>
|
||||
<XIcon className="size-3.5" />
|
||||
</Button>
|
||||
</div>
|
||||
<div className="min-h-0 flex-1 overflow-hidden">{body}</div>
|
||||
</section>
|
||||
);
|
||||
|
||||
if (isSqueezed) {
|
||||
return (
|
||||
<>
|
||||
<div className="hidden" aria-hidden="true" />
|
||||
<Sheet
|
||||
open={open}
|
||||
onOpenChange={(next) => {
|
||||
if (!next) {
|
||||
close();
|
||||
}
|
||||
}}
|
||||
>
|
||||
<SheetContent
|
||||
side="right"
|
||||
showCloseButton={false}
|
||||
overlayClassName="bg-background/35 supports-backdrop-filter:backdrop-blur-[1px]"
|
||||
className="preview-sheet-content p-0 font-heading data-[side=right]:w-full data-[side=right]:sm:max-w-md"
|
||||
>
|
||||
<SheetHeader className="sr-only">
|
||||
<SheetTitle>Document preview</SheetTitle>
|
||||
<SheetDescription>
|
||||
Preview of the active document citation
|
||||
</SheetDescription>
|
||||
</SheetHeader>
|
||||
<div className="flex h-full flex-col">{renderedContent}</div>
|
||||
</SheetContent>
|
||||
</Sheet>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
return renderedContent;
|
||||
};
|
||||
579
studio/frontend/src/features/rag/components/preview-pdf-view.tsx
Normal file
579
studio/frontend/src/features/rag/components/preview-pdf-view.tsx
Normal file
|
|
@ -0,0 +1,579 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import { cn } from "@/lib/utils";
|
||||
import {
|
||||
ChevronLeftIcon,
|
||||
ChevronRightIcon,
|
||||
CopyIcon,
|
||||
LoaderIcon,
|
||||
RotateCcwIcon,
|
||||
SearchIcon,
|
||||
ZoomInIcon,
|
||||
ZoomOutIcon,
|
||||
} from "lucide-react";
|
||||
import {
|
||||
type CSSProperties,
|
||||
type FC,
|
||||
useCallback,
|
||||
useEffect,
|
||||
useId,
|
||||
useMemo,
|
||||
useRef,
|
||||
useState,
|
||||
} from "react";
|
||||
import { Document, Page, pdfjs } from "react-pdf";
|
||||
import "react-pdf/dist/Page/AnnotationLayer.css";
|
||||
import "react-pdf/dist/Page/TextLayer.css";
|
||||
import type { PreviewPdfRegion, PreviewTarget } from "../api/rag-api";
|
||||
import { PreviewUnavailable } from "./preview-unavailable";
|
||||
|
||||
// Configure pdfjs worker in the same module where react-pdf is used,
|
||||
// per the react-pdf README. `import.meta.url` resolves to the JS bundle
|
||||
// containing this module, and Vite (+ Tauri) rewrites the URL during
|
||||
// build so the worker is co-located with the chunk.
|
||||
pdfjs.GlobalWorkerOptions.workerSrc = new URL(
|
||||
"pdfjs-dist/build/pdf.worker.min.mjs",
|
||||
import.meta.url,
|
||||
).toString();
|
||||
|
||||
type PreviewPdfFile = Blob | string;
|
||||
|
||||
interface PreviewPdfViewProps {
|
||||
target: PreviewTarget;
|
||||
file: PreviewPdfFile;
|
||||
}
|
||||
|
||||
type LoadSuccess = { numPages: number };
|
||||
type PdfLightThemeStyle = CSSProperties & Record<`--${string}`, string>;
|
||||
|
||||
const RESIZE_DEBOUNCE_MS = 100;
|
||||
const MIN_PDF_WIDTH = 280;
|
||||
// Body has p-2 (8px each side) + stable scrollbar gutter (~10px) + a tiny
|
||||
// breathing margin so the page render doesn't kiss the scrollbar.
|
||||
const PDF_BODY_GUTTER_PX = 28;
|
||||
const PDF_THUMBNAIL_WIDTH = 64;
|
||||
|
||||
const PDF_LIGHT_THEME_STYLE: PdfLightThemeStyle = {
|
||||
"--background": "oklch(1 0 0)",
|
||||
"--foreground": "oklch(0.2686 0 0)",
|
||||
"--card": "oklch(1 0 0)",
|
||||
"--card-foreground": "oklch(0.1281 0.0179 169.2764)",
|
||||
"--popover": "oklch(1 0 0)",
|
||||
"--popover-foreground": "oklch(0.1281 0.0179 169.2764)",
|
||||
"--primary": "#17b88b",
|
||||
"--primary-foreground": "oklch(1 0 0)",
|
||||
"--secondary": "oklch(0.9596 0.0275 167.8295)",
|
||||
"--secondary-foreground": "oklch(0.2868 0.0649 159.9823)",
|
||||
"--muted": "oklch(0.9702 0 0)",
|
||||
"--muted-foreground": "oklch(0.5486 0 0)",
|
||||
"--accent": "oklch(0.9596 0.0275 167.8295)",
|
||||
"--accent-foreground": "oklch(0.2868 0.0649 159.9823)",
|
||||
"--border": "oklch(0.9208 0.0101 164.8536)",
|
||||
"--input": "oklch(0.9208 0.0101 164.8536)",
|
||||
"--ring": "#17b88b",
|
||||
colorScheme: "light",
|
||||
};
|
||||
|
||||
function escapeHtml(value: string): string {
|
||||
return value
|
||||
.replace(/&/g, "&")
|
||||
.replace(/</g, "<")
|
||||
.replace(/>/g, ">")
|
||||
.replace(/"/g, """);
|
||||
}
|
||||
|
||||
function markFirstMatch(text: string, needle: string): string | null {
|
||||
const trimmed = needle.trim();
|
||||
if (trimmed.length < 2) {
|
||||
return null;
|
||||
}
|
||||
const lower = text.toLowerCase();
|
||||
const start = lower.indexOf(trimmed.toLowerCase());
|
||||
if (start < 0) {
|
||||
return null;
|
||||
}
|
||||
const end = start + trimmed.length;
|
||||
return `${escapeHtml(text.slice(0, start))}<mark>${escapeHtml(
|
||||
text.slice(start, end),
|
||||
)}</mark>${escapeHtml(text.slice(end))}`;
|
||||
}
|
||||
|
||||
// Keep text-layer highlighting opt-in. Citation snippets render in the card
|
||||
// below; using them here would mark common words across unrelated PDF text.
|
||||
function highlightPdfText(text: string, searchTerm: string): string {
|
||||
const trimmed = searchTerm.trim();
|
||||
if (trimmed.length < 2) {
|
||||
return escapeHtml(text);
|
||||
}
|
||||
const searchHit = markFirstMatch(text, trimmed);
|
||||
if (searchHit) {
|
||||
return searchHit;
|
||||
}
|
||||
return escapeHtml(text);
|
||||
}
|
||||
|
||||
function regionIsOnPage(region: PreviewPdfRegion, pageNumber: number): boolean {
|
||||
if (region.confidence !== "exact") {
|
||||
return false;
|
||||
}
|
||||
if (region.pageNumber != null) {
|
||||
return region.pageNumber === pageNumber;
|
||||
}
|
||||
return region.pageIndex === pageNumber - 1;
|
||||
}
|
||||
|
||||
interface PdfThumbnailProps {
|
||||
pageNumber: number;
|
||||
active: boolean;
|
||||
onSelect: (pageNumber: number) => void;
|
||||
}
|
||||
|
||||
/** Lazy thumbnail rendered via IntersectionObserver — only mounts the
|
||||
* inner <Page> when scrolled into view (or close to it), so large PDFs
|
||||
* stay responsive even when the rail caps at 80 buttons. */
|
||||
const PdfThumbnail: FC<PdfThumbnailProps> = ({
|
||||
pageNumber,
|
||||
active,
|
||||
onSelect,
|
||||
}) => {
|
||||
const buttonRef = useRef<HTMLButtonElement | null>(null);
|
||||
const [shouldRender, setShouldRender] = useState(false);
|
||||
|
||||
useEffect(() => {
|
||||
if (shouldRender) {
|
||||
return;
|
||||
}
|
||||
const el = buttonRef.current;
|
||||
if (!el || typeof IntersectionObserver === "undefined") {
|
||||
// Fallback for jsdom / older browsers: render eagerly.
|
||||
setShouldRender(true);
|
||||
return;
|
||||
}
|
||||
const observer = new IntersectionObserver(
|
||||
(entries) => {
|
||||
for (const entry of entries) {
|
||||
if (entry.isIntersecting) {
|
||||
setShouldRender(true);
|
||||
observer.disconnect();
|
||||
return;
|
||||
}
|
||||
}
|
||||
},
|
||||
{ rootMargin: "320px" },
|
||||
);
|
||||
observer.observe(el);
|
||||
return () => observer.disconnect();
|
||||
}, [shouldRender]);
|
||||
|
||||
useEffect(() => {
|
||||
const el = buttonRef.current;
|
||||
if (!active || !el || typeof el.scrollIntoView !== "function") {
|
||||
return;
|
||||
}
|
||||
el.scrollIntoView({ block: "nearest", behavior: "smooth" });
|
||||
}, [active]);
|
||||
|
||||
return (
|
||||
<button
|
||||
ref={buttonRef}
|
||||
type="button"
|
||||
onClick={() => onSelect(pageNumber)}
|
||||
aria-label={`Go to page ${pageNumber}`}
|
||||
aria-current={active ? "page" : undefined}
|
||||
className={cn(
|
||||
"mb-1.5 flex w-full flex-col items-center gap-0.5 rounded-md p-1 outline-none transition-colors",
|
||||
"focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-1",
|
||||
active
|
||||
? "bg-secondary/70 text-secondary-foreground"
|
||||
: "hover:bg-muted/60",
|
||||
)}
|
||||
>
|
||||
<div
|
||||
className={cn(
|
||||
"overflow-hidden rounded-sm border bg-white shadow-xs",
|
||||
active
|
||||
? "border-primary/70 ring-1 ring-primary/40"
|
||||
: "border-border/60",
|
||||
)}
|
||||
style={{
|
||||
width: PDF_THUMBNAIL_WIDTH,
|
||||
minHeight: Math.round(PDF_THUMBNAIL_WIDTH * 1.3),
|
||||
}}
|
||||
>
|
||||
{shouldRender ? (
|
||||
<Page
|
||||
pageNumber={pageNumber}
|
||||
width={PDF_THUMBNAIL_WIDTH}
|
||||
renderTextLayer={false}
|
||||
renderAnnotationLayer={false}
|
||||
loading={
|
||||
<div
|
||||
className="flex h-full w-full animate-pulse items-center justify-center bg-muted/40"
|
||||
style={{
|
||||
minHeight: Math.round(PDF_THUMBNAIL_WIDTH * 1.3),
|
||||
}}
|
||||
/>
|
||||
}
|
||||
error={
|
||||
<div
|
||||
className="flex h-full w-full items-center justify-center text-[8px] text-muted-foreground"
|
||||
style={{
|
||||
minHeight: Math.round(PDF_THUMBNAIL_WIDTH * 1.3),
|
||||
}}
|
||||
>
|
||||
?
|
||||
</div>
|
||||
}
|
||||
className="pointer-events-none [&_canvas]:!h-auto [&_canvas]:!w-full"
|
||||
/>
|
||||
) : null}
|
||||
</div>
|
||||
<span
|
||||
className={cn(
|
||||
"tabular-nums text-[10px]",
|
||||
active ? "font-semibold" : "text-muted-foreground",
|
||||
)}
|
||||
>
|
||||
{pageNumber}
|
||||
</span>
|
||||
</button>
|
||||
);
|
||||
};
|
||||
|
||||
export const PreviewPdfView: FC<PreviewPdfViewProps> = ({ target, file }) => {
|
||||
const [numPages, setNumPages] = useState<number | null>(null);
|
||||
const [pageNumber, setPageNumber] = useState<number>(target.targetPage ?? 1);
|
||||
const [loadError, setLoadError] = useState<string | null>(null);
|
||||
const [zoom, setZoom] = useState(1);
|
||||
const [searchTerm, setSearchTerm] = useState("");
|
||||
const [copied, setCopied] = useState(false);
|
||||
const containerRef = useRef<HTMLDivElement | null>(null);
|
||||
const lastMeasuredWidthRef = useRef<number | null>(null);
|
||||
const lastResetKeyRef = useRef<string | null>(null);
|
||||
const [width, setWidth] = useState<number | null>(null);
|
||||
const searchInputId = useId();
|
||||
|
||||
const sourceKey =
|
||||
typeof file === "string"
|
||||
? file
|
||||
: `${target.documentId}:${target.chunkId ?? ""}:${file.size}:${file.type}`;
|
||||
|
||||
const documentFile = useMemo(() => {
|
||||
return typeof file === "string" ? { url: file } : file;
|
||||
}, [file]);
|
||||
const resetKey = `${sourceKey}:${target.targetPage ?? ""}`;
|
||||
|
||||
useEffect(() => {
|
||||
if (lastResetKeyRef.current === resetKey) {
|
||||
return;
|
||||
}
|
||||
lastResetKeyRef.current = resetKey;
|
||||
setNumPages(null);
|
||||
setLoadError(null);
|
||||
setPageNumber(target.targetPage ?? 1);
|
||||
setZoom(1);
|
||||
setSearchTerm("");
|
||||
setCopied(false);
|
||||
}, [resetKey, target.targetPage]);
|
||||
|
||||
useEffect(() => {
|
||||
const el = containerRef.current;
|
||||
if (!el) {
|
||||
return;
|
||||
}
|
||||
let timeoutId: number | null = null;
|
||||
const updateWidth = () => {
|
||||
const next = Math.max(MIN_PDF_WIDTH, el.clientWidth - PDF_BODY_GUTTER_PX);
|
||||
if (lastMeasuredWidthRef.current === next) {
|
||||
return;
|
||||
}
|
||||
lastMeasuredWidthRef.current = next;
|
||||
setWidth(next);
|
||||
};
|
||||
updateWidth();
|
||||
const observer = new ResizeObserver(() => {
|
||||
if (timeoutId !== null) {
|
||||
window.clearTimeout(timeoutId);
|
||||
}
|
||||
timeoutId = window.setTimeout(updateWidth, RESIZE_DEBOUNCE_MS);
|
||||
});
|
||||
observer.observe(el);
|
||||
return () => {
|
||||
observer.disconnect();
|
||||
if (timeoutId !== null) {
|
||||
window.clearTimeout(timeoutId);
|
||||
}
|
||||
};
|
||||
}, []);
|
||||
|
||||
useEffect(() => {
|
||||
if (!copied) {
|
||||
return;
|
||||
}
|
||||
const id = window.setTimeout(() => setCopied(false), 1200);
|
||||
return () => window.clearTimeout(id);
|
||||
}, [copied]);
|
||||
|
||||
const handleLoadSuccess = useCallback(({ numPages }: LoadSuccess) => {
|
||||
setNumPages(numPages);
|
||||
setLoadError(null);
|
||||
}, []);
|
||||
|
||||
const handleLoadError = useCallback((err: Error) => {
|
||||
setLoadError(err.message || "Failed to load PDF");
|
||||
}, []);
|
||||
|
||||
const goPrev = useCallback(() => {
|
||||
setPageNumber((p) => Math.max(1, p - 1));
|
||||
}, []);
|
||||
|
||||
const goNext = useCallback(() => {
|
||||
setPageNumber((p) =>
|
||||
numPages == null ? p + 1 : Math.min(numPages, p + 1),
|
||||
);
|
||||
}, [numPages]);
|
||||
|
||||
const textRenderer = useCallback(
|
||||
({ str }: { str: string }) => highlightPdfText(str, searchTerm),
|
||||
[searchTerm],
|
||||
);
|
||||
|
||||
const currentRegions = useMemo(
|
||||
() =>
|
||||
(target.pdfRegions ?? []).filter((region) =>
|
||||
regionIsOnPage(region, pageNumber),
|
||||
),
|
||||
[target.pdfRegions, pageNumber],
|
||||
);
|
||||
|
||||
const visiblePageNumbers = useMemo(() => {
|
||||
if (!numPages) {
|
||||
return [];
|
||||
}
|
||||
const maxButtons = 80;
|
||||
if (numPages <= maxButtons) {
|
||||
return Array.from({ length: numPages }, (_, index) => index + 1);
|
||||
}
|
||||
const half = Math.floor(maxButtons / 2);
|
||||
let start = Math.max(1, pageNumber - half);
|
||||
const end = Math.min(numPages, start + maxButtons - 1);
|
||||
start = Math.max(1, end - maxButtons + 1);
|
||||
return Array.from({ length: end - start + 1 }, (_, index) => start + index);
|
||||
}, [numPages, pageNumber]);
|
||||
|
||||
const pageWidth = width == null ? null : Math.round(width * zoom);
|
||||
const excerptKey = `${sourceKey}:${target.chunkId ?? ""}:${
|
||||
target.targetPage ?? ""
|
||||
}:${pageNumber}`;
|
||||
|
||||
const copyExcerpt = useCallback(() => {
|
||||
copyToClipboard(target.snippet ?? "").then(setCopied);
|
||||
}, [target.snippet]);
|
||||
|
||||
if (loadError) {
|
||||
return (
|
||||
<PreviewUnavailable
|
||||
filename={target.filename}
|
||||
reason={loadError}
|
||||
variant="error"
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="flex h-full min-h-0 flex-col">
|
||||
<div className="flex flex-wrap items-center justify-between gap-2 border-b border-border/60 px-3 py-2 text-xs">
|
||||
<span
|
||||
className="min-w-0 flex-1 truncate font-semibold font-heading"
|
||||
title={target.filename}
|
||||
>
|
||||
{target.filename}
|
||||
</span>
|
||||
<div className="flex shrink-0 items-center gap-2">
|
||||
{/* Navigation Pill Group */}
|
||||
<div className="flex items-center rounded-full border border-border/60 bg-muted/40 p-0.5 shadow-xs">
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
onClick={goPrev}
|
||||
disabled={pageNumber <= 1}
|
||||
aria-label="Previous page"
|
||||
className="h-7 w-7 rounded-full hover:bg-background/80"
|
||||
>
|
||||
<ChevronLeftIcon className="size-3.5" />
|
||||
</Button>
|
||||
<span className="min-w-12 text-center tabular-nums text-[10px] font-medium text-muted-foreground">
|
||||
{numPages == null
|
||||
? `${pageNumber}/?`
|
||||
: `${pageNumber}/${numPages}`}
|
||||
</span>
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
onClick={goNext}
|
||||
disabled={numPages != null && pageNumber >= numPages}
|
||||
aria-label="Next page"
|
||||
className="h-7 w-7 rounded-full hover:bg-background/80"
|
||||
>
|
||||
<ChevronRightIcon className="size-3.5" />
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
{/* Zoom Pill Group */}
|
||||
<div className="flex items-center rounded-full border border-border/60 bg-muted/40 p-0.5 shadow-xs">
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
onClick={() => setZoom((value) => Math.max(0.6, value - 0.1))}
|
||||
aria-label="Zoom out"
|
||||
className="h-7 w-7 rounded-full hover:bg-background/80"
|
||||
>
|
||||
<ZoomOutIcon className="size-3.5" />
|
||||
</Button>
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
onClick={() => setZoom(1)}
|
||||
aria-label="Reset zoom"
|
||||
className="h-7 w-7 rounded-full hover:bg-background/80"
|
||||
>
|
||||
<RotateCcwIcon className="size-3.5" />
|
||||
</Button>
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
onClick={() => setZoom((value) => Math.min(2.5, value + 0.1))}
|
||||
aria-label="Zoom in"
|
||||
className="h-7 w-7 rounded-full hover:bg-background/80"
|
||||
>
|
||||
<ZoomInIcon className="size-3.5" />
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
{/* Copy Excerpt Pill Group */}
|
||||
<div className="flex items-center rounded-full border border-border/60 bg-muted/40 p-0.5 shadow-xs">
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
onClick={copyExcerpt}
|
||||
disabled={!target.snippet}
|
||||
aria-label={
|
||||
copied ? "Copied source excerpt" : "Copy source excerpt"
|
||||
}
|
||||
className="h-7 w-7 rounded-full hover:bg-background/80"
|
||||
>
|
||||
<CopyIcon className="size-3.5" />
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
<label
|
||||
htmlFor={searchInputId}
|
||||
className="flex min-w-48 max-w-full flex-1 items-center gap-1 rounded-md border border-border/60 bg-background px-2"
|
||||
>
|
||||
<SearchIcon className="size-3.5 shrink-0 text-muted-foreground" />
|
||||
<Input
|
||||
id={searchInputId}
|
||||
value={searchTerm}
|
||||
onChange={(event) => setSearchTerm(event.target.value)}
|
||||
placeholder="Search this PDF"
|
||||
aria-label="Search this PDF"
|
||||
className="h-7 border-0 bg-transparent px-0 text-xs shadow-none focus-visible:ring-0"
|
||||
/>
|
||||
</label>
|
||||
</div>
|
||||
|
||||
{target.snippet ? (
|
||||
<div
|
||||
key={excerptKey}
|
||||
className="m-2 rounded-lg border border-border/60 bg-muted/30 p-3 shadow-xs text-[11px] leading-relaxed text-foreground/80 transition-all duration-300 animate-in fade-in"
|
||||
>
|
||||
<p className="mb-1 text-[10px] font-semibold uppercase tracking-wider text-muted-foreground/80">
|
||||
Source Excerpt
|
||||
{target.targetPage != null ? ` · Page ${target.targetPage}` : ""}
|
||||
</p>
|
||||
<p className="line-clamp-4 whitespace-pre-wrap font-sans text-muted-foreground">
|
||||
{target.snippet}
|
||||
</p>
|
||||
</div>
|
||||
) : null}
|
||||
|
||||
<Document
|
||||
file={documentFile}
|
||||
onLoadSuccess={handleLoadSuccess}
|
||||
onLoadError={handleLoadError}
|
||||
loading={
|
||||
<div className="flex h-full items-center justify-center gap-2 text-xs text-muted-foreground">
|
||||
<LoaderIcon className="size-3.5 animate-spin" />
|
||||
Loading PDF...
|
||||
</div>
|
||||
}
|
||||
error={
|
||||
<PreviewUnavailable
|
||||
filename={target.filename}
|
||||
reason="The PDF could not be opened."
|
||||
variant="error"
|
||||
/>
|
||||
}
|
||||
className="flex min-h-0 flex-1"
|
||||
>
|
||||
<div className="preview-scrollbar w-[88px] shrink-0 overflow-y-auto border-r border-border/60 bg-muted/20 p-1.5">
|
||||
{visiblePageNumbers.map((page) => (
|
||||
<PdfThumbnail
|
||||
key={page}
|
||||
pageNumber={page}
|
||||
active={page === pageNumber}
|
||||
onSelect={setPageNumber}
|
||||
/>
|
||||
))}
|
||||
</div>
|
||||
<div
|
||||
ref={containerRef}
|
||||
className="preview-scrollbar flex-1 overflow-y-scroll overflow-x-auto bg-muted/20 p-2 [scrollbar-gutter:stable]"
|
||||
>
|
||||
<div
|
||||
className="light [color-scheme:light] bg-white text-slate-900 rounded-md p-1 shadow-sm border border-border/30 [&_mark]:bg-primary/20 [&_mark]:text-slate-900 [&_mark]:ring-1 [&_mark]:ring-primary/60 [&_mark]:rounded-xs flex min-w-fit flex-col items-center"
|
||||
style={PDF_LIGHT_THEME_STYLE}
|
||||
>
|
||||
{pageWidth != null ? (
|
||||
<div
|
||||
data-testid="pdf-main-page"
|
||||
className="relative inline-block"
|
||||
>
|
||||
<Page
|
||||
pageNumber={pageNumber}
|
||||
width={pageWidth}
|
||||
customTextRenderer={textRenderer}
|
||||
renderTextLayer={true}
|
||||
renderAnnotationLayer={false}
|
||||
loading={
|
||||
<div className="py-4 text-xs text-muted-foreground">
|
||||
Rendering page...
|
||||
</div>
|
||||
}
|
||||
className="shadow-sm"
|
||||
/>
|
||||
{currentRegions.map((region, index) => (
|
||||
<div
|
||||
key={`${region.pageIndex}-${region.x}-${region.y}-${index}`}
|
||||
data-testid="pdf-region-highlight"
|
||||
className="pointer-events-none absolute rounded-sm bg-primary/20 ring-1 ring-primary/60"
|
||||
style={{
|
||||
left: `${region.x * 100}%`,
|
||||
top: `${region.y * 100}%`,
|
||||
width: `${region.width * 100}%`,
|
||||
height: `${region.height * 100}%`,
|
||||
}}
|
||||
/>
|
||||
))}
|
||||
</div>
|
||||
) : null}
|
||||
</div>
|
||||
</div>
|
||||
</Document>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
|
@ -0,0 +1,300 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { authFetch } from "@/features/auth";
|
||||
import { DownloadIcon, ExternalLinkIcon, FileTextIcon } from "lucide-react";
|
||||
import { type FC, type ReactNode, useCallback } from "react";
|
||||
import type { PreviewTarget } from "../api/rag-api";
|
||||
import { isInlineBlobAllowed } from "../stores/preview-store";
|
||||
|
||||
interface PreviewTextViewProps {
|
||||
target: PreviewTarget;
|
||||
}
|
||||
|
||||
/** Fetch the original document bytes via authFetch so the bearer
|
||||
* token rides in the Authorization header. `window.open(url)` and
|
||||
* `<a download href=url>` cannot set custom headers, so handing
|
||||
* them the raw `/file` URL gets a 401 (HTTPBearer-only backend —
|
||||
* see D1.3). */
|
||||
async function fetchOriginalBlob(target: PreviewTarget): Promise<Blob> {
|
||||
const response = await authFetch(
|
||||
`/api/rag/documents/${encodeURIComponent(target.documentId)}/file`,
|
||||
);
|
||||
if (!response.ok) {
|
||||
throw new Error(`Failed to fetch document (${response.status})`);
|
||||
}
|
||||
return response.blob();
|
||||
}
|
||||
|
||||
function clickDownloadUrl(url: string, filename: string): void {
|
||||
const anchor = document.createElement("a");
|
||||
anchor.href = url;
|
||||
anchor.download = filename;
|
||||
anchor.style.display = "none";
|
||||
document.body.appendChild(anchor);
|
||||
anchor.click();
|
||||
anchor.remove();
|
||||
}
|
||||
|
||||
async function downloadOriginal(target: PreviewTarget): Promise<void> {
|
||||
const blob = await fetchOriginalBlob(target);
|
||||
const url = URL.createObjectURL(blob);
|
||||
clickDownloadUrl(url, target.filename);
|
||||
// Defer revocation so the browser's download pipeline gets the bytes
|
||||
// before the URL goes away.
|
||||
setTimeout(() => URL.revokeObjectURL(url), 0);
|
||||
}
|
||||
|
||||
async function openOriginalInNewTab(target: PreviewTarget): Promise<void> {
|
||||
// Defense in depth: refuse to create an inline blob URL for the
|
||||
// unsafe types even if a future caller forgets the gate. The
|
||||
// browser would render an html-blob as live HTML in the new tab,
|
||||
// which is the Risk #3 / contracts §2.3 trip.
|
||||
if (!isInlineBlobAllowed(target.mediaKind)) {
|
||||
throw new Error(
|
||||
`Inline open not allowed for mediaKind "${target.mediaKind}" — use Download instead.`,
|
||||
);
|
||||
}
|
||||
const blob = await fetchOriginalBlob(target);
|
||||
const url = URL.createObjectURL(blob);
|
||||
window.open(url, "_blank", "noopener,noreferrer");
|
||||
// Defer revocation so the new tab loads the bytes first.
|
||||
setTimeout(() => URL.revokeObjectURL(url), 0);
|
||||
}
|
||||
|
||||
/** Extracted-text / snippet preview used for:
|
||||
* - `text` mediaKind (txt/md) — the snippet is the only inline
|
||||
* rendering we trust, and the original is one click away.
|
||||
* - `docx`, `html`, `unknown` — the original is NEVER rendered
|
||||
* inline from an object URL (Risk #3); we show the cited chunk
|
||||
* text plus a safe download/open action.
|
||||
*
|
||||
* When `chunk_id` was not supplied (document-row preview per
|
||||
* contracts §1.3 + decision Q2), `snippet` is `null` and we show a
|
||||
* metadata-only state instead of guessing a first chunk. */
|
||||
interface MatchRange {
|
||||
start: number;
|
||||
end: number;
|
||||
}
|
||||
|
||||
function findCharacterRange(
|
||||
snippet: string,
|
||||
target: PreviewTarget,
|
||||
): MatchRange | null {
|
||||
const { pageCharStart, pageCharEnd } = target;
|
||||
if (pageCharStart !== null && pageCharEnd !== null) {
|
||||
const start = Math.max(0, pageCharStart);
|
||||
const end = Math.min(snippet.length, pageCharEnd);
|
||||
if (start < end) {
|
||||
return { start, end };
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
function findLineRange(
|
||||
snippet: string,
|
||||
target: PreviewTarget,
|
||||
): MatchRange | null {
|
||||
const { lineStart, lineEnd } = target;
|
||||
if (lineStart !== null) {
|
||||
const lines = snippet.split("\n");
|
||||
const startLineIndex = Math.max(0, lineStart - 1);
|
||||
const endLineIndex =
|
||||
lineEnd !== null
|
||||
? Math.min(lines.length - 1, lineEnd - 1)
|
||||
: startLineIndex;
|
||||
|
||||
let charOffset = 0;
|
||||
let startChar = -1;
|
||||
let endChar = -1;
|
||||
|
||||
for (let i = 0; i < lines.length; i++) {
|
||||
if (i === startLineIndex) {
|
||||
startChar = charOffset;
|
||||
}
|
||||
charOffset += lines[i].length;
|
||||
if (i === endLineIndex) {
|
||||
endChar = charOffset;
|
||||
break;
|
||||
}
|
||||
charOffset += 1; // for '\n'
|
||||
}
|
||||
|
||||
if (startChar !== -1 && endChar !== -1 && startChar < endChar) {
|
||||
return { start: startChar, end: endChar };
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
function findDensestLineRange(snippet: string): MatchRange | null {
|
||||
const lines = snippet.split("\n");
|
||||
let bestLineIndex = -1;
|
||||
let maxAlphanumericCount = 0;
|
||||
for (let i = 0; i < lines.length; i++) {
|
||||
const alphanumericCount = lines[i].replace(/[^a-zA-Z0-9]/g, "").length;
|
||||
if (alphanumericCount > maxAlphanumericCount) {
|
||||
maxAlphanumericCount = alphanumericCount;
|
||||
bestLineIndex = i;
|
||||
}
|
||||
}
|
||||
if (bestLineIndex !== -1) {
|
||||
let charOffset = 0;
|
||||
for (let i = 0; i < bestLineIndex; i++) {
|
||||
charOffset += lines[i].length + 1;
|
||||
}
|
||||
return { start: charOffset, end: charOffset + lines[bestLineIndex].length };
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
function findFuzzyMatch(
|
||||
snippet: string,
|
||||
target: PreviewTarget,
|
||||
): MatchRange | null {
|
||||
return (
|
||||
findCharacterRange(snippet, target) ??
|
||||
findLineRange(snippet, target) ??
|
||||
findDensestLineRange(snippet)
|
||||
);
|
||||
}
|
||||
|
||||
const renderHighlightedSnippet = (
|
||||
snippet: string,
|
||||
target: PreviewTarget,
|
||||
): ReactNode => {
|
||||
const match = findFuzzyMatch(snippet, target);
|
||||
if (!match) {
|
||||
return snippet;
|
||||
}
|
||||
const before = snippet.slice(0, match.start);
|
||||
const highlighted = snippet.slice(match.start, match.end);
|
||||
const after = snippet.slice(match.end);
|
||||
|
||||
return (
|
||||
<>
|
||||
{before}
|
||||
<mark className="rounded bg-primary/20 px-0.5 text-foreground ring-1 ring-primary/60">
|
||||
{highlighted}
|
||||
</mark>
|
||||
{after}
|
||||
</>
|
||||
);
|
||||
};
|
||||
|
||||
/** Extracted-text / snippet preview used for:
|
||||
* - `text` mediaKind (txt/md) — the snippet is the only inline
|
||||
* rendering we trust, and the original is one click away.
|
||||
* - `docx`, `html`, `unknown` — the original is NEVER rendered
|
||||
* inline from an object URL (Risk #3); we show the cited chunk
|
||||
* text plus a safe download/open action.
|
||||
*
|
||||
* When `chunk_id` was not supplied (document-row preview per
|
||||
* contracts §1.3 + decision Q2), `snippet` is `null` and we show a
|
||||
* metadata-only state instead of guessing a first chunk. */
|
||||
export const PreviewTextView: FC<PreviewTextViewProps> = ({ target }) => {
|
||||
const snippet = target.snippet;
|
||||
const hasSnippet = snippet !== null && snippet.trim().length > 0;
|
||||
const hasLocator =
|
||||
target.lineStart !== null ||
|
||||
target.lineEnd !== null ||
|
||||
target.pageCharStart !== null ||
|
||||
target.pageCharEnd !== null;
|
||||
// "Open original" creates a blob: URL of the original bytes and
|
||||
// passes it to `window.open`. For `html` the new tab would render
|
||||
// it as live HTML — exactly the Risk #3 / contracts §2.3 trip
|
||||
// ("MUST refuse to create an inline object URL for mediaKind ==
|
||||
// 'html' | 'docx' | 'unknown'"). For those types the only safe
|
||||
// action is Download (backend already sets
|
||||
// Content-Disposition: attachment for those Content-Types). The
|
||||
// pdf/text/image allowlist is the same one the preview-store
|
||||
// uses to decide whether to fetch the blob at all (§5.4). */
|
||||
const canOpenInline = isInlineBlobAllowed(target.mediaKind);
|
||||
|
||||
const handleDownload = useCallback(() => {
|
||||
downloadOriginal(target).catch(() => {
|
||||
// best-effort; the user can retry the action.
|
||||
});
|
||||
}, [target]);
|
||||
|
||||
const handleOpenExternal = useCallback(() => {
|
||||
// Re-fetch through authFetch and hand the new tab a blob URL.
|
||||
// `window.open(rawApiUrl)` would send the request WITHOUT the
|
||||
// Authorization header (window.open can't set custom headers)
|
||||
// and the HTTPBearer-protected /file route would respond 401.
|
||||
// See D1.3 finding.
|
||||
openOriginalInNewTab(target).catch(() => {
|
||||
// best-effort; the user can retry.
|
||||
});
|
||||
}, [target]);
|
||||
|
||||
return (
|
||||
<div className="flex h-full flex-col gap-3 overflow-hidden p-4">
|
||||
<div className="flex items-center gap-2 text-sm">
|
||||
<FileTextIcon
|
||||
className="size-4 shrink-0 text-muted-foreground"
|
||||
aria-hidden={true}
|
||||
/>
|
||||
<span className="truncate font-medium" title={target.filename}>
|
||||
{target.filename}
|
||||
</span>
|
||||
</div>
|
||||
{target.targetPage != null ? (
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Cited from page {target.targetPage}
|
||||
</p>
|
||||
) : null}
|
||||
|
||||
<div className="flex min-h-0 flex-1 flex-col gap-2 overflow-hidden rounded-md border border-border/60 bg-muted/30 p-3">
|
||||
{hasSnippet ? (
|
||||
<>
|
||||
<p className="text-[10px] uppercase tracking-wider text-muted-foreground">
|
||||
{hasLocator ? "Highlighted source excerpt" : "Source excerpt"}
|
||||
</p>
|
||||
{hasLocator ? (
|
||||
<pre className="flex-1 overflow-auto whitespace-pre-wrap break-words text-xs leading-relaxed text-foreground/85">
|
||||
{renderHighlightedSnippet(snippet, target)}
|
||||
</pre>
|
||||
) : (
|
||||
<pre className="flex-1 overflow-auto whitespace-pre-wrap break-words text-xs leading-relaxed text-foreground/85">
|
||||
{snippet}
|
||||
</pre>
|
||||
)}
|
||||
</>
|
||||
) : (
|
||||
<p className="my-auto text-center text-xs text-muted-foreground">
|
||||
{canOpenInline
|
||||
? "No source excerpt — open the original to view this document."
|
||||
: "No source excerpt — download the original to view this document."}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="flex gap-2">
|
||||
{canOpenInline ? (
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
onClick={handleOpenExternal}
|
||||
className="flex-1"
|
||||
>
|
||||
<ExternalLinkIcon className="size-3.5" />
|
||||
Open original
|
||||
</Button>
|
||||
) : null}
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
onClick={handleDownload}
|
||||
className="flex-1"
|
||||
>
|
||||
<DownloadIcon className="size-3.5" />
|
||||
Download
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { AlertCircleIcon, FileXIcon } from "lucide-react";
|
||||
import type { FC } from "react";
|
||||
|
||||
interface PreviewUnavailableProps {
|
||||
/** Filename if known; "Document" otherwise. */
|
||||
filename?: string;
|
||||
/** One-line reason — pulled from the backend's error body when
|
||||
* available, otherwise a generic copy. */
|
||||
reason: string;
|
||||
/** "missing" → deleted/404 case; "error" → other failures. The icon
|
||||
* + tone change so a stale citation reads as "no longer available"
|
||||
* rather than a transient blip. */
|
||||
variant?: "missing" | "error";
|
||||
}
|
||||
|
||||
export const PreviewUnavailable: FC<PreviewUnavailableProps> = ({
|
||||
filename,
|
||||
reason,
|
||||
variant = "error",
|
||||
}) => {
|
||||
const Icon = variant === "missing" ? FileXIcon : AlertCircleIcon;
|
||||
const headline =
|
||||
variant === "missing" ? "Document unavailable" : "Couldn't load preview";
|
||||
|
||||
return (
|
||||
<output className="flex h-full flex-col items-center justify-center gap-3 px-6 text-center">
|
||||
<Icon
|
||||
className="size-10 text-muted-foreground"
|
||||
strokeWidth={1.5}
|
||||
aria-hidden={true}
|
||||
/>
|
||||
<div className="flex flex-col gap-1">
|
||||
<p className="text-sm font-medium">{headline}</p>
|
||||
{filename ? (
|
||||
<p className="text-xs text-muted-foreground" title={filename}>
|
||||
{filename}
|
||||
</p>
|
||||
) : null}
|
||||
<p className="mt-1 max-w-xs text-xs text-muted-foreground">{reason}</p>
|
||||
</div>
|
||||
</output>
|
||||
);
|
||||
};
|
||||
158
studio/frontend/src/features/rag/hooks/use-resizable-width.ts
Normal file
158
studio/frontend/src/features/rag/hooks/use-resizable-width.ts
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import {
|
||||
type PointerEvent as ReactPointerEvent,
|
||||
useCallback,
|
||||
useEffect,
|
||||
useRef,
|
||||
useState,
|
||||
} from "react";
|
||||
|
||||
interface UseResizableWidthOptions {
|
||||
storageKey: string;
|
||||
defaultWidth: number;
|
||||
minWidth: number;
|
||||
/** 0..1 fraction of viewport.innerWidth used as max width. Default 0.8. */
|
||||
maxWidthFraction?: number;
|
||||
/** Persist + listen for viewport-resize clamping only while true. */
|
||||
enabled?: boolean;
|
||||
}
|
||||
|
||||
interface UseResizableWidthResult {
|
||||
width: number;
|
||||
isResizing: boolean;
|
||||
startResize: (event: ReactPointerEvent<HTMLElement>) => void;
|
||||
adjustWidth: (delta: number) => void;
|
||||
resetWidth: () => void;
|
||||
}
|
||||
|
||||
function readStored(key: string, fallback: number): number {
|
||||
if (typeof window === "undefined") {
|
||||
return fallback;
|
||||
}
|
||||
try {
|
||||
const raw = window.localStorage.getItem(key);
|
||||
if (raw == null) {
|
||||
return fallback;
|
||||
}
|
||||
const parsed = Number.parseInt(raw, 10);
|
||||
return Number.isFinite(parsed) ? parsed : fallback;
|
||||
} catch {
|
||||
return fallback;
|
||||
}
|
||||
}
|
||||
|
||||
/** Drag-to-resize hook for a right-anchored panel. The handle sits on
|
||||
* the panel's LEFT edge — width grows as the pointer moves toward
|
||||
* viewport x=0. Persists to localStorage and re-clamps on viewport
|
||||
* changes so a wide panel cannot eclipse the host content. */
|
||||
export function useResizablePanelWidth({
|
||||
storageKey,
|
||||
defaultWidth,
|
||||
minWidth,
|
||||
maxWidthFraction = 0.8,
|
||||
enabled = true,
|
||||
}: UseResizableWidthOptions): UseResizableWidthResult {
|
||||
const [width, setWidth] = useState<number>(() =>
|
||||
readStored(storageKey, defaultWidth),
|
||||
);
|
||||
const [isResizing, setIsResizing] = useState(false);
|
||||
const rafRef = useRef<number | null>(null);
|
||||
|
||||
const clampWidth = useCallback(
|
||||
(next: number): number => {
|
||||
if (typeof window === "undefined") {
|
||||
return Math.max(minWidth, next);
|
||||
}
|
||||
const max = Math.floor(window.innerWidth * maxWidthFraction);
|
||||
return Math.max(minWidth, Math.min(max, next));
|
||||
},
|
||||
[minWidth, maxWidthFraction],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
if (!enabled || typeof window === "undefined") {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
window.localStorage.setItem(storageKey, String(width));
|
||||
} catch {
|
||||
// localStorage may be unavailable (private mode, quota); persist
|
||||
// is best-effort.
|
||||
}
|
||||
}, [width, storageKey, enabled]);
|
||||
|
||||
useEffect(() => {
|
||||
if (typeof window === "undefined") {
|
||||
return;
|
||||
}
|
||||
const onResize = () => {
|
||||
setWidth((w) => clampWidth(w));
|
||||
};
|
||||
window.addEventListener("resize", onResize);
|
||||
return () => window.removeEventListener("resize", onResize);
|
||||
}, [clampWidth]);
|
||||
|
||||
const startResize = useCallback(
|
||||
(event: ReactPointerEvent<HTMLElement>) => {
|
||||
if (!enabled) {
|
||||
return;
|
||||
}
|
||||
event.preventDefault();
|
||||
const target = event.currentTarget;
|
||||
const pointerId = event.pointerId;
|
||||
try {
|
||||
target.setPointerCapture(pointerId);
|
||||
} catch {
|
||||
// Pointer-capture isn't available everywhere (e.g. test envs).
|
||||
}
|
||||
setIsResizing(true);
|
||||
document.body.style.cursor = "col-resize";
|
||||
document.body.style.userSelect = "none";
|
||||
|
||||
const onMove = (e: PointerEvent) => {
|
||||
if (rafRef.current !== null) {
|
||||
cancelAnimationFrame(rafRef.current);
|
||||
}
|
||||
rafRef.current = requestAnimationFrame(() => {
|
||||
setWidth(clampWidth(window.innerWidth - e.clientX));
|
||||
});
|
||||
};
|
||||
const cleanup = (e: PointerEvent) => {
|
||||
setIsResizing(false);
|
||||
document.body.style.cursor = "";
|
||||
document.body.style.userSelect = "";
|
||||
try {
|
||||
target.releasePointerCapture(e.pointerId);
|
||||
} catch {
|
||||
// Already released or not captured.
|
||||
}
|
||||
window.removeEventListener("pointermove", onMove);
|
||||
window.removeEventListener("pointerup", cleanup);
|
||||
window.removeEventListener("pointercancel", cleanup);
|
||||
if (rafRef.current !== null) {
|
||||
cancelAnimationFrame(rafRef.current);
|
||||
rafRef.current = null;
|
||||
}
|
||||
};
|
||||
window.addEventListener("pointermove", onMove);
|
||||
window.addEventListener("pointerup", cleanup);
|
||||
window.addEventListener("pointercancel", cleanup);
|
||||
},
|
||||
[enabled, clampWidth],
|
||||
);
|
||||
|
||||
const adjustWidth = useCallback(
|
||||
(delta: number) => {
|
||||
setWidth((w) => clampWidth(w + delta));
|
||||
},
|
||||
[clampWidth],
|
||||
);
|
||||
|
||||
const resetWidth = useCallback(() => {
|
||||
setWidth(clampWidth(defaultWidth));
|
||||
}, [clampWidth, defaultWidth]);
|
||||
|
||||
return { width, isResizing, startResize, adjustWidth, resetWidth };
|
||||
}
|
||||
275
studio/frontend/src/features/rag/stores/preview-store.ts
Normal file
275
studio/frontend/src/features/rag/stores/preview-store.ts
Normal file
|
|
@ -0,0 +1,275 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { create } from "zustand";
|
||||
import {
|
||||
type PreviewMediaKind,
|
||||
type PreviewTarget,
|
||||
fetchPreviewFileBlob,
|
||||
fetchPreviewFileUrl,
|
||||
fetchPreviewTarget,
|
||||
} from "../api/rag-api";
|
||||
|
||||
/** mediaKinds that may safely back an inline object URL (e.g. PDF.js
|
||||
* worker, plain text, raster image). HTML, DOCX, and unknown are
|
||||
* forced through the extracted-text fallback per contracts §5.4 +
|
||||
* Risk #3 (unsafe HTML inline rendering). */
|
||||
const INLINE_BLOB_ALLOWLIST: ReadonlySet<PreviewMediaKind> = new Set([
|
||||
"pdf",
|
||||
"text",
|
||||
"image",
|
||||
]);
|
||||
|
||||
export function isInlineBlobAllowed(mediaKind: PreviewMediaKind): boolean {
|
||||
return INLINE_BLOB_ALLOWLIST.has(mediaKind);
|
||||
}
|
||||
|
||||
/** What the panel should mount for the current target. Computed from
|
||||
* `target.mediaKind` so the panel never has to re-derive it. */
|
||||
export type PreviewLoadStatus = "idle" | "loading" | "ready" | "error";
|
||||
|
||||
export interface PreviewRequest {
|
||||
/** Durable `rag_documents.id`. The only field required to open. */
|
||||
documentId: string;
|
||||
/** Durable `rag_chunks.id`. Optional — absence means document-row
|
||||
* preview (contracts §1.3 + decision Q2: snippet/targetPage stay
|
||||
* null, no first-chunk fallback). */
|
||||
backendChunkId?: string | null;
|
||||
}
|
||||
|
||||
interface PreviewState {
|
||||
/** Currently-open preview, or null when closed. */
|
||||
target: PreviewTarget | null;
|
||||
/** Object URL for the original file blob (PDF / text / image only).
|
||||
* Null for docx / html / unknown (extracted-text fallback) and
|
||||
* while the fetch is still in flight. */
|
||||
previewBlobUrl: string | null;
|
||||
/** Original fetched file blob for text/image fallback previews. PDFs
|
||||
* prefer `previewFileUrl` so PDF.js can issue range requests. */
|
||||
previewBlob: Blob | null;
|
||||
/** Short-lived signed URL for PDF.js range requests. */
|
||||
previewFileUrl: string | null;
|
||||
previewFileUrlExpiresAt: number | null;
|
||||
/** `target`-fetch + `blob`-fetch combined status. */
|
||||
status: PreviewLoadStatus;
|
||||
/** Last error message, if `status === "error"`. */
|
||||
error: string | null;
|
||||
/** Open key uniquely identifying the current request — used by tests
|
||||
* and by consumers that need to react to "the open call changed
|
||||
* underneath me" (e.g. re-fetch after stale closure). */
|
||||
openKey: number;
|
||||
|
||||
/** Open or replace the current preview. If a previous preview is
|
||||
* open, its object URL is revoked and its in-flight fetch is
|
||||
* aborted before the new request begins. */
|
||||
open: (req: PreviewRequest) => Promise<void>;
|
||||
/** Close the current preview. Revokes the object URL and aborts any
|
||||
* in-flight fetch. Safe to call when nothing is open. */
|
||||
close: () => void;
|
||||
}
|
||||
|
||||
// State that can't live inside the zustand object without being
|
||||
// part of the React render cycle. Kept module-scoped because the
|
||||
// preview store is a singleton.
|
||||
let activeAbortController: AbortController | null = null;
|
||||
let activeBlobUrl: string | null = null;
|
||||
let activeOpenKey = 0;
|
||||
let restoreFocusElement: HTMLElement | null = null;
|
||||
|
||||
function revokeActiveBlobUrl(): void {
|
||||
if (activeBlobUrl) {
|
||||
URL.revokeObjectURL(activeBlobUrl);
|
||||
activeBlobUrl = null;
|
||||
}
|
||||
}
|
||||
|
||||
function abortActive(): void {
|
||||
if (activeAbortController) {
|
||||
activeAbortController.abort();
|
||||
activeAbortController = null;
|
||||
}
|
||||
}
|
||||
|
||||
export const usePreviewStore = create<PreviewState>((set) => ({
|
||||
target: null,
|
||||
previewBlobUrl: null,
|
||||
previewBlob: null,
|
||||
previewFileUrl: null,
|
||||
previewFileUrlExpiresAt: null,
|
||||
status: "idle",
|
||||
error: null,
|
||||
openKey: 0,
|
||||
|
||||
async open(req) {
|
||||
// Single-slot invariant (contracts §5.1): tear down whatever was
|
||||
// there before assigning the new target. revoke → abort → reset.
|
||||
revokeActiveBlobUrl();
|
||||
abortActive();
|
||||
|
||||
activeOpenKey += 1;
|
||||
const myKey = activeOpenKey;
|
||||
const controller = new AbortController();
|
||||
activeAbortController = controller;
|
||||
const activeElement = document.activeElement;
|
||||
restoreFocusElement =
|
||||
activeElement instanceof HTMLElement ? activeElement : null;
|
||||
|
||||
set({
|
||||
target: null,
|
||||
previewBlobUrl: null,
|
||||
previewBlob: null,
|
||||
previewFileUrl: null,
|
||||
previewFileUrlExpiresAt: null,
|
||||
status: "loading",
|
||||
error: null,
|
||||
openKey: myKey,
|
||||
});
|
||||
|
||||
let target: PreviewTarget;
|
||||
try {
|
||||
target = await fetchPreviewTarget(
|
||||
req.documentId,
|
||||
req.backendChunkId ?? null,
|
||||
);
|
||||
} catch (err) {
|
||||
if (myKey !== activeOpenKey) return; // superseded
|
||||
if (activeAbortController === controller) activeAbortController = null;
|
||||
set({
|
||||
status: "error",
|
||||
error: err instanceof Error ? err.message : String(err),
|
||||
openKey: myKey,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (myKey !== activeOpenKey) {
|
||||
// The user opened a different document while we were waiting.
|
||||
return;
|
||||
}
|
||||
|
||||
// For mediaKinds outside the allowlist (docx / html / unknown),
|
||||
// skip the blob fetch entirely — the panel mounts the
|
||||
// extracted-text fallback (contracts §5.4 + Risk #3).
|
||||
if (!isInlineBlobAllowed(target.mediaKind)) {
|
||||
if (activeAbortController === controller) activeAbortController = null;
|
||||
set({
|
||||
target,
|
||||
previewBlobUrl: null,
|
||||
previewBlob: null,
|
||||
previewFileUrl: null,
|
||||
previewFileUrlExpiresAt: null,
|
||||
status: "ready",
|
||||
error: null,
|
||||
openKey: myKey,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (target.mediaKind === "pdf") {
|
||||
try {
|
||||
const previewFile = await fetchPreviewFileUrl(
|
||||
req.documentId,
|
||||
controller.signal,
|
||||
);
|
||||
if (myKey !== activeOpenKey) return;
|
||||
if (activeAbortController === controller) activeAbortController = null;
|
||||
set({
|
||||
target,
|
||||
previewBlobUrl: null,
|
||||
previewBlob: null,
|
||||
previewFileUrl: previewFile.url,
|
||||
previewFileUrlExpiresAt: previewFile.expiresAt,
|
||||
status: "ready",
|
||||
error: null,
|
||||
openKey: myKey,
|
||||
});
|
||||
} catch (err) {
|
||||
if (controller.signal.aborted || myKey !== activeOpenKey) return;
|
||||
if (activeAbortController === controller) activeAbortController = null;
|
||||
set({
|
||||
target,
|
||||
previewBlobUrl: null,
|
||||
previewBlob: null,
|
||||
previewFileUrl: null,
|
||||
previewFileUrlExpiresAt: null,
|
||||
status: "error",
|
||||
error: err instanceof Error ? err.message : String(err),
|
||||
openKey: myKey,
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
let blob: Blob;
|
||||
try {
|
||||
blob = await fetchPreviewFileBlob(req.documentId, controller.signal);
|
||||
} catch (err) {
|
||||
if (controller.signal.aborted || myKey !== activeOpenKey) return;
|
||||
if (activeAbortController === controller) activeAbortController = null;
|
||||
set({
|
||||
target,
|
||||
previewBlobUrl: null,
|
||||
previewBlob: null,
|
||||
previewFileUrl: null,
|
||||
previewFileUrlExpiresAt: null,
|
||||
status: "error",
|
||||
error: err instanceof Error ? err.message : String(err),
|
||||
openKey: myKey,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (myKey !== activeOpenKey) {
|
||||
// Superseded between target fetch and blob fetch — drop the bytes.
|
||||
return;
|
||||
}
|
||||
|
||||
const objectUrl = URL.createObjectURL(blob);
|
||||
activeBlobUrl = objectUrl;
|
||||
if (activeAbortController === controller) activeAbortController = null;
|
||||
set({
|
||||
target,
|
||||
previewBlobUrl: objectUrl,
|
||||
previewBlob: blob,
|
||||
previewFileUrl: null,
|
||||
previewFileUrlExpiresAt: null,
|
||||
status: "ready",
|
||||
error: null,
|
||||
openKey: myKey,
|
||||
});
|
||||
},
|
||||
|
||||
close() {
|
||||
revokeActiveBlobUrl();
|
||||
abortActive();
|
||||
activeOpenKey += 1; // poison any in-flight fetch that lands after this
|
||||
const focusTarget = restoreFocusElement;
|
||||
restoreFocusElement = null;
|
||||
set({
|
||||
target: null,
|
||||
previewBlobUrl: null,
|
||||
previewBlob: null,
|
||||
previewFileUrl: null,
|
||||
previewFileUrlExpiresAt: null,
|
||||
status: "idle",
|
||||
error: null,
|
||||
openKey: activeOpenKey,
|
||||
});
|
||||
if (focusTarget?.isConnected) {
|
||||
focusTarget.focus();
|
||||
}
|
||||
},
|
||||
}));
|
||||
|
||||
/** Test-only inspector: returns whether the module-scoped blob URL is
|
||||
* still live. Used by `preview-store.test.ts` to assert
|
||||
* URL.revokeObjectURL was paired with URL.createObjectURL. */
|
||||
export function __previewStoreInternals(): {
|
||||
activeBlobUrl: string | null;
|
||||
hasInflightController: boolean;
|
||||
} {
|
||||
return {
|
||||
activeBlobUrl,
|
||||
hasInflightController: activeAbortController !== null,
|
||||
};
|
||||
}
|
||||
|
|
@ -5,12 +5,36 @@ import { Separator } from "@/components/ui/separator";
|
|||
import type { KnowledgeBase } from "@/features/rag/api/rag-api";
|
||||
import { KBDetailPanel } from "@/features/rag/components/kb-detail-panel";
|
||||
import { KBList } from "@/features/rag/components/kb-list";
|
||||
import { PreviewPanel } from "@/features/rag/components/preview-panel";
|
||||
import { RagDefaultsSection } from "@/features/rag/components/rag-defaults-section";
|
||||
import { ThreadIndexList } from "@/features/rag/components/thread-index-list";
|
||||
import { useState } from "react";
|
||||
import { useResizablePanelWidth } from "@/features/rag/hooks/use-resizable-width";
|
||||
import { usePreviewStore } from "@/features/rag/stores/preview-store";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { type CSSProperties, useState } from "react";
|
||||
|
||||
export function KnowledgeBasesTab() {
|
||||
const [selected, setSelected] = useState<KnowledgeBase | null>(null);
|
||||
const previewTarget = usePreviewStore((s) => s.target);
|
||||
const previewStatus = usePreviewStore((s) => s.status);
|
||||
const previewActive =
|
||||
previewTarget !== null ||
|
||||
previewStatus === "loading" ||
|
||||
previewStatus === "error";
|
||||
|
||||
const {
|
||||
width: previewWidth,
|
||||
isResizing: previewResizing,
|
||||
startResize: startPreviewResize,
|
||||
adjustWidth: adjustPreviewWidth,
|
||||
resetWidth: resetPreviewWidth,
|
||||
} = useResizablePanelWidth({
|
||||
storageKey: "unsloth:kb-preview-panel-width",
|
||||
defaultWidth: 560,
|
||||
minWidth: 320,
|
||||
maxWidthFraction: 0.7,
|
||||
enabled: previewActive,
|
||||
});
|
||||
|
||||
return (
|
||||
<div className="flex h-full min-h-0 flex-col gap-4">
|
||||
|
|
@ -36,6 +60,57 @@ export function KnowledgeBasesTab() {
|
|||
</div>
|
||||
)}
|
||||
</div>
|
||||
{previewActive ? (
|
||||
<>
|
||||
<button
|
||||
type="button"
|
||||
aria-label="Resize preview panel — drag, or use arrow keys"
|
||||
aria-orientation="vertical"
|
||||
onPointerDown={startPreviewResize}
|
||||
onDoubleClick={resetPreviewWidth}
|
||||
onKeyDown={(event) => {
|
||||
if (event.key === "ArrowLeft") {
|
||||
event.preventDefault();
|
||||
adjustPreviewWidth(24);
|
||||
} else if (event.key === "ArrowRight") {
|
||||
event.preventDefault();
|
||||
adjustPreviewWidth(-24);
|
||||
} else if (event.key === "Home") {
|
||||
event.preventDefault();
|
||||
resetPreviewWidth();
|
||||
}
|
||||
}}
|
||||
className={cn(
|
||||
"group relative hidden w-1.5 cursor-col-resize touch-none select-none border-0 bg-transparent p-0 outline-none lg:block",
|
||||
"before:absolute before:inset-y-0 before:left-1/2 before:-translate-x-1/2 before:w-px before:bg-border/70 before:transition-colors",
|
||||
"hover:before:bg-primary/40 focus-visible:before:bg-primary/60",
|
||||
previewResizing && "before:bg-primary/60",
|
||||
)}
|
||||
>
|
||||
<span
|
||||
aria-hidden="true"
|
||||
className={cn(
|
||||
"absolute left-1/2 top-1/2 h-10 w-[3px] -translate-x-1/2 -translate-y-1/2 rounded-full bg-border opacity-0 transition-opacity",
|
||||
"group-hover:opacity-100 group-focus-visible:opacity-100",
|
||||
previewResizing && "opacity-100",
|
||||
)}
|
||||
/>
|
||||
</button>
|
||||
<div
|
||||
className={cn(
|
||||
"w-0 shrink-0 overflow-hidden max-lg:hidden lg:w-[var(--preview-w)]",
|
||||
!previewResizing && "transition-[width] duration-200 ease-out",
|
||||
)}
|
||||
style={
|
||||
{
|
||||
"--preview-w": `${previewWidth}px`,
|
||||
} as CSSProperties
|
||||
}
|
||||
>
|
||||
<PreviewPanel open={previewActive} />
|
||||
</div>
|
||||
</>
|
||||
) : null}
|
||||
</div>
|
||||
<Separator />
|
||||
<ThreadIndexList />
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@
|
|||
--duration-normal: 200ms;
|
||||
|
||||
/* Easing curves (Emil Kowalski) */
|
||||
--ease-out-quad: cubic-bezier(0.25, 0.46, 0.45, 0.94);
|
||||
--ease-out-quart: cubic-bezier(0.165, 0.84, 0.44, 1);
|
||||
--ease-out-cubic: cubic-bezier(0.215, 0.61, 0.355, 1);
|
||||
|
||||
|
|
@ -1154,6 +1155,55 @@
|
|||
background: oklch(0.72 0 0 / 0.25);
|
||||
}
|
||||
|
||||
.preview-scrollbar {
|
||||
scrollbar-width: thin;
|
||||
scrollbar-color: oklch(0.5 0 0 / 0.25) transparent;
|
||||
}
|
||||
|
||||
.preview-scrollbar::-webkit-scrollbar {
|
||||
width: 10px;
|
||||
height: 10px;
|
||||
}
|
||||
|
||||
.preview-scrollbar::-webkit-scrollbar-track {
|
||||
background: transparent;
|
||||
}
|
||||
|
||||
.preview-scrollbar::-webkit-scrollbar-thumb {
|
||||
background: oklch(0.5 0 0 / 0.22);
|
||||
background-clip: padding-box;
|
||||
border: 2px solid transparent;
|
||||
border-radius: 9999px;
|
||||
}
|
||||
|
||||
.preview-scrollbar:hover::-webkit-scrollbar-thumb {
|
||||
background: oklch(0.5 0 0 / 0.38);
|
||||
background-clip: padding-box;
|
||||
border: 2px solid transparent;
|
||||
}
|
||||
|
||||
.dark .preview-scrollbar {
|
||||
scrollbar-color: oklch(0.72 0 0 / 0.25) transparent;
|
||||
}
|
||||
|
||||
.dark .preview-scrollbar::-webkit-scrollbar-thumb {
|
||||
background: oklch(0.72 0 0 / 0.25);
|
||||
background-clip: padding-box;
|
||||
border: 2px solid transparent;
|
||||
}
|
||||
|
||||
.dark .preview-scrollbar:hover::-webkit-scrollbar-thumb {
|
||||
background: oklch(0.78 0 0 / 0.42);
|
||||
background-clip: padding-box;
|
||||
border: 2px solid transparent;
|
||||
}
|
||||
|
||||
.preview-sheet-content {
|
||||
transition:
|
||||
transform var(--duration-normal) var(--ease-out-quad),
|
||||
opacity var(--duration-normal) var(--ease-out-quad);
|
||||
}
|
||||
|
||||
/*---break---*/
|
||||
|
||||
@layer base {
|
||||
|
|
|
|||
1
studio/frontend/src/setupTests.ts
Normal file
1
studio/frontend/src/setupTests.ts
Normal file
|
|
@ -0,0 +1 @@
|
|||
import "@testing-library/jest-dom";
|
||||
19
studio/frontend/vitest.config.ts
Normal file
19
studio/frontend/vitest.config.ts
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
import { resolve } from "node:path";
|
||||
import react from "@vitejs/plugin-react";
|
||||
import { defineConfig } from "vitest/config";
|
||||
|
||||
export default defineConfig({
|
||||
plugins: [react()],
|
||||
test: {
|
||||
environment: "jsdom",
|
||||
globals: true,
|
||||
setupFiles: ["./src/setupTests.ts"],
|
||||
include: ["src/__tests__/**/*.{test,spec}.{ts,tsx}"],
|
||||
exclude: ["node_modules", "dist"],
|
||||
},
|
||||
resolve: {
|
||||
alias: {
|
||||
"@": resolve(__dirname, "./src"),
|
||||
},
|
||||
},
|
||||
});
|
||||
93
tests/fixtures/rag-preview/make_fixture_pdf.py
vendored
Normal file
93
tests/fixtures/rag-preview/make_fixture_pdf.py
vendored
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
"""Generate tests/fixtures/rag-preview/sample.pdf deterministically.
|
||||
|
||||
Run once: python tests/fixtures/rag-preview/make_fixture_pdf.py
|
||||
Requires no third-party deps — builds a minimal valid single-page PDF
|
||||
using only stdlib so the fixture can be regenerated in any environment.
|
||||
The output is committed alongside this script so tests load it directly.
|
||||
"""
|
||||
|
||||
import os
|
||||
import struct
|
||||
import zlib
|
||||
from pathlib import Path
|
||||
|
||||
OUTPUT = Path(__file__).parent / "sample.pdf"
|
||||
|
||||
|
||||
def _compress(data: bytes) -> bytes:
|
||||
return zlib.compress(data, level=9)
|
||||
|
||||
|
||||
def _pdf() -> bytes:
|
||||
# Minimal PDF 1.4 with one page, one text stream.
|
||||
# Structure: header, catalog, pages, page, content stream, xref, trailer.
|
||||
page_text = b"BT /F1 12 Tf 72 720 Td (RAG preview fixture - page 1) Tj ET"
|
||||
compressed = _compress(page_text)
|
||||
stream_len = len(compressed)
|
||||
|
||||
objects: list[bytes] = []
|
||||
|
||||
def obj(n: int, body: bytes) -> bytes:
|
||||
return f"{n} 0 obj\n".encode() + body + b"\nendobj\n"
|
||||
|
||||
# 1: Catalog
|
||||
objects.append(obj(1, b"<< /Type /Catalog /Pages 2 0 R >>"))
|
||||
# 2: Pages
|
||||
objects.append(obj(2, b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>"))
|
||||
# 3: Page
|
||||
objects.append(
|
||||
obj(
|
||||
3,
|
||||
(
|
||||
b"<< /Type /Page /Parent 2 0 R "
|
||||
b"/MediaBox [0 0 612 792] "
|
||||
b"/Contents 4 0 R "
|
||||
b"/Resources << /Font << /F1 5 0 R >> >> >>"
|
||||
),
|
||||
)
|
||||
)
|
||||
# 4: Content stream
|
||||
objects.append(
|
||||
obj(
|
||||
4,
|
||||
(
|
||||
f"<< /Length {stream_len} /Filter /FlateDecode >>".encode()
|
||||
+ b"\nstream\n"
|
||||
+ compressed
|
||||
+ b"\nendstream"
|
||||
),
|
||||
)
|
||||
)
|
||||
# 5: Font
|
||||
objects.append(
|
||||
obj(
|
||||
5,
|
||||
b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
|
||||
)
|
||||
)
|
||||
|
||||
header = b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n"
|
||||
body = b""
|
||||
offsets: list[int] = []
|
||||
for o in objects:
|
||||
offsets.append(len(header) + len(body))
|
||||
body += o
|
||||
|
||||
xref_offset = len(header) + len(body)
|
||||
n = len(objects)
|
||||
xref = f"xref\n0 {n + 1}\n".encode()
|
||||
xref += b"0000000000 65535 f \n"
|
||||
for off in offsets:
|
||||
xref += f"{off:010d} 00000 n \n".encode()
|
||||
trailer = (
|
||||
f"trailer\n<< /Size {n + 1} /Root 1 0 R >>\n"
|
||||
f"startxref\n{xref_offset}\n%%EOF\n"
|
||||
).encode()
|
||||
|
||||
return header + body + xref + trailer
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pdf_bytes = _pdf()
|
||||
OUTPUT.write_bytes(pdf_bytes)
|
||||
print(f"Written {len(pdf_bytes)} bytes to {OUTPUT}")
|
||||
BIN
tests/fixtures/rag-preview/sample.pdf
vendored
Normal file
BIN
tests/fixtures/rag-preview/sample.pdf
vendored
Normal file
Binary file not shown.
8
tests/fixtures/rag-preview/sample.txt
vendored
Normal file
8
tests/fixtures/rag-preview/sample.txt
vendored
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
This is a test document for RAG preview fixtures.
|
||||
|
||||
Section 1: Introduction
|
||||
The operating margin rose to 18.2% in Q3, driven by improved efficiency.
|
||||
|
||||
Section 2: Details
|
||||
Additional supporting evidence and analysis is contained here.
|
||||
Page 1 of 1.
|
||||
Loading…
Add table
Add a link
Reference in a new issue