mule-image/backend/app/services/search.py

"""
Unified search service — hybrid FTS + semantic (RRF) search.

Phase 1 (PR4): semantic-only via pgvector cosine similarity.
Phase 2 (PR5): adds FTS via tsvector, enables RRF fusion.
"""
import logging
from typing import Optional

import numpy as np
from sqlalchemy import select, text, func
from sqlalchemy.ext.asyncio import AsyncSession

from app.models import Photo
from app.models.embeddings import Embedding
from app.config import settings

logger = logging.getLogger(__name__)


async def hybrid_search(
    db: AsyncSession,
    q: Optional[str] = None,
    tag_ids: Optional[list[str]] = None,
    date_from: Optional[str] = None,
    date_to: Optional[str] = None,
    limit: int = 50,
    offset: int = 0,
) -> list[dict]:
    """Run hybrid search (FTS + semantic) with RRF fusion.

    Currently semantic-only; FTS leg added in PR5.
    """
    model_name = settings.vision.embedder.name
    results = {}

    # ── Semantic search (CLIP text → pgvector cosine) ─────────────────
    if q:
        try:
            from app.services.vision.registry import registry
            embedder = registry.get_embedder()
            query_vec = embedder.embed_text(q)

            # pgvector cosine distance: <=> returns distance (lower = closer).
            # Join photos so we can filter out discarded / hidden rows
            # inside the same query — otherwise a hidden-folder photo
            # can take a top-N rank and starve the visible results.
            vec_str = "[" + ",".join(str(float(v)) for v in query_vec) + "]"
            stmt = text("""
                SELECT e.photo_id,
                       (e.vector <=> :qvec::vector) AS distance
                FROM embeddings e
                JOIN photos p ON p.id = e.photo_id
                WHERE e.model = :model
                  AND p.is_trashed = false
                  AND p.is_hidden = false
                ORDER BY e.vector <=> :qvec::vector
                LIMIT 200
            """)
            rows = (await db.execute(stmt, {"qvec": vec_str, "model": model_name})).fetchall()

            for rank, (photo_id, distance) in enumerate(rows):
                if photo_id not in results:
                    results[photo_id] = {"semantic_rank": rank, "fts_rank": None}
                else:
                    results[photo_id]["semantic_rank"] = rank

        except Exception as e:
            logger.warning("Semantic search failed (models may not be loaded): %s", e)

    # ── FTS search (photos.search_vector + ocr_text) ────────────────
    if q:
        try:
            # Same discarded/hidden filter as the semantic leg.
            # The OCR branch joins photos (through photo_id) so we can
            # filter there too; otherwise OCR hits in hidden folders
            # would leak into results.
            fts_stmt = text("""
                SELECT id, ts_rank(search_vector, plainto_tsquery('english', :q)) AS rank
                FROM photos
                WHERE search_vector @@ plainto_tsquery('english', :q)
                  AND is_trashed = false
                  AND is_hidden = false
                UNION
                SELECT o.photo_id AS id,
                       MAX(o.confidence) AS rank
                FROM ocr_text o
                JOIN photos p ON p.id = o.photo_id
                WHERE to_tsvector('english', o.text) @@ plainto_tsquery('english', :q)
                  AND p.is_trashed = false
                  AND p.is_hidden = false
                GROUP BY o.photo_id
                ORDER BY rank DESC
                LIMIT 200
            """)
            fts_rows = (await db.execute(fts_stmt, {"q": q})).fetchall()
            for rank, (photo_id, score) in enumerate(fts_rows):
                if photo_id not in results:
                    results[photo_id] = {"semantic_rank": None, "fts_rank": rank}
                else:
                    results[photo_id]["fts_rank"] = rank
        except Exception as e:
            logger.warning("FTS search failed: %s", e)

    # ── RRF fusion ────────────────────────────────────────────────────
    k = 60
    scored = []
    for photo_id, ranks in results.items():
        score = 0.0
        if ranks["semantic_rank"] is not None:
            score += 1.0 / (k + ranks["semantic_rank"])
        if ranks.get("fts_rank") is not None:
            score += 1.0 / (k + ranks["fts_rank"])
        scored.append((photo_id, score))

    scored.sort(key=lambda x: -x[1])

    # If no text query, fall back to recent photos. Always filter out
    # discarded + hidden here — this path backs the "Tags" and "People"
    # browse views, which should honor the folder hide flag.
    if not q:
        if tag_ids:
            from app.models.tags import photo_tags
            # Subquery to get distinct photo_ids matching the tag filter
            sub = select(photo_tags.c.photo_id).where(
                photo_tags.c.tag_id.in_(tag_ids)
            ).distinct().subquery()
            stmt = select(Photo.id).join(sub, Photo.id == sub.c.photo_id)
        else:
            stmt = select(Photo.id)
        stmt = stmt.where(
            Photo.is_discarded.is_(False),
            Photo.is_hidden.is_(False),
        )
        stmt = stmt.order_by(Photo.added_at.desc())
        if date_from:
            stmt = stmt.where(Photo.taken_at >= date_from)
        if date_to:
            stmt = stmt.where(Photo.taken_at <= date_to)
        stmt = stmt.offset(offset).limit(limit)
        rows = (await db.execute(stmt)).fetchall()
        return [{"photo_id": row[0], "score": 0.0} for row in rows]

    # Apply filters to scored results
    photo_ids = [pid for pid, _ in scored]
    if not photo_ids:
        return []

    # Filter by tags if requested
    if tag_ids:
        from app.models.tags import photo_tags
        stmt = select(photo_tags.c.photo_id).where(
            photo_tags.c.photo_id.in_(photo_ids),
            photo_tags.c.tag_id.in_(tag_ids),
        ).distinct()
        valid_ids = {row[0] for row in (await db.execute(stmt)).fetchall()}
        scored = [(pid, s) for pid, s in scored if pid in valid_ids]

    # Paginate
    page = scored[offset : offset + limit]
    return [{"photo_id": pid, "score": score} for pid, score in page]