kbguard 0.1.0__py3-none-any.whl

kbguard/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

kbguard/adapters/base.py

@@ -0,0 +1,70 @@
+from collections.abc import Iterator
+from typing import Any, Literal, Protocol
+
+from kbguard.models import Chunk
+
+
+class AdapterConnectionError(Exception):
+    """Raised when an adapter can't actually reach/authenticate to its
+    store — wrong host/port, server down, bad credentials — as opposed to
+    a missing config *value* (that's a plain `ValueError`, raised before
+    any real call is attempted) or a missing *env var* (`ConfigError`).
+    Every adapter wraps its own client library's real connection failure
+    into this one shared type specifically so `cli.py` can print one clean
+    message for it; without this, a bad password or an unreachable host
+    surfaces as a full raw traceback from psycopg/qdrant_client/pinecone's
+    own exception types, confirmed by hand before this existed."""
+
+
+class VectorStoreAdapter(Protocol):
+    """The only contract kbguard knows about a vector store. Everything
+    downstream (ingest, checks, evals) works on the `Chunk` records this
+    produces — never on a store-specific client — which is what makes any
+    vector store pluggable.
+    """
+
+    name: str
+
+    # "full": iter_chunks() genuinely walks the entire collection (every
+    # native adapter). "sampled": it can't — the wrapped store exposes no
+    # enumeration API at all, only similarity search, so iter_chunks()
+    # approximates coverage via search-based sampling instead (so far,
+    # only bridge_langchain.py). SnapshotWriter records this verbatim into
+    # the snapshot manifest, so nothing downstream can mistake a sampled
+    # scan for a complete one.
+    iteration_mode: Literal["full", "sampled"]
+
+    def connect(self, **settings: Any) -> None: ...
+
+    def collections(self) -> list[str]: ...
+
+    def count(self, collection: str) -> int: ...
+
+    def iter_chunks(
+        self,
+        collection: str,
+        batch_size: int = 500,
+        with_embeddings: bool = True,
+    ) -> Iterator[Chunk]:
+        """Must page through the *entire* collection without ever holding it
+        all in memory at once — implementations page internally (keyset or
+        cursor-based, store-dependent) and yield one Chunk at a time.
+        """
+        ...
+
+    def search(
+        self,
+        collection: str,
+        query_vector: list[float],
+        k: int,
+        filters: dict[str, Any] | None = None,
+    ) -> list[tuple[Chunk, float]]:
+        """Results are returned **best-match-first**; callers rely on this
+        ordering, not on the numeric score. Scores only need to be
+        comparable *within a single call* — never across two different
+        adapters, two different collections, or against some assumed
+        absolute scale (a Chroma distance and a Qdrant similarity are not
+        the same kind of number, and callers must not assume either
+        direction — lower-is-better or higher-is-better — holds generically).
+        """
+        ...

kbguard/adapters/bridge_langchain.py

@@ -0,0 +1,151 @@
+import hashlib
+import random
+from collections.abc import Iterator
+from typing import Any
+
+from kbguard.models import Chunk
+
+# Deterministic seed: sampling coverage should be reproducible between two
+# scans of an unchanged collection (so a `diff` isn't just sampling noise),
+# not a source of real randomness — this is sampling, not security.
+_RANDOM_SEED = 0
+
+
+def _require_setting(settings: dict[str, Any], key: str) -> Any:
+    if key not in settings:
+        raise ValueError(f"store.settings.{key} is required for the bridge_langchain adapter")
+    return settings[key]
+
+
+def _stable_id(page_content: str, metadata: dict[str, Any]) -> str:
+    # Used only when the wrapped store's own Document.id is None (common —
+    # many LangChain VectorStore implementations never set it). Content-
+    # based, not random, so the same chunk gets the same id across two
+    # separate scans; two genuinely distinct chunks that happen to share
+    # identical text and metadata would collide, which is an accepted,
+    # documented gap, not a silently-assumed-safe one.
+    payload = page_content + repr(sorted(metadata.items()))
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:16]
+
+
+class LangChainBridgeAdapter:
+    """Wraps a LangChain VectorStore instance the client already has,
+    instead of connecting to a store natively. langchain_core.vectorstores
+    .VectorStore (checked against the installed library, not assumed)
+    exposes no "iterate everything" or "count" method at all — only
+    similarity search — so this adapter is deliberately less complete than
+    the native ones: iter_chunks() approximates coverage via search-based
+    sampling, reported honestly as iteration_mode = "sampled" (see
+    adapters/base.py), and count() refuses to report a number it has no
+    way to verify, rather than guessing.
+
+    A further real gap, worth naming explicitly: none of VectorStore's
+    search methods return a hit's own embedding vector, so every Chunk
+    this adapter produces has embedding=None regardless of the
+    with_embeddings argument — unlike every native adapter, which returns
+    real stored vectors. A caller relying on `embeddings.reuse_stored`
+    will recompute every chunk's embedding through the configured model
+    instead, which is a real cost difference worth knowing about, not a
+    bug in how `with_embeddings` is honored here.
+    """
+
+    name = "bridge_langchain"
+    iteration_mode = "sampled"
+
+    def __init__(self) -> None:
+        self._vectorstore: Any = None
+        self._collection_name = "default"
+        self._embedding_dim: int | None = None
+        self._sample_queries: list[str] = []
+        self._probe_count = 20
+
+    def connect(self, **settings: Any) -> None:
+        self._vectorstore = _require_setting(settings, "vectorstore")
+        # Needed to generate the random probe vectors iter_chunks() uses to
+        # sample beyond whatever sample_queries covers — there's no API on
+        # a generic VectorStore to ask its own embedding dimension.
+        self._embedding_dim = _require_setting(settings, "embedding_dim")
+        # Real queries a human knows are representative of this
+        # collection's content sample far better than random vectors do —
+        # optional, but strongly recommended; see docs/guide/adapters/
+        # bridge-langchain.md.
+        self._sample_queries = settings.get("sample_queries", [])
+        self._probe_count = settings.get("probe_count", 20)
+        self._collection_name = settings.get("collection", "default")
+
+    def _require_vectorstore(self) -> Any:
+        if self._vectorstore is None:
+            raise RuntimeError("LangChainBridgeAdapter.connect() must be called before use")
+        return self._vectorstore
+
+    def collections(self) -> list[str]:
+        # One already-constructed VectorStore instance is already scoped
+        # to whatever single collection/index it was built against — there
+        # is no "list every collection this client has" operation to call.
+        return [self._collection_name]
+
+    def count(self, collection: str) -> int:
+        raise NotImplementedError(
+            "LangChainBridgeAdapter can't report an exact count -- the wrapped VectorStore has no "
+            "universal count/enumerate method. iter_chunks() reports iteration_mode='sampled' instead "
+            "of a count it has no way to verify; see docs/guide/adapters/bridge-langchain.md"
+        )
+
+    def iter_chunks(
+        self,
+        collection: str,
+        batch_size: int = 500,
+        with_embeddings: bool = True,
+    ) -> Iterator[Chunk]:
+        vectorstore = self._require_vectorstore()
+        seen_ids: set[str] = set()
+
+        for query in self._sample_queries:
+            for doc, _score in vectorstore.similarity_search_with_score(query, k=batch_size):
+                chunk = self._to_chunk(doc)
+                if chunk.id in seen_ids:
+                    continue
+                seen_ids.add(chunk.id)
+                yield chunk
+
+        # Beyond the (optional) human-provided sample_queries, probe with
+        # random unit-ish vectors in the embedding space to pick up content
+        # those specific queries wouldn't surface — still not a guarantee
+        # of full coverage, just a broader, documented sample.
+        assert self._embedding_dim is not None  # guaranteed by connect()
+        rng = random.Random(_RANDOM_SEED)
+        for _ in range(self._probe_count):
+            probe_vector = [rng.uniform(-1.0, 1.0) for _ in range(self._embedding_dim)]
+            for doc in vectorstore.similarity_search_by_vector(probe_vector, k=batch_size):
+                chunk = self._to_chunk(doc)
+                if chunk.id in seen_ids:
+                    continue
+                seen_ids.add(chunk.id)
+                yield chunk
+
+    def search(
+        self,
+        collection: str,
+        query_vector: list[float],
+        k: int,
+        filters: dict[str, Any] | None = None,
+    ) -> list[tuple[Chunk, float]]:
+        if filters is not None:
+            raise NotImplementedError("LangChainBridgeAdapter.search() doesn't support filters yet")
+
+        vectorstore = self._require_vectorstore()
+        # similarity_search_by_vector (checked against the installed
+        # library) returns plain Documents, best-match-first, but no
+        # scores at all -- unlike every native adapter's search(). A
+        # descending rank-based placeholder satisfies the Protocol's
+        # best-match-first ordering guarantee; it is not a real similarity
+        # value and must never be compared against another adapter's
+        # score, which adapters/base.py's own docstring already rules out
+        # generically.
+        docs = vectorstore.similarity_search_by_vector(query_vector, k=k)
+        return [(self._to_chunk(doc), 1.0 - i * 0.001) for i, doc in enumerate(docs)]
+
+    def _to_chunk(self, doc: Any) -> Chunk:
+        metadata = dict(doc.metadata or {})
+        chunk_id = doc.id if getattr(doc, "id", None) is not None else _stable_id(doc.page_content, metadata)
+        return Chunk(id=chunk_id, text=doc.page_content, embedding=None, metadata=metadata)

kbguard/adapters/chroma.py

@@ -0,0 +1,125 @@
+from collections.abc import Iterator, Sequence
+from typing import Any
+
+import chromadb
+from chromadb.api import ClientAPI
+from chromadb.errors import ChromaError
+
+from kbguard.adapters.base import AdapterConnectionError
+from kbguard.models import Chunk
+
+# kbguard always supplies its own embeddings (ingest/embeddings.py — reuse
+# stored or recompute with the configured model) and never asks Chroma to
+# embed text itself. Passing embedding_function=None everywhere stops Chroma
+# from silently downloading its default embedding model on first use, which
+# would otherwise be a surprise network call this project's "runs entirely
+# in your environment" principle explicitly rules out.
+_NO_EMBEDDING_FUNCTION = None
+
+
+class ChromaAdapter:
+    name = "chroma"
+    iteration_mode = "full"
+
+    def __init__(self) -> None:
+        self._client: ClientAPI | None = None
+
+    def connect(self, **settings: Any) -> None:
+        path = settings.get("path")
+        # No `path` -> in-process, in-memory client: this is what makes
+        # Chroma the fastest adapter to develop and test against (no server,
+        # no Docker). Given a `path`, data persists to that local directory
+        # between runs instead.
+        try:
+            if path is not None:
+                self._client = chromadb.PersistentClient(path=str(path))
+            else:
+                self._client = chromadb.EphemeralClient()
+        except ChromaError as exc:
+            # Local-only (no remote host to misconfigure), but an
+            # unwritable/nonexistent path still fails eagerly right here
+            # (confirmed by hand) -- chromadb.errors.ChromaError is the
+            # real common base, confirmed via the real exception's own
+            # __mro__ before relying on it.
+            raise AdapterConnectionError(f"could not open the Chroma store at {path!r}: {exc}") from exc
+
+    def _require_client(self) -> ClientAPI:
+        if self._client is None:
+            raise RuntimeError("ChromaAdapter.connect() must be called before use")
+        return self._client
+
+    def collections(self) -> list[str]:
+        return [c.name for c in self._require_client().list_collections()]
+
+    def count(self, collection: str) -> int:
+        col = self._require_client().get_collection(collection, embedding_function=_NO_EMBEDDING_FUNCTION)
+        return col.count()
+
+    def iter_chunks(
+        self,
+        collection: str,
+        batch_size: int = 500,
+        with_embeddings: bool = True,
+    ) -> Iterator[Chunk]:
+        col = self._require_client().get_collection(collection, embedding_function=_NO_EMBEDDING_FUNCTION)
+        include = ["documents", "metadatas", *(["embeddings"] if with_embeddings else [])]
+
+        offset = 0
+        while True:
+            batch = col.get(limit=batch_size, offset=offset, include=include)  # type: ignore[arg-type]
+            ids = batch["ids"]
+            if not ids:
+                return
+
+            documents = batch["documents"]
+            metadatas = batch["metadatas"]
+            embeddings = batch["embeddings"]
+
+            for i, chunk_id in enumerate(ids):
+                text = documents[i] if documents is not None else None
+                metadata = metadatas[i] if metadatas is not None else None
+                embedding = [float(x) for x in embeddings[i]] if embeddings is not None else None
+                yield Chunk(
+                    id=chunk_id,
+                    text=text or "",
+                    embedding=embedding,
+                    metadata=dict(metadata or {}),
+                )
+
+            offset += len(ids)
+
+    def search(
+        self,
+        collection: str,
+        query_vector: list[float],
+        k: int,
+        filters: dict[str, Any] | None = None,
+    ) -> list[tuple[Chunk, float]]:
+        col = self._require_client().get_collection(collection, embedding_function=_NO_EMBEDDING_FUNCTION)
+        query_embeddings: list[Sequence[float] | Sequence[int]] = [query_vector]
+        result = col.query(
+            query_embeddings=query_embeddings,
+            n_results=k,
+            where=filters,
+            include=["documents", "metadatas", "embeddings", "distances"],
+        )
+
+        # query() nests one list per query vector; we only ever send one.
+        ids = result["ids"][0]
+        documents = result["documents"][0] if result["documents"] is not None else None
+        metadatas = result["metadatas"][0] if result["metadatas"] is not None else None
+        embeddings = result["embeddings"][0] if result["embeddings"] is not None else None
+        distances = result["distances"][0] if result["distances"] is not None else None
+
+        pairs: list[tuple[Chunk, float]] = []
+        for i, chunk_id in enumerate(ids):
+            text = documents[i] if documents is not None else None
+            metadata = metadatas[i] if metadatas is not None else None
+            embedding = [float(x) for x in embeddings[i]] if embeddings is not None else None
+            chunk = Chunk(id=chunk_id, text=text or "", embedding=embedding, metadata=dict(metadata or {}))
+            # Chroma's query() already returns nearest-first, satisfying the
+            # Protocol's best-match-first guarantee (adapters/base.py)
+            # without any re-sorting needed here.
+            score = float(distances[i]) if distances is not None else 0.0
+            pairs.append((chunk, score))
+        return pairs

kbguard/adapters/pgvector.py

@@ -0,0 +1,191 @@
+import re
+from collections.abc import Iterator
+from typing import Any
+
+import psycopg
+from pgvector import Vector
+from pgvector.psycopg import register_vector
+from psycopg import sql
+
+from kbguard.adapters.base import AdapterConnectionError
+from kbguard.config import resolve_env
+from kbguard.models import Chunk
+
+# Table/column names come from config (spec: pgvector adapter must work
+# against a client's existing schema, not impose its own). Identifiers
+# can't be passed as query parameters the way values can (no %s for a
+# table/column name) — psycopg.sql.Identifier() safely quotes whatever
+# string it's given, but this still validates against a strict allowlist
+# first and fails with a clear message, rather than silently accepting an
+# unusual identifier that merely happens not to be a SQL injection vector.
+_IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+
+def _safe_identifier(name: str, *, what: str) -> sql.Identifier:
+    if not _IDENTIFIER_RE.match(name):
+        raise ValueError(
+            f"{what} {name!r} is not a valid SQL identifier (letters, digits, underscore, "
+            f"can't start with a digit) — refusing to use it in a query"
+        )
+    return sql.Identifier(name)
+
+
+def _require_setting(settings: dict[str, Any], key: str) -> Any:
+    if key not in settings:
+        raise ValueError(f"store.settings.{key} is required for the pgvector adapter")
+    return settings[key]
+
+
+class PgVectorAdapter:
+    name = "pgvector"
+    iteration_mode = "full"
+
+    def __init__(self) -> None:
+        self._conn: psycopg.Connection[Any] | None = None
+        self._id_column = "id"
+        self._text_column = "text"
+        self._metadata_column = "metadata"
+        self._embedding_column = "embedding"
+
+    def connect(self, **settings: Any) -> None:
+        host = settings.get("host", "localhost")
+        port = settings.get("port", 5432)
+        user = _require_setting(settings, "user")
+        dbname = _require_setting(settings, "dbname")
+        password_env = settings.get("password_env")
+        password = resolve_env(password_env) if password_env else settings.get("password")
+
+        self._id_column = settings.get("id_column", "id")
+        self._text_column = settings.get("text_column", "text")
+        self._metadata_column = settings.get("metadata_column", "metadata")
+        self._embedding_column = settings.get("embedding_column", "embedding")
+        # Validated once here, not per-query — a bad identifier should fail
+        # at connect() time, not confusingly deep inside the first scan.
+        for value, what in [
+            (self._id_column, "id_column"),
+            (self._text_column, "text_column"),
+            (self._metadata_column, "metadata_column"),
+            (self._embedding_column, "embedding_column"),
+        ]:
+            _safe_identifier(value, what=what)
+
+        # No explicit connect_timeout, by default, means no timeout at all
+        # -- confirmed by hand with a real unresponsive port: libpq (via
+        # psycopg) took 260 SECONDS to give up on its own, not an instant
+        # connection-refused. Same class of gap as adapters/pinecone.py's
+        # own timeout fix earlier -- a misconfigured/unreachable host
+        # should fail in seconds, not hang for minutes. Note this is a
+        # PER-ADDRESS timeout, not a total deadline -- confirmed by hand
+        # that "localhost" resolves to both ::1 and 127.0.0.1, so a real
+        # worst case is roughly 2x this value, not exactly this value.
+        connect_timeout = settings.get("connect_timeout", 10)
+
+        # autocommit: kbguard never writes, so there's nothing to commit or
+        # roll back — no reason to hold an open transaction across calls.
+        try:
+            self._conn = psycopg.connect(
+                host=host,
+                port=port,
+                user=user,
+                password=password,
+                dbname=dbname,
+                autocommit=True,
+                connect_timeout=connect_timeout,
+            )
+        except psycopg.Error as exc:
+            # psycopg.connect() is eager -- a bad host/port/password/dbname
+            # raises right here, confirmed by hand, unlike qdrant/pinecone's
+            # lazy clients. psycopg.Error is the real common base for this
+            # (confirmed via OperationalError.__mro__ before relying on it).
+            raise AdapterConnectionError(f"could not connect to pgvector at {host}:{port}: {exc}") from exc
+        # Requires the vector extension already enabled on the target
+        # database — kbguard never runs CREATE EXTENSION itself (DDL is a
+        # write; read-only is a hard design constraint, not a default).
+        register_vector(self._conn)
+
+    def _require_conn(self) -> psycopg.Connection[Any]:
+        if self._conn is None:
+            raise RuntimeError("PgVectorAdapter.connect() must be called before use")
+        return self._conn
+
+    def collections(self) -> list[str]:
+        rows = self._require_conn().execute(
+            "SELECT table_name FROM information_schema.tables "
+            "WHERE table_schema = 'public' AND table_type = 'BASE TABLE' ORDER BY table_name"
+        ).fetchall()
+        return [row[0] for row in rows]
+
+    def count(self, collection: str) -> int:
+        table = _safe_identifier(collection, what="collection")
+        row = self._require_conn().execute(sql.SQL("SELECT count(*) FROM {}").format(table)).fetchone()
+        assert row is not None
+        return row[0]  # type: ignore[no-any-return]
+
+    def iter_chunks(
+        self,
+        collection: str,
+        batch_size: int = 500,
+        with_embeddings: bool = True,
+    ) -> Iterator[Chunk]:
+        conn = self._require_conn()
+        table = _safe_identifier(collection, what="collection")
+        id_col = sql.Identifier(self._id_column)
+        select_cols = [id_col, sql.Identifier(self._text_column), sql.Identifier(self._metadata_column)]
+        if with_embeddings:
+            select_cols.append(sql.Identifier(self._embedding_column))
+
+        # Keyset (id > :last), not OFFSET: a row inserted or deleted mid-scan
+        # can make OFFSET skip or repeat rows on a live, multi-writer table
+        # — a real risk for pgvector specifically, unlike Chroma, which is
+        # typically local and single-writer.
+        query = sql.SQL("SELECT {} FROM {} WHERE {} > %s ORDER BY {} LIMIT %s").format(
+            sql.SQL(", ").join(select_cols), table, id_col, id_col
+        )
+
+        last_id = ""
+        while True:
+            rows = conn.execute(query, (last_id, batch_size)).fetchall()
+            if not rows:
+                return
+            for row in rows:
+                chunk_id, text, metadata = row[0], row[1], row[2]
+                embedding = row[3].to_list() if with_embeddings and row[3] is not None else None
+                yield Chunk(id=chunk_id, text=text, embedding=embedding, metadata=dict(metadata or {}))
+            last_id = rows[-1][0]
+
+    def search(
+        self,
+        collection: str,
+        query_vector: list[float],
+        k: int,
+        filters: dict[str, Any] | None = None,
+    ) -> list[tuple[Chunk, float]]:
+        if filters is not None:
+            # Silently ignoring filters would return unfiltered results to
+            # a caller that explicitly asked for filtered ones — wrong in
+            # a way that's hard to notice. No caller passes filters yet;
+            # pgvector's filter translation lands when one actually needs it.
+            raise NotImplementedError("PgVectorAdapter.search() doesn't support filters yet")
+
+        conn = self._require_conn()
+        table = _safe_identifier(collection, what="collection")
+        id_col = sql.Identifier(self._id_column)
+        text_col = sql.Identifier(self._text_column)
+        meta_col = sql.Identifier(self._metadata_column)
+        emb_col = sql.Identifier(self._embedding_column)
+
+        # <=> is pgvector's cosine distance operator (lower = closer);
+        # ORDER BY distance ASC satisfies the Protocol's best-match-first
+        # guarantee directly — no re-sorting needed.
+        query = sql.SQL("SELECT {}, {}, {}, {}, {} <=> %s AS distance FROM {} ORDER BY distance LIMIT %s").format(
+            id_col, text_col, meta_col, emb_col, emb_col, table
+        )
+        rows = conn.execute(query, (Vector(query_vector), k)).fetchall()
+
+        results: list[tuple[Chunk, float]] = []
+        for row in rows:
+            chunk_id, text, metadata, raw_embedding, distance = row
+            embedding = raw_embedding.to_list() if raw_embedding is not None else None
+            chunk = Chunk(id=chunk_id, text=text, embedding=embedding, metadata=dict(metadata or {}))
+            results.append((chunk, float(distance)))
+        return results