feat(similarity): add precomputed similarity cache, fix fusion N+1, add 429 retry

- Add similarity/ package (compute.py, lookup.py) with numpy-based
  pairwise cosine similarity and cached lookup
- database.py: create embeddings + similarity_cache tables in _init_database(),
  add store_similarity_batch/get_cached_similarities/clear_similarity_cache helpers
- pipeline/fusion.py: replace N+1 per-motion embedding SELECT with single
  bulk JOIN using DuckDB QUALIFY window function
- ai_provider.py: retry HTTP 429 with Retry-After header support
- migrations/2026-03-22-add-similarity-cache.sql: make executable
- Add tests for similarity compute, db helpers, and 429 retry (34 pass, 2 skip)
This commit is contained in:
2026-03-22 03:02:25 +01:00
parent a248807e03
commit a78bee9b0a
11 changed files with 879 additions and 38 deletions
+38
View File
@@ -0,0 +1,38 @@
import os
import time
import ai_provider
class DummyResponse:
def __init__(self, status_code=200, json_data=None, headers=None):
self.status_code = status_code
self._json = json_data or {}
self.headers = headers or {}
def json(self):
return self._json
def test_retry_on_429_then_success(monkeypatch):
calls = {"n": 0}
def fake_post(url, json, headers, timeout):
calls["n"] += 1
if calls["n"] <= 2:
# first two calls return 429 with Retry-After: 1
return DummyResponse(
429, json_data={"error": "rate_limited"}, headers={"Retry-After": "1"}
)
return DummyResponse(200, json_data={"data": [{"embedding": [0.4, 0.5]}]})
monkeypatch.setenv("OPENROUTER_API_KEY", "sk-test")
monkeypatch.setattr("requests.post", fake_post)
start = time.time()
emb = ai_provider.get_embedding("hello")
duration = time.time() - start
# we should have waited at least ~2 seconds due to two Retry-After: 1 sleeps
assert duration >= 2
assert emb == [0.4, 0.5]
+3 -1
View File
@@ -1,8 +1,10 @@
import json
import duckdb
import pytest
# duckdb is optional for test runs; skip test if not available
duckdb = pytest.importorskip("duckdb")
from database import MotionDatabase
+65
View File
@@ -0,0 +1,65 @@
def test_similarity_compute_and_lookup(tmp_path):
import pytest
duckdb = pytest.importorskip("duckdb")
# local duckdb imported above
from database import MotionDatabase
import similarity.compute as compute
import similarity.lookup as lookup
db_path = str(tmp_path / "motions.db")
# Build MotionDatabase on tmp_path
db = MotionDatabase(db_path=db_path)
# Insert three motions directly (avoid insert_motion which expects migration-added columns)
conn = duckdb.connect(db_path)
motion_ids = []
for i in range(1, 4):
conn.execute(
"INSERT INTO motions (title, url) VALUES (?, ?)",
(f"motion {i}", f"http://example/{i}"),
)
row = conn.execute(
"SELECT id FROM motions WHERE url = ?", (f"http://example/{i}",)
).fetchone()
assert row is not None
motion_ids.append(row[0])
conn.close()
# Insert fused_embeddings for window 'W1'
vectors = [[1, 0, 0], [0, 1, 0], [1, 1, 0]]
for motion_id, vec in zip(motion_ids, vectors):
rid = db.store_fused_embedding(
motion_id=motion_id, window_id="W1", vector=vec, svd_dims=1, text_dims=2
)
assert rid != -1
# Compute similarities
inserted = compute.compute_similarities(
vector_type="fused", window_id="W1", top_k=1, db_path=db_path
)
# depending on implementation we may insert 2 or 3 rows (or more); allow 2 or 3
assert inserted in (2, 3)
# Lookup neighbors for motion 1
neighbors = lookup.get_similar_motions(
motion_id=motion_ids[0],
vector_type="fused",
window_id="W1",
top_k=2,
db_path=db_path,
)
assert len(neighbors) >= 1
# Verify ordering: motion 3 ([1,1,0]) should be closer to motion 1 ([1,0,0]) than motion 2 ([0,1,0])
if len(neighbors) >= 2:
first = neighbors[0]
second = neighbors[1]
assert first["motion_id"] == motion_ids[2]
assert first["score"] >= second["score"]
else:
# If only one neighbor returned, it should be motion 3
assert neighbors[0]["motion_id"] == motion_ids[2]
+67
View File
@@ -0,0 +1,67 @@
import json
from pathlib import Path
from database import MotionDatabase
def test_similarity_cache_roundtrip(tmp_path: Path):
db_file = tmp_path / "motions.db"
# Create MotionDatabase which should initialize schema
db = MotionDatabase(db_path=str(db_file))
# If MotionDatabase fell back to file mode, check JSON files
if getattr(db, "_file_mode", False):
emb_file = Path(str(db_file) + ".embeddings.json")
sim_file = Path(str(db_file) + ".similarity_cache.json")
assert emb_file.exists()
assert sim_file.exists()
assert json.loads(emb_file.read_text(encoding="utf-8")) == []
assert json.loads(sim_file.read_text(encoding="utf-8")) == []
else:
# Try to import duckdb only when needed
import duckdb
conn = duckdb.connect(str(db_file))
embeddings_count = conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0]
similarity_count = conn.execute(
"SELECT COUNT(*) FROM similarity_cache"
).fetchone()[0]
conn.close()
assert embeddings_count == 0
assert similarity_count == 0
# Insert two similarity rows via helper
rows = [
{
"source_motion_id": 1,
"target_motion_id": 2,
"score": 0.5,
"vector_type": "text",
"window_id": None,
},
{
"source_motion_id": 1,
"target_motion_id": 3,
"score": 0.9,
"vector_type": "text",
"window_id": None,
},
]
db.store_similarity_batch(rows)
# Read back cached similarities and verify ordering (highest score first)
results = db.get_cached_similarities(source_motion_id=1, vector_type="text")
assert len(results) == 2
# results may be dicts from DB or file-backed dicts
assert results[0]["target_motion_id"] == 3
assert abs(float(results[0]["score"]) - 0.9) < 1e-6
assert results[1]["target_motion_id"] == 2
assert abs(float(results[1]["score"]) - 0.5) < 1e-6
# Clear cache and verify it's empty
db.clear_similarity_cache(vector_type="text")
results_after_clear = db.get_cached_similarities(
source_motion_id=1, vector_type="text"
)
assert results_after_clear == []