Coverage for astrocyte/_multi_bank.py: 95%
127 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""Multi-bank recall orchestrator — parallel, cascade, and first-match strategies."""
3from __future__ import annotations
5import asyncio
6import logging
7from collections.abc import Awaitable, Callable
8from dataclasses import replace
10from astrocyte._log_safety import safe as _safe_log
11from astrocyte._recall_params import RecallParams
12from astrocyte.errors import ConfigError
13from astrocyte.policy.homeostasis import enforce_token_budget
14from astrocyte.policy.observability import MetricsCollector
15from astrocyte.types import (
16 MemoryHit,
17 MultiBankStrategy,
18 RecallRequest,
19 RecallResult,
20)
22logger = logging.getLogger("astrocyte")
24# Type aliases for the callbacks injected from Astrocyte
25RecallFn = Callable[[RecallRequest], Awaitable[RecallResult]]
26MakeRequestFn = Callable[
27 [str, str, int, int | None, list[str] | None, RecallParams],
28 Awaitable[RecallRequest],
29]
32def _bank_visit_order(bank_ids: list[str], cascade_order: list[str] | None) -> list[str]:
33 if not cascade_order:
34 return list(bank_ids)
35 out: list[str] = []
36 seen: set[str] = set()
37 for b in cascade_order:
38 if b in bank_ids and b not in seen:
39 out.append(b)
40 seen.add(b)
41 for b in bank_ids:
42 if b not in seen:
43 out.append(b)
44 seen.add(b)
45 return out
48def _dedupe_hits_by_text(hits: list[MemoryHit]) -> list[MemoryHit]:
49 """One hit per distinct text, keeping the highest-scoring instance."""
50 best: dict[str, MemoryHit] = {}
51 for h in hits:
52 prev = best.get(h.text)
53 if prev is None or h.score > prev.score:
54 best[h.text] = h
55 return sorted(best.values(), key=lambda x: x.score, reverse=True)
58def _apply_bank_weights(hits: list[MemoryHit], weights: dict[str, float] | None) -> list[MemoryHit]:
59 if not weights:
60 return list(hits)
61 out: list[MemoryHit] = []
62 for h in hits:
63 bid = h.bank_id or ""
64 w = float(weights.get(bid, 1.0))
65 out.append(replace(h, score=h.score * w))
66 return out
69def _tag_hits_with_bank(hits: list[MemoryHit], bank_id: str) -> list[MemoryHit]:
70 return [replace(h, bank_id=bank_id) if h.bank_id is None else h for h in hits]
73class MultiBankOrchestrator:
74 """Dispatches multi-bank recall across parallel, cascade, and first-match strategies."""
76 def __init__(
77 self,
78 *,
79 do_recall: RecallFn,
80 make_request: MakeRequestFn,
81 circuit_breaker_record_failure: Callable[[], None],
82 metrics: MetricsCollector,
83 provider_name: str,
84 ) -> None:
85 self._do_recall = do_recall
86 self._make_request = make_request
87 self._cb_record_failure = circuit_breaker_record_failure
88 self._metrics = metrics
89 self._provider_name = provider_name
91 async def recall(
92 self,
93 query: str,
94 bank_ids: list[str],
95 max_results: int,
96 max_tokens: int | None,
97 tags: list[str] | None,
98 params: RecallParams,
99 strategy: MultiBankStrategy,
100 ) -> RecallResult:
101 """Multi-bank recall — strategy dispatch."""
102 if strategy.mode == "parallel":
103 return await self._parallel(query, bank_ids, max_results, max_tokens, tags, params, strategy)
104 if strategy.mode == "cascade":
105 return await self._cascade(query, bank_ids, max_results, max_tokens, tags, params, strategy)
106 if strategy.mode == "first_match":
107 return await self._first_match(query, bank_ids, max_results, max_tokens, tags, params, strategy)
108 raise ConfigError(f"Unknown multi-bank mode: {strategy.mode!r}")
110 async def _parallel(
111 self,
112 query: str,
113 bank_ids: list[str],
114 max_results: int,
115 max_tokens: int | None,
116 tags: list[str] | None,
117 params: RecallParams,
118 strategy: MultiBankStrategy,
119 ) -> RecallResult:
120 reqs: list[RecallRequest] = []
121 for bid in bank_ids:
122 reqs.append(await self._make_request(query, bid, max_results, None, tags, params))
123 tasks = [self._do_recall(r) for r in reqs]
125 try:
126 results = await asyncio.wait_for(
127 asyncio.gather(*tasks, return_exceptions=True),
128 timeout=30.0,
129 )
130 except asyncio.TimeoutError:
131 logger.error(
132 "Multi-bank parallel recall timed out after 30s for banks: %s",
133 _safe_log(",".join(bank_ids)),
134 )
135 self._metrics.inc_counter(
136 "astrocyte_multi_bank_recall_timeout_total",
137 {"bank_ids": ",".join(bank_ids)},
138 )
139 return RecallResult(hits=[], total_available=0, truncated=False)
141 all_hits: list[MemoryHit] = []
142 total_available = 0
143 for bid, result in zip(bank_ids, results):
144 if isinstance(result, RecallResult):
145 all_hits.extend(_tag_hits_with_bank(result.hits, bid))
146 total_available += result.total_available
147 elif isinstance(result, BaseException):
148 logger.warning(
149 "Multi-bank recall failed for bank '%s': %s",
150 _safe_log(bid),
151 _safe_log(result),
152 )
153 self._cb_record_failure()
154 self._metrics.inc_counter(
155 "astrocyte_recall_total",
156 {"bank_id": bid, "provider": self._provider_name, "status": "error"},
157 )
159 weighted = _apply_bank_weights(all_hits, strategy.bank_weights)
160 weighted.sort(key=lambda h: h.score, reverse=True)
162 if strategy.dedup_across_banks:
163 deduped = _dedupe_hits_by_text(weighted)
164 else:
165 deduped = weighted
167 trimmed = deduped[:max_results]
168 truncated = False
169 if max_tokens:
170 trimmed, truncated = enforce_token_budget(trimmed, max_tokens)
172 return RecallResult(hits=trimmed, total_available=total_available, truncated=truncated)
174 async def _cascade(
175 self,
176 query: str,
177 bank_ids: list[str],
178 max_results: int,
179 max_tokens: int | None,
180 tags: list[str] | None,
181 params: RecallParams,
182 strategy: MultiBankStrategy,
183 ) -> RecallResult:
184 order = _bank_visit_order(bank_ids, strategy.cascade_order)
185 accumulated: list[MemoryHit] = []
186 total_available = 0
188 for bid in order:
189 result = await self._do_recall(
190 await self._make_request(query, bid, max_results, None, tags, params),
191 )
192 total_available += result.total_available
193 accumulated.extend(_tag_hits_with_bank(result.hits, bid))
195 merged_for_stop = _dedupe_hits_by_text(accumulated) if strategy.dedup_across_banks else list(accumulated)
196 if len(merged_for_stop) >= strategy.min_results_to_stop:
197 break
199 working = _dedupe_hits_by_text(accumulated) if strategy.dedup_across_banks else accumulated
200 weighted = _apply_bank_weights(working, strategy.bank_weights)
201 weighted.sort(key=lambda h: h.score, reverse=True)
202 trimmed = weighted[:max_results]
203 truncated = False
204 if max_tokens:
205 trimmed, truncated = enforce_token_budget(trimmed, max_tokens)
206 return RecallResult(hits=trimmed, total_available=total_available, truncated=truncated)
208 async def _first_match(
209 self,
210 query: str,
211 bank_ids: list[str],
212 max_results: int,
213 max_tokens: int | None,
214 tags: list[str] | None,
215 params: RecallParams,
216 strategy: MultiBankStrategy,
217 ) -> RecallResult:
218 order = _bank_visit_order(bank_ids, strategy.cascade_order)
219 total_available = 0
220 for bid in order:
221 result = await self._do_recall(
222 await self._make_request(query, bid, max_results, None, tags, params),
223 )
224 total_available += result.total_available
225 if result.hits:
226 hits = _tag_hits_with_bank(result.hits, bid)
227 hits = hits[:max_results]
228 hits = _apply_bank_weights(hits, strategy.bank_weights)
229 hits.sort(key=lambda h: h.score, reverse=True)
230 truncated = False
231 if max_tokens:
232 hits, truncated = enforce_token_budget(hits, max_tokens)
233 return RecallResult(hits=hits, total_available=total_available, truncated=truncated)
234 return RecallResult(hits=[], total_available=total_available, truncated=False)