opencode-semantic-memory 0.1.0__py3-none-any.whl

opencode_memory/storage/vectors.py

@@ -0,0 +1,199 @@
+"""Vector storage using LanceDB for semantic search."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import lancedb
+
+
+class VectorStorage:
+    """LanceDB-based vector storage for semantic search."""
+
+    def __init__(self, vectors_path: Path, embedding_dim: int = 384):
+        self.vectors_path = vectors_path
+        self.embedding_dim = embedding_dim
+        self._db: lancedb.DBConnection | None = None
+
+    @property
+    def db(self) -> Any:
+        """Lazy-load LanceDB connection."""
+        if self._db is None:
+            import lancedb
+
+            self.vectors_path.mkdir(parents=True, exist_ok=True)
+            self._db = lancedb.connect(str(self.vectors_path))
+            self._init_table()
+        return self._db
+
+    def _init_table(self) -> None:
+        """Initialize the vectors table if it doesn't exist."""
+        existing_tables = self._db.list_tables()
+        table_names = (
+            existing_tables.tables if hasattr(existing_tables, "tables") else existing_tables
+        )
+        if "memories" not in table_names:
+            self._db.create_table(
+                "memories",
+                data=[
+                    {
+                        "id": "__init__",
+                        "memory_id": 0,
+                        "text": "",
+                        "vector": [0.0] * self.embedding_dim,
+                    }
+                ],
+            )
+            self._db["memories"].delete("id = '__init__'")
+
+    def add(self, record_id: str, memory_id: int, text: str, vector: list[float]) -> None:
+        """Add a vector to the store."""
+        table = self.db["memories"]
+        table.add(
+            [
+                {
+                    "id": record_id,
+                    "memory_id": memory_id,
+                    "text": text,
+                    "vector": vector,
+                }
+            ]
+        )
+
+    def search(self, query_vector: list[float], limit: int = 20) -> list[dict[str, Any]]:
+        """Search for similar vectors."""
+        table = self.db["memories"]
+        results = table.search(query_vector).limit(limit).to_list()
+        return results
+
+    def delete(self, record_id: str) -> None:
+        """Delete a vector by ID."""
+        table = self.db["memories"]
+        table.delete(f"id = '{record_id}'")
+
+    def delete_by_memory_id(self, memory_id: int) -> None:
+        """Delete all vectors for a memory."""
+        table = self.db["memories"]
+        table.delete(f"memory_id = {memory_id}")
+
+    def delete_by_prefix(self, id_prefix: str) -> int:
+        """Delete all vectors with IDs starting with the given prefix.
+
+        Uses LanceDB's SQL-like predicate with LIKE for efficient filtering
+        without loading all records into memory.
+        """
+        import re
+
+        # Validate id_prefix to prevent SQL injection - only allow safe characters
+        if not re.match(r"^[a-zA-Z0-9_:/-]+$", id_prefix):
+            raise ValueError(f"Invalid id_prefix: {id_prefix}")
+
+        table = self.db["memories"]
+        try:
+            # Use SQL LIKE predicate - LanceDB supports this
+            # Escape any SQL special characters in the prefix
+            safe_prefix = id_prefix.replace("'", "''").replace("%", "\\%").replace("_", "\\_")
+
+            # Count matching records first (optional, for return value)
+            try:
+                matches = (
+                    table.search([0.0] * self.embedding_dim)
+                    .where(f"id LIKE '{safe_prefix}%'")
+                    .limit(10000)
+                    .to_list()
+                )
+                count = len(matches)
+            except Exception:
+                count = 0
+
+            # Delete using predicate
+            table.delete(f"id LIKE '{safe_prefix}%'")
+            return count
+        except Exception:
+            # Fallback to old method if LIKE not supported
+            try:
+                all_records = table.to_pandas()
+                matching_ids = all_records[all_records["id"].str.startswith(id_prefix)][
+                    "id"
+                ].tolist()
+                for record_id in matching_ids:
+                    # Validate each id before deletion
+                    if re.match(r"^[a-zA-Z0-9_:/-]+$", record_id):
+                        table.delete(f"id = '{record_id}'")
+                return len(matching_ids)
+            except Exception:
+                return 0
+
+    def cleanup_old_versions(self, keep_versions: int = 10) -> dict[str, Any]:
+        """Clean up old LanceDB versions to reclaim disk space.
+
+        Args:
+            keep_versions: Number of recent versions to keep (default 10)
+
+        Returns:
+            Stats about the cleanup operation
+        """
+        from datetime import timedelta
+
+        try:
+            table = self.db.open_table("memories")
+            versions = table.list_versions()
+
+            if len(versions) <= keep_versions:
+                return {
+                    "status": "skipped",
+                    "reason": f"Only {len(versions)} versions exist, keeping all",
+                    "versions_before": len(versions),
+                }
+
+            # Cleanup old versions (keep recent ones)
+            # LanceDB cleanup_old_versions uses time-based threshold
+            # We calculate how old the Nth version is and use that as threshold
+            versions_sorted = sorted(versions, key=lambda v: v["version"], reverse=True)
+            cutoff_version = versions_sorted[keep_versions - 1]
+
+            # Use the timestamp of the cutoff version
+            # Add a small buffer to ensure we keep exactly keep_versions
+            from datetime import UTC, datetime
+
+            cutoff_time = cutoff_version.get("timestamp")
+            if cutoff_time:
+                # Versions older than this will be deleted
+                now = datetime.now(UTC)
+                if cutoff_time.tzinfo is None:
+                    cutoff_time = cutoff_time.replace(tzinfo=UTC)
+                older_than = now - cutoff_time
+                if older_than.total_seconds() > 0:
+                    older_than = timedelta(seconds=max(1, older_than.total_seconds()))
+                else:
+                    older_than = timedelta(days=1)
+            else:
+                older_than = timedelta(days=1)
+
+            stats = table.cleanup_old_versions(older_than=older_than, delete_unverified=True)
+
+            # Also compact files to optimize storage
+            compact_stats = table.compact_files()
+
+            versions_after = len(table.list_versions())
+
+            return {
+                "status": "success",
+                "versions_before": len(versions),
+                "versions_after": versions_after,
+                "versions_removed": len(versions) - versions_after,
+                "bytes_freed": getattr(stats, "bytes_removed", 0),
+                "files_compacted": getattr(compact_stats, "files_removed", 0),
+            }
+        except ImportError:
+            return {
+                "status": "error",
+                "error": "pylance not installed - run: uv pip install pylance",
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "error": str(e),
+            }