Coverage for astrocyte/pipeline/tiered_retrieval.py: 97%
92 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""Adaptive tiered retrieval — progressive escalation from cache to agentic recall.
3Inspired by ByteRover's tiered approach: most queries resolve cheaply.
4Only novel/ambiguous queries escalate to full multi-strategy retrieval.
6Sync decision logic, async execution — Rust migration candidate for the sync parts.
7"""
9from __future__ import annotations
11from collections.abc import Awaitable, Callable
12from typing import TYPE_CHECKING
14from astrocyte.pipeline.provenance import stored_source
15from astrocyte.pipeline.recall_cache import RecallCache
16from astrocyte.pipeline.recent_buffer import RecentMemoryBuffer
17from astrocyte.recall.merge_result import merge_external_into_recall_result
18from astrocyte.types import MemoryHit, RecallRequest, RecallResult, RecallTrace
20if TYPE_CHECKING:
21 from astrocyte.pipeline.orchestrator import PipelineOrchestrator
23# Escalation target for tier 3+ (and tier 4 post-reformulation). Default: built-in pipeline recall.
24FullRecallFn = Callable[[RecallRequest], Awaitable[RecallResult]]
27class TieredRetriever:
28 """5-tier progressive retrieval — escalates only when needed.
30 Tier 0: Recall cache hit (~0ms, free)
31 Tier 1: Fuzzy text match on recent memories (~5ms, free)
32 Tier 2: BM25 keyword search only (~50ms, free)
33 Tier 3: Full multi-strategy parallel retrieval (~200ms, embedding cost)
34 Tier 4: Agentic recall — LLM reformulates query + retry (~3-10s, LLM cost)
35 """
37 def __init__(
38 self,
39 pipeline: PipelineOrchestrator,
40 recall_cache: RecallCache | None = None,
41 recent_buffer: RecentMemoryBuffer | None = None,
42 min_results: int = 3,
43 min_score: float = 0.3,
44 max_tier: int = 3,
45 *,
46 full_recall: FullRecallFn | None = None,
47 ) -> None:
48 self.pipeline = pipeline
49 self.recall_cache = recall_cache
50 self.recent_buffer = recent_buffer
51 self.min_results = min_results
52 self.min_score = min_score
53 self.max_tier = min(max_tier, 4)
54 # When None, tier 3+ uses ``pipeline.recall`` (pipeline-only tiering).
55 self._full_recall: FullRecallFn | None = full_recall
57 async def _invoke_full_recall(self, request: RecallRequest) -> RecallResult:
58 if self._full_recall is not None:
59 return await self._full_recall(request)
60 return await self.pipeline.recall(request)
62 def notify_retain(self, bank_id: str, memory_id: str, text: str, metadata: dict | None = None) -> None:
63 """Called after a successful retain to populate the recent buffer.
65 Should be called by the dispatcher or orchestrator for each stored chunk.
66 Invalidates the recall cache for the bank (contents changed).
67 """
68 if self.recent_buffer is not None:
69 self.recent_buffer.add(bank_id, memory_id, text, metadata)
70 if self.recall_cache is not None:
71 self.recall_cache.invalidate_bank(bank_id)
73 async def retrieve(self, request: RecallRequest) -> RecallResult:
74 """Run tiered retrieval, escalating until sufficient results found."""
75 from astrocyte.pipeline.embedding import generate_embeddings
77 # Federated / proxy hits must not use stale cache entries and are not cached themselves.
78 allow_cache = not request.external_context
80 # ── Tier 0: Cache ──
81 if self.recall_cache and self.max_tier >= 0 and allow_cache:
82 query_embeddings = await generate_embeddings([request.query], self.pipeline.llm_provider)
83 query_vector = query_embeddings[0]
85 cached = self.recall_cache.get(request.bank_id, query_vector)
86 if cached is not None:
87 return RecallResult(
88 hits=cached.hits[: request.max_results],
89 total_available=cached.total_available,
90 truncated=cached.truncated,
91 trace=RecallTrace(
92 strategies_used=["cache"],
93 total_candidates=len(cached.hits),
94 fusion_method="cache",
95 tier_used=0,
96 cache_hit=True,
97 ),
98 )
99 else:
100 query_vector = None
102 # ── Tier 1: Fuzzy text match on recent memories ──
103 if self.recent_buffer and self.max_tier >= 1:
104 fuzzy_hits = self.recent_buffer.search(
105 request.query,
106 request.bank_id,
107 limit=request.max_results * 2,
108 min_score=self.min_score,
109 )
110 if len(fuzzy_hits) >= self.min_results:
111 hits = fuzzy_hits[: request.max_results]
112 result = RecallResult(
113 hits=hits,
114 total_available=len(fuzzy_hits),
115 truncated=len(fuzzy_hits) > request.max_results,
116 trace=RecallTrace(
117 strategies_used=["fuzzy_recent"],
118 total_candidates=len(fuzzy_hits),
119 fusion_method="fuzzy_recent",
120 tier_used=1,
121 cache_hit=False,
122 ),
123 )
124 if request.external_context:
125 result = merge_external_into_recall_result(result, request.external_context, request.max_results)
126 if self.recall_cache and query_vector and allow_cache:
127 self.recall_cache.put(request.bank_id, query_vector, result)
128 return result
130 # ── Tier 2: BM25 keyword only ──
131 if self.pipeline.document_store and self.max_tier >= 2:
132 keyword_hits = await self.pipeline.document_store.search_fulltext(
133 request.query, request.bank_id, limit=request.max_results * 2
134 )
135 if len(keyword_hits) >= self.min_results:
136 avg_score = sum(h.score for h in keyword_hits) / max(len(keyword_hits), 1)
137 if avg_score >= self.min_score:
138 hits = [
139 MemoryHit(
140 text=h.text,
141 score=h.score,
142 metadata=h.metadata,
143 memory_id=h.document_id,
144 bank_id=request.bank_id,
145 source=stored_source(h.metadata),
146 )
147 for h in keyword_hits[: request.max_results]
148 ]
149 result = RecallResult(
150 hits=hits,
151 total_available=len(keyword_hits),
152 truncated=len(keyword_hits) > request.max_results,
153 trace=RecallTrace(
154 strategies_used=["keyword"],
155 total_candidates=len(keyword_hits),
156 fusion_method="bm25_only",
157 tier_used=2,
158 cache_hit=False,
159 ),
160 )
161 if request.external_context:
162 result = merge_external_into_recall_result(
163 result, request.external_context, request.max_results
164 )
165 # Cache the result for future queries
166 if self.recall_cache and query_vector and allow_cache:
167 self.recall_cache.put(request.bank_id, query_vector, result)
168 return result
170 # ── Tier 3: Full multi-strategy (pipeline recall or injected full recall, e.g. hybrid) ──
171 if self.max_tier >= 3:
172 result = await self._invoke_full_recall(request)
173 # Tag the trace
174 if result.trace:
175 result.trace.tier_used = 3
176 result.trace.cache_hit = False
178 # Check if sufficient
179 if len(result.hits) >= self.min_results or self.max_tier <= 3:
180 # Cache the result
181 if self.recall_cache and query_vector and allow_cache:
182 self.recall_cache.put(request.bank_id, query_vector, result)
183 return result
185 # ── Tier 4: Agentic recall — LLM reformulates query ──
186 if self.max_tier >= 4:
187 reformulated = await self._reformulate_query(request.query)
188 reformulated_request = RecallRequest(
189 query=reformulated,
190 bank_id=request.bank_id,
191 max_results=request.max_results,
192 max_tokens=request.max_tokens,
193 tags=request.tags,
194 fact_types=request.fact_types,
195 time_range=request.time_range,
196 include_sources=request.include_sources,
197 layer_weights=request.layer_weights,
198 detail_level=request.detail_level,
199 external_context=request.external_context,
200 )
201 result = await self._invoke_full_recall(reformulated_request)
202 if result.trace:
203 result.trace.tier_used = 4
204 result.trace.cache_hit = False
205 result.trace.strategies_used = (result.trace.strategies_used or []) + ["agentic_reformulation"]
207 if self.recall_cache and query_vector and allow_cache:
208 self.recall_cache.put(request.bank_id, query_vector, result)
209 return result
211 # Fallback: empty result (still surface federated hits if present)
212 empty = RecallResult(
213 hits=[],
214 total_available=0,
215 truncated=False,
216 trace=RecallTrace(tier_used=0, cache_hit=False),
217 )
218 if request.external_context:
219 return merge_external_into_recall_result(empty, request.external_context, request.max_results)
220 return empty
222 async def _reformulate_query(self, query: str) -> str:
223 """Use LLM to reformulate an ambiguous query for better retrieval."""
224 from astrocyte.types import Message
226 system_msg = (
227 "Reformulate the user's query to improve memory search results. "
228 "Add synonyms, expand abbreviations, and rephrase for clarity. "
229 "Return only the reformulated query, nothing else."
230 )
231 user_msg = f"<query>\n{query[:2000]}\n</query>"
232 try:
233 completion = await self.pipeline.llm_provider.complete(
234 messages=[
235 Message(role="system", content=system_msg),
236 Message(role="user", content=user_msg),
237 ],
238 max_tokens=200,
239 temperature=0.3,
240 )
241 return completion.text.strip() or query
242 except Exception:
243 return query # Fallback to original