feat(similarity): add document similarity detection with embeddings and cosine similarity
Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
@@ -19,6 +19,7 @@ from app.api.queue import router as queue_router
|
||||
from app.api.saved_searches import router as saved_searches_router
|
||||
from app.api.search import router as search_router
|
||||
from app.api.settings import router as settings_router
|
||||
from app.api.similarity import router as similarity_router
|
||||
from app.api.url_upload import router as url_upload_router
|
||||
|
||||
# Import all the individual routers
|
||||
@@ -46,3 +47,4 @@ router.include_router(url_upload_router)
|
||||
router.include_router(search_router)
|
||||
router.include_router(queue_router)
|
||||
router.include_router(saved_searches_router)
|
||||
router.include_router(similarity_router)
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Document similarity API endpoints.
|
||||
|
||||
Provides an endpoint to find documents similar to a given file based on
|
||||
text embeddings and cosine similarity scoring.
|
||||
"""
|
||||
|
||||
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.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",
|
||||
}
|
||||
|
||||
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",
|
||||
)
|
||||
@@ -67,6 +67,9 @@ class FileRecord(Base):
|
||||
# Human-readable document title from AI metadata
|
||||
document_title = Column(String, nullable=True)
|
||||
|
||||
# Pre-computed text embedding vector stored as JSON array of floats
|
||||
embedding = Column(Text, nullable=True)
|
||||
|
||||
# Timestamp when we inserted this record
|
||||
created_at = Column(DateTime(timezone=True), server_default=func.now(), index=True)
|
||||
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
"""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]
|
||||
@@ -951,7 +951,76 @@
|
||||
{% if processed_file_exists %}
|
||||
loadPDF('processed', fileId);
|
||||
{% endif %}
|
||||
|
||||
// Load similar documents
|
||||
loadSimilarDocuments(fileId);
|
||||
});
|
||||
|
||||
// Similar documents loading
|
||||
async function loadSimilarDocuments(fileId) {
|
||||
const loadingDiv = document.getElementById('similar-documents-loading');
|
||||
const contentDiv = document.getElementById('similar-documents-content');
|
||||
const emptyDiv = document.getElementById('similar-documents-empty');
|
||||
const errorDiv = document.getElementById('similar-documents-error');
|
||||
|
||||
try {
|
||||
const response = await fetch(`/api/files/${fileId}/similar?limit=5&threshold=0.3`);
|
||||
if (!response.ok) {
|
||||
throw new Error(response.statusText || 'Request failed');
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
loadingDiv.style.display = 'none';
|
||||
|
||||
if (!data.similar_documents || data.similar_documents.length === 0) {
|
||||
emptyDiv.style.display = 'block';
|
||||
if (data.message) {
|
||||
emptyDiv.querySelector('p').textContent = data.message;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Build the results HTML
|
||||
let html = '<div style="display: grid; gap: 0.75rem;">';
|
||||
for (const doc of data.similar_documents) {
|
||||
const scorePercent = Math.round(doc.similarity_score * 100);
|
||||
const title = doc.document_title || doc.original_filename || 'Untitled';
|
||||
const filename = doc.original_filename || 'Unknown';
|
||||
const createdAt = doc.created_at ? new Date(doc.created_at).toLocaleDateString() : '';
|
||||
|
||||
html += `
|
||||
<a href="/files/${doc.file_id}/detail" style="text-decoration: none; color: inherit;">
|
||||
<div style="display: flex; align-items: center; gap: 1rem; padding: 0.75rem 1rem; background-color: #f7fafc; border-radius: 0.5rem; border: 1px solid #e2e8f0; transition: border-color 0.2s;">
|
||||
<div style="flex-shrink: 0; width: 48px; height: 48px; border-radius: 50%; display: flex; align-items: center; justify-content: center; font-weight: 700; font-size: 0.875rem; color: white; background-color: ${scorePercent >= 80 ? '#48bb78' : scorePercent >= 50 ? '#ecc94b' : '#a0aec0'};">
|
||||
${scorePercent}%
|
||||
</div>
|
||||
<div style="flex: 1; min-width: 0;">
|
||||
<div style="font-weight: 600; color: #2d3748; white-space: nowrap; overflow: hidden; text-overflow: ellipsis;" title="${title}">
|
||||
${title}
|
||||
</div>
|
||||
<div style="font-size: 0.75rem; color: #718096; white-space: nowrap; overflow: hidden; text-overflow: ellipsis;">
|
||||
${filename}${createdAt ? ' · ' + createdAt : ''}
|
||||
</div>
|
||||
</div>
|
||||
<div style="flex-shrink: 0; color: #a0aec0;">
|
||||
<i class="fas fa-chevron-right" aria-hidden="true"></i>
|
||||
</div>
|
||||
</div>
|
||||
</a>
|
||||
`;
|
||||
}
|
||||
html += '</div>';
|
||||
|
||||
contentDiv.innerHTML = html;
|
||||
contentDiv.style.display = 'block';
|
||||
} catch (error) {
|
||||
console.error('Error loading similar documents:', error);
|
||||
loadingDiv.style.display = 'none';
|
||||
errorDiv.style.display = 'block';
|
||||
document.getElementById('similar-documents-error-msg').textContent =
|
||||
'Failed to load similar documents: ' + error.message;
|
||||
}
|
||||
}
|
||||
</script>
|
||||
{% endif %}
|
||||
{% endblock %}
|
||||
@@ -1527,6 +1596,24 @@
|
||||
</div>
|
||||
{% endif %}
|
||||
|
||||
<!-- Similar Documents Card -->
|
||||
<div class="detail-card" id="similar-documents-card">
|
||||
<h3><i class="fas fa-copy" aria-hidden="true"></i> Similar Documents</h3>
|
||||
<div id="similar-documents-loading" style="text-align: center; padding: 2rem; color: #718096;">
|
||||
<i class="fas fa-spinner fa-spin" aria-hidden="true" style="font-size: 1.5rem; margin-bottom: 0.5rem;"></i>
|
||||
<p>Searching for similar documents…</p>
|
||||
</div>
|
||||
<div id="similar-documents-content" style="display: none;"></div>
|
||||
<div id="similar-documents-empty" style="display: none; text-align: center; padding: 2rem; color: #718096;">
|
||||
<i class="fas fa-search" aria-hidden="true" style="font-size: 2rem; margin-bottom: 0.5rem; opacity: 0.5;"></i>
|
||||
<p>No similar documents found.</p>
|
||||
</div>
|
||||
<div id="similar-documents-error" style="display: none; text-align: center; padding: 2rem; color: #991B1B;">
|
||||
<i class="fas fa-exclamation-triangle" aria-hidden="true" style="font-size: 2rem; margin-bottom: 0.5rem;"></i>
|
||||
<p id="similar-documents-error-msg">Failed to load similar documents.</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- File Preview Card -->
|
||||
{% if original_file_exists or processed_file_exists %}
|
||||
<div class="detail-card">
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Add embedding column to files table for document similarity
|
||||
|
||||
Revision ID: 009_add_embedding_column
|
||||
Revises: 008_add_performance_indexes
|
||||
Create Date: 2026-03-01
|
||||
|
||||
"""
|
||||
|
||||
from typing import Union
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "009_add_embedding_column"
|
||||
down_revision: Union[str, None] = "008_add_performance_indexes"
|
||||
depends_on: Union[str, None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Add embedding column to files table for storing text embedding vectors."""
|
||||
op.add_column("files", sa.Column("embedding", sa.Text(), nullable=True))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Remove embedding column from files table."""
|
||||
op.drop_column("files", "embedding")
|
||||
@@ -0,0 +1,445 @@
|
||||
"""Tests for document similarity detection.
|
||||
|
||||
Tests the similarity utility functions and the API endpoint
|
||||
``GET /api/files/{file_id}/similar``.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.models import FileRecord
|
||||
from app.utils.similarity import cosine_similarity, find_similar_documents
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for cosine_similarity
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCosineSimilarity:
|
||||
"""Unit tests for the cosine_similarity function."""
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_identical_vectors_return_one(self):
|
||||
"""Identical vectors should have similarity of 1.0."""
|
||||
vec = [1.0, 2.0, 3.0]
|
||||
assert cosine_similarity(vec, vec) == pytest.approx(1.0)
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_orthogonal_vectors_return_zero(self):
|
||||
"""Orthogonal vectors should have similarity of 0.0."""
|
||||
a = [1.0, 0.0]
|
||||
b = [0.0, 1.0]
|
||||
assert cosine_similarity(a, b) == pytest.approx(0.0)
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_opposite_vectors_clamped_to_zero(self):
|
||||
"""Opposite vectors would give negative cosine; clamp to 0."""
|
||||
a = [1.0, 0.0]
|
||||
b = [-1.0, 0.0]
|
||||
assert cosine_similarity(a, b) == 0.0
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_different_length_vectors_return_zero(self):
|
||||
"""Vectors of different lengths should return 0.0."""
|
||||
a = [1.0, 2.0, 3.0]
|
||||
b = [1.0, 2.0]
|
||||
assert cosine_similarity(a, b) == 0.0
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_zero_vector_returns_zero(self):
|
||||
"""Zero-magnitude vector should return 0.0."""
|
||||
a = [0.0, 0.0, 0.0]
|
||||
b = [1.0, 2.0, 3.0]
|
||||
assert cosine_similarity(a, b) == 0.0
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_similar_vectors_high_score(self):
|
||||
"""Similar (but not identical) vectors should have a high score."""
|
||||
a = [1.0, 2.0, 3.0]
|
||||
b = [1.1, 2.1, 3.1]
|
||||
score = cosine_similarity(a, b)
|
||||
assert 0.99 < score <= 1.0
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_score_between_zero_and_one(self):
|
||||
"""All scores should be in [0, 1]."""
|
||||
a = [1.0, 0.5, 0.0]
|
||||
b = [0.0, 0.5, 1.0]
|
||||
score = cosine_similarity(a, b)
|
||||
assert 0.0 <= score <= 1.0
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_empty_vectors_return_zero(self):
|
||||
"""Empty vectors should return 0.0."""
|
||||
assert cosine_similarity([], []) == 0.0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for find_similar_documents
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFindSimilarDocuments:
|
||||
"""Unit tests for the find_similar_documents function."""
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_returns_empty_for_missing_file(self, db_session):
|
||||
"""Should return empty list when file ID does not exist."""
|
||||
result = find_similar_documents(db_session, file_id=9999)
|
||||
assert result == []
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_returns_empty_when_no_ocr_text(self, db_session):
|
||||
"""Should return empty list when target file has no OCR text."""
|
||||
file_record = FileRecord(
|
||||
filehash="abc123",
|
||||
local_filename="/tmp/test.pdf",
|
||||
file_size=1024,
|
||||
original_filename="test.pdf",
|
||||
ocr_text=None,
|
||||
)
|
||||
db_session.add(file_record)
|
||||
db_session.commit()
|
||||
|
||||
result = find_similar_documents(db_session, file_id=file_record.id)
|
||||
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
|
||||
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",
|
||||
)
|
||||
# Create a similar file
|
||||
similar = FileRecord(
|
||||
filehash="hash2",
|
||||
local_filename="/tmp/similar.pdf",
|
||||
file_size=2048,
|
||||
original_filename="similar.pdf",
|
||||
ocr_text="This is an invoice from Amazon for February 2026",
|
||||
document_title="Amazon Invoice Feb",
|
||||
mime_type="application/pdf",
|
||||
)
|
||||
# Create a different file
|
||||
different = FileRecord(
|
||||
filehash="hash3",
|
||||
local_filename="/tmp/different.pdf",
|
||||
file_size=512,
|
||||
original_filename="different.pdf",
|
||||
ocr_text="Recipe for chocolate cake with detailed instructions",
|
||||
document_title="Chocolate Cake Recipe",
|
||||
mime_type="application/pdf",
|
||||
)
|
||||
|
||||
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
|
||||
assert result[0]["file_id"] == similar.id
|
||||
assert result[0]["similarity_score"] > 0.9
|
||||
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):
|
||||
"""Should filter out documents below the threshold."""
|
||||
target = FileRecord(
|
||||
filehash="hash1",
|
||||
local_filename="/tmp/t.pdf",
|
||||
file_size=100,
|
||||
original_filename="target.pdf",
|
||||
ocr_text="target text",
|
||||
)
|
||||
candidate = FileRecord(
|
||||
filehash="hash2",
|
||||
local_filename="/tmp/c.pdf",
|
||||
file_size=100,
|
||||
original_filename="candidate.pdf",
|
||||
ocr_text="different text",
|
||||
)
|
||||
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):
|
||||
"""Should respect the limit parameter."""
|
||||
target = FileRecord(
|
||||
filehash="hash0",
|
||||
local_filename="/tmp/t.pdf",
|
||||
file_size=100,
|
||||
original_filename="target.pdf",
|
||||
ocr_text="target text",
|
||||
)
|
||||
db_session.add(target)
|
||||
|
||||
for i in range(5):
|
||||
f = FileRecord(
|
||||
filehash=f"hash{i + 1}",
|
||||
local_filename=f"/tmp/c{i}.pdf",
|
||||
file_size=100,
|
||||
original_filename=f"candidate_{i}.pdf",
|
||||
ocr_text=f"similar text {i}",
|
||||
)
|
||||
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
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_uses_cached_embedding(self, db_session):
|
||||
"""Should use cached embeddings from the database."""
|
||||
cached_embedding = [0.5, 0.5, 0.5]
|
||||
|
||||
target = FileRecord(
|
||||
filehash="hash1",
|
||||
local_filename="/tmp/t.pdf",
|
||||
file_size=100,
|
||||
original_filename="target.pdf",
|
||||
ocr_text="some text",
|
||||
embedding=json.dumps(cached_embedding),
|
||||
)
|
||||
candidate = FileRecord(
|
||||
filehash="hash2",
|
||||
local_filename="/tmp/c.pdf",
|
||||
file_size=100,
|
||||
original_filename="candidate.pdf",
|
||||
ocr_text="some text too",
|
||||
embedding=json.dumps(cached_embedding),
|
||||
)
|
||||
db_session.add_all([target, candidate])
|
||||
db_session.commit()
|
||||
|
||||
# No mock needed — cached embeddings should be used
|
||||
result = find_similar_documents(db_session, file_id=target.id, threshold=0.0)
|
||||
assert len(result) == 1
|
||||
assert result[0]["similarity_score"] == pytest.approx(1.0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration tests for the API endpoint
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSimilarDocumentsAPI:
|
||||
"""Integration tests for GET /api/files/{file_id}/similar."""
|
||||
|
||||
@pytest.mark.integration
|
||||
def test_file_not_found(self, client: TestClient):
|
||||
"""Should return 404 for non-existent file."""
|
||||
response = client.get("/api/files/9999/similar")
|
||||
assert response.status_code == 404
|
||||
|
||||
@pytest.mark.integration
|
||||
def test_no_ocr_text(self, client: TestClient, db_session):
|
||||
"""Should return empty results when file has no OCR text."""
|
||||
file_record = FileRecord(
|
||||
filehash="abc123",
|
||||
local_filename="/tmp/test.pdf",
|
||||
file_size=1024,
|
||||
original_filename="test.pdf",
|
||||
ocr_text=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 data["similar_documents"] == []
|
||||
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):
|
||||
"""Should return similar documents with scores."""
|
||||
target = FileRecord(
|
||||
filehash="hash1",
|
||||
local_filename="/tmp/target.pdf",
|
||||
file_size=1024,
|
||||
original_filename="target.pdf",
|
||||
ocr_text="Invoice from Amazon January 2026",
|
||||
)
|
||||
similar = FileRecord(
|
||||
filehash="hash2",
|
||||
local_filename="/tmp/similar.pdf",
|
||||
file_size=2048,
|
||||
original_filename="similar_invoice.pdf",
|
||||
ocr_text="Invoice from Amazon February 2026",
|
||||
document_title="Amazon Invoice",
|
||||
mime_type="application/pdf",
|
||||
)
|
||||
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()
|
||||
assert data["file_id"] == target.id
|
||||
assert data["count"] >= 1
|
||||
assert len(data["similar_documents"]) >= 1
|
||||
|
||||
doc = data["similar_documents"][0]
|
||||
assert "file_id" in doc
|
||||
assert "similarity_score" in doc
|
||||
assert 0 <= doc["similarity_score"] <= 1
|
||||
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):
|
||||
"""Should respect limit and threshold query parameters."""
|
||||
target = FileRecord(
|
||||
filehash="hash1",
|
||||
local_filename="/tmp/t.pdf",
|
||||
file_size=100,
|
||||
original_filename="t.pdf",
|
||||
ocr_text="test",
|
||||
)
|
||||
db_session.add(target)
|
||||
|
||||
for i in range(5):
|
||||
f = FileRecord(
|
||||
filehash=f"h{i}",
|
||||
local_filename=f"/tmp/c{i}.pdf",
|
||||
file_size=100,
|
||||
original_filename=f"c{i}.pdf",
|
||||
ocr_text=f"text {i}",
|
||||
)
|
||||
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()
|
||||
assert data["count"] <= 2
|
||||
|
||||
@pytest.mark.integration
|
||||
def test_invalid_limit(self, client: TestClient, db_session):
|
||||
"""Should reject invalid limit values."""
|
||||
file_record = FileRecord(
|
||||
filehash="abc",
|
||||
local_filename="/tmp/t.pdf",
|
||||
file_size=100,
|
||||
original_filename="t.pdf",
|
||||
)
|
||||
db_session.add(file_record)
|
||||
db_session.commit()
|
||||
|
||||
response = client.get(f"/api/files/{file_record.id}/similar?limit=0")
|
||||
assert response.status_code == 422
|
||||
|
||||
@pytest.mark.integration
|
||||
def test_invalid_threshold(self, client: TestClient, db_session):
|
||||
"""Should reject threshold values outside [0, 1]."""
|
||||
file_record = FileRecord(
|
||||
filehash="abc",
|
||||
local_filename="/tmp/t.pdf",
|
||||
file_size=100,
|
||||
original_filename="t.pdf",
|
||||
)
|
||||
db_session.add(file_record)
|
||||
db_session.commit()
|
||||
|
||||
response = client.get(f"/api/files/{file_record.id}/similar?threshold=1.5")
|
||||
assert response.status_code == 422
|
||||
|
||||
@pytest.mark.integration
|
||||
def test_empty_ocr_text(self, client: TestClient, db_session):
|
||||
"""Should return empty results when OCR text is empty string."""
|
||||
file_record = FileRecord(
|
||||
filehash="abc",
|
||||
local_filename="/tmp/t.pdf",
|
||||
file_size=100,
|
||||
original_filename="t.pdf",
|
||||
ocr_text="",
|
||||
)
|
||||
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
|
||||
|
||||
@pytest.mark.integration
|
||||
@patch("app.utils.similarity.generate_embedding")
|
||||
def test_response_structure(self, mock_embed, client: TestClient, db_session):
|
||||
"""Should return proper response structure for each similar document."""
|
||||
target = FileRecord(
|
||||
filehash="h1",
|
||||
local_filename="/tmp/t.pdf",
|
||||
file_size=100,
|
||||
original_filename="target.pdf",
|
||||
ocr_text="Some text content here",
|
||||
)
|
||||
other = FileRecord(
|
||||
filehash="h2",
|
||||
local_filename="/tmp/o.pdf",
|
||||
file_size=200,
|
||||
original_filename="other.pdf",
|
||||
ocr_text="Some similar text content",
|
||||
document_title="Other Doc",
|
||||
mime_type="application/pdf",
|
||||
)
|
||||
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()
|
||||
|
||||
assert "file_id" in data
|
||||
assert "similar_documents" in data
|
||||
assert "count" in data
|
||||
|
||||
if data["count"] > 0:
|
||||
doc = data["similar_documents"][0]
|
||||
assert "file_id" in doc
|
||||
assert "original_filename" in doc
|
||||
assert "document_title" in doc
|
||||
assert "similarity_score" in doc
|
||||
assert "mime_type" in doc
|
||||
assert "created_at" in doc
|
||||
Reference in New Issue
Block a user