gitea-admin/gh-christianlouis-docuelevate
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(similarity): add embedding pipeline, debug endpoints, backfill task, and scalable similarity search

Browse Source 
- Add embedding_model config setting (replaces hardcoded text-embedding-3-small)
- Add compute_document_embedding Celery task for ingestion-time embedding
- Chain embedding task into finalize_document_storage pipeline
- Add backfill_missing_embeddings periodic task (every 5 min) for legacy files
- Add debug API endpoints: embedding-status, compute-embedding, diagnostic/embeddings, diagnostic/compute-all-embeddings
- Refactor find_similar_documents to only use pre-computed embeddings (no lazy API calls)
- Use yield_per(500) and column-only queries for 100K+ scale
- Add embedding status indicator and recompute button in file detail UI

Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
2026-03-02 13:01:57 +00:00
parent b435957a9b
commit 8d7c8e7c4e
8 changed files with 964 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/api/similarity.py
+238 -2
View File
@@ -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
@@ -95,3 +98,236 @@ 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
}
]
}
```
"""
all_files = db.query(FileRecord).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,
}
app/celery_worker.py
+8
View File
@@ -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
app/config.py
+7
View File
@@ -323,6 +323,13 @@ 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."
),
)
# Text quality check - AI-based assessment of embedded PDF text
enable_text_quality_check: bool = Field(
app/tasks/compute_embedding.py
+148
View File
@@ -0,0 +1,148 @@
"""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 app.celery_app import celery
from app.database import SessionLocal
from app.models import FileRecord
from app.tasks.retry_config import BaseTaskWithRetry
from app.utils import log_task_progress
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"}
# 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,
)
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,
)
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)
if embedding:
log_task_progress(
task_id,
"compute_embedding",
"success",
f"Embedding computed ({len(embedding)} dimensions)",
file_id=file_id,
)
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,
)
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,
)
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 = 50 # max files to queue per run
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}
app/tasks/finalize_document_storage.py
+10
View File
@@ -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.
app/utils/similarity.py
+93 -30
View File
@@ -38,13 +38,14 @@ 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``.
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 +54,16 @@ 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.
"""
if model is None:
model = settings.embedding_model
# Truncate very long texts to stay within token limits (~4 chars per token)
max_chars = 30000
if len(text) > max_chars:
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 +94,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 +138,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 +204,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,
}
)