pycontextdb 0.1.0__py3-none-any.whl

contextdb/store/vector_index.py

@@ -0,0 +1,241 @@
+"""Vector indices for fast nearest-neighbor search.
+
+Two implementations ship with ContextDB:
+
+* :class:`FAISSIndex` — uses ``faiss-cpu`` when installed; suitable for
+  collections up to tens of millions of vectors.
+* :class:`NumpyIndex` — pure-numpy brute force; used as a fallback and for
+  test determinism. Fine for <10K vectors.
+
+Callers should prefer :func:`get_vector_index`, which picks the best available
+implementation at runtime.
+"""
+
+from __future__ import annotations
+
+import pickle
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def _normalize(vectors: NDArray[np.float32]) -> NDArray[np.float32]:
+    """L2-normalize rows. Zero vectors pass through unchanged."""
+    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
+    norms = np.where(norms == 0, 1.0, norms)
+    return cast("NDArray[np.float32]", (vectors / norms).astype(np.float32))
+
+
+class VectorIndex(ABC):
+    """Minimal vector index surface used by :class:`SQLiteStore`."""
+
+    @abstractmethod
+    def add(self, ids: list[str], embeddings: NDArray[np.float32]) -> None: ...
+
+    @abstractmethod
+    def search(self, query: NDArray[np.float32], top_k: int = 10) -> list[tuple[str, float]]: ...
+
+    @abstractmethod
+    def remove(self, ids: list[str]) -> None: ...
+
+    @abstractmethod
+    def save(self, path: str) -> None: ...
+
+    @abstractmethod
+    def load(self, path: str) -> None: ...
+
+    @abstractmethod
+    def __len__(self) -> int: ...
+
+
+class NumpyIndex(VectorIndex):
+    """Brute-force cosine similarity over stacked float32 vectors."""
+
+    def __init__(self, dimension: int) -> None:
+        self.dimension = dimension
+        self._ids: list[str] = []
+        self._vectors: NDArray[np.float32] = np.zeros((0, dimension), dtype=np.float32)
+
+    def add(self, ids: list[str], embeddings: NDArray[np.float32]) -> None:
+        if not ids:
+            return
+        reshaped = np.asarray(embeddings, dtype=np.float32).reshape(len(ids), self.dimension)
+        normalized = _normalize(reshaped)
+        self._ids.extend(ids)
+        self._vectors = (
+            np.vstack([self._vectors, normalized]) if len(self._vectors) else normalized
+        )
+
+    def search(self, query: NDArray[np.float32], top_k: int = 10) -> list[tuple[str, float]]:
+        if len(self._ids) == 0:
+            return []
+        q = _normalize(np.asarray(query, dtype=np.float32).reshape(1, self.dimension))[0]
+        scores = self._vectors @ q
+        k = min(top_k, len(self._ids))
+        top_idx = np.argsort(-scores)[:k]
+        return [(self._ids[i], float(scores[i])) for i in top_idx]
+
+    def remove(self, ids: list[str]) -> None:
+        drop = set(ids)
+        keep = [i for i, mid in enumerate(self._ids) if mid not in drop]
+        self._ids = [self._ids[i] for i in keep]
+        self._vectors = self._vectors[keep] if keep else np.zeros(
+            (0, self.dimension), dtype=np.float32
+        )
+
+    def save(self, path: str) -> None:
+        payload = {
+            "dimension": self.dimension,
+            "ids": self._ids,
+            "vectors": self._vectors,
+        }
+        Path(path).write_bytes(pickle.dumps(payload))
+
+    def load(self, path: str) -> None:
+        payload = pickle.loads(Path(path).read_bytes())
+        self.dimension = int(payload["dimension"])
+        self._ids = list(payload["ids"])
+        self._vectors = np.asarray(payload["vectors"], dtype=np.float32)
+
+    def __len__(self) -> int:
+        return len(self._ids)
+
+
+class FAISSIndex(VectorIndex):
+    """FAISS-backed index; requires the ``faiss-cpu`` optional dependency.
+
+    Removal is O(1): an id goes into :attr:`_removed_ids` and is filtered out
+    of every :meth:`search` call. The underlying FAISS index is rebuilt when
+    the tombstone set exceeds :attr:`_rebuild_threshold` (10% by default) so
+    search over-fetch stays bounded and index memory does not grow without
+    limit.
+    """
+
+    _rebuild_threshold: float = 0.1
+
+    def __init__(self, dimension: int, index_type: str = "flat") -> None:
+        try:
+            import faiss
+        except ImportError as exc:  # pragma: no cover - exercised only without faiss
+            raise RuntimeError(
+                "FAISS is not installed. Install with `pip install pycontextdb[faiss]` "
+                "or use NumpyIndex."
+            ) from exc
+        self._faiss = faiss
+        self.dimension = dimension
+        self.index_type = index_type
+        self._index: Any = faiss.IndexFlatIP(dimension)
+        self._ids: list[str] = []
+        self._removed_ids: set[str] = set()
+
+    def add(self, ids: list[str], embeddings: NDArray[np.float32]) -> None:
+        if not ids:
+            return
+        reshaped = np.asarray(embeddings, dtype=np.float32).reshape(len(ids), self.dimension)
+        normalized = _normalize(reshaped)
+        self._index.add(normalized)
+        self._ids.extend(ids)
+        # Re-adding a previously tombstoned id clears the tombstone.
+        if self._removed_ids:
+            self._removed_ids.difference_update(ids)
+
+    def search(self, query: NDArray[np.float32], top_k: int = 10) -> list[tuple[str, float]]:
+        if not self._ids:
+            return []
+        # Over-fetch when tombstones are present so filtered-out hits do not
+        # starve the final top-k.
+        live = len(self._ids) - len(self._removed_ids)
+        if live <= 0:
+            return []
+        fetch = min(top_k + len(self._removed_ids), len(self._ids))
+        q = _normalize(np.asarray(query, dtype=np.float32).reshape(1, self.dimension))
+        scores, indices = self._index.search(q, fetch)
+        out: list[tuple[str, float]] = []
+        for score, idx in zip(scores[0], indices[0], strict=False):
+            if idx == -1:
+                continue
+            mid = self._ids[int(idx)]
+            if mid in self._removed_ids:
+                continue
+            out.append((mid, float(score)))
+            if len(out) >= top_k:
+                break
+        return out
+
+    def remove(self, ids: list[str]) -> None:
+        """Tombstone the given ids. Amortized O(1) per id.
+
+        When tombstones exceed :attr:`_rebuild_threshold` of the live set,
+        :meth:`rebuild` is triggered to reclaim memory.
+        """
+        if not ids:
+            return
+        present = {mid for mid in ids if mid in self._ids}
+        if not present:
+            return
+        self._removed_ids.update(present)
+        total = len(self._ids)
+        if total and len(self._removed_ids) / total > self._rebuild_threshold:
+            self.rebuild()
+
+    def rebuild(self) -> None:
+        """Reconstruct the underlying FAISS index without the tombstoned ids.
+
+        Expensive (O(n) over live vectors) but amortized across removals.
+        """
+        if not self._removed_ids:
+            return
+        keep_idx = [i for i, mid in enumerate(self._ids) if mid not in self._removed_ids]
+        new_ids = [self._ids[i] for i in keep_idx]
+        if keep_idx:
+            snapshot = self._vectors_snapshot()
+            vectors = np.stack([snapshot[i] for i in keep_idx], axis=0).astype(np.float32)
+        else:
+            vectors = np.zeros((0, self.dimension), dtype=np.float32)
+        self._index = self._faiss.IndexFlatIP(self.dimension)
+        self._ids = []
+        self._removed_ids.clear()
+        if len(new_ids):
+            self._index.add(vectors)
+            self._ids = new_ids
+
+    def _vectors_snapshot(self) -> NDArray[np.float32]:
+        # Reconstruct all current vectors from the FAISS index.
+        return np.stack(
+            [self._index.reconstruct(i) for i in range(self._index.ntotal)], axis=0
+        ).astype(np.float32)
+
+    def save(self, path: str) -> None:
+        # Persisted state always reflects the post-rebuild, tombstone-free view
+        # so stale ids never leak into a reloaded index.
+        if self._removed_ids:
+            self.rebuild()
+        self._faiss.write_index(self._index, f"{path}.faiss")
+        Path(f"{path}.ids").write_bytes(pickle.dumps(self._ids))
+
+    def load(self, path: str) -> None:
+        self._index = self._faiss.read_index(f"{path}.faiss")
+        self._ids = pickle.loads(Path(f"{path}.ids").read_bytes())
+        self._removed_ids = set()
+
+    def __len__(self) -> int:
+        return len(self._ids) - len(self._removed_ids)
+
+
+def get_vector_index(dimension: int, prefer: str = "auto") -> VectorIndex:
+    """Return the best available index for ``dimension``.
+
+    ``prefer`` may be ``"faiss"``, ``"numpy"``, or ``"auto"`` (the default,
+    which uses FAISS if installed and falls back to NumPy otherwise).
+    """
+    if prefer == "numpy":
+        return NumpyIndex(dimension)
+    if prefer == "faiss":
+        return FAISSIndex(dimension)
+    try:
+        return FAISSIndex(dimension)
+    except RuntimeError:
+        return NumpyIndex(dimension)

contextdb/utils/__init__.py

@@ -0,0 +1,22 @@
+"""Utility adapters for embeddings and LLMs."""
+
+from __future__ import annotations
+
+from contextdb.utils.embeddings import (
+    EmbeddingProvider,
+    MockEmbedding,
+    OpenAIEmbedding,
+    get_embedding_provider,
+)
+from contextdb.utils.llm import LLMProvider, MockLLM, OpenAILLM, get_llm_provider
+
+__all__ = [
+    "EmbeddingProvider",
+    "LLMProvider",
+    "MockEmbedding",
+    "MockLLM",
+    "OpenAIEmbedding",
+    "OpenAILLM",
+    "get_embedding_provider",
+    "get_llm_provider",
+]

contextdb/utils/embeddings.py

@@ -0,0 +1,159 @@
+"""Embedding providers.
+
+ContextDB speaks a minimal :class:`EmbeddingProvider` protocol so swapping
+backends (OpenAI ↔ local sentence-transformers ↔ deterministic mock) is a
+one-line change. :func:`get_embedding_provider` is the factory most callers
+want.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+from abc import ABC, abstractmethod
+from typing import Any
+
+import numpy as np
+
+from contextdb.core.exceptions import ConfigError
+
+# Known OpenAI embedding dimensions; update as new models ship.
+_OPENAI_DIMS: dict[str, int] = {
+    "text-embedding-3-small": 1536,
+    "text-embedding-3-large": 3072,
+    "text-embedding-ada-002": 1536,
+}
+
+
+class EmbeddingProvider(ABC):
+    """Async embedding contract. Implementations must be batch-safe."""
+
+    @abstractmethod
+    async def embed(self, texts: list[str]) -> list[list[float]]: ...
+
+    @abstractmethod
+    def dimension(self) -> int: ...
+
+
+class OpenAIEmbedding(EmbeddingProvider):
+    """OpenAI embedding API wrapper with exponential-backoff retry."""
+
+    def __init__(
+        self,
+        model: str = "text-embedding-3-small",
+        api_key: str | None = None,
+        max_retries: int = 3,
+    ) -> None:
+        from openai import AsyncOpenAI
+
+        self.model = model
+        self.max_retries = max_retries
+        self._client = AsyncOpenAI(api_key=api_key)
+        self._dim = _OPENAI_DIMS.get(model, 1536)
+
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            return []
+        # OpenAI's per-call limit is 2048 inputs; chunk defensively.
+        chunks = [texts[i : i + 2048] for i in range(0, len(texts), 2048)]
+        out: list[list[float]] = []
+        for chunk in chunks:
+            out.extend(await self._embed_with_retry(chunk))
+        return out
+
+    async def _embed_with_retry(self, texts: list[str]) -> list[list[float]]:
+        delay = 1.0
+        for attempt in range(self.max_retries):
+            try:
+                response = await self._client.embeddings.create(
+                    model=self.model, input=texts
+                )
+                return [d.embedding for d in response.data]
+            except Exception:  # noqa: BLE001
+                if attempt == self.max_retries - 1:
+                    raise
+                await asyncio.sleep(delay)
+                delay *= 2
+        return []  # pragma: no cover
+
+    def dimension(self) -> int:
+        return self._dim
+
+
+class SentenceTransformerEmbedding(EmbeddingProvider):
+    """Local-model embeddings via the optional ``sentence-transformers`` dep."""
+
+    def __init__(self, model_name: str = "all-MiniLM-L6-v2") -> None:
+        try:
+            from sentence_transformers import SentenceTransformer
+        except ImportError as exc:  # pragma: no cover - optional dep
+            raise RuntimeError(
+                "sentence-transformers is not installed. "
+                "Install with `pip install pycontextdb[local]`."
+            ) from exc
+        self._model = SentenceTransformer(model_name)
+        self._dim = int(self._model.get_sentence_embedding_dimension())
+
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            return []
+        loop = asyncio.get_running_loop()
+
+        def _run() -> list[list[float]]:
+            return [list(map(float, v)) for v in self._model.encode(texts)]
+
+        return await loop.run_in_executor(None, _run)
+
+    def dimension(self) -> int:
+        return self._dim
+
+
+class MockEmbedding(EmbeddingProvider):
+    """Deterministic pseudo-random embeddings for tests; no network calls."""
+
+    def __init__(self, dimension: int = 384) -> None:
+        self._dim = dimension
+
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        return [self._text_to_vector(t) for t in texts]
+
+    def _text_to_vector(self, text: str) -> list[float]:
+        digest = hashlib.sha256(text.encode("utf-8")).digest()
+        seed = int.from_bytes(digest[:4], "big")
+        rng = np.random.default_rng(seed)
+        vec = rng.standard_normal(self._dim).astype(np.float32)
+        # Encourage discrimination by mixing in per-word contribution.
+        for word in text.lower().split():
+            word_digest = hashlib.md5(word.encode("utf-8"), usedforsecurity=False).digest()
+            w_seed = int.from_bytes(word_digest[:4], "big")
+            word_rng = np.random.default_rng(w_seed)
+            vec += 0.5 * word_rng.standard_normal(self._dim).astype(np.float32)
+        norm = float(np.linalg.norm(vec))
+        if norm > 0:
+            vec = vec / norm
+        return [float(x) for x in vec]
+
+    def dimension(self) -> int:
+        return self._dim
+
+
+def get_embedding_provider(
+    model: str,
+    api_key: str | None = None,
+    **kwargs: Any,
+) -> EmbeddingProvider:
+    """Pick an embedding provider from a short model string.
+
+    ``mock``/``test`` → :class:`MockEmbedding`; ``text-embedding-*`` or
+    ``openai:*`` → :class:`OpenAIEmbedding`; anything else is routed to
+    :class:`SentenceTransformerEmbedding`.
+    """
+    if model in {"mock", "test"}:
+        return MockEmbedding(dimension=kwargs.get("dimension", 384))
+    if model.startswith("text-embedding-") or model.startswith("openai:"):
+        resolved = model.replace("openai:", "", 1) if model.startswith("openai:") else model
+        return OpenAIEmbedding(model=resolved, api_key=api_key)
+    try:
+        return SentenceTransformerEmbedding(model_name=model)
+    except RuntimeError as exc:
+        raise ConfigError(str(exc)) from exc

contextdb/utils/llm.py

@@ -0,0 +1,139 @@
+"""LLM providers.
+
+Minimal async contract over an LLM chat/completion, used by extraction,
+compression, causal inference, and RL-as-policy pathways. Structured output is
+approximated by asking the model for JSON and validating at the call site.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Any
+
+from contextdb.core.exceptions import ConfigError
+
+
+class LLMProvider(ABC):
+    """Common interface for text-in / text-out LLMs."""
+
+    @abstractmethod
+    async def generate(
+        self,
+        prompt: str,
+        system: str = "",
+        temperature: float = 0.0,
+        max_tokens: int = 1000,
+        response_format: type | None = None,
+    ) -> str: ...
+
+
+class OpenAILLM(LLMProvider):
+    """OpenAI chat-completions wrapper with retry."""
+
+    def __init__(
+        self,
+        model: str = "gpt-4o-mini",
+        api_key: str | None = None,
+        max_retries: int = 3,
+    ) -> None:
+        from openai import AsyncOpenAI
+
+        self.model = model
+        self.max_retries = max_retries
+        self._client = AsyncOpenAI(api_key=api_key)
+
+    async def generate(
+        self,
+        prompt: str,
+        system: str = "",
+        temperature: float = 0.0,
+        max_tokens: int = 1000,
+        response_format: type | None = None,
+    ) -> str:
+        messages: list[dict[str, str]] = []
+        if system:
+            messages.append({"role": "system", "content": system})
+        messages.append({"role": "user", "content": prompt})
+        extra: dict[str, Any] = {}
+        if response_format is not None:
+            # Ask for JSON; callers validate against their Pydantic model.
+            extra["response_format"] = {"type": "json_object"}
+
+        delay = 1.0
+        for attempt in range(self.max_retries):
+            try:
+                response = await self._client.chat.completions.create(
+                    model=self.model,
+                    messages=messages,  # type: ignore[arg-type]
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                    **extra,
+                )
+                content = response.choices[0].message.content or ""
+                return content
+            except Exception:  # noqa: BLE001
+                if attempt == self.max_retries - 1:
+                    raise
+                await asyncio.sleep(delay)
+                delay *= 2
+        return ""  # pragma: no cover
+
+
+class MockLLM(LLMProvider):
+    """Scriptable LLM for tests.
+
+    ``responses`` maps substring keys to the string that should be returned
+    when that substring appears anywhere in the prompt. All calls are logged
+    on :attr:`calls` so tests can assert on prompt content.
+    """
+
+    def __init__(
+        self,
+        responses: dict[str, str] | None = None,
+        default: str = '{"facts": [], "entities": []}',
+    ) -> None:
+        self.responses = responses or {}
+        self.default = default
+        self.calls: list[dict[str, Any]] = []
+
+    async def generate(
+        self,
+        prompt: str,
+        system: str = "",
+        temperature: float = 0.0,
+        max_tokens: int = 1000,
+        response_format: type | None = None,
+    ) -> str:
+        self.calls.append(
+            {
+                "prompt": prompt,
+                "system": system,
+                "temperature": temperature,
+                "max_tokens": max_tokens,
+            }
+        )
+        for key, response in self.responses.items():
+            if key in prompt:
+                return response
+        return self.default
+
+
+def get_llm_provider(
+    model: str,
+    api_key: str | None = None,
+    **kwargs: Any,
+) -> LLMProvider:
+    """Route a model name to its provider.
+
+    ``mock``/``test`` → :class:`MockLLM`; ``gpt-*`` or ``o1-*`` →
+    :class:`OpenAILLM`. Anything else raises :class:`ConfigError`.
+    """
+    if model in {"mock", "test"}:
+        return MockLLM(**kwargs)
+    if model.startswith("gpt-") or model.startswith("o1-") or model.startswith("openai:"):
+        resolved = model.replace("openai:", "", 1) if model.startswith("openai:") else model
+        return OpenAILLM(model=resolved, api_key=api_key)
+    raise ConfigError(
+        f"Unknown LLM model '{model}'. Use 'mock' for tests or 'gpt-*'/'o1-*' for OpenAI."
+    )