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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mem0ai-azure-mysql
-Version: 0.1.115
+Version: 0.1.115.1
 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 → mem0ai_azure_mysql-0.1.115.1}/mem0/embeddings/azure_openai.py

@@ -1,32 +1,46 @@
 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_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")
+        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
+        else:
+            logger.info("Using Azure AD token provider 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
 
         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,
+            **auth_kwargs,
         )
 
     def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):

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

@@ -18,20 +18,30 @@ class AzureOpenAILLM(LLMBase):
         if not self.config.model:
             self.config.model = "gpt-4o"
 
+        api_key = self.config.azure_kwargs.api_key or os.getenv("LLM_AZURE_API_KEY")
         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
 
-        credential = DefaultAzureCredential()
-        token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
+        auth_kwargs = {}
+
+        if api_key:
+            auth_kwargs["azure_ad_token"] = api_key
+        else:
+            # Use Azure AD token provider if no API key is provided.
+            # This is useful for managed identities or when using Azure AD authentication.
+            credential = DefaultAzureCredential()
+            token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
+            auth_kwargs["azure_ad_token_provider"] = token_provider
+
         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,
+            **auth_kwargs,
         )
 
     def _parse_response(self, response, tools):

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

@@ -2,8 +2,11 @@ import json
 import logging
 import re
 from typing import List, Optional
+from datetime import datetime, timedelta, timezone
 
 from pydantic import BaseModel
+from azure.core.credentials import AccessToken
+from azure.core.credentials import TokenCredential
 from azure.identity import DefaultAzureCredential
 
 from mem0.memory.utils import extract_json
@@ -33,6 +36,18 @@ except ImportError:
 logger = logging.getLogger(__name__)
 
 
+class StaticTokenCredential(TokenCredential):
+    def __init__(self, token: str, expires_on: int | None = None):
+        self._token = token
+
+        if expires_on is None:
+            expires_on = int((datetime.now(timezone.utc) + timedelta(days=2)).timestamp())
+        self._expires_on = expires_on
+
+    def get_token(self, *scopes, **kwargs) -> AccessToken:
+        return AccessToken(self._token, self._expires_on)
+
+
 class OutputData(BaseModel):
     id: Optional[str]
     score: Optional[float]
@@ -77,7 +92,12 @@ class AzureAISearch(VectorStoreBase):
         self.hybrid_search = hybrid_search
         self.vector_filter_mode = vector_filter_mode
 
-        credential = DefaultAzureCredential()
+        if api_key is None:
+            logger.info("No API key provided. Using DefaultAzureCredential for authentication.")
+            credential = DefaultAzureCredential()
+        else:
+            logger.info("Using StaticTokenCredential for authentication.")
+            credential = StaticTokenCredential(api_key)
         self.search_client = SearchClient(
             endpoint=f"https://{service_name}.search.windows.net",
             index_name=self.index_name,
@@ -360,7 +380,7 @@ class AzureAISearch(VectorStoreBase):
             self.delete_col()
 
             # Reinitialize the clients
-            credential = DefaultAzureCredential()
+            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,

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

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mem0ai-azure-mysql"
-version = "0.1.115"
+version = "0.1.115.1"
 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"