Refactor tests: replace sys.modules hacks with real DI + in-memory DB
- Add db=None, embedder=None params to ai_provider_wrapper, text_pipeline, compute_similarities - New conftest.py: FakeEmbedder, mem_db (in-memory DuckDB), fake_embedder fixtures - Rewrite test_ai_provider_wrapper (4 tests), test_rerun_embeddings_retry (2 tests), test_similarity_compute_filter (1 test) with real implementations - Fix rerun_embeddings tests hanging on _get_all_windows by patching it alongside _clear_embeddings - All 53 tests pass (2 skipped), 0 sys.modules hacks in refactored files
This commit is contained in:
@@ -0,0 +1,116 @@
|
||||
"""Wrapper around ai_provider to provide retries and smaller-batch fallback.
|
||||
|
||||
Returns a list of embedding vectors aligned with inputs. For inputs that
|
||||
fail permanently the corresponding list entry will be None and an audit event
|
||||
is appended via database.db.append_audit_event.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import random
|
||||
from typing import List, Optional
|
||||
|
||||
import ai_provider
|
||||
from database import db as motion_db
|
||||
import logging
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_embeddings_with_retry(
|
||||
texts: List[str],
|
||||
motion_ids: Optional[List[Optional[int]]] = None,
|
||||
model: Optional[str] = None,
|
||||
batch_size: int = 50,
|
||||
retries: int = 3,
|
||||
db=None,
|
||||
embedder=None,
|
||||
) -> List[Optional[List[float]]]:
|
||||
"""Return embeddings aligned with `texts` or None for failed items.
|
||||
|
||||
Strategy:
|
||||
- Try batches of `batch_size` with up to `retries` attempts.
|
||||
- On persistent batch failure, fall back to per-item attempts (batch_size=1).
|
||||
- Record an audit event for items that permanently fail.
|
||||
"""
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
if motion_ids is None:
|
||||
motion_ids = [None for _ in texts]
|
||||
|
||||
results: List[Optional[List[float]]] = [None] * len(texts)
|
||||
|
||||
# resolve embedder at call time; prefer injected, otherwise use ai_provider.get_embeddings_batch
|
||||
_embedder = embedder if embedder is not None else ai_provider.get_embeddings_batch
|
||||
|
||||
def _attempt_batch(chunk_texts, start_index):
|
||||
backoff = 0.5
|
||||
last_exc = None
|
||||
for attempt in range(1, retries + 1):
|
||||
try:
|
||||
emb_chunk = _embedder(
|
||||
chunk_texts, model=model, batch_size=len(chunk_texts)
|
||||
)
|
||||
return emb_chunk, None
|
||||
except Exception as exc:
|
||||
last_exc = exc
|
||||
if attempt == retries:
|
||||
break
|
||||
sleep = backoff * (2 ** (attempt - 1))
|
||||
sleep = sleep + random.uniform(0, sleep * 0.1)
|
||||
_logger.debug(
|
||||
"Batch embedding attempt %d failed, retrying after %.2fs: %s",
|
||||
attempt,
|
||||
sleep,
|
||||
exc,
|
||||
)
|
||||
time.sleep(sleep)
|
||||
# persistent failure
|
||||
_logger.warning(
|
||||
"Batch embedding failed for texts starting at %d: %s", start_index, last_exc
|
||||
)
|
||||
return None, last_exc
|
||||
|
||||
# process in batches
|
||||
i = 0
|
||||
n = len(texts)
|
||||
while i < n:
|
||||
end = min(n, i + batch_size)
|
||||
chunk = texts[i:end]
|
||||
emb_chunk, emb_exc = _attempt_batch(chunk, i)
|
||||
if emb_chunk is not None:
|
||||
# success: assign
|
||||
for j, emb in enumerate(emb_chunk):
|
||||
results[i + j] = emb
|
||||
i = end
|
||||
continue
|
||||
|
||||
# batch failed -> fallback to per-item attempts
|
||||
for j in range(i, end):
|
||||
t = texts[j]
|
||||
mid = motion_ids[j] if j < len(motion_ids) else None
|
||||
single, single_exc = _attempt_batch([t], j)
|
||||
if single:
|
||||
results[j] = single[0]
|
||||
continue
|
||||
|
||||
# permanent failure for this item
|
||||
err_text = repr(single_exc) if single_exc is not None else "unknown"
|
||||
try:
|
||||
_db = db if db is not None else motion_db
|
||||
_db.append_audit_event(
|
||||
actor_id=None,
|
||||
action="embedding_failed",
|
||||
target_type="motion",
|
||||
target_id=str(mid) if mid is not None else None,
|
||||
metadata={"error": err_text},
|
||||
)
|
||||
except Exception:
|
||||
_logger.exception("Failed to append audit event for embedding failure")
|
||||
results[j] = None
|
||||
|
||||
i = end
|
||||
|
||||
return results
|
||||
+120
-40
@@ -2,10 +2,13 @@ import logging
|
||||
import json
|
||||
from typing import Optional, List, Tuple
|
||||
|
||||
import duckdb
|
||||
try:
|
||||
import duckdb
|
||||
except Exception:
|
||||
duckdb = None
|
||||
|
||||
from database import MotionDatabase, db as default_db
|
||||
import ai_provider
|
||||
import pipeline.ai_provider_wrapper as ai_wrapper
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -19,11 +22,14 @@ def _select_text(
|
||||
|
||||
Returns list of (motion_id, text).
|
||||
"""
|
||||
if duckdb is None:
|
||||
return []
|
||||
conn = duckdb.connect(db.db_path)
|
||||
params = [model]
|
||||
# prefer layman_explanation > description > title (keep compatibility with existing tests)
|
||||
# prefer layman_explanation > body_text > description > title
|
||||
# (adds body_text as second-priority fallback so motion HTML is used when available)
|
||||
sql = (
|
||||
"SELECT m.id, COALESCE(m.layman_explanation, m.description, m.title) AS text"
|
||||
"SELECT m.id, COALESCE(m.layman_explanation, m.body_text, m.description, m.title) AS text"
|
||||
" FROM motions m"
|
||||
" LEFT JOIN embeddings e ON e.motion_id = m.id AND e.model = ?"
|
||||
" WHERE e.id IS NULL"
|
||||
@@ -55,38 +61,49 @@ def _select_text(
|
||||
|
||||
|
||||
def ensure_text_embeddings(
|
||||
db_path: Optional[str] = None, model: Optional[str] = None, batch_size: int = 50
|
||||
) -> Tuple[int, int, int, int]:
|
||||
db_path: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
batch_size: int = 50,
|
||||
db=None,
|
||||
embedder=None,
|
||||
) -> Tuple[int, int, int, int, list]:
|
||||
"""Ensure all motions have text embeddings for `model`.
|
||||
|
||||
Uses batched API calls (batch_size texts per HTTP request) for speed.
|
||||
Returns tuple (stored_count, skipped_existing, skipped_no_text, errors).
|
||||
"""
|
||||
model = model or DEFAULT_MODEL
|
||||
db = MotionDatabase(db_path) if db_path else default_db
|
||||
if db is None:
|
||||
db = MotionDatabase(db_path) if db_path else default_db
|
||||
|
||||
# motions to process
|
||||
to_process = _select_text(db, model)
|
||||
|
||||
# how many already exist
|
||||
conn = duckdb.connect(db.db_path)
|
||||
try:
|
||||
total_motions = conn.execute("SELECT COUNT(*) FROM motions").fetchone()[0]
|
||||
except Exception:
|
||||
if duckdb is None:
|
||||
total_motions = 0
|
||||
|
||||
try:
|
||||
existing = conn.execute(
|
||||
"SELECT COUNT(DISTINCT motion_id) FROM embeddings WHERE model = ?", (model,)
|
||||
).fetchone()[0]
|
||||
except Exception:
|
||||
existing = 0
|
||||
else:
|
||||
conn = duckdb.connect(db.db_path)
|
||||
try:
|
||||
total_motions = conn.execute("SELECT COUNT(*) FROM motions").fetchone()[0]
|
||||
except Exception:
|
||||
total_motions = 0
|
||||
|
||||
conn.close()
|
||||
try:
|
||||
existing = conn.execute(
|
||||
"SELECT COUNT(DISTINCT motion_id) FROM embeddings WHERE model = ?",
|
||||
(model,),
|
||||
).fetchone()[0]
|
||||
except Exception:
|
||||
existing = 0
|
||||
|
||||
conn.close()
|
||||
|
||||
stored = 0
|
||||
skipped_no_text = 0
|
||||
errors = 0
|
||||
failed_ids: list = []
|
||||
|
||||
# Separate motions with text from those without
|
||||
with_text: List[Tuple[int, str]] = []
|
||||
@@ -111,28 +128,13 @@ def ensure_text_embeddings(
|
||||
batch_ids = [mid for mid, _ in batch]
|
||||
batch_texts = [txt for _, txt in batch]
|
||||
|
||||
try:
|
||||
vecs = ai_provider.get_embeddings_batch(
|
||||
batch_texts, model=model, batch_size=batch_size
|
||||
)
|
||||
except Exception as exc:
|
||||
_logger.error(
|
||||
"Batch embedding failed for motions %s..%s: %s",
|
||||
batch_ids[0],
|
||||
batch_ids[-1],
|
||||
exc,
|
||||
)
|
||||
errors += len(batch)
|
||||
continue
|
||||
|
||||
if len(vecs) != len(batch):
|
||||
_logger.error(
|
||||
"Batch size mismatch: expected %d, got %d embeddings",
|
||||
len(batch),
|
||||
len(vecs),
|
||||
)
|
||||
errors += len(batch)
|
||||
continue
|
||||
vecs = ai_wrapper.get_embeddings_with_retry(
|
||||
batch_texts,
|
||||
motion_ids=batch_ids,
|
||||
model=model,
|
||||
batch_size=batch_size,
|
||||
embedder=embedder,
|
||||
)
|
||||
|
||||
batch_stored = 0
|
||||
for (motion_id, _text), vec in zip(batch, vecs):
|
||||
@@ -141,6 +143,7 @@ def ensure_text_embeddings(
|
||||
"Embedding provider returned non-list for motion %s", motion_id
|
||||
)
|
||||
errors += 1
|
||||
failed_ids.append(motion_id)
|
||||
continue
|
||||
|
||||
try:
|
||||
@@ -155,11 +158,13 @@ def ensure_text_embeddings(
|
||||
res,
|
||||
)
|
||||
errors += 1
|
||||
failed_ids.append(motion_id)
|
||||
except Exception as exc:
|
||||
_logger.error(
|
||||
"Error storing embedding for motion %s: %s", motion_id, exc
|
||||
)
|
||||
errors += 1
|
||||
failed_ids.append(motion_id)
|
||||
|
||||
_logger.info(
|
||||
"Batch %d-%d: stored %d/%d (total: %d/%d)",
|
||||
@@ -172,4 +177,79 @@ def ensure_text_embeddings(
|
||||
)
|
||||
|
||||
skipped_existing = int(existing)
|
||||
# Historically some callers expected a 4-tuple; return the primary
|
||||
# metrics (stored, skipped_existing, skipped_no_text, errors).
|
||||
# The list of failed_ids is intentionally not returned here to remain
|
||||
# backward-compatible with older callers.
|
||||
return stored, skipped_existing, skipped_no_text, errors
|
||||
|
||||
|
||||
def ensure_text_embeddings_for_ids(
|
||||
db_path: Optional[str] = None,
|
||||
ids: Optional[list] = None,
|
||||
model: Optional[str] = None,
|
||||
batch_size: int = 50,
|
||||
db=None,
|
||||
embedder=None,
|
||||
) -> Tuple[int, int, int, int, list]:
|
||||
"""Ensure embeddings for a specific list of motion ids.
|
||||
|
||||
This helper selects the motion texts for the supplied ids and reuses the
|
||||
same embedding logic. Returns the same tuple shape as ensure_text_embeddings.
|
||||
"""
|
||||
model = model or DEFAULT_MODEL
|
||||
if db is None:
|
||||
db = MotionDatabase(db_path) if db_path else default_db
|
||||
|
||||
if not ids:
|
||||
return 0, 0, 0, 0, []
|
||||
|
||||
# Fetch texts for given ids
|
||||
if duckdb is None:
|
||||
return 0, 0, 0, 0, []
|
||||
conn = duckdb.connect(db.db_path)
|
||||
try:
|
||||
placeholders = ",".join("?" for _ in ids)
|
||||
rows = conn.execute(
|
||||
f"SELECT id, COALESCE(layman_explanation, body_text, description, title) AS text FROM motions WHERE id IN ({placeholders})",
|
||||
ids,
|
||||
).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
to_process = [(int(r[0]), (r[1] or "").strip() or None) for r in rows]
|
||||
|
||||
# Reuse the main loop by creating a minimal local copy of the selection
|
||||
stored = 0
|
||||
skipped_no_text = 0
|
||||
errors = 0
|
||||
failed_ids = []
|
||||
|
||||
with_text = [(mid, txt) for mid, txt in to_process if txt]
|
||||
|
||||
for batch_start in range(0, len(with_text), batch_size):
|
||||
batch = with_text[batch_start : batch_start + batch_size]
|
||||
batch_ids = [mid for mid, _ in batch]
|
||||
batch_texts = [txt for _, txt in batch]
|
||||
|
||||
vecs = ai_wrapper.get_embeddings_with_retry(
|
||||
batch_texts,
|
||||
motion_ids=batch_ids,
|
||||
model=model,
|
||||
batch_size=batch_size,
|
||||
embedder=embedder,
|
||||
)
|
||||
|
||||
for (motion_id, _text), vec in zip(batch, vecs):
|
||||
if not isinstance(vec, list):
|
||||
errors += 1
|
||||
failed_ids.append(motion_id)
|
||||
continue
|
||||
res = db.store_embedding(motion_id, model, vec)
|
||||
if res and res > 0:
|
||||
stored += 1
|
||||
else:
|
||||
errors += 1
|
||||
failed_ids.append(motion_id)
|
||||
|
||||
return stored, 0, skipped_no_text, errors, failed_ids
|
||||
|
||||
Reference in New Issue
Block a user