powermem 0.1.0__py3-none-any.whl

powermem/storage/sqlite/sqlite_vector_store.py

@@ -0,0 +1,311 @@
+"""
+SQLite vector store implementation
+
+This module provides a simple SQLite-based vector store for development and testing.
+"""
+
+import json
+import logging
+import os
+import sqlite3
+import threading
+from typing import Any, Dict, List, Optional
+
+from powermem.storage.base import VectorStoreBase, OutputData
+from powermem.utils.utils import generate_snowflake_id
+
+logger = logging.getLogger(__name__)
+
+
+class SQLiteVectorStore(VectorStoreBase):
+    """Simple SQLite-based vector store implementation."""
+    
+    def __init__(self, database_path: str = ":memory:", collection_name: str = "memories", **kwargs):
+        """
+        Initialize SQLite vector store.
+        
+        Args:
+            database_path: Path to SQLite database file
+            collection_name: Name of the collection/table
+        """
+        self.db_path = database_path
+        self.collection_name = collection_name
+        self.connection = None
+        self._lock = threading.Lock()
+        
+        # Create directory if database path is not in-memory and directory doesn't exist
+        if database_path != ":memory:":
+            db_dir = os.path.dirname(os.path.abspath(database_path))
+            if db_dir and not os.path.exists(db_dir):
+                try:
+                    os.makedirs(db_dir, exist_ok=True)
+                    logger.info(f"Created database directory: {db_dir}")
+                except OSError as e:
+                    logger.error(f"Failed to create database directory {db_dir}: {e}")
+                    raise
+        
+        # Connect to database
+        try:
+            self.connection = sqlite3.connect(database_path, check_same_thread=False)
+        except Exception as e:
+            logger.error(f"Failed to connect to SQLite database at {database_path}: {e}")
+            raise
+        
+        # Create the table
+        self.create_col()
+        
+        logger.info(f"SQLiteVectorStore initialized with db_path: {database_path}")
+    
+    def create_col(self, name=None, vector_size=None, distance=None) -> None:
+        """Create a new collection (table in SQLite)."""
+        table_name = name or self.collection_name
+        
+        with self._lock:
+            self.connection.execute(f"""
+                CREATE TABLE IF NOT EXISTS {table_name} (
+                    id INTEGER PRIMARY KEY,
+                    vector TEXT,  -- Store as JSON string
+                    payload TEXT,  -- Store as JSON string
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
+            self.connection.commit()
+    
+    def insert(self, vectors: List[List[float]], payloads=None, ids=None) -> List[int]:
+        """
+        Insert vectors into the collection.
+        
+        Args:
+            vectors: List of vectors to insert
+            payloads: List of payload dictionaries
+            ids: Deprecated parameter (ignored), IDs are now generated using Snowflake algorithm
+            
+        Returns:
+            List[int]: List of generated Snowflake IDs
+        """
+        if not vectors:
+            return []
+        
+        if payloads is None:
+            payloads = [{} for _ in vectors]
+        
+        # Generate Snowflake IDs for each vector
+        generated_ids = [generate_snowflake_id() for _ in range(len(vectors))]
+        
+        with self._lock:
+            for vector, payload, vector_id in zip(vectors, payloads, generated_ids):
+                self.connection.execute(f"""
+                    INSERT INTO {self.collection_name} 
+                    (id, vector, payload) VALUES (?, ?, ?)
+                """, (vector_id, json.dumps(vector), json.dumps(payload)))
+            
+            self.connection.commit()
+        
+        return generated_ids
+    
+    def search(self, query: str, vectors: List[List[float]] = None, limit: int = 5, filters=None) -> List[OutputData]:
+        """Search for similar vectors using simple cosine similarity."""
+        results = []
+        
+        # Extract query vector from vectors parameter (OceanBase format)
+        if vectors and len(vectors) > 0:
+            query_vector = vectors[0]
+        else:
+            # Fallback for backward compatibility
+            query_vector = query if isinstance(query, list) else [0.1] * 10
+        
+        # Build query with filters
+        query_sql = f"SELECT id, vector, payload FROM {self.collection_name}"
+        query_params = []
+        
+        # Apply filters if provided
+        if filters:
+            conditions = []
+            for key, value in filters.items():
+                # Filter by JSON field in payload
+                conditions.append(f"json_extract(payload, '$.{key}') = ?")
+                query_params.append(value)
+            
+            if conditions:
+                query_sql += " WHERE " + " AND ".join(conditions)
+                logger.info(f"SQLite search with filters: {query_sql}, params: {query_params}")
+            else:
+                logger.debug("SQLite search: filters provided but empty after processing")
+        else:
+            logger.debug("SQLite search: no filters provided")
+        
+        with self._lock:
+            if query_params:
+                cursor = self.connection.execute(query_sql, query_params)
+            else:
+                cursor = self.connection.execute(query_sql)
+            
+            row_count = 0
+            for row in cursor.fetchall():
+                row_count += 1
+                vector_id, vector_str, payload_str = row
+                vector = json.loads(vector_str)
+                payload = json.loads(payload_str)
+                
+                # Calculate cosine similarity
+                similarity = self._cosine_similarity(query_vector, vector)
+                
+                results.append(OutputData(
+                    id=vector_id,
+                    score=similarity,
+                    payload=payload
+                ))
+        
+        # Sort by similarity (descending) and return top results
+        results.sort(key=lambda x: x.score, reverse=True)
+        return results[:limit]
+    
+    def delete(self, vector_id: int) -> None:
+        """Delete a vector by ID."""
+        with self._lock:
+            self.connection.execute(f"""
+                DELETE FROM {self.collection_name} WHERE id = ?
+            """, (vector_id,))
+            self.connection.commit()
+    
+    def update(self, vector_id: int, vector=None, payload=None) -> None:
+        """Update a vector and its payload."""
+        updates = []
+        values = []
+        
+        if vector is not None:
+            updates.append("vector = ?")
+            values.append(json.dumps(vector))
+        
+        if payload is not None:
+            updates.append("payload = ?")
+            values.append(json.dumps(payload))
+        
+        if updates:
+            values.append(vector_id)
+            with self._lock:
+                self.connection.execute(f"""
+                    UPDATE {self.collection_name} 
+                    SET {', '.join(updates)}
+                    WHERE id = ?
+                """, values)
+                self.connection.commit()
+    
+    def get(self, vector_id: int) -> Optional[OutputData]:
+        """Retrieve a vector by ID."""
+        with self._lock:
+            cursor = self.connection.execute(f"""
+                SELECT id, vector, payload FROM {self.collection_name} WHERE id = ?
+            """, (vector_id,))
+            
+            row = cursor.fetchone()
+            if row:
+                vector_id, vector_str, payload_str = row
+                vector = json.loads(vector_str)
+                payload = json.loads(payload_str)
+                
+                return OutputData(
+                    id=vector_id,
+                    score=1.0,  # Exact match
+                    payload=payload
+                )
+        
+        return None
+    
+    def list_cols(self) -> List[str]:
+        """List all collections (tables)."""
+        with self._lock:
+            cursor = self.connection.execute("""
+                SELECT name FROM sqlite_master WHERE type='table'
+            """)
+            return [row[0] for row in cursor.fetchall()]
+    
+    def delete_col(self) -> None:
+        """Delete the collection (table)."""
+        with self._lock:
+            self.connection.execute(f"DROP TABLE IF EXISTS {self.collection_name}")
+            self.connection.commit()
+    
+    def col_info(self) -> Dict[str, Any]:
+        """Get information about the collection."""
+        with self._lock:
+            cursor = self.connection.execute(f"""
+                SELECT COUNT(*) FROM {self.collection_name}
+            """)
+            count = cursor.fetchone()[0]
+            
+            return {
+                "name": self.collection_name,
+                "count": count,
+                "db_path": self.db_path
+            }
+    
+    def list(self, filters=None, limit=None) -> List[OutputData]:
+        """List all memories with optional filtering."""
+        query = f"SELECT id, vector, payload FROM {self.collection_name}"
+        query_params = []
+        
+        # Apply filters if provided
+        if filters:
+            conditions = []
+            for key, value in filters.items():
+                # Filter by JSON field in payload
+                conditions.append(f"json_extract(payload, '$.{key}') = ?")
+                query_params.append(value)
+            
+            if conditions:
+                query += " WHERE " + " AND ".join(conditions)
+        
+        if limit:
+            query += f" LIMIT {limit}"
+        
+        results = []
+        with self._lock:
+            if query_params:
+                cursor = self.connection.execute(query, query_params)
+            else:
+                cursor = self.connection.execute(query)
+            
+            for row in cursor.fetchall():
+                vector_id, vector_str, payload_str = row
+                vector = json.loads(vector_str)
+                payload = json.loads(payload_str)
+                
+                results.append(OutputData(
+                    id=vector_id,
+                    score=1.0,
+                    payload=payload
+                ))
+        
+        return results
+    
+    def reset(self) -> None:
+        """Reset by deleting and recreating the collection."""
+        self.delete_col()
+        self.create_col()
+    
+    def _cosine_similarity(self, vec1: List[float], vec2: List[float]) -> float:
+        """Calculate cosine similarity between two vectors."""
+        if len(vec1) != len(vec2):
+            return 0.0
+        
+        dot_product = sum(a * b for a, b in zip(vec1, vec2))
+        magnitude1 = sum(a * a for a in vec1) ** 0.5
+        magnitude2 = sum(b * b for b in vec2) ** 0.5
+        
+        if magnitude1 == 0 or magnitude2 == 0:
+            return 0.0
+        
+        return dot_product / (magnitude1 * magnitude2)
+    
+    def close(self) -> None:
+        """Close the database connection."""
+        if hasattr(self, 'connection') and self.connection:
+            self.connection.close()
+            self.connection = None
+    
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass  # Ignore errors during cleanup

powermem/utils/__init__.py

@@ -0,0 +1,35 @@
+"""
+Utility functions and classes
+
+This module provides utility functions and helper classes.
+"""
+
+from .utils import (
+    generate_memory_id,
+    validate_memory_data,
+    sanitize_content,
+    format_memory_for_display,
+    merge_memories,
+    calculate_similarity,
+    extract_keywords,
+    format_timestamp,
+    parse_timestamp,
+    extract_json,
+    load_config_from_env,
+    convert_config_object_to_dict,
+)
+
+__all__ = [
+    "generate_memory_id",
+    "validate_memory_data",
+    "sanitize_content",
+    "format_memory_for_display",
+    "merge_memories",
+    "calculate_similarity",
+    "extract_keywords",
+    "format_timestamp",
+    "parse_timestamp",
+    "extract_json",
+    "load_config_from_env",
+    "convert_config_object_to_dict",
+]