From 649437dc851431e5eff4f4b2c6af8553c2ec6133 Mon Sep 17 00:00:00 2001 From: dtoro Date: Fri, 10 Apr 2026 09:07:32 +0200 Subject: [PATCH] feat: add embeddings pipeline and semantic search endpoint Wire the full embedding flow: - Rewrite Embedding model to use pgvector Vector(512) with HNSW index - Add embed_photo, vision_fanout, backfill_vision Celery tasks on dedicated `vision` queue - Hook vision_fanout into generate_thumbnails completion - Add POST /api/v1/photos/search with hybrid RRF ranking (semantic-only for now; FTS leg added in PR5) - Stub ocr_photo, detect_objects, extract_faces tasks for later PRs Migration 0003 drops/recreates the embeddings table (was never populated). Co-Authored-By: Claude Opus 4.6 (1M context) --- .../versions/0003_pgvector_embeddings.py | 52 ++++++ backend/app/main.py | 3 +- backend/app/models/embeddings.py | 20 +-- backend/app/routers/search.py | 74 +++++++++ backend/app/services/search.py | 115 ++++++++++++++ backend/app/tasks/celery.py | 8 +- backend/app/tasks/thumbs.py | 9 ++ backend/app/tasks/vision.py | 148 ++++++++++++++++++ 8 files changed, 418 insertions(+), 11 deletions(-) create mode 100644 backend/alembic/versions/0003_pgvector_embeddings.py create mode 100644 backend/app/routers/search.py create mode 100644 backend/app/services/search.py create mode 100644 backend/app/tasks/vision.py diff --git a/backend/alembic/versions/0003_pgvector_embeddings.py b/backend/alembic/versions/0003_pgvector_embeddings.py new file mode 100644 index 0000000..e17a311 --- /dev/null +++ b/backend/alembic/versions/0003_pgvector_embeddings.py @@ -0,0 +1,52 @@ +"""pgvector embeddings + +Revision ID: 0003_pgvector_embeddings +Revises: 0002_extend_tags +Create Date: 2026-04-10 + +Rewrite the embeddings table to use pgvector Vector(512) instead of +LargeBinary. Add composite PK (photo_id, model), created_at, and +HNSW index on vector column. +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + +revision: str = "0003_pgvector_embeddings" +down_revision: Union[str, None] = "0002_extend_tags" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # Drop the old placeholder table and recreate with pgvector types. + # No data to preserve — it was never populated. + op.execute("DROP TABLE IF EXISTS embeddings") + op.execute(""" + CREATE TABLE embeddings ( + photo_id VARCHAR NOT NULL REFERENCES photos(id) ON DELETE CASCADE, + model VARCHAR(64) NOT NULL, + vector vector(512), + created_at TIMESTAMPTZ DEFAULT now(), + PRIMARY KEY (photo_id, model) + ) + """) + # HNSW index for cosine similarity search. + # Defer creation on large backfills — drop and recreate afterward. + op.execute(""" + CREATE INDEX IF NOT EXISTS ix_embeddings_vector_hnsw + ON embeddings USING hnsw (vector vector_cosine_ops) + """) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS embeddings") + op.execute(""" + CREATE TABLE embeddings ( + photo_id VARCHAR NOT NULL REFERENCES photos(id) ON DELETE CASCADE, + model VARCHAR, + vector BYTEA, + PRIMARY KEY (photo_id) + ) + """) diff --git a/backend/app/main.py b/backend/app/main.py index 47b4499..4ceaa66 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -11,7 +11,7 @@ import os from app.config import settings from app.database import init_db -from app.routers import photos, folders, heaps, tags, discard, library +from app.routers import photos, folders, heaps, tags, discard, library, search from app.services.scanner import start_initial_scan, bootstrap_default_source_root from app.services.cleanup import cleanup_data_integrity @@ -91,6 +91,7 @@ app.include_router(heaps.router, prefix="/api/v1/heaps", tags=["heaps"]) app.include_router(tags.router, prefix="/api/v1/tags", tags=["tags"]) app.include_router(discard.router, prefix="/api/v1/discard", tags=["discard"]) app.include_router(library.router, prefix="/api/v1/library", tags=["library"]) +app.include_router(search.router, prefix="/api/v1/photos/search", tags=["search"]) @app.get("/") async def root(): diff --git a/backend/app/models/embeddings.py b/backend/app/models/embeddings.py index 74d7e1b..4368f0e 100644 --- a/backend/app/models/embeddings.py +++ b/backend/app/models/embeddings.py @@ -1,17 +1,19 @@ """ -Embedding model definition (placeholder for AI features) +Embedding model — stores CLIP/SigLIP image embeddings via pgvector. + +Composite PK (photo_id, model) allows re-embedding with newer models +without clobbering old vectors. """ -from sqlalchemy import Column, String, ForeignKey, LargeBinary -import uuid +from sqlalchemy import Column, String, ForeignKey, DateTime, func +from pgvector.sqlalchemy import Vector from app.database import Base + class Embedding(Base): - """ - Placeholder table for future AI embeddings (CLIP, face recognition, etc.) - """ __tablename__ = 'embeddings' - + photo_id = Column(String, ForeignKey('photos.id', ondelete='CASCADE'), primary_key=True) - model = Column(String) # e.g., 'clip-vit-b32', 'face-recognition', etc. - vector = Column(LargeBinary) # raw float32 bytes for embedding vector \ No newline at end of file + model = Column(String(64), primary_key=True) # e.g. 'openclip_vitb32' + vector = Column(Vector(512)) # OpenCLIP ViT-B/32 → 512-d + created_at = Column(DateTime(timezone=True), server_default=func.now()) diff --git a/backend/app/routers/search.py b/backend/app/routers/search.py new file mode 100644 index 0000000..3954fed --- /dev/null +++ b/backend/app/routers/search.py @@ -0,0 +1,74 @@ +""" +Search API router — unified hybrid search endpoint. +""" +from typing import Optional +from fastapi import APIRouter, Depends +from pydantic import BaseModel +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy import select + +from app.database import get_db +from app.models import Photo +from app.services.search import hybrid_search + +router = APIRouter() + + +class SearchRequest(BaseModel): + q: Optional[str] = None + filters: Optional[dict] = None + limit: int = 50 + offset: int = 0 + + +@router.post("") +async def search_photos(body: SearchRequest, db: AsyncSession = Depends(get_db)): + """Unified search endpoint. Every query runs hybrid (FTS + semantic) + by default — the user never picks a mode. + + Filters: + - tag_ids: list of tag IDs (any kind: user, object, face_cluster) + - date_from / date_to: ISO date strings + """ + filters = body.filters or {} + + results = await hybrid_search( + db=db, + q=body.q, + tag_ids=filters.get("tag_ids"), + date_from=filters.get("date_from"), + date_to=filters.get("date_to"), + limit=body.limit, + offset=body.offset, + ) + + if not results: + return {"results": [], "total": 0} + + # Hydrate with photo data + photo_ids = [r["photo_id"] for r in results] + stmt = select(Photo).where(Photo.id.in_(photo_ids)) + rows = (await db.execute(stmt)).scalars().all() + photo_map = {p.id: p for p in rows} + + hydrated = [] + for r in results: + photo = photo_map.get(r["photo_id"]) + if not photo: + continue + hydrated.append({ + "id": photo.id, + "filename": photo.filename, + "filepath": photo.filepath, + "media_type": photo.media_type, + "width": photo.width, + "height": photo.height, + "taken_at": photo.taken_at.isoformat() if photo.taken_at else None, + "rating": photo.rating, + "color_label": photo.color_label, + "thumb_small": photo.thumb_small, + "thumb_medium": photo.thumb_medium, + "score": r["score"], + }) + + return {"results": hydrated, "total": len(hydrated)} diff --git a/backend/app/services/search.py b/backend/app/services/search.py new file mode 100644 index 0000000..1f7ef1c --- /dev/null +++ b/backend/app/services/search.py @@ -0,0 +1,115 @@ +""" +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) + 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 + WHERE e.model = :model + 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 (placeholder for PR5) ────────────────────────────── + # Will be: tsvector @@ plainto_tsquery(q), ranked by ts_rank. + # For now, skip. + + # ── 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 + if not q: + stmt = select(Photo.id).order_by(Photo.created_at.desc()) + if tag_ids: + from app.models.tags import photo_tags + stmt = stmt.join(photo_tags, Photo.id == photo_tags.c.photo_id).where( + photo_tags.c.tag_id.in_(tag_ids) + ).distinct() + 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] diff --git a/backend/app/tasks/celery.py b/backend/app/tasks/celery.py index 52a1798..33df67a 100644 --- a/backend/app/tasks/celery.py +++ b/backend/app/tasks/celery.py @@ -9,7 +9,7 @@ celery_app = Celery( 'mulita', broker=settings.celery_broker_url, backend=settings.celery_result_backend, - include=['app.tasks.scan', 'app.tasks.thumbs'] + include=['app.tasks.scan', 'app.tasks.thumbs', 'app.tasks.vision'] ) # Configure Celery @@ -22,6 +22,12 @@ celery_app.conf.update( task_routes={ 'app.tasks.thumbs.*': {'queue': 'high'}, 'app.tasks.scan.*': {'queue': 'low'}, + 'app.tasks.vision.*': {'queue': 'vision'}, + 'embed_photo': {'queue': 'vision'}, + 'ocr_photo': {'queue': 'vision'}, + 'detect_objects': {'queue': 'vision'}, + 'extract_faces': {'queue': 'vision'}, + 'vision_fanout': {'queue': 'vision'}, }, task_default_queue='default', task_default_exchange='default', diff --git a/backend/app/tasks/thumbs.py b/backend/app/tasks/thumbs.py index a25bf7e..41c8961 100644 --- a/backend/app/tasks/thumbs.py +++ b/backend/app/tasks/thumbs.py @@ -317,6 +317,15 @@ async def _generate_thumbnails_async(photo_id: str, task): await session.commit() logger.info(f"Thumbnails generated for photo {photo_id}") + + # Dispatch vision pipeline (embedding, OCR, detection, faces) + # after thumbs are ready so vision tasks have images to read. + try: + from app.tasks.vision import vision_fanout + vision_fanout.delay(photo_id) + except Exception as e: + logger.warning(f"Could not dispatch vision_fanout for {photo_id}: {e}") + return {'status': 'success', 'photo_id': photo_id} except Exception as e: diff --git a/backend/app/tasks/vision.py b/backend/app/tasks/vision.py new file mode 100644 index 0000000..aa1d09b --- /dev/null +++ b/backend/app/tasks/vision.py @@ -0,0 +1,148 @@ +""" +Celery tasks for the vision pipeline — embedding, OCR, object detection, +face recognition. + +All tasks run on the dedicated `vision` queue with limited concurrency +(memory-bound CPU inference). They read thumbnails generated by +generate_thumbnails, so they MUST run after thumbs complete. +""" +import asyncio +import logging +from pathlib import Path + +import numpy as np +from celery import shared_task +from sqlalchemy import select, delete, text +from PIL import Image + +from app.database import AsyncSessionLocal +from app.models import Photo +from app.models.embeddings import Embedding +from app.config import settings + +logger = logging.getLogger(__name__) + + +def _load_thumb(photo_id: str, size: str = "medium") -> np.ndarray | None: + """Load a thumbnail as an RGB numpy array.""" + thumb_path = Path(f"/data/thumbs/{photo_id}/{size}.webp") + if not thumb_path.exists(): + logger.warning("Thumbnail not found: %s", thumb_path) + return None + img = Image.open(thumb_path).convert("RGB") + return np.array(img) + + +@shared_task(name='embed_photo', queue='vision') +def embed_photo(photo_id: str): + """Generate CLIP embedding for a photo and store in pgvector.""" + if not settings.vision.enabled: + return {'status': 'skipped', 'reason': 'vision disabled'} + return asyncio.run(_embed_photo_async(photo_id)) + + +async def _embed_photo_async(photo_id: str): + image = _load_thumb(photo_id, "medium") # 640px + if image is None: + return {'status': 'error', 'message': 'thumbnail not found'} + + from app.services.vision.registry import registry + embedder = registry.get_embedder() + vector = embedder.embed_image(image) + + model_name = settings.vision.embedder.name + + async with AsyncSessionLocal() as session: + # Upsert: delete existing then insert + await session.execute( + delete(Embedding).where( + Embedding.photo_id == photo_id, + Embedding.model == model_name, + ) + ) + emb = Embedding( + photo_id=photo_id, + model=model_name, + vector=vector.tolist(), + ) + session.add(emb) + await session.commit() + + logger.info("Embedded photo %s with %s", photo_id, model_name) + return {'status': 'success', 'photo_id': photo_id} + + +@shared_task(name='vision_fanout', queue='vision') +def vision_fanout(photo_id: str): + """Dispatch all enabled vision tasks for a photo.""" + if not settings.vision.enabled: + return {'status': 'skipped', 'reason': 'vision disabled'} + + embed_photo.delay(photo_id) + + if settings.vision.ocr.enabled: + ocr_photo.delay(photo_id) + if settings.vision.detector.enabled: + detect_objects.delay(photo_id) + if settings.vision.faces.enabled: + extract_faces.delay(photo_id) + + return {'status': 'dispatched', 'photo_id': photo_id} + + +@shared_task(name='ocr_photo', queue='vision') +def ocr_photo(photo_id: str): + """OCR a photo — implemented in PR5.""" + return {'status': 'not_implemented'} + + +@shared_task(name='detect_objects', queue='vision') +def detect_objects(photo_id: str): + """Detect objects in a photo — implemented in PR6.""" + return {'status': 'not_implemented'} + + +@shared_task(name='extract_faces', queue='vision') +def extract_faces(photo_id: str): + """Detect faces and extract embeddings — implemented in PR7.""" + return {'status': 'not_implemented'} + + +@shared_task(name='backfill_vision') +def backfill_vision(task: str | None = None, limit: int | None = None): + """Queue vision tasks for photos that haven't been processed yet.""" + return asyncio.run(_backfill_vision_async(task, limit)) + + +async def _backfill_vision_async(task: str | None, limit: int | None): + model_name = settings.vision.embedder.name + + async with AsyncSessionLocal() as session: + # Find photos without embeddings + stmt = text(""" + SELECT p.id FROM photos p + LEFT JOIN embeddings e ON e.photo_id = p.id AND e.model = :model + WHERE e.photo_id IS NULL + AND p.processing_status = 'completed' + ORDER BY p.created_at DESC + """) + if limit: + stmt = text(str(stmt) + f" LIMIT {limit}") + + result = await session.execute(stmt, {"model": model_name}) + photo_ids = [row[0] for row in result.fetchall()] + + count = 0 + for pid in photo_ids: + if task == 'embed' or task is None: + embed_photo.delay(pid) + if task == 'ocr' or task is None: + ocr_photo.delay(pid) + if task == 'detect' or task is None: + detect_objects.delay(pid) + if task == 'faces' or task is None: + extract_faces.delay(pid) + count += 1 + + logger.info("Backfill queued %d photos for vision processing", count) + return {'status': 'queued', 'count': count}