diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 30213f0..922e92f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -28,7 +28,7 @@ jobs: with: path: | ${{ github.workspace }}/.cache/huggingface - key: ${{ runner.os }}-sentence-transformers-${{ hashFiles('pyproject.toml', 'app/startup/settings.py') }} + key: ${{ runner.os }}-onnx-${{ hashFiles('pyproject.toml', 'app/infrastructure/embeddings/local_provider.py') }} - name: Install dependencies run: | diff --git a/.github/workflows/real-embedding-smoke.yml b/.github/workflows/real-embedding-smoke.yml index ab39da8..0c3a808 100644 --- a/.github/workflows/real-embedding-smoke.yml +++ b/.github/workflows/real-embedding-smoke.yml @@ -41,7 +41,7 @@ jobs: with: path: | ${{ github.workspace }}/.cache/huggingface - key: ${{ runner.os }}-sentence-transformers-${{ hashFiles('pyproject.toml', 'app/startup/settings.py') }} + key: ${{ runner.os }}-onnx-${{ hashFiles('pyproject.toml', 'app/infrastructure/embeddings/local_provider.py') }} - name: Install dependencies run: | diff --git a/app/core/ports/db/concept_repositories.py b/app/core/ports/db/concept_repositories.py index 2150f86..1ec1749 100644 --- a/app/core/ports/db/concept_repositories.py +++ b/app/core/ports/db/concept_repositories.py @@ -93,6 +93,16 @@ def get_concept_bundle( ) -> dict[str, Any] | None: """This method returns one concept plus directly related graph records.""" + @abstractmethod + def get_concept_bundles( + self, + *, + repo_id: str, + concept_ids: Sequence[str], + include_lifecycle_events: bool = False, + ) -> dict[str, dict[str, Any]]: + """Return bundles keyed by concept ID; omit missing concepts.""" + @abstractmethod def find_concepts_for_memory_ids( self, *, repo_id: str, memory_ids: Sequence[str] diff --git a/app/core/use_cases/admin/initialize_runtime.py b/app/core/use_cases/admin/initialize_runtime.py index f3d1e91..f5737ae 100644 --- a/app/core/use_cases/admin/initialize_runtime.py +++ b/app/core/use_cases/admin/initialize_runtime.py @@ -230,20 +230,21 @@ def run_initialize_runtime( schema_changed = ports.apply_schema_migrations(machine_config) mutated_machine = mutated_machine or schema_changed - if not maintenance_only: - machine_config = ports.update_bootstrap_state( - machine_config, - bootstrap_state=ports.bootstrap_state_provisioning, - current_step="embeddings", - last_error=None, - ) - ports.save_machine_config(machine_config) - embedding_changed, machine_config = ports.prewarm_embeddings( - machine_config, - skip_model_download=skip_model_download, - ) - ports.save_machine_config(machine_config) - mutated_machine = mutated_machine or embedding_changed + machine_config = ports.update_bootstrap_state( + machine_config, + bootstrap_state=ports.bootstrap_state_provisioning, + current_step="embeddings", + last_error=None, + ) + ports.save_machine_config(machine_config) + embedding_changed, machine_config = ports.prewarm_embeddings( + machine_config, + skip_model_download=skip_model_download, + ) + ports.save_machine_config(machine_config) + mutated_machine = mutated_machine or embedding_changed + if machine_config.embeddings.readiness_state == "failed": + raise InitDependencyError(machine_config.embeddings.last_error) machine_config = ports.update_bootstrap_state( machine_config, diff --git a/app/core/use_cases/retrieval/deterministic_graph_recall.py b/app/core/use_cases/retrieval/deterministic_graph_recall.py index 88c54b6..ba0167e 100644 --- a/app/core/use_cases/retrieval/deterministic_graph_recall.py +++ b/app/core/use_cases/retrieval/deterministic_graph_recall.py @@ -140,14 +140,19 @@ def build_deterministic_graph_pack( uow=uow, thresholds=thresholds, ) + concept_bundles = uow.concepts.get_concept_bundles( + repo_id=request.repo_id, + concept_ids=list(concept_candidates), + ) selected_concepts, concept_trace = _select_concepts( request=request, concept_candidates=concept_candidates, - uow=uow, + concept_bundles=concept_bundles, ) traversal = _traverse_selected_concepts( request=request, selected_concepts=selected_concepts, + concept_bundles=concept_bundles, memory_candidates=memory_candidates, uow=uow, ) @@ -484,8 +489,14 @@ def _expand_structural_memory_relations( predicates=predicates, ) for row in rows: - key = (row["subject_memory_id"], row["predicate"], row["object_memory_id"]) - memory_relations[key] = row | {"validated_at": _iso(row["validated_at"])} + key = ( + row["subject_memory_id"], + row["predicate"], + row["object_memory_id"], + ) + memory_relations[key] = row | { + "validated_at": _iso(row["validated_at"]) + } del memory_relations[key]["visible_memory_ids"] for neighbor in select_structural_memory_relation_neighbors( rows, anchor_memory_id=anchor_memory_id @@ -605,14 +616,14 @@ def _select_concepts( *, request: MemoryReadRequest, concept_candidates: dict[str, dict[str, Any]], - uow: IUnitOfWork, + concept_bundles: dict[str, dict[str, Any]], ) -> tuple[list[dict[str, Any]], dict[str, Any]]: selected: list[tuple[float, str, dict[str, Any]]] = [] rejected_count = 0 + query_terms = _tokenize(request.query) + identifiers = _extract_identifiers(request.query) for concept_id, candidate in concept_candidates.items(): - bundle = uow.concepts.get_concept_bundle( - repo_id=request.repo_id, concept_ref=concept_id - ) + bundle = concept_bundles.get(concept_id) if bundle is None: rejected_count += 1 continue @@ -621,8 +632,8 @@ def _select_concepts( continue score = float(candidate["score"]) + _bundle_signal_score( bundle=bundle, - query_terms=_tokenize(request.query), - identifiers=_extract_identifiers(request.query), + query_terms=query_terms, + identifiers=identifiers, ) if score <= 0: rejected_count += 1 @@ -653,6 +664,7 @@ def _traverse_selected_concepts( *, request: MemoryReadRequest, selected_concepts: list[dict[str, Any]], + concept_bundles: dict[str, dict[str, Any]], memory_candidates: dict[str, dict[str, Any]], uow: IUnitOfWork, ) -> dict[str, Any]: @@ -711,12 +723,22 @@ def _traverse_selected_concepts( ] neighbors = [] selected_ids = {entry["bundle"]["concept"].id for entry in selected_concepts} + missing_ids = [ + key + for key in unique_neighbor_ids + if key not in concept_bundles and key not in selected_ids + ] + if missing_ids: + concept_bundles.update( + uow.concepts.get_concept_bundles( + repo_id=request.repo_id, + concept_ids=missing_ids, + ) + ) for neighbor_id in unique_neighbor_ids: if neighbor_id in selected_ids: continue - bundle = uow.concepts.get_concept_bundle( - repo_id=request.repo_id, concept_ref=neighbor_id - ) + bundle = concept_bundles.get(neighbor_id) if bundle is None or not _active_concept(bundle["concept"]): continue neighbors.append( @@ -1214,9 +1236,13 @@ def _brief_case_texts( for target in memories: linked = relations_by_target[target["id"]] if not linked: - cases.extend(_brief_memory_texts( - [target], kinds={"solution"}, link_roles={"solution_for", "example_of"} - )) + cases.extend( + _brief_memory_texts( + [target], + kinds={"solution"}, + link_roles={"solution_for", "example_of"}, + ) + ) for relation in linked: subject = memories_by_id[relation["subject_memory_id"]] qualifiers = [relation["status"]] diff --git a/app/infrastructure/db/admin/provisioning/external_postgres.py b/app/infrastructure/db/admin/provisioning/external_postgres.py index 7e1b9c5..3d4a7a6 100644 --- a/app/infrastructure/db/admin/provisioning/external_postgres.py +++ b/app/infrastructure/db/admin/provisioning/external_postgres.py @@ -2,7 +2,6 @@ from __future__ import annotations -import importlib.metadata import secrets import psycopg @@ -43,7 +42,6 @@ def build_fresh_machine_config( raise RuntimeError("runtime.embeddings must be configured") app_dsn = _provision_external_app_role(admin_dsn=admin_dsn) instance_id = dsn_fingerprint(admin_dsn) - backend_version = _sentence_transformers_version() return MachineConfig( config_version=CONFIG_VERSION, bootstrap_version=BOOTSTRAP_VERSION, @@ -56,10 +54,10 @@ def build_fresh_machine_config( managed=None, backups=BackupState(root=str(get_machine_backups_dir()), mirror_root=None), embeddings=EmbeddingRuntimeState( - provider=str(embeddings.get("provider") or "sentence_transformers"), + provider=str(embeddings.get("provider") or "onnxruntime"), model=str(embeddings.get("model") or "all-MiniLM-L6-v2"), model_revision=None, - backend_version=backend_version, + backend_version=None, cache_path=str(get_machine_models_dir()), readiness_state="pending", last_error=None, @@ -205,12 +203,3 @@ def _app_dsn_from_admin_dsn(*, admin_dsn: str, app_user: str, app_password: str) url = make_url(admin_dsn) return str(url.set(username=app_user, password=app_password)) - - -def _sentence_transformers_version() -> str | None: - """Return the installed sentence-transformers version when present.""" - - try: - return importlib.metadata.version("sentence-transformers") - except importlib.metadata.PackageNotFoundError: - return None diff --git a/app/infrastructure/db/admin/provisioning/managed_local.py b/app/infrastructure/db/admin/provisioning/managed_local.py index 2995894..73c392a 100644 --- a/app/infrastructure/db/admin/provisioning/managed_local.py +++ b/app/infrastructure/db/admin/provisioning/managed_local.py @@ -89,7 +89,7 @@ def build_fresh_machine_config(*, embeddings: dict[str, object]) -> MachineConfi ), backups=BackupState(root=str(get_machine_backups_dir()), mirror_root=None), embeddings=EmbeddingRuntimeState( - provider=str(embeddings.get("provider") or "sentence_transformers"), + provider=str(embeddings.get("provider") or "onnxruntime"), model=str(embeddings.get("model") or "all-MiniLM-L6-v2"), model_revision=None, backend_version=None, @@ -288,7 +288,7 @@ def recover_machine_config_from_docker( ), backups=BackupState(root=str(get_machine_backups_dir()), mirror_root=None), embeddings=EmbeddingRuntimeState( - provider="sentence_transformers", + provider="onnxruntime", model=str(embeddings.get("model") or "all-MiniLM-L6-v2"), model_revision=None, backend_version=None, diff --git a/app/infrastructure/db/runtime/repos/relational/concepts_repo.py b/app/infrastructure/db/runtime/repos/relational/concepts_repo.py index 8068683..36ebba8 100644 --- a/app/infrastructure/db/runtime/repos/relational/concepts_repo.py +++ b/app/infrastructure/db/runtime/repos/relational/concepts_repo.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections import defaultdict from datetime import datetime, timezone from typing import Any, Sequence @@ -373,9 +374,7 @@ def get_lifecycle_target( table, converter = _lifecycle_target_table(target_type) row = ( self._session.execute( - select(table).where( - table.c.repo_id == repo_id, table.c.id == target_id - ) + select(table).where(table.c.repo_id == repo_id, table.c.id == target_id) ) .mappings() .first() @@ -443,131 +442,113 @@ def get_concept_bundle( concept = self.get_concept_by_ref(repo_id=repo_id, concept_ref=concept_ref) if concept is None: return None - relation_rows = ( - self._session.execute( - select(concept_relations).where( - concept_relations.c.repo_id == repo_id, - ( - (concept_relations.c.subject_concept_id == concept.id) - | (concept_relations.c.object_concept_id == concept.id) - ), - ) - ) - .mappings() - .all() - ) - claim_rows = ( - self._session.execute( - select(concept_claims).where( - concept_claims.c.repo_id == repo_id, - concept_claims.c.concept_id == concept.id, - ) - ) - .mappings() - .all() - ) - grounding_rows = ( - self._session.execute( - select(concept_groundings).where( - concept_groundings.c.repo_id == repo_id, - concept_groundings.c.concept_id == concept.id, - ) - ) - .mappings() - .all() - ) - memory_link_rows = ( - self._session.execute( - select(concept_memory_links).where( - concept_memory_links.c.repo_id == repo_id, - concept_memory_links.c.concept_id == concept.id, - ) - ) - .mappings() - .all() - ) - alias_rows = ( - self._session.execute( - select(concept_aliases).where( - concept_aliases.c.repo_id == repo_id, - concept_aliases.c.concept_id == concept.id, - ) - ) - .mappings() - .all() - ) - anchor_ids = [str(row["anchor_id"]) for row in grounding_rows] - anchor_rows = [] - if anchor_ids: - anchor_rows = ( - self._session.execute( - select(anchors).where( - anchors.c.repo_id == repo_id, anchors.c.id.in_(anchor_ids) - ) - ) - .mappings() - .all() - ) - target_pairs = ( - [("relation", "concept_relation", str(row["id"])) for row in relation_rows] - + [("claim", "concept_claim", str(row["id"])) for row in claim_rows] - + [ - ("grounding", "concept_grounding", str(row["id"])) - for row in grounding_rows - ] - + [ - ("memory_link", "concept_memory_link", str(row["id"])) - for row in memory_link_rows - ] - ) - lifecycle_event_rows = [] - if include_lifecycle_events and target_pairs: - lifecycle_event_rows = ( + return self.get_concept_bundles( + repo_id=repo_id, + concept_ids=[concept.id], + include_lifecycle_events=include_lifecycle_events, + ).get(concept.id) + + def get_concept_bundles( + self, + *, + repo_id: str, + concept_ids: Sequence[str], + include_lifecycle_events: bool = False, + ) -> dict[str, dict[str, Any]]: + """Fetch candidate bundles in a fixed number of repo-scoped queries.""" + if not concept_ids: + return {} + + def rows(table, condition): + return ( self._session.execute( - select(concept_lifecycle_events).where( - concept_lifecycle_events.c.repo_id == repo_id, - or_( - *[ - and_( - concept_lifecycle_events.c.target_type - == lifecycle_target_type, - concept_lifecycle_events.c.target_id - == target_id, - ) - for lifecycle_target_type, _, target_id in target_pairs - ] - ), - ) + select(table) + .where(table.c.repo_id == repo_id, condition) + .order_by(table.c.created_at, *table.primary_key.columns) ) .mappings() .all() ) - event_pairs = [ - ("lifecycle_event", "concept_lifecycle_event", str(row["id"])) - for row in lifecycle_event_rows - ] - evidence_rows = [] - evidence_pairs = target_pairs + event_pairs - if evidence_pairs: - evidence_rows = self._unified_evidence_rows( - repo_id=repo_id, evidence_pairs=evidence_pairs - ) - return { - "concept": concept, - "aliases": [_to_alias(row) for row in alias_rows], - "relations": [_to_relation(row) for row in relation_rows], - "claims": [_to_claim(row) for row in claim_rows], - "groundings": [_to_grounding(row) for row in grounding_rows], - "memory_links": [_to_memory_link(row) for row in memory_link_rows], - "lifecycle_events": [ - _to_lifecycle_event(row) for row in lifecycle_event_rows - ], - "anchors": [_to_anchor(row) for row in anchor_rows], - "evidence": [_to_evidence(row) for row in evidence_rows], + + bundles = { + str(row["id"]): { + "concept": _to_concept(row), + "relations": [], + "claims": [], + "groundings": [], + "memory_links": [], + "aliases": [], + "anchors": [], + "lifecycle_events": [], + "evidence": [], + } + for row in rows(concepts, concepts.c.id.in_(concept_ids)) } + if not bundles: + return {} + ids = list(bundles) + owners = defaultdict(set) + for row in rows( + concept_relations, + or_( + concept_relations.c.subject_concept_id.in_(ids), + concept_relations.c.object_concept_id.in_(ids), + ), + ): + relation = _to_relation(row) + for concept_id in { + relation.subject_concept_id, + relation.object_concept_id, + } & bundles.keys(): + bundles[concept_id]["relations"].append(relation) + owners[("relation", relation.id)].add(concept_id) + for field, table, convert, kind in ( + ("claims", concept_claims, _to_claim, "claim"), + ("groundings", concept_groundings, _to_grounding, "grounding"), + ("memory_links", concept_memory_links, _to_memory_link, "memory_link"), + ("aliases", concept_aliases, _to_alias, None), + ): + for row in rows(table, table.c.concept_id.in_(ids)): + item = convert(row) + concept_id = str(row["concept_id"]) + bundles[concept_id][field].append(item) + if kind: + owners[(kind, item.id)].add(concept_id) + anchor_owners = defaultdict(set) + for concept_id, bundle in bundles.items(): + for grounding in bundle["groundings"]: + anchor_owners[grounding.anchor_id].add(concept_id) + if anchor_owners: + for row in rows(anchors, anchors.c.id.in_(anchor_owners)): + anchor = _to_anchor(row) + for concept_id in anchor_owners[anchor.id]: + bundles[concept_id]["anchors"].append(anchor) + if include_lifecycle_events and owners: + for row in rows( + concept_lifecycle_events, + _target_filter(concept_lifecycle_events, owners), + ): + event = _to_lifecycle_event(row) + concept_ids_for_event = owners[ + (event.target_type.value, event.target_id) + ] + for concept_id in concept_ids_for_event: + bundles[concept_id]["lifecycle_events"].append(event) + owners[("lifecycle_event", event.id)].update(concept_ids_for_event) + if owners: + for row in self._unified_evidence_rows( + repo_id=repo_id, + targets=[(f"concept_{kind}", target_id) for kind, target_id in owners], + ): + evidence = _to_evidence(row) + for concept_id in owners[ + (evidence.target_type.value, evidence.target_id) + ]: + bundles[concept_id]["evidence"].append(evidence) + return bundles def _unified_evidence_rows( - self, *, repo_id: str, evidence_pairs: Sequence[tuple[str, str, str]] + self, *, repo_id: str, targets: Sequence[tuple[str, str]] ) -> list[dict[str, Any]]: """Return concept evidence rows from unified evidence storage.""" @@ -596,15 +577,7 @@ def _unified_evidence_rows( ) .where( evidence_links.c.repo_id == repo_id, - or_( - *[ - and_( - evidence_links.c.target_type == unified_target_type, - evidence_links.c.target_id == target_id, - ) - for _, unified_target_type, target_id in evidence_pairs - ] - ), + _target_filter(evidence_links, targets), ) .order_by(evidence_links.c.created_at, evidence_links.c.id) ) @@ -836,7 +809,9 @@ def _unified_evidence_to_concept_row(row) -> dict[str, Any]: "repo_id": row["repo_id"], "target_type": _CONCEPT_TARGET_TYPE_BY_UNIFIED[str(row["target_type"])], "target_id": row["target_id"], - "evidence_kind": "transcript" if source_kind == "episode_event" else source_kind, + "evidence_kind": "transcript" + if source_kind == "episode_event" + else source_kind, "anchor_id": row["anchor_id"], "memory_id": row["memory_id"], "commit_ref": row["commit_ref"], @@ -903,3 +878,16 @@ def _normalize_text(value: str) -> str: """Normalize natural keys for aliases and claim text.""" return " ".join(value.strip().lower().split()) + + +def _target_filter(table, targets): + """Match typed IDs without one SQL OR expression per record.""" + ids_by_type = defaultdict(list) + for kind, target_id in targets: + ids_by_type[kind].append(target_id) + return or_( + *( + and_(table.c.target_type == kind, table.c.target_id.in_(ids)) + for kind, ids in ids_by_type.items() + ) + ) diff --git a/app/infrastructure/embeddings/local_provider.py b/app/infrastructure/embeddings/local_provider.py index e8da431..e96ce99 100644 --- a/app/infrastructure/embeddings/local_provider.py +++ b/app/infrastructure/embeddings/local_provider.py @@ -1,49 +1,68 @@ -"""This module defines a sentence-transformers-backed local embedding provider.""" +"""Run the pinned MiniLM embedding model locally with ONNX Runtime.""" -from typing import Sequence +from pathlib import Path from app.core.ports.embeddings.provider import IEmbeddingProvider +MODEL_REPOSITORY = "sentence-transformers/all-MiniLM-L6-v2" +MODEL_REVISION = "c9745ed1d9f207416be6d2e6f8de32d1f16199bf" +MODEL_FILES = ("onnx/model.onnx", "tokenizer.json") -class SentenceTransformersEmbeddingProvider(IEmbeddingProvider): - """This class generates embeddings with a local sentence-transformers model.""" - def __init__( - self, - *, - model: str, - cache_folder: str | None = None, - local_files_only: bool = False, - ) -> None: - """This method stores sentence-transformers model configuration for lazy loading.""" +def model_directory(cache_folder: str) -> Path: + """Locate the pinned Hugging Face snapshot without importing its downloader.""" + return ( + Path(cache_folder) + / "models--sentence-transformers--all-MiniLM-L6-v2" + / "snapshots" + / MODEL_REVISION + ) - self._model_name = model - self._cache_folder = cache_folder - self._local_files_only = local_files_only - self._model = None - def _get_model(self): - """This method lazily loads the configured sentence-transformers model.""" +class OnnxEmbeddingProvider(IEmbeddingProvider): + """Preserve MiniLM tokenization, masked mean pooling, and normalization.""" - if self._model is not None: - return self._model - try: - from sentence_transformers import SentenceTransformer + def __init__(self, *, cache_folder: str) -> None: + self._directory = model_directory(cache_folder) + self._session = None - self._model = SentenceTransformer( - self._model_name, - cache_folder=self._cache_folder, - local_files_only=self._local_files_only, - ) - except Exception as exc: + def _load_model(self) -> None: + """Keep database-only units of work free of model loading.""" + if self._session is not None: + return + directory = self._directory + if not all((directory / name).is_file() for name in MODEL_FILES): raise RuntimeError( - "sentence-transformers is unavailable for local embedding generation" - ) from exc - return self._model + "Shellbrain embedding files are missing. Run `shellbrain upgrade` to download them." + ) + import numpy as np + import onnxruntime as ort + from tokenizers import Tokenizer - def embed(self, text: str) -> Sequence[float]: - """This method returns a dense embedding vector from the local sentence-transformers model.""" + self._np = np + self._tokenizer = Tokenizer.from_file(str(directory / "tokenizer.json")) + self._tokenizer.enable_truncation(max_length=256) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + options.inter_op_num_threads = 1 + self._session = ort.InferenceSession( + str(directory / "onnx/model.onnx"), + sess_options=options, + providers=["CPUExecutionProvider"], + ) - model = self._get_model() - vector = model.encode(text, convert_to_numpy=False, normalize_embeddings=False) - return [float(value) for value in vector] + def embed(self, text: str) -> list[float]: + """Return the same normalized 384-dimensional vector as MiniLM's encoder.""" + self._load_model() + np = self._np + encoded = self._tokenizer.encode(text.strip()) + inputs = { + "input_ids": np.array([encoded.ids], dtype=np.int64), + "attention_mask": np.array([encoded.attention_mask], dtype=np.int64), + "token_type_ids": np.array([encoded.type_ids], dtype=np.int64), + } + vectors = self._session.run(None, inputs)[0] + mask = inputs["attention_mask"][..., None].astype(np.float32) + pooled = (vectors * mask).sum(axis=1) / np.maximum(mask.sum(axis=1), 1e-9) + pooled /= np.maximum(np.linalg.norm(pooled, axis=1, keepdims=True), 1e-12) + return pooled[0].tolist() diff --git a/app/infrastructure/embeddings/prewarm.py b/app/infrastructure/embeddings/prewarm.py index 2f0250d..4591c38 100644 --- a/app/infrastructure/embeddings/prewarm.py +++ b/app/infrastructure/embeddings/prewarm.py @@ -1,98 +1,94 @@ -"""Concrete embedding prewarm mechanics for runtime initialization.""" - -from __future__ import annotations +"""Download and verify embedding artifacts during installation and upgrades.""" from dataclasses import replace -import importlib.metadata -import os +from importlib.metadata import version from pathlib import Path -from app.infrastructure.embeddings.local_provider import ( - SentenceTransformersEmbeddingProvider, -) from app.core.entities.machine_config import ( BOOTSTRAP_STATE_REPAIR_NEEDED, - EmbeddingRuntimeState, MachineConfig, ) +from app.infrastructure.embeddings.local_provider import ( + MODEL_FILES, + MODEL_REPOSITORY, + MODEL_REVISION, + OnnxEmbeddingProvider, +) def prewarm_embeddings( config: MachineConfig, *, skip_model_download: bool ) -> tuple[bool, MachineConfig]: - """Prewarm the configured embedding backend and pin its runtime metadata.""" - - backend_version = _sentence_transformers_version() + """Verify the local model before marking an installation ready.""" if skip_model_download: updated = replace( config, - embeddings=_embedding_state( - config, - backend_version=backend_version, + embeddings=replace( + config.embeddings, readiness_state="skipped", last_error="Model prewarm was skipped during init.", ), ) - return True, updated - - os.environ["HF_HOME"] = config.embeddings.cache_path - Path(config.embeddings.cache_path).mkdir(parents=True, exist_ok=True) - provider = SentenceTransformersEmbeddingProvider( - model=config.embeddings.model, - cache_folder=config.embeddings.cache_path, - ) + return updated != config, updated try: - provider.embed("shellbrain upgrade warmup") + _verify_model_revision(config) + try: + OnnxEmbeddingProvider(cache_folder=config.embeddings.cache_path).embed( + "shellbrain upgrade warmup" + ) + except Exception: + # Only bootstrap may fetch or replace missing/corrupt model files. + from huggingface_hub import snapshot_download + + snapshot_download( + MODEL_REPOSITORY, + revision=MODEL_REVISION, + cache_dir=config.embeddings.cache_path, + allow_patterns=list(MODEL_FILES), + force_download=True, + ) + OnnxEmbeddingProvider(cache_folder=config.embeddings.cache_path).embed( + "shellbrain upgrade warmup" + ) + updated = replace( + config, + embeddings=replace( + config.embeddings, + provider="onnxruntime", + model_revision=MODEL_REVISION, + backend_version=version("onnxruntime"), + readiness_state="ready", + last_error=None, + ), + ) except Exception as exc: updated = replace( config, bootstrap_state=BOOTSTRAP_STATE_REPAIR_NEEDED, current_step="embeddings", last_error=str(exc), - embeddings=_embedding_state( - config, - backend_version=backend_version, - readiness_state="failed", - last_error=str(exc), + embeddings=replace( + config.embeddings, readiness_state="failed", last_error=str(exc) ), ) - return True, updated - updated = replace( - config, - embeddings=_embedding_state( - config, - backend_version=backend_version, - readiness_state="ready", - last_error=None, - ), - ) - return ( - config.embeddings.readiness_state != "ready" - or config.embeddings.backend_version != backend_version, - updated, - ) - + return updated != config, updated -def _embedding_state( - config: MachineConfig, - *, - backend_version: str | None, - readiness_state: str, - last_error: str | None, -) -> EmbeddingRuntimeState: - return EmbeddingRuntimeState( - provider=config.embeddings.provider, - model=config.embeddings.model, - model_revision=config.embeddings.model_revision, - backend_version=backend_version, - cache_path=config.embeddings.cache_path, - readiness_state=readiness_state, - last_error=last_error, - ) - -def _sentence_transformers_version() -> str | None: - try: - return importlib.metadata.version("sentence-transformers") - except importlib.metadata.PackageNotFoundError: - return None +def _verify_model_revision(config: MachineConfig) -> None: + """Do not mix existing vectors with an unverified model revision.""" + revision = config.embeddings.model_revision + if revision is None: + cached_ref = ( + Path(config.embeddings.cache_path) + / "models--sentence-transformers--all-MiniLM-L6-v2" + / "refs/main" + ) + if cached_ref.is_file(): + revision = cached_ref.read_text().strip() + if config.embeddings.model != "all-MiniLM-L6-v2" or revision not in ( + None, + MODEL_REVISION, + ): + raise RuntimeError( + "The existing embedding model differs from the pinned MiniLM revision. Re-embed stored memories and concepts before changing models." + ) diff --git a/app/infrastructure/host_apps/inner_agents/codex_cli.py b/app/infrastructure/host_apps/inner_agents/codex_cli.py index 8df3ea3..6743595 100644 --- a/app/infrastructure/host_apps/inner_agents/codex_cli.py +++ b/app/infrastructure/host_apps/inner_agents/codex_cli.py @@ -256,7 +256,7 @@ def _codex_exec_args( ] for feature in _CODEX_DISABLED_FEATURES: args.extend(["--disable", feature]) - if model == "gpt-5.6-luna": + if model in {"gpt-5.6-luna", "gpt-6-luna"}: args.extend(["-c", 'service_tier="fast"']) args.extend( [ diff --git a/app/startup/embeddings.py b/app/startup/embeddings.py index 206e321..02d8d3f 100644 --- a/app/startup/embeddings.py +++ b/app/startup/embeddings.py @@ -5,7 +5,7 @@ from app.core.ports.embeddings.provider import IEmbeddingProvider from app.infrastructure.local_state.machine_config_store import load_machine_config from app.infrastructure.embeddings.local_provider import ( - SentenceTransformersEmbeddingProvider, + OnnxEmbeddingProvider, ) @@ -17,19 +17,14 @@ def get_embedding_model_name() -> str: def get_embedding_provider() -> IEmbeddingProvider: """This function constructs the configured local embedding provider.""" - model = get_embedding_model_name() machine_config = load_machine_config() cache_folder = str(get_machine_models_dir()) - local_files_only = False if machine_config is not None: cache_folder = machine_config.embeddings.cache_path if machine_config.embeddings.readiness_state != "ready": raise RuntimeError( "Shellbrain embeddings are not ready. Rerun `shellbrain upgrade` to finish model setup." ) - local_files_only = True - return SentenceTransformersEmbeddingProvider( - model=model, + return OnnxEmbeddingProvider( cache_folder=cache_folder, - local_files_only=local_files_only, ) diff --git a/app/startup/internal_agent_config.py b/app/startup/internal_agent_config.py index 5f1ef11..4aa5d3b 100644 --- a/app/startup/internal_agent_config.py +++ b/app/startup/internal_agent_config.py @@ -73,7 +73,7 @@ def default_internal_agents_config() -> InternalAgentsConfig: ), build_knowledge=BuildKnowledgeSettings( provider="auto", - model="gpt-5.6-luna", + model="gpt-6-luna", reasoning="xhigh", timeout_seconds=600, ), diff --git a/app/startup/runtime_admin.py b/app/startup/runtime_admin.py index 3ea29e6..2e6d368 100644 --- a/app/startup/runtime_admin.py +++ b/app/startup/runtime_admin.py @@ -337,4 +337,4 @@ def _recover_machine_config() -> MachineConfig | None: def _runtime_embeddings_config() -> dict[str, object]: """Return runtime embedding config for adapter construction.""" - return {"provider": "sentence_transformers", "model": DEFAULT_EMBEDDING_MODEL} + return {"provider": "onnxruntime", "model": DEFAULT_EMBEDDING_MODEL} diff --git a/pyproject.toml b/pyproject.toml index 97f8e5e..b5743ed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "shellbrain" -version = "0.1.61" +version = "0.1.62" description = "Repo-scoped Shellbrain CLI with explicit evidence-backed writes." readme = "README.md" requires-python = ">=3.11" @@ -14,7 +14,10 @@ dependencies = [ "pydantic>=2.7,<3.0", "psycopg[binary]>=3.1,<4.0", "pgvector>=0.3,<1.0", - "sentence-transformers>=3.0,<4.0", + "onnxruntime>=1.20,<2.0", + "tokenizers>=0.20,<1.0", + "numpy>=1.26,<3.0", + "huggingface-hub>=0.26,<2.0", ] [project.scripts] diff --git a/pytest.ini b/pytest.ini index b467ca1..822835c 100644 --- a/pytest.ini +++ b/pytest.ini @@ -2,4 +2,4 @@ testpaths = tests markers = - real_embedding: runs tests that require sentence-transformers model loading + real_embedding: runs tests that require ONNX model loading diff --git a/tests/config/test_loader.py b/tests/config/test_loader.py index 82ed9f3..8ee0ae7 100644 --- a/tests/config/test_loader.py +++ b/tests/config/test_loader.py @@ -34,7 +34,7 @@ def test_packaged_defaults_preserve_inner_agent_settings() -> None: assert "fallback" not in settings["build_context"] assert "enabled" not in settings["build_knowledge"] assert "fallback" not in settings["build_knowledge"] - assert settings["build_knowledge"]["model"] == "gpt-5.6-luna" + assert settings["build_knowledge"]["model"] == "gpt-6-luna" assert settings["build_knowledge"]["reasoning"] == "xhigh" assert settings["build_knowledge"]["timeout_seconds"] == 600 assert settings["build_knowledge"]["max_shellbrain_reads"] == 8 @@ -163,7 +163,7 @@ def test_startup_wires_codex_non_recall_runners(monkeypatch, runner_getter) -> N @pytest.mark.parametrize( ("settings_getter", "expected_model"), - ((get_build_knowledge_settings, "gpt-5.6-luna"),), + ((get_build_knowledge_settings, "gpt-6-luna"),), ) def test_startup_preserves_codex_workflow_models( monkeypatch, settings_getter, expected_model: str diff --git a/tests/config/test_onboarding_assets.py b/tests/config/test_onboarding_assets.py index 6e9ab50..38d84de 100644 --- a/tests/config/test_onboarding_assets.py +++ b/tests/config/test_onboarding_assets.py @@ -25,14 +25,9 @@ def test_readme_should_teach_the_installer_first_happy_path() -> None: assert "curl -L shellbrain.ai/install | bash" in readme assert "shellbrain upgrade" in readme - assert "curl -L shellbrain.ai/upgrade | bash" in readme - assert "configures the runtime automatically" in readme assert "shellbrain admin recall provider inception" in readme - assert "Use $shellbrain" in readme - assert "Use Shellbrain" in readme assert "utility_vote" not in readme assert "what should I know about this repo?" not in readme - assert "Repos register themselves on first use." in readme def test_agent_docs_should_share_the_shellbrain_protocol() -> None: diff --git a/tests/config/test_runtime_usage.py b/tests/config/test_runtime_usage.py index c96c521..36fa13b 100644 --- a/tests/config/test_runtime_usage.py +++ b/tests/config/test_runtime_usage.py @@ -179,12 +179,8 @@ def test_embedding_boot_should_use_local_only_when_machine_config_is_ready( class _FakeProvider: """Record the embedding provider constructor arguments.""" - def __init__( - self, *, model: str, cache_folder: str, local_files_only: bool - ) -> None: - self.model = model + def __init__(self, *, cache_folder: str) -> None: self.cache_folder = cache_folder - self.local_files_only = local_files_only machine_config = MachineConfig( config_version=2, @@ -226,16 +222,12 @@ def __init__( monkeypatch.setattr( "app.startup.embeddings.load_machine_config", lambda: machine_config ) - monkeypatch.setattr( - "app.startup.embeddings.SentenceTransformersEmbeddingProvider", _FakeProvider - ) + monkeypatch.setattr("app.startup.embeddings.OnnxEmbeddingProvider", _FakeProvider) provider = get_embedding_provider() assert isinstance(provider, _FakeProvider) - assert provider.model == "all-MiniLM-L6-v2" assert provider.cache_folder == "/tmp/shellbrain-models" - assert provider.local_files_only is True def test_db_boot_should_support_external_machine_config(monkeypatch) -> None: diff --git a/tests/infrastructure/test_codex_inner_agent.py b/tests/infrastructure/test_codex_inner_agent.py index af746e9..7de1ac3 100644 --- a/tests/infrastructure/test_codex_inner_agent.py +++ b/tests/infrastructure/test_codex_inner_agent.py @@ -26,12 +26,13 @@ ) -def test_codex_exec_uses_fast_service_tier_for_luna() -> None: +@pytest.mark.parametrize("model", ["gpt-5.6-luna", "gpt-6-luna"]) +def test_codex_exec_uses_fast_service_tier_for_luna(model: str) -> None: """Luna inner-agent calls should use Codex Fast mode.""" args = _codex_exec_args( command_path="/usr/bin/codex", - model="gpt-5.6-luna", + model=model, reasoning="low", workspace="/tmp/workspace", output_path="/tmp/output.json", diff --git a/tests/operations/concepts/execution/test_concept_show_use_case.py b/tests/operations/concepts/execution/test_concept_show_use_case.py index ac9bfdc..461c569 100644 --- a/tests/operations/concepts/execution/test_concept_show_use_case.py +++ b/tests/operations/concepts/execution/test_concept_show_use_case.py @@ -292,3 +292,116 @@ def _seed_deposit_addresses(uow_factory: Callable[[], PostgresUnitOfWork]) -> No uow, id_generator=_SequenceIdGenerator(), ) + + +def test_bulk_bundles_preserve_typed_evidence_lifecycle_and_repo_scope(uow_factory): + """Batching must not leak evidence when record tables reuse an ID.""" + from sqlalchemy import event, update + from app.core.entities.concepts import ( + Concept, + ConceptKind, + ConceptLifecycleEvent, + ConceptLifecycleTargetType, + ConceptLifecycleStatus, + ConceptCreatedBy, + ) + from app.infrastructure.db.runtime.models.evidence import evidence_links + + _seed_deposit_addresses(uow_factory) + with uow_factory() as uow: + first = uow.concepts.get_concept_bundle( + repo_id="repo-a", concept_ref="deposit-addresses" + ) + first_id = first["concept"].id + second_id = first["relations"][0].object_concept_id + shared_id = first["relations"][0].id + claim_id = first["claims"][0].id + uow.concepts.add_concept( + Concept( + id="third", + repo_id="repo-a", + slug="third", + name="Third", + kind=ConceptKind.DOMAIN, + ), + [], + ) + uow.concepts.add_concept( + Concept( + id="foreign", + repo_id="repo-b", + slug="foreign", + name="Foreign", + kind=ConceptKind.DOMAIN, + ), + [], + ) + uow.concepts._session.execute( + update(concept_claims) + .where(concept_claims.c.id == claim_id) + .values(id=shared_id, concept_id="third", status="archived") + ) + uow.concepts._session.execute( + update(evidence_links) + .where( + evidence_links.c.target_type == "concept_claim", + evidence_links.c.target_id == claim_id, + ) + .values(target_id=shared_id) + ) + uow.concepts.add_lifecycle_event( + ConceptLifecycleEvent( + id="event", + repo_id="repo-a", + target_type=ConceptLifecycleTargetType.CLAIM, + target_id=shared_id, + from_status=ConceptLifecycleStatus.ACTIVE, + to_status=ConceptLifecycleStatus.ARCHIVED, + rationale="Old definition", + actor=ConceptCreatedBy.MANUAL, + ) + ) + statements = [] + connection = uow.concepts._session.connection() + + def count(*args): + statements.append(args[2]) + + event.listen(connection, "before_cursor_execute", count) + try: + assert ( + uow.concepts.get_concept_bundles(repo_id="repo-a", concept_ids=[]) == {} + ) + assert statements == [] + bundles = uow.concepts.get_concept_bundles( + repo_id="repo-a", + concept_ids=["third", second_id, "missing", "foreign", first_id], + include_lifecycle_events=True, + ) + finally: + event.remove(connection, "before_cursor_execute", count) + assert len(statements) == 9 + assert set(bundles) == {first_id, second_id, "third"} + assert bundles[first_id]["relations"] == bundles[second_id]["relations"] + assert bundles[first_id]["lifecycle_events"] == [] + assert bundles["third"]["lifecycle_events"][0].id == "event" + assert bundles["third"]["claims"][0].lifecycle.status.value == "archived" + assert {item.target_type.value for item in bundles["third"]["evidence"]} == { + "claim" + } + assert {item.target_type.value for item in bundles[first_id]["evidence"]} == { + "relation" + } + assert ( + bundles[second_id]["anchors"][0].id + == bundles[second_id]["groundings"][0].anchor_id + ) + for concept_id, bundle in bundles.items(): + assert bundle == uow.concepts.get_concept_bundle( + repo_id="repo-a", concept_ref=concept_id, include_lifecycle_events=True + ) + assert bundles == uow.concepts.get_concept_bundles( + repo_id="repo-a", + concept_ids=list(reversed(bundles)), + include_lifecycle_events=True, + ) diff --git a/tests/operations/create/execution/embeddings/test_create_embedding_persistence.py b/tests/operations/create/execution/embeddings/test_create_embedding_persistence.py index e74d3fe..0576cf4 100644 --- a/tests/operations/create/execution/embeddings/test_create_embedding_persistence.py +++ b/tests/operations/create/execution/embeddings/test_create_embedding_persistence.py @@ -19,7 +19,10 @@ from tests.operations._shared.id_generators import SequenceIdGenerator from app.infrastructure.db.runtime.engine import get_engine from app.infrastructure.embeddings.local_provider import ( - SentenceTransformersEmbeddingProvider, + MODEL_REPOSITORY, + MODEL_REVISION, + MODEL_FILES, + OnnxEmbeddingProvider, ) from app.infrastructure.db.runtime.models.episodes import episode_events, episodes from app.infrastructure.db.runtime.models.memories import memory_embeddings @@ -76,7 +79,16 @@ def test_create_persists_memory_embedding_row() -> None: } ) - provider = SentenceTransformersEmbeddingProvider(model="all-MiniLM-L6-v2") + from huggingface_hub import snapshot_download + + cache = os.environ["HF_HOME"] + snapshot_download( + MODEL_REPOSITORY, + revision=MODEL_REVISION, + cache_dir=cache, + allow_patterns=list(MODEL_FILES), + ) + provider = OnnxEmbeddingProvider(cache_folder=cache) with PostgresUnitOfWork(get_session_factory(engine)) as uow: uow.episodes.create_episode( diff --git a/tests/operations/create/execution/embeddings/test_embedding_providers.py b/tests/operations/create/execution/embeddings/test_embedding_providers.py index 1a1573a..59e9321 100644 --- a/tests/operations/create/execution/embeddings/test_embedding_providers.py +++ b/tests/operations/create/execution/embeddings/test_embedding_providers.py @@ -1,82 +1,70 @@ -"""Embedding provider contracts for create execution.""" +"""Check local-only embedding execution and MiniLM pooling semantics.""" import sys -import types +from types import SimpleNamespace +import numpy as np import pytest from app.infrastructure.embeddings.local_provider import ( - SentenceTransformersEmbeddingProvider, + MODEL_FILES, + OnnxEmbeddingProvider, + model_directory, ) -def test_sentence_transformers_provider_uses_local_library_when_available( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """local embedding providers should always return embeddings when sentence-transformers is available.""" +def test_missing_artifacts_require_upgrade_without_downloading(tmp_path): + provider = OnnxEmbeddingProvider(cache_folder=str(tmp_path)) + with pytest.raises(RuntimeError, match="shellbrain upgrade"): + provider.embed("query") - captured: dict[str, object] = {} - class _FakeModel: - """This helper class returns a fixed embedding payload for test assertions.""" +def test_embedding_masks_padding_normalizes_and_preserves_tokenizer_rules( + tmp_path, monkeypatch +): + directory = model_directory(str(tmp_path)) + for name in MODEL_FILES: + path = directory / name + path.parent.mkdir(parents=True, exist_ok=True) + path.touch() + calls = [] - def encode( - self, text: str, *, convert_to_numpy: bool, normalize_embeddings: bool - ): - _ = (text, convert_to_numpy, normalize_embeddings) - return [0.1, 0.2, 0.3] + class Tokenizer: + @staticmethod + def from_file(path): + assert path == str(directory / "tokenizer.json") + return Tokenizer() - class _FakeSentenceTransformer: - """This helper class mimics the sentence-transformers model constructor.""" + def enable_truncation(self, *, max_length): + assert max_length == 256 - def __init__( - self, - model_name: str, - *, - cache_folder: str | None = None, - local_files_only: bool = False, - ) -> None: - captured["model_name"] = model_name - captured["cache_folder"] = cache_folder - captured["local_files_only"] = local_files_only - self._model = _FakeModel() - - def encode( - self, text: str, *, convert_to_numpy: bool, normalize_embeddings: bool - ): - return self._model.encode( - text, - convert_to_numpy=convert_to_numpy, - normalize_embeddings=normalize_embeddings, + def encode(self, text): + calls.append(text) + return SimpleNamespace( + ids=[101, 42, 0], attention_mask=[1, 1, 0], type_ids=[0, 0, 0] ) - monkeypatch.setitem( - sys.modules, - "sentence_transformers", - types.SimpleNamespace(SentenceTransformer=_FakeSentenceTransformer), - ) - - provider = SentenceTransformersEmbeddingProvider( - model="all-MiniLM-L6-v2", - cache_folder="/tmp/shellbrain-models", - local_files_only=True, - ) - assert provider.embed("hello") == [0.1, 0.2, 0.3] - assert captured == { - "model_name": "all-MiniLM-L6-v2", - "cache_folder": "/tmp/shellbrain-models", - "local_files_only": True, - } - + class Session: + def __init__(self, path, *, sess_options, providers): + assert path == str(directory / "onnx/model.onnx") + assert ( + sess_options.intra_op_num_threads + == sess_options.inter_op_num_threads + == 1 + ) + assert providers == ["CPUExecutionProvider"] -def test_sentence_transformers_provider_raises_without_library( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """local embedding providers should always fail fast when sentence-transformers is unavailable.""" + def run(self, outputs, inputs): + assert all(value.dtype == np.int64 for value in inputs.values()) + return [np.array([[[3, 0], [0, 4], [999, 999]]], dtype=np.float32)] + monkeypatch.setitem(sys.modules, "tokenizers", SimpleNamespace(Tokenizer=Tokenizer)) monkeypatch.setitem( - sys.modules, "sentence_transformers", types.ModuleType("sentence_transformers") + sys.modules, + "onnxruntime", + SimpleNamespace(SessionOptions=SimpleNamespace, InferenceSession=Session), ) - provider = SentenceTransformersEmbeddingProvider(model="all-MiniLM-L6-v2") - with pytest.raises(RuntimeError): - provider.embed("hello") + provider = OnnxEmbeddingProvider(cache_folder=str(tmp_path)) + assert provider.embed(" query \n") == pytest.approx([0.6, 0.8]) + assert provider.embed("") == pytest.approx([0.6, 0.8]) + assert calls == ["query", ""] diff --git a/tests/operations/recall/execution/test_deterministic_graph_recall.py b/tests/operations/recall/execution/test_deterministic_graph_recall.py index f1bfce3..79c190d 100644 --- a/tests/operations/recall/execution/test_deterministic_graph_recall.py +++ b/tests/operations/recall/execution/test_deterministic_graph_recall.py @@ -189,10 +189,14 @@ def test_recall_preserves_case_pairings(monkeypatch, seed_both_endpoints, exclud lambda **kw: [ { **row, - "visible_memory_ids": (row["subject_memory_id"], row["object_memory_id"]), + "visible_memory_ids": ( + row["subject_memory_id"], + row["object_memory_id"], + ), } for row in relations - if kw["anchor_memory_id"] in (row["subject_memory_id"], row["object_memory_id"]) + if kw["anchor_memory_id"] + in (row["subject_memory_id"], row["object_memory_id"]) and row["predicate"] in kw["predicates"] ], ) @@ -209,19 +213,27 @@ def test_recall_preserves_case_pairings(monkeypatch, seed_both_endpoints, exclud expected = [ {**row, "validated_at": row["validated_at"].isoformat()} for row in relations - if excluded != "budget" and (excluded != "hidden" or row["object_memory_id"] != "s2") + if excluded != "budget" + and (excluded != "hidden" or row["object_memory_id"] != "s2") ] assert sorted(pack["memory_relations"], key=str) == sorted(expected, key=str) synthesis = synthesis_pack_from_graph_pack(pack) read = execute_read_memory(request, uow).data["pack"] - assert synthesis["memory_relations"] == read["memory_relations"] == pack["memory_relations"] + assert ( + synthesis["memory_relations"] + == read["memory_relations"] + == pack["memory_relations"] + ) cases = deterministic_brief_from_graph_pack(pack)["prior_cases"] for row in expected: subject = uow.memories._memories[row["subject_memory_id"]].text target = uow.memories._memories[row["object_memory_id"]].text assert any( - subject in case and target in case and row["predicate"] in case - and row["status"] in case and row["validated_at"] in case + subject in case + and target in case + and row["predicate"] in case + and row["status"] in case + and row["validated_at"] in case for case in cases ) assert not any("lock is held" in case and "Free disk" in case for case in cases) @@ -241,15 +253,31 @@ def _request(*, query: str) -> MemoryReadRequest: @pytest.mark.parametrize("linked_solutions", [(0, 1, 2), (5,)]) def test_fast_brief_preserves_case_order_without_duplicate_solutions(linked_solutions): memories = [ - {"id": f"s{i}", "kind": "solution", "text": f"Fix {i}.", "currentness": "current"} + { + "id": f"s{i}", + "kind": "solution", + "text": f"Fix {i}.", + "currentness": "current", + } for i in range(6) ] + [ - {"id": f"p{i}", "kind": "problem", "text": f"Problem {i}.", "currentness": "current"} + { + "id": f"p{i}", + "kind": "problem", + "text": f"Problem {i}.", + "currentness": "current", + } for i in linked_solutions ] relations = [ - {"subject_memory_id": f"p{i}", "predicate": "solved_by", "object_memory_id": f"s{i}", - "status": "active", "confidence": None, "validated_at": None} + { + "subject_memory_id": f"p{i}", + "predicate": "solved_by", + "object_memory_id": f"s{i}", + "status": "active", + "confidence": None, + "validated_at": None, + } for i in reversed(linked_solutions) ] cases = deterministic_brief_from_graph_pack( @@ -373,6 +401,14 @@ def find_concepts_for_memory_ids(self, **kwargs): del kwargs return [] + def get_concept_bundles(self, *, repo_id, concept_ids): + return { + key: bundle + for key in concept_ids + if (bundle := self.get_concept_bundle(repo_id=repo_id, concept_ref=key)) + is not None + } + def get_concept_bundle(self, *, repo_id: str, concept_ref: str): if repo_id != "repo-a": return None @@ -500,7 +536,7 @@ def get_concept_bundle(self, *, repo_id: str, concept_ref: str): } -class _NoFakeConcepts: +class _NoFakeConcepts(_FakeConcepts): def find_concepts_for_memory_ids(self, **kwargs): del kwargs return [] @@ -613,7 +649,11 @@ def test_historical_facets_do_not_lower_active_concept_rank( request = _request(query="TimeoutError in app/core/settings.py") candidates = {"c-db": {"score": 1.0, "why": []}} before, _ = _select_concepts( - request=request, concept_candidates=candidates, uow=uow + request=request, + concept_candidates=candidates, + concept_bundles=uow.concepts.get_concept_bundles( + repo_id=request.repo_id, concept_ids=list(candidates) + ), ) get_bundle = uow.concepts.get_concept_bundle @@ -634,7 +674,13 @@ def with_history(**kwargs): return bundle monkeypatch.setattr(uow.concepts, "get_concept_bundle", with_history) - after, _ = _select_concepts(request=request, concept_candidates=candidates, uow=uow) + after, _ = _select_concepts( + request=request, + concept_candidates=candidates, + concept_bundles=uow.concepts.get_concept_bundles( + repo_id=request.repo_id, concept_ids=list(candidates) + ), + ) assert [entry["bundle"]["concept"].slug for entry in after] == ["db-admin"] assert after[0]["score"] == before[0]["score"] pack = build_deterministic_graph_pack(request=request, uow=uow) diff --git a/tests/operations/recovery/execution/test_init_bootstrap.py b/tests/operations/recovery/execution/test_init_bootstrap.py index d039966..20a3860 100644 --- a/tests/operations/recovery/execution/test_init_bootstrap.py +++ b/tests/operations/recovery/execution/test_init_bootstrap.py @@ -20,7 +20,9 @@ from app.infrastructure.db.admin.provisioning import managed_local as managed_runtime from app.infrastructure.db.admin.provisioning import docker_prerequisites from app.infrastructure.local_state.repo_registration_store import RepoRegistration -from app.entrypoints.cli.presenters.init import render_success_lines as render_init_lines +from app.entrypoints.cli.presenters.init import ( + render_success_lines as render_init_lines, +) @pytest.fixture(autouse=True) @@ -54,9 +56,7 @@ def test_run_init_should_block_when_corrupt_config_cannot_be_recovered( init_module, "try_load_machine_config", lambda: (None, "corrupt toml") ) monkeypatch.setattr(init_module, "backup_corrupt_machine_config", lambda: preserved) - monkeypatch.setattr( - init_module, "_recover_machine_config", lambda: None - ) + monkeypatch.setattr(init_module, "_recover_machine_config", lambda: None) monkeypatch.setattr( "app.startup.runtime_admin.save_recovery_stub", lambda **kwargs: captured_stub.update(kwargs), @@ -434,8 +434,9 @@ def test_run_init_should_run_maintenance_when_machine_already_ready( monkeypatch.setattr( init_module, "prewarm_embeddings", - lambda config, skip_model_download: (_ for _ in ()).throw( - AssertionError("ready maintenance should not prewarm embeddings") + lambda config, skip_model_download: ( + call_order.append("embeddings") or False, + config, ), ) monkeypatch.setattr( @@ -448,9 +449,7 @@ def test_run_init_should_run_maintenance_when_machine_already_ready( monkeypatch.setattr( init_module, "install_host_assets", - lambda host_mode, force: ( - call_order.append("assets") or _FakeHostAssetsResult() - ), + lambda host_mode, force: call_order.append("assets") or _FakeHostAssetsResult(), ) result = init_module.run_init( @@ -471,12 +470,14 @@ def test_run_init_should_run_maintenance_when_machine_already_ready( "wait", "reconcile", "migrate", + "embeddings", "assets", ] +@pytest.mark.parametrize("embedding_failed", [False, True]) def test_ready_external_init_should_run_maintenance_without_docker( - tmp_path: Path, monkeypatch + tmp_path: Path, monkeypatch, embedding_failed ) -> None: """ready external-postgres init should migrate the DB without managed Docker setup.""" @@ -486,6 +487,22 @@ def test_ready_external_init_should_run_maintenance_without_docker( ready_config = _external_machine_config( bootstrap_state="ready", readiness_state="ready", last_error=None ) + from dataclasses import replace + + warmed_config = ( + replace( + ready_config, + bootstrap_state="repair_needed", + embeddings=replace( + ready_config.embeddings, + readiness_state="failed", + last_error="download failed", + ), + ) + if embedding_failed + else ready_config + ) + saved = [] call_order: list[str] = [] monkeypatch.setattr(init_module, "get_shellbrain_home", lambda: home_root) @@ -504,7 +521,7 @@ def test_ready_external_init_should_run_maintenance_without_docker( AssertionError("ready maintenance should not ask storage questions") ), ) - monkeypatch.setattr(init_module, "save_machine_config", lambda config: None) + monkeypatch.setattr(init_module, "save_machine_config", saved.append) monkeypatch.setattr( init_module, "_migrate_machine_config", @@ -540,8 +557,9 @@ def test_ready_external_init_should_run_maintenance_without_docker( monkeypatch.setattr( init_module, "prewarm_embeddings", - lambda config, skip_model_download: (_ for _ in ()).throw( - AssertionError("ready maintenance should not prewarm embeddings") + lambda config, skip_model_download: ( + call_order.append("embeddings") or embedding_failed, + warmed_config, ), ) monkeypatch.setattr( @@ -561,8 +579,13 @@ def test_ready_external_init_should_run_maintenance_without_docker( render_success_lines=_render_success_lines, ) - assert result.outcome == init_module.INIT_OUTCOME_NOOP - assert call_order == ["config", "wait", "reconcile", "migrate"] + if embedding_failed: + assert result.outcome == init_module.INIT_OUTCOME_BLOCKED_DEPENDENCY + assert saved[-1].bootstrap_state == "repair_needed" + assert result.lines == ["download failed"] + else: + assert result.outcome == init_module.INIT_OUTCOME_NOOP + assert call_order == ["config", "wait", "reconcile", "migrate", "embeddings"] def test_ready_init_should_propagate_schema_migration_failure( @@ -1107,3 +1130,100 @@ def __enter__(self): def __exit__(self, exc_type, exc, tb) -> bool: return False + + +@pytest.mark.parametrize("existing_ready", [False, True]) +def test_prewarm_installs_onnx_for_fresh_and_existing_installations( + tmp_path, monkeypatch, existing_ready +): + from dataclasses import replace + from app.infrastructure.embeddings import prewarm + from app.infrastructure.embeddings.local_provider import ( + MODEL_FILES, + MODEL_REVISION, + model_directory, + ) + import huggingface_hub + + config = _machine_config( + bootstrap_state="ready" if existing_ready else "provisioning" + ) + config = replace( + config, + embeddings=replace( + config.embeddings, + cache_path=str(tmp_path), + readiness_state="ready" if existing_ready else "pending", + ), + ) + downloads = [] + + def download(repo_id, *, revision, cache_dir, allow_patterns, force_download): + downloads.append(repo_id) + assert force_download is True + assert revision == MODEL_REVISION + assert allow_patterns == list(MODEL_FILES) + for name in allow_patterns: + path = model_directory(cache_dir) / name + path.parent.mkdir(parents=True, exist_ok=True) + path.touch() + + monkeypatch.setattr(huggingface_hub, "snapshot_download", download) + + def embed(self, text): + if not all( + (model_directory(str(tmp_path)) / name).is_file() for name in MODEL_FILES + ): + raise RuntimeError("missing model files") + return [0.0] * 384 + + monkeypatch.setattr(prewarm.OnnxEmbeddingProvider, "embed", embed) + changed, ready = prewarm.prewarm_embeddings(config, skip_model_download=False) + assert changed and ready.embeddings.readiness_state == "ready" + assert ready.embeddings.provider == "onnxruntime" + assert ready.embeddings.model_revision == MODEL_REVISION + assert len(downloads) == 1 + assert prewarm.prewarm_embeddings(ready, skip_model_download=False) == ( + False, + ready, + ) + assert len(downloads) == 1 + + +def test_prewarm_failure_keeps_installation_unready(tmp_path, monkeypatch): + from dataclasses import replace + import huggingface_hub + from app.infrastructure.embeddings.prewarm import prewarm_embeddings + + config = _machine_config( + bootstrap_state="ready", readiness_state="ready", last_error=None + ) + config = replace( + config, embeddings=replace(config.embeddings, cache_path=str(tmp_path)) + ) + + def interrupted(*args, **kwargs): + raise OSError("interrupted model download") + + monkeypatch.setattr(huggingface_hub, "snapshot_download", interrupted) + changed, failed = prewarm_embeddings(config, skip_model_download=False) + assert changed + assert failed.bootstrap_state == "repair_needed" + assert failed.embeddings.readiness_state == "failed" + assert failed.embeddings.last_error == "interrupted model download" + + +def test_prewarm_rejects_unverified_cached_revision(tmp_path): + from dataclasses import replace + from app.infrastructure.embeddings.prewarm import prewarm_embeddings + + config = _machine_config(bootstrap_state="ready") + config = replace( + config, embeddings=replace(config.embeddings, cache_path=str(tmp_path)) + ) + ref = tmp_path / "models--sentence-transformers--all-MiniLM-L6-v2/refs/main" + ref.parent.mkdir(parents=True) + ref.write_text("different-revision") + _, failed = prewarm_embeddings(config, skip_model_download=False) + assert failed.embeddings.readiness_state == "failed" + assert "Re-embed" in failed.embeddings.last_error