From 68029da2ecb433e79c415575d33753c04cbe8e63 Mon Sep 17 00:00:00 2001 From: Isha Zaka Date: Sat, 19 Sep 2026 22:51:56 -0400 Subject: [PATCH] feat: add search explain output --- src/dynavec/__init__.py | 12 ++- src/dynavec/client.py | 152 ++++++++++++++++++++++++++++++-- src/dynavec/models.py | 18 ++++ src/dynavec/namespace.py | 6 +- tests/test_client_inmemory.py | 157 +++++++++++++++++++++++++++++++++- 5 files changed, 332 insertions(+), 13 deletions(-) diff --git a/src/dynavec/__init__.py b/src/dynavec/__init__.py index 53a1133..18e93b5 100644 --- a/src/dynavec/__init__.py +++ b/src/dynavec/__init__.py @@ -50,7 +50,14 @@ XlsxSource, ingest, ) -from .models import Document, IndexInfo, SearchResult, UpsertResult +from .models import ( + Document, + ExplainedSearchResult, + IndexInfo, + SearchExplanation, + SearchResult, + UpsertResult, +) from .namespace import NamespaceView from .quantization import ( OPQRotation, @@ -84,6 +91,8 @@ "Document", "IndexInfo", "SearchResult", + "SearchExplanation", + "ExplainedSearchResult", "UpsertResult", "NamespaceView", "ProductQuantizer", @@ -133,4 +142,3 @@ "ItemTooLargeError", "MissingDependencyError", ] - diff --git a/src/dynavec/client.py b/src/dynavec/client.py index 9dfce97..9ff27fb 100644 --- a/src/dynavec/client.py +++ b/src/dynavec/client.py @@ -52,7 +52,14 @@ from .metrics import normalize_scores as normalize_metric_scores from .metrics import rescore as metric_rescore from .metrics import score as metric_score -from .models import Document, IndexInfo, SearchResult, UpsertResult +from .models import ( + Document, + ExplainedSearchResult, + IndexInfo, + SearchExplanation, + SearchResult, + UpsertResult, +) from .namespace import NamespaceView from .provisioning import provision_all from .retrieval import distance_to_score, maximal_marginal_relevance, reciprocal_rank_fusion @@ -344,7 +351,8 @@ def search( include_vectors: bool = False, use_cache: bool | None = None, normalize_scores: bool = False, - ) -> list[SearchResult]: + explain: bool = False, + ) -> list[SearchResult] | ExplainedSearchResult: """Semantic search. Provide ``query`` (embedded) or a raw ``vector``. ``rescore`` re-orders the ANN candidates with a client-side metric @@ -355,11 +363,20 @@ def search( If a cache is configured, repeated/similar queries are served from it (set ``use_cache=False`` to force a fresh search). + Set ``explain=True`` to return results together with per-stage timings + and candidate counts for debugging. """ t0 = time.perf_counter() tel = self._telemetry + explanation = SearchExplanation() if explain else None + try: + stage_t0 = time.perf_counter() query_vector = self._resolve_query_vector(query, vector) + if explanation is not None: + explanation.timings_ms["query_vector"] = round( + (time.perf_counter() - stage_t0) * 1000, 3 + ) # cache key includes ranking options so different ranking != same entry cache_on = self._cache is not None if use_cache is None else use_cache @@ -374,10 +391,43 @@ def search( "normalize_scores": normalize_scores, }, } + stage_t0 = time.perf_counter() cached = self._cache.get(namespace, query_vector, top_k, cache_filter) + + if explanation is not None: + explanation.timings_ms["cache_lookup"] = round( + (time.perf_counter() - stage_t0) * 1000, 3 + ) + explanation.candidate_counts["cached"] = ( + len(cached) if cached is not None else 0 + ) + if cached is not None: - self._record_search(tel, t0, namespace, top_k, cached, True, - filter, rescore, rerank, query) + if explanation is not None: + explanation.candidate_counts["final"] = len(cached) + explanation.timings_ms["total"] = round( + (time.perf_counter() - t0) * 1000, 3 + ) + + self._record_search( + tel, + t0, + namespace, + top_k, + cached, + True, + filter, + rescore, + rerank, + query, + ) + + if explanation is not None: + return ExplainedSearchResult( + results=cached, + explanation=explanation, + ) + return cached results = self._search_core( @@ -391,13 +441,43 @@ def search( mmr_lambda=mmr_lambda, include_vectors=include_vectors, normalize_scores=normalize_scores, + explanation=explanation, ) if cache_on and self._cache is not None and results: + stage_t0 = time.perf_counter() self._cache.put(namespace, query_vector, top_k, cache_filter, results) - self._record_search(tel, t0, namespace, top_k, results, - (False if cache_on else None), - filter, rescore, rerank, query) + + if explanation is not None: + explanation.timings_ms["cache_write"] = round( + (time.perf_counter() - stage_t0) * 1000, 3 + ) + + if explanation is not None: + explanation.candidate_counts["final"] = len(results) + explanation.timings_ms["total"] = round( + (time.perf_counter() - t0) * 1000, 3 + ) + + self._record_search( + tel, + t0, + namespace, + top_k, + results, + (False if cache_on else None), + filter, + rescore, + rerank, + query, + ) + + if explanation is not None: + return ExplainedSearchResult( + results=results, + explanation=explanation, + ) + return results except Exception as exc: # noqa: BLE001 - record then re-raise if tel is not None: @@ -422,6 +502,7 @@ def _search_core( mmr_lambda, include_vectors, normalize_scores, + explanation: SearchExplanation | None = None, ) -> list[SearchResult]: needs_vectors = rerank == "mmr" or rescore is not None or include_vectors fetch_k = top_k * self.config.over_fetch if (rerank or rescore) else top_k @@ -432,9 +513,18 @@ def _search_core( # transparently fall back to S3 (hot tier can only speed up, never break). results: list[SearchResult] | None = None if self._hot is not None: + stage_t0 = time.perf_counter() results = self._hot.search(namespace, query_vector, fetch_k, filter) + if explanation is not None: + explanation.timings_ms["hot_lookup"] = round( + (time.perf_counter() - stage_t0) * 1000, 3 + ) + if results is not None: + explanation.candidate_counts["retrieved"] = len(results) + if results is None: + stage_t0 = time.perf_counter() raw = self._vectors.query( query_vector=query_vector, top_k=fetch_k, @@ -442,19 +532,38 @@ def _search_core( return_metadata=True, return_distance=True, ) + if explanation is not None: + explanation.timings_ms["vector_search"] = round( + (time.perf_counter() - stage_t0) * 1000, 3 + ) + explanation.candidate_counts["retrieved"] = len(raw) if not raw: return [] hits = [(self._split_key(v["key"])[1], v.get("distance")) for v in raw] ids = [h[0] for h in hits] + stage_t0 = time.perf_counter() hydrated = self._docs.get_many(namespace, ids) + if explanation is not None: + explanation.timings_ms["hydration"] = round( + (time.perf_counter() - stage_t0) * 1000, 3 + ) + explanation.candidate_counts["hydrated"] = len(hydrated) + vec_by_key = {} if needs_vectors: + stage_t0 = time.perf_counter() vec_by_key = self._vectors.get_vectors( [self._s3_key(namespace, doc_id) for doc_id in ids] ) + if explanation is not None: + explanation.timings_ms["vector_fetch"] = round( + (time.perf_counter() - stage_t0) * 1000, 3 + ) + explanation.candidate_counts["vectors_fetched"] = len(vec_by_key) + results = [] for doc_id, distance in hits: doc = hydrated.get(doc_id, {}) @@ -475,28 +584,57 @@ def _search_core( ) if rescore is not None: + stage_t0 = time.perf_counter() results = self._apply_rescore(query_vector, results, rescore) + + if explanation is not None: + explanation.timings_ms["rescore"] = round( + (time.perf_counter() - stage_t0) * 1000, 3 + ) + explanation.candidate_counts["rescored"] = len(results) if rerank == "mmr": + stage_t0 = time.perf_counter() results = maximal_marginal_relevance( results, query_vector, top_k=top_k, lambda_mult=mmr_lambda, ) + + if explanation is not None: + explanation.timings_ms["rerank"] = round( + (time.perf_counter() - stage_t0) * 1000, 3 + ) + explanation.candidate_counts["reranked"] = len(results) + elif rerank == "cross-encoder": + stage_t0 = time.perf_counter() results = self._cross_encoder_rerank( query, results, top_k=top_k, ) + + if explanation is not None: + explanation.timings_ms["rerank"] = round( + (time.perf_counter() - stage_t0) * 1000, 3 + ) + explanation.candidate_counts["reranked"] = len(results) + else: results = results[:top_k] if normalize_scores and results: + stage_t0 = time.perf_counter() normalized = normalize_metric_scores(np.asarray([r.score for r in results])) for result, normalized_score in zip(results, normalized): result.score = float(normalized_score) + if explanation is not None: + explanation.timings_ms["normalize_scores"] = round( + (time.perf_counter() - stage_t0) * 1000, 3 + ) + if not include_vectors: for r in results: r.vector = None diff --git a/src/dynavec/models.py b/src/dynavec/models.py index c90e003..9f240f5 100644 --- a/src/dynavec/models.py +++ b/src/dynavec/models.py @@ -52,6 +52,24 @@ def to_dict(self) -> dict[str, Any]: "text": self.text, "metadata": self.metadata, } + + +@dataclass +class SearchExplanation: + """Debug information collected for an explained search.""" + + timings_ms: dict[str, float] = field(default_factory=dict) + candidate_counts: dict[str, int] = field(default_factory=dict) + + +@dataclass +class ExplainedSearchResult: + """Search results together with per-stage debug information.""" + + results: list[SearchResult] = field(default_factory=list) + explanation: SearchExplanation = field(default_factory=SearchExplanation) + + @dataclass class IndexInfo: """Snapshot of the provisioned S3 Vectors index + DynamoDB table.""" diff --git a/src/dynavec/namespace.py b/src/dynavec/namespace.py index 22b3a2f..6be2782 100644 --- a/src/dynavec/namespace.py +++ b/src/dynavec/namespace.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any -from .models import SearchResult +from .models import ExplainedSearchResult, SearchResult if TYPE_CHECKING: from .client import Dynavec @@ -35,7 +35,9 @@ def upsert(self, documents, **kw) -> Any: def update(self, *args, **kw) -> Any: return self._db.update(*args, namespace=self._ns, **kw) - def search(self, query: str | None = None, **kw) -> list[SearchResult]: + def search( + self, query: str | None = None, **kw + ) -> list[SearchResult] | ExplainedSearchResult: return self._db.search(query, namespace=self._ns, **kw) def search_stream(self, query: str | None = None, **kw): diff --git a/tests/test_client_inmemory.py b/tests/test_client_inmemory.py index f8ec948..a0fa61f 100644 --- a/tests/test_client_inmemory.py +++ b/tests/test_client_inmemory.py @@ -11,7 +11,13 @@ import pytest import dynavec.client as client_mod -from dynavec import Document, Dynavec, DynavecConfig, SemanticCache +from dynavec import ( + Document, + Dynavec, + DynavecConfig, + ExplainedSearchResult, + SemanticCache, +) from dynavec.config import NS_METADATA_KEY from dynavec.embeddings.base import Embedder from dynavec.exceptions import ConfigurationError @@ -218,6 +224,103 @@ def test_upsert_and_search_roundtrip(db): assert hits[0].score >= hits[1].score +def test_search_without_explain_returns_results_list(db): + db.upsert([Document(id="1", text="apple pie")]) + + result = db.search("apple") + + assert isinstance(result, list) + assert len(result) == 1 + + +def test_search_explain_handles_empty_results(db): + result = db.search("missing", explain=True) + + assert isinstance(result, ExplainedSearchResult) + assert result.results == [] + + explanation = result.explanation + + assert explanation.candidate_counts["retrieved"] == 0 + assert explanation.candidate_counts["final"] == 0 + assert explanation.timings_ms["vector_search"] >= 0 + assert explanation.timings_ms["total"] >= 0 + + +def test_search_explain_reports_cache_hit(db): + db._cache = SemanticCache(threshold=0.99) + db.upsert([Document(id="1", text="apple pie")]) + + db.search("apple pie", top_k=3) + + result = db.search("apple pie", top_k=3, explain=True) + + assert isinstance(result, ExplainedSearchResult) + assert len(result.results) == 1 + + explanation = result.explanation + + assert explanation.candidate_counts["cached"] == 1 + assert explanation.candidate_counts["final"] == 1 + assert explanation.timings_ms["cache_lookup"] >= 0 + assert explanation.timings_ms["total"] >= 0 + + assert "vector_search" not in explanation.timings_ms + assert "hydration" not in explanation.timings_ms + + +def test_search_explain_records_rescore_stage(db): + db.upsert( + [ + Document(id="1", text="apple pie recipe"), + Document(id="2", text="apple orchard tour"), + Document(id="3", text="rocket launch"), + ] + ) + + result = db.search( + "apple", + top_k=2, + rescore="cosine", + explain=True, + ) + + assert isinstance(result, ExplainedSearchResult) + assert len(result.results) == 2 + + explanation = result.explanation + + assert explanation.timings_ms["rescore"] >= 0 + assert explanation.candidate_counts["rescored"] >= 2 + assert explanation.candidate_counts["final"] == 2 + + +def test_search_explain_returns_structured_debug_result(db): + db.upsert( + [ + Document(id="1", text="apple pie recipe"), + Document(id="2", text="rocket launch schedule"), + Document(id="3", text="apple orchard tour"), + ] + ) + + result = db.search("apple", top_k=2, explain=True) + + assert isinstance(result, ExplainedSearchResult) + assert len(result.results) == 2 + + explanation = result.explanation + + assert explanation.timings_ms["query_vector"] >= 0 + assert explanation.timings_ms["vector_search"] >= 0 + assert explanation.timings_ms["hydration"] >= 0 + assert explanation.timings_ms["total"] >= 0 + + assert explanation.candidate_counts["retrieved"] >= 2 + assert explanation.candidate_counts["hydrated"] >= 2 + assert explanation.candidate_counts["final"] == 2 + + def test_metadata_filter_scopes_results(db): db.upsert( [ @@ -610,6 +713,56 @@ def predict(self, pairs): assert results[0].score == 0.95 assert FakeCrossEncoder.instance.model_name == "toy-cross-encoder" + +def test_search_explain_records_cross_encoder_rerank_stage( + monkeypatch, cross_encoder_db +): + class FakeCrossEncoder: + def __init__(self, model_name): + self.model_name = model_name + + def predict(self, pairs): + return [ + 0.95 if "target document" in document else 0.10 + for _, document in pairs + ] + + monkeypatch.setitem( + sys.modules, + "sentence_transformers", + SimpleNamespace(CrossEncoder=FakeCrossEncoder), + ) + + cross_encoder_db.upsert([ + Document( + id="d1", + text="ordinary document", + vector=[0.1] * 8, + ), + Document( + id="d2", + text="target document", + vector=[0.1] * 8, + ), + ]) + + result = cross_encoder_db.search( + "find the target", + top_k=1, + rerank="cross-encoder", + explain=True, + ) + + assert isinstance(result, ExplainedSearchResult) + assert result.results[0].id == "d2" + + explanation = result.explanation + + assert explanation.timings_ms["rerank"] >= 0 + assert explanation.candidate_counts["reranked"] == 1 + assert explanation.candidate_counts["final"] == 1 + + def test_cross_encoder_rerank_requires_text_query(cross_encoder_db): cross_encoder_db.upsert([ Document( @@ -650,4 +803,4 @@ def __init__(self, model_name): "find something", top_k=1, rerank="cross-encoder", - ) \ No newline at end of file + )