voidaccess 1.3.0__py3-none-any.whl

vector/__init__.py

@@ -0,0 +1,39 @@
+"""
+vector — Phase 4 embedding storage and semantic search.
+
+Re-exports the public API from embedder, store, and search.
+"""
+
+from vector.embedder import embed_batch, embed_text, get_embedder
+from vector.search import (
+    cross_investigation_recall,
+    find_pages_similar_to,
+    find_related_pages,
+)
+from vector.store import (
+    bulk_check_cache,
+    get_cached_page,
+    get_collection,
+    get_collection_stats,
+    is_duplicate,
+    search_similar,
+    store_page,
+    upsert_page,
+)
+
+__all__ = [
+    "get_embedder",
+    "embed_text",
+    "embed_batch",
+    "get_collection",
+    "upsert_page",
+    "search_similar",
+    "is_duplicate",
+    "get_collection_stats",
+    "find_related_pages",
+    "find_pages_similar_to",
+    "cross_investigation_recall",
+    "get_cached_page",
+    "store_page",
+    "bulk_check_cache",
+]

vector/embedder.py

@@ -0,0 +1,100 @@
+"""
+Local embedding generation via sentence-transformers (all-MiniLM-L6-v2).
+
+Uses the shared model singleton from vector.model_singleton.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from vector.model_singleton import get_embedding_model
+
+logger = logging.getLogger(__name__)
+
+_MAX_TOKENS = 512
+
+
+def get_embedder():
+    """Return the shared SentenceTransformer instance, or None if unavailable."""
+    return get_embedding_model()
+
+
+def _truncate_to_model_limit(text: str, model) -> str:
+    if not text:
+        return text
+    try:
+        tok = getattr(model, "tokenizer", None)
+        if tok is None:
+            return text[:50000]
+        encoded = tok.encode(
+            text,
+            add_special_tokens=True,
+            truncation=True,
+            max_length=_MAX_TOKENS,
+        )
+        return tok.decode(encoded, skip_special_tokens=True)
+    except Exception:
+        return text[:50000]
+
+
+def embed_text(text: str) -> list[float] | None:
+    """
+    Return a 384-dim embedding as a plain Python list, or None if unavailable
+    or text is empty.
+    """
+    if not (text and str(text).strip()):
+        return None
+    model = get_embedder()
+    if model is None:
+        return None
+    try:
+        truncated = _truncate_to_model_limit(str(text), model)
+        vec = model.encode(
+            truncated,
+            convert_to_numpy=True,
+            show_progress_bar=False,
+        )
+        return [float(x) for x in vec.tolist()]
+    except Exception as exc:
+        logger.warning("embed_text failed: %s", exc)
+        return None
+
+
+def embed_batch(texts: list[str]) -> list[list[float] | None]:
+    """
+    Batch embedding. Returns a list parallel to *texts*; any failure becomes None.
+    """
+    if not texts:
+        return []
+    model = get_embedder()
+    if model is None:
+        return [None for _ in texts]
+    prepared: list[str] = []
+    empty_indices: set[int] = set()
+    for i, t in enumerate(texts):
+        if not (t and str(t).strip()):
+            empty_indices.add(i)
+            prepared.append("")
+        else:
+            prepared.append(_truncate_to_model_limit(str(t), model))
+    out: list[list[float] | None] = [None] * len(texts)
+    for i in empty_indices:
+        out[i] = None
+    to_encode_idx = [i for i in range(len(texts)) if i not in empty_indices]
+    if not to_encode_idx:
+        return out
+    try:
+        batch_in = [prepared[i] for i in to_encode_idx]
+        encoded = model.encode(
+            batch_in,
+            convert_to_numpy=True,
+            show_progress_bar=False,
+        )
+        for j, row_idx in enumerate(to_encode_idx):
+            out[row_idx] = [float(x) for x in encoded[j].tolist()]
+    except Exception as exc:
+        logger.warning("embed_batch failed: %s", exc)
+        for i in to_encode_idx:
+            out[i] = None
+    return out  # type: ignore[return-value]

vector/model_singleton.py

@@ -0,0 +1,49 @@
+"""
+Thread-safe singleton for the SentenceTransformer embedding model.
+
+Loads all-MiniLM-L6-v2 lazily on first use, then reuses the same instance
+across all consumers. This eliminates the ~80 MB duplicate model weight
+problem when multiple modules each instantiate their own model at load time.
+
+Import torch and SentenceTransformer INSIDE the getter function to avoid
+the 2-5 second startup delay from torch enumerating CUDA devices.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+
+logger = logging.getLogger(__name__)
+
+_model: "SentenceTransformer | None" = None
+_lock = threading.Lock()
+
+
+def get_embedding_model() -> "SentenceTransformer | None":
+    """
+    Return the shared SentenceTransformer instance.
+
+    Lazy-loads the model on first call, then caches it for all subsequent calls.
+    Thread-safe: uses a lock to prevent race conditions during init.
+    """
+    global _model
+
+    if _model is not None:
+        return _model
+
+    with _lock:
+        if _model is not None:
+            return _model
+
+        import torch  # noqa: PLC0415 - imported inside function per M-6
+        from sentence_transformers import SentenceTransformer  # noqa: PLC0415
+
+        try:
+            _model = SentenceTransformer("all-MiniLM-L6-v2")
+            logger.info("Loaded embedding model all-MiniLM-L6-v2 (singleton)")
+        except Exception as exc:
+            logger.warning("Failed to load embedding model: %s", exc)
+            _model = None
+
+    return _model

vector/search.py

@@ -0,0 +1,87 @@
+"""
+Higher-level semantic search built on vector/store.py.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from . import store
+
+logger = logging.getLogger(__name__)
+
+
+def find_related_pages(query: str, n_results: int = 10) -> list[dict]:
+    """Semantic search over stored pages (metadata + distance)."""
+    return store.search_similar(query, n_results=n_results)
+
+
+def find_pages_similar_to(reference_url: str, n_results: int = 10) -> list[dict]:
+    """
+    Find pages similar to *reference_url* using its stored embedding.
+    Returns [] if the URL is not in the collection.
+    """
+    col = store.get_collection()
+    if col is None:
+        return []
+    try:
+        import hashlib
+
+        pid = hashlib.sha256(reference_url.encode("utf-8")).hexdigest()
+        got = col.get(ids=[pid], include=["embeddings"])
+        embs = got.get("embeddings") or []
+        if not embs or embs[0] is None:
+            return []
+        emb = list(embs[0])
+        n = max(1, int(n_results))
+        res = col.query(
+            query_embeddings=[emb],
+            n_results=n + 1,
+            include=["distances", "metadatas"],
+        )
+        ids = (res.get("ids") or [[]])[0]
+        dists = (res.get("distances") or [[]])[0]
+        metas = (res.get("metadatas") or [[]])[0]
+        out: list[dict] = []
+        for i, _eid in enumerate(ids):
+            if _eid == pid:
+                continue
+            m = metas[i] if i < len(metas) and metas[i] else {}
+            md = dict(m) if isinstance(m, dict) else {}
+            raw_pid = md.get("page_id")
+            page_id_out: int | None = None
+            if raw_pid is not None:
+                try:
+                    page_id_out = int(raw_pid)
+                except (TypeError, ValueError):
+                    page_id_out = None
+            dist_f = float(dists[i]) if i < len(dists) else 0.0
+            out.append(
+                {
+                    "url": md.get("url", ""),
+                    "page_id": page_id_out,
+                    "distance": dist_f,
+                    "metadata": md,
+                }
+            )
+            if len(out) >= n:
+                break
+        return out[:n]
+    except Exception as exc:
+        logger.warning("find_pages_similar_to failed: %s", exc)
+        return []
+
+
+def cross_investigation_recall(
+    query: str,
+    exclude_investigation_id: int | None = None,
+) -> list[dict]:
+    """
+    Similar pages across investigations; optionally exclude one investigation_id.
+    """
+    where: dict[str, Any] | None = None
+    if exclude_investigation_id is not None:
+        ex = str(exclude_investigation_id)
+        where = {"investigation_id": {"$ne": ex}}
+    return store.search_similar(query, n_results=10, where=where)