embedflow 0.1.0__py3-none-any.whl

embedflow/indexes/qdrant_backend.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+import os
+import uuid
+from typing import Any
+
+import numpy as np
+
+from .base import SearchHit, VectorIndex, validate_k, validate_query_vector
+
+
+class QdrantIndex(VectorIndex):
+    """Qdrant adapter. The dependency is optional and imported lazily."""
+
+    def __init__(self, client: Any, collection: str, dimension: int,
+                 documents: dict[str, Any] | None = None, metric: str = "cosine",
+                 vector_name: str | None = None):
+        try:
+            dimension_value = int(dimension)
+            dimension_exact = float(dimension) == dimension_value
+        except (TypeError, ValueError, OverflowError) as exc:
+            raise ValueError("Qdrant vector dimension must be a positive integer") from exc
+        if isinstance(dimension, bool) or not dimension_exact or dimension_value < 1:
+            raise ValueError("Qdrant vector dimension must be a positive integer")
+        metric = str(metric).lower()
+        if metric not in {"cosine", "dot", "inner_product"}:
+            raise ValueError("metric must be cosine, dot, or inner_product")
+        if not str(collection).strip():
+            raise ValueError("Qdrant collection must be non-empty")
+        self.client, self.collection, self.dimension = client, str(collection), dimension_value
+        self.documents, self.metric, self.vector_name = documents or {}, metric, vector_name
+        # Qdrant accepts integer IDs or UUIDs, while document stores commonly
+        # use IDs such as ``doc-17``.  Keep a deterministic mapping for rows we
+        # write so arbitrary application IDs round-trip through search/retrieve.
+        self._id_map: dict[str, str] = {}
+        self._closed = False
+
+    @staticmethod
+    def _storage_id(document_id: str) -> str:
+        return str(uuid.uuid5(uuid.NAMESPACE_URL, f"embedflow:{document_id}"))
+
+    @classmethod
+    def connect(cls, path_or_url: str, collection: str, dimension: int,
+                documents: dict[str, Any] | None = None, metric: str = "cosine",
+                api_key_env: str | None = "QDRANT_API_KEY", vector_name: str | None = None) -> QdrantIndex:
+        try:
+            dimension_value = int(dimension)
+            dimension_exact = float(dimension) == dimension_value
+        except (TypeError, ValueError, OverflowError) as exc:
+            raise ValueError("Qdrant vector dimension must be a non-negative integer") from exc
+        if isinstance(dimension, bool) or not dimension_exact or dimension_value < 0:
+            raise ValueError("Qdrant vector dimension must be a non-negative integer")
+        try:
+            from qdrant_client import QdrantClient
+        except ImportError as exc:
+            raise RuntimeError("Qdrant backend requires qdrant-client; install it with `pip install qdrant-client`") from exc
+        path_or_url = str(path_or_url)
+        if "://" not in path_or_url and not path_or_url.startswith("http"):
+            client = QdrantClient(path=path_or_url)
+        else:
+            api_key = os.environ.get(api_key_env) if api_key_env else None
+            client = QdrantClient(url=path_or_url, api_key=api_key)
+        if dimension_value == 0:
+            try:
+                info = client.get_collection(collection)
+                vectors_config = getattr(getattr(info, "config", None), "params", None)
+                vectors_config = getattr(vectors_config, "vectors", None)
+                if isinstance(vectors_config, dict):
+                    vectors_config = vectors_config.get(vector_name) if vector_name else next(iter(vectors_config.values()), None)
+                dimension = int(getattr(vectors_config, "size", 0) or 0)
+                if dimension <= 0:
+                    raise ValueError("could not infer Qdrant collection vector dimension; set source.dimension")
+            except Exception:
+                close = getattr(client, "close", None)
+                if callable(close):
+                    close()
+                raise
+        return cls(client, collection, dimension, documents, metric, vector_name)
+
+    @classmethod
+    def build(cls, path_or_url: str, collection: str, vectors: np.ndarray, ids: list[str],
+              documents: dict[str, Any] | None = None, metric: str = "cosine",
+              api_key_env: str | None = "QDRANT_API_KEY", vector_name: str | None = None) -> QdrantIndex:
+        values = np.asarray(vectors, dtype="float32")
+        metric = str(metric).lower()
+        if metric not in {"cosine", "dot", "inner_product"}:
+            raise ValueError("metric must be cosine, dot, or inner_product")
+        if values.ndim != 2 or values.shape[0] == 0 or values.shape[1] < 1 or values.shape[0] != len(ids):
+            raise ValueError("vectors must be a non-empty 2-D matrix aligned with IDs")
+        if len(set(map(str, ids))) != len(ids):
+            raise ValueError("duplicate document IDs are not allowed")
+        if not np.isfinite(values).all():
+            raise ValueError("vectors must be finite")
+        if metric == "cosine" and np.any(np.linalg.norm(values, axis=1) <= 1e-12):
+            raise ValueError("cosine vectors must be non-zero")
+        result = cls.connect(path_or_url, collection, int(values.shape[1]), documents, metric,
+                             api_key_env=api_key_env, vector_name=vector_name)
+        try:
+            from qdrant_client.models import Distance, VectorParams
+            distance = Distance.COSINE if metric == "cosine" else Distance.DOT
+            vectors_config = (VectorParams(size=result.dimension, distance=distance) if not vector_name
+                              else {vector_name: VectorParams(size=result.dimension, distance=distance)})
+            if hasattr(result.client, "collection_exists"):
+                if result.client.collection_exists(collection):
+                    result.client.delete_collection(collection)
+                result.client.create_collection(collection_name=collection, vectors_config=vectors_config)
+            else:  # qdrant-client versions before collection_exists
+                result.client.recreate_collection(collection_name=collection, vectors_config=vectors_config)
+            payloads = [{"text": documents[str(i)]} for i in ids] if documents else None
+            result.upsert(ids, values, payloads)
+        except ImportError as exc:
+            result.close()
+            raise RuntimeError("qdrant-client is required") from exc
+        except Exception:
+            result.close()
+            raise
+        return result
+
+    def search(self, query_vector: np.ndarray, k: int) -> list[SearchHit]:
+        if self._closed:
+            raise RuntimeError("Qdrant index is closed")
+        k = validate_k(k)
+        value = validate_query_vector(query_vector, self.dimension).tolist()
+        # qdrant-client renamed search -> query_points; support both APIs.
+        if hasattr(self.client, "query_points"):
+            kwargs = {"collection_name": self.collection, "query": value, "limit": k}
+            if self.vector_name:
+                kwargs["using"] = self.vector_name
+            result = self.client.query_points(**kwargs).points
+        else:
+            kwargs = {"collection_name": self.collection, "query_vector": value, "limit": k}
+            if self.vector_name:
+                try:
+                    result = self.client.search(**kwargs, using=self.vector_name)
+                except TypeError:
+                    result = self.client.search(collection_name=self.collection,
+                                                query_vector=(self.vector_name, value), limit=k)
+            else:
+                result = self.client.search(**kwargs)
+        hits = []
+        for rank, p in enumerate(result):
+            payload = p.payload or {}
+            document_id = str(payload.get("_embedflow_document_id", p.id))
+            self._id_map[document_id] = str(p.id)
+            hits.append(SearchHit(document_id, float(p.score), rank))
+        return hits
+
+    def fetch_documents(self, ids: list[str]) -> dict[str, Any]:
+        if self._closed:
+            raise RuntimeError("Qdrant index is closed")
+        if self.documents:
+            return {str(i): self.documents[str(i)] for i in ids if str(i) in self.documents}
+        if not ids:
+            return {}
+        storage_ids = [self._id_map.get(str(x), self._storage_id(str(x))) for x in ids]
+        points = self.client.retrieve(collection_name=self.collection, ids=storage_ids, with_payload=True)
+        output = {}
+        for p in points:
+            payload = p.payload or {}
+            document_id = str(payload.get("_embedflow_document_id", p.id))
+            output[document_id] = payload.get("text", payload)
+        return output
+
+    def size(self) -> int:
+        if self._closed:
+            raise RuntimeError("Qdrant index is closed")
+        info = self.client.get_collection(self.collection)
+        return int(getattr(info, "points_count", 0) or 0)
+
+    def metadata(self) -> dict[str, Any]:
+        if self._closed:
+            raise RuntimeError("Qdrant index is closed")
+        return {"backend": "qdrant", "collection": self.collection, "dimension": self.dimension,
+                "metric": self.metric, "vector_name": self.vector_name, "size": self.size()}
+
+    def close(self) -> None:
+        """Release a local file lock or remote client connection.
+
+        ``QdrantClient.close`` is available in current qdrant-client releases;
+        older clients simply have no close method, so teardown remains safe.
+        The method is idempotent to support engine and application shutdown
+        paths both calling it defensively.
+        """
+        if self._closed:
+            return
+        self._closed = True
+        close = getattr(self.client, "close", None)
+        if callable(close):
+            close()
+
+    def upsert(self, ids: list[str], vectors: np.ndarray, payloads: list[dict[str, Any]] | None = None) -> None:
+        try:
+            from qdrant_client.models import PointStruct
+        except ImportError as exc:
+            raise RuntimeError("qdrant-client is required") from exc
+        ids = [str(value) for value in ids]
+        values = np.asarray(vectors, dtype="float32")
+        if values.ndim != 2 or values.shape != (len(ids), self.dimension):
+            raise ValueError(f"Qdrant vectors must have shape ({len(ids)}, {self.dimension})")
+        if not ids or len(set(ids)) != len(ids):
+            raise ValueError("Qdrant upsert requires non-empty unique document IDs")
+        if not np.isfinite(values).all():
+            raise ValueError("Qdrant vectors must be finite")
+        if self.metric == "cosine" and np.any(np.linalg.norm(values, axis=1) <= 1e-12):
+            raise ValueError("Qdrant cosine vectors must be non-zero")
+        if payloads is not None and len(payloads) != len(ids):
+            raise ValueError("Qdrant payload count must match document IDs")
+        payloads = payloads or [{} for _ in ids]
+        storage_ids = []
+        normalized_payloads = []
+        for document_id, payload in zip(ids, payloads):
+            original_id = str(document_id)
+            storage_id = self._storage_id(original_id)
+            self._id_map[original_id] = storage_id
+            storage_ids.append(storage_id)
+            if payload is not None and not isinstance(payload, dict):
+                raise ValueError("Qdrant payloads must be objects")
+            item = dict(payload or {})
+            item.setdefault("_embedflow_document_id", original_id)
+            normalized_payloads.append(item)
+        self.client.upsert(collection_name=self.collection,
+                           points=[PointStruct(id=storage_id,
+                                               vector=({self.vector_name: v.tolist()} if self.vector_name else v.tolist()),
+                                               payload=payload)
+                                   for storage_id, v, payload in zip(storage_ids, values, normalized_payloads)])

embedflow/metrics/__init__.py

@@ -0,0 +1,3 @@
+from .latency import STAGES, aggregate_records, summarize
+
+__all__ = ["STAGES", "aggregate_records", "summarize"]

embedflow/metrics/latency.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Iterable
+from typing import Any
+
+import numpy as np
+
+STAGES = ("source_query_encode_ms", "source_ann_ms", "target_query_encode_ms", "cache_lookup_ms",
+          "synchronous_target_encode_ms", "target_score_ms", "topk_ms", "total_ms")
+
+
+def summarize(values: Iterable[float]) -> dict[str, float | int]:
+    arr = np.asarray(list(values), dtype="float64")
+    if arr.size == 0: return {"count": 0}
+    if not np.isfinite(arr).all() or (arr < 0).any(): raise ValueError("latency values must be finite and non-negative")
+    return {"count": int(arr.size), "mean_ms": float(np.mean(arr)), "p50_ms": float(np.quantile(arr, .50)),
+            "p90_ms": float(np.quantile(arr, .90)), "p95_ms": float(np.quantile(arr, .95)),
+            "p99_ms": float(np.quantile(arr, .99)), "min_ms": float(np.min(arr)), "max_ms": float(np.max(arr)),
+            "std_ms": float(np.std(arr, ddof=1)) if arr.size > 1 else 0.0,
+            # A zero-duration synthetic fixture has no meaningful finite QPS;
+            # return 0 rather than leaking infinity into JSON reports.
+            "qps": float(1000.0 / np.mean(arr)) if float(np.mean(arr)) > 0 else 0.0}
+
+
+def aggregate_records(records: list[dict[str, Any]], group_keys: tuple[str, ...] = ("mode", "K", "nprobe")) -> list[dict[str, Any]]:
+    grouped: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list)
+    malformed = 0
+    for row in records:
+        if not isinstance(row, dict):
+            malformed += 1
+            continue
+        try:
+            values = [float(row.get(stage, 0.0)) for stage in STAGES]
+            if not np.isfinite(values).all() or (np.asarray(values) < 0).any():
+                raise ValueError
+        except (TypeError, ValueError):
+            # Telemetry is append-only and may contain a partially written or
+            # hand-edited line. Skip that row while making the omission visible.
+            malformed += 1
+            continue
+        grouped[tuple(row.get(k) for k in group_keys)].append(row)
+    out = []
+    for key, rows in sorted(grouped.items(), key=lambda x: tuple(str(v) for v in x[0])):
+        result = dict(zip(group_keys, key))
+        for stage in STAGES: result.update({f"{stage}_{metric}": value for metric, value in summarize(float(r.get(stage, 0.0)) for r in rows).items()})
+        if malformed:
+            result["malformed_rows_skipped"] = malformed
+        out.append(result)
+    return out

embedflow/migration/__init__.py

@@ -0,0 +1,3 @@
+from .state import DocumentStore, MigrationState
+
+__all__ = ["DocumentStore", "MigrationState"]

embedflow/migration/compatibility.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import json
+import random
+from collections.abc import Iterable, Mapping
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+K_VALUES = (10, 20, 50, 100, 200, 500)
+
+
+def _research_features(records: list[tuple[str, float, int]], target_scores: dict[str, float], ks: tuple[int, ...]) -> dict[str, float]:
+    """Delegate the frozen feature implementation when this checkout has it."""
+    try:
+        if max(ks) < 500:
+            raise ImportError("the frozen helper requires the complete K<=500 feature set")
+        from src.probe_features import build_features
+        return build_features(records, target_scores, ks=ks)
+    except (ImportError, ModuleNotFoundError):
+        final_ids = [x[0] for x in _rerank(records, target_scores, max(ks))[:10]]
+        top50 = [x[0] for x in _rerank(records, target_scores, min(50, len(records)))[:10]]
+        top500 = [x[0] for x in _rerank(records, target_scores, min(500, len(records)))[:10]]
+        stability = len(set(top50) & set(final_ids)) / max(1, len(final_ids))
+        return {"probe_residual_tail_50_mean": 1 - stability, "stability_to_500_50_mean": stability,
+                "deepest_p90": float(max((x[2] for x in _rerank(records, target_scores, min(500, len(records)))[:10]), default=0)),
+                "late_tail_area": float(len(set(top500) - set(top50)) / 10), "last_shell_any_rate": float(bool(set(top500) - set(top50))),
+                "fraction_margin_nonpositive": 0.0}
+
+
+def _rerank(records: list[tuple[str, float, int]], scores: dict[str, float], k: int) -> list[tuple[str, float, int]]:
+    subset = records[:min(k, len(records))]
+    return sorted(subset, key=lambda x: (-float(scores[x[0]]), int(x[2])))
+
+
+def _mean_features(rows: list[dict[str, float]]) -> dict[str, float]:
+    if not rows:
+        raise ValueError("cannot aggregate empty feature rows")
+    keys = sorted({k for r in rows for k in r})
+    output: dict[str, float] = {}
+    for key in keys:
+        values = [float(r.get(key, 0.0)) for r in rows]
+        if not np.isfinite(values).all():
+            raise ValueError(f"feature {key!r} contains non-finite values")
+        output[key] = float(np.mean(values))
+    return output
+
+
+def recommend_k(per_query: list[dict[str, Any]], kmax: int) -> tuple[int, str]:
+    if isinstance(kmax, bool) or int(kmax) != kmax or int(kmax) < 10:
+        raise ValueError("kmax must be an integer >= 10")
+    if not per_query: return min(50, int(kmax)), "No probe rows; using the configured conservative default."
+    for k in K_VALUES:
+        if k > kmax: continue
+        try:
+            vals = [float(row.get(f"stability_{k}", row.get("stability_to_500_50_mean", 0.0))) for row in per_query]
+        except (TypeError, ValueError) as exc:
+            raise ValueError(f"probe stability at K={k} must be numeric") from exc
+        if not np.isfinite(vals).all() or any(value < 0 or value > 1 for value in vals):
+            raise ValueError(f"probe stability at K={k} must be finite and in [0, 1]")
+        if vals and float(np.mean(vals)) >= 0.90:
+            return k, f"Mean finite-pool top-10 stability at K={k} is {float(np.mean(vals)):.3f}."
+    return min(200, int(kmax)), "No tested K reached the 0.90 finite-pool stability heuristic; review before deployment."
+
+
+def run_probe(source_model: Any, target_model: Any, index: Any, documents: Any,
+              queries: Iterable[tuple[str, str]], kmax: int = 500, seed: int = 42,
+              limit: int | None = None,
+              target_document_vectors: Mapping[str, np.ndarray] | None = None) -> dict[str, Any]:
+    """Run a finite candidate compatibility probe with frozen T2-v1 decision logic."""
+    if isinstance(kmax, bool) or int(kmax) != kmax or int(kmax) < 10:
+        raise ValueError("probe kmax must be an integer >= 10")
+    if limit is not None and (isinstance(limit, bool) or int(limit) != limit or int(limit) < 1):
+        raise ValueError("probe limit must be a positive integer when provided")
+    rows = list(queries)
+    normalized_rows: list[tuple[str, str]] = []
+    seen_query_ids: set[str] = set()
+    for item in rows:
+        if not isinstance(item, (tuple, list)) or len(item) != 2:
+            raise ValueError("probe queries must be (query_id, text) pairs")
+        query_id, text = str(item[0]), item[1]
+        if not isinstance(text, str) or not text.strip():
+            raise ValueError(f"probe query {query_id!r} must contain non-empty text")
+        if query_id in seen_query_ids:
+            raise ValueError(f"duplicate probe query ID {query_id!r}")
+        seen_query_ids.add(query_id)
+        normalized_rows.append((query_id, text))
+    rows = normalized_rows
+    if not rows:
+        raise ValueError("probe queries must contain at least one query")
+    rng = random.Random(seed); rng.shuffle(rows)
+    rows = rows[:limit] if limit else rows
+    kmax = min(int(kmax), index.size())
+    if kmax < 10: raise ValueError("compatibility probe needs at least 10 index candidates")
+    try:
+        from src.t2_v1 import verify_t2_hash
+        root = Path(__file__).resolve().parents[2]
+        package_root = Path(__file__).resolve().parents[1]
+        if (root / "frozen").exists():
+            verify_t2_hash(root)
+        elif (package_root / "frozen").exists():
+            verify_t2_hash(package_root)
+    except (ImportError, FileNotFoundError):
+        pass
+    per_query = []
+    for query_id, text in rows:
+        sq = np.asarray(source_model.encode_queries([text])[0], dtype="float32")
+        if sq.ndim != 1 or sq.shape[0] != int(source_model.dimension) or not np.isfinite(sq).all():
+            raise ValueError("source query encoder returned an invalid vector")
+        candidates = index.search(sq, kmax)
+        ids = [x.document_id for x in candidates]
+        if not ids:
+            raise ValueError("source index returned no candidates")
+        if target_document_vectors is None:
+            docs = documents.get(ids)
+            tv = np.asarray(target_model.encode_documents([docs[x] for x in ids]), dtype="float32")
+        else:
+            missing_vectors = [document_id for document_id in ids if document_id not in target_document_vectors]
+            if missing_vectors:
+                raise ValueError(f"target probe vectors are missing {len(missing_vectors)} candidate IDs (e.g. {missing_vectors[:3]})")
+            tv = np.asarray([target_document_vectors[document_id] for document_id in ids], dtype="float32")
+        tq = np.asarray(target_model.encode_queries([text])[0], dtype="float32")
+        if tv.ndim != 2 or tv.shape != (len(ids), int(target_model.dimension)) or not np.isfinite(tv).all():
+            raise ValueError("target document dimension or finiteness mismatch in probe")
+        if tq.ndim != 1 or tq.shape[0] != int(target_model.dimension) or not np.isfinite(tq).all():
+            raise ValueError("target query encoder returned an invalid vector")
+        scores = {did: float(vec @ tq) for did, vec in zip(ids, tv)}
+        features = _research_features([(x.document_id, x.score, x.source_rank) for x in candidates], scores, tuple(k for k in K_VALUES if k <= kmax))
+        features["query_id"] = str(query_id)
+        features["target_scores"] = scores
+        # Preserve per-K finite-pool stability for K recommendation.
+        final = [x[0] for x in _rerank([(x.document_id, x.score, x.source_rank) for x in candidates], scores, kmax)[:10]]
+        for k in K_VALUES:
+            if k <= kmax:
+                got = [x[0] for x in _rerank([(x.document_id, x.score, x.source_rank) for x in candidates], scores, k)[:10]]
+                features[f"stability_{k}"] = len(set(got) & set(final)) / max(1, len(final))
+        per_query.append(features)
+    scalar_rows = [{k: v for k, v in row.items() if isinstance(v, (int, float, np.number))} for row in per_query]
+    aggregate = _mean_features(scalar_rows)
+    try:
+        from src.t2_v1 import decide
+        diagnostic = decide(aggregate)
+        implementation = "src.t2_v1.decide"
+    except (ImportError, KeyError) as exc:
+        raise RuntimeError("frozen T2-v1 implementation is unavailable; refusing to invent a rule") from exc
+    recommended, rationale = recommend_k(per_query, kmax)
+    return {"diagnostic": diagnostic, "recommended_k": recommended, "rationale": rationale,
+            "features": aggregate, "per_query": per_query, "queries": len(per_query),
+            "kmax": kmax, "seed": seed, "implementation": implementation,
+            "warning": "T2-v1 is an empirical finite-tail diagnostic, not a mathematical guarantee."}
+
+
+def save_probe(result: dict[str, Any], path: str | Path) -> None:
+    path = Path(path); path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(result, indent=2, ensure_ascii=False, default=float))