search-as-code 0.0.1__py3-none-any.whl

search_as_code/__init__.py

@@ -0,0 +1,89 @@
+"""search-as-code — one API, any vector database.
+
+A unified "search as code" agentic harness: agents write portable Python against
+a single primitive API, executed in a sandbox with intermediate state kept out
+of the model context, regardless of which vector DB is underneath.
+
+Quickstart:
+
+    import search_as_code as sac
+
+    s = sac.Session("memory")                       # any backend, same call
+    s.add([{"id": "1", "text": "hello world"}])
+    hits = s.search("hello", top_k=3)
+    print(hits.to_evidence())
+"""
+
+from .adapters import MemoryStore, VectorStore, available, connect, register
+from .embeddings import Embedder, HashEmbedder, as_embedder, get_embedder
+from .explore import ProfilePack, explore
+from .errors import (
+    BackendError,
+    BackendNotFoundError,
+    ConfigurationError,
+    DimensionMismatchError,
+    EmbeddingError,
+    ExtractorRequiredError,
+    GeneratorRequiredError,
+    InvalidArgumentError,
+    InvalidEmbedderError,
+    InvalidFilterError,
+    InvalidModeError,
+    MissingDependencyError,
+    SacError,
+)
+from .primitives import (
+    abstain,
+    auto_filter,
+    confidence,
+    decompose,
+    dedup,
+    diversity_quota,
+    expand,
+    extract,
+    fan_out,
+    freshness,
+    fuse,
+    mmr,
+    normalize_query,
+    normalize_scores,
+    quality_filter,
+    rare_terms,
+    relative_score_fusion,
+    rephrase,
+    rerank,
+    rrf,
+    score_cutoff,
+    topics,
+)
+from .rerankers import CrossEncoderReranker, QwenReranker
+from .sandbox import ExecResult, LocalExecutor, Sandbox
+from .session import Session, route
+from .types import Capabilities, Document, Hit, ResultSet
+
+__version__ = "0.0.1"
+
+__all__ = [
+    # data model
+    "Document", "Hit", "ResultSet", "Capabilities",
+    # backends
+    "connect", "register", "available", "VectorStore", "MemoryStore",
+    # embeddings
+    "Embedder", "HashEmbedder", "as_embedder", "get_embedder",
+    # harness
+    "Session", "route", "Sandbox", "LocalExecutor", "ExecResult",
+    # exploration / onboarding
+    "explore", "ProfilePack",
+    # primitives
+    "fan_out", "fuse", "dedup", "rerank", "freshness", "extract",
+    "mmr", "expand", "decompose", "rephrase", "rrf", "topics", "auto_filter", "score_cutoff",
+    "normalize_scores", "relative_score_fusion", "diversity_quota", "confidence", "abstain",
+    "normalize_query", "rare_terms", "quality_filter",
+    # rerankers
+    "CrossEncoderReranker", "QwenReranker",
+    # errors
+    "SacError", "ConfigurationError", "BackendNotFoundError", "MissingDependencyError",
+    "InvalidArgumentError", "InvalidModeError", "InvalidFilterError",
+    "DimensionMismatchError", "InvalidEmbedderError", "GeneratorRequiredError",
+    "ExtractorRequiredError", "EmbeddingError", "BackendError",
+]

search_as_code/_resilience.py

@@ -0,0 +1,99 @@
+"""Retry / batching helpers for network-backed adapters.
+
+Small, dependency-free utilities so every real adapter gets consistent,
+testable resilience: exponential backoff with jitter around transient backend
+calls, and chunking for bulk upserts. The in-memory backend needs none of this;
+the network adapters (OpenSearch/Qdrant/Chroma/pgvector) opt in.
+"""
+
+from __future__ import annotations
+
+import functools
+import time
+from typing import Any, Callable, Iterable, Iterator, Sequence, TypeVar
+
+from .errors import BackendError, SacError
+
+T = TypeVar("T")
+
+DEFAULT_ATTEMPTS = 3
+DEFAULT_BACKOFF = 0.5  # seconds; grows 0.5, 1.0, 2.0, ...
+DEFAULT_BATCH_SIZE = 500
+
+
+def with_retry(
+    fn: Callable[..., T],
+    *args: Any,
+    attempts: int = DEFAULT_ATTEMPTS,
+    backoff: float = DEFAULT_BACKOFF,
+    exceptions: tuple[type[BaseException], ...] = (Exception,),
+    backend: str = "backend",
+    op: str = "call",
+    sleep: Callable[[float], None] = time.sleep,
+    **kwargs: Any,
+) -> T:
+    """Call ``fn(*args, **kwargs)``, retrying transient ``exceptions``.
+
+    Retries ``attempts`` times with exponential backoff (deterministic multiples
+    of ``backoff``; ``sleep`` is injectable so tests run instantly). On final
+    failure the underlying error is wrapped in :class:`BackendError` (which is
+    also a ``RuntimeError``) with the original attached via ``__cause__``.
+    ``SacError`` is never retried or re-wrapped — those are our own typed,
+    non-transient errors.
+    """
+    if attempts < 1:
+        raise ValueError("attempts must be >= 1")
+    last: BaseException | None = None
+    for i in range(1, attempts + 1):
+        try:
+            return fn(*args, **kwargs)
+        except SacError:
+            raise  # our own typed, non-transient errors: never retry or re-wrap
+        except exceptions as e:  # noqa: BLE001 - deliberately broad, caller-scoped
+            last = e
+            if i >= attempts:
+                break
+            sleep(backoff * (2 ** (i - 1)))
+    raise BackendError(
+        f"{backend} {op} failed after {attempts} attempt(s)",
+        backend=backend,
+        op=op,
+        cause=type(last).__name__ if last else None,
+    ) from last
+
+
+def retry(
+    *,
+    attempts: int = DEFAULT_ATTEMPTS,
+    backoff: float = DEFAULT_BACKOFF,
+    exceptions: tuple[type[BaseException], ...] = (Exception,),
+    backend: str = "backend",
+) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    """Decorator form of :func:`with_retry` (``op`` defaults to the func name)."""
+
+    def deco(fn: Callable[..., T]) -> Callable[..., T]:
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> T:
+            return with_retry(
+                fn, *args,
+                attempts=attempts, backoff=backoff, exceptions=exceptions,
+                backend=backend, op=fn.__name__, **kwargs,
+            )
+
+        return wrapper
+
+    return deco
+
+
+def chunked(seq: Sequence[T] | Iterable[T], size: int = DEFAULT_BATCH_SIZE) -> Iterator[list[T]]:
+    """Yield ``seq`` in lists of at most ``size`` (for batched upserts)."""
+    if size < 1:
+        raise ValueError("size must be >= 1")
+    batch: list[T] = []
+    for item in seq:
+        batch.append(item)
+        if len(batch) >= size:
+            yield batch
+            batch = []
+    if batch:
+        yield batch

search_as_code/adapters/__init__.py

@@ -0,0 +1,5 @@
+from .base import VectorStore
+from .memory import MemoryStore
+from .registry import available, connect, register
+
+__all__ = ["VectorStore", "MemoryStore", "connect", "register", "available"]

search_as_code/adapters/base.py

@@ -0,0 +1,108 @@
+"""The one interface every vector database must satisfy.
+
+This ABC is the contract that makes "no separate SDK per DB" possible: agent
+code and the primitive layer only ever touch these methods.  Adapters translate
+them to their backend's native client.
+
+Design rules for adapters:
+* Return :class:`Hit` with **larger-is-better** scores (convert distances).
+* Accept the portable filter dialect from :mod:`search_as_code.filters`.
+* Declare honestly via :meth:`capabilities`; the primitive layer emulates any
+  capability you report as ``False`` (e.g. client-side rerank or keyword search)
+  so agent code behaves the same everywhere.
+"""
+
+from __future__ import annotations
+
+import abc
+from typing import Any, Optional, Sequence
+
+from ..types import Capabilities, Document, ResultSet
+
+
+class VectorStore(abc.ABC):
+    #: registry key, e.g. "memory", "qdrant"
+    backend: str = "base"
+
+    @abc.abstractmethod
+    def capabilities(self) -> Capabilities: ...
+
+    @abc.abstractmethod
+    def upsert(self, docs: Sequence[Document]) -> None: ...
+
+    @abc.abstractmethod
+    def query_vector(
+        self,
+        vector: Sequence[float],
+        top_k: int = 10,
+        flt: Optional[dict[str, Any]] = None,
+    ) -> ResultSet:
+        """Dense nearest-neighbour search."""
+
+    # ---- optional capabilities: default to NotImplemented so the primitive
+    # layer knows to emulate. Adapters override when the backend supports them.
+
+    def query_keyword(
+        self,
+        text: str,
+        top_k: int = 10,
+        flt: Optional[dict[str, Any]] = None,
+    ) -> ResultSet:
+        raise NotImplementedError
+
+    def query_hybrid(
+        self,
+        vector: Sequence[float],
+        text: str,
+        top_k: int = 10,
+        flt: Optional[dict[str, Any]] = None,
+        alpha: float = 0.5,
+    ) -> ResultSet:
+        raise NotImplementedError
+
+    def query_regex(
+        self,
+        pattern: str,
+        top_k: int = 10,
+        flt: Optional[dict[str, Any]] = None,
+    ) -> ResultSet:
+        """Exact / regex / operator search over document text.
+
+        The "search as code for code" primitive (Hornet, Chroma): agents need
+        exact, case-sensitive matching that semantic search blurs away.
+        """
+        raise NotImplementedError
+
+    def embed_query(self, text: str) -> Optional[list[float]]:
+        """Server-side embedding, if the backend does it. Returns None otherwise."""
+        return None
+
+    @abc.abstractmethod
+    def get(self, ids: Sequence[str]) -> list[Document]: ...
+
+    def delete(self, ids: Sequence[str]) -> None:
+        raise NotImplementedError
+
+    def count(self) -> int:
+        raise NotImplementedError
+
+    # ---- introspection: let the agent discover the DATA SHAPE before querying.
+    # "The first call tells the agent what data exists and what queries are safe to
+    # construct" (schema-first agentic retrieval). Adapters override with real detail.
+
+    def describe_schema(self) -> dict[str, Any]:
+        """Field names -> types (+ sample values where cheap), doc count, backend.
+        Feed this to the LLM so it knows the fields/types before writing retrieval code."""
+        try:
+            n = self.count()
+        except Exception:
+            n = None
+        return {"backend": self.backend, "count": n, "fields": {}, "note": "override for full schema"}
+
+    def sample(self, n: int = 5) -> list[Document]:
+        """A few representative documents so the agent sees the data shape (prose vs
+        table vs fact-card, typical length, metadata)."""
+        raise NotImplementedError
+
+    def close(self) -> None:
+        pass

search_as_code/adapters/chroma.py

@@ -0,0 +1,95 @@
+"""Chroma adapter. ``pip install 'search-as-code[chroma]'``.
+
+Chroma returns L2/cosine *distances* (smaller is better); we convert to a
+larger-is-better similarity so scores are comparable across every backend.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Sequence
+
+from .._resilience import DEFAULT_BATCH_SIZE, chunked
+from ..errors import MissingDependencyError
+from ..filters import normalize
+from ..types import Capabilities, Document, Hit, ResultSet
+from .base import VectorStore
+
+_OP_MAP = {"$eq": "$eq", "$ne": "$ne", "$gt": "$gt", "$gte": "$gte",
+           "$lt": "$lt", "$lte": "$lte", "$in": "$in", "$nin": "$nin"}
+
+
+class ChromaStore(VectorStore):
+    backend = "chroma"
+
+    def __init__(self, collection: str = "sac", persist_path: Optional[str] = None,
+                 batch_size: int = DEFAULT_BATCH_SIZE, **_: Any):
+        try:
+            import chromadb
+        except ImportError as e:  # pragma: no cover - optional dep
+            raise MissingDependencyError("chromadb", extra="search-as-code[chroma]") from e
+        self._client = chromadb.PersistentClient(path=persist_path) if persist_path else chromadb.Client()
+        self._col = self._client.get_or_create_collection(collection)
+        self.batch_size = batch_size
+
+    def capabilities(self) -> Capabilities:
+        return Capabilities(dense=True, keyword=False, hybrid=False, metadata_filter=True)
+
+    def upsert(self, docs: Sequence[Document]) -> None:
+        docs = [d for d in docs if d.vector is not None]
+        for batch in chunked(docs, self.batch_size):
+            self._col.upsert(
+                ids=[d.id for d in batch],
+                embeddings=[d.vector for d in batch],
+                documents=[d.text or "" for d in batch],
+                metadatas=[d.metadata or {"_": ""} for d in batch],
+            )
+
+    def _to_where(self, flt: Optional[dict]) -> Optional[dict]:
+        if not flt:
+            return None
+        clauses = []
+        for field_name, cond in normalize(flt).items():
+            if field_name.startswith("$"):
+                continue
+            clauses.append({field_name: {_OP_MAP[op]: v for op, v in cond.items() if op in _OP_MAP}})
+        if not clauses:
+            return None
+        return clauses[0] if len(clauses) == 1 else {"$and": clauses}
+
+    def query_vector(self, vector, top_k=10, flt=None) -> ResultSet:
+        res = self._col.query(
+            query_embeddings=[list(vector)],
+            n_results=top_k,
+            where=self._to_where(flt),
+            include=["documents", "metadatas", "distances"],
+        )
+        hits = []
+        ids = res["ids"][0]
+        for i, _id in enumerate(ids):
+            dist = res["distances"][0][i]
+            hits.append(
+                Hit(
+                    id=_id,
+                    score=1.0 / (1.0 + float(dist)),  # distance -> larger-is-better
+                    document=Document(
+                        id=_id,
+                        text=res["documents"][0][i],
+                        metadata=res["metadatas"][0][i] or {},
+                    ),
+                    store=self.backend,
+                )
+            )
+        return ResultSet(hits)
+
+    def get(self, ids: Sequence[str]) -> list[Document]:
+        res = self._col.get(ids=list(ids), include=["documents", "metadatas"])
+        return [
+            Document(id=_id, text=res["documents"][i], metadata=res["metadatas"][i] or {})
+            for i, _id in enumerate(res["ids"])
+        ]
+
+    def delete(self, ids: Sequence[str]) -> None:
+        self._col.delete(ids=list(ids))
+
+    def count(self) -> int:
+        return self._col.count()

search_as_code/adapters/faiss_store.py

@@ -0,0 +1,90 @@
+"""FAISS adapter — in-process ANN, no server.
+
+FAISS handles dense (exact inner-product on L2-normalized vectors == cosine); the
+keyword/hybrid/regex capabilities are served client-side by a composed
+:class:`MemoryStore` so agent code behaves identically to a full backend. This is
+the "one primitive API, any vector DB" thesis for an embedded index.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+import numpy as np
+
+from ..types import Capabilities, Document, Hit, ResultSet
+from .base import VectorStore
+from .memory import MemoryStore
+
+
+def _norm(mat: np.ndarray) -> np.ndarray:
+    n = np.linalg.norm(mat, axis=1, keepdims=True)
+    n[n == 0] = 1.0
+    return mat / n
+
+
+class FaissStore(VectorStore):
+    backend = "faiss"
+
+    def __init__(self, dim: int, **_: Any):
+        import faiss
+
+        self.dim = int(dim)
+        self.index = faiss.IndexFlatIP(self.dim)  # exact cosine (vectors normalized)
+        self._ids: list[str] = []
+        self._mem = MemoryStore()  # holds docs + keyword/regex/text
+
+    def capabilities(self) -> Capabilities:
+        return Capabilities(dense=True, keyword=True, hybrid=True, regex=True,
+                            server_side_embedding=False, native_rerank=False,
+                            metadata_filter=True)
+
+    def upsert(self, docs: Sequence[Document]) -> None:
+        vecs, keep = [], []
+        for d in docs:
+            if d.vector is None:
+                continue
+            vecs.append(d.vector)
+            keep.append(d)
+        if vecs:
+            self.index.add(_norm(np.asarray(vecs, dtype=np.float32)))
+            self._ids.extend(d.id for d in keep)
+        self._mem.upsert(docs)
+
+    def query_vector(self, vector, top_k=10, flt=None) -> ResultSet:
+        if self.index.ntotal == 0:
+            return ResultSet()
+        q = _norm(np.asarray([vector], dtype=np.float32))
+        # over-fetch when filtering, then apply the portable filter client-side
+        k = min(self.index.ntotal, top_k * (8 if flt else 1))
+        scores, idx = self.index.search(q, k)
+        from ..filters import matches
+        hits = []
+        for s, i in zip(scores[0], idx[0]):
+            if i < 0:
+                continue
+            doc = self._mem._docs.get(self._ids[i])
+            if doc is None or (flt and not matches(doc.metadata, flt)):
+                continue
+            hits.append(Hit(id=self._ids[i], score=float(s), document=doc, store=self.backend))
+            if len(hits) >= top_k:
+                break
+        return ResultSet(hits)
+
+    def query_keyword(self, text, top_k=10, flt=None) -> ResultSet:
+        return self._mem.query_keyword(text, top_k=top_k, flt=flt)
+
+    def query_regex(self, pattern, top_k=10, flt=None) -> ResultSet:
+        return self._mem.query_regex(pattern, top_k=top_k, flt=flt)
+
+    def query_hybrid(self, vector, text, top_k=10, flt=None, alpha=0.5) -> ResultSet:
+        dense = self.query_vector(vector, top_k=top_k * 4, flt=flt)
+        kw = self.query_keyword(text, top_k=top_k * 4, flt=flt)
+        from ..primitives import fuse
+        return fuse([dense, kw], weights=[alpha, 1 - alpha]).top(top_k)
+
+    def get(self, ids: Sequence[str]) -> list[Document]:
+        return self._mem.get(ids)
+
+    def count(self) -> int:
+        return int(self.index.ntotal)

search_as_code/adapters/memory.py

@@ -0,0 +1,148 @@
+"""Zero-dependency reference adapter.
+
+Brute-force cosine over an in-process dict plus a naive BM25-ish keyword score.
+It exists so the entire harness — primitives, session, sandbox, tests, demos —
+runs with nothing installed but numpy.  Treat it as the executable spec every
+real adapter must match.
+"""
+
+from __future__ import annotations
+
+import math
+from collections import Counter
+from typing import Any, Optional, Sequence
+
+import numpy as np
+
+from ..embeddings import _tokenize
+from ..filters import matches
+from ..types import Capabilities, Document, Hit, ResultSet
+from .base import VectorStore
+
+
+class MemoryStore(VectorStore):
+    backend = "memory"
+
+    def __init__(self, **_: Any):
+        self._docs: dict[str, Document] = {}
+        self._matrix: Optional[np.ndarray] = None
+        self._ids: list[str] = []
+        self._dirty = True
+
+    def capabilities(self) -> Capabilities:
+        return Capabilities(
+            dense=True,
+            keyword=True,
+            hybrid=True,
+            regex=True,
+            server_side_embedding=False,
+            native_rerank=False,
+            metadata_filter=True,
+        )
+
+    def upsert(self, docs: Sequence[Document]) -> None:
+        for d in docs:
+            self._docs[d.id] = d
+        self._dirty = True
+
+    def _rebuild(self) -> None:
+        self._ids = [d.id for d in self._docs.values() if d.vector is not None]
+        if self._ids:
+            mat = np.array([self._docs[i].vector for i in self._ids], dtype=np.float32)
+            norms = np.linalg.norm(mat, axis=1, keepdims=True)
+            norms[norms == 0] = 1.0
+            self._matrix = mat / norms
+        else:
+            self._matrix = None
+        self._dirty = False
+
+    def _candidates(self, flt: Optional[dict[str, Any]]) -> list[int]:
+        if not flt:
+            return list(range(len(self._ids)))
+        return [i for i, did in enumerate(self._ids) if matches(self._docs[did].metadata, flt)]
+
+    def query_vector(self, vector, top_k=10, flt=None) -> ResultSet:
+        if self._dirty:
+            self._rebuild()
+        if self._matrix is None:
+            return ResultSet()
+        q = np.asarray(vector, dtype=np.float32)
+        qn = np.linalg.norm(q) or 1.0
+        sims = self._matrix @ (q / qn)
+        cand = self._candidates(flt)
+        cand.sort(key=lambda i: sims[i], reverse=True)
+        hits = [
+            Hit(id=self._ids[i], score=float(sims[i]), document=self._docs[self._ids[i]], store=self.backend)
+            for i in cand[:top_k]
+        ]
+        return ResultSet(hits)
+
+    def query_keyword(self, text, top_k=10, flt=None) -> ResultSet:
+        terms = Counter(_tokenize(text))
+        if not terms:
+            return ResultSet()
+        docs = [d for d in self._docs.values() if d.text and matches(d.metadata, flt)]
+        n = len(docs) or 1
+        df: Counter = Counter()
+        toks_by_id: dict[str, Counter] = {}
+        for d in docs:
+            toks = Counter(_tokenize(d.text or ""))
+            toks_by_id[d.id] = toks
+            for t in set(toks):
+                df[t] += 1
+        scored = []
+        for d in docs:
+            toks = toks_by_id[d.id]
+            dl = sum(toks.values()) or 1
+            score = 0.0
+            for term, qf in terms.items():
+                if term in toks:
+                    idf = math.log(1 + n / (1 + df[term]))
+                    score += qf * idf * (toks[term] / dl)
+            if score > 0:
+                scored.append(Hit(id=d.id, score=score, document=d, store=self.backend))
+        scored.sort(key=lambda h: h.score, reverse=True)
+        return ResultSet(scored[:top_k])
+
+    def query_hybrid(self, vector, text, top_k=10, flt=None, alpha=0.5) -> ResultSet:
+        dense = self.query_vector(vector, top_k=top_k * 4, flt=flt)
+        kw = self.query_keyword(text, top_k=top_k * 4, flt=flt)
+        from ..primitives import fuse  # local import to avoid cycle
+
+        return fuse([dense, kw], weights=[alpha, 1 - alpha]).top(top_k)
+
+    def query_regex(self, pattern, top_k=10, flt=None) -> ResultSet:
+        import re
+
+        rx = re.compile(pattern)
+        hits = []
+        for d in self._docs.values():
+            if not d.text or not matches(d.metadata, flt):
+                continue
+            found = rx.findall(d.text)
+            if found:
+                hits.append(Hit(id=d.id, score=float(len(found)), document=d, store=self.backend))
+        hits.sort(key=lambda h: h.score, reverse=True)
+        return ResultSet(hits[:top_k])
+
+    def get(self, ids: Sequence[str]) -> list[Document]:
+        return [self._docs[i] for i in ids if i in self._docs]
+
+    def delete(self, ids: Sequence[str]) -> None:
+        for i in ids:
+            self._docs.pop(i, None)
+        self._dirty = True
+
+    def count(self) -> int:
+        return len(self._docs)
+
+    def sample(self, n: int = 5) -> list[Document]:
+        return list(self._docs.values())[:n]
+
+    def describe_schema(self) -> dict:
+        docs = self.sample(3)
+        keys: set = set()
+        for d in docs:
+            keys |= set((d.metadata or {}).keys())
+        return {"backend": self.backend, "count": len(self._docs),
+                "metadata_keys": sorted(keys), "sample_text": [(d.text or "")[:200] for d in docs]}