Coverage for astrocyte/pipeline/reflect_stage.py: 73%
170 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1from __future__ import annotations
3import logging
4from typing import TYPE_CHECKING
6from astrocyte.pipeline.cross_encoder_rerank import (
7 cross_encoder_rerank,
8)
9from astrocyte.pipeline.embedding import generate_embeddings
10from astrocyte.pipeline.fusion import (
11 ScoredItem,
12)
13from astrocyte.pipeline.query_plan import build_query_plan
14from astrocyte.pipeline.reflect import synthesize
15from astrocyte.pipeline.reranking import (
16 cross_encoder_like_rerank,
17)
18from astrocyte.policy.homeostasis import enforce_token_budget
19from astrocyte.recall.authority import apply_recall_authority
20from astrocyte.types import (
21 MemoryHit,
22 RecallRequest,
23 ReflectRequest,
24 ReflectResult,
25 VectorItem,
26)
28if TYPE_CHECKING:
29 pass
33from astrocyte.pipeline._orchestrator_common import (
34 _abstention_floor_for_skepticism,
35 _resolve_skepticism_for_abstention,
36 _source_ids_from_metadata,
37)
39_logger = logging.getLogger("astrocyte.mip")
42class ReflectStageMixin:
43 """Reflect pipeline: retrieve context → synthesize an answer.
45 A behavior slice of PipelineOrchestrator, composed into it as a mixin.
46 ``self`` is the full orchestrator at runtime (attributes come from its
47 ``__init__`` / ``apply_config``); shared helpers live on sibling mixins.
48 """
50 async def reflect(self, request: ReflectRequest) -> ReflectResult:
51 """Reflect pipeline: recall → LLM synthesis."""
52 query_plan = build_query_plan(request.query)
54 # 1. Run recall with larger result set. Aggregate/multi-hop queries need
55 # more candidate memories so synthesis can combine facts instead of
56 # answering from the first plausible hit.
57 recall_request = RecallRequest(
58 query=request.query,
59 bank_id=request.bank_id,
60 max_results=query_plan.recall_max_results,
61 max_tokens=request.max_tokens,
62 tags=request.tags,
63 as_of=request.as_of,
64 query_reference_date=request.query_reference_date,
65 )
66 recall_result = await self.recall(recall_request)
67 expanded_hits = await self._expand_reflect_sources(
68 request.bank_id,
69 self._rank_reflect_context(
70 request.query,
71 recall_result.hits,
72 limit=query_plan.reflect_rank_limit,
73 ),
74 limit=query_plan.reflect_expand_limit,
75 tags=request.tags,
76 )
77 recall_result.hits = expanded_hits
78 if request.max_tokens:
79 recall_result.hits, expanded_truncated = enforce_token_budget(
80 recall_result.hits,
81 request.max_tokens,
82 )
83 recall_result.truncated = recall_result.truncated or expanded_truncated
85 path_ctx = self._entity_path_authority_context(recall_result.hits)
86 auth_ctx: str | None = path_ctx
87 ra = self.recall_authority
88 if ra and ra.enabled and ra.apply_to_reflect:
89 recall_result = apply_recall_authority(recall_result, ra)
90 auth_ctx = "\n\n".join(part for part in (path_ctx, recall_result.authority_context) if part)
92 # 2. Resolve per-bank ReflectSpec from MIP (Phase 2, Step 9)
93 mip_reflect = None
94 if self.mip_router is not None:
95 bank_pipeline = self.mip_router.resolve_pipeline_for_bank(request.bank_id)
96 if bank_pipeline is not None:
97 mip_reflect = bank_pipeline.reflect
99 # 2b. Auto-select a prompt when no MIP prompt override is set. Priority
100 # (highest → lowest):
101 # 1. MIP explicit prompt (always wins — never overridden here)
102 # 2. Evidence-strict gate: when top raw semantic score < threshold,
103 # retrieval is uncertain. Force citation to prevent the LLM from
104 # constructing answers from tangential memories — the primary
105 # adversarial failure mode in open-domain benchmarks. Uses cosine
106 # similarity from recall(), which is set to 0.0 when no semantic
107 # results were found.
108 # 3. query_plan prompt variant (pre-retrieval query-shape routing)
109 # 4. _auto_prompt_variant fallback (legacy temporal/inference heuristic)
110 _EVIDENCE_STRICT_THRESHOLD = 0.5
111 if mip_reflect is None or mip_reflect.prompt is None:
112 from astrocyte.mip.schema import ReflectSpec
113 from astrocyte.pipeline.reflect import _auto_prompt_variant
115 prompt_variant = query_plan.prompt_variant or _auto_prompt_variant(request.query)
117 # Evidence-strict override: weak retrieval scores mean the top
118 # semantic match is only tangentially related — upgrade to citation
119 # mode so the LLM admits uncertainty instead of hallucinating.
120 # Skip if query_plan already chose evidence_strict (adversarial
121 # query shape) — no need to double-set.
122 if recall_result.top_semantic_score < _EVIDENCE_STRICT_THRESHOLD and prompt_variant != "evidence_strict":
123 prompt_variant = "evidence_strict"
125 if prompt_variant is not None:
126 if mip_reflect is None:
127 mip_reflect = ReflectSpec(prompt=prompt_variant)
128 else:
129 # Preserve any other MIP settings (e.g. promote_metadata)
130 mip_reflect = ReflectSpec(
131 prompt=prompt_variant,
132 promote_metadata=mip_reflect.promote_metadata,
133 )
135 # 2b1. Adversarial-defense gate: premise verification.
136 # Decompose the question into atomic claims, verify each
137 # against memory. Short-circuit to "insufficient evidence"
138 # when ANY presupposition fails. Targets false-premise and
139 # negative-existence adversarial questions.
140 if self.adversarial_premise_verification_enabled and self.llm_provider is not None:
141 try:
142 from astrocyte.pipeline.premise_verification import verify_question
144 async def _premise_recall(claim: str, max_results: int) -> list[MemoryHit]:
145 sub_request = RecallRequest(
146 query=claim,
147 bank_id=request.bank_id,
148 max_results=max_results,
149 max_tokens=request.max_tokens,
150 tags=request.tags,
151 as_of=request.as_of,
152 query_reference_date=request.query_reference_date,
153 )
154 sub_result = await self.recall(sub_request)
155 return sub_result.hits
157 verification = await verify_question(
158 request.query,
159 recall_fn=_premise_recall,
160 llm_provider=self.llm_provider,
161 min_confidence=self.adversarial_premise_min_confidence,
162 )
163 short_circuit = verification.short_circuit_message()
164 if short_circuit is not None:
165 _logger.info(
166 "reflect: premise verification short-circuit — %s",
167 short_circuit,
168 )
169 return ReflectResult(answer=short_circuit, sources=[])
170 except Exception as exc:
171 _logger.warning(
172 "premise verification failed (%s); continuing without the guard.",
173 exc,
174 )
176 # 2c. Adversarial-defense gate: score-floor abstention.
177 # Distinct from the evidence-strict prompt switch above.
178 # ``evidence_strict`` keeps invoking the LLM with a tighter
179 # prompt; the abstention floor short-circuits BEFORE any LLM
180 # call when retrieval is so weak that the question almost
181 # certainly has no answer in memory. Targets adversarial
182 # questions (negative existence, false premise, time-shift)
183 # where the LLM left to its own devices invents an answer.
184 #
185 # The effective floor is now derived from the request's
186 # ``dispositions.skepticism`` (1=trust→never abstain,
187 # 5=skeptical→aggressive abstention) rather than the legacy
188 # global ``abstention_enabled`` bool. This lets one deployment
189 # serve both adversarial-resistant agents and answer-everything
190 # assistants from a single config — see
191 # ``_abstention_floor_for_skepticism`` for the mapping. Legacy
192 # ``abstention_enabled`` remains as a fallback when no
193 # per-call dispositions are supplied.
194 skepticism = _resolve_skepticism_for_abstention(
195 request.dispositions,
196 self.adversarial_abstention_enabled,
197 )
198 effective_floor = _abstention_floor_for_skepticism(
199 skepticism,
200 self.adversarial_abstention_floor,
201 )
202 if (
203 effective_floor > 0.0
204 and recall_result.top_semantic_score < effective_floor
205 and (not recall_result.hits or all((h.score or 0.0) < effective_floor for h in recall_result.hits[:5]))
206 ):
207 _logger.info(
208 "reflect: abstention floor triggered (top_semantic=%.3f < "
209 "%.3f, skepticism=%d); returning 'insufficient evidence' "
210 "without LLM call.",
211 recall_result.top_semantic_score,
212 effective_floor,
213 skepticism,
214 )
215 return ReflectResult(
216 answer="insufficient evidence: no memory supports this question.",
217 sources=[],
218 )
220 # 3. Synthesize. Two paths:
221 #
222 # (a) Agentic loop (Hindsight parity) — when configured, the LLM
223 # selects between ``recall`` and ``done`` over up to N
224 # iterations. Targets multi-hop / open-domain where a
225 # single-shot recall misses bridge memories. Each loop
226 # ``recall`` reuses the full upgraded recall pipeline (RRF
227 # + spread + cross-encoder rerank + scope tags).
228 #
229 # (b) Single-shot synthesis — original path. Faster, simpler.
230 synthesize_kwargs = dict(
231 dispositions=request.dispositions,
232 max_tokens=request.max_tokens or 2048,
233 authority_context=auth_ctx,
234 mip_reflect=mip_reflect,
235 # Forward the relative-phrase anchor so the synthesis
236 # prompt's ``<reference_date>`` block reflects the
237 # question's contemporaneous date, not the run wall-clock.
238 query_reference_date=request.query_reference_date,
239 )
240 if self.agentic_reflect_params is not None:
241 from astrocyte.pipeline.agentic_reflect import agentic_reflect
243 async def _loop_recall(query: str, max_results: int) -> list[MemoryHit]:
244 # Reuses everything we just shipped: spread, cross-
245 # encoder rerank, tag scoping, RRF — same recall the
246 # outer step ran, parameterized by the agent's refined
247 # query and its requested max_results. Forwards the
248 # outer ``as_of`` (audit replay) AND
249 # ``query_reference_date`` (relative-phrase anchor) so
250 # every sub-recall the agent makes sees the same
251 # temporal context as the seed reflect.
252 sub_request = RecallRequest(
253 query=query,
254 bank_id=request.bank_id,
255 max_results=max_results,
256 max_tokens=request.max_tokens,
257 tags=request.tags,
258 as_of=request.as_of,
259 query_reference_date=request.query_reference_date,
260 )
261 sub_result = await self.recall(sub_request)
262 return sub_result.hits
264 # ``search_observations`` tool — searches the consolidated
265 # observation layer when the consolidator is wired up. The
266 # ::obs bank scope is reused to mirror how observations are
267 # stored at retain time.
268 observations_fn = None
269 if self._observation_consolidator is not None:
271 async def _loop_observations(query: str, max_results: int) -> list[MemoryHit]:
272 qvec_batch = await generate_embeddings([query], self.llm_provider)
273 qvec = qvec_batch[0] if qvec_batch else []
274 # ``request.as_of`` (when set) propagates through to
275 # observation search so time-travel reflect sees the
276 # same observation snapshot as raw recall. When
277 # unset, search uses the bank's current state.
278 obs_results = await self._retrieve_observations(
279 qvec,
280 request.bank_id,
281 max_results,
282 request.as_of,
283 )
284 # Convert ScoredItem → MemoryHit so the loop can
285 # cite IDs uniformly across tools.
286 return [
287 MemoryHit(
288 text=item.text,
289 score=item.score,
290 fact_type=item.fact_type,
291 metadata=item.metadata,
292 tags=item.tags,
293 memory_id=item.id,
294 bank_id=request.bank_id,
295 memory_layer="observation",
296 )
297 for item in (obs_results or [])
298 ]
300 observations_fn = _loop_observations
302 # ``expand`` tool — fetch source memories cited by a
303 # compiled fact / wiki / observation. Reuses the existing
304 # source-expansion path with a single-id seed.
305 async def _loop_expand(memory_id: str, max_sources: int) -> list[MemoryHit]:
306 # Find the seed hit in the running pool to expand from.
307 seed: MemoryHit | None = next(
308 (h for h in recall_result.hits if h.memory_id == memory_id),
309 None,
310 )
311 if seed is None:
312 return []
313 expanded = await self._expand_reflect_sources(
314 request.bank_id,
315 [seed],
316 limit=max(1, max_sources) + 1,
317 tags=request.tags,
318 )
319 # Drop the seed itself so the model sees only NEW evidence.
320 return [h for h in expanded if h.memory_id != memory_id]
322 # ``search_mental_models`` tool — top of the hierarchy.
323 # Forwards to ``MentalModelStore.list(bank_id, scope=...)``,
324 # converts each MentalModel to a MemoryHit so the agent can
325 # cite ``model_id`` like any other ``memory_id``. Returns
326 # all models in scope (no semantic ranking — mental models
327 # are usually a handful per bank, the LLM picks).
328 mental_models_fn = None
329 if self.mental_model_service is not None:
331 async def _loop_mental_models(
332 query: str,
333 scope: str | None,
334 ) -> list[MemoryHit]:
335 try:
336 models = await self.mental_model_service.list( # type: ignore[union-attr]
337 request.bank_id,
338 scope=scope,
339 )
340 except Exception as exc:
341 _logger.warning(
342 "agentic_reflect.mental_models list failed (%s)",
343 exc,
344 )
345 return []
346 return [
347 MemoryHit(
348 text=f"# {m.title}\n\n{m.content}",
349 score=1.0,
350 fact_type="mental_model",
351 metadata={
352 "scope": m.scope,
353 "revision": m.revision,
354 "refreshed_at": m.refreshed_at.isoformat(),
355 },
356 memory_id=m.model_id,
357 bank_id=m.bank_id,
358 memory_layer="mental_model",
359 )
360 for m in models
361 ]
363 mental_models_fn = _loop_mental_models
365 return await agentic_reflect(
366 request.query,
367 initial_hits=recall_result.hits,
368 recall_fn=_loop_recall,
369 observations_fn=observations_fn,
370 expand_fn=_loop_expand,
371 mental_models_fn=mental_models_fn,
372 llm_provider=self.llm_provider,
373 params=self.agentic_reflect_params,
374 final_synthesize_fn=synthesize,
375 final_synthesize_kwargs=synthesize_kwargs,
376 )
377 return await synthesize(
378 query=request.query,
379 hits=recall_result.hits,
380 llm_provider=self.llm_provider,
381 **synthesize_kwargs,
382 )
384 def _rank_reflect_context(
385 self,
386 query: str,
387 hits: list[MemoryHit],
388 *,
389 limit: int,
390 ) -> list[MemoryHit]:
391 """Apply final precision rerank and hierarchy before synthesis.
393 Uses the configured cross-encoder reranker (Hindsight parity) when
394 ``self.cross_encoder`` is set; otherwise falls back to the
395 deterministic heuristic ``cross_encoder_like_rerank``. The
396 heuristic remains the default so reflect stays dependency-free
397 unless the operator opts into the cross-encoder backend.
398 """
399 if not hits:
400 return hits
402 items = [
403 ScoredItem(
404 id=h.memory_id or f"hit-{idx}",
405 text=h.text,
406 score=h.score,
407 fact_type=h.fact_type,
408 metadata=h.metadata,
409 tags=h.tags,
410 memory_layer=h.memory_layer,
411 occurred_at=h.occurred_at,
412 retained_at=h.retained_at,
413 )
414 for idx, h in enumerate(hits)
415 ]
417 if self.cross_encoder is not None:
418 try:
419 scored = cross_encoder_rerank(
420 items,
421 query,
422 model=self.cross_encoder,
423 top_k=self.cross_encoder_top_k,
424 )
425 except Exception as exc: # pragma: no cover — defensive
426 _logger.warning(
427 "cross-encoder rerank failed (%s); falling back to heuristic.",
428 exc,
429 )
430 scored = cross_encoder_like_rerank(items, query)
431 else:
432 scored = cross_encoder_like_rerank(items, query)
434 hit_by_id = {h.memory_id or f"hit-{idx}": h for idx, h in enumerate(hits)}
435 return [hit_by_id[item.id] for item in scored[:limit] if item.id in hit_by_id]
437 def _entity_path_authority_context(self, hits: list[MemoryHit]) -> str | None:
438 path_lines: list[str] = []
439 direct_lines: list[str] = []
440 for idx, hit in enumerate(hits, 1):
441 path = (hit.metadata or {}).get("_entity_path") if hit.metadata else None
442 if path:
443 path_lines.append(f"- Memory {idx}: entity_path={path}")
444 elif hit.score >= 0.5:
445 direct_lines.append(f"- Memory {idx}: direct_facts")
446 if not path_lines:
447 return None
448 sections = ["entity_path_evidence:", *path_lines]
449 if direct_lines:
450 sections.extend(["direct_facts:", *direct_lines[:8]])
451 sections.append("supporting_context: use entity-path evidence before unrelated semantic matches.")
452 return "\n".join(sections)
454 async def _expand_reflect_sources(
455 self,
456 bank_id: str,
457 hits: list[MemoryHit],
458 *,
459 limit: int,
460 tags: list[str] | None = None,
461 ) -> list[MemoryHit]:
462 """Append raw sources cited by top wiki/observation hits.
464 This mirrors Hindsight's reflect loop in a bounded, non-agentic form:
465 start from compiled/observation evidence, then expand to raw facts for
466 grounding before synthesis.
468 ``tags`` (optional): when set, only fetched memories carrying every
469 listed tag are appended. Closes the leak where a tag-scoped reflect
470 could otherwise pull cross-scope raw memories via a wiki page's
471 ``_wiki_source_ids`` metadata.
472 """
473 if not hits:
474 return hits
476 source_ids: list[str] = []
477 seen_sources: set[str] = set()
478 for hit in hits:
479 for sid in _source_ids_from_metadata(hit.metadata):
480 if sid not in seen_sources:
481 seen_sources.add(sid)
482 source_ids.append(sid)
483 if not source_ids:
484 return hits[:limit]
486 raw_hits = await self._fetch_memory_hits_by_id(bank_id, source_ids, tags=tags)
487 existing_ids = {h.memory_id for h in hits if h.memory_id}
488 expanded = list(hits)
489 for raw in raw_hits:
490 if raw.memory_id and raw.memory_id not in existing_ids:
491 expanded.append(raw)
492 existing_ids.add(raw.memory_id)
493 if len(expanded) >= limit:
494 break
495 return expanded[:limit]
497 async def _fetch_memory_hits_by_id(
498 self,
499 bank_id: str,
500 ids: list[str],
501 *,
502 tags: list[str] | None = None,
503 ) -> list[MemoryHit]:
504 target = set(ids)
505 found: dict[str, VectorItem] = {}
506 offset = 0
507 batch = 100
508 # Tag scoping: a fetched memory is kept only if it carries every
509 # tag in ``tags``. ``None``/empty disables the filter (legacy
510 # behavior). Comparison is case-insensitive to match recall's
511 # tag-filter convention.
512 required_tags = {str(t).lower() for t in tags} if tags else None
513 while target:
514 chunk = await self.vector_store.list_vectors(bank_id, offset=offset, limit=batch)
515 if not chunk:
516 break
517 for item in chunk:
518 if item.id not in target:
519 continue
520 if required_tags is not None:
521 item_tags = {str(t).lower() for t in (item.tags or [])}
522 if not required_tags.issubset(item_tags):
523 # ID matches but scope tags missing — drop and
524 # mark as resolved so we don't keep scanning.
525 target.discard(item.id)
526 continue
527 found[item.id] = item
528 target.discard(item.id)
529 if len(chunk) < batch:
530 break
531 offset += batch
533 return [
534 MemoryHit(
535 text=item.text,
536 score=0.0,
537 fact_type=item.fact_type,
538 metadata=item.metadata,
539 tags=item.tags,
540 memory_id=item.id,
541 bank_id=bank_id,
542 memory_layer=item.memory_layer or "raw",
543 occurred_at=item.occurred_at,
544 retained_at=item.retained_at,
545 )
546 for sid in ids
547 if (item := found.get(sid)) is not None
548 ]