diff --git a/studio/backend/core/rag/authorization.py b/studio/backend/core/rag/authorization.py index ffb787a95f..dbd1787785 100644 --- a/studio/backend/core/rag/authorization.py +++ b/studio/backend/core/rag/authorization.py @@ -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: diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index 90c0d20316..857d8e54c6 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -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" diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 767bf65941..f8c286d165 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -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( """ diff --git a/studio/backend/core/rag/locators.py b/studio/backend/core/rag/locators.py index b01bdb3bce..db0dae319a 100644 --- a/studio/backend/core/rag/locators.py +++ b/studio/backend/core/rag/locators.py @@ -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: diff --git a/studio/backend/core/rag/parsers/pdf.py b/studio/backend/core/rag/parsers/pdf.py index 0f04def09b..4d087df929 100644 --- a/studio/backend/core/rag/parsers/pdf.py +++ b/studio/backend/core/rag/parsers/pdf.py @@ -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) diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index 0ca61b93d7..f92507b434 100644 --- a/studio/backend/core/rag/retrieval.py +++ b/studio/backend/core/rag/retrieval.py @@ -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] = [] diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 50fd283791..53ba6e2f05 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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" ] diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 5b0ebe5d27..47bd5ae436 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -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( diff --git a/studio/backend/tests/test_rag_authorization.py b/studio/backend/tests/test_rag_authorization.py index 7492bd295b..e16865588f 100644 --- a/studio/backend/tests/test_rag_authorization.py +++ b/studio/backend/tests/test_rag_authorization.py @@ -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 diff --git a/studio/backend/tests/test_rag_chunk_locators.py b/studio/backend/tests/test_rag_chunk_locators.py index f42fc31cf2..1d20d30eda 100644 --- a/studio/backend/tests/test_rag_chunk_locators.py +++ b/studio/backend/tests/test_rag_chunk_locators.py @@ -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) diff --git a/studio/backend/tests/test_rag_locator_backfill.py b/studio/backend/tests/test_rag_locator_backfill.py index 989dac14e5..2503ab1d53 100644 --- a/studio/backend/tests/test_rag_locator_backfill.py +++ b/studio/backend/tests/test_rag_locator_backfill.py @@ -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: diff --git a/studio/backend/tests/test_rag_locator_migration.py b/studio/backend/tests/test_rag_locator_migration.py index a9d368e3dd..627af23d05 100644 --- a/studio/backend/tests/test_rag_locator_migration.py +++ b/studio/backend/tests/test_rag_locator_migration.py @@ -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", diff --git a/studio/backend/tests/test_rag_preview_routes.py b/studio/backend/tests/test_rag_preview_routes.py index 2d568f6769..22979e917a 100644 --- a/studio/backend/tests/test_rag_preview_routes.py +++ b/studio/backend/tests/test_rag_preview_routes.py @@ -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= 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"") 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") diff --git a/studio/backend/tests/test_rag_preview_target_locators.py b/studio/backend/tests/test_rag_preview_target_locators.py index 97bb192776..e653461ad7 100644 --- a/studio/backend/tests/test_rag_preview_target_locators.py +++ b/studio/backend/tests/test_rag_preview_target_locators.py @@ -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: diff --git a/studio/backend/tests/test_rag_source_identity.py b/studio/backend/tests/test_rag_source_identity.py index 88ac4109c6..b949612bac 100644 --- a/studio/backend/tests/test_rag_source_identity.py +++ b/studio/backend/tests/test_rag_source_identity.py @@ -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 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) diff --git a/tests/fixtures/rag-preview/make_fixture_pdf.py b/tests/fixtures/rag-preview/make_fixture_pdf.py index e7c2c615c6..b40cd01907 100644 --- a/tests/fixtures/rag-preview/make_fixture_pdf.py +++ b/tests/fixtures/rag-preview/make_fixture_pdf.py @@ -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: