diff --git a/.env.demo b/.env.demo index e06ab45b..3158ab8c 100644 --- a/.env.demo +++ b/.env.demo @@ -346,4 +346,10 @@ ENABLE_DEDUPLICATION=True SHOW_DEDUPLICATION_STEP=True # Minimum cosine similarity score (0–1) for two documents to be flagged as # near-duplicates. 0.85 means 85 % semantic overlap. Lower = more matches. -NEAR_DUPLICATE_THRESHOLD=0.85 +NEAR_DUPLICATE_THRESHOLD=0.85 +# Model used to generate text embeddings for document similarity. +# Must be supported by your OpenAI-compatible API endpoint. +EMBEDDING_MODEL=text-embedding-3-small +# Maximum tokens to send to the embedding model. Set below the model's +# context window (e.g. 8000 for an 8192-token model). +EMBEDDING_MAX_TOKENS=8000 diff --git a/app/api/similarity.py b/app/api/similarity.py index 9c14ce1b..e7b5f844 100644 --- a/app/api/similarity.py +++ b/app/api/similarity.py @@ -1,9 +1,11 @@ """Document similarity API endpoints. -Provides an endpoint to find documents similar to a given file based on -text embeddings and cosine similarity scoring. +Provides endpoints to find documents similar to a given file based on +text embeddings and cosine similarity scoring, plus debug/diagnostic +endpoints for inspecting and triggering embedding computation. """ +import json import logging from typing import Annotated @@ -11,6 +13,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from sqlalchemy.orm import Session from app.auth import require_login +from app.config import settings from app.database import get_db from app.models import FileRecord @@ -79,6 +82,19 @@ def get_similar_documents( "message": "No OCR text available for similarity comparison", } + # Check whether an embedding has been computed yet + if not file_record.embedding: + return { + "file_id": file_id, + "similar_documents": [], + "count": 0, + "message": ( + "Embedding not yet computed for this file. " + "It will be generated automatically during processing or via the backfill task. " + "You can also trigger it manually with POST /api/files/{file_id}/compute-embedding." + ), + } + try: from app.utils.similarity import find_similar_documents @@ -95,3 +111,363 @@ def get_similar_documents( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to compute document similarity", ) + + +# --------------------------------------------------------------------------- +# Debug / diagnostic endpoints +# --------------------------------------------------------------------------- + + +@router.get("/files/{file_id}/embedding-status") +@require_login +def get_embedding_status( + request: Request, + file_id: int, + db: DbSession, +): + """Return the embedding status for a single file. + + Useful for debugging whether the embedding has been computed + and cached for a given document. + + Response: + ```json + { + "file_id": 42, + "has_embedding": true, + "embedding_dimensions": 1536, + "has_ocr_text": true, + "ocr_text_length": 4200, + "embedding_model": "text-embedding-3-small" + } + ``` + """ + file_record = db.query(FileRecord).filter(FileRecord.id == file_id).first() + if not file_record: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="File not found") + + has_embedding = False + embedding_dimensions = None + if file_record.embedding: + try: + parsed = json.loads(file_record.embedding) + has_embedding = True + embedding_dimensions = len(parsed) + except (json.JSONDecodeError, TypeError): + pass + + has_ocr_text = bool(file_record.ocr_text and file_record.ocr_text.strip()) + + return { + "file_id": file_id, + "has_embedding": has_embedding, + "embedding_dimensions": embedding_dimensions, + "has_ocr_text": has_ocr_text, + "ocr_text_length": len(file_record.ocr_text) if file_record.ocr_text else 0, + "embedding_model": settings.embedding_model, + } + + +@router.post("/files/{file_id}/compute-embedding") +@require_login +def trigger_compute_embedding( + request: Request, + file_id: int, + db: DbSession, +): + """Trigger embedding computation for a single file. + + If the file already has a cached embedding it will be recomputed. + The computation happens synchronously so the caller receives the + result immediately. + + Response: + ```json + { + "file_id": 42, + "status": "success", + "embedding_dimensions": 1536 + } + ``` + """ + file_record = db.query(FileRecord).filter(FileRecord.id == file_id).first() + if not file_record: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="File not found") + + if not file_record.ocr_text or not file_record.ocr_text.strip(): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="File has no OCR text — cannot generate embedding", + ) + + try: + from app.utils.similarity import generate_embedding + + # Clear cached embedding to force recomputation + file_record.embedding = None + db.flush() + + embedding = generate_embedding(file_record.ocr_text) + file_record.embedding = json.dumps(embedding) + db.commit() + + return { + "file_id": file_id, + "status": "success", + "embedding_dimensions": len(embedding), + } + except Exception as e: + db.rollback() + logger.error(f"Failed to compute embedding for file {file_id}: {e}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Embedding computation failed: {e}", + ) + + +@router.get("/diagnostic/embeddings") +@require_login +def get_embeddings_overview( + request: Request, + db: DbSession, +): + """Return an overview of embedding status across all files. + + Provides aggregate counts as well as a per-file breakdown so an + administrator can quickly identify documents that are missing + embeddings. + + Response: + ```json + { + "total_files": 120, + "files_with_ocr_text": 95, + "files_with_embedding": 42, + "files_missing_embedding": 53, + "embedding_model": "text-embedding-3-small", + "files": [ + { + "file_id": 1, + "original_filename": "invoice.pdf", + "has_ocr_text": true, + "has_embedding": true, + "embedding_dimensions": 1536 + } + ] + } + ``` + """ + # Use column-only query to avoid loading full ORM objects into memory + all_files = ( + db.query( + FileRecord.id, + FileRecord.original_filename, + FileRecord.ocr_text, + FileRecord.embedding, + ) + .order_by(FileRecord.id.desc()) + .all() + ) + + files_info = [] + total_with_ocr = 0 + total_with_embedding = 0 + + for f in all_files: + has_ocr = bool(f.ocr_text and f.ocr_text.strip()) + has_emb = False + emb_dims = None + + if f.embedding: + try: + parsed = json.loads(f.embedding) + has_emb = True + emb_dims = len(parsed) + except (json.JSONDecodeError, TypeError): + pass + + if has_ocr: + total_with_ocr += 1 + if has_emb: + total_with_embedding += 1 + + files_info.append( + { + "file_id": f.id, + "original_filename": f.original_filename, + "has_ocr_text": has_ocr, + "has_embedding": has_emb, + "embedding_dimensions": emb_dims, + } + ) + + return { + "total_files": len(all_files), + "files_with_ocr_text": total_with_ocr, + "files_with_embedding": total_with_embedding, + "files_missing_embedding": total_with_ocr - total_with_embedding, + "embedding_model": settings.embedding_model, + "files": files_info, + } + + +@router.post("/diagnostic/compute-all-embeddings") +@require_login +def trigger_compute_all_embeddings( + request: Request, + db: DbSession, +): + """Queue embedding computation for all files that have OCR text but no embedding. + + Each file is processed as a separate Celery task so the endpoint + returns immediately. + + Response: + ```json + { + "status": "queued", + "files_queued": 53 + } + ``` + """ + candidates = ( + db.query(FileRecord) + .filter( + FileRecord.ocr_text.isnot(None), + FileRecord.ocr_text != "", + (FileRecord.embedding.is_(None)) | (FileRecord.embedding == ""), + ) + .all() + ) + + queued = 0 + for f in candidates: + try: + from app.tasks.compute_embedding import compute_document_embedding + + compute_document_embedding.delay(f.id) + queued += 1 + except Exception as e: + logger.warning(f"Could not queue embedding for file {f.id}: {e}") + + return { + "status": "queued", + "files_queued": queued, + } + + +@router.get("/similarity/pairs") +@require_login +def get_similarity_pairs( + request: Request, + db: DbSession, + threshold: float = Query(0.7, ge=0.0, le=1.0, description="Minimum similarity score for a pair"), + limit: int = Query(50, ge=1, le=200, description="Maximum number of pairs to return"), + page: int = Query(1, ge=1, description="Page number"), +): + """Return pairs of documents with high similarity across the entire corpus. + + Unlike the per-file ``/files/{id}/similar`` endpoint, this scans every + document that has a pre-computed embedding and returns **all** pairs + whose cosine similarity exceeds ``threshold``, sorted by descending + score. + + To keep memory bounded the query loads only the columns needed for + scoring and streams results in chunks. + + Response: + ```json + { + "pairs": [ + { + "file_a": {"file_id": 1, "original_filename": "invoice_jan.pdf", ...}, + "file_b": {"file_id": 5, "original_filename": "invoice_feb.pdf", ...}, + "similarity_score": 0.94 + } + ], + "total_pairs": 12, + "threshold": 0.7, + "page": 1, + "pages": 1, + "embedding_coverage": {"total_files": 120, "files_with_embedding": 95} + } + ``` + """ + from app.utils.similarity import cosine_similarity + + # Load all files that have embeddings (columns only for efficiency) + rows = ( + db.query( + FileRecord.id, + FileRecord.original_filename, + FileRecord.document_title, + FileRecord.mime_type, + FileRecord.created_at, + FileRecord.embedding, + ) + .filter( + FileRecord.embedding.isnot(None), + FileRecord.embedding != "", + ) + .order_by(FileRecord.id) + .all() + ) + + # Parse embeddings upfront + parsed: list[tuple] = [] + for row in rows: + try: + vec = json.loads(row.embedding) + parsed.append((row, vec)) + except (json.JSONDecodeError, TypeError): + continue + + # Pairwise comparison (triangle: i < j avoids duplicating A↔B / B↔A) + all_pairs: list[dict] = [] + for i in range(len(parsed)): + row_a, vec_a = parsed[i] + for j in range(i + 1, len(parsed)): + row_b, vec_b = parsed[j] + score = cosine_similarity(vec_a, vec_b) + if score >= threshold: + all_pairs.append( + { + "file_a": _row_to_dict(row_a), + "file_b": _row_to_dict(row_b), + "similarity_score": round(score, 4), + } + ) + + # Sort by score descending + all_pairs.sort(key=lambda p: p["similarity_score"], reverse=True) + + total_pairs = len(all_pairs) + total_pages = max(1, (total_pairs + limit - 1) // limit) + offset = (page - 1) * limit + page_pairs = all_pairs[offset : offset + limit] + + total_files = db.query(FileRecord).count() + + return { + "pairs": page_pairs, + "total_pairs": total_pairs, + "threshold": threshold, + "page": page, + "pages": total_pages, + "per_page": limit, + "embedding_coverage": { + "total_files": total_files, + "files_with_embedding": len(parsed), + }, + } + + +def _row_to_dict(row) -> dict: + """Serialise a column-only query row to a dict for JSON responses.""" + return { + "file_id": row.id, + "original_filename": row.original_filename, + "document_title": row.document_title, + "mime_type": row.mime_type, + "created_at": row.created_at.isoformat() if row.created_at else None, + } diff --git a/app/celery_worker.py b/app/celery_worker.py index d6bd83d4..d04f5916 100644 --- a/app/celery_worker.py +++ b/app/celery_worker.py @@ -9,6 +9,7 @@ from app import tasks # noqa: F401 - Imports app/tasks.py so Celery can registe from app.celery_app import celery from app.config import settings from app.tasks.check_credentials import check_credentials +from app.tasks.compute_embedding import backfill_missing_embeddings, compute_document_embedding # noqa: F401 from app.tasks.convert_to_pdf import convert_to_pdf # noqa: F401 from app.tasks.embed_metadata_into_pdf import embed_metadata_into_pdf # noqa: F401 from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt # noqa: F401 @@ -91,6 +92,13 @@ celery.conf.beat_schedule = { "schedule": crontab(minute="*/1"), # Every minute "options": {"expires": 55}, # Must complete within 55 seconds }, + # Backfill embeddings for files that were processed before the + # embedding pipeline was enabled, or where the embedding task failed. + "backfill-missing-embeddings": { + "task": "backfill_missing_embeddings", + "schedule": crontab(minute="*/5"), # Every 5 minutes + "options": {"expires": 240}, # 4 minutes expiry + }, } # Remove None entries from beat_schedule diff --git a/app/config.py b/app/config.py index 5a4c2fd1..91f2237d 100644 --- a/app/config.py +++ b/app/config.py @@ -323,6 +323,29 @@ class Settings(BaseSettings): "them near-duplicates. Higher values require closer content matches. Default: 0.85." ), ) + embedding_model: str = Field( + default="text-embedding-3-small", + description=( + "Model name used for generating text embeddings via the OpenAI-compatible API. " + "Embeddings drive the document similarity feature. Default: text-embedding-3-small." + ), + ) + embedding_max_tokens: int = Field( + default=8000, + description=( + "Maximum number of tokens to send to the embedding model. " + "Text is truncated to approximately this many tokens (using a " + "conservative 3-chars-per-token estimate) before calling the API. " + "Set this below the model's context window (e.g. 8000 for an 8192-token model)." + ), + ) + embedding_backfill_batch_size: int = Field( + default=50, + description=( + "Maximum number of files to queue for embedding computation per " + "backfill run. Keeps the worker and embedding API load bounded." + ), + ) # Text quality check - AI-based assessment of embedded PDF text enable_text_quality_check: bool = Field( diff --git a/app/tasks/compute_embedding.py b/app/tasks/compute_embedding.py new file mode 100644 index 00000000..50fc3cd6 --- /dev/null +++ b/app/tasks/compute_embedding.py @@ -0,0 +1,174 @@ +"""Celery task for pre-computing document text embeddings. + +Runs after document processing to ensure embeddings are available for +the similarity feature without requiring a user to trigger them on first +access. +""" + +import logging +from datetime import datetime, timezone + +from app.celery_app import celery +from app.config import settings +from app.database import SessionLocal +from app.models import FileRecord +from app.tasks.retry_config import BaseTaskWithRetry +from app.utils import log_task_progress +from app.utils.step_manager import update_step_status + +logger = logging.getLogger(__name__) + + +@celery.task(base=BaseTaskWithRetry, bind=True, name="compute_document_embedding") +def compute_document_embedding(self, file_id: int) -> dict: + """Compute and cache the text embedding for a single document. + + Skips silently when the file has no OCR text or already has a cached + embedding. The result is stored in ``FileRecord.embedding`` for + subsequent similarity queries. + + Args: + file_id: Primary key of the :class:`~app.models.FileRecord`. + + Returns: + A dict with ``status`` (``"success"`` / ``"skipped"`` / ``"error"``) + and optional ``detail`` message. + """ + task_id = self.request.id + logger.info("[%s] Computing embedding for file %s", task_id, file_id) + log_task_progress( + task_id, + "compute_embedding", + "in_progress", + f"Computing text embedding for file {file_id}", + file_id=file_id, + ) + + with SessionLocal() as db: + file_record = db.query(FileRecord).filter(FileRecord.id == file_id).first() + if not file_record: + logger.warning("[%s] File %s not found, skipping embedding", task_id, file_id) + return {"status": "skipped", "detail": "File not found"} + + now = datetime.now(timezone.utc) + update_step_status(db, file_id, "compute_embedding", "in_progress", started_at=now) + + # Already has a cached embedding – nothing to do + if file_record.embedding: + logger.info("[%s] File %s already has a cached embedding", task_id, file_id) + log_task_progress( + task_id, + "compute_embedding", + "success", + "Embedding already cached", + file_id=file_id, + ) + update_step_status(db, file_id, "compute_embedding", "success", completed_at=now) + return {"status": "skipped", "detail": "Embedding already cached"} + + if not file_record.ocr_text or not file_record.ocr_text.strip(): + logger.info("[%s] File %s has no OCR text, skipping embedding", task_id, file_id) + log_task_progress( + task_id, + "compute_embedding", + "skipped", + "No OCR text available", + file_id=file_id, + ) + update_step_status(db, file_id, "compute_embedding", "skipped", completed_at=now) + return {"status": "skipped", "detail": "No OCR text available"} + + try: + from app.utils.similarity import compute_and_store_embedding + + embedding = compute_and_store_embedding(db, file_record) + completed = datetime.now(timezone.utc) + if embedding: + log_task_progress( + task_id, + "compute_embedding", + "success", + f"Embedding computed ({len(embedding)} dimensions)", + file_id=file_id, + ) + update_step_status(db, file_id, "compute_embedding", "success", completed_at=completed) + return { + "status": "success", + "detail": f"Embedding computed ({len(embedding)} dimensions)", + } + else: + log_task_progress( + task_id, + "compute_embedding", + "failure", + "Embedding computation returned None", + file_id=file_id, + ) + update_step_status( + db, + file_id, + "compute_embedding", + "failure", + error_message="Embedding computation returned None", + completed_at=completed, + ) + return {"status": "error", "detail": "Embedding computation returned None"} + except Exception as exc: + logger.exception("[%s] Embedding computation failed for file %s: %s", task_id, file_id, exc) + log_task_progress( + task_id, + "compute_embedding", + "failure", + f"Exception: {exc}", + file_id=file_id, + ) + update_step_status( + db, + file_id, + "compute_embedding", + "failure", + error_message=str(exc), + completed_at=datetime.now(timezone.utc), + ) + return {"status": "error", "detail": str(exc)} + + +@celery.task(bind=True, name="backfill_missing_embeddings") +def backfill_missing_embeddings(self) -> dict: + """Periodic task that computes embeddings for documents that lack them. + + Iterates over all ``FileRecord`` rows that have OCR text but no + cached embedding and queues a :func:`compute_document_embedding` + task for each one. A configurable ``batch_size`` caps the number + of tasks queued per run to avoid overwhelming the worker or the + embedding API. + + Returns: + A dict with the number of tasks ``queued``. + """ + batch_size = settings.embedding_backfill_batch_size + task_id = self.request.id + logger.info("[%s] Backfill: scanning for files missing embeddings (batch_size=%d)", task_id, batch_size) + + with SessionLocal() as db: + candidates = ( + db.query(FileRecord.id) + .filter( + FileRecord.ocr_text.isnot(None), + FileRecord.ocr_text != "", + (FileRecord.embedding.is_(None)) | (FileRecord.embedding == ""), + ) + .limit(batch_size) + .all() + ) + + queued = 0 + for (file_id,) in candidates: + try: + compute_document_embedding.delay(file_id) + queued += 1 + except Exception as exc: + logger.warning("[%s] Could not queue embedding for file %s: %s", task_id, file_id, exc) + + logger.info("[%s] Backfill: queued %d embedding tasks", task_id, queued) + return {"queued": queued} diff --git a/app/tasks/finalize_document_storage.py b/app/tasks/finalize_document_storage.py index 1876e73d..228b9b09 100644 --- a/app/tasks/finalize_document_storage.py +++ b/app/tasks/finalize_document_storage.py @@ -76,6 +76,16 @@ def finalize_document_storage(self, original_file: str, processed_file: str, met # We pass 'True' (delete_after) and 'file_id' as per Main branch requirements send_to_all_destinations.delay(processed_file, True, file_id) + # 3a. Queue embedding computation so similarity scores are ready for queries + if file_id is not None: + try: + from app.tasks.compute_embedding import compute_document_embedding + + compute_document_embedding.delay(file_id) + logger.info(f"[{task_id}] Queued embedding computation for file {file_id}") + except Exception as e: + logger.warning(f"[{task_id}] Could not queue embedding task: {e}") + # 4. Send Notification (From Copilot) # Note: This notification is sent after processing is complete but while uploads # are being queued. diff --git a/app/utils/similarity.py b/app/utils/similarity.py index 617f1976..590084e1 100644 --- a/app/utils/similarity.py +++ b/app/utils/similarity.py @@ -38,13 +38,15 @@ def _get_embedding_client() -> Any: ) -def generate_embedding(text: str, model: str = "text-embedding-3-small") -> list[float]: +def generate_embedding(text: str, model: str | None = None) -> list[float]: """Generate a text embedding vector using the OpenAI-compatible API. Args: - text: The input text to embed. Truncated to ~8000 tokens worth of - characters to stay within model limits. - model: The embedding model to use. Defaults to ``text-embedding-3-small``. + text: The input text to embed. Truncated to stay within the + model's context window based on ``settings.embedding_max_tokens`` + (default 8 000 tokens ≈ 24 000 characters). + model: The embedding model to use. When ``None`` (the default), the + value of ``settings.embedding_model`` is used. Returns: A list of floats representing the embedding vector. @@ -53,12 +55,25 @@ def generate_embedding(text: str, model: str = "text-embedding-3-small") -> list RuntimeError: If the OpenAI client cannot be created. Exception: If the API call fails. """ - # Truncate very long texts to stay within token limits (~4 chars per token) - max_chars = 30000 + if model is None: + model = settings.embedding_model + + # Truncate to stay within the model's context window. + # Conservative 3 chars/token estimate (actual ratio varies by language; + # English averages ~4 chars/token but 3 gives a safety margin). + max_chars = settings.embedding_max_tokens * 3 if len(text) > max_chars: + logger.debug( + "Truncating text from %d to %d chars (~%d tokens) for model %s", + len(text), + max_chars, + settings.embedding_max_tokens, + model, + ) text = text[:max_chars] client = _get_embedding_client() + logger.debug("Generating embedding for %d chars using model=%s", len(text), model) response = client.embeddings.create(input=text, model=model) return response.data[0].embedding @@ -89,8 +104,41 @@ def cosine_similarity(vec_a: list[float], vec_b: list[float]) -> float: return max(0.0, min(1.0, similarity)) -def _get_or_compute_embedding(db: Session, file_record: Any) -> list[float] | None: - """Retrieve a cached embedding or compute and store a new one. +def _get_cached_embedding(file_record: Any) -> list[float] | None: + """Return the cached embedding for a file record, or ``None``. + + This is a **read-only** helper — it never triggers an API call. Use + :func:`compute_and_store_embedding` when you need to generate a new + embedding. + + Args: + file_record: A ``FileRecord`` instance (or any object with ``id`` + and ``embedding`` attributes). + + Returns: + The parsed embedding vector, or ``None`` if no valid cached + embedding exists. + """ + raw = file_record.embedding if hasattr(file_record, "embedding") else None + if not raw: + return None + try: + cached = json.loads(raw) + logger.debug("Using cached embedding for file %s (%d dimensions)", file_record.id, len(cached)) + return cached + except (json.JSONDecodeError, TypeError): + logger.warning("Invalid cached embedding for file %s", file_record.id) + return None + + +def compute_and_store_embedding(db: Session, file_record: Any) -> list[float] | None: + """Generate an embedding for a file and persist it in the database. + + Called during document ingestion (Celery task) or via the manual + ``POST /api/files/{id}/compute-embedding`` debug endpoint. The + similarity query path (:func:`find_similar_documents`) intentionally + does **not** call this — it only reads pre-computed embeddings so + that it returns instantly without blocking on external API calls. Args: db: Active database session. @@ -100,40 +148,54 @@ def _get_or_compute_embedding(db: Session, file_record: Any) -> list[float] | No The embedding vector, or ``None`` if the document has no OCR text or embedding generation fails. """ - # Return cached embedding if available + # Return cached embedding if already present if file_record.embedding: try: - return json.loads(file_record.embedding) + cached = json.loads(file_record.embedding) + logger.debug("Embedding already cached for file %s (%d dims)", file_record.id, len(cached)) + return cached except (json.JSONDecodeError, TypeError): - logger.warning(f"Invalid cached embedding for file {file_record.id}, recomputing") + logger.warning("Invalid cached embedding for file %s, recomputing", file_record.id) # Need OCR text to generate an embedding if not file_record.ocr_text or not file_record.ocr_text.strip(): + logger.debug("No OCR text for file %s, cannot generate embedding", file_record.id) return None try: + logger.info("Computing embedding for file %s (%d chars of OCR text)", file_record.id, len(file_record.ocr_text)) embedding = generate_embedding(file_record.ocr_text) - # Cache the embedding in the database + # Persist in the database file_record.embedding = json.dumps(embedding) db.commit() + logger.info("Embedding computed and cached for file %s (%d dimensions)", file_record.id, len(embedding)) return embedding except Exception as e: db.rollback() - logger.error(f"Failed to generate embedding for file {file_record.id}: {e}") + logger.error("Failed to generate embedding for file %s: %s", file_record.id, e) return None +# Keep the legacy alias so that existing callers (e.g. tests) keep working. +_get_or_compute_embedding = compute_and_store_embedding + + def find_similar_documents( db: Session, file_id: int, limit: int = 5, threshold: float = 0.3, ) -> list[dict[str, Any]]: - """Find documents similar to the given file. + """Find documents similar to the given file using **pre-computed** embeddings. - Computes cosine similarity between the target document's embedding and - all other documents that have OCR text. Results are sorted by descending - similarity score. + Only documents whose embeddings were already generated (during + ingestion or via the debug endpoint) are considered. No external API + calls are made — the function reads cached vectors from the database + and computes cosine similarity in-process. + + To keep memory usage bounded for large corpora (100 k+ documents) the + candidate query fetches only the columns needed for scoring and + iterates in chunks via ``yield_per``. Args: db: Active database session. @@ -152,43 +214,54 @@ def find_similar_documents( """ from app.models import FileRecord - # Get the target document + # Get the target document's cached embedding (read-only, no API call) target = db.query(FileRecord).filter(FileRecord.id == file_id).first() if not target: return [] - # Get the target embedding - target_embedding = _get_or_compute_embedding(db, target) + target_embedding = _get_cached_embedding(target) if not target_embedding: + logger.info("No cached embedding for target file %s — skipping similarity search", file_id) return [] - # Get candidate documents (those with OCR text, excluding the target) + # Query only candidates that already have a pre-computed embedding. + # Fetch only the columns needed for scoring to minimise memory use. + # yield_per streams rows in chunks so we never materialise all 100k+ + # records at once. candidates = ( - db.query(FileRecord) + db.query( + FileRecord.id, + FileRecord.original_filename, + FileRecord.document_title, + FileRecord.mime_type, + FileRecord.created_at, + FileRecord.embedding, + ) .filter( FileRecord.id != file_id, - FileRecord.ocr_text.isnot(None), - FileRecord.ocr_text != "", + FileRecord.embedding.isnot(None), + FileRecord.embedding != "", ) - .all() + .yield_per(500) ) - results = [] - for candidate in candidates: - candidate_embedding = _get_or_compute_embedding(db, candidate) - if not candidate_embedding: + results: list[dict[str, Any]] = [] + for row in candidates: + try: + candidate_embedding: list[float] = json.loads(row.embedding) + except (json.JSONDecodeError, TypeError): continue score = cosine_similarity(target_embedding, candidate_embedding) if score >= threshold: results.append( { - "file_id": candidate.id, - "original_filename": candidate.original_filename, - "document_title": candidate.document_title, + "file_id": row.id, + "original_filename": row.original_filename, + "document_title": row.document_title, "similarity_score": round(score, 4), - "mime_type": candidate.mime_type, - "created_at": candidate.created_at.isoformat() if candidate.created_at else None, + "mime_type": row.mime_type, + "created_at": row.created_at.isoformat() if row.created_at else None, } ) diff --git a/app/utils/step_manager.py b/app/utils/step_manager.py index 9e4449a3..ecbf2482 100644 --- a/app/utils/step_manager.py +++ b/app/utils/step_manager.py @@ -23,6 +23,7 @@ BASE_MAIN_PROCESSING_STEPS = [ "embed_metadata_into_pdf", "finalize_document_storage", "send_to_all_destinations", + "compute_embedding", ] OPTIONAL_PROCESSING_STEPS = { @@ -217,6 +218,7 @@ def get_file_overall_status(db: Session, file_id: int) -> Dict: "embed_metadata_into_pdf", "finalize_document_storage", "send_to_all_destinations", + "compute_embedding", } # Add check_for_duplicates if deduplication is enabled @@ -325,6 +327,7 @@ def get_step_summary(db: Session, file_id: int) -> Dict: "embed_metadata_into_pdf", "finalize_document_storage", "send_to_all_destinations", + "compute_embedding", } # Add check_for_duplicates if deduplication is enabled diff --git a/app/views/files.py b/app/views/files.py index dad67680..c1a2c7da 100644 --- a/app/views/files.py +++ b/app/views/files.py @@ -389,8 +389,12 @@ def _compute_processing_flow(logs): }, "extract_metadata_with_gpt": {"label": "Extract Metadata (GPT)", "next": ["embed_metadata_into_pdf"]}, "embed_metadata_into_pdf": {"label": "Embed Metadata into PDF", "next": ["finalize_document_storage"]}, - "finalize_document_storage": {"label": "Finalize & Queue Distribution", "next": ["send_to_all_destinations"]}, + "finalize_document_storage": { + "label": "Finalize & Queue Distribution", + "next": ["send_to_all_destinations", "compute_embedding"], + }, "send_to_all_destinations": {"label": "Upload to Destinations", "next": [], "has_branches": True}, + "compute_embedding": {"label": "Compute Embedding", "next": []}, } # Filter out deduplication step if not enabled or if not showing it @@ -815,3 +819,54 @@ def duplicates_page( "error": str(e), }, ) + + +@router.get("/similarity") +@require_login +def similarity_dashboard_page( + request: Request, + db: Session = Depends(get_db), +): + """Render the corpus-wide similarity dashboard. + + Passes the configured threshold and embedding coverage stats so the + template can display them immediately while the JS fetches the actual + pairs from the API asynchronously. + """ + from app.config import settings + from app.models import FileRecord + + try: + total_files = db.query(FileRecord).count() + files_with_embedding = ( + db.query(FileRecord).filter(FileRecord.embedding.isnot(None), FileRecord.embedding != "").count() + ) + files_with_ocr = db.query(FileRecord).filter(FileRecord.ocr_text.isnot(None), FileRecord.ocr_text != "").count() + + return templates.TemplateResponse( + "similarity_dashboard.html", + { + "request": request, + "default_threshold": settings.near_duplicate_threshold, + "embedding_model": settings.embedding_model, + "total_files": total_files, + "files_with_embedding": files_with_embedding, + "files_with_ocr": files_with_ocr, + "files_missing_embedding": files_with_ocr - files_with_embedding, + }, + ) + except Exception as e: + logger.error(f"Error rendering similarity dashboard: {e}") + return templates.TemplateResponse( + "similarity_dashboard.html", + { + "request": request, + "default_threshold": 0.85, + "embedding_model": "text-embedding-3-small", + "total_files": 0, + "files_with_embedding": 0, + "files_with_ocr": 0, + "files_missing_embedding": 0, + "error": str(e), + }, + ) diff --git a/docs/API.md b/docs/API.md index 78356dab..304345b0 100644 --- a/docs/API.md +++ b/docs/API.md @@ -670,7 +670,7 @@ curl -OJ "http:///api/files/123/download?version=original" **GET** `/api/files/{file_id}/similar` -Find documents similar to the specified file using text embeddings and cosine similarity. Similarity scores range from 0 (completely different) to 1 (identical content). Embeddings are generated from OCR-extracted text and cached for subsequent requests. +Find documents similar to the specified file using pre-computed text embeddings and cosine similarity. Similarity scores range from 0 (completely different) to 1 (identical content). Embeddings are computed automatically during document ingestion and cached in the database. **Parameters**: - `limit` (optional, default: `5`, max: `20`): Maximum number of similar documents to return @@ -706,9 +706,94 @@ curl "http:///api/files/42/similar?limit=10&threshold=0.5" **Error Responses**: - `404`: File not found - `422`: Invalid query parameters (limit or threshold out of range) -- `500`: Embedding generation failed +- `500`: Internal error -> **Note:** Documents without OCR text are excluded from similarity comparisons. The response includes a `message` field when the target file has no OCR text available. +> **Note:** Only pre-computed embeddings are used — no API calls are made during the query. If a file's embedding has not been computed yet, the response includes a `message` field explaining this. Documents without OCR text are excluded from similarity comparisons. + +### Similarity Pairs (Corpus-Wide) + +**GET** `/api/similarity/pairs` + +Scan the entire document corpus for pairs of highly similar documents, ranked by score. Unlike the per-file `/files/{id}/similar` endpoint, this discovers all matching pairs across all files. + +**Parameters**: +- `threshold` (optional, default: `0.7`, range: `0.0–1.0`): Minimum similarity score for a pair +- `limit` (optional, default: `50`, max: `200`): Maximum pairs per page +- `page` (optional, default: `1`): Page number + +**Response**: +```json +{ + "pairs": [ + { + "file_a": { + "file_id": 1, + "original_filename": "invoice_jan.pdf", + "document_title": "January Invoice", + "mime_type": "application/pdf", + "created_at": "2026-01-15T10:30:00+00:00" + }, + "file_b": { + "file_id": 5, + "original_filename": "invoice_feb.pdf", + "document_title": "February Invoice", + "mime_type": "application/pdf", + "created_at": "2026-02-15T10:30:00+00:00" + }, + "similarity_score": 0.94 + } + ], + "total_pairs": 12, + "threshold": 0.7, + "page": 1, + "pages": 1, + "per_page": 50, + "embedding_coverage": { + "total_files": 120, + "files_with_embedding": 95 + } +} +``` + +**Example**: +```bash +# Find all document pairs above 90% similarity +curl "http:///api/similarity/pairs?threshold=0.9" +``` + +### Embedding Diagnostics + +**GET** `/api/files/{file_id}/embedding-status` + +Check the embedding status for a specific file: whether OCR text is available, whether an embedding has been computed, and how many dimensions it has. + +```bash +curl "http:///api/files/42/embedding-status" +``` + +**POST** `/api/files/{file_id}/compute-embedding` + +Manually trigger embedding computation for a single file. Useful for debugging or re-computing after configuration changes. Requires OCR text to be available. + +```bash +curl -X POST "http:///api/files/42/compute-embedding" +``` + +**GET** `/api/diagnostic/embeddings` + +Get an overview of embedding coverage across all files: total files, how many have OCR text, how many have embeddings, and per-file status. + +```bash +curl "http:///api/diagnostic/embeddings" +``` + +**POST** `/api/diagnostic/compute-all-embeddings` + +Queue embedding computation for all files that have OCR text but no embedding yet. Each file is processed as a separate background task. + +```bash +curl -X POST "http:///api/diagnostic/compute-all-embeddings" +``` ### Batch Processing diff --git a/docs/ConfigurationGuide.md b/docs/ConfigurationGuide.md index 1d60ee4a..e7fd4120 100644 --- a/docs/ConfigurationGuide.md +++ b/docs/ConfigurationGuide.md @@ -868,10 +868,15 @@ After OCR processes a document, its extracted text is converted to a vector embe | Variable | Description | Default | |---|---|---| | `NEAR_DUPLICATE_THRESHOLD` | Minimum cosine similarity (0–1) for two documents to be considered near-duplicates. `0.85` means ≥ 85 % semantic overlap. | `0.85` | +| `EMBEDDING_MODEL` | Model name for generating text embeddings via the OpenAI-compatible API. Must be supported by the endpoint configured with `OPENAI_BASE_URL`. | `text-embedding-3-small` | +| `EMBEDDING_MAX_TOKENS` | Maximum tokens to send to the embedding model. Text is truncated to approximately this many tokens before calling the API. Set below the model's context window (e.g. 8 000 for an 8 192-token model). | `8000` | Near-duplicate detection: -- Is performed **on demand** via `GET /api/files/{id}/duplicates` — not automatically during ingest (OCR text is required). -- Is exposed in the **Duplicates** management page (`/duplicates` → "Near-Duplicate Finder" tab). +- Embeddings are computed **automatically during document ingestion** as a processing step ("Compute Embedding"). +- A periodic **backfill task** (every 5 minutes) picks up any files that were processed before the embedding pipeline was enabled. +- The **Similarity dashboard** (`/similarity`) shows all pairs of documents above the threshold, ranked by score. +- The **Duplicates** management page (`/duplicates` → "Near-Duplicate Finder" tab) allows per-file lookup. +- Debug endpoints are available to inspect embedding status and trigger recomputation (see API docs). - Documents without OCR text cannot be compared and are excluded from results. A score of **≥ 0.90** reliably identifies the same document scanned twice. A score of **0.70–0.90** suggests partial content overlap. Adjust `NEAR_DUPLICATE_THRESHOLD` to tune sensitivity. diff --git a/frontend/templates/base.html b/frontend/templates/base.html index ef42f9cc..44f537fe 100644 --- a/frontend/templates/base.html +++ b/frontend/templates/base.html @@ -104,6 +104,9 @@ Duplicates + + Similarity + Queue Monitor @@ -184,6 +187,9 @@ Duplicates + + Similarity + Queue Monitor diff --git a/frontend/templates/file_detail.html b/frontend/templates/file_detail.html index e57cd92d..c9d34f57 100644 --- a/frontend/templates/file_detail.html +++ b/frontend/templates/file_detail.html @@ -952,10 +952,75 @@ loadPDF('processed', fileId); {% endif %} - // Load similar documents + // Load similar documents and embedding status loadSimilarDocuments(fileId); + loadEmbeddingStatus(fileId); }); + // Fetch and display embedding debug info + async function loadEmbeddingStatus(fileId) { + const statusDiv = document.getElementById('embedding-status'); + const actionsDiv = document.getElementById('embedding-actions'); + try { + const response = await fetch(`/api/files/${fileId}/embedding-status`); + if (!response.ok) return; + const data = await response.json(); + + let html = ''; + if (data.has_embedding) { + html = ` Embedding: ${data.embedding_dimensions} dimensions (model: ${data.embedding_model})`; + statusDiv.style.backgroundColor = '#f0fff4'; + statusDiv.style.color = '#276749'; + statusDiv.style.border = '1px solid #c6f6d5'; + } else if (data.has_ocr_text) { + html = ` Embedding: not yet computed — OCR text available (${data.ocr_text_length} chars). Click "Recompute Embedding" to generate.`; + statusDiv.style.backgroundColor = '#fffff0'; + statusDiv.style.color = '#975a16'; + statusDiv.style.border = '1px solid #fefcbf'; + } else { + html = ` Embedding: unavailable — no OCR text extracted for this file.`; + statusDiv.style.backgroundColor = '#fff5f5'; + statusDiv.style.color = '#9b2c2c'; + statusDiv.style.border = '1px solid #fed7d7'; + } + statusDiv.innerHTML = html; + statusDiv.style.display = 'block'; + actionsDiv.style.display = data.has_ocr_text ? 'block' : 'none'; + } catch (error) { + console.error('Error loading embedding status:', error); + } + } + + // Trigger embedding recomputation + async function recomputeEmbedding(fileId) { + const btn = document.getElementById('recompute-embedding-btn'); + btn.disabled = true; + btn.innerHTML = ' Computing…'; + + try { + const response = await fetch(`/api/files/${fileId}/compute-embedding`, { method: 'POST' }); + const data = await response.json(); + if (!response.ok) { + throw new Error(data.detail || 'Request failed'); + } + btn.innerHTML = ' Done'; + btn.style.backgroundColor = '#48bb78'; + // Refresh both embedding status and similar documents + loadEmbeddingStatus(fileId); + loadSimilarDocuments(fileId); + } catch (error) { + btn.innerHTML = ' Failed'; + btn.style.backgroundColor = '#e53e3e'; + console.error('Recompute embedding failed:', error); + } finally { + setTimeout(() => { + btn.disabled = false; + btn.innerHTML = ' Recompute Embedding'; + btn.style.backgroundColor = '#4299e1'; + }, 3000); + } + } + // Similar documents loading async function loadSimilarDocuments(fileId) { const loadingDiv = document.getElementById('similar-documents-loading'); @@ -1602,6 +1667,8 @@

Similar Documents

+ +

Searching for similar documents…

@@ -1615,6 +1682,12 @@

Failed to load similar documents.

+ +
diff --git a/frontend/templates/similarity_dashboard.html b/frontend/templates/similarity_dashboard.html new file mode 100644 index 00000000..1090ce3a --- /dev/null +++ b/frontend/templates/similarity_dashboard.html @@ -0,0 +1,368 @@ +{% extends "base.html" %} +{% block title %}Document Similarity - DocuElevate{% endblock %} + +{% block head_extra %} + +{% endblock %} + +{% block content %} +
+

+ + Document Similarity +

+

+ Pairs of documents with high semantic similarity, ranked by score. + Embeddings are computed during document ingestion; a background task + also backfills any files that were processed before this feature was enabled. +

+ + +
+
+
{{ total_files }}
+
Total Files
+
+
+
{{ files_with_embedding }}
+
With Embedding
+
+
+
+ {{ files_missing_embedding }} +
+
Missing Embedding
+
+
+
{{ embedding_model }}
+
Embedding Model
+
+
+ + {% if files_missing_embedding > 0 %} +
+ + {{ files_missing_embedding }} file(s) have OCR text but no embedding yet. + The background task will compute them automatically every 5 minutes, or you can + . +
+ {% endif %} + + +
+
+
+ + +
+
+ + +
+ +
+
+ + +
+ +

Scanning for similar document pairs…

+
+ + + + + + +
+{% endblock %} + +{% block scripts %} + +{% endblock %} diff --git a/tests/test_duplicates.py b/tests/test_duplicates.py index 6dd84b2d..45486168 100644 --- a/tests/test_duplicates.py +++ b/tests/test_duplicates.py @@ -7,6 +7,7 @@ Covers: - ``GET /duplicates`` — duplicate management UI page """ +import json from unittest.mock import patch import pytest @@ -19,7 +20,7 @@ from app.models import FileRecord # --------------------------------------------------------------------------- -def _make_file(db, *, filehash, filename, is_duplicate=False, duplicate_of_id=None, ocr_text=None): +def _make_file(db, *, filehash, filename, is_duplicate=False, duplicate_of_id=None, ocr_text=None, embedding=None): """Insert a FileRecord and return it.""" record = FileRecord( filehash=filehash, @@ -30,6 +31,7 @@ def _make_file(db, *, filehash, filename, is_duplicate=False, duplicate_of_id=No is_duplicate=is_duplicate, duplicate_of_id=duplicate_of_id, ocr_text=ocr_text, + embedding=embedding, ) db.add(record) db.commit() @@ -195,25 +197,24 @@ class TestGetFileDuplicates: assert data["duplicate_of"]["id"] == orig.id @pytest.mark.integration - @patch("app.utils.similarity.generate_embedding") - def test_near_duplicates_returned(self, mock_embed, client: TestClient, db_session): + def test_near_duplicates_returned(self, client: TestClient, db_session): """Near-duplicates found via embedding similarity should appear in results.""" + embedding = json.dumps([1.0, 0.0, 0.0]) target = _make_file( db_session, filehash="th1", filename="target.pdf", ocr_text="Invoice from Acme Corp for January services rendered", + embedding=embedding, ) similar = _make_file( db_session, filehash="th2", # different hash — same content (re-scan) filename="rescan.pdf", ocr_text="Invoice from Acme Corp for January services rendered", + embedding=embedding, ) - # Same embedding → cosine similarity = 1.0 - mock_embed.return_value = [1.0, 0.0, 0.0] - response = client.get(f"/api/files/{target.id}/duplicates?near_duplicate_threshold=0.8") assert response.status_code == 200 data = response.json() diff --git a/tests/test_similarity.py b/tests/test_similarity.py index b81fd053..45357747 100644 --- a/tests/test_similarity.py +++ b/tests/test_similarity.py @@ -108,18 +108,23 @@ class TestFindSimilarDocuments: assert result == [] @pytest.mark.unit - @patch("app.utils.similarity.generate_embedding") - def test_finds_similar_documents(self, mock_embed, db_session): - """Should find similar documents based on embedding similarity.""" - # Create a target file with OCR text + def test_finds_similar_documents(self, db_session): + """Should find similar documents based on pre-computed embedding similarity.""" + # Pre-computed embeddings that reflect similarity + target_embedding = [1.0, 0.0, 0.0] + similar_embedding = [0.95, 0.05, 0.0] + different_embedding = [0.0, 0.0, 1.0] + + # Create a target file with pre-computed embedding target = FileRecord( filehash="hash1", local_filename="/tmp/target.pdf", file_size=1024, original_filename="target.pdf", ocr_text="This is an invoice from Amazon for January 2026", + embedding=json.dumps(target_embedding), ) - # Create a similar file + # Create a similar file with pre-computed embedding similar = FileRecord( filehash="hash2", local_filename="/tmp/similar.pdf", @@ -128,8 +133,9 @@ class TestFindSimilarDocuments: ocr_text="This is an invoice from Amazon for February 2026", document_title="Amazon Invoice Feb", mime_type="application/pdf", + embedding=json.dumps(similar_embedding), ) - # Create a different file + # Create a different file with pre-computed embedding different = FileRecord( filehash="hash3", local_filename="/tmp/different.pdf", @@ -138,26 +144,12 @@ class TestFindSimilarDocuments: ocr_text="Recipe for chocolate cake with detailed instructions", document_title="Chocolate Cake Recipe", mime_type="application/pdf", + embedding=json.dumps(different_embedding), ) db_session.add_all([target, similar, different]) db_session.commit() - # Mock embeddings that reflect similarity - target_embedding = [1.0, 0.0, 0.0] - similar_embedding = [0.95, 0.05, 0.0] - different_embedding = [0.0, 0.0, 1.0] - - def mock_embed_side_effect(text): - if "January" in text or "invoice" in text.lower()[:30]: - return target_embedding - elif "February" in text: - return similar_embedding - else: - return different_embedding - - mock_embed.side_effect = mock_embed_side_effect - result = find_similar_documents(db_session, file_id=target.id, threshold=0.3) assert len(result) == 1 @@ -166,8 +158,7 @@ class TestFindSimilarDocuments: assert result[0]["original_filename"] == "similar.pdf" @pytest.mark.unit - @patch("app.utils.similarity.generate_embedding") - def test_respects_threshold(self, mock_embed, db_session): + def test_respects_threshold(self, db_session): """Should filter out documents below the threshold.""" target = FileRecord( filehash="hash1", @@ -175,6 +166,7 @@ class TestFindSimilarDocuments: file_size=100, original_filename="target.pdf", ocr_text="target text", + embedding=json.dumps([1.0, 0.0]), ) candidate = FileRecord( filehash="hash2", @@ -182,26 +174,25 @@ class TestFindSimilarDocuments: file_size=100, original_filename="candidate.pdf", ocr_text="different text", + embedding=json.dumps([0.1, 0.99]), ) db_session.add_all([target, candidate]) db_session.commit() - # Return nearly orthogonal vectors -> low similarity - mock_embed.side_effect = lambda text: [1.0, 0.0] if "target" in text else [0.1, 0.99] - result = find_similar_documents(db_session, file_id=target.id, threshold=0.9) assert len(result) == 0 @pytest.mark.unit - @patch("app.utils.similarity.generate_embedding") - def test_respects_limit(self, mock_embed, db_session): + def test_respects_limit(self, db_session): """Should respect the limit parameter.""" + embedding = [1.0, 0.0, 0.0] target = FileRecord( filehash="hash0", local_filename="/tmp/t.pdf", file_size=100, original_filename="target.pdf", ocr_text="target text", + embedding=json.dumps(embedding), ) db_session.add(target) @@ -212,12 +203,11 @@ class TestFindSimilarDocuments: file_size=100, original_filename=f"candidate_{i}.pdf", ocr_text=f"similar text {i}", + embedding=json.dumps(embedding), ) db_session.add(f) db_session.commit() - mock_embed.return_value = [1.0, 0.0, 0.0] - result = find_similar_documents(db_session, file_id=target.id, limit=2, threshold=0.0) assert len(result) <= 2 @@ -286,15 +276,16 @@ class TestSimilarDocumentsAPI: assert "message" in data @pytest.mark.integration - @patch("app.utils.similarity.generate_embedding") - def test_returns_similar_documents(self, mock_embed, client: TestClient, db_session): + def test_returns_similar_documents(self, client: TestClient, db_session): """Should return similar documents with scores.""" + embedding = [1.0, 0.0, 0.0] target = FileRecord( filehash="hash1", local_filename="/tmp/target.pdf", file_size=1024, original_filename="target.pdf", ocr_text="Invoice from Amazon January 2026", + embedding=json.dumps(embedding), ) similar = FileRecord( filehash="hash2", @@ -304,12 +295,11 @@ class TestSimilarDocumentsAPI: ocr_text="Invoice from Amazon February 2026", document_title="Amazon Invoice", mime_type="application/pdf", + embedding=json.dumps(embedding), ) db_session.add_all([target, similar]) db_session.commit() - mock_embed.return_value = [1.0, 0.0, 0.0] - response = client.get(f"/api/files/{target.id}/similar") assert response.status_code == 200 data = response.json() @@ -324,15 +314,16 @@ class TestSimilarDocumentsAPI: assert "original_filename" in doc @pytest.mark.integration - @patch("app.utils.similarity.generate_embedding") - def test_query_parameters(self, mock_embed, client: TestClient, db_session): + def test_query_parameters(self, client: TestClient, db_session): """Should respect limit and threshold query parameters.""" + embedding = [1.0, 0.0] target = FileRecord( filehash="hash1", local_filename="/tmp/t.pdf", file_size=100, original_filename="t.pdf", ocr_text="test", + embedding=json.dumps(embedding), ) db_session.add(target) @@ -343,12 +334,11 @@ class TestSimilarDocumentsAPI: file_size=100, original_filename=f"c{i}.pdf", ocr_text=f"text {i}", + embedding=json.dumps(embedding), ) db_session.add(f) db_session.commit() - mock_embed.return_value = [1.0, 0.0] - response = client.get(f"/api/files/{target.id}/similar?limit=2&threshold=0.0") assert response.status_code == 200 data = response.json() @@ -403,15 +393,36 @@ class TestSimilarDocumentsAPI: assert data["count"] == 0 @pytest.mark.integration - @patch("app.utils.similarity.generate_embedding") - def test_response_structure(self, mock_embed, client: TestClient, db_session): + def test_embedding_not_computed_message(self, client: TestClient, db_session): + """Should return a message when OCR text exists but no embedding yet.""" + file_record = FileRecord( + filehash="noembhash", + local_filename="/tmp/noemb.pdf", + file_size=100, + original_filename="noemb.pdf", + ocr_text="Some OCR text content", + embedding=None, + ) + db_session.add(file_record) + db_session.commit() + + response = client.get(f"/api/files/{file_record.id}/similar") + assert response.status_code == 200 + data = response.json() + assert data["count"] == 0 + assert "message" in data + assert "not yet computed" in data["message"].lower() + + def test_response_structure(self, client: TestClient, db_session): """Should return proper response structure for each similar document.""" + embedding = [1.0, 0.0] target = FileRecord( filehash="h1", local_filename="/tmp/t.pdf", file_size=100, original_filename="target.pdf", ocr_text="Some text content here", + embedding=json.dumps(embedding), ) other = FileRecord( filehash="h2", @@ -421,12 +432,11 @@ class TestSimilarDocumentsAPI: ocr_text="Some similar text content", document_title="Other Doc", mime_type="application/pdf", + embedding=json.dumps(embedding), ) db_session.add_all([target, other]) db_session.commit() - mock_embed.return_value = [1.0, 0.0] - response = client.get(f"/api/files/{target.id}/similar") assert response.status_code == 200 data = response.json() @@ -443,3 +453,562 @@ class TestSimilarDocumentsAPI: assert "similarity_score" in doc assert "mime_type" in doc assert "created_at" in doc + + +# --------------------------------------------------------------------------- +# Tests for embedding status endpoint +# --------------------------------------------------------------------------- + + +class TestEmbeddingStatusAPI: + """Integration tests for GET /api/files/{file_id}/embedding-status.""" + + @pytest.mark.integration + def test_file_not_found(self, client: TestClient): + """Should return 404 for non-existent file.""" + response = client.get("/api/files/9999/embedding-status") + assert response.status_code == 404 + + @pytest.mark.integration + def test_file_without_embedding_or_ocr(self, client: TestClient, db_session): + """Should report no embedding and no OCR text.""" + file_record = FileRecord( + filehash="abc1", + local_filename="/tmp/test.pdf", + file_size=100, + original_filename="test.pdf", + ) + db_session.add(file_record) + db_session.commit() + + response = client.get(f"/api/files/{file_record.id}/embedding-status") + assert response.status_code == 200 + data = response.json() + assert data["file_id"] == file_record.id + assert data["has_embedding"] is False + assert data["embedding_dimensions"] is None + assert data["has_ocr_text"] is False + assert data["ocr_text_length"] == 0 + assert "embedding_model" in data + + @pytest.mark.integration + def test_file_with_ocr_text_no_embedding(self, client: TestClient, db_session): + """Should report OCR text present but no embedding.""" + file_record = FileRecord( + filehash="abc2", + local_filename="/tmp/test2.pdf", + file_size=100, + original_filename="test2.pdf", + ocr_text="Some OCR text content", + ) + db_session.add(file_record) + db_session.commit() + + response = client.get(f"/api/files/{file_record.id}/embedding-status") + assert response.status_code == 200 + data = response.json() + assert data["has_embedding"] is False + assert data["has_ocr_text"] is True + assert data["ocr_text_length"] == 21 + + @pytest.mark.integration + def test_file_with_cached_embedding(self, client: TestClient, db_session): + """Should report embedding present with correct dimensions.""" + embedding = [0.1, 0.2, 0.3, 0.4, 0.5] + file_record = FileRecord( + filehash="abc3", + local_filename="/tmp/test3.pdf", + file_size=100, + original_filename="test3.pdf", + ocr_text="Some text", + embedding=json.dumps(embedding), + ) + db_session.add(file_record) + db_session.commit() + + response = client.get(f"/api/files/{file_record.id}/embedding-status") + assert response.status_code == 200 + data = response.json() + assert data["has_embedding"] is True + assert data["embedding_dimensions"] == 5 + assert data["has_ocr_text"] is True + + +# --------------------------------------------------------------------------- +# Tests for compute-embedding endpoint +# --------------------------------------------------------------------------- + + +class TestComputeEmbeddingAPI: + """Integration tests for POST /api/files/{file_id}/compute-embedding.""" + + @pytest.mark.integration + def test_file_not_found(self, client: TestClient): + """Should return 404 for non-existent file.""" + response = client.post("/api/files/9999/compute-embedding") + assert response.status_code == 404 + + @pytest.mark.integration + def test_no_ocr_text(self, client: TestClient, db_session): + """Should return 400 when file has no OCR text.""" + file_record = FileRecord( + filehash="emb1", + local_filename="/tmp/emb1.pdf", + file_size=100, + original_filename="emb1.pdf", + ) + db_session.add(file_record) + db_session.commit() + + response = client.post(f"/api/files/{file_record.id}/compute-embedding") + assert response.status_code == 400 + + @pytest.mark.integration + @patch("app.utils.similarity.generate_embedding") + def test_computes_embedding(self, mock_embed, client: TestClient, db_session): + """Should compute and store an embedding.""" + mock_embed.return_value = [0.1, 0.2, 0.3] + + file_record = FileRecord( + filehash="emb2", + local_filename="/tmp/emb2.pdf", + file_size=100, + original_filename="emb2.pdf", + ocr_text="Some document text", + ) + db_session.add(file_record) + db_session.commit() + + response = client.post(f"/api/files/{file_record.id}/compute-embedding") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "success" + assert data["embedding_dimensions"] == 3 + + # Verify embedding is stored + db_session.refresh(file_record) + assert file_record.embedding is not None + stored = json.loads(file_record.embedding) + assert len(stored) == 3 + + @pytest.mark.integration + @patch("app.utils.similarity.generate_embedding") + def test_recomputes_existing_embedding(self, mock_embed, client: TestClient, db_session): + """Should overwrite existing embedding when recomputing.""" + mock_embed.return_value = [0.9, 0.8, 0.7] + + file_record = FileRecord( + filehash="emb3", + local_filename="/tmp/emb3.pdf", + file_size=100, + original_filename="emb3.pdf", + ocr_text="Some text", + embedding=json.dumps([0.1, 0.2, 0.3]), + ) + db_session.add(file_record) + db_session.commit() + + response = client.post(f"/api/files/{file_record.id}/compute-embedding") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "success" + + db_session.refresh(file_record) + stored = json.loads(file_record.embedding) + assert stored == [0.9, 0.8, 0.7] + + +# --------------------------------------------------------------------------- +# Tests for diagnostic embeddings overview endpoint +# --------------------------------------------------------------------------- + + +class TestEmbeddingsOverviewAPI: + """Integration tests for GET /api/diagnostic/embeddings.""" + + @pytest.mark.integration + def test_empty_database(self, client: TestClient): + """Should return zero counts on empty database.""" + response = client.get("/api/diagnostic/embeddings") + assert response.status_code == 200 + data = response.json() + assert data["total_files"] == 0 + assert data["files_with_ocr_text"] == 0 + assert data["files_with_embedding"] == 0 + assert data["files_missing_embedding"] == 0 + assert "embedding_model" in data + assert data["files"] == [] + + @pytest.mark.integration + def test_mixed_files(self, client: TestClient, db_session): + """Should report correct counts for mixed embedding states.""" + # File with both OCR text and embedding + f1 = FileRecord( + filehash="diag1", + local_filename="/tmp/d1.pdf", + file_size=100, + original_filename="d1.pdf", + ocr_text="Some text", + embedding=json.dumps([0.1, 0.2]), + ) + # File with OCR text but no embedding + f2 = FileRecord( + filehash="diag2", + local_filename="/tmp/d2.pdf", + file_size=100, + original_filename="d2.pdf", + ocr_text="More text", + ) + # File with no OCR text + f3 = FileRecord( + filehash="diag3", + local_filename="/tmp/d3.pdf", + file_size=100, + original_filename="d3.pdf", + ) + db_session.add_all([f1, f2, f3]) + db_session.commit() + + response = client.get("/api/diagnostic/embeddings") + assert response.status_code == 200 + data = response.json() + assert data["total_files"] == 3 + assert data["files_with_ocr_text"] == 2 + assert data["files_with_embedding"] == 1 + assert data["files_missing_embedding"] == 1 + assert len(data["files"]) == 3 + + # Check per-file info + files_by_id = {f["file_id"]: f for f in data["files"]} + assert files_by_id[f1.id]["has_embedding"] is True + assert files_by_id[f1.id]["embedding_dimensions"] == 2 + assert files_by_id[f2.id]["has_embedding"] is False + assert files_by_id[f2.id]["has_ocr_text"] is True + assert files_by_id[f3.id]["has_ocr_text"] is False + + +# --------------------------------------------------------------------------- +# Tests for compute-all-embeddings endpoint +# --------------------------------------------------------------------------- + + +class TestComputeAllEmbeddingsAPI: + """Integration tests for POST /api/diagnostic/compute-all-embeddings.""" + + @pytest.mark.integration + @patch("app.tasks.compute_embedding.compute_document_embedding.delay") + def test_queues_tasks_for_files_missing_embeddings(self, mock_delay, client: TestClient, db_session): + """Should queue embedding tasks for files with OCR text but no embedding.""" + # File with OCR text but no embedding -> should be queued + f1 = FileRecord( + filehash="all1", + local_filename="/tmp/a1.pdf", + file_size=100, + original_filename="a1.pdf", + ocr_text="Text for embedding", + ) + # File already with embedding -> should NOT be queued + f2 = FileRecord( + filehash="all2", + local_filename="/tmp/a2.pdf", + file_size=100, + original_filename="a2.pdf", + ocr_text="More text", + embedding=json.dumps([0.1, 0.2]), + ) + # File without OCR text -> should NOT be queued + f3 = FileRecord( + filehash="all3", + local_filename="/tmp/a3.pdf", + file_size=100, + original_filename="a3.pdf", + ) + db_session.add_all([f1, f2, f3]) + db_session.commit() + + response = client.post("/api/diagnostic/compute-all-embeddings") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "queued" + assert data["files_queued"] == 1 + mock_delay.assert_called_once_with(f1.id) + + @pytest.mark.integration + @patch("app.tasks.compute_embedding.compute_document_embedding.delay") + def test_empty_database_queues_nothing(self, mock_delay, client: TestClient): + """Should queue nothing when database is empty.""" + response = client.post("/api/diagnostic/compute-all-embeddings") + assert response.status_code == 200 + data = response.json() + assert data["files_queued"] == 0 + mock_delay.assert_not_called() + + +# --------------------------------------------------------------------------- +# Tests for compute_document_embedding Celery task +# --------------------------------------------------------------------------- + + +class TestComputeDocumentEmbeddingTask: + """Unit tests for the compute_document_embedding Celery task.""" + + @pytest.mark.unit + @patch("app.utils.similarity.generate_embedding") + def test_computes_embedding_for_file(self, mock_embed, db_session): + """Should compute and store embedding when file has OCR text.""" + mock_embed.return_value = [0.1, 0.2, 0.3] + + file_record = FileRecord( + filehash="task1", + local_filename="/tmp/task1.pdf", + file_size=100, + original_filename="task1.pdf", + ocr_text="Some document text", + ) + db_session.add(file_record) + db_session.commit() + + from app.tasks.compute_embedding import compute_document_embedding + + # Patch SessionLocal to return our test session + with patch("app.tasks.compute_embedding.SessionLocal") as mock_session_local: + mock_session_local.return_value.__enter__ = lambda self: db_session + mock_session_local.return_value.__exit__ = lambda self, *args: None + + result = compute_document_embedding(file_record.id) + + assert result["status"] == "success" + assert "dimensions" in result["detail"] + + @pytest.mark.unit + def test_skips_missing_file(self, db_session): + """Should skip when file ID does not exist.""" + from app.tasks.compute_embedding import compute_document_embedding + + with patch("app.tasks.compute_embedding.SessionLocal") as mock_session_local: + mock_session_local.return_value.__enter__ = lambda self: db_session + mock_session_local.return_value.__exit__ = lambda self, *args: None + + result = compute_document_embedding(9999) + + assert result["status"] == "skipped" + + @pytest.mark.unit + def test_skips_file_without_ocr_text(self, db_session): + """Should skip when file has no OCR text.""" + file_record = FileRecord( + filehash="task2", + local_filename="/tmp/task2.pdf", + file_size=100, + original_filename="task2.pdf", + ) + db_session.add(file_record) + db_session.commit() + + from app.tasks.compute_embedding import compute_document_embedding + + with patch("app.tasks.compute_embedding.SessionLocal") as mock_session_local: + mock_session_local.return_value.__enter__ = lambda self: db_session + mock_session_local.return_value.__exit__ = lambda self, *args: None + + result = compute_document_embedding(file_record.id) + + assert result["status"] == "skipped" + + @pytest.mark.unit + def test_skips_file_with_existing_embedding(self, db_session): + """Should skip when file already has a cached embedding.""" + file_record = FileRecord( + filehash="task3", + local_filename="/tmp/task3.pdf", + file_size=100, + original_filename="task3.pdf", + ocr_text="Some text", + embedding=json.dumps([0.1, 0.2]), + ) + db_session.add(file_record) + db_session.commit() + + from app.tasks.compute_embedding import compute_document_embedding + + with patch("app.tasks.compute_embedding.SessionLocal") as mock_session_local: + mock_session_local.return_value.__enter__ = lambda self: db_session + mock_session_local.return_value.__exit__ = lambda self, *args: None + + result = compute_document_embedding(file_record.id) + + assert result["status"] == "skipped" + assert "already cached" in result["detail"] + + +# --------------------------------------------------------------------------- +# Tests for similarity pairs endpoint +# --------------------------------------------------------------------------- + + +class TestSimilarityPairsAPI: + """Integration tests for GET /api/similarity/pairs.""" + + @pytest.mark.integration + def test_empty_database(self, client: TestClient): + """Should return zero pairs on empty database.""" + response = client.get("/api/similarity/pairs") + assert response.status_code == 200 + data = response.json() + assert data["total_pairs"] == 0 + assert data["pairs"] == [] + assert "embedding_coverage" in data + + @pytest.mark.integration + def test_finds_similar_pairs(self, client: TestClient, db_session): + """Should find and return pairs of similar files.""" + emb_a = [1.0, 0.0, 0.0] + emb_b = [0.98, 0.02, 0.0] # Very similar to A + emb_c = [0.0, 0.0, 1.0] # Different from A and B + + f1 = FileRecord( + filehash="pairA", + local_filename="/tmp/pA.pdf", + file_size=100, + original_filename="fileA.pdf", + ocr_text="text A", + embedding=json.dumps(emb_a), + ) + f2 = FileRecord( + filehash="pairB", + local_filename="/tmp/pB.pdf", + file_size=100, + original_filename="fileB.pdf", + ocr_text="text B", + embedding=json.dumps(emb_b), + ) + f3 = FileRecord( + filehash="pairC", + local_filename="/tmp/pC.pdf", + file_size=100, + original_filename="fileC.pdf", + ocr_text="text C", + embedding=json.dumps(emb_c), + ) + db_session.add_all([f1, f2, f3]) + db_session.commit() + + response = client.get("/api/similarity/pairs?threshold=0.9") + assert response.status_code == 200 + data = response.json() + + # Only A-B pair should be above 0.9 + assert data["total_pairs"] == 1 + pair = data["pairs"][0] + assert pair["similarity_score"] > 0.9 + pair_ids = {pair["file_a"]["file_id"], pair["file_b"]["file_id"]} + assert pair_ids == {f1.id, f2.id} + + @pytest.mark.integration + def test_respects_threshold(self, client: TestClient, db_session): + """Should filter pairs below threshold.""" + emb = [1.0, 0.0] + different_emb = [0.0, 1.0] + + f1 = FileRecord( + filehash="thA", + local_filename="/tmp/thA.pdf", + file_size=100, + original_filename="thA.pdf", + ocr_text="a", + embedding=json.dumps(emb), + ) + f2 = FileRecord( + filehash="thB", + local_filename="/tmp/thB.pdf", + file_size=100, + original_filename="thB.pdf", + ocr_text="b", + embedding=json.dumps(different_emb), + ) + db_session.add_all([f1, f2]) + db_session.commit() + + response = client.get("/api/similarity/pairs?threshold=0.9") + assert response.status_code == 200 + data = response.json() + assert data["total_pairs"] == 0 + + @pytest.mark.integration + def test_pagination(self, client: TestClient, db_session): + """Should respect pagination parameters.""" + emb = [1.0, 0.0, 0.0] + for i in range(5): + f = FileRecord( + filehash=f"pg{i}", + local_filename=f"/tmp/pg{i}.pdf", + file_size=100, + original_filename=f"pg{i}.pdf", + ocr_text=f"text {i}", + embedding=json.dumps(emb), + ) + db_session.add(f) + db_session.commit() + + response = client.get("/api/similarity/pairs?threshold=0.0&limit=2&page=1") + assert response.status_code == 200 + data = response.json() + assert len(data["pairs"]) <= 2 + assert data["per_page"] == 2 + + +# --------------------------------------------------------------------------- +# Tests for backfill_missing_embeddings task +# --------------------------------------------------------------------------- + + +class TestBackfillMissingEmbeddingsTask: + """Unit tests for the backfill_missing_embeddings Celery task.""" + + @pytest.mark.unit + @patch("app.tasks.compute_embedding.compute_document_embedding.delay") + def test_queues_files_without_embeddings(self, mock_delay, db_session): + """Should queue tasks for files with OCR text but no embedding.""" + f1 = FileRecord( + filehash="bf1", + local_filename="/tmp/bf1.pdf", + file_size=100, + original_filename="bf1.pdf", + ocr_text="Some text", + ) + f2 = FileRecord( + filehash="bf2", + local_filename="/tmp/bf2.pdf", + file_size=100, + original_filename="bf2.pdf", + ocr_text="More text", + embedding=json.dumps([0.1]), + ) + db_session.add_all([f1, f2]) + db_session.commit() + + from app.tasks.compute_embedding import backfill_missing_embeddings + + with patch("app.tasks.compute_embedding.SessionLocal") as mock_session_local: + mock_session_local.return_value.__enter__ = lambda self: db_session + mock_session_local.return_value.__exit__ = lambda self, *args: None + + result = backfill_missing_embeddings() + + assert result["queued"] == 1 + mock_delay.assert_called_once_with(f1.id) + + @pytest.mark.unit + @patch("app.tasks.compute_embedding.compute_document_embedding.delay") + def test_empty_database(self, mock_delay, db_session): + """Should queue nothing when no files need embeddings.""" + from app.tasks.compute_embedding import backfill_missing_embeddings + + with patch("app.tasks.compute_embedding.SessionLocal") as mock_session_local: + mock_session_local.return_value.__enter__ = lambda self: db_session + mock_session_local.return_value.__exit__ = lambda self, *args: None + + result = backfill_missing_embeddings() + + assert result["queued"] == 0 + mock_delay.assert_not_called()