Coverage for astrocyte/pipeline/orchestrator.py: 83%
221 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""Pipeline orchestrator — coordinates Tier 1 retain/recall/reflect flows.
3Async (coordinates I/O stages). See docs/_design/built-in-pipeline.md.
4"""
6from __future__ import annotations
8import asyncio
9import collections
10import inspect
11import logging
12import os
13import time
14from collections import defaultdict
15from contextlib import asynccontextmanager
16from typing import TYPE_CHECKING, Any
18from astrocyte.config import DedupConfig, ExtractionProfileConfig, RecallAuthorityConfig
20# Re-exported for backward compatibility: these pure module helpers moved to
21# _orchestrator_common during the retain/recall/reflect stage extraction but
22# remain importable from this module (some tests import them here directly).
23from astrocyte.pipeline._orchestrator_common import ( # noqa: E402,F401
24 _abstention_floor_for_skepticism,
25 _build_cooccurrence_pairs,
26 _deterministic_names,
27 _entities_from_metadata,
28 _resolve_skepticism_for_abstention,
29 _source_ids_from_metadata,
30 _warn_on_version_drift,
31)
32from astrocyte.pipeline.agentic_reflect import AgenticReflectParams
33from astrocyte.pipeline.chunking import DEFAULT_CHUNK_SIZE
34from astrocyte.pipeline.cross_encoder_rerank import (
35 CrossEncoderProtocol,
36)
37from astrocyte.pipeline.fusion import (
38 DEFAULT_RRF_K,
39)
40from astrocyte.pipeline.link_expansion import LinkExpansionParams
41from astrocyte.pipeline.recall_stage import RecallStageMixin
42from astrocyte.pipeline.reflect_stage import ReflectStageMixin
43from astrocyte.pipeline.retain_stage import RetainStageMixin
44from astrocyte.pipeline.sync_push import SyncPushStageMixin
45from astrocyte.policy.signal_quality import DedupDetector
46from astrocyte.types import (
47 Completion,
48 Message,
49)
51if TYPE_CHECKING:
52 from astrocyte.mip.router import MipRouter
53 from astrocyte.pipeline.entity_resolution import EntityResolver
54 from astrocyte.pipeline.pipeline_config import PipelineConfig
55 from astrocyte.provider import DocumentStore, GraphStore, LLMProvider, VectorStore, WikiStore
58# Public surface of this module. The _orchestrator_common helpers are
59# deliberate re-exports (moved during the stage-mixin extraction; tests and
60# older callers import them from here) — listing them in __all__ marks the
61# re-export as intentional for linters and CodeQL alike.
62__all__ = [ # noqa: RUF022 — grouped: class first, then re-exported helpers
63 "PipelineOrchestrator",
64 "_abstention_floor_for_skepticism",
65 "_build_cooccurrence_pairs",
66 "_deterministic_names",
67 "_entities_from_metadata",
68 "_resolve_skepticism_for_abstention",
69 "_source_ids_from_metadata",
70 "_warn_on_version_drift",
71]
73_logger = logging.getLogger("astrocyte.mip")
76class _RetainProfiler:
77 """Aggregate per-stage timings during retain for evidence-driven
78 bottleneck identification.
80 Disabled unless ``ASTROCYTE_RETAIN_PROFILE=1`` (env var) — the cost
81 of the time.monotonic() calls is small but non-zero and we don't
82 want it on by default in production. When enabled, the orchestrator
83 captures wall time for each suspected hot path (SFE LLM, embedding
84 generation, vector insert, entity merge, entity resolution) and
85 emits an aggregated p50/p95/max breakdown via :meth:`report`.
87 The samples accumulate across the lifetime of the orchestrator so a
88 single bench run produces one breakdown that covers all retain
89 calls. Call :meth:`reset` between batches if you want per-batch
90 isolation.
92 Why not Prometheus / OTel: those exist (``observability.*`` config),
93 but require running collectors and a separate analysis stack just
94 to find a bottleneck. This is a stop-gap for dev/bench
95 investigation — write data once, print at end of run, done.
96 """
98 def __init__(self) -> None:
99 self.enabled = os.environ.get("ASTROCYTE_RETAIN_PROFILE") == "1"
100 self.samples: dict[str, list[float]] = defaultdict(list)
102 @asynccontextmanager
103 async def time(self, stage: str):
104 """Async context manager that records elapsed wall time (ms)
105 under ``stage`` if profiling is enabled. No-op otherwise so the
106 production hot path stays clean."""
107 if not self.enabled:
108 yield
109 return
110 t0 = time.monotonic()
111 try:
112 yield
113 finally:
114 self.samples[stage].append((time.monotonic() - t0) * 1000.0)
116 def reset(self) -> None:
117 self.samples.clear()
119 def report(self, prefix: str = "[retain.profile]") -> None:
120 """Print p50/p95/max for every recorded stage, ordered by total
121 time (descending). The dominant stage is first — that's what
122 you want to optimize.
124 Uses ``print`` rather than the module logger so the breakdown
125 always reaches stdout regardless of how the caller configured
126 logging — this is dev/bench tooling output, not production
127 telemetry, and we want it to be impossible to lose."""
128 if not self.enabled:
129 print(f"{prefix} (profiler disabled — set ASTROCYTE_RETAIN_PROFILE=1)")
130 return
131 if not self.samples:
132 print(f"{prefix} (profiler enabled but captured no samples — instrumentation unwired?)")
133 return
134 import statistics
136 rows: list[tuple[str, int, float, float, float, float]] = []
137 for stage, samples in self.samples.items():
138 if not samples:
139 continue
140 rows.append(
141 (
142 stage,
143 len(samples),
144 sum(samples),
145 statistics.median(samples),
146 statistics.quantiles(samples, n=20)[18] if len(samples) >= 20 else max(samples),
147 max(samples),
148 ),
149 )
150 # Sort by total descending — biggest cost first.
151 rows.sort(key=lambda r: r[2], reverse=True)
152 print(f"{prefix} aggregate breakdown (sorted by total wall time):")
153 print(
154 f"{prefix} {'stage':<22} {'n':<7} {'total_ms':<12} {'p50_ms':<10} {'p95_ms':<10} {'max_ms':<10}",
155 )
156 for stage, n, total, p50, p95, mx in rows:
157 print(
158 f"{prefix} {stage:<22} {n:<7d} {total:<12.0f} {p50:<10.1f} {p95:<10.1f} {mx:<10.1f}",
159 )
162class _TrackingLLMProvider:
163 """Transparent wrapper that accumulates token usage from an LLMProvider.
165 All calls are forwarded to the underlying provider. After each ``complete()``
166 call, ``Completion.usage`` (if present) is added to ``tokens_used``.
167 Embedding calls do not consume completion tokens and are forwarded as-is.
169 This is used internally by :class:`PipelineOrchestrator` so that the
170 evaluation framework can report ``total_tokens_used`` per run without
171 modifying individual pipeline modules. See ``evaluation.md`` §2.2.
172 """
174 SPI_VERSION: int = 1
176 def __init__(self, inner: LLMProvider) -> None:
177 self._inner = inner
178 self.tokens_used: int = 0
180 async def complete(
181 self,
182 messages: list[Message],
183 model: str | None = None,
184 max_tokens: int = 1024,
185 temperature: float = 0.0,
186 tools: list | None = None, # list[ToolDefinition] — kept loose to avoid import cycle
187 tool_choice: str | None = None,
188 response_format: dict | None = None,
189 ) -> Completion:
190 # Forward native function-calling kwargs through to the underlying
191 # provider so the agentic reflect loop (Hindsight parity) can use
192 # ``tools=``/``tool_choice=`` end-to-end. Without this, the loop
193 # silently fell back to forced single-shot synthesis on every
194 # call — observed in 1986/1986 questions on the 2026-05-01 bench.
195 #
196 # Backward compat: only thread the new kwargs when the caller
197 # actually supplied them. Legacy providers / test fakes whose
198 # ``complete()`` signatures predate the tools/response_format
199 # extensions keep working when invoked via the old text-only
200 # path.
201 extra_kwargs: dict = {}
202 if tools is not None:
203 extra_kwargs["tools"] = tools
204 if tool_choice is not None:
205 extra_kwargs["tool_choice"] = tool_choice
206 if response_format is not None:
207 extra_kwargs["response_format"] = response_format
208 result = await self._inner.complete(
209 messages,
210 model=model,
211 max_tokens=max_tokens,
212 temperature=temperature,
213 **extra_kwargs,
214 )
215 if result.usage:
216 self.tokens_used += result.usage.input_tokens + result.usage.output_tokens
217 # Tier-3 observability: record cost + per-phase tokens when a
218 # benchmark collector is attached. Silent no-op otherwise.
219 collector = getattr(self, "_metrics_collector", None)
220 if collector is not None:
221 try:
222 collector.record_completion_call(
223 model=getattr(result, "model", None) or "",
224 input_tokens=result.usage.input_tokens,
225 output_tokens=result.usage.output_tokens,
226 )
227 except Exception:
228 pass # never let metrics break a real call
229 return result
231 async def embed(
232 self,
233 texts: list[str],
234 model: str | None = None,
235 ) -> list[list[float]]:
236 result = await self._inner.embed(texts, model=model)
237 # Embeddings don't return usage objects; estimate input tokens from
238 # the rough heuristic of ~1 token per 4 chars per text (conservative).
239 collector = getattr(self, "_metrics_collector", None)
240 if collector is not None:
241 try:
242 est_tokens = sum(max(1, len(t) // 4) for t in texts)
243 collector.record_embedding_call(
244 model=model or "text-embedding-3-small",
245 tokens=est_tokens,
246 )
247 except Exception:
248 pass # metrics are best-effort; never block the embedding call
249 return result
251 def capabilities(self) -> Any:
252 return self._inner.capabilities()
254 def reset_tokens(self) -> int:
255 """Return accumulated tokens and reset counter to zero."""
256 total = self.tokens_used
257 self.tokens_used = 0
258 return total
261class PipelineOrchestrator(RetainStageMixin, RecallStageMixin, ReflectStageMixin, SyncPushStageMixin):
262 """Orchestrates the Tier 1 built-in intelligence pipeline.
264 Coordinates async stages: chunk → embed → store → retrieve → fuse → rerank.
265 """
267 def __init__(
268 self,
269 vector_store: VectorStore,
270 llm_provider: LLMProvider,
271 graph_store: GraphStore | None = None,
272 document_store: DocumentStore | None = None,
273 chunk_strategy: str = "sentence",
274 max_chunk_size: int = DEFAULT_CHUNK_SIZE,
275 rrf_k: int = DEFAULT_RRF_K,
276 semantic_overfetch: int = 5,
277 extraction_profiles: dict[str, ExtractionProfileConfig] | None = None,
278 *,
279 enable_temporal_retrieval: bool = True,
280 temporal_scan_cap: int = 500,
281 temporal_half_life_days: float = 7.0,
282 enable_intent_aware_recall: bool = True,
283 enable_multi_query_expansion: bool = True,
284 enable_hyde: bool = False,
285 wiki_store: WikiStore | None = None,
286 wiki_confidence_threshold: float = 0.7,
287 entity_resolver: EntityResolver | None = None,
288 enable_observation_consolidation: bool = True,
289 max_pending_consolidations: int = 32,
290 max_deferred_consolidations: int = 1000,
291 observation_weight: float = 0.0,
292 observation_injection_weight: float = 1.5,
293 multi_query_confidence_threshold: float = 0.72,
294 final_rerank_mode: str = "heuristic",
295 final_rerank_top_n: int = 30,
296 final_rerank_keep_n: int | None = None,
297 ) -> None:
298 self.vector_store = vector_store
299 self._tracker = _TrackingLLMProvider(llm_provider)
300 self.llm_provider: LLMProvider = self._tracker # type: ignore[assignment]
301 self.graph_store = graph_store
302 self.document_store = document_store
303 self.chunk_strategy = chunk_strategy
304 self.max_chunk_size = max_chunk_size
305 self.extraction_profiles = extraction_profiles
306 self.rrf_k = rrf_k
307 self.semantic_overfetch = semantic_overfetch
308 # Temporal retrieval knobs — see astrocyte.pipeline.retrieval for full
309 # semantics. Enabled by default: the strategy no-ops gracefully on
310 # banks whose vectors have no timestamps, so enabling it is safe.
311 self.enable_temporal_retrieval = enable_temporal_retrieval
312 self.temporal_scan_cap = temporal_scan_cap
313 self.temporal_half_life_days = temporal_half_life_days
314 # M9 BM25-IDF keyword strategy. When True (and the document store
315 # advertises ``search_fulltext_bm25``), the keyword leg routes
316 # through the materialized-view path with corpus IDF + length
317 # normalisation instead of the classic ``ts_rank_cd``. Wired by
318 # ``Astrocyte.set_pipeline`` from ``bm25_idf.enabled`` config.
319 self.bm25_idf_enabled: bool = False
320 #: M10 source-aware retain + recall. ``source_store`` is the
321 #: SourceStore instance (or ``None``); the three flags below
322 #: control behaviour. All wired by ``Astrocyte.set_pipeline`` from
323 #: the ``source_aware_retrieval`` config block.
324 self.source_store: object | None = None
325 self.source_retain_provenance: bool = False
326 self.source_chunk_expansion: bool = False
327 self.source_expansion_score_multiplier: float = 0.5
328 self.source_expansion_max_per_hit: int = 4
329 # Intent-aware recall: heuristic query classifier biases RRF
330 # weights per strategy. Conservative (always fuses all strategies
331 # even under bias), so enabling is safe — a misclassification
332 # degrades to a soft lean rather than a strategy drop.
333 self.enable_intent_aware_recall = enable_intent_aware_recall
334 # Multi-query expansion: decompose complex questions into sub-questions,
335 # recall for each independently, and merge via RRF. Improves multi-hop
336 # coverage at the cost of N-1 extra embedding + retrieval passes per query.
337 # Disabled by default; enable for multi-hop-heavy workloads.
338 # When enabled, the confidence gate (multi_query_confidence_threshold)
339 # suppresses expansion when the top raw semantic score already exceeds the
340 # threshold — avoiding costly sub-query decomposition when a direct answer
341 # is already retrieved with high confidence. Cosine similarity is used as
342 # the gate signal (not RRF scores, which are rank-based and not comparable
343 # across query runs). Default threshold 0.72 passes ~20-30% of queries.
344 self.enable_multi_query_expansion = enable_multi_query_expansion
345 self.multi_query_confidence_threshold: float = multi_query_confidence_threshold
346 # HyDE (R1): generate a hypothetical answer, embed it, and run a second
347 # semantic search pass with that vector. Disabled by default — adds one
348 # LLM call per recall. Enable for multi-hop / paraphrase-heavy workloads.
349 self.enable_hyde = enable_hyde
350 self._dedup = DedupDetector(similarity_threshold=0.95)
351 #: Retain-time near-duplicate check on/off; set from
352 #: ``signal_quality.dedup.enabled``. Off means no chunk is ever dropped
353 #: as a duplicate, MIP ``dedup`` rules included (they tune the check,
354 #: they cannot turn it on).
355 self.dedup_enabled: bool = True
356 #: What a duplicate chunk does to the retain when no matched MIP rule
357 #: sets ``dedup.action``; set from ``signal_quality.dedup.action``,
358 #: mapped to the pipeline's vocabulary (``DedupConfig.pipeline_action``).
359 self.dedup_action: str = "skip_chunk"
360 #: Also check each chunk against its nearest neighbours in the vector
361 #: store when the in-process cache has no match, so dedup holds across
362 #: processes (see ``RetainStageMixin._find_duplicate_chunks``). One
363 #: extra ``search_similar`` per chunk; set from
364 #: ``signal_quality.dedup.consult_store``.
365 self.dedup_consult_store: bool = True
366 #: Per-bank dedup settings from ``banks.<id>.signal_quality.dedup``,
367 #: replacing the four attributes above for that bank. Each entry is
368 #: already merged over the top-level block by the config loader.
369 self.dedup_by_bank: dict[str, DedupConfig] = {}
370 # Forget-cache invalidation: ``Astrocyte.forget`` calls this hook so
371 # the in-memory dedup cache doesn't keep matching against memories
372 # that are gone from the vector store. Without this, re-retain after
373 # forget silently returns ``stored=False, error="All chunks are
374 # near-duplicates"`` for similar content (see invalidate_dedup_cache
375 # below).
376 #: Per-stage retain timing aggregator. No-op unless
377 #: ``ASTROCYTE_RETAIN_PROFILE=1`` is set in the environment;
378 #: when enabled, retain_many wraps suspect call sites and the
379 #: caller can inspect the breakdown via ``profiler.report()``
380 #: or read raw samples from ``profiler.samples``.
381 self._profiler = _RetainProfiler()
382 #: Set by :meth:`astrocyte._astrocyte.Astrocyte.set_pipeline` when ``recall_authority`` is configured.
383 self.recall_authority: RecallAuthorityConfig | None = None
384 #: Set by :meth:`astrocyte._astrocyte.Astrocyte.set_pipeline` when
385 #: ``cross_encoder_rerank.enabled`` is true. ``None`` falls back to
386 #: the heuristic ``cross_encoder_like_rerank`` in ``_rank_reflect_context``.
387 self.cross_encoder: CrossEncoderProtocol | None = None
388 #: When ``cross_encoder`` is set, only the first ``cross_encoder_top_k``
389 #: candidates are scored to bound inference cost. Default mirrors the
390 #: config default (30).
391 self.cross_encoder_top_k: int = 30
392 #: Set by :meth:`astrocyte._astrocyte.Astrocyte.set_pipeline` when
393 #: ``link_expansion.enabled`` is true. ``None`` skips the
394 #: post-fusion 3-signal expansion. Replaced the old BFS-hop
395 #: ``spreading_activation_params`` per Hindsight C3 rewrite.
396 self.link_expansion_params: LinkExpansionParams | None = None
397 #: Adversarial-defense score-floor abstention. When the top
398 #: recall hit's score is below this floor, reflect short-circuits
399 #: to "insufficient evidence" without invoking the LLM. Targets
400 #: adversarial questions where the LLM left to its own devices
401 #: would hallucinate an answer from weak hits. Wired by
402 #: ``Astrocyte.set_pipeline`` from ``adversarial_defense`` config.
403 self.adversarial_abstention_enabled: bool = False
404 self.adversarial_abstention_floor: float = 0.2
405 #: Pre-loop premise verification — decompose question into atomic
406 #: claims, verify each. Wired below in the reflect path.
407 self.adversarial_premise_verification_enabled: bool = False
408 self.adversarial_premise_min_confidence: float = 0.6
409 #: Tighten the agentic-reflect system prompt with explicit
410 #: adversarial-defense rules ("insufficient evidence is always
411 #: a valid answer", premise-check, etc.).
412 self.adversarial_prompt_enabled: bool = False
413 #: Hindsight-parity causal-link extraction at retain time. When
414 #: enabled, one extra LLM call per record produces directional
415 #: ``causes`` edges. Wired by ``Astrocyte.set_pipeline``.
416 self.causal_links_enabled: bool = False
417 self.causal_max_pairs_per_memory: int = 4
418 self.causal_min_confidence: float = 0.7
419 #: Hindsight-parity semantic-kNN graph (C3a). When enabled, each
420 #: new memory at retain time gets ``MemoryLink(link_type="semantic")``
421 #: edges to its top-K most-similar existing memories above the
422 #: similarity threshold. Wired by ``Astrocyte.set_pipeline``.
423 self.semantic_link_graph_enabled: bool = False
424 self.semantic_link_graph_top_k: int = 5
425 self.semantic_link_graph_threshold: float = 0.7
426 #: Structured fact extraction at retain time (5-dim
427 #: what/when/where/who/why with embedded entities + caused_by
428 #: relations). When enabled, replaces chunk_text +
429 #: extract_entities + fact_causal_extraction with a single
430 #: LLM call producing structured facts. Each fact becomes one
431 #: memory.
432 self.structured_fact_extraction_enabled: bool = False
433 self.structured_fact_extraction_max_facts: int = 30
434 #: Mode: "verbatim" stores raw chunk text + metadata (preserves
435 #: vocabulary for embedding-match), "concise" stores LLM-paraphrased
436 #: structured facts. Default verbatim because of the recall_hit_rate
437 #: regression that concise paraphrasing causes.
438 self.structured_fact_extraction_mode: str = "verbatim"
439 #: Chunking strategy used by verbatim SFE pre-chunking. Defaults
440 #: to "paragraph" (which gives the LLM full session context for
441 #: metadata extraction); LoCoMo benchmark losses 2.5 pts overall
442 #: when set to "dialogue".
443 self.structured_fact_extraction_chunk_strategy: str = "paragraph"
444 #: Per-chunk character budget for verbatim SFE pre-chunking.
445 #: When ``None`` the SFE path falls through to the same
446 #: chunk-size resolver legacy retain uses (orchestrator default
447 #: 512 chars). Bumping this is the dominant retain throughput
448 #: lever for SFE — fewer chunks = fewer ``facts[]`` entries the
449 #: LLM has to emit per session, and gpt-4o-mini latency is
450 #: roughly linear in output tokens. 2048 measured ~4× on LME-
451 #: shaped traffic without accuracy regression on LoCoMo.
452 self.structured_fact_extraction_chunk_max_size: int | None = None
453 #: Per-chunk parallel verbatim extraction (Phase 3 of cost-
454 #: control port). When True, the SFE path sends one LLM call
455 #: per chunk in parallel rather than one batched call. Drops
456 #: cross-chunk causal_relations — pair with
457 #: ``causal_links.enabled=false``.
458 self.structured_fact_extraction_parallel_chunks: bool = False
459 #: Max in-flight LLM calls per session when
460 #: ``parallel_chunks`` is True.
461 self.structured_fact_extraction_parallel_chunks_max_concurrency: int = 6
462 #: Entity co-occurrence link cap (2026-05-06 retain-profile fix).
463 #: When ``enabled``, retain creates ``co_occurs`` links between
464 #: at most ``max_entities`` entities per memory — bounding the
465 #: Cartesian product to ``C(K,2)`` per retain regardless of N.
466 #: Profiling on LME measured the unbounded path at 34% of
467 #: retain wall with O(N²) drift; capping at K=5 brings it to
468 #: <1% steady state.
469 self.entity_cooccurrence_enabled: bool = True
470 self.entity_cooccurrence_max_entities: int = 5
471 #: Query-level temporal constraint extraction. When enabled,
472 #: recall parses temporal expressions in the query into a
473 #: time_range filter applied to retrieval. Regex pre-pass is
474 #: free; ``allow_llm_fallback`` opt-in adds 1 LLM call per
475 #: temporal-marker query.
476 self.query_analyzer_enabled: bool = False
477 self.query_analyzer_allow_llm_fallback: bool = True
478 #: M18a-1 — extended temporal-expansion pattern set in the regex
479 #: pre-pass (word-numbers, "a few X ago", "the other day", "this/
480 #: earlier this <unit>", "recently/lately"). Default False.
481 #: Flipped via per-bank config or env override
482 #: ``ASTROCYTE_M18_ENABLE_TEMPORAL_EXPANSION=1`` for bench runs.
483 self.query_analyzer_enable_temporal_expansion: bool = False
484 #: Hindsight-parity agentic reflect loop. ``None`` = single-shot
485 #: synthesis (legacy path). Set by ``Astrocyte.set_pipeline``
486 #: when ``agentic_reflect.enabled`` is true.
487 self.agentic_reflect_params: AgenticReflectParams | None = None
488 #: Set by :meth:`astrocyte._astrocyte.Astrocyte.set_pipeline` when MIP is configured.
489 #: Used by :meth:`recall` to resolve per-bank rerank/reflect overrides (P3).
490 self.mip_router: MipRouter | None = None
491 # M8 W5 — wiki tier precedence. When a WikiStore is wired up and a
492 # compiled wiki page scores above ``wiki_confidence_threshold``, recall
493 # returns the wiki hit + raw-memory citations instead of running the full
494 # parallel-retrieve / RRF pipeline.
495 self.wiki_store: WikiStore | None = wiki_store
496 self.wiki_confidence_threshold: float = wiki_confidence_threshold
497 # M11: entity resolution — alias-of links between entities.
498 # None means the stage is skipped (opt-in, no cost when disabled).
499 self.entity_resolver: EntityResolver | None = entity_resolver
500 # Observation consolidation — post-retain background LLM pass that
501 # maintains a deduplicated observations layer. Disabled by default
502 # (opt-in). When enabled, recall also runs a separate "observation"
503 # strategy that searches the observations layer and fuses results into
504 # the main RRF pipeline with a configurable weight boost.
505 self.enable_observation_consolidation: bool = enable_observation_consolidation
506 self.observation_weight: float = observation_weight
507 # Weight applied to the ::obs bank when intent-gated injection fires
508 # (EXPLORATORY / RELATIONAL queries). Kept separate from observation_weight
509 # so callers can disable global injection (observation_weight=0.0) while
510 # still enabling the intent-gated path.
511 self.observation_injection_weight: float = observation_injection_weight
512 self.final_rerank_mode = final_rerank_mode
513 self.final_rerank_top_n = final_rerank_top_n
514 self.final_rerank_keep_n = final_rerank_keep_n
515 if enable_observation_consolidation:
516 from astrocyte.pipeline.observation import ObservationConsolidator
518 self._observation_consolidator: ObservationConsolidator | None = ObservationConsolidator()
519 else:
520 self._observation_consolidator = None
521 self._background_tasks: set[asyncio.Task[None]] = set()
523 # Ceiling on outstanding fire-and-forget consolidations.
524 #
525 # Without one, ``_consolidate_observations`` is an unbounded producer
526 # feeding a bounded consumer: every retain() spawns a task, each task
527 # calls ``llm_provider``, and providers bound their own concurrency
528 # (ClaudeCliProvider's semaphore defaults to 4). Under sustained Add
529 # load the queue grows without limit and *foreground* retain() calls
530 # starve behind it — measured at 421s of saturated backlog still
531 # draining after the client had disconnected, which surfaced as
532 # client-side ReadTimeouts on Add, not as any server-side error.
533 #
534 # Consolidation is best-effort by contract (the raw memories are
535 # already stored and searchable), so the correct response to
536 # saturation is to SHED, not to queue. Shedding keeps Add latency
537 # bounded; queueing trades it away for derived data nobody is waiting
538 # on. ``consolidations_shed`` makes the trade observable rather than
539 # silent — a persistently rising count means the deployment is
540 # under-provisioned for its ingest rate.
541 self.max_pending_consolidations: int = max_pending_consolidations
542 self.consolidations_shed: int = 0
544 # Overflow beyond the in-flight cap is deferred, not dropped (see
545 # ``RetainStageMixin._submit_consolidation``). Bounded on purpose — an
546 # unbounded queue would only move the original defect from task count
547 # to memory. Each entry holds one embedding and the first chunk's text;
548 # a 1536-dim vector as a Python float list is ~49 KB, so the default
549 # costs ~50 MB at worst. 0 restores shed-on-overflow.
550 self.max_deferred_consolidations: int = max_deferred_consolidations
551 self._deferred_consolidations: collections.deque = collections.deque()
552 self.consolidations_deferred: int = 0
553 self._consolidation_closed: bool = False
555 # Mental-model service — wires the agentic reflect loop to the
556 # configured ``MentalModelStore`` (typically ``PostgresMentalModelStore``).
557 # ``None`` when no store is configured; ``set_mental_model_service``
558 # is called from ``Astrocyte.set_pipeline`` when a store is present.
559 # When set, the agent gets ``search_mental_models`` as a tool —
560 # the highest-quality tier in the hierarchical priority order
561 # (mental_models → observations → recall → expand → done).
562 self.mental_model_service: object | None = None
564 def apply_config(self, cfg: PipelineConfig) -> None:
565 """Apply a resolved :class:`PipelineConfig` to this orchestrator.
567 Replaces the old inline attribute-poking in ``Astrocyte.set_pipeline``:
568 the config is derived once (see :meth:`PipelineConfig.from_config`) and
569 applied here in a single place. Field names on ``PipelineConfig`` mirror
570 the orchestrator attributes exactly, so this is a flat assignment — any
571 drift between the two surfaces is a loud ``AttributeError`` at wiring
572 time rather than a silently-ignored flag.
573 """
574 for name, value in cfg.as_orchestrator_attrs().items():
575 if not hasattr(self, name):
576 raise AttributeError(f"PipelineConfig field {name!r} has no matching orchestrator attribute")
577 setattr(self, name, value)
579 @property
580 def dedup_similarity_threshold(self) -> float:
581 """Default cosine threshold of the retain-time dedup check.
583 Set from ``signal_quality.dedup.similarity_threshold``; a matched MIP
584 rule's ``dedup.threshold`` still overrides it per retain. Stored on the
585 detector itself so the cache and store checks share one value.
586 """
587 return self._dedup.threshold
589 @dedup_similarity_threshold.setter
590 def dedup_similarity_threshold(self, value: float) -> None:
591 self._dedup.threshold = value
593 @property
594 def tokens_used(self) -> int:
595 """Total LLM tokens consumed through this orchestrator since last reset."""
596 return self._tracker.tokens_used
598 def reset_token_counter(self) -> int:
599 """Return accumulated token count and reset to zero."""
600 return self._tracker.reset_tokens()
602 def invalidate_dedup_cache(
603 self,
604 bank_id: str,
605 memory_ids: list[str] | None = None,
606 ) -> None:
607 """Drop forgotten memories from the in-memory dedup cache.
609 Called by ``Astrocyte.forget`` after the underlying store has
610 accepted the forget. Without this, a re-retain of similar content
611 after forget hits the in-memory ``DedupDetector`` cache (still
612 holding the forgotten memory's embedding) and the pipeline
613 short-circuits with ``RetainResult(stored=False, deduplicated=True,
614 error="All chunks are near-duplicates")`` — a silent no-op from
615 the caller's perspective.
617 Args:
618 bank_id: The bank the forget targeted.
619 memory_ids: Specific memories to drop from the cache. ``None``
620 (the whole-bank or tag-filtered forget path) clears the
621 entire bank's cache — coarser but always-correct.
622 """
623 if memory_ids is None:
624 self._dedup.clear_bank(bank_id)
625 return
627 for mid in memory_ids:
628 self._dedup.remove(bank_id, mid)
630 async def shutdown(self) -> None:
631 """Drain background work and close provider resources owned by the pipeline."""
632 # Close the deferred queue BEFORE cancelling: a cancelled task's done
633 # callback would otherwise promote deferred work into fresh tasks
634 # mid-shutdown. Undrained entries are counted as shed, not lost silently.
635 self._consolidation_closed = True
636 if self._deferred_consolidations:
637 self.consolidations_shed += len(self._deferred_consolidations)
638 self._deferred_consolidations.clear()
639 if self._background_tasks:
640 _, pending = await asyncio.wait(self._background_tasks, timeout=2.0)
641 for task in pending:
642 task.cancel()
643 if pending:
644 await asyncio.gather(*pending, return_exceptions=True)
646 for provider in (self.vector_store, self.graph_store, self.document_store):
647 close = getattr(provider, "close", None)
648 if close is None:
649 continue
650 try:
651 result = close()
652 if inspect.isawaitable(result):
653 _ = await result
654 except Exception as exc:
655 _logger.warning("provider close failed during pipeline shutdown: %s", exc)