feat(similarity): add embedding pipeline, debug endpoints, backfill task, and scalable similarity search
- 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]
+93
-30
@@ -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,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user