embedflow 0.1.0__py3-none-any.whl

embedflow/migration/facade.py

@@ -0,0 +1,312 @@
+"""Small public facade for starting a progressive model migration.
+
+The research implementation deliberately exposes the lower-level config and
+engine objects.  This module adds the short path a user needs in an
+application: point EmbedFlow at an existing FAISS/Qdrant index, name the old
+and new embedding models, and receive a serving session.  All retrieval,
+cache, worker, and probe behavior is delegated to the existing implementation;
+this is not a second query path.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any
+
+from ..cache import SQLiteVectorCache
+from ..config import (
+    CacheConfig,
+    DocumentsConfig,
+    EmbedFlowConfig,
+    IndexConfig,
+    MigrationConfig,
+    ModelConfig,
+    hydrate_research_contract,
+    save_config,
+)
+from ..indexes.base import VectorIndex
+from ..migration.compatibility import run_probe, save_probe
+from ..migration.state import DocumentStore
+from ..models import EmbeddingModel, load_embedding_model
+from ..runtime import load_index
+from ..serving.engine import MigrationEngine
+
+_MODEL_REGISTRY = {
+    "sentence-transformers/all-MiniLM-L6-v2": "minilm_l6",
+    "Qwen/Qwen3-Embedding-0.6B": "qwen3_0_6b",
+    "Qwen/Qwen3-Embedding-4B": "qwen3_4b",
+    "Qwen/Qwen3-Embedding-8B": "qwen3_8b",
+}
+
+
+def _is_model(value: Any) -> bool:
+    return all(hasattr(value, name) for name in ("encode_queries", "encode_documents", "dimension", "fingerprint"))
+
+
+def _model_config(value: str | EmbeddingModel, model_root: str | Path | None = None) -> ModelConfig:
+    """Build a contract for a model identifier or an already-loaded model."""
+    if _is_model(value):
+        # HuggingFaceEmbeddingModel retains its exact hydrated config.  Reuse
+        # it when available so fingerprints and prompt contracts remain exact.
+        retained = getattr(value, "cfg", None)
+        if isinstance(retained, ModelConfig):
+            return retained
+        return ModelConfig(model=str(value.model_id), dimension=int(value.dimension))
+    if not isinstance(value, (str, Path)):
+        raise TypeError("old_model/new_model must be a model ID/path or EmbeddingModel instance")
+    identifier = str(value)
+    cfg = hydrate_research_contract(ModelConfig(identifier), project_root=Path(__file__).resolve().parents[2])
+    candidate = Path(identifier)
+    if candidate.exists():
+        cfg.local_path = str(candidate.resolve())
+    elif model_root is not None and identifier in _MODEL_REGISTRY:
+        staged = Path(model_root).expanduser().resolve() / _MODEL_REGISTRY[identifier]
+        if staged.exists():
+            cfg.local_path = str(staged)
+    return cfg
+
+
+def _queries(value: str | Path | Iterable[tuple[str, str]]) -> list[tuple[str, str]]:
+    if isinstance(value, (str, Path)):
+        path = Path(value)
+        if not path.exists():
+            raise FileNotFoundError(path)
+        rows: list[tuple[str, str]] = []
+        with path.open() as handle:
+            seen: set[str] = set()
+            for i, line in enumerate(handle, 1):
+                if not line.strip():
+                    continue
+                try:
+                    row = json.loads(line)
+                except json.JSONDecodeError as exc:
+                    raise ValueError(f"invalid query JSON at {path}:{i}") from exc
+                if not isinstance(row, dict):
+                    raise ValueError(f"query row {i} in {path} must be a JSON object")
+                query_id = str(row.get("id", row.get("query_id", i)))
+                text = row.get("text", row.get("query"))
+                if query_id in seen:
+                    raise ValueError(f"duplicate query ID {query_id!r} in {path}")
+                if not isinstance(text, str) or not text.strip():
+                    raise ValueError(f"query {query_id!r} has no text")
+                seen.add(query_id)
+                rows.append((query_id, text))
+        if not rows:
+            raise ValueError(f"query file is empty: {value}")
+        return rows
+    try:
+        rows = [(str(query_id), text) for query_id, text in value]
+    except (TypeError, ValueError) as exc:
+        raise ValueError("probe_queries must be an iterable of (query_id, text) pairs") from exc
+    if len({query_id for query_id, _ in rows}) != len(rows):
+        raise ValueError("probe_queries must not contain duplicate query IDs")
+    if not rows or any(not isinstance(text, str) or not text.strip() for _, text in rows):
+        raise ValueError("probe_queries must contain non-empty query text")
+    return rows
+
+
+class MigrationSession:
+    """Handle returned by :func:`migrate`.
+
+    ``search`` and ``status`` are synchronous convenience methods for an
+    application.  ``serve`` mounts the same FastAPI dashboard/API used by the
+    CLI.  Use ``close`` or a context manager to release models and the cache.
+    """
+
+    def __init__(self, engine: MigrationEngine, config: EmbedFlowConfig,
+                 config_path: Path | None, owns_source: bool, owns_target: bool, owns_index: bool):
+        self.engine = engine
+        self.config = config
+        self.config_path = config_path
+        self._owns_source = owns_source
+        self._owns_target = owns_target
+        self._owns_index = owns_index
+        self._closed = False
+
+    @property
+    def plan(self):
+        return self.engine.plan
+
+    def search(self, query: str, top_k: int = 10, candidate_depth: int | None = None,
+               max_sync_misses: int | None = None) -> dict[str, Any]:
+        return self.engine.search(query, top_k, candidate_depth, max_sync_misses)
+
+    def status(self) -> dict[str, Any]:
+        return self.engine.status()
+
+    def prewarm(self, document_ids: list[str], asynchronous: bool = True) -> dict[str, Any]:
+        return self.engine.prewarm(document_ids, asynchronous=asynchronous)
+
+    def app(self):
+        from ..serving.api import create_app
+        return create_app(self.engine)
+
+    def serve(self, host: str = "127.0.0.1", port: int = 8000, log_level: str = "info") -> None:
+        """Run the API/dashboard until interrupted."""
+        try:
+            import uvicorn
+        except ImportError as exc:
+            raise RuntimeError("serving requires uvicorn; install `embedflow[dashboard]`") from exc
+        uvicorn.run(self.app(), host=host, port=int(port), log_level=log_level)
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self.engine.close(close_models=False, close_indexes=self._owns_index)
+        if self._owns_source:
+            self.engine.source_model.close()
+        if self._owns_target:
+            self.engine.target_model.close()
+        self._closed = True
+
+    def __enter__(self) -> MigrationSession:
+        return self
+
+    def __exit__(self, *_: Any) -> None:
+        self.close()
+
+
+def migrate(
+    *,
+    index: str | Path | VectorIndex,
+    old_model: str | Path | EmbeddingModel,
+    new_model: str | Path | EmbeddingModel,
+    documents: str | Path | DocumentStore,
+    backend: str | None = None,
+    index_url: str | None = None,
+    collection: str = "embedflow",
+    vector_name: str | None = None,
+    api_key_env: str | None = "QDRANT_API_KEY",
+    metric: str = "cosine",
+    model_root: str | Path | None = None,
+    device: str | None = None,
+    cache_path: str | Path = "./embedflow_cache",
+    state_path: str | Path = "./embedflow_state.json",
+    config_path: str | Path | None = None,
+    candidate_depth: int = 50,
+    kmax_probe: int = 500,
+    max_sync_misses: int = 4,
+    background_batch_size: int = 32,
+    probe_queries: str | Path | Iterable[tuple[str, str]] | None = None,
+    probe_limit: int | None = None,
+    start_worker: bool = True,
+) -> MigrationSession:
+    """Start progressive migration over an existing index.
+
+    ``index`` may be an existing EmbedFlow ``VectorIndex`` instance or a path
+    to a FAISS index (or a Qdrant path/URL when ``backend="qdrant"``).  Model
+    strings are revision-hydrated when they are registered by the research
+    contract; model objects can be supplied by an application directly.
+    ``probe_queries`` is optional so serving can start immediately.  When
+    supplied, the existing frozen T2-v1 implementation is run before the
+    session is returned.
+    """
+    if isinstance(documents, DocumentStore):
+        document_store = documents
+    else:
+        document_store = DocumentStore(str(documents))
+    source_cfg = _model_config(old_model, model_root)
+    target_cfg = _model_config(new_model, model_root)
+    selected_device = device or target_cfg.device or source_cfg.device or "cpu"
+
+    owns_source = not _is_model(old_model)
+    owns_target = not _is_model(new_model)
+    owns_index = not (isinstance(index, VectorIndex) or all(hasattr(index, name) for name in ("search", "size", "dimension", "metadata")))
+    source_model = None
+    target_model = None
+    source_index = None
+    cache = None
+    engine = None
+    try:
+        source_model = old_model if not owns_source else load_embedding_model(source_cfg, model_root=model_root, device=selected_device)
+        target_model = new_model if not owns_target else load_embedding_model(target_cfg, model_root=model_root, device=selected_device)
+        source_cfg.dimension = int(source_model.dimension)
+        target_cfg.dimension = int(target_model.dimension)
+
+        if not owns_index:
+            source_index = index
+            index_backend = backend or ("qdrant" if str(source_index.metadata().get("backend", "")).lower() == "qdrant" else "faiss")
+            index_path = str(getattr(source_index, "path", "./legacy.index") or "./legacy.index")
+        else:
+            index_path = str(index_url or index)
+            index_backend = (backend or ("qdrant" if "://" in index_path else "faiss")).lower()
+            if index_backend not in {"faiss", "qdrant"}:
+                raise ValueError("backend must be faiss or qdrant")
+
+        cfg = EmbedFlowConfig(
+            source=source_cfg,
+            target=target_cfg,
+            index=IndexConfig(backend=index_backend, path=index_path, collection=collection,
+                              url=index_url, metric=str(metric), vector_name=vector_name,
+                              api_key_env=api_key_env),
+            documents=DocumentsConfig(path=str(document_store.path)),
+            migration=MigrationConfig(candidate_depth=int(candidate_depth), kmax_probe=max(int(kmax_probe), int(candidate_depth)),
+                                      max_sync_misses=int(max_sync_misses), background_batch_size=int(background_batch_size)),
+            cache=CacheConfig(path=str(cache_path)),
+            state_path=str(state_path),
+            dashboard_title=f"EmbedFlow — {source_cfg.model} → {target_cfg.model}",
+        )
+        base = Path(config_path).expanduser().resolve().parent if config_path else Path.cwd()
+        cfg.resolve_paths(base)
+        cfg.validate()
+
+        if owns_index:
+            source_index = load_index(cfg, document_store)
+        if int(source_index.dimension) != int(source_model.dimension):
+            raise ValueError(f"source model dimension {source_model.dimension} != existing index dimension {source_index.dimension}")
+        stored_fingerprint = source_index.metadata().get("model_fingerprint")
+        if stored_fingerprint and stored_fingerprint != source_model.fingerprint:
+            raise ValueError("source model fingerprint does not match the existing index contract")
+
+        cache = SQLiteVectorCache(cfg.cache.path, target_model.fingerprint, target_model.dimension)
+        probe: dict[str, Any] = {}
+        if probe_queries is not None:
+            probe = run_probe(source_model, target_model, source_index, document_store, _queries(probe_queries),
+                              kmax=cfg.migration.kmax_probe, seed=42, limit=probe_limit)
+            if config_path:
+                save_probe(probe, Path(config_path).expanduser().resolve().with_name("probe_result.json"))
+        if config_path:
+            config_file = Path(config_path).expanduser().resolve()
+            save_config(cfg, config_file)
+        else:
+            config_file = None
+        engine = MigrationEngine(cfg, source_model, target_model, source_index, cache, document_store,
+                                 probe=probe, start_worker=start_worker)
+        return MigrationSession(engine, cfg, config_file, owns_source, owns_target, owns_index)
+    except Exception:
+        # A failed setup must not leak a loaded model, SQLite connection, or a
+        # local Qdrant file lock. Only close resources created by this call;
+        # caller-owned model/index objects remain their caller's responsibility.
+        if engine is not None:
+            try:
+                engine.close(close_models=False, close_indexes=owns_index)
+            except Exception:
+                pass
+        elif cache is not None:
+            try:
+                cache.close()
+            except Exception:
+                pass
+        if owns_index and source_index is not None:
+            try:
+                close_index = getattr(source_index, "close", None)
+                if callable(close_index):
+                    close_index()
+            except Exception:
+                pass
+        if owns_source and source_model is not None:
+            try:
+                source_model.close()
+            except Exception:
+                pass
+        if owns_target and target_model is not None:
+            try:
+                target_model.close()
+            except Exception:
+                pass
+        raise
+
+
+__all__ = ["MigrationSession", "migrate"]

embedflow/migration/materializer.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import sqlite3
+import threading
+import time
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from ..cache import SQLiteVectorCache
+
+
+class PersistentWorkQueue:
+    def __init__(self, path: str | Path, max_retries: int = 3):
+        p = Path(path); self.path = p if p.suffix else p / "materialization_queue.sqlite3"
+        try:
+            retries = int(max_retries)
+            exact = float(max_retries) == retries
+        except (TypeError, ValueError, OverflowError):
+            retries, exact = 0, False
+        if isinstance(max_retries, bool) or not exact or retries < 1:
+            raise ValueError("max_retries must be a positive integer")
+        self.path.parent.mkdir(parents=True, exist_ok=True); self.max_retries = retries
+        self.db = None; self.lock = threading.RLock()
+        try:
+            self.db = sqlite3.connect(str(self.path), check_same_thread=False, timeout=30)
+            self.db.execute("PRAGMA busy_timeout=30000")
+            self.db.execute("PRAGMA journal_mode=WAL")
+            self.db.execute("""CREATE TABLE IF NOT EXISTS work (
+                document_id TEXT PRIMARY KEY, status TEXT NOT NULL, attempts INTEGER NOT NULL DEFAULT 0,
+                error TEXT, enqueued_at REAL NOT NULL, updated_at REAL NOT NULL)""")
+            self.db.execute("UPDATE work SET status='pending',updated_at=? WHERE status='processing'", (time.time(),)); self.db.commit()
+        except sqlite3.DatabaseError as exc:
+            if self.db is not None:
+                self.db.close()
+            self.db = None
+            raise ValueError(f"invalid materialization queue database: {self.path}") from exc
+
+    def _ensure_open(self) -> None:
+        if self.db is None:
+            raise RuntimeError("materialization queue is closed")
+
+    def enqueue(self, ids: Iterable[str]) -> int:
+        now = time.time(); n = 0
+        with self.lock:
+            self._ensure_open()
+            for document_id in dict.fromkeys(str(x) for x in ids):
+                cur = self.db.execute("INSERT OR IGNORE INTO work(document_id,status,enqueued_at,updated_at) VALUES(?,?,?,?)",
+                                      (document_id, "pending", now, now)); n += cur.rowcount
+            self.db.commit()
+        return n
+
+    def claim(self, limit: int) -> list[str]:
+        if isinstance(limit, bool) or int(limit) != limit or int(limit) < 1:
+            raise ValueError("queue claim limit must be a positive integer")
+        with self.lock:
+            self._ensure_open()
+            # A write transaction makes SELECT+UPDATE atomic across multiple
+            # worker processes sharing the same SQLite queue.  Without it,
+            # two workers can claim the same pending document before either
+            # marks it processing.
+            self.db.execute("BEGIN IMMEDIATE")
+            try:
+                rows = self.db.execute("SELECT document_id FROM work WHERE status='pending' AND attempts<? ORDER BY enqueued_at,document_id LIMIT ?",
+                                       (self.max_retries, int(limit))).fetchall()
+                ids = [str(x[0]) for x in rows]; now = time.time()
+                self.db.executemany("UPDATE work SET status='processing',attempts=attempts+1,updated_at=? WHERE document_id=?",
+                                    [(now, x) for x in ids]); self.db.commit(); return ids
+            except Exception:
+                self.db.rollback()
+                raise
+
+    def complete(self, ids: Iterable[str]) -> None:
+        with self.lock:
+            self._ensure_open()
+            self.db.executemany("UPDATE work SET status='done',updated_at=? WHERE document_id=?", [(time.time(), str(x)) for x in ids]); self.db.commit()
+
+    def fail(self, ids: Iterable[str], error: str) -> None:
+        if not str(error).strip():
+            error = "unknown materialization failure"
+        with self.lock:
+            self._ensure_open()
+            self.db.executemany("UPDATE work SET status=CASE WHEN attempts<? THEN 'pending' ELSE 'error' END,error=?,updated_at=? WHERE document_id=?",
+                                [(self.max_retries, str(error), time.time(), str(x)) for x in ids]); self.db.commit()
+
+    def stats(self) -> dict[str, int]:
+        with self.lock:
+            self._ensure_open()
+            rows = self.db.execute("SELECT status,COUNT(*) FROM work GROUP BY status").fetchall()
+        out = {"pending": 0, "processing": 0, "done": 0, "error": 0}; out.update({str(k): int(v) for k, v in rows}); return out
+
+    def close(self) -> None:
+        with self.lock:
+            if self.db is not None:
+                self.db.close()
+                self.db = None
+
+
+class MaterializationWorker:
+    def __init__(self, target_model: Any, documents: Any, cache: SQLiteVectorCache,
+                 queue_path: str | Path, batch_size: int = 32, max_retries: int = 3,
+                 state: Any | None = None):
+        self.target_model, self.documents, self.cache = target_model, documents, cache
+        try:
+            parsed_batch_size = int(batch_size)
+            exact = float(batch_size) == parsed_batch_size
+        except (TypeError, ValueError, OverflowError):
+            parsed_batch_size, exact = 0, False
+        if isinstance(batch_size, bool) or not exact or parsed_batch_size < 1:
+            raise ValueError("batch_size must be a positive integer")
+        self.queue = PersistentWorkQueue(queue_path, max_retries); self.batch_size = parsed_batch_size; self.state = state
+        self.stop_event = threading.Event(); self.thread: threading.Thread | None = None
+        self._closed = False
+        self._stats = {"materialized": 0, "errors": 0, "started_at": None, "last_throughput_docs_sec": 0.0}
+
+    def enqueue(self, ids: Iterable[str]) -> int:
+        if self._closed:
+            raise RuntimeError("materialization worker is closed")
+        return self.queue.enqueue(ids)
+
+    def start(self) -> None:
+        if self._closed:
+            raise RuntimeError("materialization worker is closed")
+        if self.thread and self.thread.is_alive(): return
+        self.stop_event.clear(); self._stats["started_at"] = time.time(); self.thread = threading.Thread(target=self._run, name="embedflow-materializer", daemon=True); self.thread.start()
+
+    def _run(self) -> None:
+        while not self.stop_event.is_set():
+            try:
+                ids = self.queue.claim(self.batch_size)
+            except RuntimeError:
+                if self.stop_event.is_set() or self._closed:
+                    return
+                raise
+            except Exception as exc:
+                self._stats["errors"] += 1
+                if self.state:
+                    self.state.add_error(f"materialization queue failed: {exc}")
+                time.sleep(0.1)
+                continue
+            if not ids:
+                time.sleep(0.05)
+                continue
+            try:
+                # Resolve documents individually so one missing/deleted row is
+                # marked as an error without poisoning the rest of a batch.
+                available: list[str] = []
+                texts: list[str] = []
+                missing: list[str] = []
+                for document_id in ids:
+                    try:
+                        value = self.documents.get([document_id]).get(document_id)
+                    except Exception:
+                        value = None
+                    if value is None:
+                        missing.append(document_id)
+                    else:
+                        available.append(document_id); texts.append(str(value))
+                if missing:
+                    self.queue.fail(missing, "document text is unavailable")
+                    self._stats["errors"] += len(missing)
+                if not available:
+                    continue
+                vectors = self.target_model.encode_documents(texts, batch_size=self.batch_size)
+                self.cache.put(available, np.asarray(vectors, dtype="float32")); self.queue.complete(available)
+                self._stats["materialized"] += len(available)
+                elapsed = max(time.time() - float(self._stats["started_at"] or time.time()), 1e-6)
+                self._stats["last_throughput_docs_sec"] = self._stats["materialized"] / elapsed
+                if self.state: self.state.update(materialized_documents=self._stats["materialized"], materializer_throughput_docs_sec=self._stats["last_throughput_docs_sec"], queue=self.queue.stats())
+            except Exception as exc:
+                self.queue.fail(ids, repr(exc)); self._stats["errors"] += len(ids)
+                if self.state: self.state.add_error(f"materialization failed for {len(ids)} docs: {exc}")
+
+    def stop(self, timeout: float = 5.0) -> None:
+        self.stop_event.set()
+        if self.thread:
+            self.thread.join(timeout=timeout)
+            if self.thread.is_alive():
+                raise RuntimeError("materialization worker did not stop before timeout")
+
+    def stats(self) -> dict[str, Any]: return {**self._stats, "queue": self.queue.stats()}
+
+    def close(self, timeout: float = 5.0) -> None:
+        if self._closed:
+            return
+        self.stop(timeout=timeout)
+        self.queue.close()
+        self._closed = True

embedflow/migration/planner.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import Any
+
+
+@dataclass
+class MigrationPlan:
+    status: str
+    target_model: str
+    candidate_depth: int
+    diagnostic: str
+    ann_status: str
+    corpus_size: int
+    cached_target_vectors: int
+    cache_fraction: float
+    migration_strategy: str
+    warnings: list[str]
+    rationale: str = ""
+    source_model: str = ""
+
+    def to_dict(self) -> dict[str, Any]: return asdict(self)
+
+
+def make_plan(*, target_model: str, corpus_size: int, diagnostic: str,
+              recommended_k: int | None, ann_status: str, cached_target_vectors: int,
+              default_k: int = 50, throughput_docs_per_second: float | None = None,
+              gpu_price_per_hour: float | None = None, source_model: str = "") -> MigrationPlan:
+    diagnostic = str(diagnostic or "UNKNOWN").upper()
+    if diagnostic not in {"SAFE", "EXPAND", "UNSAFE_OR_UNCERTAIN", "UNKNOWN"}:
+        raise ValueError(f"unsupported diagnostic: {diagnostic}")
+    ann_value = str(ann_status or "UNKNOWN").upper()
+    if ann_value not in {"PASS", "WARNING", "UNKNOWN"}:
+        raise ValueError(f"unsupported ANN status: {ann_status}")
+    try:
+        corpus_int = int(corpus_size)
+        corpus_exact = float(corpus_size) == corpus_int
+        cached_int = int(cached_target_vectors)
+        cached_exact = float(cached_target_vectors) == cached_int
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise ValueError("corpus_size and cached_target_vectors must be integers") from exc
+    if isinstance(corpus_size, bool) or not corpus_exact or corpus_int < 0:
+        raise ValueError("corpus_size must be a non-negative integer")
+    if isinstance(cached_target_vectors, bool) or not cached_exact or cached_int < 0:
+        raise ValueError("cached_target_vectors must be a non-negative integer")
+    # ``candidate_depth: auto`` is a configuration convenience.  The actual
+    # serving plan always resolves to an integer, using the probe recommendation
+    # when available and the conservative default otherwise.
+    if str(default_k).lower() == "auto":
+        fallback_k = 50
+    else:
+        try:
+            fallback_k = int(default_k)
+            if isinstance(default_k, bool) or float(default_k) != fallback_k or fallback_k < 1:
+                raise ValueError
+        except (TypeError, ValueError, OverflowError) as exc:
+            raise ValueError("default_k must be auto or a positive integer") from exc
+    raw_k = fallback_k if recommended_k is None else recommended_k
+    try:
+        k = int(raw_k)
+        exact_k = float(raw_k) == k
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise ValueError("recommended_k must be a positive integer") from exc
+    if isinstance(raw_k, bool) or not exact_k or k < 1:
+        raise ValueError("recommended_k must be a positive integer")
+    warnings: list[str] = []
+    if diagnostic != "SAFE":
+        warnings.append("Compatibility is not a SAFE T2-v1 diagnostic; validate manually before relying on progressive migration.")
+    if ann_value in {"WARNING", "UNKNOWN"}:
+        warnings.append("ANN fidelity is not confirmed; T2-v1 does not diagnose ANN approximation error.")
+    if cached_int > corpus_int:
+        warnings.append("Cache count exceeds corpus size; check document identity and cache contract.")
+    fraction = (cached_int / corpus_int) if corpus_int else 0.0
+    strategy = "PROGRESSIVE" if diagnostic == "SAFE" else "PROGRESSIVE_WITH_REVIEW"
+    status = "READY" if diagnostic == "SAFE" and ann_value in {"PASS", "UNKNOWN"} else "REVIEW"
+    rationale = "Finite-pool diagnostic supports the selected starting K; this is an empirical recommendation, not a guarantee." if diagnostic == "SAFE" else "Use source-only/full-backfill fallback until compatibility and ANN health are reviewed."
+    return MigrationPlan(status, target_model, k, diagnostic, ann_value, corpus_int,
+                         cached_int, float(fraction), strategy, warnings, rationale, source_model)

embedflow/migration/state.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import json
+import threading
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any
+
+from ..config import EmbedFlowConfig
+
+
+class DocumentStore:
+    """Memory-indexed JSONL document store for progressive materialization."""
+
+    def __init__(self, path: str | Path, id_field: str = "id", text_field: str = "text"):
+        self.path = Path(path)
+        if not self.path.exists(): raise FileNotFoundError(self.path)
+        if not isinstance(id_field, str) or not id_field.strip() or not isinstance(text_field, str) or not text_field.strip():
+            raise ValueError("document id_field and text_field must be non-empty strings")
+        self.documents: dict[str, str] = {}
+        self.id_field, self.text_field = id_field, text_field
+        with self.path.open() as handle:
+            for line_no, line in enumerate(handle, 1):
+                if not line.strip(): continue
+                try: row = json.loads(line)
+                except json.JSONDecodeError as exc: raise ValueError(f"invalid JSON at {self.path}:{line_no}") from exc
+                if not isinstance(row, dict): raise ValueError(f"document row {line_no} in {self.path} must be a JSON object")
+                if id_field not in row or text_field not in row: raise ValueError(f"missing {id_field}/{text_field} at {self.path}:{line_no}")
+                document_id = str(row[id_field])
+                if not document_id.strip(): raise ValueError(f"empty document ID at {self.path}:{line_no}")
+                if document_id in self.documents: raise ValueError(f"duplicate document ID {document_id!r}")
+                text = row[text_field]
+                if text is None or not str(text).strip(): raise ValueError(f"missing document text for {document_id!r}")
+                self.documents[document_id] = str(text)
+        if not self.documents: raise ValueError(f"document file is empty: {self.path}")
+
+    def get(self, document_ids: Iterable[str]) -> dict[str, str]:
+        ids = [str(x) for x in document_ids]
+        missing = [x for x in ids if x not in self.documents]
+        if missing: raise KeyError(f"document text missing for IDs: {missing[:5]}")
+        return {x: self.documents[x] for x in ids}
+
+    def size(self) -> int: return len(self.documents)
+
+
+class MigrationState:
+    def __init__(self, path: str | Path, cfg: EmbedFlowConfig):
+        self.path = Path(path); self._lock = threading.RLock()
+        self.data: dict[str, Any] = {
+            "source_model": cfg.source.model, "target_model": cfg.target.model,
+            "target_model_fingerprint": cfg.target.fingerprint, "status": "INITIALIZING",
+            "diagnostic": "UNKNOWN", "ann_status": "UNKNOWN", "candidate_depth": cfg.migration.candidate_depth,
+            "migration_strategy": "PROGRESSIVE", "queries": 0, "errors": [],
+        }
+        if self.path.exists():
+            try:
+                loaded = json.loads(self.path.read_text())
+            except (OSError, json.JSONDecodeError) as exc: raise ValueError(f"invalid state file {self.path}") from exc
+            if not isinstance(loaded, dict):
+                raise ValueError(f"state file {self.path} must contain a JSON object")
+            stored_fingerprint = loaded.get("target_model_fingerprint")
+            if stored_fingerprint and str(stored_fingerprint) != str(cfg.target.fingerprint):
+                raise ValueError("state file target model fingerprint does not match the configured target contract")
+            self.data.update(loaded)
+        self.save()
+
+    def update(self, **values: Any) -> None:
+        with self._lock:
+            self.data.update(values); self.save()
+
+    def add_error(self, message: str) -> None:
+        with self._lock:
+            errors = self.data.setdefault("errors", []); errors.append(str(message)); self.data["errors"] = errors[-100:]; self.save()
+
+    def save(self) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = self.path.with_suffix(self.path.suffix + ".tmp")
+        tmp.write_text(json.dumps(self.data, indent=2, sort_keys=True, ensure_ascii=False)); tmp.replace(self.path)
+
+    def snapshot(self) -> dict[str, Any]:
+        with self._lock: return json.loads(json.dumps(self.data))

embedflow/models/__init__.py

@@ -0,0 +1,4 @@
+from .base import EmbeddingModel
+from .huggingface import HashEmbeddingModel, HuggingFaceEmbeddingModel, load_embedding_model
+
+__all__ = ["EmbeddingModel", "HashEmbeddingModel", "HuggingFaceEmbeddingModel", "load_embedding_model"]