Add bounded read recovery and reproducible robustness experiments

benchmark/run.py

@@ -228,6 +228,7 @@ async def _run_one(
             "surfaces_used": (result or {}).get("surfaces_used", []),
             "tool_trace": (result or {}).get("tool_trace", []),
             "evidence": (result or {}).get("evidence", []),
+            "recovery": (result or {}).get("recovery", {}),
             "completed_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
         }
         for passthrough in ("reference_answer", "metadata"):

datamind/agent/options.py

@@ -9,6 +9,8 @@
 """
 from __future__ import annotations
 
+import asyncio
+import time
 from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Any
@@ -42,6 +44,7 @@
 from datamind.core.model_clients import build_model_client
 from datamind.core.protocols import EmbeddingProvider, TextModelClient, ToolCallingModelClient
 from datamind.core.tools import ToolRegistry
+from datamind.core.recovery import recovery_summary
 
 from .base import AgentLoopConfig, AgentLoopProtocol
 from .loop_native import NativeAgentLoop
@@ -137,9 +140,22 @@ async def query(
         history: list[dict] | None = None,
         final_contract: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
-        return await self.loop.run_turn(
-            user_message=message, history=history, final_contract=final_contract,
-        )
+        parent = current_context()
+        ctx = (replace(parent, extra=dict(parent.extra)) if parent else
+               RequestContext.new(profile=self.ledger._profile if self.ledger else "default"))
+        ctx.extra["recovery_events"] = []
+        deadline = time.monotonic() + getattr(getattr(self.loop, "_cfg", None), "wall_clock_timeout_s", 300.)
+        ctx.extra["query_deadline"] = min(deadline, ctx.extra.get("query_deadline", deadline))
+        with bind_context(ctx):
+            try:
+                async with asyncio.timeout(max(0., ctx.extra["query_deadline"] - time.monotonic())):
+                    result = await self.loop.run_turn(
+                        user_message=message, history=history, final_contract=final_contract,
+                    )
+            finally:
+                if parent is not None:
+                    parent.extra.setdefault("recovery_events", []).extend(ctx.extra["recovery_events"])
+        return {**result, "recovery": recovery_summary(ctx.extra["recovery_events"])}
 
 
 @dataclass

datamind/capabilities/embedding/providers/openai_compatible.py

@@ -23,6 +23,7 @@
 from datamind.core.errors import ExternalServiceError
 from datamind.core.logging import get_logger
 from datamind.core.registry import embedding_registry
+from datamind.core.recovery import read_retry_budget_active
 
 _log = get_logger("embedding.openai_compatible")
 
@@ -212,6 +213,17 @@ async def _call(self, inputs: list[str]) -> list[list[float]]:
                         "embedding_dimension_detected",
                         extra={"model": self._model, "dim": self.dimension},
                     )
+                # Provider-reported usage only. Never infer tokens from text
+                # length or label an unavailable usage report as zero cost.
+                from datamind.core.logging import current_context
+                ctx = current_context()
+                usage = body.get("usage")
+                if ctx is not None and isinstance(usage, dict):
+                    tokens = usage.get("total_tokens")
+                    if isinstance(tokens, int) and not isinstance(tokens, bool) and tokens >= 0:
+                        ctx.extra.setdefault("embedding_usage", []).append({
+                            "model": self._model, "total_tokens": tokens,
+                        })
                 return vecs
             except httpx.HTTPStatusError as exc:
                 # Ordinary 4xx (especially 401/403) are deterministic.
@@ -229,6 +241,10 @@ async def _call(self, inputs: list[str]) -> list[list[float]]:
                 )
                 if not retriable:
                     raise
+                if read_retry_budget_active():
+                    # Let the enclosing reader/retriever spend its one budget;
+                    # don't multiply HTTP retries by tool/vector retries.
+                    raise
                 if attempt < self._max_retries:
                     retry_after = None
                     if resp is not None and resp.headers.get("retry-after"):

datamind/capabilities/kb/providers/hybrid_retriever.py

@@ -18,6 +18,7 @@
 from datamind.core.logging import get_logger
 from datamind.core.protocols import EmbeddingProvider, RetrievedChunk, VectorStore
 from datamind.core.registry import retriever_registry
+from datamind.core.recovery import recover_read
 from datamind.capabilities.kb.filters import matches_metadata, validate_metadata_filter
 
 _log = get_logger("retriever.hybrid")
@@ -54,13 +55,25 @@ def __init__(
         vector_weight: float = 1.0,
         bm25_weight: float = 1.0,
         candidate_multiplier: int = 3,
+        vector_failure_fallback: str = "error",
+        vector_retries: int = 0,
+        vector_timeout_s: float | None = None,
     ) -> None:
         self._store = vector_store
         self._embed = embedding
         self._rrf_k = rrf_k
         self._vw = vector_weight
         self._bw = bm25_weight
         self._cm = candidate_multiplier
+        if vector_failure_fallback not in {"error", "bm25"}:
+            raise ValueError("vector_failure_fallback must be error or bm25")
+        if not 0 <= vector_retries <= 5:
+            raise ValueError("vector_retries must be between 0 and 5")
+        if vector_timeout_s is not None and vector_timeout_s <= 0:
+            raise ValueError("vector_timeout_s must be positive")
+        self._fallback = vector_failure_fallback
+        self._vector_retries = vector_retries
+        self._vector_timeout = vector_timeout_s
         # Lazy-built on first call; rebuildable.
         self._bm25: BM25Okapi | None = None
         self._ids: list[str] = []
@@ -109,157 +122,85 @@ async def rebuild_lexical(self) -> None:
         await self._ensure_lexical()
 
     async def aretrieve(
-        self,
-        query: str,
-        *,
-        top_k: int = 5,
+        self, query: str, *, top_k: int = 5,
         filters: dict[str, Any] | None = None,
     ) -> list[RetrievedChunk]:
-        validate_metadata_filter(filters)
-        await self._ensure_lexical()
-        k_inner = top_k * self._cm
-
-        # Vector side
-        vec = await self._embed.embed_query(query)
-        vec_hits = await self._store.query(vec, top_k=k_inner, where=filters)
-        vec_ranked = {ch.id: (rank, ch) for rank, ch in enumerate(vec_hits)}
-
-        # BM25 must use the exact same candidate scope as the vector branch.
-        # Filtering after fusion would allow an out-of-scope lexical hit to
-        # displace a valid vector candidate.
-        if self._bm25 is not None and self._ids:
-            scores = self._bm25.get_scores(_tokenize(query))
-            eligible = [
-                i for i in range(len(scores))
-                if matches_metadata(self._metas[i], filters)
-            ]
-            order = sorted(eligible, key=lambda i: -scores[i])[:k_inner]
-            bm25_ranked = {
-                self._ids[i]: (
-                    rank,
-                    RetrievedChunk(
-                        id=self._ids[i],
-                        text=self._texts[i],
-                        score=float(scores[i]),
-                        source=self._metas[i].get("source"),
-                        metadata=self._metas[i],
-                    ),
-                )
-                for rank, i in enumerate(order)
-                if scores[i] > 0
-            }
-        else:
-            bm25_ranked = {}
-
-        # Reciprocal Rank Fusion
-        fused: dict[str, tuple[float, RetrievedChunk]] = {}
-        for idmap, weight in ((vec_ranked, self._vw), (bm25_ranked, self._bw)):
-            for cid, (rank, ch) in idmap.items():
-                contrib = weight / (self._rrf_k + rank + 1)
-                if cid in fused:
-                    prev_score, prev_ch = fused[cid]
-                    fused[cid] = (prev_score + contrib, prev_ch)
-                else:
-                    fused[cid] = (contrib, ch)
-
-        top = sorted(fused.values(), key=lambda sc: -sc[0])[:top_k]
-        # Preserve fused score so callers can reason about it.
-        out = [
-            RetrievedChunk(
-                id=ch.id,
-                text=ch.text,
-                score=float(score),
-                source=ch.source,
-                metadata=ch.metadata,
-            )
-            for score, ch in top
-        ]
-        _log.info(
-            "retrieved",
-            extra={
-                "top_k": top_k,
-                "vec_hits": len(vec_ranked),
-                "bm25_hits": len(bm25_ranked),
-                "fused": len(out),
-            },
-        )
-        return out
+        return await self._retrieve(query, top_k=top_k, filters=filters, version=None)
 
     async def aretrieve_at_version(
-        self,
-        query: str,
-        *,
-        top_k: int = 5,
-        filters: dict[str, Any] | None = None,
-        version: int = None,
+        self, query: str, *, top_k: int = 5,
+        filters: dict[str, Any] | None = None, version: int | None = None,
     ) -> list[RetrievedChunk]:
-        validate_metadata_filter(filters)
-        _bm25, _ids, _texts, _metas = await self._ensure_lexical_at_version(version=version)
-        k_inner = top_k * self._cm
-
-        # Vector side
-        vec = await self._embed.embed_query(query)
-        vec_hits = await self._store.query(vec, top_k=k_inner, where=filters, version=version)
-        vec_ranked = {ch.id: (rank, ch) for rank, ch in enumerate(vec_hits)}
+        return await self._retrieve(query, top_k=top_k, filters=filters, version=version)
 
-        # BM25 must use the exact same candidate scope as the vector branch.
-        # Filtering after fusion would allow an out-of-scope lexical hit to
-        # displace a valid vector candidate.
-        if _bm25 is not None and _ids:
-            scores = _bm25.get_scores(_tokenize(query))
-            eligible = [
-                i for i in range(len(scores))
-                if matches_metadata(_metas[i], filters)
-            ]
-            order = sorted(eligible, key=lambda i: -scores[i])[:k_inner]
-            bm25_ranked = {
-                _ids[i]: (
-                    rank,
-                    RetrievedChunk(
-                        id=_ids[i],
-                        text=_texts[i],
-                        score=float(scores[i]),
-                        source=_metas[i].get("source"),
-                        metadata=_metas[i],
-                    ),
-                )
-                for rank, i in enumerate(order)
-                if scores[i] > 0
-            }
+    async def _retrieve(self, query, *, top_k, filters, version):
+        validate_metadata_filter(filters)
+        if top_k < 1:
+            raise ValueError("top_k must be positive")
+        if version is None:
+            await self._ensure_lexical()
+            bm25, ids, texts, metas = self._bm25, self._ids, self._texts, self._metas
         else:
-            bm25_ranked = {}
-
-        # Reciprocal Rank Fusion
+            # Lexical and vector reads must use exactly the same historical view.
+            bm25, ids, texts, metas = await self._ensure_lexical_at_version(version)
+        k_inner = top_k * self._cm
+        stage = "kb.embedding"
+
+        async def vector_search():
+            nonlocal stage
+            stage = "kb.embedding"
+            vec = await self._embed.embed_query(query)
+            stage = "kb.vector_query"
+            kwargs = {"top_k": k_inner, "where": filters}
+            if version is not None:
+                kwargs["version"] = version
+            return await self._store.query(vec, **kwargs)
+
+        def rank_lexical():
+            if bm25 is not None and ids:
+                scores = bm25.get_scores(_tokenize(query))
+                eligible = [i for i in range(len(scores)) if matches_metadata(metas[i], filters)]
+                order = sorted(eligible, key=lambda i: -scores[i])[:k_inner]
+                return {
+                    ids[i]: (rank, RetrievedChunk(
+                        id=ids[i], text=texts[i], score=float(scores[i]),
+                        source=metas[i].get("source"), metadata=metas[i],
+                    )) for rank, i in enumerate(order) if scores[i] > 0
+                }
+            else:
+                return {}
+
+        fallback_ranked = {}
+        async def lexical_fallback():
+            nonlocal fallback_ranked
+            fallback_ranked = rank_lexical()
+            return None  # Distinct from a successful vector search returning [].
+
+        vec_hits = await recover_read(
+            vector_search, component=lambda: stage, retries=self._vector_retries,
+            attempt_timeout_s=self._vector_timeout,
+            fallback=lexical_fallback if self._fallback == "bm25" else None,
+            fallback_method="bm25",
+        )
+        degraded = vec_hits is None
+        vec_ranked = {ch.id: (rank, ch) for rank, ch in enumerate(vec_hits or [])}
+        bm25_ranked = fallback_ranked if degraded else rank_lexical()
         fused: dict[str, tuple[float, RetrievedChunk]] = {}
-        for idmap, weight in ((vec_ranked, self._vw), (bm25_ranked, self._bw)):
-            for cid, (rank, ch) in idmap.items():
-                contrib = weight / (self._rrf_k + rank + 1)
+        # Lexical fallback remains useful even if lexical fusion was disabled.
+        lexical_weight = 1.0 if degraded else self._bw
+        for ranked, weight in ((vec_ranked, self._vw), (bm25_ranked, lexical_weight)):
+            for cid, (rank, ch) in ranked.items():
+                score = weight / (self._rrf_k + rank + 1)
                 if cid in fused:
-                    prev_score, prev_ch = fused[cid]
-                    fused[cid] = (prev_score + contrib, prev_ch)
+                    previous, old = fused[cid]
+                    fused[cid] = (previous + score, old)
                 else:
-                    fused[cid] = (contrib, ch)
-
-        top = sorted(fused.values(), key=lambda sc: -sc[0])[:top_k]
-        # Preserve fused score so callers can reason about it.
-        out = [
-            RetrievedChunk(
-                id=ch.id,
-                text=ch.text,
-                score=float(score),
-                source=ch.source,
-                metadata=ch.metadata,
-            )
-            for score, ch in top
-        ]
-        _log.info(
-            "retrieved",
-            extra={
-                "top_k": top_k,
-                "vec_hits": len(vec_ranked),
-                "bm25_hits": len(bm25_ranked),
-                "fused": len(out),
-                "version": version,
-            },
-        )
+                    fused[cid] = (score, ch)
+        out = [RetrievedChunk(
+            id=ch.id, text=ch.text, score=float(score), source=ch.source,
+            metadata={**ch.metadata, **({"degraded": True, "retrieval_method": "bm25"} if degraded else {})},
+        ) for score, ch in sorted(fused.values(), key=lambda item: -item[0])[:top_k]]
+        _log.info("retrieved", extra={"top_k": top_k, "vec_hits": len(vec_ranked),
+                  "bm25_hits": len(bm25_ranked), "fused": len(out),
+                  "version": version, "degraded": degraded})
         return out

datamind/capabilities/kb/service.py

@@ -309,6 +309,9 @@ def build_kb_service(
     elif strategy == "hybrid":
         retriever = retriever_registry.create(
             "hybrid", vector_store=vector_store, embedding=embedding,
+            vector_failure_fallback=settings.retrieval.vector_failure_fallback,
+            vector_retries=settings.retrieval.vector_retries,
+            vector_timeout_s=settings.retrieval.vector_timeout_s,
         )
     else:
         raise ConfigError(