beaver-db 0.4.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

beaver/__init__.py

@@ -1 +1,2 @@
-from .core import BeaverDB, Document
+from .core import BeaverDB
+from .collections import Document, WalkDirection

beaver/collections.py

@@ -0,0 +1,327 @@
+import json
+import sqlite3
+import uuid
+from enum import Enum
+from typing import Any, List, Literal, Set
+
+import numpy as np
+from scipy.spatial import cKDTree
+
+
+class WalkDirection(Enum):
+    OUTGOING = "outgoing"
+    INCOMING = "incoming"
+
+
+class Document:
+    """A data class representing a single item in a collection."""
+
+    def __init__(
+        self, embedding: list[float] | None = None, id: str | None = None, **metadata
+    ):
+        self.id = id or str(uuid.uuid4())
+
+        if embedding is None:
+            self.embedding = None
+        else:
+            if not isinstance(embedding, list) or not all(
+                isinstance(x, (int, float)) for x in embedding
+            ):
+                raise TypeError("Embedding must be a list of numbers.")
+            self.embedding = np.array(embedding, dtype=np.float32)
+
+        for key, value in metadata.items():
+            setattr(self, key, value)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serializes the document's metadata to a dictionary."""
+        metadata = self.__dict__.copy()
+        metadata.pop("embedding", None)
+        metadata.pop("id", None)
+        return metadata
+
+    def __repr__(self):
+        metadata_str = ", ".join(f"{k}={v!r}" for k, v in self.to_dict().items())
+        return f"Document(id='{self.id}', {metadata_str})"
+
+
+class CollectionWrapper:
+    """
+    A wrapper for multi-modal collection operations with an in-memory ANN index,
+    FTS, and graph capabilities.
+    """
+
+    def __init__(self, name: str, conn: sqlite3.Connection):
+        self._name = name
+        self._conn = conn
+        self._kdtree: cKDTree | None = None
+        self._doc_ids: List[str] = []
+        self._local_index_version = -1  # Version of the in-memory index
+
+    def _flatten_metadata(self, metadata: dict, prefix: str = "") -> dict[str, str]:
+        """Flattens a nested dictionary and filters for string values."""
+        flat_dict = {}
+        for key, value in metadata.items():
+            new_key = f"{prefix}__{key}" if prefix else key
+            if isinstance(value, dict):
+                flat_dict.update(self._flatten_metadata(value, new_key))
+            elif isinstance(value, str):
+                flat_dict[new_key] = value
+        return flat_dict
+
+    def _get_db_version(self) -> int:
+        """Gets the current version of the collection from the database."""
+        cursor = self._conn.cursor()
+        cursor.execute(
+            "SELECT version FROM beaver_collection_versions WHERE collection_name = ?",
+            (self._name,),
+        )
+        result = cursor.fetchone()
+        return result[0] if result else 0
+
+    def _is_index_stale(self) -> bool:
+        """Checks if the in-memory index is out of sync with the DB."""
+        if self._local_index_version == -1:
+            return True
+        return self._local_index_version < self._get_db_version()
+
+    def index(self, document: Document, *, fts: bool = True):
+        """Indexes a Document, performing an upsert and updating the FTS index."""
+        with self._conn:
+            if fts:
+                self._conn.execute(
+                    "DELETE FROM beaver_fts_index WHERE collection = ? AND item_id = ?",
+                    (self._name, document.id),
+                )
+                string_fields = self._flatten_metadata(document.to_dict())
+                if string_fields:
+                    fts_data = [
+                        (self._name, document.id, path, content)
+                        for path, content in string_fields.items()
+                    ]
+                    self._conn.executemany(
+                        "INSERT INTO beaver_fts_index (collection, item_id, field_path, field_content) VALUES (?, ?, ?, ?)",
+                        fts_data,
+                    )
+
+            self._conn.execute(
+                "INSERT OR REPLACE INTO beaver_collections (collection, item_id, item_vector, metadata) VALUES (?, ?, ?, ?)",
+                (
+                    self._name,
+                    document.id,
+                    (
+                        document.embedding.tobytes()
+                        if document.embedding is not None
+                        else None
+                    ),
+                    json.dumps(document.to_dict()),
+                ),
+            )
+            # Atomically increment the collection's version number
+            self._conn.execute(
+                """
+                INSERT INTO beaver_collection_versions (collection_name, version) VALUES (?, 1)
+                ON CONFLICT(collection_name) DO UPDATE SET version = version + 1
+                """,
+                (self._name,),
+            )
+
+    def drop(self, document: Document):
+        """Removes a document and all its associated data from the collection."""
+        if not isinstance(document, Document):
+            raise TypeError("Item to drop must be a Document object.")
+        with self._conn:
+            self._conn.execute(
+                "DELETE FROM beaver_collections WHERE collection = ? AND item_id = ?",
+                (self._name, document.id),
+            )
+            self._conn.execute(
+                "DELETE FROM beaver_fts_index WHERE collection = ? AND item_id = ?",
+                (self._name, document.id),
+            )
+            self._conn.execute(
+                "DELETE FROM beaver_edges WHERE collection = ? AND (source_item_id = ? OR target_item_id = ?)",
+                (self._name, document.id, document.id),
+            )
+            self._conn.execute(
+                """
+                INSERT INTO beaver_collection_versions (collection_name, version) VALUES (?, 1)
+                ON CONFLICT(collection_name) DO UPDATE SET version = version + 1
+                """,
+                (self._name,),
+            )
+
+    def refresh(self):
+        """Forces a rebuild of the in-memory ANN index from data in SQLite."""
+        cursor = self._conn.cursor()
+        cursor.execute(
+            "SELECT item_id, item_vector FROM beaver_collections WHERE collection = ? AND item_vector IS NOT NULL",
+            (self._name,),
+        )
+        vectors, self._doc_ids = [], []
+        for row in cursor.fetchall():
+            self._doc_ids.append(row["item_id"])
+            vectors.append(np.frombuffer(row["item_vector"], dtype=np.float32))
+
+        self._kdtree = cKDTree(vectors) if vectors else None
+        self._local_index_version = self._get_db_version()
+
+    def search(
+        self, vector: list[float], top_k: int = 10
+    ) -> list[tuple[Document, float]]:
+        """Performs a fast approximate nearest neighbor search."""
+        if self._is_index_stale():
+            self.refresh()
+        if not self._kdtree:
+            return []
+
+        distances, indices = self._kdtree.query(
+            np.array(vector, dtype=np.float32), k=top_k
+        )
+        if top_k == 1:
+            distances, indices = [distances], [indices]
+
+        result_ids = [self._doc_ids[i] for i in indices]
+        placeholders = ",".join("?" for _ in result_ids)
+        sql = f"SELECT item_id, item_vector, metadata FROM beaver_collections WHERE collection = ? AND item_id IN ({placeholders})"
+
+        cursor = self._conn.cursor()
+        rows = cursor.execute(sql, (self._name, *result_ids)).fetchall()
+        row_map = {row["item_id"]: row for row in rows}
+
+        results = []
+        for i, doc_id in enumerate(result_ids):
+            row = row_map.get(doc_id)
+            if row:
+                embedding = np.frombuffer(row["item_vector"], dtype=np.float32).tolist()
+                doc = Document(
+                    id=doc_id, embedding=embedding, **json.loads(row["metadata"])
+                )
+                results.append((doc, float(distances[i])))
+        return results
+
+    def match(
+        self, query: str, on_field: str | None = None, top_k: int = 10
+    ) -> list[tuple[Document, float]]:
+        """Performs a full-text search on indexed string fields."""
+        cursor = self._conn.cursor()
+        sql_query = """
+            SELECT t1.item_id, t1.item_vector, t1.metadata, fts.rank
+            FROM beaver_collections AS t1 JOIN (
+                SELECT DISTINCT item_id, rank FROM beaver_fts_index
+                WHERE beaver_fts_index MATCH ? {} ORDER BY rank LIMIT ?
+            ) AS fts ON t1.item_id = fts.item_id
+            WHERE t1.collection = ? ORDER BY fts.rank
+        """
+        params, field_filter_sql = [], ""
+        if on_field:
+            field_filter_sql = "AND field_path = ?"
+            params.extend([query, on_field])
+        else:
+            params.append(query)
+        params.extend([top_k, self._name])
+
+        rows = cursor.execute(
+            sql_query.format(field_filter_sql), tuple(params)
+        ).fetchall()
+        results = []
+        for row in rows:
+            embedding = (
+                np.frombuffer(row["item_vector"], dtype=np.float32).tolist()
+                if row["item_vector"]
+                else None
+            )
+            doc = Document(
+                id=row["item_id"], embedding=embedding, **json.loads(row["metadata"])
+            )
+            results.append((doc, row["rank"]))
+        return results
+
+    def connect(
+        self, source: Document, target: Document, label: str, metadata: dict = None
+    ):
+        """Creates a directed edge between two documents."""
+        if not isinstance(source, Document) or not isinstance(target, Document):
+            raise TypeError("Source and target must be Document objects.")
+        with self._conn:
+            self._conn.execute(
+                "INSERT OR REPLACE INTO beaver_edges (collection, source_item_id, target_item_id, label, metadata) VALUES (?, ?, ?, ?, ?)",
+                (
+                    self._name,
+                    source.id,
+                    target.id,
+                    label,
+                    json.dumps(metadata) if metadata else None,
+                ),
+            )
+
+    def neighbors(self, doc: Document, label: str | None = None) -> list[Document]:
+        """Retrieves the neighboring documents connected to a given document."""
+        sql = "SELECT t1.item_id, t1.item_vector, t1.metadata FROM beaver_collections AS t1 JOIN beaver_edges AS t2 ON t1.item_id = t2.target_item_id AND t1.collection = t2.collection WHERE t2.collection = ? AND t2.source_item_id = ?"
+        params = [self._name, doc.id]
+        if label:
+            sql += " AND t2.label = ?"
+            params.append(label)
+
+        rows = self._conn.cursor().execute(sql, tuple(params)).fetchall()
+        return [
+            Document(
+                id=row["item_id"],
+                embedding=(
+                    np.frombuffer(row["item_vector"], dtype=np.float32).tolist()
+                    if row["item_vector"]
+                    else None
+                ),
+                **json.loads(row["metadata"]),
+            )
+            for row in rows
+        ]
+
+    def walk(
+        self,
+        source: Document,
+        labels: List[str],
+        depth: int,
+        *,
+        direction: Literal[
+            WalkDirection.OUTGOING, WalkDirection.INCOMING
+        ] = WalkDirection.OUTGOING,
+    ) -> List[Document]:
+        """Performs a graph traversal (BFS) from a starting document."""
+        if not isinstance(source, Document):
+            raise TypeError("The starting point must be a Document object.")
+        if depth <= 0:
+            return []
+
+        source_col, target_col = (
+            ("source_item_id", "target_item_id")
+            if direction == WalkDirection.OUTGOING
+            else ("target_item_id", "source_item_id")
+        )
+        sql = f"""
+            WITH RECURSIVE walk_bfs(item_id, current_depth) AS (
+                SELECT ?, 0
+                UNION ALL
+                SELECT edges.{target_col}, bfs.current_depth + 1
+                FROM beaver_edges AS edges JOIN walk_bfs AS bfs ON edges.{source_col} = bfs.item_id
+                WHERE edges.collection = ? AND bfs.current_depth < ? AND edges.label IN ({','.join('?' for _ in labels)})
+            )
+            SELECT DISTINCT t1.item_id, t1.item_vector, t1.metadata
+            FROM beaver_collections AS t1 JOIN walk_bfs AS bfs ON t1.item_id = bfs.item_id
+            WHERE t1.collection = ? AND bfs.current_depth > 0
+        """
+        params = [source.id, self._name, depth] + labels + [self._name]
+
+        rows = self._conn.cursor().execute(sql, tuple(params)).fetchall()
+        return [
+            Document(
+                id=row["item_id"],
+                embedding=(
+                    np.frombuffer(row["item_vector"], dtype=np.float32).tolist()
+                    if row["item_vector"]
+                    else None
+                ),
+                **json.loads(row["metadata"]),
+            )
+            for row in rows
+        ]