powermem 0.1.0__py3-none-any.whl

powermem/core/setup.py

@@ -0,0 +1,162 @@
+"""
+Setup utilities for powermem
+
+This module provides setup functions compatible with initialization,
+"""
+
+import json
+import os
+import uuid
+from typing import Optional, Dict, Any
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Set up the directory path
+VECTOR_ID = str(uuid.uuid4())
+home_dir = os.path.expanduser("~")
+powermem_dir = os.environ.get("POWERMEM_DIR") or os.path.join(home_dir, ".powermem")
+os.makedirs(powermem_dir, exist_ok=True)
+
+
+def setup_config():
+    """Setup configuration file."""
+    config_path = os.path.join(powermem_dir, "config.json")
+    if not os.path.exists(config_path):
+        user_id = str(uuid.uuid4())
+        config = {"user_id": user_id}
+        with open(config_path, "w") as config_file:
+            json.dump(config, config_file, indent=4)
+
+
+def get_user_id() -> str:
+    """Get or create user ID."""
+    config_path = os.path.join(powermem_dir, "config.json")
+    if not os.path.exists(config_path):
+        setup_config()
+    
+    try:
+        with open(config_path, "r") as config_file:
+            config = json.load(config_file)
+            user_id = config.get("user_id")
+            return user_id
+    except Exception:
+        return "anonymous_user"
+
+
+def from_config(config: Optional[Dict[str, Any]] = None, **kwargs):
+    """
+    Create Memory instance from configuration.
+    
+    powermem now uses field names natively: 'embedder' and 'vector_store'.
+    
+    Args:
+        config: Configuration dictionary
+               - llm: LLM provider configuration
+               - embedder: Embedder configuration
+               - vector_store: Vector store config (uses OceanBase)
+        **kwargs: Additional parameters
+    
+    Returns:
+        Memory instance
+        
+    Example:
+        ```python
+        from powermem import from_config
+        
+        memory = from_config({
+            "llm": {"provider": "openai", "config": {"api_key": "..."}},
+            "embedder": {"provider": "openai", "config": {"api_key": "..."}},
+            "vector_store": {"provider": "oceanbase", "config": {...}},
+        })
+        ```
+    """
+    from ..core.memory import Memory
+    
+    if config is None:
+        # Use auto config from environment
+        from ..config_loader import auto_config
+        config = auto_config()
+    
+    converted_config = _convert_legacy_to_mem_config(config)
+    
+    return Memory(config=converted_config, **kwargs)
+
+
+def _convert_legacy_to_mem_config(config: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Convert legacy powermem config format.
+    
+    Now powermem uses field names natively, so we only convert legacy format.
+    
+    Args:
+        config: Legacy powermem configuration dictionary
+        
+    Returns:
+        configuration dictionary
+    """
+    if "embedder" in config or "vector_store" in config:
+        return config
+    
+    converted = {}
+    
+    # LLM stays the same
+    if "llm" in config:
+        converted["llm"] = config["llm"]
+    
+    # Convert embedding to embedder
+    if "embedding" in config:
+        converted["embedder"] = config["embedding"]
+    
+    # Convert database to vector_store
+    if "database" in config:
+        db_config = config["database"]
+        converted["vector_store"] = {
+            "provider": db_config.get("provider", "oceanbase"),
+            "config": db_config.get("config", {})
+        }
+    else:
+        converted["vector_store"] = {
+            "provider": "oceanbase",
+            "config": {}
+        }
+    
+    return converted
+
+
+def get_or_create_user_id(vector_store=None) -> str:
+    """
+    Store user_id in vector store and return it.
+    
+    Args:
+        vector_store: Optional vector store instance
+        
+    Returns:
+        User ID
+    """
+    user_id = get_user_id()
+    
+    if vector_store is None:
+        return user_id
+    
+    # Try to get existing user_id from vector store
+    try:
+        existing = vector_store.get(vector_id=user_id)
+        if existing and hasattr(existing, "payload") and existing.payload and "user_id" in existing.payload:
+            return existing.payload["user_id"]
+    except Exception:
+        pass
+    
+    # If we get here, we need to insert the user_id
+    try:
+        dims = getattr(vector_store, "embedding_model_dims", 1536)
+        vector_store.insert(
+            vectors=[[0.1] * dims], 
+            payloads=[{"user_id": user_id, "type": "user_identity"}], 
+            ids=[user_id]
+        )
+    except Exception:
+        pass
+    
+    return user_id
+

powermem/core/telemetry.py

@@ -0,0 +1,215 @@
+"""
+Telemetry management for memory operations
+
+This module handles telemetry data collection and reporting.
+"""
+
+import logging
+import json
+import time
+from typing import Any, Dict, Optional
+from datetime import datetime
+import httpx
+
+logger = logging.getLogger(__name__)
+
+
+class TelemetryManager:
+    """
+    Manages telemetry data collection and reporting.
+    """
+    
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        """
+        Initialize telemetry manager.
+        
+        Args:
+            config: Configuration dictionary
+        """
+        self.config = config or {}
+        self.enabled = self.config.get("enable_telemetry", False)
+        self.endpoint = self.config.get("telemetry_endpoint", "https://telemetry.powermem.ai")
+        self.api_key = self.config.get("telemetry_api_key")
+        self.batch_size = self.config.get("telemetry_batch_size", 100)
+        self.flush_interval = self.config.get("telemetry_flush_interval", 30)
+        
+        self.events = []
+        self.last_flush = time.time()
+        
+        logger.info(f"TelemetryManager initialized - enabled: {self.enabled}")
+    
+    def capture_event(
+        self,
+        event_name: str,
+        properties: Optional[Dict[str, Any]] = None,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+    ) -> None:
+        """
+        Capture a telemetry event.
+        
+        Args:
+            event_name: Name of the event
+            properties: Event properties
+            user_id: User ID associated with the event
+            agent_id: Agent ID associated with the event
+        """
+        if not self.enabled:
+            return
+        
+        try:
+            event = {
+                "event_name": event_name,
+                "properties": properties or {},
+                "user_id": user_id,
+                "agent_id": agent_id,
+                "timestamp": datetime.utcnow().isoformat(),
+                "version": "0.1.0",
+            }
+            
+            self.events.append(event)
+            
+            # Flush if batch size reached
+            if len(self.events) >= self.batch_size:
+                self._flush_events()
+            
+        except Exception as e:
+            logger.error(f"Failed to capture telemetry event: {e}")
+    
+    def _flush_events(self) -> None:
+        """Flush events to the telemetry endpoint."""
+        if not self.events or not self.enabled:
+            return
+        
+        try:
+            if self.api_key:
+                headers = {
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                }
+            else:
+                headers = {"Content-Type": "application/json"}
+            
+            payload = {
+                "events": self.events.copy(),
+                "timestamp": datetime.utcnow().isoformat(),
+            }
+            
+            # Send events asynchronously to avoid blocking
+            self._send_events_async(payload, headers)
+            
+            # Clear events after sending
+            self.events.clear()
+            self.last_flush = time.time()
+            
+        except Exception as e:
+            logger.error(f"Failed to flush telemetry events: {e}")
+    
+    def _send_events_async(self, payload: Dict[str, Any], headers: Dict[str, str]) -> None:
+        """Send events asynchronously."""
+        try:
+            # Try to get current event loop
+            import asyncio
+            try:
+                loop = asyncio.get_running_loop()
+                # If we're in an async context, schedule the task
+                asyncio.create_task(self._send_request(payload, headers))
+            except RuntimeError:
+                # No running event loop, use httpx in sync mode for now
+                import httpx
+                try:
+                    with httpx.Client(timeout=10.0) as client:
+                        response = client.post(
+                            f"{self.endpoint}/events",
+                            json=payload,
+                            headers=headers,
+                            timeout=10.0
+                        )
+                        response.raise_for_status()
+                except Exception as sync_e:
+                    logger.debug(f"Failed to send telemetry events synchronously: {sync_e}")
+            
+        except Exception as e:
+            logger.debug(f"Failed to send telemetry events: {e}")
+    
+    async def _send_request(self, payload: Dict[str, Any], headers: Dict[str, str]) -> None:
+        """Helper method to send HTTP request asynchronously."""
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.endpoint}/events",
+                json=payload,
+                headers=headers,
+                timeout=10.0
+            )
+            response.raise_for_status()
+    
+    def flush(self) -> None:
+        """Manually flush all pending events."""
+        self._flush_events()
+    
+    def set_user_properties(self, user_id: str, properties: Dict[str, Any]) -> None:
+        """
+        Set user properties for telemetry.
+        
+        Args:
+            user_id: User ID
+            properties: User properties
+        """
+        if not self.enabled:
+            return
+        
+        try:
+            event = {
+                "event_name": "user_properties",
+                "properties": properties,
+                "user_id": user_id,
+                "timestamp": datetime.utcnow().isoformat(),
+                "version": "0.1.0",
+            }
+            
+            self.events.append(event)
+            
+        except Exception as e:
+            logger.error(f"Failed to set user properties: {e}")
+    
+    def track_performance(self, operation: str, duration: float, metadata: Optional[Dict[str, Any]] = None) -> None:
+        """
+        Track performance metrics.
+        
+        Args:
+            operation: Operation name
+            duration: Duration in seconds
+            metadata: Additional metadata
+        """
+        if not self.enabled:
+            return
+        
+        self.capture_event(
+            "performance_metric",
+            {
+                "operation": operation,
+                "duration": duration,
+                "metadata": metadata or {},
+            }
+        )
+    
+    def track_error(self, error_type: str, error_message: str, context: Optional[Dict[str, Any]] = None) -> None:
+        """
+        Track error events.
+        
+        Args:
+            error_type: Type of error
+            error_message: Error message
+            context: Additional context
+        """
+        if not self.enabled:
+            return
+        
+        self.capture_event(
+            "error",
+            {
+                "error_type": error_type,
+                "error_message": error_message,
+                "context": context or {},
+            }
+        )

powermem/integrations/__init__.py

@@ -0,0 +1,17 @@
+"""
+Integration layer for external services
+
+This module provides integrations with LLMs, embeddings, rerank, and other services.
+"""
+
+from .llm.factory import LLMFactory
+from .embeddings.factory import EmbedderFactory
+from .rerank.factory import RerankFactory
+from .rerank.configs import RerankConfig
+
+__all__ = [
+    "LLMFactory",
+    "EmbedderFactory",
+    "RerankFactory",
+    "RerankConfig",
+]

powermem/integrations/embeddings/__init__.py

@@ -0,0 +1,13 @@
+"""
+Embeddings integration module
+
+This module provides embeddings integrations and factory.
+"""
+
+from .factory import EmbedderFactory
+
+EmbedderFactory = EmbedderFactory
+
+__all__ = [
+    "EmbedderFactory",
+]

powermem/integrations/embeddings/aws_bedrock.py

@@ -0,0 +1,100 @@
+import json
+import os
+from typing import Literal, Optional
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+try:
+    import boto3
+except ImportError:
+    raise ImportError("The 'boto3' library is required. Please install it using 'pip install boto3'.")
+
+import numpy as np
+
+
+class AWSBedrockEmbedding(EmbeddingBase):
+    """AWS Bedrock embedding implementation.
+
+    This class uses AWS Bedrock's embedding models.
+    """
+
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "amazon.titan-embed-text-v1"
+
+        # Get AWS config from environment variables or use defaults
+        aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID", "")
+        aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY", "")
+        aws_session_token = os.environ.get("AWS_SESSION_TOKEN", "")
+
+        # Check if AWS config is provided in the config
+        if hasattr(self.config, "aws_access_key_id"):
+            aws_access_key = self.config.aws_access_key_id
+        if hasattr(self.config, "aws_secret_access_key"):
+            aws_secret_key = self.config.aws_secret_access_key
+        
+        # AWS region is always set in config - see BaseEmbedderConfig
+        aws_region = self.config.aws_region or "us-west-2"
+
+        self.client = boto3.client(
+            "bedrock-runtime",
+            region_name=aws_region,
+            aws_access_key_id=aws_access_key if aws_access_key else None,
+            aws_secret_access_key=aws_secret_key if aws_secret_key else None,
+            aws_session_token=aws_session_token if aws_session_token else None,
+        )
+
+    def _normalize_vector(self, embeddings):
+        """Normalize the embedding to a unit vector."""
+        emb = np.array(embeddings)
+        norm_emb = emb / np.linalg.norm(emb)
+        return norm_emb.tolist()
+
+    def _get_embedding(self, text):
+        """Call out to Bedrock embedding endpoint."""
+
+        # Format input body based on the provider
+        provider = self.config.model.split(".")[0]
+        input_body = {}
+
+        if provider == "cohere":
+            input_body["input_type"] = "search_document"
+            input_body["texts"] = [text]
+        else:
+            # Amazon and other providers
+            input_body["inputText"] = text
+
+        body = json.dumps(input_body)
+
+        try:
+            response = self.client.invoke_model(
+                body=body,
+                modelId=self.config.model,
+                accept="application/json",
+                contentType="application/json",
+            )
+
+            response_body = json.loads(response.get("body").read())
+
+            if provider == "cohere":
+                embeddings = response_body.get("embeddings")[0]
+            else:
+                embeddings = response_body.get("embedding")
+
+            return embeddings
+        except Exception as e:
+            raise ValueError(f"Error getting embedding from AWS Bedrock: {e}")
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using AWS Bedrock.
+
+        Args:
+            text (str): The text to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
+        Returns:
+            list: The embedding vector.
+        """
+        return self._get_embedding(text)

powermem/integrations/embeddings/azure_openai.py

@@ -0,0 +1,55 @@
+import os
+from typing import Literal, Optional
+
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+from openai import AzureOpenAI
+
+from powermem.integrations.embeddings.base import EmbeddingBase
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+SCOPE = "https://cognitiveservices.azure.com/.default"
+
+
+class AzureOpenAIEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        api_key = self.config.azure_kwargs.api_key or os.getenv("EMBEDDING_AZURE_OPENAI_API_KEY")
+        azure_deployment = self.config.azure_kwargs.azure_deployment or os.getenv("EMBEDDING_AZURE_DEPLOYMENT")
+        azure_endpoint = self.config.azure_kwargs.azure_endpoint or os.getenv("EMBEDDING_AZURE_ENDPOINT")
+        api_version = self.config.azure_kwargs.api_version or os.getenv("EMBEDDING_AZURE_API_VERSION")
+        default_headers = self.config.azure_kwargs.default_headers
+
+        # If the API key is not provided or is a placeholder, use DefaultAzureCredential.
+        if api_key is None or api_key == "" or api_key == "your-api-key":
+            self.credential = DefaultAzureCredential()
+            azure_ad_token_provider = get_bearer_token_provider(
+                self.credential,
+                SCOPE,
+            )
+            api_key = None
+        else:
+            azure_ad_token_provider = None
+
+        self.client = AzureOpenAI(
+            azure_deployment=azure_deployment,
+            azure_endpoint=azure_endpoint,
+            azure_ad_token_provider=azure_ad_token_provider,
+            api_version=api_version,
+            api_key=api_key,
+            http_client=self.config.http_client,
+            default_headers=default_headers,
+        )
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using OpenAI.
+
+        Args:
+            text (str): The text to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
+        Returns:
+            list: The embedding vector.
+        """
+        text = text.replace("\n", " ")
+        return self.client.embeddings.create(input=[text], model=self.config.model).data[0].embedding

powermem/integrations/embeddings/base.py

@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+from typing import Literal, Optional
+
+from powermem.integrations.embeddings.config.base import BaseEmbedderConfig
+
+
+class EmbeddingBase(ABC):
+    """Initialized a base embedding class
+
+    :param config: Embedding configuration option class, defaults to None
+    :type config: Optional[BaseEmbedderConfig], optional
+    """
+
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        if config is None:
+            self.config = BaseEmbedderConfig()
+        else:
+            self.config = config
+
+    @abstractmethod
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]]):
+        """
+        Get the embedding for the given text.
+
+        Args:
+            text (str): The text to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
+        Returns:
+            list: The embedding vector.
+        """
+        pass