copilot-swe-agent[bot]
55543be3b0
- Use column-only query in embeddings overview to reduce memory for 100K+ files - Add embedding_backfill_batch_size config setting (default 50) - Fix WCAG touch target on backfill button (min-height/min-width 44px) - Add inline comment explaining 3 chars/token truncation estimate - Import settings in compute_embedding task for configurable batch size Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
474 lines
14 KiB
Python
474 lines
14 KiB
Python
"""Document similarity API endpoints.
|
||
|
||
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
|
||
|
||
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
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
router = APIRouter()
|
||
|
||
DbSession = Annotated[Session, Depends(get_db)]
|
||
|
||
|
||
@router.get("/files/{file_id}/similar")
|
||
@require_login
|
||
def get_similar_documents(
|
||
request: Request,
|
||
file_id: int,
|
||
db: DbSession,
|
||
limit: int = Query(5, ge=1, le=20, description="Maximum number of similar documents to return"),
|
||
threshold: float = Query(0.3, ge=0.0, le=1.0, description="Minimum similarity score (0–1)"),
|
||
):
|
||
"""Find documents similar to the specified file.
|
||
|
||
Uses text embeddings generated from OCR-extracted text and cosine
|
||
similarity to rank documents by relevance. Similarity scores range
|
||
from 0 (completely different) to 1 (identical content).
|
||
|
||
Embeddings are generated on first access and cached for subsequent
|
||
requests. Documents without OCR text are excluded.
|
||
|
||
Query Parameters:
|
||
- limit: Maximum results to return (default: 5, max: 20)
|
||
- threshold: Minimum similarity score to include (default: 0.3)
|
||
|
||
Example:
|
||
```
|
||
GET /api/files/42/similar?limit=5&threshold=0.5
|
||
```
|
||
|
||
Response:
|
||
```json
|
||
{
|
||
"file_id": 42,
|
||
"similar_documents": [
|
||
{
|
||
"file_id": 15,
|
||
"original_filename": "Invoice_2026-01.pdf",
|
||
"document_title": "January Invoice",
|
||
"similarity_score": 0.8934,
|
||
"mime_type": "application/pdf",
|
||
"created_at": "2026-01-15T10:30:00+00:00"
|
||
}
|
||
],
|
||
"count": 1
|
||
}
|
||
```
|
||
"""
|
||
# Verify the file exists
|
||
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():
|
||
return {
|
||
"file_id": file_id,
|
||
"similar_documents": [],
|
||
"count": 0,
|
||
"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
|
||
|
||
similar = find_similar_documents(db, file_id, limit=limit, threshold=threshold)
|
||
|
||
return {
|
||
"file_id": file_id,
|
||
"similar_documents": similar,
|
||
"count": len(similar),
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"Error finding similar documents for file {file_id}: {e}")
|
||
raise HTTPException(
|
||
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,
|
||
}
|