diff --git a/CHANGELOG.md b/CHANGELOG.md index de9c5a1..e4643d2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,9 @@ All notable changes to dynavec are documented here. This project adheres to ## [Unreleased] ### Added +- **`MultiQueryRetriever` and `HyDERetriever`** (#215) — first-class query-expansion + retrievers that fan out LLM-generated reformulations (or a hypothetical answer passage) + and fuse the ranked lists with RRF. Take plain callables; no new dependencies. - **`warm_cache()`** (#194) — pre-populate the query cache from a list of common queries. - **Learned RRF fusion weights** (#204) — `RRFWeightFitter` fits per-retriever RRF weights by maximizing nDCG over labeled queries. diff --git a/README.md b/README.md index 7867e20..d730efe 100644 --- a/README.md +++ b/README.md @@ -179,6 +179,34 @@ hits = db.search(vector=my_query_vector, top_k=5) Control the split with `DynavecConfig.filterable_keys` (allowlist of keys pushed to S3 Vectors for filtering) — keep it small; S3 Vectors caps filterable metadata size per vector. +### Query expansion: Multi-Query and HyDE + +Short or oddly-worded queries often miss documents that use different vocabulary. +Two adapters widen the candidate pool and fuse the ranked lists with RRF. They take +plain callables, so there is no LLM dependency: + +```python +from dynavec import MultiQueryRetriever, HyDERetriever + +multi = MultiQueryRetriever( + db.namespace("kb"), + generate_queries=lambda q: my_llm_variations(q, n=3), # -> list[str] + top_k=4, +) +hyde = HyDERetriever( + db.namespace("kb"), + generate_hypothetical=lambda q: my_llm_answer(q), # -> str + top_k=4, +) +hits = multi.search("how does serverless vector storage work?") +``` + +Fused results carry the **RRF score** (not cosine), and each sub-search costs one S3 +Vectors query plus one DynamoDB hydration, so a call runs roughly `1 + n_queries` +searches. If the LLM callable fails, the original query is searched alone (set +`on_generate_error="raise"` to fail loudly). See `examples/query_expansion.py` +(runs offline). + --- ## Framework integrations @@ -358,6 +386,7 @@ hit-rate, and a filterable **traces** table with per-trace drill-down. | **Distance metrics** | Index on cosine/euclidean (S3 Vectors native); client-side rescore in cosine / dot / euclidean / manhattan or a **weighted combination**, with optional result-set normalization | `search(..., rescore="dot", normalize_scores=True)` | | **Concurrency** | GIL-aware thread pool — real parallelism for I/O-bound AWS calls; parallel batched writes + `search_many`; tunable botocore connection pool (default 10, raise for high concurrency) | `DynavecConfig(max_workers=8, max_pool_connections=50)`, `db.search_many([...])` | | **Streaming** | Results yielded page-by-page as S3 Vectors paginates, so agents start consuming early | `for hit in db.search_stream(q): ...` | +| **Query expansion** | LLM-driven Multi-Query and HyDE retrieval fused with RRF; plain callables, no LLM dependency | `MultiQueryRetriever(kb, gen)`, `HyDERetriever(kb, gen)` | | **Namespace RAG** | Per-tenant/collection handles; isolation + even partitioning | `kb = db.namespace("kb"); kb.search(...)` | | **Product quantization** | Compress cached/hot-tier vectors up to 32× (ADC distance) | `ProductQuantizer(m=96).fit(X)` | | **Knowledge graph / ER** | Entities + relations in DynamoDB linked to embeddings; traverse to scope/guide vector search (GraphRAG) | `db.graph_add_edge(...)`, `db.graph_search(q, seed_entities=[...])` | diff --git a/examples/query_expansion.py b/examples/query_expansion.py new file mode 100644 index 0000000..88e5543 --- /dev/null +++ b/examples/query_expansion.py @@ -0,0 +1,128 @@ +"""Multi-Query and HyDE retrieval -- runs offline, no AWS or LLM needed. + + python examples/query_expansion.py + +The corpus below never uses the words in the user's query ("car won't start"), +which is the vocabulary-mismatch case where single-query search struggles. Both +retrievers take plain callables for the LLM step; here they are canned stand-ins. +In production swap in a real LLM, e.g.: + + def generate_queries(q: str) -> list[str]: + reply = my_llm(f"Give 3 alternative search queries for: {q}. One per line.") + return [line.strip() for line in reply.splitlines() if line.strip()] + + def generate_hypothetical(q: str) -> str: + return my_llm(f"Write a short passage that answers: {q}") + + retriever = MultiQueryRetriever(db.namespace("kb"), generate_queries, top_k=3) +""" + +from __future__ import annotations + +import math +import re +import zlib + +import dynavec.client as cm +from dynavec import Document, Dynavec, DynavecConfig, HyDERetriever, MultiQueryRetriever +from dynavec.embeddings.base import Embedder + +DIM = 128 + + +class BagOfWordsEmbedder(Embedder): + """Toy lexical embedder: hashed bag of words (so it needs word overlap).""" + + dimension = DIM + + def embed_documents(self, texts): + out = [] + for text in texts: + vec = [0.0] * DIM + for word in re.findall(r"[a-z']+", text.lower()): + vec[zlib.crc32(word.encode()) % DIM] += 1.0 + out.append(vec) + return out + + +# --- tiny in-memory stand-ins for S3 Vectors + DynamoDB (demo only) ----------- +class _S3(cm.S3VectorsStore): + def __init__(self, config, boto_session=None): + self.config, self._store = config, {} + + def put_vectors(self, vectors): + for key, vec, meta in vectors: + self._store[key] = (list(vec), dict(meta)) + + def query(self, query_vector, top_k, filter=None, **_): + def dist(v): + dot = sum(a * b for a, b in zip(query_vector, v)) + norm = (math.sqrt(sum(a * a for a in query_vector)) or 1e-9) * ( + math.sqrt(sum(b * b for b in v)) or 1e-9 + ) + return 1 - dot / norm + + rows = sorted(((k, dist(v)) for k, (v, _) in self._store.items()), key=lambda r: r[1]) + return [{"key": k, "distance": d} for k, d in rows[:top_k]] + + +class _DDB(cm.DynamoDBStore): + def __init__(self, config, boto_session=None): + self.config, self._store = config, {} + + def put_many(self, namespace, items): + for doc_id, text, meta in items: + self._store[(namespace, doc_id)] = {"text": text, "metadata": dict(meta)} + + def get_many(self, namespace, ids): + return {i: self._store[(namespace, i)] for i in ids if (namespace, i) in self._store} + + +CORPUS = { + "ignition": "Diagnosing engine ignition failure: check the battery, starter motor and spark plugs.", + "brakes": "Replacing worn brake pads and rotors on a front disc brake assembly.", + "tires": "Rotating tires and checking tread depth improves handling and fuel economy.", + "oil": "Changing engine oil and the oil filter every five thousand miles.", + "coolant": "Flushing the radiator coolant prevents the engine from overheating in summer.", +} + + +def show(title, hits): + print(f"\n{title}") + for h in hits: + print(f" {h.id:<9} {h.text[:70]}") + + +def main() -> None: + cm.S3VectorsStore, cm.DynamoDBStore = _S3, _DDB # demo only: no AWS + cfg = DynavecConfig(vector_bucket="demo", index="demo", table="demo", dimension=DIM) + db = Dynavec(cfg, embedder=BagOfWordsEmbedder()) + db.upsert([Document(id=k, text=v) for k, v in CORPUS.items()], namespace="kb") + kb = db.namespace("kb") + + query = "car won't start" + show("plain search (top 2):", kb.search(query, top_k=2)) + + multi = MultiQueryRetriever( + kb, + generate_queries=lambda q: [ + "engine ignition failure starter motor", + "dead battery spark plugs", + ], + top_k=2, + ) + show("MultiQueryRetriever (top 2):", multi.search(query)) + + hyde = HyDERetriever( + kb, + generate_hypothetical=lambda q: ( + "When a car will not start, the cause is usually a dead battery, " + "a failing starter motor, or worn spark plugs in the engine ignition." + ), + top_k=2, + ) + show("HyDERetriever (top 2):", hyde.search(query)) + + +if __name__ == "__main__": + main() diff --git a/src/dynavec/__init__.py b/src/dynavec/__init__.py index 103518a..f9731a0 100644 --- a/src/dynavec/__init__.py +++ b/src/dynavec/__init__.py @@ -62,6 +62,7 @@ maximal_marginal_relevance, reciprocal_rank_fusion, ) +from .retrievers import HyDERetriever, MultiQueryRetriever from .spfresh import ( Partition, SPFreshConfig, @@ -92,6 +93,8 @@ "warm_cache", "reciprocal_rank_fusion", "maximal_marginal_relevance", + "MultiQueryRetriever", + "HyDERetriever", "RRFWeightFitter", "HotTier", "Partition", @@ -123,4 +126,3 @@ "ItemTooLargeError", "MissingDependencyError", ] - diff --git a/src/dynavec/integrations/tools.py b/src/dynavec/integrations/tools.py index 50736d3..11402be 100644 --- a/src/dynavec/integrations/tools.py +++ b/src/dynavec/integrations/tools.py @@ -12,10 +12,11 @@ from ..client import Dynavec from ..namespace import NamespaceView +from ..retrievers import QueryExpansionRetriever def make_retriever_fn( - source: Dynavec | NamespaceView, + source: Dynavec | NamespaceView | QueryExpansionRetriever, *, top_k: int = 4, namespace: str = "default", @@ -27,10 +28,16 @@ def make_retriever_fn( """Return ``fn(query: str) -> str`` — the lowest common denominator tool. Works as-is in LangGraph nodes, CrewAI tools, Strands tools, or any - function-calling agent. + function-calling agent. ``source`` may also be a + :class:`~dynavec.retrievers.MultiQueryRetriever` or + :class:`~dynavec.retrievers.HyDERetriever` (``rescore`` is not supported there). """ + if isinstance(source, QueryExpansionRetriever) and rescore is not None: + raise ValueError("rescore is not supported with query-expansion retrievers") def _search(query: str): + if isinstance(source, QueryExpansionRetriever): + return source.search(query, top_k=top_k, filter=filter) if isinstance(source, NamespaceView): return source.search(query, top_k=top_k, filter=filter, rescore=rescore) return source.search( @@ -63,9 +70,7 @@ def as_langchain_tool(source, *, name: str = "dynavec_search", **kw) -> Any: raise MissingDependencyError("as_langchain_tool", "langchain-core", "langchain") from exc fn = make_retriever_fn(source, **kw) - return StructuredTool.from_function( - func=fn, name=name, description=fn.__doc__ - ) + return StructuredTool.from_function(func=fn, name=name, description=fn.__doc__) def as_crewai_tool(source, *, name: str = "dynavec_search", **kw) -> Any: diff --git a/src/dynavec/retrievers.py b/src/dynavec/retrievers.py new file mode 100644 index 0000000..5878a25 --- /dev/null +++ b/src/dynavec/retrievers.py @@ -0,0 +1,341 @@ +"""Query-expansion retrievers: Multi-Query and HyDE, fused with RRF. + +Single-query vector search misses documents when the query is short, colloquial +or uses different vocabulary than the corpus. These two adapters widen the +candidate pool *before* fusion: + +* :class:`MultiQueryRetriever` asks an LLM for a few reformulations of the query, + searches with the original and every reformulation in parallel, and fuses the + ranked lists with :func:`~dynavec.retrieval.reciprocal_rank_fusion`. +* :class:`HyDERetriever` (Hypothetical Document Embeddings) asks an LLM to write a + short *answer passage*, embeds that passage as a **document**, and searches with + the resulting vector -- optionally fused with the plain query search. + +Both take plain callables, so dynavec stays free of LLM dependencies:: + + from dynavec import MultiQueryRetriever + + retriever = MultiQueryRetriever( + db.namespace("docs"), + generate_queries=lambda q: my_llm_variations(q, n=3), + top_k=4, + ) + hits = retriever.search("how does serverless vector storage work?") + +Notes +----- +* ``SearchResult.score`` on fused results is the **RRF score** (roughly 0.01-0.03), + *not* a cosine similarity, and it is on that scale even when the LLM call fails + and only the original query is searched. ``distance`` comes from the first list + a document appears in; the original query's list is always first. +* Each sub-search is a normal :meth:`Dynavec.search` call, so each one pays for + one S3 Vectors query plus one DynamoDB hydration of ``per_query_k`` documents. + Cost per retriever call is therefore roughly ``(1 + n_queries)`` searches. +* Do not call ``retriever.search`` from a task already running on the client's + internal executor: the fan-out submits into that same pool. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Callable, Sequence +from typing import TYPE_CHECKING, Literal + +from .exceptions import ConfigurationError +from .models import Metadata, SearchResult +from .namespace import NamespaceView +from .retrieval import reciprocal_rank_fusion + +if TYPE_CHECKING: # pragma: no cover - typing only (client imports retrieval) + from .client import Dynavec + +logger = logging.getLogger(__name__) + +OnGenerateError = Literal["fallback", "raise"] + +__all__ = ["QueryExpansionRetriever", "MultiQueryRetriever", "HyDERetriever"] + + +def _clean_texts(original: str, candidates: object, limit: int) -> list[str]: + """Normalize LLM output into a small list of distinct, non-blank strings. + + Accepts a list/tuple of strings (or a single string), strips whitespace, + drops blanks and non-strings, removes case-insensitive duplicates (including + duplicates of ``original``), and keeps at most ``limit`` entries in order. + """ + if isinstance(candidates, str): + candidates = [candidates] + if not isinstance(candidates, (list, tuple)): + return [] + seen = {original.strip().casefold()} + out: list[str] = [] + for item in candidates: + if not isinstance(item, str): + continue + text = item.strip() + key = text.casefold() + if not text or key in seen: + continue + seen.add(key) + out.append(text) + if len(out) >= limit: + break + return out + + +class QueryExpansionRetriever: + """Shared plumbing: source resolution, fan-out, RRF fusion, error policy.""" + + def __init__( + self, + source: Dynavec | NamespaceView, + *, + namespace: str = "default", + top_k: int = 4, + per_query_k: int | None = None, + rrf_k: int = 60, + include_original: bool = True, + original_weight: float = 1.0, + on_generate_error: OnGenerateError = "fallback", + ) -> None: + if top_k < 1: + raise ValueError("top_k must be >= 1") + if per_query_k is not None and per_query_k < 1: + raise ValueError("per_query_k must be >= 1") + if rrf_k < 1: + raise ValueError("rrf_k must be >= 1") + if original_weight <= 0: + raise ValueError("original_weight must be > 0") + if on_generate_error not in ("fallback", "raise"): + raise ValueError("on_generate_error must be 'fallback' or 'raise'") + + if isinstance(source, NamespaceView): + # Same package: unwrap the view so we can reach the client's thread + # pool and embedder. The view's namespace always wins. + self._db = source._db + self._namespace = source.namespace + else: + self._db = source + self._namespace = namespace + + self.top_k = top_k + self.per_query_k = per_query_k + self.rrf_k = rrf_k + self.include_original = include_original + self.original_weight = original_weight + self.on_generate_error = on_generate_error + + # ------------------------------------------------------------- public API + def search( + self, + query: str, + *, + top_k: int | None = None, + filter: Metadata | None = None, + use_cache: bool | None = None, + ) -> list[SearchResult]: + """Expand ``query``, search, fuse with RRF, and return the top results.""" + if not isinstance(query, str) or not query.strip(): + raise ValueError("query must be a non-empty string") + k = self.top_k if top_k is None else top_k + if k < 1: + raise ValueError("top_k must be >= 1") + depth = self.per_query_k if self.per_query_k is not None else 2 * k + depth = max(depth, k) + + # Each entry is (weight, zero-arg callable returning one ranked list). + plan = self._plan(query, depth, filter, use_cache) + lists = self._fan_out([call for _, call in plan]) + weights = [w for w, _ in plan] + return reciprocal_rank_fusion(lists, k=self.rrf_k, weights=weights)[:k] + + async def asearch( + self, + query: str, + *, + top_k: int | None = None, + filter: Metadata | None = None, + use_cache: bool | None = None, + ) -> list[SearchResult]: + """Async wrapper (runs :meth:`search` in a worker thread).""" + return await asyncio.to_thread( + self.search, query, top_k=top_k, filter=filter, use_cache=use_cache + ) + + # ------------------------------------------------------------- internals + def _plan( + self, + query: str, + depth: int, + filter: Metadata | None, + use_cache: bool | None, + ) -> list[tuple[float, Callable[[], list[SearchResult]]]]: + raise NotImplementedError + + def _text_search( + self, text: str, depth: int, filter: Metadata | None, use_cache: bool | None + ) -> Callable[[], list[SearchResult]]: + def call() -> list[SearchResult]: + return self._db.search( + text, + top_k=depth, + namespace=self._namespace, + filter=filter, + use_cache=use_cache, + ) + + return call + + def _vector_search( + self, + vector: Sequence[float], + depth: int, + filter: Metadata | None, + use_cache: bool | None, + ) -> Callable[[], list[SearchResult]]: + def call() -> list[SearchResult]: + return self._db.search( + vector=list(vector), + top_k=depth, + namespace=self._namespace, + filter=filter, + use_cache=use_cache, + ) + + return call + + def _fan_out(self, calls: list[Callable[[], list[SearchResult]]]) -> list[list[SearchResult]]: + """Run sub-searches concurrently; results keep submission order. + + Collecting futures in submission order (not completion order) is what + makes fusion deterministic: RRF breaks score ties by first appearance. + """ + if len(calls) == 1: + return [calls[0]()] + futures = [self._db._executor.submit(call) for call in calls] + return [f.result() for f in futures] + + def _generate(self, fn: Callable[[str], object], query: str) -> object | None: + """Call the user's LLM callable under the configured error policy.""" + try: + return fn(query) + except Exception as exc: # noqa: BLE001 - user callable, any failure + if self.on_generate_error == "raise": + raise + logger.warning( + "%s: generation failed (%s: %s); falling back to the original query only", + type(self).__name__, + type(exc).__name__, + exc, + ) + return None + + +class MultiQueryRetriever(QueryExpansionRetriever): + """Search with an LLM's reformulations of the query and fuse with RRF. + + Parameters + ---------- + source: + A :class:`~dynavec.client.Dynavec` client (use ``namespace=``) or a + :class:`~dynavec.namespace.NamespaceView` (its namespace is used). + generate_queries: + ``Callable[[str], list[str]]`` returning alternative phrasings. Output + is stripped, de-duplicated (case-insensitive, also against the original) + and capped at ``n_queries``. + n_queries: + Maximum number of reformulations to use (extras are dropped). + include_original: + Also search with the original query (recommended: guards against poor + reformulations). If generation fails, the original is always searched. + original_weight: + RRF weight of the original query's list; reformulations weigh ``1.0``. + top_k / per_query_k / rrf_k: + Final result count, candidates fetched per sub-search (default + ``2 * top_k``; never less than ``top_k``), and the RRF constant. + on_generate_error: + ``"fallback"`` (default) logs a warning and searches the original query + only; ``"raise"`` propagates the callable's exception. Errors from the + stores themselves always propagate. + """ + + def __init__( + self, + source: Dynavec | NamespaceView, + generate_queries: Callable[[str], Sequence[str]], + *, + n_queries: int = 3, + **kwargs, + ) -> None: + if not callable(generate_queries): + raise TypeError("generate_queries must be callable") + if n_queries < 1: + raise ValueError("n_queries must be >= 1") + super().__init__(source, **kwargs) + self._generate_queries = generate_queries + self.n_queries = n_queries + + def _plan(self, query, depth, filter, use_cache): + raw = self._generate(self._generate_queries, query) + expansions = _clean_texts(query, raw, self.n_queries) + + plan: list[tuple[float, Callable[[], list[SearchResult]]]] = [] + # Original goes first so `distance` and tie-breaks favor it. + if self.include_original or not expansions: + plan.append((self.original_weight, self._text_search(query, depth, filter, use_cache))) + for text in expansions: + plan.append((1.0, self._text_search(text, depth, filter, use_cache))) + return plan + + +class HyDERetriever(QueryExpansionRetriever): + """Hypothetical Document Embeddings retrieval, optionally fused with the query. + + The hypothetical passage is embedded with ``embedder.embed_documents`` (the + *document* side of asymmetric models) and searched by vector. + + Parameters + ---------- + source: + A ``Dynavec`` client or ``NamespaceView`` that has an embedder configured. + generate_hypothetical: + ``Callable[[str], str]`` returning a short passage that would answer the + query. Blank output is treated as a generation failure. + include_original: + Fuse in the plain query search (default ``True``). Recommended: a + hallucinated passage can point away from the right documents. + original_weight, top_k, per_query_k, rrf_k, on_generate_error: + As for :class:`MultiQueryRetriever`. + """ + + def __init__( + self, + source: Dynavec | NamespaceView, + generate_hypothetical: Callable[[str], str], + **kwargs, + ) -> None: + if not callable(generate_hypothetical): + raise TypeError("generate_hypothetical must be callable") + super().__init__(source, **kwargs) + if self._db.embedder is None: + raise ConfigurationError( + "HyDERetriever needs an embedder to embed the hypothetical passage. " + "Pass one to Dynavec(..., embedder=...)." + ) + self._generate_hypothetical = generate_hypothetical + + def _plan(self, query, depth, filter, use_cache): + raw = self._generate(self._generate_hypothetical, query) + cleaned = _clean_texts("", raw, 1) + hypothetical = cleaned[0] if cleaned else None + + plan: list[tuple[float, Callable[[], list[SearchResult]]]] = [] + if self.include_original or hypothetical is None: + plan.append((self.original_weight, self._text_search(query, depth, filter, use_cache))) + if hypothetical is not None: + # Embedding happens here (calling thread) so embedder errors surface + # directly instead of inside a pool future. + vector = self._db.embedder.embed_documents([hypothetical])[0] + plan.append((1.0, self._vector_search(vector, depth, filter, use_cache))) + return plan diff --git a/tests/test_retrievers.py b/tests/test_retrievers.py new file mode 100644 index 0000000..e4a2ffb --- /dev/null +++ b/tests/test_retrievers.py @@ -0,0 +1,326 @@ +"""Offline tests for MultiQueryRetriever and HyDERetriever (#215). + +Reuses the in-memory FakeS3 / FakeDDB harness from test_client_inmemory and a +tiny mapping embedder so rankings are fully controlled. +""" + +import logging +import threading +import time + +import pytest +from test_client_inmemory import FakeDDB, FakeGraph, FakeS3 + +import dynavec.client as client_mod +from dynavec import Document, Dynavec, DynavecConfig, HyDERetriever, MultiQueryRetriever +from dynavec.embeddings.base import Embedder +from dynavec.exceptions import ConfigurationError +from dynavec.integrations.tools import make_retriever_fn + +E0, E1, E2, E3 = ([1.0, 0, 0, 0], [0, 1.0, 0, 0], [0, 0, 1.0, 0], [0, 0, 0, 1.0]) + +# query text -> vector (anything unknown embeds to a neutral vector) +VECTORS = { + "orig": E0, + "alt": E1, + "alt2": E2, + "hypo": E1, +} + + +class MapEmbedder(Embedder): + dimension = 4 + + def __init__(self): + self.doc_calls: list[list[str]] = [] + self.query_calls: list[str] = [] + + def embed_documents(self, texts): + self.doc_calls.append(list(texts)) + return [list(VECTORS.get(t, [0.5, 0.5, 0.5, 0.5])) for t in texts] + + def embed_query(self, text): + self.query_calls.append(text) + return list(VECTORS.get(text, [0.5, 0.5, 0.5, 0.5])) + + +@pytest.fixture +def db(monkeypatch): + monkeypatch.setattr(client_mod, "S3VectorsStore", FakeS3) + monkeypatch.setattr(client_mod, "DynamoDBStore", FakeDDB) + monkeypatch.setattr(client_mod, "GraphStore", FakeGraph) + cfg = DynavecConfig(vector_bucket="b", index="i", table="t", dimension=4) + d = Dynavec(cfg, embedder=MapEmbedder()) + d.upsert( + [ + Document(id="d0", text="zero", vector=[1.0, 0.0, 0.0, 0.0], metadata={"lang": "en"}), + Document(id="d1", text="one", vector=[0.9, 0.1, 0.0, 0.0], metadata={"lang": "en"}), + Document(id="d2", text="two", vector=[0.0, 1.0, 0.0, 0.0], metadata={"lang": "fr"}), + Document(id="d3", text="three", vector=[0.0, 0.9, 0.1, 0.0], metadata={"lang": "en"}), + Document(id="d4", text="four", vector=[0.0, 0.0, 1.0, 0.0], metadata={"lang": "en"}), + ] + ) + return d + + +def ids(results): + return [r.id for r in results] + + +# ----------------------------------------------------------------- multi-query +def test_multiquery_fuses_lists_and_keeps_original_first_on_ties(db): + r = MultiQueryRetriever(db, lambda q: ["alt"], top_k=2, per_query_k=2) + # orig -> [d0, d1], alt -> [d2, d3]; d0 and d2 tie on RRF, original list wins. + assert ids(r.search("orig")) == ["d0", "d2"] + + +def test_per_query_k_is_floored_at_top_k(db): + seen = [] + real = db.search + + def spy(query=None, **kw): + seen.append(kw["top_k"]) + return real(query, **kw) + + db.search = spy + MultiQueryRetriever(db, lambda q: ["alt"], top_k=5, per_query_k=2).search("orig") + MultiQueryRetriever(db, lambda q: ["alt"], top_k=3).search("orig") + assert seen == [5, 5, 6, 6] + + +def test_multiquery_original_weight_reorders(db): + r = MultiQueryRetriever(db, lambda q: ["alt"], top_k=2, per_query_k=2, original_weight=3.0) + assert ids(r.search("orig")) == ["d0", "d1"] # unweighted this would be d0, d2 + + +def test_multiquery_dedupes_documents_across_lists(db): + r = MultiQueryRetriever(db, lambda q: ["orig-ish", "alt"], top_k=10, per_query_k=3) + out = ids(r.search("orig")) + assert len(out) == len(set(out)) + + +def test_multiquery_include_original_false_skips_original(db): + calls = [] + real = db.search + + def spy(query=None, **kw): + calls.append(query) + return real(query, **kw) + + db.search = spy + MultiQueryRetriever(db, lambda q: ["alt"], include_original=False).search("orig") + assert calls == ["alt"] + + +def test_multiquery_sanitizes_generated_queries(db): + calls = [] + real = db.search + + def spy(query=None, **kw): + calls.append(query) + return real(query, **kw) + + db.search = spy + gen = lambda q: [" alt ", "ALT", "", " ", "ORIG", None, 7, "alt2", "extra"] # noqa: E731 + MultiQueryRetriever(db, gen, n_queries=2).search("orig") + # original first, then blanks/dupes/non-strings dropped, capped at n_queries=2 + assert sorted(calls) == ["alt", "alt2", "orig"] # thread order is not asserted + + +def test_multiquery_accepts_single_string_output(db): + r = MultiQueryRetriever(db, lambda q: "alt", top_k=2, per_query_k=2) + assert ids(r.search("orig")) == ["d0", "d2"] + + +def test_multiquery_generation_failure_falls_back_and_warns(db, caplog): + def boom(q): + raise RuntimeError("llm down") + + r = MultiQueryRetriever(db, boom, top_k=2) + with caplog.at_level(logging.WARNING, logger="dynavec.retrievers"): + out = r.search("orig") + assert ids(out) == ["d0", "d1"] + assert "falling back" in caplog.text + + +def test_multiquery_empty_generation_searches_original_even_without_include_original(db): + r = MultiQueryRetriever(db, lambda q: [], include_original=False, top_k=2) + assert ids(r.search("orig")) == ["d0", "d1"] + + +def test_multiquery_on_generate_error_raise(db): + def boom(q): + raise RuntimeError("llm down") + + with pytest.raises(RuntimeError): + MultiQueryRetriever(db, boom, on_generate_error="raise").search("orig") + + +def test_multiquery_store_errors_propagate(db): + def broken_search(*a, **kw): + raise OSError("s3 vectors unavailable") + + db.search = broken_search + with pytest.raises(OSError): + MultiQueryRetriever(db, lambda q: ["alt"]).search("orig") + + +def test_multiquery_scores_are_rrf_scale(db): + out = MultiQueryRetriever(db, lambda q: ["alt"], rrf_k=60).search("orig") + assert all(0 < r.score <= 2 / 61 for r in out) + + +def test_multiquery_result_order_is_independent_of_thread_timing(db): + real = db.search + + def make_slow(delays): + def slow(query=None, **kw): + time.sleep(delays.get(query, 0)) + return real(query, **kw) + + return slow + + r = MultiQueryRetriever(db, lambda q: ["alt", "alt2"], top_k=5, per_query_k=3) + db.search = make_slow({"orig": 0.15}) + a = ids(r.search("orig")) + db.search = make_slow({"alt2": 0.15}) + b = ids(r.search("orig")) + assert a == b + + +def test_multiquery_runs_subsearches_concurrently(db): + active, peak, lock = 0, 0, threading.Lock() + real = db.search + + def tracked(query=None, **kw): + nonlocal active, peak + with lock: + active += 1 + peak = max(peak, active) + time.sleep(0.05) + try: + return real(query, **kw) + finally: + with lock: + active -= 1 + + db.search = tracked + MultiQueryRetriever(db, lambda q: ["alt", "alt2"]).search("orig") + assert peak >= 2 + + +def test_filter_is_passed_to_every_subsearch(db): + r = MultiQueryRetriever(db, lambda q: ["alt"], top_k=10, per_query_k=10) + out = r.search("orig", filter={"lang": "en"}) + assert "d2" not in ids(out) # d2 is lang=fr + assert set(ids(out)) <= {"d0", "d1", "d3", "d4"} + + +def test_namespace_view_and_namespace_kwarg(db): + db.upsert([Document(id="k0", text="kb doc", vector=E0)], namespace="kb") + via_view = MultiQueryRetriever(db.namespace("kb"), lambda q: ["alt"]) + via_kwarg = MultiQueryRetriever(db, lambda q: ["alt"], namespace="kb") + assert ids(via_view.search("orig")) == ["k0"] + assert ids(via_kwarg.search("orig")) == ["k0"] + # default namespace is untouched by the kb doc + assert "k0" not in ids(MultiQueryRetriever(db, lambda q: ["alt"]).search("orig")) + + +def test_per_call_top_k_override(db): + r = MultiQueryRetriever(db, lambda q: ["alt"], top_k=4) + assert len(r.search("orig", top_k=1)) == 1 + + +async def test_asearch_matches_search(db): + r = MultiQueryRetriever(db, lambda q: ["alt"], top_k=2, per_query_k=2) + assert ids(await r.asearch("orig")) == ids(r.search("orig")) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"top_k": 0}, + {"per_query_k": 0}, + {"rrf_k": 0}, + {"original_weight": 0}, + {"n_queries": 0}, + {"on_generate_error": "ignore"}, + ], +) +def test_multiquery_validates_arguments(db, kwargs): + with pytest.raises(ValueError): + MultiQueryRetriever(db, lambda q: [], **kwargs) + + +def test_rejects_blank_query_and_non_callable(db): + r = MultiQueryRetriever(db, lambda q: []) + with pytest.raises(ValueError): + r.search(" ") + with pytest.raises(TypeError): + MultiQueryRetriever(db, "not callable") + + +# ------------------------------------------------------------------------ HyDE +def test_hyde_embeds_hypothetical_as_document_and_fuses_with_query(db): + emb = db.embedder + r = HyDERetriever(db, lambda q: "hypo", top_k=2, per_query_k=2) + out = r.search("orig") + # "hypo" is embedded on the document side, the raw query on the query side + assert ["hypo"] in emb.doc_calls + assert "orig" in emb.query_calls and "hypo" not in emb.query_calls + # orig -> [d0, d1]; hypo (E1) -> [d2, d3]; tie between d0 and d2 goes to original + assert ids(out) == ["d0", "d2"] + + +def test_hyde_without_original_uses_only_hypothetical(db): + r = HyDERetriever(db, lambda q: "hypo", include_original=False, top_k=2) + assert ids(r.search("orig")) == ["d2", "d3"] + assert "orig" not in db.embedder.query_calls + + +@pytest.mark.parametrize("bad", ["", " ", None, []]) +def test_hyde_blank_generation_falls_back_to_original(db, bad): + r = HyDERetriever(db, lambda q: bad, include_original=False, top_k=2) + assert ids(r.search("orig")) == ["d0", "d1"] + + +def test_hyde_generation_exception_falls_back(db): + def boom(q): + raise TimeoutError("slow llm") + + assert ids(HyDERetriever(db, boom, top_k=2).search("orig")) == ["d0", "d1"] + with pytest.raises(TimeoutError): + HyDERetriever(db, boom, on_generate_error="raise").search("orig") + + +def test_hyde_requires_embedder(monkeypatch): + monkeypatch.setattr(client_mod, "S3VectorsStore", FakeS3) + monkeypatch.setattr(client_mod, "DynamoDBStore", FakeDDB) + monkeypatch.setattr(client_mod, "GraphStore", FakeGraph) + bare = Dynavec(DynavecConfig(vector_bucket="b", index="i", table="t", dimension=4)) + with pytest.raises(ConfigurationError): + HyDERetriever(bare, lambda q: "x") + + +def test_hyde_honors_filter_and_namespace_view(db): + db.upsert([Document(id="k0", text="kb", vector=E1, metadata={"lang": "en"})], namespace="kb") + r = HyDERetriever(db.namespace("kb"), lambda q: "hypo", top_k=5) + assert ids(r.search("orig", filter={"lang": "en"})) == ["k0"] + + +async def test_hyde_asearch(db): + r = HyDERetriever(db, lambda q: "hypo", top_k=2, per_query_k=2) + assert ids(await r.asearch("orig")) == ids(r.search("orig")) + + +# ------------------------------------------------------------- tool factory +def test_make_retriever_fn_accepts_retrievers(db): + r = MultiQueryRetriever(db, lambda q: ["alt"], top_k=4, per_query_k=2) + fn = make_retriever_fn(r, top_k=2, join=" | ") + assert fn("orig") == "zero | two" + + +def test_make_retriever_fn_rejects_rescore_for_retrievers(db): + r = MultiQueryRetriever(db, lambda q: ["alt"]) + with pytest.raises(ValueError): + make_retriever_fn(r, rescore="cosine")