bvandeusen
79cd1234e2
GalleryService.similar() ranks images by pgvector cosine distance to a source image's precomputed SigLIP embedding — no query-time ML inference. Composes with the Phase-1/2 scope filters (AND) but replaces the date sort (always nearest-first, bounded top-N, no cursor). Returns None for a missing source (→404), [] for a source with no embedding (video / pending ML); excludes self and NULL-embedding rows. New GET /api/gallery/similar?similar_to=<id>&limit=N. Image-detail payload gains has_embedding so the UI can hide the surface. Alembic 0036 adds an HNSW vector_cosine_ops index on siglip_embedding (1152<2000 dims) so the search is sub-50ms ANN instead of a full scan; one-time ~30-60s build over existing embeddings on deploy. Shared _gallery_images/_image_json helpers de-dup the scroll/similar builders. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
135 lines
4.4 KiB
Python
135 lines
4.4 KiB
Python
"""Phase-3 visual "more like this" — pgvector cosine ranking over the
|
|
precomputed SigLIP image embeddings. No query-time ML inference."""
|
|
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest
|
|
|
|
from backend.app.models import ImageRecord, Tag, TagKind
|
|
from backend.app.models.tag import image_tag
|
|
from backend.app.services.gallery_service import GalleryService
|
|
|
|
pytestmark = pytest.mark.integration
|
|
|
|
|
|
def _vec(*head):
|
|
"""A 1152-dim embedding with the given leading values, rest zero. Cosine
|
|
distance to _vec(1,0) grows as the 2nd component grows, so callers can
|
|
order fixtures deterministically by direction."""
|
|
v = [0.0] * 1152
|
|
for i, x in enumerate(head):
|
|
v[i] = float(x)
|
|
return v
|
|
|
|
|
|
async def _img(db, n, emb):
|
|
rec = ImageRecord(
|
|
path=f"/images/sim/{n}.jpg", sha256=f"e{n:063d}",
|
|
size_bytes=1, mime="image/jpeg", width=1, height=1,
|
|
origin="imported_filesystem", integrity_status="unknown",
|
|
siglip_embedding=emb,
|
|
)
|
|
base = datetime(2026, 1, 1, 12, 0, tzinfo=UTC)
|
|
rec.created_at = base - timedelta(minutes=n)
|
|
rec.effective_date = rec.created_at
|
|
db.add(rec)
|
|
await db.flush()
|
|
return rec
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_similar_ranks_nearest_first(db):
|
|
src = await _img(db, 1, _vec(1, 0))
|
|
near = await _img(db, 2, _vec(1, 0.05)) # almost same direction
|
|
mid = await _img(db, 3, _vec(1, 1)) # 45°
|
|
far = await _img(db, 4, _vec(0, 1)) # orthogonal
|
|
svc = GalleryService(db)
|
|
res = await svc.similar(src.id, limit=10)
|
|
assert [i.id for i in res] == [near.id, mid.id, far.id] # self excluded
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_similar_excludes_null_embeddings(db):
|
|
src = await _img(db, 1, _vec(1, 0))
|
|
have = await _img(db, 2, _vec(1, 0.1))
|
|
await _img(db, 3, None) # un-embedded (e.g. a video) → excluded
|
|
svc = GalleryService(db)
|
|
res = await svc.similar(src.id, limit=10)
|
|
assert [i.id for i in res] == [have.id]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_similar_source_without_embedding_returns_empty(db):
|
|
src = await _img(db, 1, None)
|
|
await _img(db, 2, _vec(1, 0))
|
|
svc = GalleryService(db)
|
|
assert await svc.similar(src.id, limit=10) == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_similar_missing_source_returns_none(db):
|
|
svc = GalleryService(db)
|
|
assert await svc.similar(99999, limit=10) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_similar_composes_with_tag_filter(db):
|
|
src = await _img(db, 1, _vec(1, 0))
|
|
await _img(db, 2, _vec(1, 0.02)) # nearest, but untagged
|
|
tagged = await _img(db, 3, _vec(1, 0.6)) # farther, but carries the tag
|
|
tag = Tag(name="t", kind=TagKind.general)
|
|
db.add(tag)
|
|
await db.flush()
|
|
await db.execute(image_tag.insert().values(
|
|
image_record_id=tagged.id, tag_id=tag.id, source="manual"))
|
|
svc = GalleryService(db)
|
|
res = await svc.similar(src.id, limit=10, tag_ids=[tag.id])
|
|
assert [i.id for i in res] == [tagged.id] # scope AND-narrows the ranked set
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_similar_respects_limit(db):
|
|
src = await _img(db, 1, _vec(1, 0))
|
|
for n in range(2, 7):
|
|
await _img(db, n, _vec(1, 0.1 * n))
|
|
svc = GalleryService(db)
|
|
res = await svc.similar(src.id, limit=2)
|
|
assert len(res) == 2
|
|
|
|
|
|
# --- API ---
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_similar_endpoint(client, db):
|
|
src = await _img(db, 1, _vec(1, 0))
|
|
near = await _img(db, 2, _vec(1, 0.05))
|
|
await db.commit()
|
|
resp = await client.get(f"/api/gallery/similar?similar_to={src.id}&limit=10")
|
|
assert resp.status_code == 200
|
|
body = await resp.get_json()
|
|
assert [i["id"] for i in body["images"]] == [near.id]
|
|
assert body["next_cursor"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_similar_404_when_source_missing(client):
|
|
resp = await client.get("/api/gallery/similar?similar_to=99999")
|
|
assert resp.status_code == 404
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_similar_requires_param(client):
|
|
resp = await client.get("/api/gallery/similar")
|
|
assert resp.status_code == 400
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_detail_reports_has_embedding(client, db):
|
|
embedded = await _img(db, 1, _vec(1, 0))
|
|
plain = await _img(db, 2, None)
|
|
await db.commit()
|
|
e = await (await client.get(f"/api/gallery/image/{embedded.id}")).get_json()
|
|
p = await (await client.get(f"/api/gallery/image/{plain.id}")).get_json()
|
|
assert e["has_embedding"] is True
|
|
assert p["has_embedding"] is False
|