diff --git a/tests/test_hybrid.py b/tests/test_hybrid.py new file mode 100644 index 0000000..23c7b66 --- /dev/null +++ b/tests/test_hybrid.py @@ -0,0 +1,272 @@ +""" +Contract tests for :func:`vd.hybrid_search` and the :class:`vd.SupportsHybrid` protocol. + +Every backend the suite can reach runs the same parametrized hybrid contract. +The runtime ``isinstance(collection, SupportsHybrid)`` discovery splits the +sweep into two paths automatically — native adapters use their own +``hybrid_search`` method; everything else uses the client-side BM25 + RRF +fallback. The contract assertions are the same either way. + +What's tested +------------- +- ``vd.hybrid_search`` returns a non-empty iterator for a non-trivial query + on a populated collection. +- Each result is a dict with ``id``, ``text``, ``score`` (fused), ``metadata``. +- Results are ordered by fused score, descending. +- A query that matches *only* lexically still ranks the matching doc above + noise, proving the lexical signal flows through the fusion. Similarly for + a query that matches *only* densely. +- The native vs fallback split is observable via ``isinstance(c, + SupportsHybrid)``; both code paths produce the contract above. +- ``bm25_lexical_search`` ranks documents in expected order on a small + hand-built corpus (no fusion, no embedder needed). +""" + +import pytest + +import vd + +# ---------- BM25-only unit tests (no fusion, no embedder, no server) ------- # + + +def test_bm25_basic_ordering(): + """BM25 ranks term-matching docs above non-matching docs.""" + client = vd.connect("memory") + col = client.create_collection("bm25", dimension=2) + col["a"] = vd.Document(id="a", text="the quick brown fox", vector=[1.0, 0.0]) + col["b"] = vd.Document(id="b", text="lazy dog sleeps", vector=[0.0, 1.0]) + col["c"] = vd.Document(id="c", text="quick fox runs", vector=[0.5, 0.5]) + + hits = vd.bm25_lexical_search(col, "quick fox", limit=3) + ids = [h["id"] for h in hits] + assert ids[:2] == ["c", "a"] or ids[:2] == ["a", "c"], ( + f"both 'a' and 'c' contain 'quick fox'; got {ids}" + ) + assert "b" not in ids[:2] + + +def test_bm25_empty_query_returns_empty_list(): + client = vd.connect("memory") + col = client.create_collection("bm25e", dimension=2) + col["a"] = vd.Document(id="a", text="hello world", vector=[1.0, 0.0]) + assert vd.bm25_lexical_search(col, "", limit=10) == [] + assert vd.bm25_lexical_search(col, " ", limit=10) == [] + + +def test_bm25_skips_docs_with_empty_text(): + client = vd.connect("memory") + col = client.create_collection("bm25t", dimension=2) + col["a"] = vd.Document(id="a", text="", vector=[1.0, 0.0]) + col["b"] = vd.Document(id="b", text="match this", vector=[0.0, 1.0]) + hits = vd.bm25_lexical_search(col, "match", limit=10) + assert [h["id"] for h in hits] == ["b"] + + +def test_bm25_filter_is_applied(): + client = vd.connect("memory") + col = client.create_collection("bm25f", dimension=2) + col["a"] = vd.Document( + id="a", text="cats purr softly", vector=[1.0, 0.0], metadata={"kind": "x"} + ) + col["b"] = vd.Document( + id="b", text="cats meow loudly", vector=[0.0, 1.0], metadata={"kind": "y"} + ) + hits = vd.bm25_lexical_search(col, "cats", limit=10, filter={"kind": "x"}) + assert [h["id"] for h in hits] == ["a"] + + +# ---------- Hybrid contract: parametrized over every reachable backend ----- # + + +# Vectors are chosen so dense-only and lexical-only signals are separable. +# Dimension matches the test embedder (EMBED_DIM=16 from conftest.py). +_DOCS = [ + # Dense-only winner for the query vector below: doc 'dense_top' has a vector + # close to the query vector but text that has nothing to do with the query. + { + "id": "dense_top", + "text": "completely unrelated text about geology and rocks", + "vec_seed": 0, # makes it match the query vector closely + }, + # Lexical-only winner: text matches but vector is orthogonal. + { + "id": "lex_top", + "text": "machine learning embeddings power retrieval systems", + "vec_seed": 1, + }, + # Both — should top the fused ranking. + { + "id": "both", + "text": "machine learning models for retrieval", + "vec_seed": 0, + }, + # Noise documents. + { + "id": "noise_1", + "text": "cooking recipes for autumn", + "vec_seed": 2, + }, + { + "id": "noise_2", + "text": "history of medieval trade routes", + "vec_seed": 3, + }, +] + + +def _seeded_vector(seed: int, dim: int) -> list[float]: + """Deterministic small vector — used so we can craft known-similar pairs.""" + import math + + base = [math.sin(seed + i) for i in range(dim)] + norm = math.sqrt(sum(x * x for x in base)) or 1.0 + return [x / norm for x in base] + + +@pytest.fixture +def populated_collection(client): + """A collection populated with the dense+lexical-separable test docs.""" + dim = 16 # conftest.EMBED_DIM + col = client.create_collection("hybrid_test", dimension=dim) + for d in _DOCS: + col[d["id"]] = vd.Document( + id=d["id"], + text=d["text"], + vector=_seeded_vector(d["vec_seed"], dim), + ) + return col + + +def test_hybrid_search_returns_contract_shape(populated_collection): + """Each hit is a dict with id, text, score, metadata; sorted desc by score.""" + query_vec = _seeded_vector(0, 16) + hits = list( + vd.hybrid_search( + populated_collection, + query_vec, + query_text="machine learning retrieval", + limit=4, + ) + ) + assert len(hits) > 0, "hybrid_search returned no results" + assert len(hits) <= 4 + for h in hits: + assert {"id", "text", "score", "metadata"} <= set(h.keys()) + # Descending score order. + scores = [h["score"] for h in hits] + assert scores == sorted(scores, reverse=True) + + +def test_hybrid_search_fuses_dense_and_lexical(populated_collection): + """ + The 'both' doc (top dense AND top lexical) ranks above 'dense_top' + (dense-only) and above 'lex_top' (lexical-only). RRF makes this true on + both the native and fallback paths. + """ + query_vec = _seeded_vector(0, 16) + hits = list( + vd.hybrid_search( + populated_collection, + query_vec, + query_text="machine learning retrieval", + limit=5, + ) + ) + ids = [h["id"] for h in hits] + assert "both" in ids, f"'both' missing from fused results: {ids}" + both_idx = ids.index("both") + # 'both' must be at or near the top — fused signal beats single-signal docs. + assert both_idx <= 1, ( + f"'both' should rank top-2; got position {both_idx} in {ids}" + ) + # And both single-signal docs should also appear above noise. + assert "noise_1" not in ids[:2] and "noise_2" not in ids[:2] + + +def test_hybrid_search_string_query_uses_text_for_both_sides( + populated_collection, embedder +): + """When query is a string, it's auto-used for both dense (via embedder) + lexical.""" + hits = list( + vd.hybrid_search(populated_collection, "machine learning retrieval", limit=5) + ) + ids = [h["id"] for h in hits] + # 'both' has the literal lexical match; should be top-ranked. + assert "both" in ids[:2] + + +def test_hybrid_search_requires_query_text_when_query_is_vector(populated_collection): + """A vector query without query_text is a clear error (no lexical side).""" + query_vec = _seeded_vector(0, 16) + with pytest.raises(ValueError, match="query_text"): + list(vd.hybrid_search(populated_collection, query_vec, limit=3)) + + +def test_hybrid_search_honors_filter(populated_collection): + """Filter is applied on both sub-searches; results are restricted accordingly.""" + # Re-tag one doc so we can filter it in. + populated_collection["both"] = vd.Document( + id="both", + text="machine learning models for retrieval", + vector=_seeded_vector(0, 16), + metadata={"tag": "keep"}, + ) + query_vec = _seeded_vector(0, 16) + hits = list( + vd.hybrid_search( + populated_collection, + query_vec, + query_text="machine learning retrieval", + limit=5, + filter={"tag": "keep"}, + ) + ) + ids = [h["id"] for h in hits] + assert ids == ["both"], f"only the tagged doc should pass the filter; got {ids}" + + +def test_native_vs_fallback_path_is_observable(populated_collection, backend_name): + """isinstance(c, SupportsHybrid) cleanly splits the two paths.""" + is_native = isinstance(populated_collection, vd.SupportsHybrid) + if backend_name in {"weaviate", "elasticsearch", "redis"}: + assert is_native, ( + f"{backend_name} should be SupportsHybrid in this PR; got {is_native}" + ) + else: + assert not is_native, ( + f"{backend_name} should NOT be SupportsHybrid yet; got {is_native}" + ) + + +def test_hybrid_search_custom_lexical_callable(populated_collection): + """Users can supply their own lexical_search callable (fallback path only).""" + if isinstance(populated_collection, vd.SupportsHybrid): + pytest.skip("lexical_search= override only applies on the fallback path") + + calls = {"count": 0} + + def my_lex(collection, query_text, *, limit, filter, **kwargs): + calls["count"] += 1 + # Trivial: return doc 'noise_1' as a sentinel so we can assert it was called. + doc = collection["noise_1"] + return [ + { + "id": doc.id, + "text": doc.text, + "score": 99.0, + "metadata": dict(doc.metadata), + } + ] + + hits = list( + vd.hybrid_search( + populated_collection, + _seeded_vector(0, 16), + query_text="anything", + limit=5, + lexical_search=my_lex, + ) + ) + assert calls["count"] == 1 + ids = [h["id"] for h in hits] + assert "noise_1" in ids, "the custom lexical search's hit should fuse in" diff --git a/vd/__init__.py b/vd/__init__.py index 629dee1..2a3a298 100644 --- a/vd/__init__.py +++ b/vd/__init__.py @@ -162,7 +162,9 @@ def skills_dir() -> _Path: # ----- advanced search ----------------------------------------------------- # from vd.search import ( # noqa: E402 + bm25_lexical_search, deduplicate_results, + hybrid_search, multi_query_search, reciprocal_rank_fusion, search_similar_to_document, @@ -284,6 +286,8 @@ def skills_dir() -> _Path: "search_similar_to_document", "reciprocal_rank_fusion", "deduplicate_results", + "hybrid_search", + "bm25_lexical_search", # configuration "load_config", "save_config", diff --git a/vd/backends/elasticsearch.py b/vd/backends/elasticsearch.py index 7de6575..49578af 100644 --- a/vd/backends/elasticsearch.py +++ b/vd/backends/elasticsearch.py @@ -365,6 +365,70 @@ def _query( raw_results = [_hit_to_result(h) for h in hits] return apply_client_filter(raw_results, filter, limit=limit) + # ----- native hybrid via ES BM25 + dense, fused client-side ----------- # + + def _lexical_query( + self, + text: str, + *, + limit: int, + filter: Optional[Filter], + **kwargs, + ) -> list[SearchResult]: + """ + BM25 lexical search over the ES ``text`` field. + + Used as the lexical side of :meth:`hybrid_search`. Metadata is filtered + client-side (same approach as :meth:`_query`). + """ + del kwargs + if not _index_exists(self._es, self.name): + return [] + k = overfetch_limit(limit, filter) + response = self._es.search( + index=self.name, + query={"match": {"text": text}}, + size=k, + ) + hits = response["hits"]["hits"] + raw_results = [_hit_to_result(h) for h in hits] + return list(apply_client_filter(raw_results, filter, limit=limit)) + + def hybrid_search( + self, + query, + *, + query_text=None, + limit: int = 10, + filter: Optional[Filter] = None, + k_dense: Optional[int] = None, + k_lexical: Optional[int] = None, + rrf_k: int = 60, + egress=None, + **kwargs, + ): + """ + Hybrid (kNN + BM25) search via Elasticsearch, fused client-side with RRF. + + See :class:`vd.SupportsHybrid` for the canonical contract. Backend + notes: ES 8.x has a server-side RRF retriever, but we deliberately run + ``knn`` and ``match`` separately and fuse client-side so the fused + score is uniform across vd backends. Pass ``query_text=...`` + explicitly when ``query`` is a vector. + """ + return self._hybrid_via_rrf( + query, + self._lexical_query, + query_text=query_text, + limit=limit, + filter=filter, + k_dense=k_dense, + k_lexical=k_lexical, + rrf_k=rrf_k, + egress=egress, + **kwargs, + ) + # --------------------------------------------------------------------------- # # Client diff --git a/vd/backends/redis.py b/vd/backends/redis.py index f4d2f5a..42f119a 100644 --- a/vd/backends/redis.py +++ b/vd/backends/redis.py @@ -365,12 +365,129 @@ def _query( return apply_client_filter(results, filter, limit=limit) + # ----- native hybrid via RediSearch BM25 + dense, fused client-side --- # + + def _lexical_query( + self, + text: str, + *, + limit: int, + filter: Optional[Filter], + **kwargs, + ) -> list[SearchResult]: + """ + BM25 lexical search via ``FT.SEARCH`` on the ``text`` field. + + Used as the lexical side of :meth:`hybrid_search`. Metadata filtering + is applied client-side (vd stores metadata as a JSON blob, not as + individually indexed RediSearch fields). + """ + del kwargs + if not _index_exists(self._redis, self._index_name): + return [] + fetch = overfetch_limit(limit, filter) + # Escape RediSearch query-language special characters; keep tokens. + escaped = _escape_redisearch_query(text) + if not escaped: + return [] + q = ( + Query(f"@text:({escaped})") + .return_fields("text", "metadata", "vd_id") + .paging(0, fetch) + .dialect(2) + ) + try: + response = self._redis.ft(self._index_name).search(q) + except Exception: + return [] + results: list[SearchResult] = [] + for doc in response.docs: + try: + metadata = json.loads(getattr(doc, "metadata", "{}") or "{}") + except (json.JSONDecodeError, TypeError): + metadata = {} + results.append( + { + "id": getattr(doc, "vd_id", ""), + "text": getattr(doc, "text", ""), + "score": 0.0, # RRF fuser uses ranks, not raw scores. + "metadata": metadata, + } + ) + return list(apply_client_filter(results, filter, limit=limit)) + + def hybrid_search( + self, + query, + *, + query_text=None, + limit: int = 10, + filter: Optional[Filter] = None, + k_dense: Optional[int] = None, + k_lexical: Optional[int] = None, + rrf_k: int = 60, + egress=None, + **kwargs, + ): + """ + Hybrid (KNN + BM25) search via RediSearch, fused client-side with RRF. + + See :class:`vd.SupportsHybrid` for the canonical contract. Backend + notes: Redis 8.4+ has a native ``HybridQuery`` that fuses BM25 and + KNN server-side; we deliberately run them separately and fuse + client-side so the fused score is uniform across vd backends and so + the adapter works on Redis < 8.4 as well. Pass ``query_text=...`` + explicitly when ``query`` is a vector. + """ + return self._hybrid_via_rrf( + query, + self._lexical_query, + query_text=query_text, + limit=limit, + filter=filter, + k_dense=k_dense, + k_lexical=k_lexical, + rrf_k=rrf_k, + egress=egress, + **kwargs, + ) + # --------------------------------------------------------------------------- # # Module-level helper used by _read (defined after the class for readability) # --------------------------------------------------------------------------- # +_REDISEARCH_SPECIAL_CHARS = set(",.<>{}[]\"':;!@#$%^&*()-+=~|\\/") + + +def _escape_redisearch_query(text: str) -> str: + """ + Escape RediSearch query-language special characters in ``text``. + + Returns a space-joined OR-able token string. Returns an empty string when + no usable tokens remain. + """ + cleaned: list[str] = [] + current: list[str] = [] + for ch in text: + if ch in _REDISEARCH_SPECIAL_CHARS: + if current: + cleaned.append("".join(current)) + current = [] + elif ch.isspace(): + if current: + cleaned.append("".join(current)) + current = [] + else: + current.append(ch) + if current: + cleaned.append("".join(current)) + # Filter out one-char tokens that RediSearch treats as stopwords. + tokens = [t for t in cleaned if len(t) >= 2] + return " | ".join(tokens) + + def _raw_to_document(raw: dict) -> Document: """ Rebuild a :class:`~vd.base.Document` from a Redis HGETALL response. diff --git a/vd/backends/weaviate.py b/vd/backends/weaviate.py index c4fdb15..6aa5587 100644 --- a/vd/backends/weaviate.py +++ b/vd/backends/weaviate.py @@ -388,6 +388,86 @@ def _query( ] return apply_client_filter(raw_results, filter, limit=limit) + # ----- native hybrid via Weaviate BM25 + dense, fused client-side ------ # + + def _lexical_query( + self, + text: str, + *, + limit: int, + filter: Optional[Filter], + **kwargs, + ) -> list[SearchResult]: + """ + BM25 lexical search over the Weaviate ``text`` property. + + Used as the lexical side of :meth:`hybrid_search`. Metadata filtering + is applied client-side (Weaviate's typed filter builder cannot operate + on vd's JSON-stringified metadata field). + """ + del kwargs # unused; bm25 call below has no pass-through knobs + fetch = overfetch_limit(limit, filter) + wcol = self._wcol() + response = wcol.query.bm25( + query=text, + limit=fetch, + ) + raw_results: list[SearchResult] = [] + for obj in response.objects: + props = obj.properties or {} + metadata = {} + try: + metadata = json.loads(props.get(_META_PROP, "") or "{}") + except (TypeError, ValueError): + metadata = {} + raw_results.append( + { + "id": props.get(_VD_ID_PROP, str(obj.uuid)), + "text": props.get(_TEXT_PROP, ""), + # BM25 score available via metadata.score if requested; + # the fuser uses ranks not raw scores, so leave as 0. + "score": 0.0, + "metadata": metadata, + } + ) + return list(apply_client_filter(raw_results, filter, limit=limit)) + + def hybrid_search( + self, + query, + *, + query_text=None, + limit: int = 10, + filter: Optional[Filter] = None, + k_dense: Optional[int] = None, + k_lexical: Optional[int] = None, + rrf_k: int = 60, + egress=None, + **kwargs, + ): + """ + Hybrid (dense + BM25) search via Weaviate, fused client-side with RRF. + + See :class:`vd.SupportsHybrid` for the canonical contract. Backend + notes: Weaviate's first-class ``query.hybrid(...)`` is **not** used + here — we run ``near_vector`` and ``bm25`` separately and fuse them + client-side so the result is independent of Weaviate's internal score + normalization and works the same way across native and fallback paths. + Pass ``query_text=...`` explicitly when ``query`` is a vector. + """ + return self._hybrid_via_rrf( + query, + self._lexical_query, + query_text=query_text, + limit=limit, + filter=filter, + k_dense=k_dense, + k_lexical=k_lexical, + rrf_k=rrf_k, + egress=egress, + **kwargs, + ) + # --------------------------------------------------------------------------- # # WeaviateClient diff --git a/vd/base.py b/vd/base.py index 483fb07..522324f 100644 --- a/vd/base.py +++ b/vd/base.py @@ -289,24 +289,56 @@ class SupportsHybrid(Protocol): A collection that supports native hybrid (dense + lexical) search. Hybrid search has no syntactic convergence across vector databases, so it - is an opt-in capability, never baseline. Feature-discover before calling:: + is an opt-in capability, never baseline. Prefer the top-level + :func:`vd.hybrid_search` — it dispatches to this protocol when the + collection implements it and falls back to a pure-Python BM25 + RRF + fusion otherwise. Feature-discover directly only when you specifically + need to refuse the fallback path:: if isinstance(collection, SupportsHybrid): - hits = collection.hybrid_search(vec, "query text", alpha=0.5) + hits = collection.hybrid_search("query text", limit=20) - Backends without native hybrid do not implement this; combine separate - dense + lexical result lists with :func:`vd.reciprocal_rank_fusion`. + The portable contract is **Reciprocal Rank Fusion** (every native backend + supports it). Weighted-blend (``alpha``) and other backend-specific + fusion variants are accepted via ``**kwargs`` and documented per adapter + — they are not portable across backends. + + Parameters + ---------- + query : str or list[float] + Query text (embedded via the collection's embedder if configured) + or a pre-computed query vector for the dense side. + query_text : str, optional + Explicit text for the lexical side. Defaults to ``query`` when + ``query`` is a string. **Required** when ``query`` is a vector. + limit : int + Number of fused results to return. + filter : dict, optional + Canonical ``vd`` metadata filter applied to both sub-searches. + k_dense, k_lexical : int, optional + How many results to fetch from each sub-search before fusion. + Both default to ``max(4 * limit, 50)``. Widen for higher recall. + rrf_k : int + Reciprocal Rank Fusion constant (typically 60). + egress : callable, optional + Transform applied to each fused result before it is yielded. + **kwargs + Backend-specific knobs (e.g. ``alpha=0.7`` on weaviate, + ``ranker="weighted"`` on milvus). Documented per adapter. """ def hybrid_search( self, - query: Vector, - query_text: str, + query: Union[str, Vector], *, + query_text: Optional[str] = None, limit: int = 10, filter: Optional[Filter] = None, - alpha: float = 0.5, - fusion: str = "rrf", + k_dense: Optional[int] = None, + k_lexical: Optional[int] = None, + rrf_k: int = 60, + egress: Optional[Callable[[SearchResult], Any]] = None, + **kwargs, ) -> Iterator[SearchResult]: ... @@ -490,6 +522,38 @@ def _resolve_query(self, query: Union[str, Vector]) -> Vector: vector = self.embed(query) if isinstance(query, str) else query return self._vet_vector(vector) + def _resolve_hybrid_inputs( + self, + query: Union[str, Vector], + query_text: Optional[str], + ) -> tuple[Vector, str]: + """ + Normalize ``(query, query_text)`` for ``hybrid_search`` into ``(vec, text)``. + + - ``query`` may be a string (used for both the dense and lexical sides + when ``query_text`` is omitted) or a pre-computed vector (then + ``query_text`` is **required**). + - The returned ``vec`` is dimension-vetted; the returned ``text`` is + guaranteed to be a non-empty string for the lexical side. + """ + if isinstance(query, str): + text = query_text if query_text is not None else query + vec = self._vet_vector(self.embed(query)) + else: + if query_text is None: + raise ValueError( + "hybrid_search needs a `query_text` for the lexical side " + "when `query` is a vector. Either pass query_text=..., or " + "pass `query` as a string and let the embedder handle both." + ) + text = query_text + vec = self._vet_vector(query) + if not text: + raise ValueError( + "hybrid_search needs a non-empty lexical query string." + ) + return vec, text + # ----- MutableMapping interface --------------------------------------- # def __setitem__(self, key: str, value: Union[str, tuple, Document]) -> None: @@ -562,6 +626,58 @@ def search( for result in self._query(query_vector, limit=limit, filter=filter, **kwargs): yield egress(result) if egress is not None else result + # ----- hybrid orchestration (used by SupportsHybrid adapters) --------- # + + def _hybrid_via_rrf( + self, + query: Union[str, Vector], + lexical_query: Callable[..., Iterable[SearchResult]], + *, + query_text: Optional[str] = None, + limit: int = 10, + filter: Optional[Filter] = None, + k_dense: Optional[int] = None, + k_lexical: Optional[int] = None, + rrf_k: int = 60, + egress: Optional[Callable[[SearchResult], Any]] = None, + **kwargs, + ) -> Iterator[SearchResult]: + """ + Run dense ``_query`` + a backend-specific ``lexical_query`` and fuse with RRF. + + The pattern every adapter's :meth:`hybrid_search` calls. The adapter + passes in its lexical primitive (``self._lexical_query``); this method + handles input resolution, filter validation, over-fetch defaults, and + Reciprocal Rank Fusion. + + ``lexical_query`` must accept ``(query_text, *, limit, filter, **kwargs)`` + and return an iterable of result dicts shaped like ``_query`` results. + """ + from vd.filters import validate_filter + from vd.search import _HYBRID_OVERFETCH_FLOOR, _rrf_fuse + + validate_filter(filter, supported=self.supported_filter_operators) + vec, text = self._resolve_hybrid_inputs(query, query_text) + k_dense_eff = ( + k_dense if k_dense is not None else max(4 * limit, _HYBRID_OVERFETCH_FLOOR) + ) + k_lex_eff = ( + k_lexical + if k_lexical is not None + else max(4 * limit, _HYBRID_OVERFETCH_FLOOR) + ) + # NOTE: ``**kwargs`` carry backend-specific native-fusion knobs (e.g. + # ``alpha=0.7``). This client-RRF orchestration cannot honor them and + # deliberately drops them rather than risk leaking them into the dense + # or lexical sub-query calls (which would raise TypeError on most + # backends). Adapters wanting native fusion should override + # ``hybrid_search`` and call their backend's fused API directly. + del kwargs # unused on this path + dense = list(self._query(vec, limit=k_dense_eff, filter=filter)) + lex = list(lexical_query(text, limit=k_lex_eff, filter=filter)) + for hit in _rrf_fuse([dense, lex], rrf_k=rrf_k, limit=limit): + yield egress(hit) if egress is not None else hit + # ----- batch convenience (also satisfies SupportsBatch) --------------- # def add_documents( diff --git a/vd/search.py b/vd/search.py index 6a26530..ae70d22 100644 --- a/vd/search.py +++ b/vd/search.py @@ -5,9 +5,11 @@ other advanced search patterns. """ -from typing import Any, Callable, Iterator, Optional, Union +import math +import re +from typing import Any, Callable, Iterable, Iterator, Optional, Union -from vd.base import Collection, SearchResult +from vd.base import Collection, SearchResult, SupportsHybrid, Vector def multi_query_search( @@ -346,3 +348,286 @@ def deduplicate_results( if result["score"] > seen[key_value]["score"]: seen[key_value] = result yield result + + +# --------------------------------------------------------------------------- # +# Hybrid search — top-level entry + client-side BM25 fallback +# --------------------------------------------------------------------------- # + +#: Minimum default for `k_dense`/`k_lexical` when callers don't override. +_HYBRID_OVERFETCH_FLOOR = 50 + +#: Simple ASCII-word tokenizer for the built-in BM25 fallback. Adapters that +#: want better tokenization (stemming, CJK, etc.) should pass a custom +#: ``lexical_search`` callable. +_TOKEN_RE = re.compile(r"\w+", re.UNICODE) + + +def _tokenize(text: str) -> list[str]: + """Lowercased word tokens — used by the built-in BM25 fallback.""" + return _TOKEN_RE.findall(text.lower()) + + +def bm25_lexical_search( + collection: Collection, + query_text: str, + *, + limit: int = 10, + filter: Optional[dict] = None, + k1: float = 1.5, + b: float = 0.75, +) -> list[SearchResult]: + """ + Brute-force BM25 lexical search over a vd collection's stored ``text``. + + Iterates every document in ``collection``, tokenizes its ``text`` field, + and computes Okapi BM25 scores against ``query_text``. Used as the default + lexical side of :func:`hybrid_search` when a collection does not implement + :class:`SupportsHybrid`. + + Cost is **O(N)** in the collection size — fine for prototypes and + collections up to ~100k documents. For larger workloads, either switch to + a backend with native hybrid search (weaviate, elasticsearch, redis, …) + or pass a custom ``lexical_search`` callable to :func:`hybrid_search` that + consults a real text index. + + Parameters + ---------- + collection : Collection + Any vd Collection. Documents whose ``text`` is empty contribute zero + score and are filtered out of the result. + query_text : str + The lexical query. + limit : int + Maximum number of results. + filter : dict, optional + Canonical ``vd`` metadata filter. Applied client-side via + :func:`vd.filters.matches_filter`. + k1, b : float + BM25 hyperparameters. Defaults match the standard Okapi BM25. + + Returns + ------- + list[dict] + Result dicts in the same shape as :meth:`Collection.search` — + ``{"id", "text", "score", "metadata"}`` — sorted by descending score. + + Examples + -------- + >>> import vd + >>> c = vd.connect('memory').create_collection('t', dimension=2) + >>> c['a'] = vd.Document(id='a', text='the quick brown fox', vector=[1.0, 0.0]) + >>> c['b'] = vd.Document(id='b', text='lazy dog sleeps', vector=[0.0, 1.0]) + >>> hits = bm25_lexical_search(c, 'quick fox', limit=1) + >>> hits[0]['id'] + 'a' + """ + from vd.filters import matches_filter + + query_tokens = _tokenize(query_text) + if not query_tokens: + return [] + + # Pass 1: tokenize once, compute document frequencies and lengths. + docs: list[tuple[str, str, dict, list[str]]] = [] + df: dict[str, int] = {} + for doc_id in collection: + doc = collection[doc_id] + if filter is not None and not matches_filter(doc.metadata or {}, filter): + continue + tokens = _tokenize(doc.text or "") + if not tokens: + continue + docs.append((doc_id, doc.text, dict(doc.metadata or {}), tokens)) + for term in set(tokens): + df[term] = df.get(term, 0) + 1 + + if not docs: + return [] + + n_docs = len(docs) + avg_len = sum(len(toks) for _, _, _, toks in docs) / n_docs + + # Pass 2: BM25 scoring (Okapi). + query_terms = set(query_tokens) + idf = { + term: math.log(1 + (n_docs - df[term] + 0.5) / (df[term] + 0.5)) + for term in query_terms + if term in df + } + + scored: list[SearchResult] = [] + for doc_id, text, metadata, tokens in docs: + score = 0.0 + doc_len = len(tokens) + tf: dict[str, int] = {} + for tok in tokens: + if tok in idf: + tf[tok] = tf.get(tok, 0) + 1 + if not tf: + continue + for term, freq in tf.items(): + numer = freq * (k1 + 1) + denom = freq + k1 * (1 - b + b * doc_len / avg_len) + score += idf[term] * numer / denom + scored.append( + {"id": doc_id, "text": text, "score": score, "metadata": metadata} + ) + + scored.sort(key=lambda r: r["score"], reverse=True) + return scored[:limit] + + +def _rrf_fuse( + result_lists: Iterable[list[SearchResult]], + *, + rrf_k: int = 60, + limit: int = 10, +) -> list[SearchResult]: + """Reciprocal Rank Fusion over result lists, returning the top ``limit``.""" + scores: dict[str, dict[str, Any]] = {} + for results in result_lists: + for rank, item in enumerate(results, 1): + doc_id = item["id"] + contribution = 1.0 / (rrf_k + rank) + if doc_id not in scores: + # Keep the first occurrence's payload for text/metadata. + scores[doc_id] = {"item": dict(item), "score": 0.0} + scores[doc_id]["score"] += contribution + fused = [] + for doc_id, entry in scores.items(): + item = entry["item"] + # Replace the source score with the fused RRF score so downstream + # consumers can rely on result["score"] = fused score. + item["score"] = entry["score"] + fused.append(item) + fused.sort(key=lambda r: r["score"], reverse=True) + return fused[:limit] + + +def hybrid_search( + collection: Collection, + query: Union[str, Vector], + *, + query_text: Optional[str] = None, + limit: int = 10, + filter: Optional[dict] = None, + k_dense: Optional[int] = None, + k_lexical: Optional[int] = None, + rrf_k: int = 60, + lexical_search: Optional[Callable[..., list[SearchResult]]] = None, + egress: Optional[Callable[[SearchResult], Any]] = None, + **kwargs, +) -> Iterator[SearchResult]: + """ + Hybrid (dense + lexical) search that works on any vd Collection. + + Dispatches to the collection's native ``hybrid_search`` when it implements + :class:`~vd.SupportsHybrid` (efficient, server-side). Otherwise fuses the + collection's own dense :meth:`~vd.Collection.search` with a client-side + lexical scan (default: :func:`bm25_lexical_search`) via **Reciprocal Rank + Fusion**. + + The portable contract is RRF. Backend-specific knobs (weighted blend + ``alpha``, fusion-type variants, native ranker choices) are accepted via + ``**kwargs`` and forwarded to the adapter when it has a native + implementation; they are ignored by the client-side fallback. + + Parameters + ---------- + collection : Collection + Any vd Collection — native-hybrid or not. + query : str or list[float] + Query text (embedded by the collection if it has an embedder) or a + pre-computed query vector. When ``query`` is a vector, ``query_text`` + is **required**. + query_text : str, optional + Explicit text for the lexical side. Defaults to ``query`` when + ``query`` is a string. + limit : int + Number of fused results to return. + filter : dict, optional + Canonical ``vd`` metadata filter, applied to both sub-searches. + k_dense, k_lexical : int, optional + How many results to fetch from each sub-search before fusion. Default + is ``max(4 * limit, 50)`` for each side. Widen for higher recall. + rrf_k : int + Reciprocal Rank Fusion constant (typically 60). + lexical_search : callable, optional + Custom ``lexical_search(collection, query_text, *, limit, filter, + **kwargs) -> list[SearchResult]``. Defaults to + :func:`bm25_lexical_search`. Used only on the fallback path. + egress : callable, optional + Per-result transform applied before yielding. + **kwargs + Extra options. On the native path they are forwarded to the adapter + (e.g. ``alpha=0.7`` on weaviate). On the fallback path they are + ignored. + + Yields + ------ + dict + Fused result dicts. ``score`` is the RRF score on the fallback path, + or the adapter's fused score on the native path. + + Examples + -------- + >>> import vd + >>> client = vd.connect('memory') + >>> col = client.create_collection('docs', dimension=2) + >>> col['a'] = vd.Document(id='a', text='cats purr', + ... vector=[1.0, 0.0]) + >>> col['b'] = vd.Document(id='b', text='dogs bark', + ... vector=[0.0, 1.0]) + >>> hits = list(vd.hybrid_search(col, [0.9, 0.1], query_text='cats', + ... limit=1)) + >>> hits[0]['id'] + 'a' + """ + k_dense_eff = k_dense if k_dense is not None else max(4 * limit, _HYBRID_OVERFETCH_FLOOR) + k_lexical_eff = ( + k_lexical if k_lexical is not None else max(4 * limit, _HYBRID_OVERFETCH_FLOOR) + ) + + # Native path. + if isinstance(collection, SupportsHybrid): + for hit in collection.hybrid_search( + query, + query_text=query_text, + limit=limit, + filter=filter, + k_dense=k_dense_eff, + k_lexical=k_lexical_eff, + rrf_k=rrf_k, + egress=egress, + **kwargs, + ): + yield hit + return + + # Fallback: dense via collection.search() + lexical via callable, fused by RRF. + # Resolve the text for the lexical side. (The dense side accepts the original + # `query` directly — the collection's search() does its own embed/vet.) + if isinstance(query, str): + text = query_text if query_text is not None else query + else: + if query_text is None: + raise ValueError( + "hybrid_search needs a `query_text` for the lexical side when " + "`query` is a vector. Either pass query_text=..., or pass " + "`query` as a string and let the embedder handle both." + ) + text = query_text + if not text: + raise ValueError("hybrid_search needs a non-empty lexical query string.") + + dense_hits = list( + collection.search(query, limit=k_dense_eff, filter=filter) + ) + + lex_fn = lexical_search if lexical_search is not None else bm25_lexical_search + lex_hits = lex_fn(collection, text, limit=k_lexical_eff, filter=filter) + + fused = _rrf_fuse([dense_hits, lex_hits], rrf_k=rrf_k, limit=limit) + for hit in fused: + yield egress(hit) if egress is not None else hit