powermem 0.1.0__py3-none-any.whl

powermem/storage/base.py

@@ -0,0 +1,109 @@
+"""
+Abstract base class for storage implementations
+
+This module defines the storage interface that all implementations must follow.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, Any, List
+
+from pydantic import BaseModel
+
+
+class OutputData(BaseModel):
+    id: Optional[int]  # memory id (Snowflake ID - 64-bit integer)
+    score: Optional[float]  # distance
+    payload: Optional[Dict]  # metadata
+
+class VectorStoreBase(ABC):
+    """
+    Abstract base class for storage implementations.
+    
+    This class defines the interface that all storage backends must implement.
+    """
+
+    @abstractmethod
+    def create_col(self, name, vector_size, distance):
+        """Create a new collection."""
+        pass
+
+    @abstractmethod
+    def insert(self, vectors, payloads=None, ids=None):
+        """Insert vectors into a collection."""
+        pass
+
+    @abstractmethod
+    def search(self, query, vectors, limit=5, filters=None):
+        """Search for similar vectors."""
+        pass
+
+    @abstractmethod
+    def delete(self, vector_id):
+        """Delete a vector by ID."""
+        pass
+
+    @abstractmethod
+    def update(self, vector_id, vector=None, payload=None):
+        """Update a vector and its payload."""
+        pass
+
+    @abstractmethod
+    def get(self, vector_id):
+        """Retrieve a vector by ID."""
+        pass
+
+    @abstractmethod
+    def list_cols(self):
+        """List all collections."""
+        pass
+
+    @abstractmethod
+    def delete_col(self):
+        """Delete a collection."""
+        pass
+
+    @abstractmethod
+    def col_info(self):
+        """Get information about a collection."""
+        pass
+
+    @abstractmethod
+    def list(self, filters=None, limit=None):
+        """List all memories."""
+        pass
+
+    @abstractmethod
+    def reset(self):
+        """Reset by delete the collection and recreate it."""
+        pass
+
+class GraphStoreBase(ABC):
+    """
+    Abstract base class for graph storage implementations.
+
+    This class defines the interface that all graph storage backends must implement.
+    """
+    @abstractmethod
+    def add(self, data: str, filters: Dict[str, Any]) -> Dict[str, Any]:
+        """Add data to the graph."""
+        pass
+
+    @abstractmethod
+    def search(self, query: str, filters: Dict[str, Any], limit: int = 10) -> List[Dict[str, Any]]:
+        """Search for memories."""
+        pass
+
+    @abstractmethod
+    def delete_all(self, filters: Dict[str, Any]) -> None:
+        """Delete all graph data for the given filters."""
+        pass
+
+    @abstractmethod
+    def get_all(self, filters: Dict[str, Any], limit: int = 100) -> List[Dict[str, str]]:
+        """Retrieve all nodes and relationships from the graph database."""
+        pass
+
+    @abstractmethod
+    def reset(self) -> None:
+        """Reset the graph by clearing all nodes and relationships."""
+        pass

powermem/storage/config/base.py

@@ -0,0 +1,13 @@
+from abc import ABC
+from typing import Any, Dict
+
+from pydantic import BaseModel, model_validator
+
+
+class BaseVectorStoreConfig(BaseModel, ABC):
+    """
+    Base configuration class for all vector store providers.
+    
+    This class provides common validation logic that is shared
+    across all vector store implementations.
+    """

powermem/storage/config/oceanbase.py

@@ -0,0 +1,58 @@
+from typing import Any, ClassVar, Dict, Optional
+
+from pydantic import Field, model_validator
+
+from powermem.storage.config.base import BaseVectorStoreConfig
+
+
+class OceanBaseConfig(BaseVectorStoreConfig):
+    try:
+        from pyobvector import ObVecClient
+    except ImportError:
+        raise ImportError("The 'pyobvector' library is required. Please install it using 'pip install pyobvector'.")
+    ObVecClient: ClassVar[type] = ObVecClient
+
+    collection_name: str = Field("power_mem", description="Default name for the collection")
+
+    # Connection parameters
+    host: str = Field("localhost", description="OceanBase server host")
+    port: str = Field("2881", description="OceanBase server port")
+    user: str = Field("root@test", description="OceanBase username")
+    password: str = Field("", description="OceanBase password")
+    db_name: str = Field("test", description="OceanBase database name")
+
+    # Vector index parameters
+    index_type: str = Field("HNSW", description="Type of vector index (HNSW, IVF, FLAT, etc.)")
+    vidx_metric_type: str = Field("l2", description="Distance metric (l2, inner_product, cosine)")
+    embedding_model_dims: Optional[int] = Field(None, description="Dimension of vectors")
+
+    # Advanced parameters
+    vidx_algo_params: Optional[Dict[str, Any]] = Field(None, description="Index algorithm parameters")
+    normalize: bool = Field(False, description="Whether to normalize vectors")
+    include_sparse: bool = Field(False, description="Whether to include sparse vector support")
+    hybrid_search: bool = Field(True, description="Whether to enable hybrid search")
+    auto_configure_vector_index: bool = Field(True,
+                                              description="Whether to automatically configure vector index settings")
+
+    # Fulltext search parameters
+    fulltext_parser: str = Field("ik", description="Fulltext parser type (ik, ngram, ngram2, beng, space)")
+
+    # Field names
+    primary_field: str = Field("id", description="Primary key field name")
+    vector_field: str = Field("embedding", description="Vector field name")
+    text_field: str = Field("document", description="Text field name")
+    metadata_field: str = Field("metadata", description="Metadata field name")
+    vidx_name: str = Field("vidx", description="Vector index name")
+
+    vector_weight: float = Field(0.5, description="Weight for vector search")
+    fts_weight: float = Field(0.5, description="Weight for fulltext search")
+
+    model_config = {
+        "arbitrary_types_allowed": True,
+    }
+
+
+
+class OceanBaseGraphConfig(OceanBaseConfig):
+    # Graph search parameters
+    max_hops: int = Field(3, description="Maximum number of hops for multi-hop graph search")

powermem/storage/config/pgvector.py

@@ -0,0 +1,52 @@
+from typing import Any, Optional
+
+from pydantic import Field, model_validator
+
+from powermem.storage.config.base import BaseVectorStoreConfig
+
+
+class PGVectorConfig(BaseVectorStoreConfig):
+    dbname: str = Field("postgres", description="Default name for the database")
+    collection_name: str = Field("power_mem", description="Default name for the collection")
+    embedding_model_dims: Optional[int] = Field(1536, description="Dimensions of the embedding model")
+    user: Optional[str] = Field(None, description="Database user")
+    password: Optional[str] = Field(None, description="Database password")
+    host: Optional[str] = Field(None, description="Database host. Default is localhost")
+    port: Optional[int] = Field(None, description="Database port. Default is 1536")
+    diskann: Optional[bool] = Field(False, description="Use diskann for approximate nearest neighbors search")
+    hnsw: Optional[bool] = Field(True, description="Use hnsw for faster search")
+    minconn: Optional[int] = Field(1, description="Minimum number of connections in the pool")
+    maxconn: Optional[int] = Field(5, description="Maximum number of connections in the pool")
+    # New SSL and connection options
+    sslmode: Optional[str] = Field(None,
+                                   description="SSL mode for PostgreSQL connection (e.g., 'require', 'prefer', 'disable')")
+    connection_string: Optional[str] = Field(None,
+                                             description="PostgreSQL connection string (overrides individual connection parameters)")
+    connection_pool: Optional[Any] = Field(None,
+                                           description="psycopg connection pool object (overrides connection string and individual parameters)")
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_auth_and_connection(cls, values):
+        # If connection_pool is provided, skip validation of individual connection parameters
+        if values.get("connection_pool") is not None:
+            return values
+
+        # If connection_string is provided, skip validation of individual connection parameters
+        if values.get("connection_string") is not None:
+            return values
+
+        # Otherwise, validate individual connection parameters
+        user, password = values.get("user"), values.get("password")
+        host, port = values.get("host"), values.get("port")
+        
+        # Only validate if user explicitly provided values (not using defaults)
+        if user is not None or password is not None:
+            if not user or not password:
+                raise ValueError("Both 'user' and 'password' must be provided when not using connection_string.")
+        
+        if host is not None or port is not None:
+            if not host or not port:
+                raise ValueError("Both 'host' and 'port' must be provided when not using connection_string.")
+        
+        return values

powermem/storage/config/sqlite.py

@@ -0,0 +1,27 @@
+from typing import Optional
+
+from pydantic import Field
+
+from powermem.storage.config.base import BaseVectorStoreConfig
+
+
+class SQLiteConfig(BaseVectorStoreConfig):
+    """Configuration for SQLite vector store."""
+    
+    database_path: str = Field(
+        default="./data/powermem_dev.db",
+        description="Path to SQLite database file"
+    )
+    collection_name: str = Field(
+        default="memories",
+        description="Name of the collection/table"
+    )
+    enable_wal: bool = Field(
+        default=True,
+        description="Enable Write-Ahead Logging for better concurrency"
+    )
+    timeout: int = Field(
+        default=30,
+        description="Connection timeout in seconds"
+    )
+

powermem/storage/configs.py

@@ -0,0 +1,159 @@
+"""
+Storage configuration management
+
+This module handles storage configuration and validation.
+"""
+
+from typing import Dict, Optional, Union
+
+from pydantic import BaseModel, Field, model_validator
+
+from powermem.integrations.llm.configs import LLMConfig
+from powermem.storage.config.oceanbase import OceanBaseGraphConfig
+
+
+class VectorStoreConfig(BaseModel):
+    provider: str = Field(
+        description="Provider of the vector store (e.g., 'oceanbase', 'pgvector')",
+        default="oceanbase",
+    )
+    config: Optional[Dict] = Field(
+        description="Configuration for the specific vector store",
+        default=None
+    )
+
+    _provider_configs: Dict[str, str] = {
+        "oceanbase": "OceanBaseConfig",
+        "pgvector": "PGVectorConfig",
+        "sqlite": "SQLiteConfig",
+    }
+
+    @model_validator(mode="after")
+    def validate_config(self) -> "VectorStoreConfig":
+        """
+        Validate the configuration without converting to provider-specific config class.
+        The conversion is handled by VectorStoreFactory.create() when needed.
+        """
+        provider = self.provider
+        config = self.config
+
+        if provider is not None and provider == "postgres":
+            provider = "pgvector"
+
+        if provider not in self._provider_configs:
+            raise ValueError(f"Unsupported vector store provider: {provider}")
+
+        if config is None:
+            self.config = {}
+            return self
+
+        if not isinstance(config, dict):
+            raise ValueError(f"Config must be a dictionary, got {type(config)}")
+
+        # Handle connection_args for backward compatibility
+        # If connection_args exists, flatten it into the main config
+        if "connection_args" in config:
+            connection_args = config.pop("connection_args")
+            if isinstance(connection_args, dict):
+                # Merge connection_args into config (connection_args values take precedence)
+                for key, value in connection_args.items():
+                    if key not in config:
+                        # Convert port to string if it's an int (for OceanBase compatibility)
+                        if key == "port" and isinstance(value, int):
+                            config[key] = str(value)
+                        else:
+                            config[key] = value
+                self.config = config
+
+        # Convert port to string if it's an int (for OceanBase compatibility)
+        # This handles both direct port field and port from connection_args
+        if "port" in config and isinstance(config["port"], int):
+            config["port"] = str(config["port"])
+            self.config = config
+
+        # Validate config by attempting to create provider-specific config instance
+        # This ensures the config has valid fields, but we don't store the converted object
+        module = __import__(
+            f"powermem.storage.config.{provider}",
+            fromlist=[self._provider_configs[provider]],
+        )
+        config_class = getattr(module, self._provider_configs[provider])
+
+        # Add default path if needed
+        if "path" not in config and "path" in config_class.__annotations__:
+            config["path"] = f"/tmp/{provider}"
+            self.config = config
+
+        # Validate by creating instance (throws error if invalid)
+        try:
+            config_class(**config)
+        except Exception as e:
+            raise ValueError(f"Invalid configuration for {provider}: {e}")
+
+        # Keep config as dict, don't convert to config_class instance
+        return self
+
+class GraphStoreConfig(BaseModel):
+    enabled: bool = Field(
+        description="Whether to enable graph store",
+        default=False,
+    )
+    provider: str = Field(
+        description="Provider of the data store (e.g., 'oceanbase')",
+        default="oceanbase",
+    )
+    config: Optional[Union[Dict, OceanBaseGraphConfig]] = Field(
+        description="Configuration for the specific data store",
+        default=None
+    )
+    llm: Optional[LLMConfig] = Field(
+        description="LLM configuration for querying the graph store",
+        default=None
+    )
+    custom_prompt: Optional[str] = Field(
+        description="Custom prompt to fetch entities from the given text",
+        default=None
+    )
+    custom_extract_relations_prompt: Optional[str] = Field(
+        description="Custom prompt for extracting relations from text",
+        default=None
+    )
+    custom_update_graph_prompt: Optional[str] = Field(
+        description="Custom prompt for updating graph memories",
+        default=None
+    )
+    custom_delete_relations_prompt: Optional[str] = Field(
+        description="Custom prompt for deleting relations",
+        default=None
+    )
+
+    @model_validator(mode="after")
+    def validate_config(self) -> "GraphStoreConfig":
+        """
+        Validate the configuration without converting to provider-specific config class.
+        Keep config as dict for consistency.
+        """
+        if self.config is None:
+            self.config = {}
+            return self
+
+        # If config is a Pydantic BaseModel instance, convert it to dict
+        if isinstance(self.config, BaseModel):
+            self.config = self.config.model_dump()
+
+        if not isinstance(self.config, dict):
+            raise ValueError(f"Config must be a dictionary or BaseModel instance, got {type(self.config)}")
+
+
+        # Validate config based on provider
+        provider = self.provider
+        if provider == "oceanbase":
+            try:
+                OceanBaseGraphConfig(**self.config)
+            except Exception as e:
+                raise ValueError(f"Invalid configuration for {provider}: {e}")
+        else:
+            raise ValueError(f"Unsupported graph store provider: {provider}")
+
+        # Keep config as dict, don't convert
+        return self

powermem/storage/factory.py

@@ -0,0 +1,59 @@
+"""
+Storage factory for creating storage instances
+
+This module provides a factory for creating different storage backends.
+"""
+
+import importlib
+
+
+def load_class(class_type):
+    module_path, class_name = class_type.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)
+
+class VectorStoreFactory:
+    provider_to_class = {
+        "oceanbase": "powermem.storage.oceanbase.oceanbase.OceanBaseVectorStore",
+        "sqlite": "powermem.storage.sqlite.sqlite_vector_store.SQLiteVectorStore",
+        "pgvector": "powermem.storage.pgvector.pgvector.PGVectorStore",
+        "postgres": "powermem.storage.pgvector.pgvector.PGVectorStore",  # Alias for pgvector
+    }
+
+    @classmethod
+    def create(cls, provider_name, config):
+        class_type = cls.provider_to_class.get(provider_name)
+        if class_type:
+            if not isinstance(config, dict):
+                config = config.model_dump()
+            vector_store_instance = load_class(class_type)
+            return vector_store_instance(**config)
+        else:
+            raise ValueError(f"Unsupported VectorStore provider: {provider_name}")
+
+    @classmethod
+    def reset(cls, instance):
+        instance.reset()
+        return instance
+
+
+class GraphStoreFactory:
+    """
+    Factory for creating MemoryGraph instances for different graph store providers.
+    Usage: GraphStoreFactory.create(provider_name, config)
+    """
+
+    provider_to_class = {
+        "oceanbase": "powermem.storage.oceanbase.oceanbase_graph.MemoryGraph",
+        "default": "powermem.storage.oceanbase.oceanbase_graph.MemoryGraph",
+    }
+
+    @classmethod
+    def create(cls, provider_name, config):
+        class_type = cls.provider_to_class.get(provider_name, cls.provider_to_class["default"])
+        try:
+            GraphClass = load_class(class_type)
+        except (ImportError, AttributeError) as e:
+            raise ImportError(f"Could not import MemoryGraph for provider '{provider_name}': {e}")
+        return GraphClass(config)
+