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) <noreply@anthropic.com>
This commit is contained in:
2026-04-10 09:07:32 +02:00
parent b1c2bdf7f0
commit 649437dc85
8 changed files with 418 additions and 11 deletions

View File

@@ -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)
)
""")

View File

@@ -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():

View File

@@ -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
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())

View File

@@ -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)}

View File

@@ -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]

View File

@@ -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',

View File

@@ -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:

148
backend/app/tasks/vision.py Normal file
View File

@@ -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}