diff --git a/app/config.py b/app/config.py index 8031763a..6c141533 100644 --- a/app/config.py +++ b/app/config.py @@ -90,6 +90,45 @@ def get_env_variable( env_value = get_env_variable("PDF_EXTRACT_IMAGES", "False").lower() PDF_EXTRACT_IMAGES = True if env_value == "true" else False +# Optional pre-extraction webhook: when text extraction returns pages averaging +# fewer than PRE_EXTRACTION_WEBHOOK_MIN_CHARS characters, the file is POSTed to +# PRE_EXTRACTION_WEBHOOK_URL. The webhook is expected to respond with +# `{"text": "...", "provider": "..."}` and that text replaces the original +# extraction. Disabled when the URL is empty (the default). +PRE_EXTRACTION_WEBHOOK_URL = get_env_variable("PRE_EXTRACTION_WEBHOOK_URL", "") +PRE_EXTRACTION_WEBHOOK_MIN_CHARS = int( + get_env_variable("PRE_EXTRACTION_WEBHOOK_MIN_CHARS", "100") +) +PRE_EXTRACTION_WEBHOOK_TIMEOUT = int( + get_env_variable("PRE_EXTRACTION_WEBHOOK_TIMEOUT", "60") +) + +# Optional multimodal-RAG visual pipeline. When VISUAL_EMBED_URL is set, each +# ingested PDF is rendered to page PNGs via pdftoppm and each page is sent to +# the CLIP embed sidecar; the resulting vectors are stored in the +# `visual_chunks` pgvector table for retrieval. All steps soft-fail: if +# pdftoppm, the sidecar, or the DB is unreachable, the text ingest path still +# succeeds. +VISUAL_EMBED_URL = get_env_variable("VISUAL_EMBED_URL", "") +# Text-query endpoint of the same CLIP sidecar. Empty → derive from +# VISUAL_EMBED_URL by swapping /embed/image for /embed/text. +VISUAL_TEXT_EMBED_URL = get_env_variable("VISUAL_TEXT_EMBED_URL", "") +VISUAL_PAGE_DPI = int(get_env_variable("VISUAL_PAGE_DPI", "100")) +VISUAL_STORAGE_ROOT = get_env_variable("VISUAL_STORAGE_ROOT", "/var/rag-visual") +VISUAL_EMBED_TIMEOUT = int(get_env_variable("VISUAL_EMBED_TIMEOUT", "30")) +VISUAL_SCORE_THRESHOLD = float(get_env_variable("VISUAL_SCORE_THRESHOLD", "0.25")) +VISUAL_QUERY_TOP_K = int(get_env_variable("VISUAL_QUERY_TOP_K", "3")) +# Text-to-visual-page-coupling (Phase 4). When True, /query uses the pages of +# the text-search hits as the primary source of visual_matches, with CLIP +# cross-modal retrieval as a secondary signal. Rationale: CLIP-style models +# weight text-inside-images, so a page with a matching headline can outrank +# the page with the actually-relevant photo. The text pipeline already knows +# which pages are relevant — couple the visuals to that decision. +VISUAL_TEXT_COUPLED = get_env_variable("VISUAL_TEXT_COUPLED", "true").lower() == "true" +VISUAL_TEXT_COUPLED_MAX_PAGES = int( + get_env_variable("VISUAL_TEXT_COUPLED_MAX_PAGES", "4") +) + if POSTGRES_USE_UNIX_SOCKET: connection_suffix = f"{urllib.parse.quote_plus(POSTGRES_USER)}:{urllib.parse.quote_plus(POSTGRES_PASSWORD)}@/{urllib.parse.quote_plus(POSTGRES_DB)}?host={urllib.parse.quote_plus(DB_HOST)}" else: diff --git a/app/models.py b/app/models.py index 835764f9..45b5e7c4 100644 --- a/app/models.py +++ b/app/models.py @@ -31,6 +31,11 @@ class QueryRequestBody(BaseModel): file_id: str k: int = 4 entity_id: Optional[str] = None + # When True, the /query response is wrapped as + # {"chunks": [...], "visual_matches": [...]} and the visual pipeline + # (multimodal-RAG) is consulted in addition to the text index. + # Default False keeps the legacy flat-list shape. + include_visual: bool = False class CleanupMethod(str, Enum): @@ -42,3 +47,11 @@ class QueryMultipleBody(BaseModel): query: str file_ids: List[str] k: int = 4 + include_visual: bool = False + + +class VisualMatch(BaseModel): + file_id: str + page_number: int + image_path: str + score: float diff --git a/app/routes/document_routes.py b/app/routes/document_routes.py index c5ebd3f5..435afc30 100644 --- a/app/routes/document_routes.py +++ b/app/routes/document_routes.py @@ -52,6 +52,19 @@ clean_text, process_documents, cleanup_temp_encoding_file, + maybe_enrich_with_webhook, +) +from app.utils.visual_embed import ( + embed_text_query, + fetch_visual_chunks_for_pages, + maybe_embed_visuals, + similarity_search_visual, +) +from app.config import ( + VISUAL_EMBED_URL, + VISUAL_QUERY_TOP_K, + VISUAL_TEXT_COUPLED, + VISUAL_TEXT_COUPLED_MAX_PAGES, ) from app.utils.health import is_health_ok @@ -147,6 +160,11 @@ async def load_file_content( loader, known_type, file_ext = get_loader(filename, content_type, file_path) loop = asyncio.get_running_loop() data = await loop.run_in_executor(executor, lambda: list(loader.lazy_load())) + # Optional pre-extraction webhook (e.g. OCR sidecar) — no-op when + # PRE_EXTRACTION_WEBHOOK_URL is unset. + data = await loop.run_in_executor( + executor, maybe_enrich_with_webhook, file_path, data + ) return data, known_type, file_ext finally: # Clean up temporary UTF-8 file if it was created for encoding conversion @@ -311,6 +329,126 @@ def get_cached_query_embedding(query: str): return vector_store.embedding_function.embed_query(query) +def _pages_by_file_from_text_docs(text_documents) -> dict: + """Extract {file_id: [page_number, ...]} from text similarity_search results. + + Preserves rank order (earliest text hit → first page in the list) and + dedupes within each file. Reads ``page_number`` first, falling back to + ``page`` — PyPDFLoader writes ``page`` while custom loaders and older + chunks may use ``page_number``. Chunks with neither are silently + skipped (e.g. text-only formats like .txt that have no page concept). + """ + pages_by_file: dict = {} + if not text_documents: + return pages_by_file + for entry in text_documents: + doc = entry[0] if isinstance(entry, tuple) else entry + metadata = getattr(doc, "metadata", None) or {} + file_id = metadata.get("file_id") + page_number = metadata.get("page_number") + if page_number is None: + page_number = metadata.get("page") + if not file_id or page_number is None: + continue + try: + page_number = int(page_number) + except (TypeError, ValueError): + continue + bucket = pages_by_file.setdefault(file_id, []) + if page_number not in bucket: + bucket.append(page_number) + return pages_by_file + + +async def _fetch_visual_matches_for_file_ids( + request: Request, + query: str, + file_ids: List[str], + text_documents: Optional[list] = None, +) -> list: + """Cross-modal retrieval for /query's visual_matches field. + + Two signals are merged when ``VISUAL_TEXT_COUPLED`` is enabled: + + 1. **Primary** — text-page-coupling: use the pages of the text-search + hits (per file_id) to look up visual_chunks directly. Fixes the + CLIP-weights-text-inside-images bias: a page whose headline matches + the query outranks the page with the actually-relevant photo. + 2. **Secondary** — CLIP cross-modal: text query embedded via the CLIP + text encoder, nearest neighbours in ``visual_chunks``. Keeps + visually-relevant pages that the text pipeline missed. + + Dedup is by (file_id, page_number); primary wins. Soft-fails on any + error (returns []) so a broken sidecar never breaks the main query path. + """ + if not VISUAL_EMBED_URL or not file_ids: + return [] + + results: list = [] + seen: set = set() + + # Primary signal: text-page-coupling. + if VISUAL_TEXT_COUPLED: + try: + from app.services.database import PSQLDatabase + + pool = await PSQLDatabase.get_pool() + pages_by_file = _pages_by_file_from_text_docs(text_documents) + for file_id, pages in pages_by_file.items(): + if file_id not in file_ids: + continue + capped = pages[:VISUAL_TEXT_COUPLED_MAX_PAGES] + rows = await fetch_visual_chunks_for_pages( + pool=pool, file_id=file_id, page_numbers=capped + ) + # Preserve text-rank order by reindexing from ``capped``. + by_page = {r["page_number"]: r for r in rows} + for page in capped: + row = by_page.get(page) + if not row: + continue + key = (row["file_id"], row["page_number"]) + if key in seen: + continue + seen.add(key) + # Synthetic score above VISUAL_SCORE_THRESHOLD so downstream + # filters treat text-coupled hits as top-tier. + results.append({**row, "score": 1.0, "source": "text_coupled"}) + except Exception as exc: + logger.warning("visual retrieval: text-page-coupling failed: %s", exc) + + # Secondary signal: CLIP cross-modal. + try: + loop = asyncio.get_running_loop() + query_embedding = await loop.run_in_executor( + request.app.state.thread_pool, embed_text_query, query + ) + except Exception as exc: + logger.warning("visual retrieval: text embed failed: %s", exc) + return results + try: + from app.services.database import PSQLDatabase + + pool = await PSQLDatabase.get_pool() + clip_hits = await similarity_search_visual( + pool=pool, + query_embedding=query_embedding, + file_ids=file_ids, + k=VISUAL_QUERY_TOP_K, + ) + except Exception as exc: + logger.warning("visual retrieval: similarity search failed: %s", exc) + return results + + for hit in clip_hits: + key = (hit["file_id"], hit["page_number"]) + if key in seen: + continue + seen.add(key) + results.append({**hit, "source": "clip"}) + return results + + @router.post("/query") async def query_embeddings_by_file_id( body: QueryRequestBody, @@ -341,6 +479,11 @@ async def query_embeddings_by_file_id( ) if not documents: + if body.include_visual: + visual = await _fetch_visual_matches_for_file_ids( + request, body.query, [body.file_id], text_documents=documents + ) + return {"chunks": authorized_documents, "visual_matches": visual} return authorized_documents document, score = documents[0] @@ -369,6 +512,14 @@ async def query_embeddings_by_file_id( f"Unauthorized access attempt by user {user_authorized} to a document with user_id {doc_user_id}" ) + if body.include_visual: + visual = await _fetch_visual_matches_for_file_ids( + request, + body.query, + [body.file_id], + text_documents=authorized_documents, + ) + return {"chunks": authorized_documents, "visual_matches": visual} return authorized_documents except HTTPException as http_exc: @@ -756,6 +907,12 @@ async def embed_local_file( data = await loop.run_in_executor( request.app.state.thread_pool, lambda: list(loader.lazy_load()) ) + data = await loop.run_in_executor( + request.app.state.thread_pool, + maybe_enrich_with_webhook, + file_path, + data, + ) result = await store_data_in_vector_db( data, @@ -765,6 +922,15 @@ async def embed_local_file( executor=request.app.state.thread_pool, ) + # Visual pipeline (multimodal-RAG). Soft-fails by design; never raises. + await maybe_embed_visuals( + file_path=file_path, + file_id=document.file_id, + file_ext=file_ext, + user_id=user_id, + executor=request.app.state.thread_pool, + ) + if result: return { "status": True, @@ -841,6 +1007,15 @@ async def embed_file( executor=request.app.state.thread_pool, ) + # Visual pipeline (multimodal-RAG). Soft-fails by design; never raises. + await maybe_embed_visuals( + file_path=validated_file_path, + file_id=file_id, + file_ext=file_ext, + user_id=user_id, + executor=request.app.state.thread_pool, + ) + if not result: response_status = False response_message = "Failed to process/store the file data." @@ -977,6 +1152,15 @@ async def embed_file_upload( executor=request.app.state.thread_pool, ) + # Visual pipeline (multimodal-RAG). Soft-fails by design; never raises. + await maybe_embed_visuals( + file_path=validated_temp_file_path, + file_id=file_id, + file_ext=file_ext, + user_id=user_id, + executor=request.app.state.thread_pool, + ) + if not result: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -1031,6 +1215,14 @@ async def query_embeddings_by_file_ids(request: Request, body: QueryMultipleBody embedding, k=body.k, filter={"file_id": {"$in": body.file_ids}} ) + if body.include_visual: + visual = await _fetch_visual_matches_for_file_ids( + request, body.query, body.file_ids, text_documents=documents + ) + # When include_visual is True, we must not 404 just because the + # text path came back empty — visual matches might still be useful. + return {"chunks": documents, "visual_matches": visual} + # Ensure documents list is not empty if not documents: raise HTTPException( diff --git a/app/utils/document_loader.py b/app/utils/document_loader.py index 900b1561..e589c5b1 100644 --- a/app/utils/document_loader.py +++ b/app/utils/document_loader.py @@ -9,7 +9,15 @@ from langchain_core.documents import Document -from app.config import known_source_ext, PDF_EXTRACT_IMAGES, CHUNK_OVERLAP, logger +from app.config import ( + known_source_ext, + PDF_EXTRACT_IMAGES, + CHUNK_OVERLAP, + PRE_EXTRACTION_WEBHOOK_URL, + PRE_EXTRACTION_WEBHOOK_MIN_CHARS, + PRE_EXTRACTION_WEBHOOK_TIMEOUT, + logger, +) from langchain_community.document_loaders import ( TextLoader, PyPDFLoader, @@ -188,6 +196,83 @@ def remove_non_utf8(text: str) -> str: return text +def maybe_enrich_with_webhook( + file_path: str, documents: List[Document] +) -> List[Document]: + """ + Optional hook: when PRE_EXTRACTION_WEBHOOK_URL is set and the current + extraction returned pages averaging fewer than + PRE_EXTRACTION_WEBHOOK_MIN_CHARS characters, POST the original file to the + webhook and substitute its returned text. On any failure the original + documents are returned unchanged (soft-fail by design). + """ + url = PRE_EXTRACTION_WEBHOOK_URL + if not url: + return documents + + # Compute average characters per extracted page; empty list counts as 0. + if documents: + avg_chars = sum( + len((doc.page_content or "").strip()) for doc in documents + ) / len(documents) + else: + avg_chars = 0 + + if avg_chars >= PRE_EXTRACTION_WEBHOOK_MIN_CHARS: + return documents + + # Import lazily so module import works in environments without `requests` + # (e.g. test collection phase when the feature is disabled). + import requests + + try: + with open(file_path, "rb") as f: + response = requests.post( + url, + files={"file": (os.path.basename(file_path), f)}, + timeout=PRE_EXTRACTION_WEBHOOK_TIMEOUT, + ) + response.raise_for_status() + payload = response.json() + except Exception as exc: # broad by intent — never break ingest + logger.warning( + "pre-extraction webhook failed, falling back to original text: %s", exc + ) + return documents + + text = (payload.get("text") or "").strip() + if not text: + logger.warning( + "pre-extraction webhook returned empty text, keeping original extraction" + ) + return documents + + provider = payload.get("provider", "unknown") + source = ( + documents[0].metadata.get("source") + if documents and isinstance(documents[0].metadata, dict) + else file_path + ) + + logger.info( + "pre-extraction webhook enriched %s with %d chars from provider %s", + file_path, + len(text), + provider, + ) + + return [ + Document( + page_content=text, + metadata={ + "source": source, + "ocr_used": True, + "ocr_provider": provider, + }, + ) + ] + + def process_documents(documents: List[Document]) -> str: processed_text = "" last_page: Optional[int] = None diff --git a/app/utils/visual_embed.py b/app/utils/visual_embed.py new file mode 100644 index 00000000..3560f341 --- /dev/null +++ b/app/utils/visual_embed.py @@ -0,0 +1,311 @@ +"""Optional visual ingest pipeline for Multimodal-RAG. + +Sits alongside the existing text pipeline. When VISUAL_EMBED_URL is set +and the uploaded file is a PDF, this module: + + 1. Uses PyMuPDF (fitz) to render each page as a PNG at VISUAL_PAGE_DPI. + 2. POSTs each PNG to the CLIP embed service (``VISUAL_EMBED_URL``). + 3. Persists (file_id, page_number, image_path, embedding) into the + ``visual_chunks`` pgvector table. + +Everything is soft-fail by design — a broken sidecar, PyMuPDF raising +on a malformed PDF, or unreachable database must NEVER break the text +ingest path. When disabled (URL empty) the helper is a no-op and +imports are lazy. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import shutil +from pathlib import Path +from typing import List, Optional + +from app.config import ( + VISUAL_EMBED_URL, + VISUAL_EMBED_TIMEOUT, + VISUAL_PAGE_DPI, + VISUAL_SCORE_THRESHOLD, + VISUAL_STORAGE_ROOT, + VISUAL_TEXT_EMBED_URL, + logger, +) + + +def _visuals_enabled() -> bool: + return bool(VISUAL_EMBED_URL) + + +def _page_number_from_path(p: Path) -> Optional[int]: + """Parse ``prefix-3.png`` or ``prefix-003.png`` → 3. Returns None if unparseable.""" + try: + return int(p.stem.rsplit("-", 1)[-1]) + except (ValueError, IndexError): + return None + + +def _vector_literal(embedding: List[float]) -> str: + """pgvector's text input format: '[0.1,0.2,...]' (no spaces to keep it compact).""" + return "[" + ",".join(f"{v:.6f}" for v in embedding) + "]" + + +def render_pdf_pages(pdf_path: str, out_dir: Path, dpi: int) -> List[Path]: + """Render each PDF page to ``out_dir/page-N.png`` via PyMuPDF. + + Switched away from `pdftoppm` so the container image doesn't need + poppler-utils installed — PyMuPDF ships as a self-contained wheel. + Raises RuntimeError on any failure so the caller can soft-fail. + Kept as a sync function: called from a ThreadPoolExecutor. + """ + out_dir.mkdir(parents=True, exist_ok=True) + try: + import fitz # PyMuPDF + except ImportError as exc: + raise RuntimeError("PyMuPDF not installed") from exc + + # Zoom factor maps DPI onto the PDF's default 72-dpi coordinate space. + zoom = dpi / 72.0 + matrix = fitz.Matrix(zoom, zoom) + page_paths: List[Path] = [] + try: + with fitz.open(pdf_path) as doc: + # Zero-pad page numbers to match the pdftoppm naming convention + # our callers (and tests) expect: page-01.png … page-NN.png + width = max(2, len(str(doc.page_count))) + for i in range(doc.page_count): + page = doc.load_page(i) + pix = page.get_pixmap(matrix=matrix, alpha=False) + out = out_dir / f"page-{str(i + 1).zfill(width)}.png" + pix.save(str(out)) + page_paths.append(out) + except Exception as exc: + raise RuntimeError(f"PyMuPDF render failed: {exc}") from exc + + return sorted(page_paths) + + +def embed_image(image_path: Path) -> List[float]: + """Synchronous HTTP call to the CLIP sidecar. Raises on error.""" + import requests + + with open(image_path, "rb") as f: + resp = requests.post( + VISUAL_EMBED_URL, + files={"file": (image_path.name, f, "image/png")}, + timeout=VISUAL_EMBED_TIMEOUT, + ) + resp.raise_for_status() + payload = resp.json() + embedding = payload.get("embedding") + if not isinstance(embedding, list) or not embedding: + raise RuntimeError(f"clip sidecar returned malformed payload: {payload!r}") + return embedding + + +def embed_text_query(query: str) -> List[float]: + """Call the CLIP text endpoint so the returned vector lives in the + same space as image embeddings. Used by the /query visual path.""" + import requests + + url = VISUAL_TEXT_EMBED_URL or ( + VISUAL_EMBED_URL.replace("/embed/image", "/embed/text") + if VISUAL_EMBED_URL and "/embed/image" in VISUAL_EMBED_URL + else None + ) + if not url: + raise RuntimeError("VISUAL_TEXT_EMBED_URL not configured") + resp = requests.post(url, json={"text": query}, timeout=VISUAL_EMBED_TIMEOUT) + resp.raise_for_status() + payload = resp.json() + embedding = payload.get("embedding") + if not isinstance(embedding, list) or not embedding: + raise RuntimeError(f"clip sidecar returned malformed payload: {payload!r}") + return embedding + + +async def persist_visual_chunk( + pool, + file_id: str, + page_number: int, + image_path: str, + embedding: List[float], + cmetadata: dict, +) -> None: + """Upsert one row into ``visual_chunks``. Unique on (file_id, page_number).""" + vec = _vector_literal(embedding) + async with pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO visual_chunks (file_id, page_number, image_path, embedding, cmetadata) + VALUES ($1, $2, $3, $4::vector, $5::jsonb) + ON CONFLICT (file_id, page_number) DO UPDATE + SET image_path = EXCLUDED.image_path, + embedding = EXCLUDED.embedding, + cmetadata = EXCLUDED.cmetadata + """, + file_id, + page_number, + image_path, + vec, + json.dumps(cmetadata), + ) + + +async def fetch_visual_chunks_for_pages( + pool, + file_id: str, + page_numbers: List[int], +) -> List[dict]: + """Lookup rows from ``visual_chunks`` by explicit (file_id, page_number). + + Used by the text-page-coupling path: the text pipeline has already + identified the relevant pages, and we attach the visuals of those pages + regardless of CLIP score. Missing pages are simply absent from the result. + """ + if not page_numbers: + return [] + async with pool.acquire() as conn: + rows = await conn.fetch( + """ + SELECT file_id, page_number, image_path + FROM visual_chunks + WHERE file_id = $1 + AND page_number = ANY($2::int[]) + ORDER BY page_number + """, + file_id, + list(page_numbers), + ) + return [ + { + "file_id": r["file_id"], + "page_number": r["page_number"], + "image_path": r["image_path"], + } + for r in rows + ] + + +async def similarity_search_visual( + pool, + query_embedding: List[float], + file_ids: List[str], + k: int, +) -> List[dict]: + """Cosine-similarity search over ``visual_chunks`` filtered by file_ids. + + Returns a list of dicts shaped as the /query response expects. + Results below VISUAL_SCORE_THRESHOLD are dropped. + """ + if not file_ids: + return [] + vec = _vector_literal(query_embedding) + async with pool.acquire() as conn: + rows = await conn.fetch( + """ + SELECT file_id, page_number, image_path, + 1 - (embedding <=> $1::vector) AS score + FROM visual_chunks + WHERE file_id = ANY($2::text[]) + ORDER BY embedding <=> $1::vector + LIMIT $3 + """, + vec, + file_ids, + k, + ) + return [ + { + "file_id": r["file_id"], + "page_number": r["page_number"], + "image_path": r["image_path"], + "score": float(r["score"]), + } + for r in rows + if float(r["score"]) >= VISUAL_SCORE_THRESHOLD + ] + + +async def maybe_embed_visuals( + file_path: str, + file_id: str, + file_ext: str, + user_id: str, + executor, +) -> int: + """Entry point — call after the text pipeline has completed. + + Returns the number of pages successfully embedded (0 if disabled or + soft-failed). Never raises. + """ + if not _visuals_enabled(): + return 0 + if file_ext.lower() != "pdf": + return 0 + + out_dir = Path(VISUAL_STORAGE_ROOT) / file_id + loop = asyncio.get_running_loop() + + # 1. Render pages + try: + page_paths = await loop.run_in_executor( + executor, render_pdf_pages, file_path, out_dir, VISUAL_PAGE_DPI + ) + except RuntimeError as exc: + logger.warning("visual ingest: pdftoppm step failed for %s: %s", file_id, exc) + return 0 + + if not page_paths: + logger.info("visual ingest: 0 pages rendered for %s", file_id) + return 0 + + # 2. Lazily grab a DB pool. If pgvector is unavailable (mongo deploy) + # we still keep the PNGs — but skip persistence. + try: + from app.services.database import PSQLDatabase + + pool = await PSQLDatabase.get_pool() + except Exception as exc: + logger.warning("visual ingest: cannot get DB pool, skipping persistence: %s", exc) + return 0 + + # 3. Embed + persist each page + persisted = 0 + for page_path in page_paths: + page_num = _page_number_from_path(page_path) + if page_num is None: + continue + try: + embedding = await loop.run_in_executor(executor, embed_image, page_path) + except Exception as exc: + logger.warning( + "visual ingest: embed failed for %s p%s: %s", file_id, page_num, exc + ) + continue + try: + await persist_visual_chunk( + pool, + file_id=file_id, + page_number=page_num, + image_path=str(page_path), + embedding=embedding, + cmetadata={"user_id": user_id, "source": os.path.basename(file_path)}, + ) + except Exception as exc: + logger.warning( + "visual ingest: persist failed for %s p%s: %s", file_id, page_num, exc + ) + continue + persisted += 1 + + logger.info("visual ingest: %d/%d pages embedded for %s", persisted, len(page_paths), file_id) + return persisted + + +def cleanup_visual_storage(file_id: str) -> None: + """Optional: delete the on-disk page PNGs for a file_id. Used by tests and rollback.""" + target = Path(VISUAL_STORAGE_ROOT) / file_id + if target.exists() and target.is_dir(): + shutil.rmtree(target, ignore_errors=True) diff --git a/requirements.lite.txt b/requirements.lite.txt index b3454c06..5f20afbc 100644 --- a/requirements.lite.txt +++ b/requirements.lite.txt @@ -35,3 +35,5 @@ chardet==5.2.0 langchain-ollama==1.0.1 tenacity>=9.0.0 msoffcrypto-tool>=6.0.0,<7 +requests>=2.31.0,<3 +pymupdf>=1.24.0,<2 diff --git a/requirements.txt b/requirements.txt index 18cfaf06..307ddeab 100644 --- a/requirements.txt +++ b/requirements.txt @@ -39,3 +39,5 @@ pydantic>=2.10.6,<3 chardet==5.2.0 tenacity>=9.0.0 msoffcrypto-tool>=6.0.0,<7 +requests>=2.31.0,<3 +pymupdf>=1.24.0,<2 diff --git a/test_requirements.txt b/test_requirements.txt index dcf35571..906cfe52 100644 --- a/test_requirements.txt +++ b/test_requirements.txt @@ -2,4 +2,6 @@ pytest==8.3.4 pytest-asyncio==0.26.0 pytest-postgresql==7.0.1 mongomock==4.3.0 -httpx==0.27.0 \ No newline at end of file +httpx==0.27.0 +responses==0.25.3 +requests>=2.31.0,<3 diff --git a/tests/test_main.py b/tests/test_main.py index a7fce89a..f7b7d406 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -173,6 +173,317 @@ def test_query_embeddings_by_file_id(auth_headers): assert doc["page_content"] == "Queried content" +def test_query_include_visual_wraps_response(auth_headers, monkeypatch): + """include_visual=True returns {chunks, visual_matches} and calls the visual helper.""" + + async def fake_visual(request, query, file_ids, text_documents=None): + assert file_ids == ["testid1"] + return [ + { + "file_id": "testid1", + "page_number": 2, + "image_path": "/var/rag-visual/testid1/page-2.png", + "score": 0.87, + } + ] + + monkeypatch.setattr( + document_routes, "_fetch_visual_matches_for_file_ids", fake_visual + ) + + data = { + "query": "Wie ist das Layout auf Seite 2?", + "file_id": "testid1", + "k": 4, + "entity_id": "testuser", + "include_visual": True, + } + response = client.post("/query", json=data, headers=auth_headers) + assert response.status_code == 200 + body = response.json() + assert isinstance(body, dict) + assert "chunks" in body and "visual_matches" in body + assert body["visual_matches"][0]["page_number"] == 2 + assert body["visual_matches"][0]["score"] == 0.87 + + +def test_query_include_visual_when_text_empty(auth_headers, monkeypatch): + """When the text index returns nothing but include_visual=True, + the response still wraps and ships the visual hits (no 404).""" + from app.services.vector_store.async_pg_vector import AsyncPgVector + + async def no_docs(self, embedding, k, filter=None, executor=None): + return [] + + monkeypatch.setattr( + AsyncPgVector, "asimilarity_search_with_score_by_vector", no_docs + ) + + async def fake_visual(request, query, file_ids, text_documents=None): + return [ + { + "file_id": "testid1", + "page_number": 1, + "image_path": "/x/p1.png", + "score": 0.42, + } + ] + + monkeypatch.setattr( + document_routes, "_fetch_visual_matches_for_file_ids", fake_visual + ) + + response = client.post( + "/query", + json={ + "query": "q", + "file_id": "testid1", + "k": 4, + "entity_id": "testuser", + "include_visual": True, + }, + headers=auth_headers, + ) + assert response.status_code == 200 + body = response.json() + assert body["chunks"] == [] + assert body["visual_matches"][0]["score"] == 0.42 + + +def _patch_visual_merge_deps( + monkeypatch, + *, + text_pages_for_file: dict, + clip_hits: list, + coupled_visuals: dict = None, + text_coupled: bool = True, + max_pages: int = 4, +): + """Helper: wire the dependencies of ``_fetch_visual_matches_for_file_ids`` + for the merge-logic tests below. + + ``text_pages_for_file``: {file_id: [page, …]} — what text-search returns + as (Document, score) with page_number metadata. + ``coupled_visuals``: {file_id: [{page_number, image_path}, …]} — what the + visual_chunks DB has for those pages. If None, derive from text_pages. + ``clip_hits``: list of CLIP-retrieved dicts (pages, scores). + """ + from app.services.vector_store.async_pg_vector import AsyncPgVector + from app.routes import document_routes + from app.utils import visual_embed + + if coupled_visuals is None: + coupled_visuals = { + fid: [ + {"file_id": fid, "page_number": p, "image_path": f"/x/{fid}/p-{p}.png"} + for p in pages + ] + for fid, pages in text_pages_for_file.items() + } + + async def fake_text_search(self, embedding, k, filter=None, executor=None): + wanted = None + if filter: + eq = filter.get("file_id", {}) + if isinstance(eq, dict): + wanted = eq.get("$eq") or eq.get("$in") + else: + wanted = eq + file_ids = [wanted] if isinstance(wanted, str) else (wanted or []) + out = [] + for fid in file_ids: + for page in text_pages_for_file.get(fid, []): + out.append( + ( + Document( + page_content=f"chunk from {fid} p{page}", + metadata={ + "file_id": fid, + "page_number": page, + "user_id": "testuser", + }, + ), + 0.9, + ) + ) + return out + + monkeypatch.setattr( + AsyncPgVector, "asimilarity_search_with_score_by_vector", fake_text_search + ) + monkeypatch.setattr(document_routes, "VISUAL_EMBED_URL", "http://clip.test/embed/image") + monkeypatch.setattr(document_routes, "VISUAL_TEXT_COUPLED", text_coupled) + monkeypatch.setattr(document_routes, "VISUAL_TEXT_COUPLED_MAX_PAGES", max_pages) + + class _FakePool: + pass + + class _DB: + @classmethod + async def get_pool(cls): + return _FakePool() + + import app.services.database as db_mod + + monkeypatch.setattr(db_mod, "PSQLDatabase", _DB) + + async def fake_fetch_coupled(pool, file_id, page_numbers): + rows = coupled_visuals.get(file_id, []) + wanted = set(page_numbers) + return [r for r in rows if r["page_number"] in wanted] + + monkeypatch.setattr(document_routes, "fetch_visual_chunks_for_pages", fake_fetch_coupled) + monkeypatch.setattr(visual_embed, "fetch_visual_chunks_for_pages", fake_fetch_coupled) + + def fake_embed_text(query): + return [0.1] * 768 + + monkeypatch.setattr(document_routes, "embed_text_query", fake_embed_text) + + async def fake_clip(pool, query_embedding, file_ids, k): + return [h for h in clip_hits if h["file_id"] in file_ids][:k] + + monkeypatch.setattr(document_routes, "similarity_search_visual", fake_clip) + + +def test_pages_by_file_from_text_docs_accepts_page_or_page_number(): + """PyPDFLoader writes ``page``, custom loaders write ``page_number``. + The helper must accept both (prefers page_number), preserve rank order + within a file, dedupe, and drop chunks without either key.""" + from app.routes.document_routes import _pages_by_file_from_text_docs + + docs = [ + (Document(page_content="a", metadata={"file_id": "f1", "page": 75}), 0.9), + (Document(page_content="b", metadata={"file_id": "f1", "page": 76}), 0.8), + (Document(page_content="c", metadata={"file_id": "f1", "page": 75}), 0.7), # dup + (Document(page_content="d", metadata={"file_id": "f2", "page_number": 3}), 0.6), + (Document(page_content="e", metadata={"file_id": "f2"}), 0.5), # no page → skip + (Document(page_content="f", metadata={"file_id": "f3", "page_number": "not-int"}), 0.4), + ] + out = _pages_by_file_from_text_docs(docs) + assert out == {"f1": [75, 76], "f2": [3]} + + +def test_query_with_text_coupling_returns_text_pages_first(auth_headers, monkeypatch): + """Text pipeline finds pages [5, 8, 12]. visual_chunks also has 50, 51 + (CLIP top hits). Expect visual_matches to contain {5, 8, 12} with + source='text_coupled' regardless of CLIP score, plus the CLIP hits that + don't collide, deduped.""" + _patch_visual_merge_deps( + monkeypatch, + text_pages_for_file={"testid1": [5, 8, 12]}, + coupled_visuals={ + "testid1": [ + {"file_id": "testid1", "page_number": p, "image_path": f"/x/p-{p}.png"} + for p in (5, 8, 12, 50, 51) + ] + }, + clip_hits=[ + {"file_id": "testid1", "page_number": 50, "image_path": "/x/p-50.png", "score": 0.42}, + {"file_id": "testid1", "page_number": 51, "image_path": "/x/p-51.png", "score": 0.40}, + # This one collides with a text-coupled page — must be deduped. + {"file_id": "testid1", "page_number": 5, "image_path": "/x/p-5.png", "score": 0.38}, + ], + ) + + response = client.post( + "/query", + json={ + "query": "Vietnam group photo", + "file_id": "testid1", + "k": 4, + "entity_id": "testuser", + "include_visual": True, + }, + headers=auth_headers, + ) + assert response.status_code == 200 + body = response.json() + matches = body["visual_matches"] + pages = [m["page_number"] for m in matches] + # Text-coupled first (in text-rank order), then CLIP, with no dups. + assert pages[:3] == [5, 8, 12] + assert set(pages) == {5, 8, 12, 50, 51} + assert all(m["source"] == "text_coupled" for m in matches if m["page_number"] in (5, 8, 12)) + assert all(m["source"] == "clip" for m in matches if m["page_number"] in (50, 51)) + assert all(m["score"] == 1.0 for m in matches if m["source"] == "text_coupled") + + +def test_query_with_text_coupling_disabled_falls_back_to_clip(auth_headers, monkeypatch): + """With VISUAL_TEXT_COUPLED=False, behaviour is Phase-3: only CLIP hits + are returned, no synthetic 1.0 scores, no 'text_coupled' source.""" + _patch_visual_merge_deps( + monkeypatch, + text_pages_for_file={"testid1": [5, 8, 12]}, + clip_hits=[ + {"file_id": "testid1", "page_number": 50, "image_path": "/x/p-50.png", "score": 0.42}, + {"file_id": "testid1", "page_number": 51, "image_path": "/x/p-51.png", "score": 0.40}, + ], + text_coupled=False, + ) + + response = client.post( + "/query", + json={ + "query": "anything", + "file_id": "testid1", + "k": 4, + "entity_id": "testuser", + "include_visual": True, + }, + headers=auth_headers, + ) + assert response.status_code == 200 + body = response.json() + pages = [m["page_number"] for m in body["visual_matches"]] + assert pages == [50, 51] + assert all(m["source"] == "clip" for m in body["visual_matches"]) + + +def test_query_with_text_coupling_caps_at_max_pages(auth_headers, monkeypatch): + """Text pipeline finds 10 pages, but VISUAL_TEXT_COUPLED_MAX_PAGES=4 + caps the text-coupled signal at 4. The remaining pages (ranks 5..10) + are NOT attached just because the text found them — but CLIP hits for + pages outside the capped set are still included, deduped.""" + text_pages = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19] + all_visuals = [ + {"file_id": "testid1", "page_number": p, "image_path": f"/x/p-{p}.png"} + for p in text_pages + ] + _patch_visual_merge_deps( + monkeypatch, + text_pages_for_file={"testid1": text_pages}, + coupled_visuals={"testid1": all_visuals}, + clip_hits=[ + # CLIP returns a page OUTSIDE the capped-4 text set — must survive. + {"file_id": "testid1", "page_number": 77, "image_path": "/x/p-77.png", "score": 0.31}, + # And a page INSIDE the capped set — must be deduped away (primary wins). + {"file_id": "testid1", "page_number": 10, "image_path": "/x/p-10.png", "score": 0.29}, + ], + max_pages=4, + ) + + response = client.post( + "/query", + json={ + "query": "q", + "file_id": "testid1", + "k": 4, + "entity_id": "testuser", + "include_visual": True, + }, + headers=auth_headers, + ) + assert response.status_code == 200 + body = response.json() + matches = body["visual_matches"] + text_coupled_pages = [m["page_number"] for m in matches if m["source"] == "text_coupled"] + clip_pages = [m["page_number"] for m in matches if m["source"] == "clip"] + assert text_coupled_pages == [10, 11, 12, 13] # capped at 4, rank order preserved + assert 77 in clip_pages + assert 10 not in clip_pages # deduped — text-coupled won + + def test_embed_local_file(tmp_path, auth_headers, monkeypatch): # Monkeypatch RAG_UPLOAD_DIR so the file is within the allowed directory. monkeypatch.setattr(document_routes, "RAG_UPLOAD_DIR", str(tmp_path)) diff --git a/tests/utils/test_pre_extraction_webhook.py b/tests/utils/test_pre_extraction_webhook.py new file mode 100644 index 00000000..1d890912 --- /dev/null +++ b/tests/utils/test_pre_extraction_webhook.py @@ -0,0 +1,158 @@ +"""Tests for the optional pre-extraction webhook in document_loader. + +The webhook is enabled only when PRE_EXTRACTION_WEBHOOK_URL is set. These tests +exercise the helper directly so we can cover the branching logic without +standing up FastAPI. +""" + +import importlib + +import pytest +import responses +from langchain_core.documents import Document + + +WEBHOOK_URL = "http://ocr-sidecar.test/extract" + + +@pytest.fixture +def enabled_webhook(monkeypatch): + """Reload config + document_loader so the module-level constants pick up + the new environment variables.""" + monkeypatch.setenv("PRE_EXTRACTION_WEBHOOK_URL", WEBHOOK_URL) + monkeypatch.setenv("PRE_EXTRACTION_WEBHOOK_MIN_CHARS", "100") + monkeypatch.setenv("PRE_EXTRACTION_WEBHOOK_TIMEOUT", "5") + + import app.config as config_module + import app.utils.document_loader as loader_module + + importlib.reload(config_module) + importlib.reload(loader_module) + yield loader_module + + # Reset after the test so other tests see the default (disabled) state. + monkeypatch.delenv("PRE_EXTRACTION_WEBHOOK_URL", raising=False) + importlib.reload(config_module) + importlib.reload(loader_module) + + +@pytest.fixture +def disabled_webhook(monkeypatch): + monkeypatch.delenv("PRE_EXTRACTION_WEBHOOK_URL", raising=False) + import app.config as config_module + import app.utils.document_loader as loader_module + + importlib.reload(config_module) + importlib.reload(loader_module) + yield loader_module + + +def _make_pdf(tmp_path, name="input.pdf"): + p = tmp_path / name + p.write_bytes(b"%PDF-1.4\n%fake\n") + return str(p) + + +def test_webhook_disabled_is_noop(disabled_webhook, tmp_path): + """When the feature flag is empty, the helper must return the input + unchanged without issuing any HTTP requests.""" + docs = [Document(page_content="", metadata={"source": "a.pdf"})] + file_path = _make_pdf(tmp_path) + out = disabled_webhook.maybe_enrich_with_webhook(file_path, docs) + assert out is docs + + +def test_above_threshold_skips_webhook(enabled_webhook, tmp_path): + """If average chars per page already exceeds the threshold we trust the + existing extraction and do not call the webhook.""" + docs = [ + Document(page_content="a" * 150, metadata={"source": "a.pdf"}), + Document(page_content="b" * 150, metadata={"source": "a.pdf"}), + ] + file_path = _make_pdf(tmp_path) + + # responses active but no registered endpoints — an unexpected call would + # raise ConnectionError, which is exactly what we want to assert. + with responses.RequestsMock() as rsps: + out = enabled_webhook.maybe_enrich_with_webhook(file_path, docs) + assert out == docs + assert len(rsps.calls) == 0 + + +def test_below_threshold_triggers_webhook_and_replaces_documents( + enabled_webhook, tmp_path +): + """Pages with effectively no text should be sent to the webhook and the + result should replace the documents, annotated with ocr metadata.""" + docs = [ + Document(page_content="", metadata={"source": "scanned.pdf"}), + Document(page_content=" ", metadata={"source": "scanned.pdf"}), + ] + file_path = _make_pdf(tmp_path, name="scanned.pdf") + + with responses.RequestsMock() as rsps: + rsps.add( + responses.POST, + WEBHOOK_URL, + json={ + "text": "OCRed contract content", + "provider": "azure-di", + "pages_processed": 2, + }, + status=200, + ) + out = enabled_webhook.maybe_enrich_with_webhook(file_path, docs) + + assert len(out) == 1 + assert out[0].page_content == "OCRed contract content" + assert out[0].metadata["ocr_used"] is True + assert out[0].metadata["ocr_provider"] == "azure-di" + assert out[0].metadata["source"] == "scanned.pdf" + + +def test_webhook_failure_falls_back_to_original(enabled_webhook, tmp_path): + """HTTP errors must never break ingest; we keep the original documents.""" + docs = [Document(page_content="", metadata={"source": "scanned.pdf"})] + file_path = _make_pdf(tmp_path, name="scanned.pdf") + + with responses.RequestsMock() as rsps: + rsps.add(responses.POST, WEBHOOK_URL, status=500) + out = enabled_webhook.maybe_enrich_with_webhook(file_path, docs) + + assert out == docs + + +def test_webhook_empty_text_falls_back(enabled_webhook, tmp_path): + """A 200 response with empty text means OCR had nothing to offer; keep + whatever PyPDF returned rather than destroy the extraction.""" + docs = [Document(page_content="", metadata={"source": "scanned.pdf"})] + file_path = _make_pdf(tmp_path, name="scanned.pdf") + + with responses.RequestsMock() as rsps: + rsps.add( + responses.POST, + WEBHOOK_URL, + json={"text": "", "provider": "azure-di"}, + status=200, + ) + out = enabled_webhook.maybe_enrich_with_webhook(file_path, docs) + + assert out == docs + + +def test_webhook_handles_empty_document_list(enabled_webhook, tmp_path): + """Empty list: avg_chars is 0 → webhook is called.""" + file_path = _make_pdf(tmp_path, name="scanned.pdf") + + with responses.RequestsMock() as rsps: + rsps.add( + responses.POST, + WEBHOOK_URL, + json={"text": "rescued text", "provider": "azure-di"}, + status=200, + ) + out = enabled_webhook.maybe_enrich_with_webhook(file_path, []) + + assert len(out) == 1 + assert out[0].page_content == "rescued text" + assert out[0].metadata["ocr_used"] is True diff --git a/tests/utils/test_visual_embed.py b/tests/utils/test_visual_embed.py new file mode 100644 index 00000000..7ac0a98c --- /dev/null +++ b/tests/utils/test_visual_embed.py @@ -0,0 +1,418 @@ +"""Tests for the optional multimodal-RAG visual ingest pipeline. + +The pipeline is feature-flagged via VISUAL_EMBED_URL. Tests reload the +config module with env vars set so module-level constants pick up. + +We never call real pdftoppm or real clip-embed-service here — both are +mocked. The pgvector pool is replaced by a FakePool that records calls. +""" + +from __future__ import annotations + +import asyncio +import importlib +import json +from pathlib import Path + +import pytest + + +VISUAL_URL = "http://clip.test/embed/image" +VISUAL_TEXT_URL = "http://clip.test/embed/text" + + +class FakeConn: + def __init__(self, parent: "FakePool"): + self.parent = parent + + async def execute(self, sql, *args): + self.parent.executed.append({"sql": sql, "args": args}) + return "INSERT 0 1" + + async def fetch(self, sql, *args): + self.parent.fetched.append({"sql": sql, "args": args}) + return self.parent.fetch_result + + +class _Ctx: + def __init__(self, conn): + self.conn = conn + + async def __aenter__(self): + return self.conn + + async def __aexit__(self, *exc): + return False + + +class FakePool: + def __init__(self): + self.executed: list = [] + self.fetched: list = [] + self.fetch_result: list = [] + + def acquire(self): + return _Ctx(FakeConn(self)) + + +@pytest.fixture +def enabled_visual(monkeypatch, tmp_path): + storage = tmp_path / "rag-visual" + monkeypatch.setenv("VISUAL_EMBED_URL", VISUAL_URL) + monkeypatch.setenv("VISUAL_TEXT_EMBED_URL", VISUAL_TEXT_URL) + monkeypatch.setenv("VISUAL_PAGE_DPI", "100") + monkeypatch.setenv("VISUAL_STORAGE_ROOT", str(storage)) + monkeypatch.setenv("VISUAL_SCORE_THRESHOLD", "0.25") + + import app.config as config_module + import app.utils.visual_embed as visual_module + + importlib.reload(config_module) + importlib.reload(visual_module) + yield visual_module + + monkeypatch.delenv("VISUAL_EMBED_URL", raising=False) + importlib.reload(config_module) + importlib.reload(visual_module) + + +@pytest.fixture +def disabled_visual(monkeypatch): + monkeypatch.delenv("VISUAL_EMBED_URL", raising=False) + import app.config as config_module + import app.utils.visual_embed as visual_module + + importlib.reload(config_module) + importlib.reload(visual_module) + yield visual_module + + +def _async(coro): + return asyncio.get_event_loop().run_until_complete(coro) + + +def test_disabled_feature_is_noop(disabled_visual, tmp_path): + pdf = tmp_path / "any.pdf" + pdf.write_bytes(b"%PDF-1.4\n") + + async def run(): + return await disabled_visual.maybe_embed_visuals( + file_path=str(pdf), + file_id="file-1", + file_ext="pdf", + user_id="u1", + executor=None, + ) + + result = asyncio.run(run()) + assert result == 0 + + +def test_non_pdf_is_noop(enabled_visual, tmp_path): + f = tmp_path / "notes.txt" + f.write_text("hi") + + async def run(): + return await enabled_visual.maybe_embed_visuals( + file_path=str(f), + file_id="file-2", + file_ext="txt", + user_id="u1", + executor=None, + ) + + assert asyncio.run(run()) == 0 + + +def test_vector_literal_format(enabled_visual): + assert enabled_visual._vector_literal([0.1, 0.2]) == "[0.100000,0.200000]" + + +def test_page_number_from_path(enabled_visual): + assert enabled_visual._page_number_from_path(Path("/x/page-3.png")) == 3 + assert enabled_visual._page_number_from_path(Path("/x/page-042.png")) == 42 + assert enabled_visual._page_number_from_path(Path("/x/bad.png")) is None + + +def test_render_pdf_pages_raises_on_fitz_failure( + enabled_visual, tmp_path, monkeypatch +): + """If PyMuPDF.open throws (corrupt PDF, missing file, etc.) the + helper must re-raise RuntimeError so the pipeline can soft-fail + cleanly.""" + import sys + import types + + fake_fitz = types.ModuleType("fitz") + + def _raise(*a, **kw): + raise FileNotFoundError("no such file") + + fake_fitz.open = _raise + fake_fitz.Matrix = lambda *a, **kw: None + monkeypatch.setitem(sys.modules, "fitz", fake_fitz) + + out = tmp_path / "pages" + with pytest.raises(RuntimeError, match="PyMuPDF"): + enabled_visual.render_pdf_pages("/non/existent.pdf", out, 100) + + +def test_embed_image_calls_sidecar(enabled_visual, tmp_path, monkeypatch): + png = tmp_path / "p.png" + png.write_bytes(b"\x89PNG\r\n\x1a\n") + + captured = {} + + class FakeResp: + def raise_for_status(self): + pass + + def json(self): + return {"embedding": [0.1] * 768, "dim": 768, "model": "nomic"} + + def fake_post(url, **kw): + captured["url"] = url + captured["timeout"] = kw.get("timeout") + captured["files"] = kw.get("files") + return FakeResp() + + import requests + + monkeypatch.setattr(requests, "post", fake_post) + + emb = enabled_visual.embed_image(png) + assert emb == [0.1] * 768 + assert captured["url"] == VISUAL_URL + assert captured["files"]["file"][0] == "p.png" + + +def test_embed_text_query_calls_text_endpoint(enabled_visual, monkeypatch): + captured = {} + + class FakeResp: + def raise_for_status(self): + pass + + def json(self): + return {"embedding": [0.2] * 768} + + def fake_post(url, **kw): + captured["url"] = url + captured["json"] = kw.get("json") + return FakeResp() + + import requests + + monkeypatch.setattr(requests, "post", fake_post) + + emb = enabled_visual.embed_text_query("layout seite 2") + assert emb == [0.2] * 768 + assert captured["url"] == VISUAL_TEXT_URL + assert captured["json"] == {"text": "layout seite 2"} + + +def test_persist_visual_chunk_builds_pg_insert(enabled_visual): + pool = FakePool() + asyncio.run( + enabled_visual.persist_visual_chunk( + pool=pool, + file_id="f1", + page_number=2, + image_path="/var/rag-visual/f1/page-2.png", + embedding=[0.3] * 768, + cmetadata={"user_id": "u1", "source": "x.pdf"}, + ) + ) + assert len(pool.executed) == 1 + sql = pool.executed[0]["sql"] + assert "INSERT INTO visual_chunks" in sql + assert "ON CONFLICT (file_id, page_number) DO UPDATE" in sql + args = pool.executed[0]["args"] + assert args[0] == "f1" + assert args[1] == 2 + assert args[2] == "/var/rag-visual/f1/page-2.png" + assert args[3].startswith("[0.3000") and args[3].endswith("]") + assert json.loads(args[4]) == {"user_id": "u1", "source": "x.pdf"} + + +def test_similarity_search_visual_drops_low_scores(enabled_visual): + pool = FakePool() + pool.fetch_result = [ + {"file_id": "f1", "page_number": 1, "image_path": "/x/p1.png", "score": 0.9}, + {"file_id": "f1", "page_number": 2, "image_path": "/x/p2.png", "score": 0.1}, + {"file_id": "f2", "page_number": 5, "image_path": "/x/p5.png", "score": 0.3}, + ] + results = asyncio.run( + enabled_visual.similarity_search_visual( + pool=pool, query_embedding=[0.1] * 768, file_ids=["f1", "f2"], k=10 + ) + ) + assert [r["page_number"] for r in results] == [1, 5] # 0.1 dropped by threshold + assert results[0]["score"] == pytest.approx(0.9) + + +def test_fetch_visual_chunks_for_pages_returns_only_requested_pages(enabled_visual): + """Direct lookup by (file_id, page_number). DB returns subset; helper + passes them through unchanged (order is SQL-imposed — ORDER BY page_number).""" + pool = FakePool() + pool.fetch_result = [ + {"file_id": "f1", "page_number": 5, "image_path": "/x/f1/p-05.png"}, + {"file_id": "f1", "page_number": 8, "image_path": "/x/f1/p-08.png"}, + {"file_id": "f1", "page_number": 12, "image_path": "/x/f1/p-12.png"}, + ] + results = asyncio.run( + enabled_visual.fetch_visual_chunks_for_pages( + pool=pool, file_id="f1", page_numbers=[5, 8, 12] + ) + ) + assert [r["page_number"] for r in results] == [5, 8, 12] + assert all(r["file_id"] == "f1" for r in results) + assert "score" not in results[0] # helper does not synthesize a score + + # Verify the SQL params: file_id + int[] of page numbers. + assert len(pool.fetched) == 1 + args = pool.fetched[0]["args"] + assert args[0] == "f1" + assert list(args[1]) == [5, 8, 12] + + +def test_fetch_visual_chunks_for_pages_returns_empty_for_missing_pages(enabled_visual): + """If the DB has no rows for the requested pages (e.g. visual ingest + soft-failed for that file), the helper returns an empty list — the + caller treats that as 'no visuals for these pages' and moves on.""" + pool = FakePool() + pool.fetch_result = [] + results = asyncio.run( + enabled_visual.fetch_visual_chunks_for_pages( + pool=pool, file_id="f-missing", page_numbers=[1, 2, 3] + ) + ) + assert results == [] + + # Empty input short-circuits — no DB call at all. + pool2 = FakePool() + results2 = asyncio.run( + enabled_visual.fetch_visual_chunks_for_pages( + pool=pool2, file_id="f1", page_numbers=[] + ) + ) + assert results2 == [] + assert pool2.fetched == [] + + +def test_similarity_search_visual_empty_files_returns_empty(enabled_visual): + pool = FakePool() + results = asyncio.run( + enabled_visual.similarity_search_visual( + pool=pool, query_embedding=[0.1] * 768, file_ids=[], k=10 + ) + ) + assert results == [] + assert pool.fetched == [] + + +def test_full_pipeline_happy_path(enabled_visual, tmp_path, monkeypatch): + """End-to-end flow (pdftoppm + HTTP + DB are all mocked).""" + pdf = tmp_path / "flyer.pdf" + pdf.write_bytes(b"%PDF-1.4\n") + + import app.config as cfg + + out_dir = Path(cfg.VISUAL_STORAGE_ROOT) / "file-7" + + # Mock pdftoppm: instead of rendering, drop two fake PNGs. + def fake_render(pdf_path, out, dpi): + out.mkdir(parents=True, exist_ok=True) + p1 = out / "page-1.png" + p2 = out / "page-2.png" + p1.write_bytes(b"\x89PNG1") + p2.write_bytes(b"\x89PNG2") + return [p1, p2] + + monkeypatch.setattr(enabled_visual, "render_pdf_pages", fake_render) + + # Mock embed_image + def fake_embed(image_path): + return [0.5] * 768 + + monkeypatch.setattr(enabled_visual, "embed_image", fake_embed) + + # Mock DB pool + pool = FakePool() + + class _DB: + @classmethod + async def get_pool(cls): + return pool + + # Inject fake PSQLDatabase + import app.services.database as db_mod + + monkeypatch.setattr(db_mod, "PSQLDatabase", _DB) + + async def run(): + return await enabled_visual.maybe_embed_visuals( + file_path=str(pdf), + file_id="file-7", + file_ext="pdf", + user_id="u1", + executor=None, + ) + + persisted = asyncio.run(run()) + assert persisted == 2 + assert len(pool.executed) == 2 + assert pool.executed[0]["args"][0] == "file-7" + assert pool.executed[0]["args"][1] in (1, 2) + + +def test_full_pipeline_continues_when_single_page_embed_fails( + enabled_visual, tmp_path, monkeypatch +): + pdf = tmp_path / "flyer.pdf" + pdf.write_bytes(b"%PDF-1.4\n") + + def fake_render(pdf_path, out, dpi): + out.mkdir(parents=True, exist_ok=True) + p1 = out / "page-1.png" + p2 = out / "page-2.png" + p1.write_bytes(b"\x89PNG1") + p2.write_bytes(b"\x89PNG2") + return [p1, p2] + + monkeypatch.setattr(enabled_visual, "render_pdf_pages", fake_render) + + calls = {"n": 0} + + def flaky_embed(image_path): + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("sidecar down for page 1") + return [0.5] * 768 + + monkeypatch.setattr(enabled_visual, "embed_image", flaky_embed) + + pool = FakePool() + + class _DB: + @classmethod + async def get_pool(cls): + return pool + + import app.services.database as db_mod + + monkeypatch.setattr(db_mod, "PSQLDatabase", _DB) + + async def run(): + return await enabled_visual.maybe_embed_visuals( + file_path=str(pdf), + file_id="file-8", + file_ext="pdf", + user_id="u1", + executor=None, + ) + + persisted = asyncio.run(run()) + # Page 1 failed, page 2 succeeded + assert persisted == 1 + assert len(pool.executed) == 1