Cache the fused RRF union per query fingerprint

radis/pgsearch/providers.py

@@ -3,6 +3,7 @@
 import logging
 import unicodedata
 from collections.abc import Iterator
+from dataclasses import asdict
 from datetime import date, datetime, time, timedelta
 from typing import Literal, NamedTuple, cast
 
@@ -400,14 +401,101 @@ class _FusedHybrid(NamedTuple):
     query_str: str
 
 
+# Bumped whenever the cached payload's shape changes, so entries written by an
+# older code version are missed rather than misread.
+_FUSED_CACHE_SCHEMA = 1
+
+
+def _fused_cache_key(search: Search) -> str:
+    """Fingerprint of everything that determines the fused union.
+
+    Same discipline as `_embed_query_cached`: a shared cache backend (production
+    uses the database) outlives process restarts and config changes, so the
+    fusion/beam settings and the embedding configuration are part of the key,
+    not just the query and filters."""
+    spec = settings.EMBEDDINGS_MODEL
+    fingerprint = json.dumps(
+        {
+            "schema": _FUSED_CACHE_SCHEMA,
+            "query": QueryParser.unparse(search.query),
+            "filters": asdict(search.filters),
+            "fts_max": settings.HYBRID_FTS_MAX_RESULTS,
+            "rrf_k": settings.HYBRID_RRF_K,
+            "ef_search": settings.HYBRID_HNSW_EF_SEARCH,
+            "embeddings": None
+            if spec is None
+            else [
+                settings.EMBEDDINGS_BASE_URL,
+                spec.model,
+                spec.params,
+                settings.EMBEDDINGS_QUERY_INSTRUCTION,
+                settings.EMBEDDINGS_DIM,
+            ],
+        },
+        sort_keys=True,
+        default=str,  # date/datetime filter values
+    )
+    return "pgsearch-fused-" + hashlib.sha256(fingerprint.encode()).hexdigest()
+
+
+def _fused_to_cached(fused: _FusedHybrid) -> dict:
+    """Plain-dict payload instead of pickling the NamedTuple, so a cached entry
+    never carries a stale class shape across deploys."""
+    return {
+        "ids": list(fused.rrf_score_by_id.keys()),
+        "scores": list(fused.rrf_score_by_id.values()),
+        "dist_ids": list(fused.vec_distance.keys()),
+        "dists": list(fused.vec_distance.values()),
+        "total_relation": fused.total_relation,
+        "configs": fused.configs,
+        "query_str": fused.query_str,
+    }
+
+
+def _fused_from_cached(cached: dict) -> _FusedHybrid:
+    rrf_score_by_id = dict(zip(cached["ids"], cached["scores"], strict=True))
+    return _FusedHybrid(
+        ordered_ids=list(rrf_score_by_id),
+        rrf_score_by_id=rrf_score_by_id,
+        vec_distance=dict(zip(cached["dist_ids"], cached["dists"], strict=True)),
+        total_relation=cached["total_relation"],
+        configs=[(config, list(codes)) for config, codes in cached["configs"]],
+        query_str=cached["query_str"],
+    )
+
+
 def _fuse_hybrid(search: Search, caller: str) -> _FusedHybrid:
+    """Cached front of `_fuse_hybrid_uncached`.
+
+    Pagination re-runs the whole fusion for every page, and ranking every FTS
+    match of a common term costs seconds on a large corpus, so the fused union
+    is cached per query fingerprint and page 2..n (and repeated searches) reuse
+    it. Degraded FTS-only results — a configured embedding model that produced
+    no vector — are never cached, so an embedding outage cannot pin searches to
+    FTS-only for the TTL. New or updated reports enter a cached query's results
+    only after the TTL expires."""
+    timeout = settings.HYBRID_FUSED_CACHE_TIMEOUT_SECONDS
+    if timeout <= 0:
+        return _fuse_hybrid_uncached(search, caller)[0]
+    key = _fused_cache_key(search)
+    cached = cache.get(key)
+    if cached is not None:
+        return _fused_from_cached(cached)
+    fused, degraded = _fuse_hybrid_uncached(search, caller)
+    if not degraded:
+        cache.set(key, _fused_to_cached(fused), timeout=timeout)
+    return fused
+
+
+def _fuse_hybrid_uncached(search: Search, caller: str) -> tuple[_FusedHybrid, bool]:
     """Run both retrievers and fuse them with RRF.
 
     Shared by search(), retrieve() and count() so all three describe the same
     union. The vector half embeds only the positive branches (§7.8 strips
     ``NOT``) and enforces the top-level negations on its candidates; the FTS
     half consumes the full boolean tsquery. Returns the fused, ordered report
-    ids plus the per-id scores/distances and the query metadata callers reuse."""
+    ids plus the per-id scores/distances and the query metadata callers reuse,
+    and whether the result is degraded (vector half skipped by a failure)."""
     query_str = _build_query_string(search.query)
     configs = _language_configs(search.filters)
     filter_query = _build_filter_query(search.filters)
@@ -419,6 +507,11 @@ def _fuse_hybrid(search: Search, caller: str) -> _FusedHybrid:
     query_vec: list[float] | None = None
     if settings.EMBEDDINGS_MODEL is not None and query_text.strip():
         query_vec = _embed_query_cached(query_text, caller)
+    # A configured model that produced no vector means the embedding service is
+    # failing or rate limited right now; the caller must not cache this result.
+    degraded = (
+        settings.EMBEDDINGS_MODEL is not None and bool(query_text.strip()) and query_vec is None
+    )
 
     vec_rank: dict[int, int] = {}
     vec_distance: dict[int, float] = {}
@@ -468,7 +561,7 @@ def _fuse_hybrid(search: Search, caller: str) -> _FusedHybrid:
         total_relation=total_relation,
         configs=configs,
         query_str=query_str,
-    )
+    ), degraded
 
 
 def search(search: Search) -> SearchResult:

radis/pgsearch/tests/test_fused_cache.py

@@ -0,0 +1,185 @@
+from unittest.mock import patch
+
+import pytest
+from django.contrib.auth.models import Group
+from django.core.cache import cache
+from django.db import connection
+from django.test.utils import CaptureQueriesContext
+
+from radis.core.utils.embedding_client import EmbeddingClientError
+from radis.core.utils.model_spec import parse_model_spec
+from radis.pgsearch.models import ReportSearchIndex
+from radis.pgsearch.providers import search
+from radis.reports.factories import ReportFactory
+from radis.reports.models import Report
+from radis.search.site import Search, SearchFilters
+from radis.search.utils.query_parser import QueryParser
+
+pytestmark = pytest.mark.django_db
+
+
+@pytest.fixture(autouse=True)
+def _clear_cache():
+    """The fused cache and the query-embedding cache share the process-local test
+    cache backend, which outlives individual tests."""
+    cache.clear()
+    yield
+    cache.clear()
+
+
+@pytest.fixture(autouse=True)
+def _embeddings_model_configured(settings):
+    settings.EMBEDDINGS_MODEL = parse_model_spec("qwen3")
+
+
+def _unit_vec(idx: int, dim: int) -> list[float]:
+    v = [0.0] * dim
+    v[idx % dim] = 1.0
+    return v
+
+
+def _make_report(body: str, **overrides) -> Report:
+    """Pin language to ``en`` so index-side and query-side tsquery configs match
+    (see test_provider_hybrid._make_report)."""
+    return ReportFactory.create(body=body, language__code="en", **overrides)
+
+
+def _make_search(query_str: str, group_id: int, offset: int = 0, limit: int = 25) -> Search:
+    node, _ = QueryParser().parse(query_str)
+    assert node is not None
+    return Search(
+        query=node,
+        filters=SearchFilters(group=group_id, language="en"),
+        offset=offset,
+        limit=limit,
+    )
+
+
+@pytest.fixture
+def group(db):
+    return Group.objects.create(name="radiology")
+
+
+@pytest.fixture
+def reports(group, settings):
+    dim = settings.EMBEDDINGS_DIM
+    r_fts = _make_report(body="Findings: pneumothorax on the left.")
+    r_fts.groups.add(group)
+    # Doesn't lexically match "pneumothorax": reachable via the vector side only.
+    r_vec = _make_report(body="Lungs are clear bilaterally.")
+    r_vec.groups.add(group)
+    ReportSearchIndex.objects.filter(report=r_fts).update(embedding=_unit_vec(99, dim))
+    ReportSearchIndex.objects.filter(report=r_vec).update(embedding=_unit_vec(0, dim))
+    return r_fts, r_vec
+
+
+def _mock_embedding(dim: int, idx: int = 0):
+    mocked = patch("radis.pgsearch.providers.EmbeddingClient")
+    MockClient = mocked.start()
+    MockClient.return_value.__enter__.return_value = MockClient.return_value
+    MockClient.return_value.__exit__.return_value = None
+    MockClient.return_value.embed_query.return_value = _unit_vec(idx, dim)
+    return mocked, MockClient
+
+
+def _ran_fusion(queries) -> bool:
+    """Whether any captured query ran a fusion retriever.
+
+    The FTS half ranks every match (ts_rank with no report_id list); the vector
+    half is the only user of the <=> operator. The page-document fetch also
+    annotates ts_rank/ts_headline, but only for an explicit ``report_id IN``
+    page slice, and must not count as fusion work."""
+    for q in queries:
+        sql = q["sql"]
+        if "<=>" in sql:
+            return True
+        if "ts_rank" in sql and '"report_id" IN' not in sql:
+            return True
+    return False
+
+
+def test_second_identical_search_reuses_cached_fusion(group, reports, settings):
+    mocked, _ = _mock_embedding(settings.EMBEDDINGS_DIM)
+    try:
+        with CaptureQueriesContext(connection) as first:
+            result1 = search(_make_search("pneumothorax", group.pk))
+        with CaptureQueriesContext(connection) as second:
+            result2 = search(_make_search("pneumothorax", group.pk))
+    finally:
+        mocked.stop()
+
+    assert _ran_fusion(first.captured_queries)
+    # The fusion (FTS ranking + vector scan) must be served from cache; only the
+    # page-document fetch may hit the database.
+    assert not _ran_fusion(second.captured_queries)
+    assert [d.document_id for d in result1.documents] == [d.document_id for d in result2.documents]
+    assert result1.total_count == result2.total_count
+
+
+def test_pagination_is_served_from_the_cached_union(group, reports, settings):
+    mocked, _ = _mock_embedding(settings.EMBEDDINGS_DIM)
+    try:
+        page1 = search(_make_search("pneumothorax", group.pk, offset=0, limit=1))
+        with CaptureQueriesContext(connection) as ctx:
+            page2 = search(_make_search("pneumothorax", group.pk, offset=1, limit=1))
+    finally:
+        mocked.stop()
+
+    assert not _ran_fusion(ctx.captured_queries)
+    ids1 = [d.document_id for d in page1.documents]
+    ids2 = [d.document_id for d in page2.documents]
+    assert len(ids1) == len(ids2) == 1
+    assert ids1 != ids2  # genuinely the next page of the same union
+
+
+def test_filters_participate_in_the_cache_key(reports, settings):
+    other = Group.objects.create(name="cardiology")
+    r_other = _make_report(body="No pneumothorax after drainage.")
+    r_other.groups.add(other)
+    ReportSearchIndex.objects.filter(report=r_other).update(
+        embedding=_unit_vec(1, settings.EMBEDDINGS_DIM)
+    )
+    r_fts, _ = reports
+
+    mocked, _ = _mock_embedding(settings.EMBEDDINGS_DIM)
+    try:
+        first_group = reports[0].groups.first()
+        result_first = search(_make_search("pneumothorax", first_group.pk))
+        result_other = search(_make_search("pneumothorax", other.pk))
+    finally:
+        mocked.stop()
+
+    ids_first = {d.document_id for d in result_first.documents}
+    ids_other = {d.document_id for d in result_other.documents}
+    assert r_fts.document_id in ids_first
+    assert ids_other == {r_other.document_id}
+
+
+def test_degraded_fts_only_result_is_not_cached(group, reports, settings):
+    r_fts, r_vec = reports
+    mocked, MockClient = _mock_embedding(settings.EMBEDDINGS_DIM)
+    try:
+        MockClient.return_value.embed_query.side_effect = EmbeddingClientError("down")
+        degraded = search(_make_search("pneumothorax", group.pk))
+        assert {d.document_id for d in degraded.documents} == {r_fts.document_id}
+
+        # Service recovers: the next search must not be pinned to the degraded union.
+        MockClient.return_value.embed_query.side_effect = None
+        recovered = search(_make_search("pneumothorax", group.pk))
+    finally:
+        mocked.stop()
+
+    assert r_vec.document_id in {d.document_id for d in recovered.documents}
+
+
+def test_zero_timeout_disables_the_cache(group, reports, settings):
+    settings.HYBRID_FUSED_CACHE_TIMEOUT_SECONDS = 0
+    mocked, _ = _mock_embedding(settings.EMBEDDINGS_DIM)
+    try:
+        search(_make_search("pneumothorax", group.pk))
+        with CaptureQueriesContext(connection) as ctx:
+            search(_make_search("pneumothorax", group.pk))
+    finally:
+        mocked.stop()
+
+    assert _ran_fusion(ctx.captured_queries)

radis/pgsearch/tests/test_providers.py

@@ -21,6 +21,7 @@
 
 import pytest
 from adit_radis_shared.accounts.factories import GroupFactory
+from django.core.cache import cache
 from django.utils import timezone
 
 from radis.pgsearch import providers
@@ -113,6 +114,10 @@ def test_updating_report_body_refreshes_search_vector():
 
     report.body = "revised findings about fracture"
     report.save()
+    # The fused union of the "pneumonia" search above is cached for
+    # HYBRID_FUSED_CACHE_TIMEOUT_SECONDS and would still list the report; this
+    # test is about the signal, not the cache's documented staleness window.
+    cache.clear()
 
     # The post_save signal re-saves the search vector with the new body.
     assert run_search("fracture") == [report.document_id]

radis/settings/base.py

@@ -565,6 +565,14 @@ def _resolve_embeddings_model() -> ModelSpec | None:
 DATABASES["default"].setdefault("OPTIONS", {})["options"] = (
     f"-c hnsw.ef_search={HYBRID_HNSW_EF_SEARCH}"
 )
+# How long a query's fused RRF union (ordered ids and scores) stays cached, in
+# seconds. Pagination re-runs the whole fusion for every page, and ranking every
+# FTS match of a common term costs seconds on a large corpus, so page 2..n and
+# repeated searches reuse the cached union instead. The trade is freshness:
+# reports created or updated within the TTL appear in a cached query's results
+# only after it expires. 0 disables the cache. Degraded FTS-only results
+# (configured embedding model, no vector) are never cached.
+HYBRID_FUSED_CACHE_TIMEOUT_SECONDS = _optional_env("HYBRID_FUSED_CACHE_TIMEOUT_SECONDS", int, 300)
 
 # Chat
 CHAT_GENERATE_TITLE_SYSTEM_PROMPT = """