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:
@@ -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]
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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]
|
||||
@@ -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 == []
|
||||
Reference in New Issue
Block a user