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

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):

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):

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.dist-info → mem0ai_azure_mysql-0.1.115.1.dist-info}/METADATA

@@ -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.dist-info → mem0ai_azure_mysql-0.1.115.1.dist-info}/RECORD

@@ -38,7 +38,7 @@ mem0/dbs/configs.py,sha256=EoWnTNM_YPR39VgDDv8wS0KfA1nJ8Q0bpGtJdx1wu84,624
 mem0/dbs/mysql.py,sha256=AdrJQPrzqAN1cDjYgT-tA5n7CLPjJwR8LPlBhIJbT80,12731
 mem0/embeddings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mem0/embeddings/aws_bedrock.py,sha256=MvlmXfUxua2sIEoxzJKN9RPA5Xk9pN9peeDr3QBjuGw,3531
-mem0/embeddings/azure_openai.py,sha256=WU65mN6VQPROwmBguVwF8o8ydLFktZEjLmy9m6dQ660,1843
+mem0/embeddings/azure_openai.py,sha256=rJqPzA4XdLgZrEuEhk4UI2NPfCg_0w82wIa6x5RdrAw,2390
 mem0/embeddings/base.py,sha256=O-oFuizpx9tyL8qjhRu0l42GDkjgORy_UMtMCkRfc4k,973
 mem0/embeddings/configs.py,sha256=ucqSXa_Wiiy_LQJRpulD_-VFQQc2tWOOzEV3pLrHMow,881
 mem0/embeddings/gemini.py,sha256=6WPlY1GRYSfX2aD0ywou3uM4XenEQm6oB1c7x7R5oOk,1509
@@ -60,7 +60,7 @@ mem0/graphs/neptune/main.py,sha256=JoyMpFqRe9RNLnh4f_HyfUaxwHtaDC6ZCdr0hqR8Vcw,1
 mem0/llms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mem0/llms/anthropic.py,sha256=LtKWVTAQxQnpiMrB0jwihn4mOO9RBjUT_vEjeXPBYhE,2250
 mem0/llms/aws_bedrock.py,sha256=awRKZvsYpab3jYFu0e3_gTRxAToF3-M1nm9Sbr1QwzA,10520
-mem0/llms/azure_openai.py,sha256=6l-sTZ6RvFbMyLrxLdDJFoanaM0lsTQn4X9nvpcz5aI,4170
+mem0/llms/azure_openai.py,sha256=t1gZsMX48xYjAbG20kQtBZ7dWiZGW5oPDQqTOw5X14Q,4579
 mem0/llms/azure_openai_structured.py,sha256=gWFt73wn7okFcJAI8_a7SSxcnoV8_Qrf9axiV5qKNf8,2722
 mem0/llms/base.py,sha256=Uit80Gs0MNKKH2z9h0Fkmr-4u4QATCpKFlPXj_8DBuo,1061
 mem0/llms/configs.py,sha256=yvtU5WDBsiygCiMHjhG9ZCb1NIAMv8uXYD_KL8_X54w,997
@@ -90,7 +90,7 @@ mem0/proxy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mem0/proxy/main.py,sha256=DhGH9dvVd3GTS1yyZzTBpyYvgArDRcYIgKLBCcosCLE,7398
 mem0/utils/factory.py,sha256=FZ-8qKK4K7unnEfcmcYAY77i2IVDK1nbcsB1WeGlR2I,5448
 mem0/vector_stores/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mem0/vector_stores/azure_ai_search.py,sha256=pnTzzYGfeU2Q_pm8D_o_uJu-uRDNMgoYERGu7HZZb6M,14804
+mem0/vector_stores/azure_ai_search.py,sha256=AXcuvzll3dMWiisATZRIsEygc1tMVvAp7xdy8ei7rCE,15706
 mem0/vector_stores/baidu.py,sha256=Cp7TGtMnZRkuDAnRw0F0OjNg229dH31dXYZL1JmPQes,12683
 mem0/vector_stores/base.py,sha256=yb9TCcxximzdJd8vfNJRROHBrUyYLbH26-E63fclcDo,1344
 mem0/vector_stores/chroma.py,sha256=fHGbvecYPvlhHrS4lkSWt4S1EnxicEWuqG6qOqDleQA,7539
@@ -109,8 +109,8 @@ mem0/vector_stores/supabase.py,sha256=RCegadYW4Q8CT_NSeCywUVW01TLJQT4mDzeStNHEJH
 mem0/vector_stores/upstash_vector.py,sha256=dAIef-bIB6AE_dVR8UpVWZUPQs-wLhqfxEHyyXO5kQU,9073
 mem0/vector_stores/vertex_ai_vector_search.py,sha256=RaaPLTTbTPftlg7nm9WTROOhurSxqnAyx-g3tYoxv0I,24659
 mem0/vector_stores/weaviate.py,sha256=v91CA79Z_zJBohwaXlOW6jtudDlGq0kYKQFHMHGqUy4,11833
-mem0ai_azure_mysql-0.1.115.data/data/README.md,sha256=HhoTn6hBxTIsFwbm5zHT7aNmQOa5yaTgwdRIF6V7hX4,6525
-mem0ai_azure_mysql-0.1.115.dist-info/METADATA,sha256=BTfI2Osc41_2mBpXlJAynVJsWOtUvq1MbPMILl52D5M,8948
-mem0ai_azure_mysql-0.1.115.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mem0ai_azure_mysql-0.1.115.dist-info/licenses/LICENSE,sha256=C7y-kxw1MpOi-vzggyYYHf7qDlaMVmr9TOgzenD14hk,11349
-mem0ai_azure_mysql-0.1.115.dist-info/RECORD,,
+mem0ai_azure_mysql-0.1.115.1.data/data/README.md,sha256=HhoTn6hBxTIsFwbm5zHT7aNmQOa5yaTgwdRIF6V7hX4,6525
+mem0ai_azure_mysql-0.1.115.1.dist-info/METADATA,sha256=onCN-qJDuXyPZkn5b99HeNv_BPR_r_pubXYvp4sJjYA,8950
+mem0ai_azure_mysql-0.1.115.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mem0ai_azure_mysql-0.1.115.1.dist-info/licenses/LICENSE,sha256=C7y-kxw1MpOi-vzggyYYHf7qDlaMVmr9TOgzenD14hk,11349
+mem0ai_azure_mysql-0.1.115.1.dist-info/RECORD,,