Coverage for astrocyte/wiring.py: 98%
88 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""Shared provider resolution — one implementation for every deployment.
3Astrocyte is provider-agnostic by design: `astrocyte.yaml` names providers and
4the entry-point SPI resolves them. That guarantee only holds if every service
5resolves a given config *the same way*, and for a while they did not — the
6gateway ignored ``embedding_provider`` while the AML adapter honoured it, so
7one YAML file meant different things depending on which process loaded it. A
8key that parses and silently does nothing is worse than an unsupported one.
10This module is the single resolver for providers *and* stores. Store wiring
11used to be left to each caller, and the same failure recurred one level up:
12the gateway and AML adapter each grew their own copy, while ``astrocyte-mcp``
13— the entry point coding agents actually launch — got none, so every
14``memory_retain`` it served returned ``stored: false`` under any config.
15:func:`build_astrocyte` is now the one way to turn a config into a working
16brain.
17"""
19from __future__ import annotations
21import logging
22import os
23from typing import TYPE_CHECKING, Any
25from ._discovery import resolve_provider
26from .errors import ConfigError
28if TYPE_CHECKING: # pragma: no cover
29 from .config import AstrocyteConfig
31logger = logging.getLogger("astrocyte.wiring")
33__all__ = [
34 "build_astrocyte",
35 "build_pipeline",
36 "config_kwargs",
37 "instantiate_provider",
38 "resolve_llm_provider",
39 "resolve_store",
40 "wants_storage_pipeline",
41 "wire_astrocyte",
42]
44# Config key -> entry-point group. Each store is named by ``<key>`` and
45# configured by ``<key>_config``, with ``ASTROCYTE_<KEY>`` as the env fallback.
46_STORE_GROUPS = {
47 "vector_store": "vector_stores",
48 "graph_store": "graph_stores",
49 "document_store": "document_stores",
50 "wiki_store": "wiki_stores",
51 "mental_model_store": "mental_model_stores",
52 "source_store": "source_stores",
53}
56def config_kwargs(cfg: Any) -> dict[str, Any]:
57 """Provider config as kwargs, dropping unset keys.
59 ``None`` values are stripped so a partially-specified block falls through
60 to the provider's own defaults rather than overriding them with ``None``.
61 """
62 if cfg is None:
63 return {}
64 if not isinstance(cfg, dict):
65 cfg = {k: v for k, v in vars(cfg).items() if not k.startswith("_")}
66 return {k: v for k, v in cfg.items() if v is not None}
69def instantiate_provider(name: str, group: str, cfg: Any = None, *, label: str | None = None) -> Any:
70 """Resolve ``name`` in entry-point ``group`` and construct it."""
71 what = label or f"{group} {name!r}"
72 try:
73 cls = resolve_provider(name, group)
74 except LookupError as exc:
75 raise ConfigError(
76 f"{what} not found. Install a provider package or use a 'package.module:ClassName' path. ({exc})"
77 ) from exc
78 kwargs = config_kwargs(cfg)
79 try:
80 return cls(**kwargs) if kwargs else cls()
81 except TypeError as exc:
82 raise ConfigError(f"Invalid configuration for {what}: {exc}") from exc
85def resolve_llm_provider(config: AstrocyteConfig) -> Any:
86 """The LLM provider for ``config``, composing a split backend if configured.
88 When ``embedding_provider`` is set it is resolved separately and paired
89 with the completion provider through :class:`CompositeLLMProvider`. This is
90 not a niche path: some completion providers cannot embed at all —
91 ``ClaudeCliProvider`` raises ``NotImplementedError`` from ``embed()`` — so
92 ignoring the key means ``retain()`` fails at the embedding step with no
93 indication that the configured embedder was never consulted.
94 """
95 name = config.llm_provider or os.environ.get("ASTROCYTE_LLM_PROVIDER")
96 if not name:
97 # Kept as a default for tests and dev servers, but never silently: mock
98 # embeddings make every retain "succeed" and every recall return noise.
99 logger.warning(
100 "No llm_provider configured; using the mock provider. Memories will not be "
101 "meaningfully recalled. Set llm_provider in astrocyte.yaml (or run `astrocyte setup`)."
102 )
103 name = "mock"
104 completion = instantiate_provider(name, "llm_providers", config.llm_provider_config, label=f"llm_provider {name!r}")
106 embedder_name = getattr(config, "embedding_provider", None)
107 if not embedder_name or embedder_name == name:
108 # Same name means one provider serves both roles; composing it with
109 # itself would only add a layer of indirection.
110 return completion
112 from .providers.composite import CompositeLLMProvider
114 embedder = instantiate_provider(
115 embedder_name,
116 "llm_providers",
117 getattr(config, "embedding_provider_config", None),
118 label=f"embedding_provider {embedder_name!r}",
119 )
120 return CompositeLLMProvider(completion_provider=completion, embedding_provider=embedder)
123def resolve_store(config: AstrocyteConfig, kind: str) -> Any | None:
124 """The configured store of ``kind`` (e.g. ``"vector_store"``), or None.
126 Resolved from ``config.<kind>`` or ``ASTROCYTE_<KIND>``, constructed with
127 ``config.<kind>_config``.
128 """
129 group = _STORE_GROUPS.get(kind)
130 if group is None:
131 raise ValueError(f"unknown store kind {kind!r}; expected one of {sorted(_STORE_GROUPS)}")
132 name = getattr(config, kind, None) or os.environ.get(f"ASTROCYTE_{kind.upper()}")
133 if not name:
134 return None
135 return instantiate_provider(name, group, getattr(config, f"{kind}_config", None), label=f"{kind} {name!r}")
138def build_pipeline(config: AstrocyteConfig, *, wiki_store: Any | None = None, **overrides: Any) -> Any:
139 """A storage-tier :class:`PipelineOrchestrator` built from ``config``.
141 A vector store is required, and deliberately has no silent default: an
142 unconfigured memory server that falls back to an in-process store reports
143 every retain as a success and forgets it all on restart.
145 ``overrides`` are passed to the orchestrator for callers with their own
146 latency contract — e.g. the agent hooks disable recall-time query
147 expansion, an LLM call that costs 5–9 s through a CLI provider.
148 """
149 from .pipeline.entity_resolution import EntityResolver
150 from .pipeline.orchestrator import PipelineOrchestrator
152 vector_store = resolve_store(config, "vector_store")
153 if vector_store is None:
154 raise ConfigError(
155 "No vector_store configured. Run `astrocyte setup` for a local SQLite store, "
156 "or set `vector_store` in astrocyte.yaml (e.g. sqlite, postgres)."
157 )
158 llm = resolve_llm_provider(config)
159 graph_store = resolve_store(config, "graph_store")
160 document_store = resolve_store(config, "document_store")
161 # A vector store that also implements DocumentStore (Postgres via tsvector,
162 # SQLite via FTS5) enables the keyword retrieval leg with no extra config.
163 if document_store is None and hasattr(vector_store, "search_fulltext"):
164 document_store = vector_store
166 entity_resolver = None
167 if config.entity_resolution.enabled:
168 if graph_store is None:
169 raise ConfigError("entity_resolution.enabled requires a graph_store provider")
170 entity_resolver = EntityResolver(
171 similarity_threshold=config.entity_resolution.similarity_threshold,
172 confirmation_threshold=config.entity_resolution.confirmation_threshold,
173 max_candidates_per_entity=config.entity_resolution.max_candidates_per_entity,
174 )
176 return PipelineOrchestrator(
177 vector_store=vector_store,
178 llm_provider=llm,
179 graph_store=graph_store,
180 document_store=document_store,
181 wiki_store=wiki_store,
182 entity_resolver=entity_resolver,
183 **overrides,
184 )
187def wants_storage_pipeline(config: AstrocyteConfig) -> bool:
188 """Does ``config`` name a vector store (in the file or via env)?
190 That is the signal a caller expects a working storage-tier brain; configs
191 without one (engine tier, or a caller wiring stores by hand) are left as is.
192 """
193 return bool(getattr(config, "vector_store", None) or os.environ.get("ASTROCYTE_VECTOR_STORE"))
196def build_astrocyte(config: AstrocyteConfig) -> Any:
197 """A fully wired :class:`Astrocyte`: pipeline plus every configured store."""
198 from ._astrocyte import Astrocyte
200 return wire_astrocyte(Astrocyte(config), config)
203def wire_astrocyte(brain: Any, config: AstrocyteConfig) -> Any:
204 """Attach the pipeline and every configured store to ``brain``.
206 Optional stores (wiki, mental model, source) are attached only when
207 configured, matching the gateway's behaviour, so their tools report
208 "not configured" rather than failing obscurely.
209 """
210 from .config import access_grants_for_astrocyte
212 config.provider_tier = "storage"
213 wiki_store = resolve_store(config, "wiki_store")
214 pipeline = build_pipeline(config, wiki_store=wiki_store)
215 brain.set_pipeline(pipeline)
217 if wiki_store is not None:
218 brain.set_wiki_store(wiki_store)
219 if config.wiki_compile.auto_start:
220 from .pipeline.compile import CompileEngine
221 from .pipeline.compile_trigger import CompileQueue, CompileTriggerConfig
223 engine = CompileEngine(
224 vector_store=pipeline.vector_store,
225 llm_provider=pipeline.llm_provider,
226 wiki_store=wiki_store,
227 )
228 brain.set_compile_queue(
229 CompileQueue(
230 engine,
231 CompileTriggerConfig(
232 size_threshold=config.wiki_compile.size_threshold,
233 staleness_days=config.wiki_compile.staleness_days,
234 staleness_min_memories=config.wiki_compile.staleness_min_memories,
235 ),
236 max_queue_size=config.wiki_compile.max_queue_size,
237 )
238 )
239 if (mental_model_store := resolve_store(config, "mental_model_store")) is not None:
240 brain.set_mental_model_store(mental_model_store)
241 if (source_store := resolve_store(config, "source_store")) is not None:
242 brain.set_source_store(source_store)
243 if grants := access_grants_for_astrocyte(config):
244 brain.set_access_grants(grants)
245 return brain