gh-christianlouis-docuelevate/app/utils/similarity.py

"""Document similarity detection using text embeddings and cosine similarity.

Provides functions to generate text embeddings via the configured AI provider
(OpenAI-compatible) and compute cosine similarity scores between documents.
Embeddings are cached in the ``FileRecord.embedding`` column to avoid
redundant API calls.
"""

import json
import logging
import math
from typing import Any

from sqlalchemy.orm import Session

from app.config import settings

logger = logging.getLogger(__name__)


def _get_embedding_client() -> Any:
    """Create an OpenAI client for embedding generation.

    Returns:
        An ``openai.OpenAI`` client instance configured from application settings.

    Raises:
        RuntimeError: If the ``openai`` package is not installed.
    """
    try:
        import openai
    except ImportError as exc:
        raise RuntimeError("The 'openai' package is required for embedding generation") from exc

    return openai.OpenAI(
        api_key=settings.openai_api_key,
        base_url=settings.openai_base_url,
    )


def generate_embedding(text: str, model: str = "text-embedding-3-small") -> 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``.

    Returns:
        A list of floats representing the embedding vector.

    Raises:
        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 len(text) > max_chars:
        text = text[:max_chars]

    client = _get_embedding_client()
    response = client.embeddings.create(input=text, model=model)
    return response.data[0].embedding


def cosine_similarity(vec_a: list[float], vec_b: list[float]) -> float:
    """Compute cosine similarity between two vectors.

    Args:
        vec_a: First embedding vector.
        vec_b: Second embedding vector.

    Returns:
        A similarity score between 0 and 1.  Returns 0.0 if either vector
        has zero magnitude.
    """
    if len(vec_a) != len(vec_b):
        return 0.0

    dot_product = sum(a * b for a, b in zip(vec_a, vec_b, strict=True))
    magnitude_a = math.sqrt(sum(a * a for a in vec_a))
    magnitude_b = math.sqrt(sum(b * b for b in vec_b))

    if magnitude_a == 0.0 or magnitude_b == 0.0:
        return 0.0

    similarity = dot_product / (magnitude_a * magnitude_b)
    # Clamp to [0, 1] to handle floating-point drift
    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.

    Args:
        db: Active database session.
        file_record: A ``FileRecord`` instance.

    Returns:
        The embedding vector, or ``None`` if the document has no OCR text
        or embedding generation fails.
    """
    # Return cached embedding if available
    if file_record.embedding:
        try:
            return json.loads(file_record.embedding)
        except (json.JSONDecodeError, TypeError):
            logger.warning(f"Invalid cached embedding for file {file_record.id}, recomputing")

    # Need OCR text to generate an embedding
    if not file_record.ocr_text or not file_record.ocr_text.strip():
        return None

    try:
        embedding = generate_embedding(file_record.ocr_text)
        # Cache the embedding in the database
        file_record.embedding = json.dumps(embedding)
        db.commit()
        return embedding
    except Exception as e:
        db.rollback()
        logger.error(f"Failed to generate embedding for file {file_record.id}: {e}")
        return None


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.

    Computes cosine similarity between the target document's embedding and
    all other documents that have OCR text.  Results are sorted by descending
    similarity score.

    Args:
        db: Active database session.
        file_id: The ID of the target ``FileRecord``.
        limit: Maximum number of similar documents to return.
        threshold: Minimum similarity score (0–1) to include in results.

    Returns:
        A list of dicts, each containing:
        - ``file_id``: The similar document's ID.
        - ``original_filename``: The document's original filename.
        - ``document_title``: The document's AI-extracted title (may be None).
        - ``similarity_score``: Cosine similarity (0–1, rounded to 4 decimals).
        - ``mime_type``: The document's MIME type.
        - ``created_at``: ISO-formatted creation timestamp.
    """
    from app.models import FileRecord

    # Get the target document
    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)
    if not target_embedding:
        return []

    # Get candidate documents (those with OCR text, excluding the target)
    candidates = (
        db.query(FileRecord)
        .filter(
            FileRecord.id != file_id,
            FileRecord.ocr_text.isnot(None),
            FileRecord.ocr_text != "",
        )
        .all()
    )

    results = []
    for candidate in candidates:
        candidate_embedding = _get_or_compute_embedding(db, candidate)
        if not candidate_embedding:
            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,
                    "similarity_score": round(score, 4),
                    "mime_type": candidate.mime_type,
                    "created_at": candidate.created_at.isoformat() if candidate.created_at else None,
                }
            )

    # Sort by similarity score descending
    results.sort(key=lambda x: x["similarity_score"], reverse=True)
    return results[:limit]