diff --git a/README.md b/README.md index d1af575c..25e2b299 100644 --- a/README.md +++ b/README.md @@ -347,8 +347,17 @@ Measured on an Apple M2 Max (64 GB). The harnesses are the proof — run them yo | `IAI_MCP_PYTHON` | — | Absolute path to the venv Python (for the MCP host config) | | `IAI_MCP_RECALL_CONCURRENCY` | `2` | Maximum cued `memory_recall` calls dispatched concurrently by the socket daemon | | `IAI_MCP_RECALL_SLOT_WAIT_SEC` | `0.25` | How long an overflow cued recall waits for a slot before returning `_degraded: recall_busy` | - -The old `IAI_MCP_EMBED_MODEL` knob is gone — the embedder is a single built-in English-only model. There are many internal tuning knobs (`IAI_MCP_*`), but you shouldn't need to touch them. +| `IAI_MCP_EMBED_PROVIDER` | `native` | `native` for built-in BGE or `http` for a replaceable loopback provider | +| `IAI_MCP_EMBED_URL` | — | Loopback endpoint or base URL for the `http` provider | +| `IAI_MCP_EMBED_DIM` | `384` | Vector dimension; required for the `http` provider | +| `IAI_MCP_EMBED_MODEL_ID` | — | Model identifier; required for the `http` provider | +| `IAI_MCP_EMBED_TIMEOUT_SEC` | `30` | Local provider request timeout | + +The built-in Rust BGE model remains the zero-configuration default. Setting the +provider to `http` replaces it completely: the native model is not constructed, +downloaded, or run. This makes multilingual and domain-specific embedders +possible without adding a Python ML stack to iai-mcp. See +[`docs/EMBEDDERS.md`](docs/EMBEDDERS.md) for the protocol and migration steps. --- diff --git a/docs/DEPLOYMENT.md b/docs/DEPLOYMENT.md index ef45677f..82670ef7 100644 --- a/docs/DEPLOYMENT.md +++ b/docs/DEPLOYMENT.md @@ -11,8 +11,8 @@ per-release log see [`CHANGELOG.md`](../CHANGELOG.md). |---|---| | Python | 3.11–3.12 (CPython) | | OS | macOS or Linux. The daemon uses `fcntl.flock` and Unix-socket IPC. Windows support is in beta. WSL2 works as a Linux target. | -| RAM | 8+ GB comfortable. The `bge-small-en-v1.5` embedder occupies ~600 MB resident once loaded. | -| Disk | ~5 GB free for model weights + store + WAL. Model weights live in `~/.cache/huggingface/` (~130 MB). | +| RAM | 8+ GB comfortable. The default `bge-small-en-v1.5` embedder occupies ~600 MB resident once loaded; external-provider usage depends on the selected model. | +| Disk | ~5 GB free for model weights + store + WAL. Default model weights live in `~/.cache/huggingface/` (~130 MB); external-provider usage varies. | | Toolchain (source build only) | A Rust toolchain is needed when compiling the native extension from source. On Linux, `libssl-dev` and `pkg-config` (or your distro's equivalents). | The native extension (`iai_mcp_native` — the embedder, graph algorithms, and diff --git a/docs/EMBEDDERS.md b/docs/EMBEDDERS.md new file mode 100644 index 00000000..d27351a0 --- /dev/null +++ b/docs/EMBEDDERS.md @@ -0,0 +1,74 @@ +# Replaceable embedding providers + +iai-mcp uses the native Rust `bge-small-en-v1.5` embedder by default. A local +HTTP provider can replace it completely for multilingual, domain-specific, or +shared-model deployments. iai-mcp does not import or construct the native +embedder when the HTTP provider is selected. + +## Configuration + +```bash +export IAI_MCP_EMBED_PROVIDER=http +export IAI_MCP_EMBED_URL=http://127.0.0.1:4488 +export IAI_MCP_EMBED_DIM=1024 +export IAI_MCP_EMBED_MODEL_ID=your-model-id +export IAI_MCP_EMBED_TIMEOUT_SEC=30 +``` + +Only unauthenticated loopback HTTP URLs are accepted. The model service stays +local, and iai-mcp adds no ML framework dependency. Multiple iai-mcp processes +can share one warm model service instead of loading one model copy per store. + +## Protocol + +iai-mcp sends `POST /embed` unless `IAI_MCP_EMBED_URL` already ends in +`/embed`: + +```json +{ + "texts": ["What should be recalled?"], + "input_type": "query" +} +``` + +`input_type` is either `query` for retrieval cues or `document` for memories. +The provider owns tokenization, prefixes, pooling, normalization, batching, and +model loading. This distinction supports asymmetric models without teaching +iai-mcp about any specific model family. + +The response is: + +```json +{ + "model": "your-model-id", + "dimensions": 1024, + "vectors": [[0.01, -0.02, 0.03]] +} +``` + +iai-mcp rejects a wrong model identifier, vector count, dimension, non-numeric +value, or non-finite value. + +## Migrating an existing store + +Changing a model invalidates every stored vector, even when the old and new +models use the same dimension. Stop the daemon, make a store backup, configure +the new provider, and run: + +```bash +iai-mcp migrate --reembed-to-configured-provider --dry-run +iai-mcp migrate --reembed-to-configured-provider +``` + +The migration flushes pending in-process writes, stages a complete replacement +table, preserves storage-only fields and encrypted payloads, keeps the previous +records table, and updates the persisted dimension. Restart iai-mcp before +recall so its vector indexes reopen with the new dimension, then run: + +```bash +iai-mcp doctor +``` + +Starting a populated store with a provider whose dimension differs from the +store fails fast and points to the migration command. This prevents mixed-model +or mixed-dimension indexes from being used silently. diff --git a/mcp-wrapper/src/tools.ts b/mcp-wrapper/src/tools.ts index 319ec0ff..dcb9f159 100644 --- a/mcp-wrapper/src/tools.ts +++ b/mcp-wrapper/src/tools.ts @@ -52,7 +52,7 @@ export const toolSchemas: Record = { type: "string", description: "Natural-language query to match against stored memories. " + - "Embedded server-side via bge-small-en-v1.5 (384d) unless " + + "Embedded server-side by the configured provider unless " + "`cue_embedding` is supplied.", }, budget_tokens: { @@ -74,7 +74,7 @@ export const toolSchemas: Record = { items: { type: "number" }, description: "Optional pre-computed embedding vector for the cue " + - "(EMBED_DIM=384 floats; bge-small-en-v1.5). " + + "(its dimension must match the current store). " + "When omitted, the daemon embeds the cue server-side. " + "Used by memory_contradict and tests that need byte-stable embeddings.", }, @@ -167,7 +167,7 @@ export const toolSchemas: Record = { items: { type: "number" }, description: "Optional pre-computed embedding vector for the contradicting " + - "fact (EMBED_DIM=384 floats; bge-small-en-v1.5). When omitted, " + + "fact (its dimension must match the current store). When omitted, " + "the daemon embeds new_fact server-side.", }, }, diff --git a/src/iai_mcp/brainview.py b/src/iai_mcp/brainview.py index fe6a3c0c..0078b18e 100644 --- a/src/iai_mcp/brainview.py +++ b/src/iai_mcp/brainview.py @@ -1189,9 +1189,9 @@ def search_direct(self, query: str, k: int = 12) -> dict: if not query: return {"hits": []} try: - from iai_mcp.embed import embedder_for_store + from iai_mcp.embed import embed_query, embedder_for_store - emb = embedder_for_store(self.store).embed(query[:512]) + emb = embed_query(embedder_for_store(self.store), query[:512]) raw = self.store.query_similar(list(emb), k=max(1, min(int(k), 24))) except Exception as exc: # noqa: BLE001 -- search is navigation, degrade to empty logger.warning("brainview search failed: %s", exc) diff --git a/src/iai_mcp/cli/__init__.py b/src/iai_mcp/cli/__init__.py index 32e115e7..046a91c5 100644 --- a/src/iai_mcp/cli/__init__.py +++ b/src/iai_mcp/cli/__init__.py @@ -402,13 +402,22 @@ def _build_parser() -> argparse.ArgumentParser: "embedding column is rewritten -- literal_surface is never touched." ), ) + m.add_argument( + "--reembed-to-configured-provider", + action="store_true", + help=( + "Re-embed every record with the configured embedding provider, " + "including dimension changes. Uses the crash-safe staging-table " + "migration and retains the previous table until cleanup." + ), + ) m.add_argument( "--reembed-batch-size", type=int, default=256, help=( - "Records per id-ordered window for --reembed-from-text. Bounds " - "memory; embed calls are streamed within each window. Default 256." + "Records per batch for --reembed-from-text or " + "--reembed-to-configured-provider. Default 256." ), ) m.add_argument( diff --git a/src/iai_mcp/cli/_analytics.py b/src/iai_mcp/cli/_analytics.py index b7fe2ec7..ebe3881d 100644 --- a/src/iai_mcp/cli/_analytics.py +++ b/src/iai_mcp/cli/_analytics.py @@ -99,6 +99,11 @@ def cmd_migrate(args: argparse.Namespace) -> int: from iai_mcp import cli as _cli from iai_mcp.store import MemoryStore store = MemoryStore() + verbose = bool(getattr(args, "verbose", False)) + + def _progress(i: int, n: int) -> None: + if verbose: + print(f"[{i + 1}/{n}] migrating...") # --reembed-from-text owns --resume / --rollback when combined with it: # resume continues the reembed from its last committed window, rollback @@ -124,13 +129,40 @@ def cmd_migrate(args: argparse.Namespace) -> int: ) return 0 + if bool(getattr(args, "reembed_to_configured_provider", False)): + from iai_mcp.embed import Embedder + from iai_mcp.migrate import migrate_reembed_to_current_dim + + target = Embedder() + dry_run = bool(getattr(args, "dry_run", False)) + batch_size = int(getattr(args, "reembed_batch_size", 256)) + result = migrate_reembed_to_current_dim( + store, + target, + dry_run=dry_run, + progress=_progress, + force=True, + batch_size=batch_size, + ) + prefix = "[dry-run] would re-embed" if dry_run else "re-embedded" + count = result.get("would_update", result.get("updated", 0)) + print( + f"{prefix} {count} records from {result['source_dim']}d to " + f"{result['target_dim']}d with {target.model_key}" + ) + if result.get("restart_required"): + print( + "restart IAE before recall so the vector indexes reopen at the new dimension" + ) + return 0 + if bool(getattr(args, "rollback", False)): from iai_mcp import migrate return migrate._rollback(store.db, store) if bool(getattr(args, "resume", False)): from iai_mcp import migrate - from iai_mcp.embed import embedder_for_store - target = embedder_for_store(store) + from iai_mcp.embed import Embedder + target = Embedder() return migrate._resume(store.db, store, target) if bool(getattr(args, "rederive_timestamps", False)): @@ -171,11 +203,6 @@ def cmd_migrate(args: argparse.Namespace) -> int: from_v = int(getattr(args, "from_", 1)) to_v = int(getattr(args, "to", 2)) dry_run = bool(getattr(args, "dry_run", False)) - verbose = bool(getattr(args, "verbose", False)) - - def _progress(i: int, n: int) -> None: - if verbose: - print(f"[{i + 1}/{n}] migrating...") if from_v == 1 and to_v == 2: from iai_mcp.migrate import migrate_v1_to_v2 diff --git a/src/iai_mcp/core/__init__.py b/src/iai_mcp/core/__init__.py index 9e8c795c..073f59cf 100644 --- a/src/iai_mcp/core/__init__.py +++ b/src/iai_mcp/core/__init__.py @@ -101,7 +101,7 @@ def _crisis_degraded_recall(store: MemoryStore, params: dict) -> dict: """ try: from iai_mcp.cue_router import _classify_cue - from iai_mcp.embed import embedder_for_store + from iai_mcp.embed import embed_query, embedder_for_store from iai_mcp.pipeline import K_CANDIDATES from iai_mcp.core._serializers import _hit_to_json from iai_mcp.types import MemoryHit @@ -109,7 +109,7 @@ def _crisis_degraded_recall(store: MemoryStore, params: dict) -> dict: cue_mode, _cue_intent, _triggered_pattern = _classify_cue(params.get("cue", "")) embedder = embedder_for_store(store) - _cue_vec = embedder.embed(params["cue"]) + _cue_vec = embed_query(embedder, params["cue"]) _ann_pairs = store.query_similar(_cue_vec, k=K_CANDIDATES) _candidate_recs: dict = {_r.id: _r for _r, _s in _ann_pairs} @@ -368,7 +368,7 @@ def _trace_mark(_name: str) -> None: mode=cue_mode, ) else: - from iai_mcp.embed import embedder_for_store + from iai_mcp.embed import embed_query, embedder_for_store from iai_mcp.pipeline import recall_for_response try: from iai_mcp.daemon_state import load_state as _ds_load @@ -400,7 +400,7 @@ def _trace_mark(_name: str) -> None: _encode_ms: "float | None" = None _encode_t0 = _time.perf_counter() try: - _cue_vec = embedder.embed(params["cue"]) + _cue_vec = embed_query(embedder, params["cue"]) _encode_ms = (_time.perf_counter() - _encode_t0) * 1000.0 _trace_mark("encode") except Exception as _emb_exc: @@ -960,9 +960,9 @@ def _trace_mark(_name: str) -> None: except Exception as exc: # noqa: BLE001 -- one lane failing must not blank the other logger.debug("memory_search lexical lane failed: %s", exc) try: - from iai_mcp.embed import embedder_for_store + from iai_mcp.embed import embed_query, embedder_for_store - vec = embedder_for_store(store).embed(query[:512]) + vec = embed_query(embedder_for_store(store), query[:512]) for sem_rank, (rec, cos) in enumerate(store.query_similar(list(vec), k=k)): rid = str(rec.id) if rid in merged: @@ -1248,7 +1248,7 @@ def _trace_mark(_name: str) -> None: if method == "memory_temporal_recall": from iai_mcp.events import flush_event_buffer, query_events - from iai_mcp.embed import embedder_for_store + from iai_mcp.embed import embed_query, embedder_for_store from iai_mcp.store._store import _normalize_ts_for_compare try: @@ -1284,7 +1284,7 @@ def _trace_mark(_name: str) -> None: cue_vec: list[float] | None = None if cue: embedder = embedder_for_store(store) - cue_vec = embedder.embed(cue) + cue_vec = embed_query(embedder, cue) record_hits = store.query_similar_temporal( vec=cue_vec, as_of=as_of_norm, k=limit, ) diff --git a/src/iai_mcp/daemon/__init__.py b/src/iai_mcp/daemon/__init__.py index 11c67d31..3336d12c 100644 --- a/src/iai_mcp/daemon/__init__.py +++ b/src/iai_mcp/daemon/__init__.py @@ -761,7 +761,6 @@ def _install_warm_embedder_override(store) -> tuple[object, bool]: orig_efs = _embed_mod.embedder_for_store try: warm = orig_efs(store) - def _held_embedder_for_store(_store): return warm diff --git a/src/iai_mcp/daemon/_boot_warmup.py b/src/iai_mcp/daemon/_boot_warmup.py index aeabdaec..d2f48c62 100644 --- a/src/iai_mcp/daemon/_boot_warmup.py +++ b/src/iai_mcp/daemon/_boot_warmup.py @@ -137,11 +137,11 @@ def warm_dispatch_surface(store: Any) -> dict: from iai_mcp import runtime_graph_cache # noqa: F401 from iai_mcp.core import _serializers # noqa: F401 from iai_mcp.cue_router import _classify_cue # noqa: F401 - from iai_mcp.embed import embedder_for_store + from iai_mcp.embed import embed_query, embedder_for_store from iai_mcp.pipeline import K_CANDIDATES # noqa: F401 embedder = embedder_for_store(store) - embedder.embed("boot warm-up probe cue") + embed_query(embedder, "boot warm-up probe cue") runtime_graph_cache.load_recall_structural(store) return {"elapsed_ms": (time.perf_counter() - _t0) * 1000.0} except Exception as exc: # noqa: BLE001 -- warm-up must never break boot diff --git a/src/iai_mcp/doctor/_storage_checks.py b/src/iai_mcp/doctor/_storage_checks.py index b37f46af..b3aed15e 100644 --- a/src/iai_mcp/doctor/_storage_checks.py +++ b/src/iai_mcp/doctor/_storage_checks.py @@ -621,30 +621,32 @@ def check_u_recall_centrality_regression() -> CheckResult: def check_v_native_embedder() -> CheckResult: import math + backend = "unknown" try: - import iai_mcp_native # noqa: F401 - from iai_mcp.embed import Embedder + from iai_mcp.embed import Embedder, embed_query emb = Embedder() - assert emb._backend == "rust", f"backend={emb._backend!r}" - vec = emb.embed("smoke") - assert len(vec) == 384, f"expected 384 dims, got {len(vec)}" + backend = emb._backend + vec = embed_query(emb, "smoke") + assert len(vec) == emb.DIM, f"expected {emb.DIM} dims, got {len(vec)}" assert all(math.isfinite(float(x)) for x in vec[:3]), ( "non-finite values in output" ) except Exception as exc: # noqa: BLE001 + remedy = ( + "rebuild with: cd rust/iai_mcp_native && maturin develop --release" + if backend in {"unknown", "rust"} + else "check IAI_MCP_EMBED_URL and the loopback embedder service" + ) return CheckResult( - name="(v) native Rust embedder", + name="(v) configured embedder", passed=False, - detail=( - f"{type(exc).__name__}: {exc} — rebuild with: " - "cd rust/iai_mcp_native && maturin develop --release" - ), + detail=f"{type(exc).__name__}: {exc} — {remedy}", ) return CheckResult( - name="(v) native Rust embedder", + name="(v) configured embedder", passed=True, - detail="encode ok, backend=rust, 384-dim", + detail=f"encode ok, backend={backend}, {emb.DIM}-dim, model={emb.model_key}", ) diff --git a/src/iai_mcp/embed.py b/src/iai_mcp/embed.py index 8a07cf40..15fe9049 100644 --- a/src/iai_mcp/embed.py +++ b/src/iai_mcp/embed.py @@ -2,12 +2,16 @@ import logging import os +import json +import math from dataclasses import dataclass +from typing import Literal +from urllib.error import HTTPError, URLError +from urllib.parse import urlparse +from urllib.request import Request, urlopen import numpy as np -from iai_mcp_native import embed as _rust - logger = logging.getLogger(__name__) @@ -21,6 +25,10 @@ embed_failure_total: int = 0 +InputType = Literal["query", "document"] +VALID_PROVIDERS: set[str] = {"native", "http"} +MAX_HTTP_RESPONSE_BYTES = 32 * 1024 * 1024 + def _resolve_model_key(model_key: str | None = None) -> str: if model_key is not None: @@ -44,9 +52,129 @@ def _resolve_quantize_mode() -> str | None: return raw +def _resolve_provider() -> str: + provider = os.environ.get("IAI_MCP_EMBED_PROVIDER", "native").strip().lower() + if provider not in VALID_PROVIDERS: + raise ValueError( + f"IAI_MCP_EMBED_PROVIDER={provider!r} is not valid; " + f"valid: {sorted(VALID_PROVIDERS)}" + ) + return provider + + +def _resolve_http_config() -> tuple[str, int, float, str]: + raw_url = os.environ.get("IAI_MCP_EMBED_URL", "").strip() + if not raw_url: + raise ValueError("IAI_MCP_EMBED_URL is required for the http provider") + parsed = urlparse(raw_url) + if ( + parsed.scheme != "http" + or parsed.hostname not in {"127.0.0.1", "localhost", "::1"} + or parsed.username is not None + or parsed.password is not None + or parsed.query + or parsed.fragment + ): + raise ValueError( + "IAI_MCP_EMBED_URL must be an unauthenticated loopback http URL" + ) + if parsed.path in {"", "/"}: + raw_url = raw_url.rstrip("/") + "/embed" + elif parsed.path != "/embed": + raise ValueError("IAI_MCP_EMBED_URL path must be /embed or empty") + + raw_dim = os.environ.get("IAI_MCP_EMBED_DIM", "").strip() + try: + dim = int(raw_dim) + except ValueError as exc: + raise ValueError( + "IAI_MCP_EMBED_DIM must be a positive integer for the http provider" + ) from exc + if dim <= 0: + raise ValueError( + "IAI_MCP_EMBED_DIM must be a positive integer for the http provider" + ) + + raw_timeout = os.environ.get("IAI_MCP_EMBED_TIMEOUT_SEC", "30").strip() + try: + timeout = float(raw_timeout) + except ValueError as exc: + raise ValueError("IAI_MCP_EMBED_TIMEOUT_SEC must be positive") from exc + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError("IAI_MCP_EMBED_TIMEOUT_SEC must be positive") + + model_id = os.environ.get("IAI_MCP_EMBED_MODEL_ID", "").strip() + if not model_id: + raise ValueError("IAI_MCP_EMBED_MODEL_ID is required for the http provider") + return raw_url, dim, timeout, model_id + + +class _HttpEmbeddingBackend: + def __init__(self, url: str, dim: int, timeout: float, model_id: str) -> None: + self.url = url + self.dim = dim + self.timeout = timeout + self.model_id = model_id + + def encode_batch( + self, texts: list[str], *, input_type: InputType + ) -> list[list[float]]: + request = Request( + self.url, + data=json.dumps({"texts": texts, "input_type": input_type}).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urlopen(request, timeout=self.timeout) as response: + raw = response.read(MAX_HTTP_RESPONSE_BYTES + 1) + if len(raw) > MAX_HTTP_RESPONSE_BYTES: + raise RuntimeError("http embed provider response is too large") + payload = json.loads(raw) + except ( + HTTPError, + URLError, + TimeoutError, + UnicodeDecodeError, + json.JSONDecodeError, + ) as exc: + raise RuntimeError(f"http embed provider failed: {exc}") from exc + + vectors = payload.get("vectors") if isinstance(payload, dict) else None + dimensions = payload.get("dimensions") if isinstance(payload, dict) else None + model = payload.get("model") if isinstance(payload, dict) else None + if model != self.model_id: + raise ValueError( + f"http embed provider returned model {model!r}; " + f"expected {self.model_id!r}" + ) + if dimensions != self.dim: + raise ValueError( + f"http embed provider returned dimension {dimensions!r}; " + f"expected {self.dim}" + ) + if not isinstance(vectors, list) or len(vectors) != len(texts): + raise ValueError( + "http embed provider returned an unexpected number of vectors" + ) + + checked: list[list[float]] = [] + for vector in vectors: + if not isinstance(vector, list) or len(vector) != self.dim: + raise ValueError(f"http embed provider vector must be {self.dim}d") + if not all( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and math.isfinite(value) + for value in vector + ): + raise ValueError("http embed provider returned a non-finite vector") + checked.append([float(value) for value in vector]) + return checked + + @dataclass(frozen=True) class QuantizedVector: - values: list[int] scale: float zero_point: int @@ -74,7 +202,6 @@ def _quantize_int8(vec: list[float]) -> QuantizedVector: class Embedder: - DEFAULT_MODEL_KEY: str = DEFAULT_MODEL_KEY DEFAULT_DIM: int = MODEL_REGISTRY[DEFAULT_MODEL_KEY]["dim"] DEFAULT_MODEL: str = MODEL_REGISTRY[DEFAULT_MODEL_KEY]["hf"] @@ -86,6 +213,22 @@ def __init__( *, model_name: str | None = None, ) -> None: + self._quantize_mode: str | None = _resolve_quantize_mode() + provider = _resolve_provider() + if provider == "http": + if model_key is not None or model_name is not None: + raise ValueError( + "model_key and model_name only apply to the native provider" + ) + url, dim, timeout, model_id = _resolve_http_config() + self.model_key = model_id + self.model_name = model_id + self.DIM = dim + self._model = _HttpEmbeddingBackend(url, dim, timeout, model_id) + self._backend = "http" + self.supports_batch = True + return + if model_key is None and model_name is not None: match = next( (k for k, v in MODEL_REGISTRY.items() if v["hf"] == model_name), @@ -104,37 +247,71 @@ def __init__( self.model_name: str = spec["hf"] self.DIM: int = int(spec["dim"]) - self._model = _rust.Embedder() - self._backend: str = "rust" + from iai_mcp_native import embed as rust - self._quantize_mode: str | None = _resolve_quantize_mode() + self._model = rust.Embedder() + self._backend: str = "rust" + self.supports_batch = False - def _encode_one(self, text: str) -> list[float]: + def _encode_batch( + self, texts: list[str], *, input_type: InputType + ) -> list[list[float]]: global embed_failure_total + if input_type not in {"query", "document"}: + raise ValueError("input_type must be 'query' or 'document'") + if not texts: + return [] try: - return self._model.encode(text) + if self._backend == "http": + return self._model.encode_batch(texts, input_type=input_type) + return [self._model.encode(text) for text in texts] except Exception as exc: embed_failure_total += 1 logger.error( - "native embed encode failed: %s: %s", + "%s embed encode failed: %s: %s", + "native" if self._backend == "rust" else self._backend, type(exc).__name__, exc, ) raise - def embed(self, text: str) -> list[float]: - return self._encode_one(text) + def embed(self, text: str, *, input_type: InputType = "document") -> list[float]: + return self._encode_batch([text], input_type=input_type)[0] - def embed_batch(self, texts: list[str]) -> list[list[float]]: - return [self._encode_one(t) for t in texts] + def embed_query(self, text: str) -> list[float]: + return self.embed(text, input_type="query") - def embed_quantized(self, text: str) -> QuantizedVector: - fp32 = self.embed(text) + def embed_batch( + self, texts: list[str], *, input_type: InputType = "document" + ) -> list[list[float]]: + return self._encode_batch(texts, input_type=input_type) + + def embed_quantized( + self, text: str, *, input_type: InputType = "document" + ) -> QuantizedVector: + fp32 = self.embed(text, input_type=input_type) return _quantize_int8(fp32) +def embed_query(embedder, text: str) -> list[float]: + """Embed a retrieval cue while preserving compatibility with test doubles.""" + method = getattr(embedder, "embed_query", None) + if callable(method): + return method(text) + return embedder.embed(text) + + def embedder_for_store(store) -> "Embedder": target_dim = getattr(store, "embed_dim", None) + if _resolve_provider() == "http": + embedder = Embedder() + if target_dim is not None and int(target_dim) != embedder.DIM: + raise ValueError( + f"store uses {target_dim}d embeddings but the configured http " + f"provider uses {embedder.DIM}d; run the re-embedding migration " + "before starting IAE with this provider" + ) + return embedder if target_dim is None: return Embedder() preferred = {384: "bge-small-en-v1.5"} diff --git a/src/iai_mcp/foresight.py b/src/iai_mcp/foresight.py index 0f99ffb0..e3f11385 100644 --- a/src/iai_mcp/foresight.py +++ b/src/iai_mcp/foresight.py @@ -210,9 +210,9 @@ def _blended_cue(store: Any, cue_embedding: "list[float]") -> "list[float]": goal = (entry.goal or "").strip() if entry is not None else "" if len(goal) < 12: return list(cue_embedding) - from iai_mcp.embed import embedder_for_store # noqa: PLC0415 + from iai_mcp.embed import embed_query, embedder_for_store # noqa: PLC0415 - goal_vec = embedder_for_store(store).embed(goal[:512]) + goal_vec = embed_query(embedder_for_store(store), goal[:512]) blended = [ (1.0 - weight) * a + weight * b for a, b in zip(cue_embedding, goal_vec) @@ -249,10 +249,12 @@ def refresh_from_anchor(store: Any, embedder: Any) -> bool: store._foresight_anchor = None try: _ts, text, session_id = anchor + from iai_mcp.embed import embed_query + refresh_pack( store, cue_text=text, - cue_embedding=embedder.embed(text[:512]), + cue_embedding=embed_query(embedder, text[:512]), session_id=session_id, ) return True diff --git a/src/iai_mcp/hippo/_db.py b/src/iai_mcp/hippo/_db.py index 4d4334d2..4d878e7c 100644 --- a/src/iai_mcp/hippo/_db.py +++ b/src/iai_mcp/hippo/_db.py @@ -939,12 +939,19 @@ def _initialize_hnsw_index(self) -> None: active_label_count = len(self._label_map) hnsw_loaded_count = self._hnsw.get_current_count() - if active_label_count != sqlite_count or hnsw_loaded_count != sqlite_count: + vectors_match = self._loaded_hnsw_matches_sqlite_sample() + if ( + active_label_count != sqlite_count + or hnsw_loaded_count != sqlite_count + or not vectors_match + ): _log.info( - "Boot integrity check: labels=%d hnsw=%d sqlite=%d — rebuilding", + "Boot integrity check: labels=%d hnsw=%d sqlite=%d " + "vectors_match=%s — rebuilding", active_label_count, hnsw_loaded_count, sqlite_count, + vectors_match, ) self._rebuild_index_from_sqlite() @@ -954,6 +961,37 @@ def _initialize_hnsw_index(self) -> None: # fresh-alloc fallback path (the standby is intentionally absent then). self._allocate_standby_index(cap) + def _loaded_hnsw_matches_sqlite_sample(self) -> bool: + """Reject a stale on-disk ANN after a model or dimension migration.""" + rows = self._conn.execute( + "SELECT vec_label, embedding FROM records" + " WHERE tombstoned_at IS NULL" + " AND COALESCE(embedding_pending, 0) = 0" + " ORDER BY vec_label LIMIT 3" + ).fetchall() + rows += self._conn.execute( + "SELECT vec_label, embedding FROM records" + " WHERE tombstoned_at IS NULL" + " AND COALESCE(embedding_pending, 0) = 0" + " ORDER BY vec_label DESC LIMIT 3" + ).fetchall() + try: + for row in rows: + stored = np.frombuffer(row["embedding"], dtype=np.float32) + indexed = self._hnsw.get_items([int(row["vec_label"])])[0] + norm = float(np.linalg.norm(stored)) + expected = stored if norm == 0.0 else stored / norm + if indexed.shape != expected.shape or not np.allclose( + indexed, + expected, + rtol=1e-4, + atol=1e-5, + ): + return False + except (RuntimeError, ValueError, IndexError): + return False + return True + def _allocate_standby_index(self, cap: int) -> None: from iai_mcp.hippo import ( HNSW_EF_CONSTRUCTION, diff --git a/src/iai_mcp/lilli/brain.py b/src/iai_mcp/lilli/brain.py index d12b487b..223ded9e 100644 --- a/src/iai_mcp/lilli/brain.py +++ b/src/iai_mcp/lilli/brain.py @@ -65,11 +65,11 @@ def recall(self, cue: str, *, limit: int = 5, session_id: str = "brain-recall") raise RuntimeError( "Brain.recall requires hippo_conn (MemoryStore-like instance)" ) - from iai_mcp.embed import embedder_for_store + from iai_mcp.embed import embed_query, embedder_for_store from iai_mcp.retrieve import recall as _retrieve_recall embedder = embedder_for_store(self.hippo_conn) - cue_embedding = embedder.embed(cue) + cue_embedding = embed_query(embedder, cue) response = _retrieve_recall( self.hippo_conn, cue_embedding=cue_embedding, diff --git a/src/iai_mcp/migrate/_reembed.py b/src/iai_mcp/migrate/_reembed.py index 1f2b1403..8e32e8c6 100644 --- a/src/iai_mcp/migrate/_reembed.py +++ b/src/iai_mcp/migrate/_reembed.py @@ -89,35 +89,18 @@ def _stage_record_to_table( rec: MemoryRecord, new_embedding: list[float], ) -> None: - if not rec.structure_hv: - from iai_mcp.tem import bind_structure - rec.structure_hv = bind_structure(rec) - new_rec = MemoryRecord( - id=rec.id, - tier=rec.tier, - literal_surface=rec.literal_surface, - aaak_index=rec.aaak_index, - embedding=new_embedding, - structure_hv=rec.structure_hv, - community_id=rec.community_id, - centrality=rec.centrality, - detail_level=rec.detail_level, - pinned=rec.pinned, - stability=rec.stability, - difficulty=rec.difficulty, - last_reviewed=rec.last_reviewed, - never_decay=rec.never_decay, - never_merge=rec.never_merge, - provenance=rec.provenance, - created_at=rec.created_at, - updated_at=rec.updated_at, - tags=rec.tags, - language=rec.language, - s5_trust_score=rec.s5_trust_score, - profile_modulation_gain=rec.profile_modulation_gain, - schema_version=rec.schema_version, - ) - target_tbl.add([store._to_row(new_rec)]) + # Copy the authoritative SQL row byte-for-byte and replace only the vector. + # Rebuilding from MemoryRecord would drop storage-only fields such as + # vec_label, tombstones and pending state, and would decrypt/re-encrypt text. + with store.db.ro_conn() as conn: + source_row = conn.execute( + "SELECT * FROM records WHERE id = ?", (str(rec.id),) + ).fetchone() + if source_row is None: + raise ValueError(f"source record disappeared during re-embedding: {rec.id}") + staged_row = dict(source_row) + staged_row["embedding"] = new_embedding + target_tbl.add([staged_row]) def _stage_loop( @@ -132,56 +115,99 @@ def _stage_loop( started_idx: int = 0, already_staged_ids: Optional[set[str]] = None, progress: Optional[Callable[[int, int], None]] = None, + batch_size: int = 256, ) -> tuple[int, list[str]]: staged_count = 0 failures: list[str] = [] - staged_ids: list[str] = list(already_staged_ids or []) - skipped_set: set[str] = set(staged_ids) + skipped_set: set[str] = set(already_staged_ids or []) idx = started_idx - for rec in source_iter: - rec_id_str = str(rec.id) - if rec_id_str in skipped_set: - continue - if progress is not None: - try: - progress(idx, total) - except (TypeError, ValueError): - pass + pending_batch: list[MemoryRecord] = [] + + def flush_batch(batch: list[MemoryRecord]) -> None: + nonlocal idx, staged_count + if not batch: + return try: - new_embedding = target_embedder.embed(rec.literal_surface) - _stage_record_to_table(store, target_tbl, rec, new_embedding) + batch_method = getattr(target_embedder, "embed_batch", None) + if callable(batch_method): + vectors = batch_method([rec.literal_surface for rec in batch]) + else: + vectors = [target_embedder.embed(rec.literal_surface) for rec in batch] + if len(vectors) != len(batch): + raise ValueError( + f"embed_batch returned {len(vectors)} vectors for {len(batch)} records" + ) except (KeyboardInterrupt, SystemExit): raise except (OSError, ValueError, RuntimeError) as exc: + batch_ids = [str(rec.id) for rec in batch] log.warning( - "migrate_reembed_per_row_failed", - extra={ - "record_id": rec_id_str, - "error": str(exc)[:160], - }, + "migrate_reembed_batch_failed", + extra={"record_ids": batch_ids, "error": str(exc)[:160]}, ) - failures.append(rec_id_str) + failures.extend(batch_ids) + idx += len(batch) + return + + last_staged_id: str | None = None + for rec, new_embedding in zip(batch, vectors, strict=True): + rec_id_str = str(rec.id) + if progress is not None: + try: + progress(idx, total) + except (TypeError, ValueError): + pass + try: + _stage_record_to_table(store, target_tbl, rec, new_embedding) + except (KeyboardInterrupt, SystemExit): + raise + except (OSError, ValueError, RuntimeError) as exc: + log.warning( + "migrate_reembed_per_row_failed", + extra={"record_id": rec_id_str, "error": str(exc)[:160]}, + ) + failures.append(rec_id_str) + idx += 1 + continue + + staged_count += 1 + last_staged_id = rec_id_str idx += 1 - continue - staged_count += 1 - staged_ids.append(rec_id_str) - _progress_write( - store, - { - "started_at": started_at_iso, - "ts": int(time.time()), - "row_index": idx, - "last_rid": rec_id_str, - "total": total, - "target_dim": target_dim, - "target_model_key": getattr(target_embedder, "model_key", "unknown"), - "staged_ids": staged_ids, - "failures": failures, - }, - ) - idx += 1 + # One fsynced checkpoint per provider batch keeps the migration + # crash-safe without rewriting a growing JSON list for every row. + # Resume reconstructs the authoritative staged-id set from the table. + if last_staged_id is not None: + _progress_write( + store, + { + "started_at": started_at_iso, + "ts": int(time.time()), + "row_index": idx - 1, + "last_rid": last_staged_id, + "total": total, + "target_dim": target_dim, + "target_model_key": getattr( + target_embedder, "model_key", "unknown" + ), + "failures": failures, + }, + ) + + effective_batch_size = ( + max(1, int(batch_size)) + if bool(getattr(target_embedder, "supports_batch", False)) + else 1 + ) + for rec in source_iter: + if str(rec.id) in skipped_set: + continue + pending_batch.append(rec) + if len(pending_batch) >= effective_batch_size: + flush_batch(pending_batch) + pending_batch = [] + flush_batch(pending_batch) return staged_count, failures @@ -211,7 +237,7 @@ def _validate_and_swap( ) -> dict: orig = store.db.open_table(RECORDS_TABLE).count_rows() staged = store.db.open_table(STAGING_TABLE).count_rows() - if orig > 0 and staged < orig * 0.99: + if staged != orig or failures: log.error( "migrate_reembed_validate_failed", extra={ @@ -222,8 +248,8 @@ def _validate_and_swap( }, ) raise RuntimeError( - f"reembed staging produced {staged}/{orig} rows " - f"({staged/max(orig,1):.3%}); refusing to swap. Inspect tables " + f"reembed staging produced {staged}/{orig} rows with " + f"{len(failures)} failures; refusing to swap. Inspect tables " f"manually or run `iai-mcp migrate --rollback`." ) @@ -244,11 +270,29 @@ def _validate_and_swap( except (OSError, ValueError, RuntimeError) as exc: log.error("migration_reembed event write failed: %s", exc) + from iai_mcp.hippo import HippoDB, _txn + + if not isinstance(store.db, HippoDB): + raise RuntimeError( + f"re-embed table swap requires a Hippo store, got {type(store.db).__name__}" + ) ts = int(time.time()) old_name = f"{OLD_TABLE_PREFIX}{ts}" - _swap_tables_filesystem(store.db, source=RECORDS_TABLE, dest=old_name) - _swap_tables_filesystem(store.db, source=STAGING_TABLE, dest=RECORDS_TABLE) - + # The two renames and dimension metadata are one storage transition. A + # crash can therefore expose either the complete old layout or the + # complete new layout, never a 1024d table advertised as 384d. + with store.db._conn_lock: + with _txn(store.db._conn): + store.db._conn.execute( + f"ALTER TABLE [{RECORDS_TABLE}] RENAME TO [{old_name}]" + ) + store.db._conn.execute( + f"ALTER TABLE [{STAGING_TABLE}] RENAME TO [{RECORDS_TABLE}]" + ) + store.db._conn.execute( + "UPDATE [_hippo_meta] SET value = ? WHERE key = 'embed_dim'", + (str(target_dim),), + ) store._embed_dim = target_dim _progress_clear(store) @@ -261,6 +305,7 @@ def _validate_and_swap( "failures": len(failures), "duration_sec": duration_sec, "old_table": old_name, + "restart_required": True, } @@ -269,14 +314,24 @@ def migrate_reembed_to_current_dim( target_embedder, dry_run: bool = False, progress: Optional[Callable[[int, int], None]] = None, + *, + force: bool = False, + batch_size: int = 256, ) -> dict: t0 = time.time() + # A caller may migrate in the same process that just captured records. + # Make the buffered writes authoritative before counting or staging them; + # otherwise a small corpus can be swapped as an apparently empty table. + from iai_mcp.store import flush_record_buffer + + flush_record_buffer(store) + source_dim = int(store.embed_dim) target_dim = int(target_embedder.DIM) started_at_iso = datetime.now(timezone.utc).isoformat() - if source_dim == target_dim: + if source_dim == target_dim and not force: try: write_event( store, @@ -329,6 +384,7 @@ def migrate_reembed_to_current_dim( total=total, started_at_iso=started_at_iso, progress=progress, + batch_size=batch_size, ) duration_sec = time.time() - t0 @@ -504,7 +560,11 @@ def _resume(db, store: MemoryStore, target_embedder) -> int: already_staged: set[str] = set() else: target_tbl = db.open_table(STAGING_TABLE) - already_staged = set(progress_state.get("staged_ids") or []) + with db.ro_conn() as conn: + already_staged = { + str(row[0]) + for row in conn.execute(f"SELECT id FROM [{STAGING_TABLE}]").fetchall() + } source_dim = int(store.embed_dim) started_at_iso = progress_state.get( diff --git a/src/iai_mcp/pipeline.py b/src/iai_mcp/pipeline.py index 561e8c6f..e3e7b991 100644 --- a/src/iai_mcp/pipeline.py +++ b/src/iai_mcp/pipeline.py @@ -14,7 +14,7 @@ import numpy as np from iai_mcp.community import CommunityAssignment -from iai_mcp.embed import Embedder +from iai_mcp.embed import Embedder, embed_query from iai_mcp.events import TELEMETRY_EMBED_NATIVE_FAILURE, write_event from iai_mcp.exceptions import ( NativeError, @@ -739,7 +739,7 @@ def _recall_core( ) try: - cue_emb = embedder.embed(cue) + cue_emb = embed_query(embedder, cue) except Exception as exc: write_event( store, diff --git a/src/iai_mcp/retrieve.py b/src/iai_mcp/retrieve.py index b0213e5c..f643aaf0 100644 --- a/src/iai_mcp/retrieve.py +++ b/src/iai_mcp/retrieve.py @@ -154,10 +154,10 @@ def recall( # vector (degraded, never a crash into recall). if cue_text and (not cue_embedding or not any(cue_embedding)): try: - from iai_mcp.embed import embedder_for_store + from iai_mcp.embed import embed_query, embedder_for_store # Full cue, same as the primary path — the encoder truncates at # its own token limit; a char slice here would rank differently. - cue_embedding = list(embedder_for_store(store).embed(cue_text)) + cue_embedding = list(embed_query(embedder_for_store(store), cue_text)) except Exception as exc: # noqa: BLE001 -- degraded beats dead log.warning("recall cue re-embed failed, using caller vector: %s", exc) diff --git a/src/iai_mcp/semantic_recall.py b/src/iai_mcp/semantic_recall.py index 5e11c652..11f54575 100644 --- a/src/iai_mcp/semantic_recall.py +++ b/src/iai_mcp/semantic_recall.py @@ -141,10 +141,12 @@ def _recall_daemon_down_post_warm( degraded_semantic_recall as _degrade, EMBED_DIM, ) + from iai_mcp.embed import embed_query try: - _test_vec = embedder.embed(cue) - if not isinstance(_test_vec, (list, tuple)) or len(_test_vec) != EMBED_DIM: + _test_vec = embed_query(embedder, cue) + expected_dim = int(getattr(embedder, "DIM", EMBED_DIM)) + if not isinstance(_test_vec, (list, tuple)) or len(_test_vec) != expected_dim: raise ValueError(f"embed returned unexpected dim {len(_test_vec) if hasattr(_test_vec, '__len__') else '?'}") except Exception as exc: # noqa: BLE001 logger.debug("daemon_down_local_embed_failed: %s", exc) diff --git a/tests/conftest.py b/tests/conftest.py index 36dd2e44..e4d312ee 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -131,6 +131,11 @@ def _crypto_passphrase_env(monkeypatch: pytest.MonkeyPatch) -> None: "IAI_MCP_CRYPTO_PASSPHRASE", "LILLI_FSYNC_MODE", "IAI_MCP_EMBED_MODEL", + "IAI_MCP_EMBED_PROVIDER", + "IAI_MCP_EMBED_URL", + "IAI_MCP_EMBED_DIM", + "IAI_MCP_EMBED_TIMEOUT_SEC", + "IAI_MCP_EMBED_MODEL_ID", ) diff --git a/tests/test_embedder_http_provider.py b/tests/test_embedder_http_provider.py new file mode 100644 index 00000000..f7950895 --- /dev/null +++ b/tests/test_embedder_http_provider.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +import builtins +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from types import SimpleNamespace + +import pytest + + +@pytest.fixture +def embed_server(): + requests: list[dict] = [] + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *_args) -> None: + pass + + def do_POST(self) -> None: + size = int(self.headers["Content-Length"]) + payload = json.loads(self.rfile.read(size)) + requests.append(payload) + vectors = [[float(len(text)), 1.0, 0.0] for text in payload["texts"]] + body = json.dumps( + { + "model": "test-multilingual", + "dimensions": 3, + "vectors": vectors, + } + ).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield server.server_port, requests + finally: + server.shutdown() + server.server_close() + thread.join(timeout=2) + + +def _configure(monkeypatch: pytest.MonkeyPatch, port: int) -> None: + monkeypatch.setenv("IAI_MCP_EMBED_PROVIDER", "http") + monkeypatch.setenv("IAI_MCP_EMBED_URL", f"http://127.0.0.1:{port}") + monkeypatch.setenv("IAI_MCP_EMBED_DIM", "3") + monkeypatch.setenv("IAI_MCP_EMBED_MODEL_ID", "test-multilingual") + + +def test_http_provider_replaces_native_and_preserves_input_type( + monkeypatch: pytest.MonkeyPatch, embed_server +) -> None: + from iai_mcp.embed import Embedder, embed_query + + port, requests = embed_server + _configure(monkeypatch, port) + + embedder = Embedder() + assert embedder._backend == "http" + assert embedder.model_key == "test-multilingual" + assert embedder.DIM == 3 + assert embedder.embed("memory") == [6.0, 1.0, 0.0] + assert embed_query(embedder, "question") == [8.0, 1.0, 0.0] + assert embedder.embed_batch(["a", "bb"]) == [ + [1.0, 1.0, 0.0], + [2.0, 1.0, 0.0], + ] + assert [request["input_type"] for request in requests] == [ + "document", + "query", + "document", + ] + + +def test_http_provider_preserves_empty_batch_without_request( + monkeypatch: pytest.MonkeyPatch, embed_server +) -> None: + from iai_mcp.embed import Embedder + + port, requests = embed_server + _configure(monkeypatch, port) + assert Embedder().embed_batch([]) == [] + assert requests == [] + + +def test_http_provider_caps_response_size(monkeypatch: pytest.MonkeyPatch) -> None: + from iai_mcp import embed as embed_module + + class OversizedResponse: + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + def read(self, size: int) -> bytes: + return b"x" * size + + _configure(monkeypatch, 4488) + monkeypatch.setattr(embed_module, "urlopen", lambda *_a, **_kw: OversizedResponse()) + + with pytest.raises(RuntimeError, match="too large"): + embed_module.Embedder().embed("memory") + + +def test_http_provider_does_not_import_native_backend( + monkeypatch: pytest.MonkeyPatch, embed_server +) -> None: + from iai_mcp.embed import Embedder + + port, _requests = embed_server + _configure(monkeypatch, port) + real_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + if name.startswith("iai_mcp_native"): + raise AssertionError("native backend must not load for the http provider") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", guarded_import) + assert Embedder()._backend == "http" + + +def test_http_provider_rejects_non_loopback_url( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from iai_mcp.embed import Embedder + + monkeypatch.setenv("IAI_MCP_EMBED_PROVIDER", "http") + monkeypatch.setenv("IAI_MCP_EMBED_URL", "https://example.com/embed") + monkeypatch.setenv("IAI_MCP_EMBED_DIM", "3") + with pytest.raises(ValueError, match="loopback"): + Embedder() + + +@pytest.mark.parametrize( + "url", + [ + "http://user:pass@127.0.0.1:4488/embed", + "http://127.0.0.1:4488/embed?model=other", + "http://127.0.0.1:4488/embed#fragment", + "http://127.0.0.1:4488/other", + ], +) +def test_http_provider_rejects_ambiguous_loopback_urls( + monkeypatch: pytest.MonkeyPatch, url: str +) -> None: + from iai_mcp.embed import Embedder + + monkeypatch.setenv("IAI_MCP_EMBED_PROVIDER", "http") + monkeypatch.setenv("IAI_MCP_EMBED_URL", url) + monkeypatch.setenv("IAI_MCP_EMBED_DIM", "3") + monkeypatch.setenv("IAI_MCP_EMBED_MODEL_ID", "test-multilingual") + with pytest.raises(ValueError): + Embedder() + + +def test_http_provider_requires_model_id( + monkeypatch: pytest.MonkeyPatch, embed_server +) -> None: + from iai_mcp.embed import Embedder + + port, _requests = embed_server + _configure(monkeypatch, port) + monkeypatch.delenv("IAI_MCP_EMBED_MODEL_ID") + with pytest.raises(ValueError, match="MODEL_ID is required"): + Embedder() + + +def test_http_provider_refuses_store_dimension_mismatch( + monkeypatch: pytest.MonkeyPatch, embed_server +) -> None: + from iai_mcp.embed import embedder_for_store + + port, _requests = embed_server + _configure(monkeypatch, port) + with pytest.raises(ValueError, match="re-embedding migration"): + embedder_for_store(SimpleNamespace(embed_dim=384)) + + +def test_http_provider_refuses_unexpected_model( + monkeypatch: pytest.MonkeyPatch, embed_server +) -> None: + from iai_mcp.embed import Embedder + + port, _requests = embed_server + _configure(monkeypatch, port) + monkeypatch.setenv("IAI_MCP_EMBED_MODEL_ID", "another-model") + with pytest.raises(ValueError, match="returned model"): + Embedder().embed("memory") + + +def test_legacy_test_double_still_embeds_queries() -> None: + from iai_mcp.embed import embed_query + + class LegacyEmbedder: + def embed(self, text: str) -> list[float]: + return [float(len(text))] + + assert embed_query(LegacyEmbedder(), "cue") == [3.0] + + +def test_foresight_anchor_is_embedded_as_a_query(monkeypatch) -> None: + from iai_mcp import foresight + + class RecordingEmbedder: + def __init__(self) -> None: + self.calls: list[tuple[str, str]] = [] + + def embed_query(self, text: str) -> list[float]: + self.calls.append((text, "query")) + return [1.0] + + store = SimpleNamespace( + _foresight_anchor=("now", "the next-turn retrieval cue", "session") + ) + embedder = RecordingEmbedder() + monkeypatch.setattr(foresight, "refresh_pack", lambda *_a, **_kw: None) + + assert foresight.refresh_from_anchor(store, embedder) + assert embedder.calls == [("the next-turn retrieval cue", "query")] diff --git a/tests/test_migrate_reembed_crash_safe.py b/tests/test_migrate_reembed_crash_safe.py index 23cb75a9..3fd92b83 100644 --- a/tests/test_migrate_reembed_crash_safe.py +++ b/tests/test_migrate_reembed_crash_safe.py @@ -107,7 +107,7 @@ def test_successful_migration_promotes_old_to_records(tmp_path, monkeypatch): f"exactly one records_old_ expected (deferred cleanup); got {old_tables}" ) - assert store.db.open_table("records").count_rows() >= 19 + assert store.db.open_table("records").count_rows() == 20 def test_mid_migration_kill_preserves_old_table(tmp_path, monkeypatch): diff --git a/tests/test_migrate_reembed_to_current_dim.py b/tests/test_migrate_reembed_to_current_dim.py index edea8334..237947d0 100644 --- a/tests/test_migrate_reembed_to_current_dim.py +++ b/tests/test_migrate_reembed_to_current_dim.py @@ -1,18 +1,22 @@ from __future__ import annotations +import json +import sqlite3 from datetime import datetime, timezone from uuid import UUID, uuid4 +import numpy as np +import pytest class _DimEmbedder: - def __init__(self, dim: int): self.DIM = dim self.model_key = f"fake-dim-{dim}" def embed(self, text: str) -> list[float]: import math + vec = [0.0] * self.DIM for i, ch in enumerate(text or ""): vec[i % self.DIM] += ord(ch) / 256.0 @@ -27,11 +31,13 @@ def _fresh_store(tmp_path, dim: int, monkeypatch): monkeypatch.setenv("IAI_MCP_STORE", str(tmp_path / "iai")) monkeypatch.setenv("IAI_MCP_EMBED_DIM", str(dim)) from iai_mcp.store import MemoryStore + return MemoryStore() def _seed_records(store, embedder, n: int = 3) -> list[UUID]: from iai_mcp.types import MemoryRecord + ids = [] now = datetime.now(timezone.utc) for i in range(n): @@ -53,7 +59,13 @@ def _seed_records(store, embedder, n: int = 3) -> list[UUID]: last_reviewed=now, never_decay=False, never_merge=False, - provenance=[{"ts": "2026-04-17T00:00:00+00:00", "cue": f"seed-{i}", "session_id": "seed"}], + provenance=[ + { + "ts": "2026-04-17T00:00:00+00:00", + "cue": f"seed-{i}", + "session_id": "seed", + } + ], created_at=now, updated_at=now, tags=["test", "migration"], @@ -67,7 +79,9 @@ def _seed_records(store, embedder, n: int = 3) -> list[UUID]: return ids -def test_reembed_upgrades_dim_and_preserves_all_non_embedding_fields(tmp_path, monkeypatch): +def test_reembed_upgrades_dim_and_preserves_all_non_embedding_fields( + tmp_path, monkeypatch +): src_embedder = _DimEmbedder(384) target_embedder = _DimEmbedder(1024) @@ -75,19 +89,47 @@ def test_reembed_upgrades_dim_and_preserves_all_non_embedding_fields(tmp_path, m assert store.embed_dim == 384 seeded_ids = _seed_records(store, src_embedder, n=3) pre = {rid: store.get(rid) for rid in seeded_ids} + with store.db.ro_conn() as conn: + stored_before = dict( + conn.execute( + "SELECT * FROM records WHERE id = ?", + (str(seeded_ids[0]),), + ).fetchone() + ) + encrypted_before = stored_before["literal_surface"] + from iai_mcp.crypto import is_encrypted + + assert is_encrypted(encrypted_before) from iai_mcp.migrate import migrate_reembed_to_current_dim + result = migrate_reembed_to_current_dim(store, target_embedder) assert result["target_dim"] == 1024 assert result["source_dim"] == 384 assert result["updated"] == 3 assert store.embed_dim == 1024 + with store.db.ro_conn() as conn: + stored_after = dict( + conn.execute( + "SELECT * FROM records WHERE id = ?", + (str(seeded_ids[0]),), + ).fetchone() + ) + encrypted_after = stored_after["literal_surface"] + assert encrypted_after == encrypted_before, ( + "migration must preserve encrypted storage bytes exactly" + ) + assert { + key: value for key, value in stored_after.items() if key != "embedding" + } == {key: value for key, value in stored_before.items() if key != "embedding"} for rid in seeded_ids: post = store.get(rid) assert post is not None - assert post.literal_surface == pre[rid].literal_surface, "literal_surface byte-identical" + assert post.literal_surface == pre[rid].literal_surface, ( + "literal_surface byte-identical" + ) assert post.tier == pre[rid].tier assert post.tags == pre[rid].tags assert post.language == pre[rid].language @@ -103,24 +145,229 @@ def test_reembed_upgrades_dim_and_preserves_all_non_embedding_fields(tmp_path, m assert post.embedding != pre[rid].embedding, "embedding must be re-computed" +def test_reembed_table_swap_and_dimension_metadata_are_atomic(tmp_path, monkeypatch): + store = _fresh_store(tmp_path, 384, monkeypatch) + _seed_records(store, _DimEmbedder(384), n=2) + with store.db._conn_lock: + store.db._conn.execute( + """ + CREATE TRIGGER reject_embed_dim_update + BEFORE UPDATE ON _hippo_meta + WHEN NEW.key = 'embed_dim' + BEGIN + SELECT RAISE(ABORT, 'simulated metadata failure'); + END + """ + ) + + from iai_mcp.migrate import migrate_reembed_to_current_dim + + with pytest.raises(sqlite3.IntegrityError, match="simulated metadata failure"): + migrate_reembed_to_current_dim(store, _DimEmbedder(1024)) + + assert store.db.open_table("records").count_rows() == 2 + assert store.db.open_table("records_v_new").count_rows() == 2 + with store.db.ro_conn() as conn: + stored_dim = conn.execute( + "SELECT value FROM _hippo_meta WHERE key = 'embed_dim'" + ).fetchone()[0] + assert stored_dim == "384" + assert store.embed_dim == 384 + + +def test_reopen_rebuilds_stale_hnsw_after_dimension_migration(tmp_path, monkeypatch): + source = _DimEmbedder(384) + target = _DimEmbedder(1024) + store = _fresh_store(tmp_path, 384, monkeypatch) + seeded = _seed_records(store, source, n=3) + + from iai_mcp.migrate import migrate_reembed_to_current_dim + + migrate_reembed_to_current_dim(store, target) + store.db.close() # persists the still-live pre-restart 384d HNSW buffer + + from iai_mcp.store import MemoryStore + + reopened = MemoryStore() + try: + hits = reopened.query_similar( + target.embed( + "Record #0 with literal surface content that must survive migration." + ), + k=3, + ) + assert seeded[0] in {record.id for record, _score in hits} + with reopened.db.ro_conn() as conn: + row = conn.execute( + "SELECT vec_label, embedding FROM records WHERE id = ?", + (str(seeded[0]),), + ).fetchone() + stored = np.frombuffer(row["embedding"], dtype=np.float32) + indexed = reopened.db._hnsw.get_items([int(row["vec_label"])])[0] + assert indexed.shape == (1024,) + assert np.allclose(indexed, stored, rtol=1e-4, atol=1e-5) + finally: + reopened.db.close() + + +def test_reembed_uses_real_http_provider_batch_path(tmp_path, monkeypatch): + store = _fresh_store(tmp_path, 384, monkeypatch) + _seed_records(store, _DimEmbedder(384), n=3) + + class Response: + def __init__(self, body: bytes): + self.body = body + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + def read(self, _size: int) -> bytes: + return self.body + + requests: list[dict] = [] + + def fake_urlopen(request, *, timeout): + assert timeout == 30.0 + payload = json.loads(request.data) + requests.append(payload) + vectors = [[1.0, 0.0, 0.0] for _ in payload["texts"]] + return Response( + json.dumps( + { + "model": "test-http-model", + "dimensions": 3, + "vectors": vectors, + } + ).encode() + ) + + monkeypatch.setenv("IAI_MCP_EMBED_PROVIDER", "http") + monkeypatch.setenv("IAI_MCP_EMBED_URL", "http://127.0.0.1:4488/embed") + monkeypatch.setenv("IAI_MCP_EMBED_DIM", "3") + monkeypatch.setenv("IAI_MCP_EMBED_MODEL_ID", "test-http-model") + monkeypatch.setattr("iai_mcp.embed.urlopen", fake_urlopen) + + from iai_mcp.embed import Embedder + from iai_mcp.migrate import migrate_reembed_to_current_dim + + result = migrate_reembed_to_current_dim(store, Embedder()) + + assert result["updated"] == 3 + assert store.embed_dim == 3 + assert len(requests) == 1 + assert requests[0]["input_type"] == "document" + assert len(requests[0]["texts"]) == 3 + + +def test_reembed_flushes_records_buffered_in_same_process(tmp_path, monkeypatch): + source = _DimEmbedder(384) + store = _fresh_store(tmp_path, 384, monkeypatch) + seeded_ids = _seed_records(store, source, n=3) + + from iai_mcp.migrate import migrate_reembed_to_current_dim + + result = migrate_reembed_to_current_dim(store, _DimEmbedder(1024)) + + assert result["updated"] == 3 + assert {record.id for record in store.all_records()} == set(seeded_ids) + + def test_reembed_idempotent_same_dim_no_op(tmp_path, monkeypatch): src = _DimEmbedder(384) store = _fresh_store(tmp_path, 384, monkeypatch) _seed_records(store, src, n=2) from iai_mcp.migrate import migrate_reembed_to_current_dim + result = migrate_reembed_to_current_dim(store, _DimEmbedder(384)) assert result["updated"] == 0 assert result["skipped"] == 2 or result.get("no_op") is True assert store.embed_dim == 384 +def test_reembed_can_force_model_change_at_same_dimension(tmp_path, monkeypatch): + source = _DimEmbedder(384) + store = _fresh_store(tmp_path, 384, monkeypatch) + seeded = _seed_records(store, source, n=2) + before = {rid: store.get(rid).embedding for rid in seeded} + + class Replacement(_DimEmbedder): + model_key = "test-model" + + def embed(self, text: str) -> list[float]: + vector = super().embed(text) + return list(reversed(vector)) + + from iai_mcp.migrate import migrate_reembed_to_current_dim + + result = migrate_reembed_to_current_dim(store, Replacement(384), force=True) + + assert result["updated"] == 2 + assert all(store.get(rid).embedding != before[rid] for rid in seeded) + + +def test_reembed_batches_provider_calls(tmp_path, monkeypatch): + source = _DimEmbedder(384) + store = _fresh_store(tmp_path, 384, monkeypatch) + _seed_records(store, source, n=10) + + class BatchSpy(_DimEmbedder): + def __init__(self, dim: int): + super().__init__(dim) + self.batch_sizes: list[int] = [] + self.supports_batch = True + + def embed_batch(self, texts: list[str]) -> list[list[float]]: + self.batch_sizes.append(len(texts)) + return super().embed_batch(texts) + + target = BatchSpy(1024) + from iai_mcp.migrate import migrate_reembed_to_current_dim + + result = migrate_reembed_to_current_dim(store, target, batch_size=4) + + assert result["updated"] == 10 + assert target.batch_sizes == [4, 4, 2] + + +def test_reembed_checkpoints_once_per_provider_batch(tmp_path, monkeypatch): + source = _DimEmbedder(384) + store = _fresh_store(tmp_path, 384, monkeypatch) + _seed_records(store, source, n=10) + + class BatchEmbedder(_DimEmbedder): + supports_batch = True + + from iai_mcp.migrate import _reembed + + writes: list[dict] = [] + real_write = _reembed._progress_write + + def spy_write(store, state): + writes.append(state) + real_write(store, state) + + monkeypatch.setattr(_reembed, "_progress_write", spy_write) + result = _reembed.migrate_reembed_to_current_dim( + store, BatchEmbedder(1024), batch_size=4 + ) + + assert result["updated"] == 10 + assert len(writes) == 3 + assert [state["row_index"] for state in writes] == [3, 7, 9] + assert all("staged_ids" not in state for state in writes) + + def test_reembed_dry_run_reports_without_mutating(tmp_path, monkeypatch): src = _DimEmbedder(384) store = _fresh_store(tmp_path, 384, monkeypatch) seeded = _seed_records(store, src, n=2) from iai_mcp.migrate import migrate_reembed_to_current_dim + result = migrate_reembed_to_current_dim(store, _DimEmbedder(1024), dry_run=True) assert result["would_update"] == 2 assert store.embed_dim == 384 @@ -128,13 +375,36 @@ def test_reembed_dry_run_reports_without_mutating(tmp_path, monkeypatch): assert len(post.embedding) == 384 +def test_reembed_refuses_swap_when_any_record_fails(tmp_path, monkeypatch): + source = _DimEmbedder(384) + store = _fresh_store(tmp_path, 384, monkeypatch) + _seed_records(store, source, n=10) + + class OneFailure(_DimEmbedder): + def embed(self, text: str) -> list[float]: + if "#4 " in text: + raise ValueError("simulated record failure") + return super().embed(text) + + from iai_mcp.migrate import migrate_reembed_to_current_dim + + with pytest.raises(RuntimeError, match="refusing to swap"): + migrate_reembed_to_current_dim(store, OneFailure(1024)) + + assert store.db.open_table("records").count_rows() == 10 + assert store.db.open_table("records_v_new").count_rows() == 9 + assert store.embed_dim == 384 + + def test_reembed_emits_migration_event(tmp_path, monkeypatch): from iai_mcp.events import query_events + src = _DimEmbedder(384) store = _fresh_store(tmp_path, 384, monkeypatch) _seed_records(store, src, n=1) from iai_mcp.migrate import migrate_reembed_to_current_dim + migrate_reembed_to_current_dim(store, _DimEmbedder(1024)) events = query_events(store, kind="migration_reembed", limit=5) diff --git a/tests/test_native_fail_loud.py b/tests/test_native_fail_loud.py index 67c1a3ce..4bb086e9 100644 --- a/tests/test_native_fail_loud.py +++ b/tests/test_native_fail_loud.py @@ -102,7 +102,6 @@ def test_recall_cue_encode_failure_emits_store_event_and_raises( def _broken_embed(self, text: str) -> list[float]: raise RuntimeError("boom native encode") - monkeypatch.setattr(Embedder, "_encode_one", lambda self, t: (_ for _ in ()).throw(RuntimeError("boom native encode"))) monkeypatch.setattr(Embedder, "embed", _broken_embed) monkeypatch.setattr(