mem0ai-azure-mysql 0.1.115.1__tar.gz → 0.1.115.2__tar.gz

{mem0ai_azure_mysql-0.1.115.1 → mem0ai_azure_mysql-0.1.115.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mem0ai-azure-mysql
-Version: 0.1.115.1
+Version: 0.1.115.2
 Summary: Long-term memory for AI Agents with Azure DefaultAzureCredential authentication and MySQL history database support
 License-File: LICENSE
 Requires-Python: <4.0,>=3.9

{mem0ai_azure_mysql-0.1.115.1 → mem0ai_azure_mysql-0.1.115.2}/mem0/configs/base.py

@@ -78,13 +78,27 @@ class AzureConfig(BaseModel):
         default_headers (Dict[str, str]): Headers to include in requests to the Azure API.
     """
 
-    api_key: str = Field(
+    api_key: str | None = Field(
         description="The API key used for authenticating with the Azure service.",
         default=None,
     )
-    azure_deployment: str = Field(description="The name of the Azure deployment.", default=None)
-    azure_endpoint: str = Field(description="The endpoint URL for the Azure service.", default=None)
-    api_version: str = Field(description="The version of the Azure API being used.", default=None)
+    azure_ad_token: str | None = Field(
+        description="The Azure AD token used for authentication.",
+        default=None,
+    )
+    azure_deployment: str | None = Field(
+        description="The name of the Azure deployment.",
+        default=None,
+    )
+    azure_endpoint: str | None = Field(
+        description="The endpoint URL for the Azure service.",
+        default=None,
+    )
+    api_version: str | None = Field(
+        description="The version of the Azure API being used.",
+        default=None
+    )
     default_headers: Optional[Dict[str, str]] = Field(
-        description="Headers to include in requests to the Azure API.", default=None
+        description="Headers to include in requests to the Azure API.",
+        default=None
     )

{mem0ai_azure_mysql-0.1.115.1 → mem0ai_azure_mysql-0.1.115.2}/mem0/configs/vector_stores/azure_ai_search.py

@@ -5,8 +5,9 @@ from pydantic import BaseModel, Field, model_validator
 
 class AzureAISearchConfig(BaseModel):
     collection_name: str = Field("mem0", description="Name of the collection")
-    service_name: str = Field(None, description="Azure AI Search service name")
-    api_key: str = Field(None, description="API key for the Azure AI Search service")
+    service_name: str | None = Field(None, description="Azure AI Search service name")
+    api_key: str | None = Field(None, description="API key for the Azure AI Search service")
+    azure_ad_token: str | None = Field(None, description="Azure AD token for authentication")
     embedding_model_dims: int = Field(1536, description="Dimension of the embedding vector")
     compression_type: Optional[str] = Field(
         None, description="Type of vector compression to use. Options: 'scalar', 'binary', or None"

{mem0ai_azure_mysql-0.1.115.1 → mem0ai_azure_mysql-0.1.115.2}/mem0/embeddings/azure_openai.py

@@ -1,35 +1,36 @@
 import os
-import logging
 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
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
 
 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_API_KEY")
+        azure_ad_token = self.config.azure_kwargs.azure_ad_token or os.getenv("EMBEDDING_AZURE_AD_TOKEN")
         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
 
-        auth_kwargs = {}
+        if azure_endpoint is None:
+            raise ValueError("Azure endpoint must be provided either in config or as an environment variable.")
 
+        auth_kwargs = {}
         if api_key:
-            logger.info(f"Using API key for Azure OpenAI embedding authentication. {api_key}")
-            auth_kwargs["azure_ad_token"] = api_key
+            print("Using API key for Azure OpenAI Embedding authentication.")
+            auth_kwargs["api_key"] = api_key
+        elif azure_ad_token:
+            print("Using Azure AD token for Azure OpenAI Embedding authentication.")
+            auth_kwargs["azure_ad_token"] = azure_ad_token
         else:
-            logger.info("Using Azure AD token provider for Azure OpenAI embedding authentication.")
+            print("Using DefaultAzureCredential for Azure OpenAI Embedding authentication.")
             credential = DefaultAzureCredential()
             token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
             auth_kwargs["azure_ad_token_provider"] = token_provider

{mem0ai_azure_mysql-0.1.115.1 → mem0ai_azure_mysql-0.1.115.2}/mem0/llms/azure_openai.py

@@ -19,18 +19,24 @@ class AzureOpenAILLM(LLMBase):
             self.config.model = "gpt-4o"
 
         api_key = self.config.azure_kwargs.api_key or os.getenv("LLM_AZURE_API_KEY")
+        azure_ad_token = self.config.azure_kwargs.azure_ad_token or os.getenv("LLM_AZURE_AD_TOKEN")
         azure_deployment = self.config.azure_kwargs.azure_deployment or os.getenv("LLM_AZURE_DEPLOYMENT")
         azure_endpoint = self.config.azure_kwargs.azure_endpoint or os.getenv("LLM_AZURE_ENDPOINT")
         api_version = self.config.azure_kwargs.api_version or os.getenv("LLM_AZURE_API_VERSION")
         default_headers = self.config.azure_kwargs.default_headers
 
-        auth_kwargs = {}
+        if azure_endpoint is None:
+            raise ValueError("Azure endpoint must be provided either in config or as an environment variable.")
 
+        auth_kwargs = {}
         if api_key:
-            auth_kwargs["azure_ad_token"] = api_key
+            print("Using API key for Azure OpenAI authentication.")
+            auth_kwargs["api_key"] = api_key
+        elif azure_ad_token:
+            print("Using Azure AD token for Azure OpenAI authentication.")
+            auth_kwargs["azure_ad_token"] = azure_ad_token
         else:
-            # Use Azure AD token provider if no API key is provided.
-            # This is useful for managed identities or when using Azure AD authentication.
+            print("Using DefaultAzureCredential for Azure OpenAI authentication.")
             credential = DefaultAzureCredential()
             token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
             auth_kwargs["azure_ad_token_provider"] = token_provider

{mem0ai_azure_mysql-0.1.115.1 → mem0ai_azure_mysql-0.1.115.2}/mem0/vector_stores/azure_ai_search.py

@@ -6,7 +6,7 @@ from datetime import datetime, timedelta, timezone
 
 from pydantic import BaseModel
 from azure.core.credentials import AccessToken
-from azure.core.credentials import TokenCredential
+from azure.core.credentials import TokenCredential, AzureKeyCredential
 from azure.identity import DefaultAzureCredential
 
 from mem0.memory.utils import extract_json
@@ -65,6 +65,7 @@ class AzureAISearch(VectorStoreBase):
         use_float16: bool = False,
         hybrid_search: bool = False,
         vector_filter_mode: Optional[str] = None,
+        azure_ad_token: Optional[str] = None
     ):
         """
         Initialize the Azure AI Search vector store.
@@ -80,6 +81,7 @@ class AzureAISearch(VectorStoreBase):
                 (Note: This flag is preserved from the initial implementation per feedback.)
             hybrid_search (bool): Whether to use hybrid search. Default is False.
             vector_filter_mode (Optional[str]): Mode for vector filtering. Default is "preFilter".
+            azure_ad_token (Optional[str]): Azure AD token for authentication.
         """
         self.service_name = service_name
         self.api_key = api_key
@@ -91,29 +93,43 @@ class AzureAISearch(VectorStoreBase):
         self.use_float16 = use_float16
         self.hybrid_search = hybrid_search
         self.vector_filter_mode = vector_filter_mode
+        self.azure_ad_token = azure_ad_token
 
-        if api_key is None:
-            logger.info("No API key provided. Using DefaultAzureCredential for authentication.")
-            credential = DefaultAzureCredential()
+        self.search_client, self.index_client = self.create_clients()
+
+        collections = self.list_cols()
+        if collection_name not in collections:
+            self.create_col()
+    
+    def create_clients(self):
+        """
+        Create the search and index clients.
+        This method is called during initialization to set up the clients.
+        """
+        if self.api_key:
+            print("Using API key for Azure AI Search authentication.")
+            credential = AzureKeyCredential(self.api_key)
+        elif self.azure_ad_token:
+            print("Using Azure AD token for Azure AI Search authentication.")
+            credential = StaticTokenCredential(self.azure_ad_token)
         else:
-            logger.info("Using StaticTokenCredential for authentication.")
-            credential = StaticTokenCredential(api_key)
-        self.search_client = SearchClient(
-            endpoint=f"https://{service_name}.search.windows.net",
+            print("Using DefaultAzureCredential for Azure AI Search authentication.")
+            credential = DefaultAzureCredential()
+
+        search_client = SearchClient(
+            endpoint=f"https://{self.service_name}.search.windows.net",
             index_name=self.index_name,
             credential=credential,
         )
-        self.index_client = SearchIndexClient(
-            endpoint=f"https://{service_name}.search.windows.net",
+        index_client = SearchIndexClient(
+            endpoint=f"https://{self.service_name}.search.windows.net",
             credential=credential,
         )
 
-        self.search_client._client._config.user_agent_policy.add_user_agent("mem0")
-        self.index_client._client._config.user_agent_policy.add_user_agent("mem0")
+        search_client._client._config.user_agent_policy.add_user_agent("mem0")
+        index_client._client._config.user_agent_policy.add_user_agent("mem0")
 
-        collections = self.list_cols()
-        if collection_name not in collections:
-            self.create_col()
+        return search_client, index_client
 
     def create_col(self):
         """Create a new index in Azure AI Search."""
@@ -379,22 +395,8 @@ class AzureAISearch(VectorStoreBase):
             # Delete the collection
             self.delete_col()
 
-            # Reinitialize the clients
-            credential = StaticTokenCredential(self.api_key) if self.api_key else DefaultAzureCredential()
-            service_endpoint = f"https://{self.service_name}.search.windows.net"
-            self.search_client = SearchClient(
-                endpoint=service_endpoint,
-                index_name=self.index_name,
-                credential=credential,
-            )
-            self.index_client = SearchIndexClient(
-                endpoint=service_endpoint,
-                credential=credential,
-            )
-
-            # Add user agent
-            self.search_client._client._config.user_agent_policy.add_user_agent("mem0")
-            self.index_client._client._config.user_agent_policy.add_user_agent("mem0")
+            # Recreate the clients
+            self.search_client, self.index_client = self.create_clients()
 
             # Create the collection
             self.create_col()

{mem0ai_azure_mysql-0.1.115.1 → mem0ai_azure_mysql-0.1.115.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mem0ai-azure-mysql"
-version = "0.1.115.1"
+version = "0.1.115.2"
 description = "Long-term memory for AI Agents with Azure DefaultAzureCredential authentication and MySQL history database support"
 readme = "README.md"
 requires-python = ">=3.9,<4.0"