mem0ai-azure-mysql 0.1.115__py3-none-any.whl

mem0/dbs/mysql.py

@@ -0,0 +1,321 @@
+import ssl
+import logging
+import threading
+import uuid
+from typing import Any, Dict, List, Optional
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import sessionmaker, Session
+
+from mem0.configs.dbs.mysql import MySQLConfig
+from mem0.dbs.base import DBBase
+
+logger = logging.getLogger(__name__)
+
+
+class MySQLManager(DBBase):
+    """MySQL implementation of DBBase for managing memory history using SQLAlchemy."""
+    
+    def __init__(self, config: Optional[MySQLConfig] = None):
+        super().__init__(config)
+        if config is None:
+            self.config = MySQLConfig()
+        else:
+            self.config = config
+            
+        self.engine: Optional[Engine] = None
+        self.Session: Optional[sessionmaker] = None
+        self._lock = threading.Lock()
+        self._connect()
+        self._migrate_history_table()
+        self._create_history_table()
+
+    def _connect(self) -> None:
+        """Establish connection to MySQL database using SQLAlchemy."""
+        try:
+            # Build connection URL
+            connection_url = self._build_connection_url()
+            
+            # Create engine with connection pooling
+            connect_args = {}
+            if hasattr(self.config, 'connection_params'):
+                # Add valid MySQL connection parameters
+                valid_params = {
+                    'ssl_ca', 'ssl_cert', 'ssl_key', 'ssl_verify_cert', 'ssl_verify_identity',
+                    'connect_timeout', 'charset', 'init_command'
+                }
+                for key, value in self.config.connection_params.items():
+                    if key in valid_params:
+                        connect_args[key] = value
+            
+            if self.config.ssl_enabled:
+                connect_args['ssl'] = ssl.create_default_context()
+            
+            self.engine = create_engine(
+                connection_url,
+                connect_args=connect_args,
+                pool_pre_ping=True,
+                pool_recycle=3600,
+                echo=False
+            )
+            
+            self.Session = sessionmaker(bind=self.engine)
+            
+            # Test connection
+            with self.engine.connect() as conn:
+                conn.execute(text("SELECT 1"))
+            
+            logger.info("Successfully connected to MySQL database using SQLAlchemy")
+            
+        except SQLAlchemyError as e:
+            logger.error(f"Failed to connect to MySQL: {e}")
+            raise
+
+    def _build_connection_url(self) -> str:
+        """Build SQLAlchemy connection URL for MySQL."""
+        # Use PyMySQL as the driver (more reliable than mysqlclient)
+        url_parts = ["mysql+pymysql://"]
+        
+        if self.config.user:
+            url_parts.append(self.config.user)
+            if self.config.password:
+                url_parts.append(f":{self.config.password}")
+            url_parts.append("@")
+        
+        url_parts.append(self.config.host or "localhost")
+        
+        if self.config.port:
+            url_parts.append(f":{self.config.port}")
+        
+        if self.config.database:
+            url_parts.append(f"/{self.config.database}")
+        
+        return "".join(url_parts)
+
+    def _migrate_history_table(self) -> None:
+        """
+        If a pre-existing history table had the old schema,
+        rename it, create the new schema, copy the intersecting data, then
+        drop the old table.
+        """
+        with self._lock:
+            if self.engine is None:
+                raise RuntimeError("Database connection is not established")
+            try:
+                with self.engine.begin() as conn:
+                    # Check if history table exists
+                    result = conn.execute(text("""
+                        SELECT COUNT(*)
+                        FROM information_schema.tables
+                        WHERE table_schema = DATABASE()
+                        AND table_name = 'history'
+                    """))
+                    
+                    count = result.scalar()
+                    if count == 0:
+                        return  # nothing to migrate
+                    
+                    # Get current table columns
+                    result = conn.execute(text("""
+                        SELECT column_name
+                        FROM information_schema.columns
+                        WHERE table_schema = DATABASE()
+                        AND table_name = 'history'
+                    """))
+                    old_cols = {row[0] for row in result.fetchall()}
+                    
+                    expected_cols = {
+                        "id",
+                        "memory_id",
+                        "old_memory",
+                        "new_memory",
+                        "event",
+                        "created_at",
+                        "updated_at",
+                        "is_deleted",
+                        "actor_id",
+                        "role",
+                    }
+                    
+                    if old_cols == expected_cols:
+                        return
+                    
+                    logger.info("Migrating history table to new schema (no convo columns).")
+                    
+                    # Clean up any existing history_old table from previous failed migration
+                    conn.execute(text("DROP TABLE IF EXISTS history_old"))
+                    
+                    # Rename the current history table
+                    conn.execute(text("ALTER TABLE history RENAME TO history_old"))
+                    
+                    # Create the new history table with updated schema
+                    conn.execute(text("""
+                        CREATE TABLE history (
+                            id           VARCHAR(36) PRIMARY KEY,
+                            memory_id    TEXT,
+                            old_memory   TEXT,
+                            new_memory   TEXT,
+                            event        TEXT,
+                            created_at   TIMESTAMP NULL,
+                            updated_at   TIMESTAMP NULL,
+                            is_deleted   INT,
+                            actor_id     TEXT,
+                            role         TEXT
+                        )
+                    """))
+                    
+                    # Copy data from old table to new table
+                    intersecting = list(expected_cols & old_cols)
+                    if intersecting:
+                        cols_str = ", ".join(f"`{col}`" for col in intersecting)
+                        query = f"INSERT INTO history ({cols_str}) SELECT {cols_str} FROM history_old"
+                        conn.execute(text(query))
+                    
+                    # Drop the old table
+                    conn.execute(text("DROP TABLE history_old"))
+                    
+                    logger.info("History table migration completed successfully.")
+                    
+            except SQLAlchemyError as e:
+                logger.error(f"History table migration failed: {e}")
+                raise
+
+    def _create_history_table(self) -> None:
+        """Create the history table if it doesn't exist."""
+        with self._lock:
+            if self.engine is None:
+                raise RuntimeError("Database connection is not established")
+            try:
+                with self.engine.begin() as conn:
+                    conn.execute(text("""
+                        CREATE TABLE IF NOT EXISTS history (
+                            id           VARCHAR(36) PRIMARY KEY,
+                            memory_id    TEXT,
+                            old_memory   TEXT,
+                            new_memory   TEXT,
+                            event        TEXT,
+                            created_at   TIMESTAMP NULL,
+                            updated_at   TIMESTAMP NULL,
+                            is_deleted   INT,
+                            actor_id     TEXT,
+                            role         TEXT
+                        )
+                    """))
+                    
+            except SQLAlchemyError as e:
+                logger.error(f"Failed to create history table: {e}")
+                raise
+
+    def add_history(
+        self,
+        memory_id: str,
+        old_memory: Optional[str],
+        new_memory: Optional[str],
+        event: str,
+        *,
+        created_at: Optional[str] = None,
+        updated_at: Optional[str] = None,
+        is_deleted: int = 0,
+        actor_id: Optional[str] = None,
+        role: Optional[str] = None,
+    ) -> None:
+        """Add a history record to the database.
+        
+        :param memory_id: The ID of the memory being tracked
+        :param old_memory: The previous memory content
+        :param new_memory: The new memory content
+        :param event: The type of event that occurred
+        :param created_at: When the record was created
+        :param updated_at: When the record was last updated
+        :param is_deleted: Whether the record is deleted (0 or 1)
+        :param actor_id: ID of the actor who made the change
+        :param role: Role of the actor
+        """
+        with self._lock:
+            if self.engine is None:
+                raise RuntimeError("Database connection is not established")
+            try:
+                with self.engine.begin() as conn:
+                    conn.execute(text("""
+                        INSERT INTO history (
+                            id, memory_id, old_memory, new_memory, event,
+                            created_at, updated_at, is_deleted, actor_id, role
+                        )
+                        VALUES (:id, :memory_id, :old_memory, :new_memory, :event,
+                                :created_at, :updated_at, :is_deleted, :actor_id, :role)
+                    """), {
+                        "id": str(uuid.uuid4()),
+                        "memory_id": memory_id,
+                        "old_memory": old_memory,
+                        "new_memory": new_memory,
+                        "event": event,
+                        "created_at": created_at,
+                        "updated_at": updated_at,
+                        "is_deleted": is_deleted,
+                        "actor_id": actor_id,
+                        "role": role,
+                    })
+                    
+            except SQLAlchemyError as e:
+                logger.error(f"Failed to add history record: {e}")
+                raise
+
+    def get_history(self, memory_id: str) -> List[Dict[str, Any]]:
+        """Retrieve history records for a given memory ID.
+        
+        :param memory_id: The ID of the memory to get history for
+        :return: List of history records as dictionaries
+        """
+        with self._lock:
+            if self.engine is None:
+                raise RuntimeError("Database connection is not established")
+            
+            with self.engine.connect() as conn:
+                result = conn.execute(text("""
+                    SELECT id, memory_id, old_memory, new_memory, event,
+                           created_at, updated_at, is_deleted, actor_id, role
+                    FROM history
+                    WHERE memory_id = :memory_id
+                    ORDER BY created_at ASC, updated_at ASC
+                """), {"memory_id": memory_id})
+                
+                rows = result.fetchall()
+
+        return [
+            {
+                "id": r.id,
+                "memory_id": r.memory_id,
+                "old_memory": r.old_memory,
+                "new_memory": r.new_memory,
+                "event": r.event,
+                "created_at": r.created_at,
+                "updated_at": r.updated_at,
+                "is_deleted": bool(r.is_deleted),
+                "actor_id": r.actor_id,
+                "role": r.role,
+            }
+            for r in rows
+        ]
+
+    def reset(self) -> None:
+        """Reset/clear all data in the database."""
+        with self._lock:
+            if self.engine is None:
+                raise RuntimeError("Database connection is not established")
+            try:
+                with self.engine.begin() as conn:
+                    conn.execute(text("DROP TABLE IF EXISTS history"))
+                self._create_history_table()
+                
+            except SQLAlchemyError as e:
+                logger.error(f"Failed to reset history table: {e}")
+                raise
+
+    def close(self) -> None:
+        """Close the database connection and clean up resources."""
+        if self.engine:
+            self.engine.dispose()
+            self.engine = None
+            self.Session = None

mem0/embeddings/aws_bedrock.py

@@ -0,0 +1,100 @@
+import json
+import os
+from typing import Literal, Optional
+
+try:
+    import boto3
+except ImportError:
+    raise ImportError("The 'boto3' library is required. Please install it using 'pip install boto3'.")
+
+import numpy as np
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+
+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", "")
+        aws_region = os.environ.get("AWS_REGION", "us-west-2")
+
+        # 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
+        if hasattr(self.config, "aws_region"):
+            aws_region = self.config.aws_region
+
+        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)

mem0/embeddings/azure_openai.py

@@ -0,0 +1,43 @@
+import os
+from typing import Literal, Optional
+
+from openai import AzureOpenAI
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+
+class AzureOpenAIEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        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
+
+        credential = DefaultAzureCredential()
+        token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
+
+        self.client = AzureOpenAI(
+            azure_deployment=azure_deployment,
+            azure_endpoint=azure_endpoint,
+            api_version=api_version,
+            azure_ad_token_provider=token_provider,
+            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

mem0/embeddings/base.py

@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+from typing import Literal, Optional
+
+from mem0.configs.embeddings.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

mem0/embeddings/configs.py

@@ -0,0 +1,30 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class EmbedderConfig(BaseModel):
+    provider: str = Field(
+        description="Provider of the embedding model (e.g., 'ollama', 'openai')",
+        default="openai",
+    )
+    config: Optional[dict] = Field(description="Configuration for the specific embedding model", default={})
+
+    @field_validator("config")
+    def validate_config(cls, v, values):
+        provider = values.data.get("provider")
+        if provider in [
+            "openai",
+            "ollama",
+            "huggingface",
+            "azure_openai",
+            "gemini",
+            "vertexai",
+            "together",
+            "lmstudio",
+            "langchain",
+            "aws_bedrock",
+        ]:
+            return v
+        else:
+            raise ValueError(f"Unsupported embedding provider: {provider}")

mem0/embeddings/gemini.py

@@ -0,0 +1,39 @@
+import os
+from typing import Literal, Optional
+
+from google import genai
+from google.genai import types
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+
+class GoogleGenAIEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "models/text-embedding-004"
+        self.config.embedding_dims = self.config.embedding_dims or self.config.output_dimensionality or 768
+
+        api_key = self.config.api_key or os.getenv("GOOGLE_API_KEY")
+
+        self.client = genai.Client(api_key=api_key)
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using Google Generative AI.
+        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", " ")
+
+        # Create config for embedding parameters
+        config = types.EmbedContentConfig(output_dimensionality=self.config.embedding_dims)
+
+        # Call the embed_content method with the correct parameters
+        response = self.client.models.embed_content(model=self.config.model, contents=text, config=config)
+
+        return response.embeddings[0].values

mem0/embeddings/huggingface.py

@@ -0,0 +1,41 @@
+import logging
+from typing import Literal, Optional
+
+from openai import OpenAI
+from sentence_transformers import SentenceTransformer
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+logging.getLogger("transformers").setLevel(logging.WARNING)
+logging.getLogger("sentence_transformers").setLevel(logging.WARNING)
+logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
+
+
+class HuggingFaceEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        if config.huggingface_base_url:
+            self.client = OpenAI(base_url=config.huggingface_base_url)
+        else:
+            self.config.model = self.config.model or "multi-qa-MiniLM-L6-cos-v1"
+
+            self.model = SentenceTransformer(self.config.model, **self.config.model_kwargs)
+
+            self.config.embedding_dims = self.config.embedding_dims or self.model.get_sentence_embedding_dimension()
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using Hugging Face.
+
+        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.
+        """
+        if self.config.huggingface_base_url:
+            return self.client.embeddings.create(input=text, model="tei").data[0].embedding
+        else:
+            return self.model.encode(text, convert_to_numpy=True).tolist()

mem0/embeddings/langchain.py

@@ -0,0 +1,35 @@
+from typing import Literal, Optional
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+try:
+    from langchain.embeddings.base import Embeddings
+except ImportError:
+    raise ImportError("langchain is not installed. Please install it using `pip install langchain`")
+
+
+class LangchainEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        if self.config.model is None:
+            raise ValueError("`model` parameter is required")
+
+        if not isinstance(self.config.model, Embeddings):
+            raise ValueError("`model` must be an instance of Embeddings")
+
+        self.langchain_model = self.config.model
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using Langchain.
+
+        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.langchain_model.embed_query(text)

mem0/embeddings/lmstudio.py

@@ -0,0 +1,29 @@
+from typing import Literal, Optional
+
+from openai import OpenAI
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+
+class LMStudioEmbedding(EmbeddingBase):
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "nomic-ai/nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf"
+        self.config.embedding_dims = self.config.embedding_dims or 1536
+        self.config.api_key = self.config.api_key or "lm-studio"
+
+        self.client = OpenAI(base_url=self.config.lmstudio_base_url, api_key=self.config.api_key)
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using LM Studio.
+        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