[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot]
parent
f25ea25570
commit
d6a7c9f8c7
16 changed files with 171 additions and 128 deletions
|
|
@ -57,7 +57,7 @@ def document_for_subject_or_404(
|
|||
`filename`, `content_type`, etc. without re-querying.
|
||||
"""
|
||||
if not document_id or not current_subject:
|
||||
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
|
||||
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)
|
||||
|
||||
with get_connection() as conn:
|
||||
row = conn.execute(
|
||||
|
|
@ -65,7 +65,7 @@ def document_for_subject_or_404(
|
|||
(document_id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
|
||||
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)
|
||||
|
||||
kb_id = row["kb_id"]
|
||||
thread_id = row["thread_id"]
|
||||
|
|
@ -76,10 +76,10 @@ def document_for_subject_or_404(
|
|||
(kb_id,),
|
||||
).fetchone()
|
||||
if owner_row is None:
|
||||
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
|
||||
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)
|
||||
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)
|
||||
return row
|
||||
|
||||
if thread_id is not None:
|
||||
|
|
@ -91,13 +91,13 @@ def document_for_subject_or_404(
|
|||
(thread_id,),
|
||||
).fetchone()
|
||||
if thread_row is None:
|
||||
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
|
||||
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)
|
||||
raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL)
|
||||
|
||||
|
||||
def chunk_belongs_to_document(chunk_id: str, document_id: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -170,12 +170,18 @@ def caption_images(
|
|||
if vlm_url and vlm_model:
|
||||
endpoint = f"{vlm_url.rstrip('/')}/v1/chat/completions"
|
||||
model_name = vlm_model
|
||||
logger.info("caption_images: using loaded chat VLM", endpoint = endpoint, model = model_name)
|
||||
logger.info(
|
||||
"caption_images: using loaded chat VLM",
|
||||
endpoint = endpoint,
|
||||
model = model_name,
|
||||
)
|
||||
else:
|
||||
logger.info("caption_images: chat VLM unavailable, loading helper")
|
||||
loaded = _load_helper_vlm()
|
||||
if loaded is None:
|
||||
logger.warning("caption_images: helper load failed, returning empty captions")
|
||||
logger.warning(
|
||||
"caption_images: helper load failed, returning empty captions"
|
||||
)
|
||||
return ["" for _ in image_bytes_list]
|
||||
helper_backend, helper_base_url, helper_model_name = loaded
|
||||
endpoint = f"{helper_base_url.rstrip('/')}/v1/chat/completions"
|
||||
|
|
|
|||
|
|
@ -590,7 +590,9 @@ def _replace_document_pages(document_id: str, pages: list[dict]) -> None:
|
|||
).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,))
|
||||
conn.execute(
|
||||
"DELETE FROM rag_document_pages WHERE document_id = ?", (document_id,)
|
||||
)
|
||||
if rows:
|
||||
conn.executemany(
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -111,7 +111,9 @@ def _find_normalized(page_text: str, needle: str) -> list[tuple[int, int]]:
|
|||
return out
|
||||
|
||||
|
||||
def _locate_unique(text: str, pages: list[ParsedPage]) -> tuple[LocatorMatch | None, str]:
|
||||
def _locate_unique(
|
||||
text: str, pages: list[ParsedPage]
|
||||
) -> tuple[LocatorMatch | None, str]:
|
||||
text = (text or "").strip()
|
||||
if not text:
|
||||
return None, "missing"
|
||||
|
|
@ -170,7 +172,9 @@ def _replace_document_pages(document_id: str, pages: list[ParsedPage]) -> None:
|
|||
for index, page in enumerate(pages)
|
||||
]
|
||||
with get_connection() as conn:
|
||||
conn.execute("DELETE FROM rag_document_pages WHERE document_id = ?", (document_id,))
|
||||
conn.execute(
|
||||
"DELETE FROM rag_document_pages WHERE document_id = ?", (document_id,)
|
||||
)
|
||||
if rows:
|
||||
conn.executemany(
|
||||
"""
|
||||
|
|
@ -340,7 +344,9 @@ def _scope_for_document(kb_id: str | None, thread_id: str | None) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
def _update_vector_payloads(scope: str | None, updates: dict[str, dict[str, Any]]) -> None:
|
||||
def _update_vector_payloads(
|
||||
scope: str | None, updates: dict[str, dict[str, Any]]
|
||||
) -> None:
|
||||
if not scope or not updates:
|
||||
return
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -105,12 +105,15 @@ def _extract_images_pymupdf(doc, pages: list[ParsedPage]) -> list[ParsedImage]:
|
|||
if union.width < _MIN_FIGURE_PT or union.height < _MIN_FIGURE_PT:
|
||||
continue
|
||||
# Expand and clip to page rect so we don't render past page edges.
|
||||
union = pymupdf.Rect(
|
||||
union.x0 - _FIGURE_MARGIN_PT,
|
||||
union.y0 - _FIGURE_MARGIN_PT,
|
||||
union.x1 + _FIGURE_MARGIN_PT,
|
||||
union.y1 + _FIGURE_MARGIN_PT,
|
||||
) & page.rect
|
||||
union = (
|
||||
pymupdf.Rect(
|
||||
union.x0 - _FIGURE_MARGIN_PT,
|
||||
union.y0 - _FIGURE_MARGIN_PT,
|
||||
union.x1 + _FIGURE_MARGIN_PT,
|
||||
union.y1 + _FIGURE_MARGIN_PT,
|
||||
)
|
||||
& page.rect
|
||||
)
|
||||
try:
|
||||
matrix = pymupdf.Matrix(_RENDER_SCALE, _RENDER_SCALE)
|
||||
pix = page.get_pixmap(clip = union, matrix = matrix, alpha = False)
|
||||
|
|
|
|||
|
|
@ -87,9 +87,7 @@ def retrieve_figure_refs(
|
|||
placeholders_docs = ""
|
||||
params: list = [scope]
|
||||
if document_ids:
|
||||
placeholders_docs = (
|
||||
f" AND document_id IN ({','.join('?' * len(document_ids))})"
|
||||
)
|
||||
placeholders_docs = f" AND document_id IN ({','.join('?' * len(document_ids))})"
|
||||
params.extend(document_ids)
|
||||
|
||||
like_clauses: list[str] = []
|
||||
|
|
|
|||
|
|
@ -337,8 +337,7 @@ def _drop_rag_tool_if_scope_empty(tools: list, rag_scope: Optional[dict]) -> lis
|
|||
if doc_count > 0:
|
||||
return tools
|
||||
return [
|
||||
t for t in tools
|
||||
if t.get("function", {}).get("name") != "search_knowledge_base"
|
||||
t for t in tools if t.get("function", {}).get("name") != "search_knowledge_base"
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1136,10 +1136,7 @@ def _content_disposition_header(filename: str, disposition: str) -> str:
|
|||
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}"
|
||||
)
|
||||
return f'{disposition}; filename="{ascii_fallback}"; ' f"filename*=UTF-8''{encoded}"
|
||||
|
||||
|
||||
def _preview_file_metadata(filename: str) -> tuple[str, str, PreviewMediaKind]:
|
||||
|
|
@ -1308,9 +1305,7 @@ def _serve_document_file_row(
|
|||
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"]
|
||||
)
|
||||
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),
|
||||
|
|
@ -1423,12 +1418,11 @@ def get_document_preview_target(
|
|||
detail = "Document not found",
|
||||
)
|
||||
|
||||
chunk_kind: PreviewChunkKind = (chunk_row["kind"] or "text") # type: ignore[assignment]
|
||||
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}"
|
||||
f"/api/rag/images/{doc_row['id']}/" f"{Path(chunk_row['image_path']).name}"
|
||||
)
|
||||
|
||||
return PreviewTargetResponse(
|
||||
|
|
|
|||
|
|
@ -24,7 +24,10 @@ 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
|
||||
from core.rag.authorization import (
|
||||
chunk_belongs_to_document,
|
||||
document_for_subject_or_404,
|
||||
)
|
||||
|
||||
|
||||
# ── Fixtures ──────────────────────────────────────────────────────────
|
||||
|
|
@ -103,7 +106,7 @@ def test_kb_doc_correct_owner_returns_row(tmp_path, monkeypatch):
|
|||
_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(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
|
||||
|
|
@ -114,7 +117,7 @@ def test_kb_doc_wrong_owner_raises_404(tmp_path, monkeypatch):
|
|||
_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(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")
|
||||
|
|
@ -127,7 +130,7 @@ def test_kb_doc_null_owner_raises_404(tmp_path, monkeypatch):
|
|||
_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(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")
|
||||
|
|
@ -146,7 +149,7 @@ def test_kb_doc_missing_kb_raises_404(tmp_path, monkeypatch):
|
|||
_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(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,))
|
||||
|
|
@ -216,7 +219,7 @@ def test_empty_subject_raises_404(tmp_path, monkeypatch):
|
|||
_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(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, "")
|
||||
|
|
@ -231,7 +234,7 @@ def test_chunk_belongs_returns_true_for_matching_doc(tmp_path, monkeypatch):
|
|||
_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(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
|
||||
|
|
@ -243,7 +246,7 @@ def test_chunk_belongs_returns_false_for_wrong_doc(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(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)
|
||||
|
|
@ -256,7 +259,7 @@ def test_chunk_belongs_returns_false_for_missing_chunk(tmp_path, monkeypatch):
|
|||
_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(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
|
||||
|
||||
|
|
|
|||
|
|
@ -29,17 +29,17 @@ def _token_count(text: str) -> int:
|
|||
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,
|
||||
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", " ", ""),
|
||||
max_tokens = 3,
|
||||
overlap_tokens = 0,
|
||||
token_counter = _token_count,
|
||||
separators = ("\n", " ", ""),
|
||||
)
|
||||
|
||||
target = next(chunk for chunk in chunks if "beta" in chunk.text)
|
||||
|
|
@ -53,16 +53,16 @@ def test_standard_chunking_records_page_local_char_and_line_spans():
|
|||
|
||||
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),
|
||||
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", " ", ""),
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ import storage.studio_db as studio_db
|
|||
from auth.authentication import get_current_subject
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
@pytest.fixture(scope = "module")
|
||||
def app():
|
||||
import sys
|
||||
|
||||
|
|
@ -38,7 +38,7 @@ def _uid() -> str:
|
|||
|
||||
def _make_client(app, subject: str = "alice"):
|
||||
app.dependency_overrides[get_current_subject] = lambda: subject
|
||||
return TestClient(app, raise_server_exceptions=True)
|
||||
return TestClient(app, raise_server_exceptions = True)
|
||||
|
||||
|
||||
def _clear_overrides(app):
|
||||
|
|
@ -76,8 +76,8 @@ def _insert_chunk(conn, chunk_id: str, doc_id: str, text: str) -> None:
|
|||
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")
|
||||
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:
|
||||
|
|
@ -111,8 +111,8 @@ def test_backfill_preserves_ids_and_updates_unique_locator(app, db_env, monkeypa
|
|||
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")
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -29,8 +29,7 @@ def test_locator_schema_is_additive_and_nullable(tmp_path, monkeypatch):
|
|||
}.issubset(chunk_cols)
|
||||
|
||||
page_cols = {
|
||||
row["name"]
|
||||
for row in conn.execute("PRAGMA table_info(rag_document_pages)")
|
||||
row["name"] for row in conn.execute("PRAGMA table_info(rag_document_pages)")
|
||||
}
|
||||
assert {
|
||||
"document_id",
|
||||
|
|
|
|||
|
|
@ -42,13 +42,15 @@ from auth.authentication import get_current_subject
|
|||
# ── App import (deferred to avoid import-time side-effects) ───────────
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
@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
|
||||
|
||||
|
||||
|
|
@ -70,7 +72,7 @@ def _uid() -> str:
|
|||
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)
|
||||
client = TestClient(app, raise_server_exceptions = True)
|
||||
return client
|
||||
|
||||
|
||||
|
|
@ -128,18 +130,20 @@ class TestPreviewTarget:
|
|||
"""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.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)
|
||||
_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}")
|
||||
resp = client.get(
|
||||
f"/api/rag/documents/{doc_id}/preview-target?chunk_id={chunk_id}"
|
||||
)
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
|
||||
|
|
@ -152,18 +156,20 @@ class TestPreviewTarget:
|
|||
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):
|
||||
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.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)
|
||||
_insert_chunk(conn, chunk_id, doc_id, page_number = 7, chunk_index = 14)
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE rag_chunks
|
||||
|
|
@ -180,7 +186,9 @@ class TestPreviewTarget:
|
|||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
resp = client.get(f"/api/rag/documents/{doc_id}/preview-target?chunk_id={chunk_id}")
|
||||
resp = client.get(
|
||||
f"/api/rag/documents/{doc_id}/preview-target?chunk_id={chunk_id}"
|
||||
)
|
||||
finally:
|
||||
_clear_overrides(app)
|
||||
|
||||
|
|
@ -199,11 +207,13 @@ class TestPreviewTarget:
|
|||
}
|
||||
]
|
||||
|
||||
def test_without_chunk_id_returns_all_null_chunk_fields(self, app, db_env, monkeypatch):
|
||||
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.parent.mkdir(parents = True, exist_ok = True)
|
||||
stored.write_bytes(b"%PDF-1.4 dummy")
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
|
||||
|
|
@ -243,12 +253,12 @@ class TestPreviewTarget:
|
|||
"""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.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_kb(conn, kb_id, owner = "alice")
|
||||
_insert_doc(conn, doc_id, kb_id, str(stored))
|
||||
|
||||
client = _make_client(app, "mallory")
|
||||
|
|
@ -265,7 +275,7 @@ class TestPreviewTarget:
|
|||
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.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))
|
||||
|
|
@ -291,7 +301,7 @@ class TestPreviewTarget:
|
|||
"""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)
|
||||
client = TestClient(app, raise_server_exceptions = False)
|
||||
resp = client.get(f"/api/rag/documents/{_uid()}/preview-target")
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
|
@ -304,14 +314,16 @@ class TestFileRoute:
|
|||
"""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)
|
||||
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")
|
||||
_insert_doc(
|
||||
conn, doc_id, kb_id, str(stored), "annual.pdf", "application/pdf"
|
||||
)
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
|
|
@ -325,11 +337,13 @@ class TestFileRoute:
|
|||
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):
|
||||
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)
|
||||
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))
|
||||
|
|
@ -337,7 +351,9 @@ class TestFileRoute:
|
|||
|
||||
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")
|
||||
_insert_doc(
|
||||
conn, doc_id, kb_id, str(stored), "annual.pdf", "application/pdf"
|
||||
)
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
|
|
@ -347,13 +363,16 @@ class TestFileRoute:
|
|||
assert "Bearer" not in signed_url
|
||||
assert "Authorization" not in signed_url
|
||||
|
||||
file_resp = client.get(signed_url, headers={"Range": "bytes=0-3"})
|
||||
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("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"
|
||||
|
||||
|
|
@ -361,7 +380,7 @@ class TestFileRoute:
|
|||
"""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)
|
||||
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))
|
||||
|
|
@ -369,17 +388,21 @@ class TestFileRoute:
|
|||
|
||||
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")
|
||||
_insert_doc(
|
||||
conn, doc_id, kb_id, str(stored), "annual.pdf", "application/pdf"
|
||||
)
|
||||
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
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):
|
||||
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)
|
||||
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))
|
||||
|
|
@ -407,15 +430,21 @@ class TestFileRoute:
|
|||
"""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)
|
||||
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")
|
||||
_insert_doc(
|
||||
conn,
|
||||
doc_id,
|
||||
kb_id,
|
||||
str(stored),
|
||||
"report.docx",
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
)
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
try:
|
||||
|
|
@ -441,13 +470,13 @@ class TestFileRoute:
|
|||
"""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)
|
||||
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_kb(conn, kb_id, owner = "alice")
|
||||
_insert_doc(conn, doc_id, kb_id, str(stored))
|
||||
|
||||
client = _make_client(app, "mallory")
|
||||
|
|
@ -457,11 +486,13 @@ class TestFileRoute:
|
|||
_clear_overrides(app)
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_deleted_file_returns_404_with_doc_file_not_found(self, app, db_env, monkeypatch):
|
||||
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)
|
||||
uploads.mkdir(parents = True, exist_ok = True)
|
||||
stored = uploads / "gone.pdf"
|
||||
stored.write_bytes(b"data")
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
|
|
@ -483,14 +514,16 @@ class TestFileRoute:
|
|||
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):
|
||||
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)
|
||||
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.parent.mkdir(parents = True, exist_ok = True)
|
||||
outside.write_bytes(b"root:x:0:0")
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
|
||||
|
|
@ -518,7 +551,7 @@ class TestFileRoute:
|
|||
"""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)
|
||||
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))
|
||||
|
|
@ -548,7 +581,7 @@ class TestImageRoute:
|
|||
doc_id, kb_id = _uid(), _uid()
|
||||
uploads = db_env / "rag" / "uploads"
|
||||
images = uploads / "images" / doc_id
|
||||
images.mkdir(parents=True, exist_ok=True)
|
||||
images.mkdir(parents = True, exist_ok = True)
|
||||
image = images / "figure.png"
|
||||
image.write_bytes(b"\x89PNG\r\n\x1a\n")
|
||||
stored = uploads / "report.pdf"
|
||||
|
|
@ -556,7 +589,7 @@ class TestImageRoute:
|
|||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id, owner="alice")
|
||||
_insert_kb(conn, kb_id, owner = "alice")
|
||||
_insert_doc(conn, doc_id, kb_id, str(stored), "report.pdf")
|
||||
|
||||
client = _make_client(app, "mallory")
|
||||
|
|
@ -572,7 +605,7 @@ class TestImageRoute:
|
|||
doc_id, kb_id = _uid(), _uid()
|
||||
uploads = db_env / "rag" / "uploads"
|
||||
images = uploads / "images" / doc_id
|
||||
images.mkdir(parents=True, exist_ok=True)
|
||||
images.mkdir(parents = True, exist_ok = True)
|
||||
image = images / "figure.png"
|
||||
image.write_bytes(b"\x89PNG\r\n\x1a\n")
|
||||
stored = uploads / "report.pdf"
|
||||
|
|
@ -580,7 +613,7 @@ class TestImageRoute:
|
|||
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
_insert_kb(conn, kb_id, owner="alice")
|
||||
_insert_kb(conn, kb_id, owner = "alice")
|
||||
_insert_doc(conn, doc_id, kb_id, str(stored), "report.pdf")
|
||||
|
||||
client = _make_client(app, "alice")
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ import storage.studio_db as studio_db
|
|||
from auth.authentication import get_current_subject
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
@pytest.fixture(scope = "module")
|
||||
def app():
|
||||
import sys
|
||||
|
||||
|
|
@ -38,7 +38,7 @@ def _uid() -> str:
|
|||
|
||||
def _make_client(app, subject: str = "alice"):
|
||||
app.dependency_overrides[get_current_subject] = lambda: subject
|
||||
return TestClient(app, raise_server_exceptions=True)
|
||||
return TestClient(app, raise_server_exceptions = True)
|
||||
|
||||
|
||||
def _clear_overrides(app):
|
||||
|
|
@ -68,7 +68,7 @@ def _seed_doc(conn, doc_id: str, kb_id: str, stored_path: str) -> None:
|
|||
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.parent.mkdir(parents = True, exist_ok = True)
|
||||
stored.write_bytes(b"%PDF-1.4")
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
|
|
@ -104,7 +104,7 @@ def test_preview_target_returns_nullable_locator_fields(app, db_env):
|
|||
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.parent.mkdir(parents = True, exist_ok = True)
|
||||
stored.write_bytes(b"%PDF-1.4")
|
||||
|
||||
with studio_db.get_connection() as conn:
|
||||
|
|
|
|||
|
|
@ -77,7 +77,7 @@ def _parse_chunks(xml_output: str) -> list[dict]:
|
|||
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)]
|
||||
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
|
||||
|
|
@ -88,8 +88,8 @@ def test_format_hits_emits_document_id_and_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)
|
||||
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
|
||||
|
|
@ -99,13 +99,13 @@ def test_citation_id_is_sequential_counter_not_uuid():
|
|||
|
||||
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_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"),
|
||||
_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)
|
||||
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)
|
||||
|
|
@ -125,8 +125,8 @@ def test_same_filename_docs_have_distinct_document_ids():
|
|||
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),
|
||||
_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)
|
||||
|
|
@ -143,8 +143,8 @@ def test_same_filename_docs_have_distinct_citation_ids():
|
|||
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),
|
||||
_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)
|
||||
|
|
@ -163,7 +163,7 @@ def test_empty_hits_returns_no_chunks_message():
|
|||
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)]
|
||||
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"
|
||||
|
|
@ -172,7 +172,7 @@ def test_page_number_attribute_present_when_page_exists():
|
|||
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)]
|
||||
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]}"
|
||||
|
|
@ -181,7 +181,7 @@ def test_page_number_attribute_absent_when_null():
|
|||
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 = _hit(chunk_id = chunk_id, document_id = doc_id, page_number = 5)
|
||||
hit.update(
|
||||
{
|
||||
"source_page_index": 4,
|
||||
|
|
@ -207,9 +207,9 @@ def test_xml_special_chars_in_filename_escaped():
|
|||
chunk_id, doc_id = _uid(), _uid()
|
||||
hits = [
|
||||
_hit(
|
||||
chunk_id=chunk_id,
|
||||
document_id=doc_id,
|
||||
filename='report <2025> "final" & draft.pdf',
|
||||
chunk_id = chunk_id,
|
||||
document_id = doc_id,
|
||||
filename = 'report <2025> "final" & draft.pdf',
|
||||
)
|
||||
]
|
||||
output = _format_hits_for_llm(hits)
|
||||
|
|
@ -229,7 +229,7 @@ def test_multiple_hits_carry_independent_ids():
|
|||
(_uid(), _uid()),
|
||||
]
|
||||
hits = [
|
||||
_hit(chunk_id=cid, document_id=did, filename=f"doc{i}.pdf")
|
||||
_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)
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ OUTPUT = Path(__file__).parent / "sample.pdf"
|
|||
|
||||
|
||||
def _compress(data: bytes) -> bytes:
|
||||
return zlib.compress(data, level=9)
|
||||
return zlib.compress(data, level = 9)
|
||||
|
||||
|
||||
def _pdf() -> bytes:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue