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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
mem0/configs/base.py CHANGED
@@ -78,13 +78,27 @@ class AzureConfig(BaseModel):
78
78
  default_headers (Dict[str, str]): Headers to include in requests to the Azure API.
79
79
  """
80
80
 
81
- api_key: str = Field(
81
+ api_key: str | None = Field(
82
82
  description="The API key used for authenticating with the Azure service.",
83
83
  default=None,
84
84
  )
85
- azure_deployment: str = Field(description="The name of the Azure deployment.", default=None)
86
- azure_endpoint: str = Field(description="The endpoint URL for the Azure service.", default=None)
87
- api_version: str = Field(description="The version of the Azure API being used.", default=None)
85
+ azure_ad_token: str | None = Field(
86
+ description="The Azure AD token used for authentication.",
87
+ default=None,
88
+ )
89
+ azure_deployment: str | None = Field(
90
+ description="The name of the Azure deployment.",
91
+ default=None,
92
+ )
93
+ azure_endpoint: str | None = Field(
94
+ description="The endpoint URL for the Azure service.",
95
+ default=None,
96
+ )
97
+ api_version: str | None = Field(
98
+ description="The version of the Azure API being used.",
99
+ default=None
100
+ )
88
101
  default_headers: Optional[Dict[str, str]] = Field(
89
- description="Headers to include in requests to the Azure API.", default=None
102
+ description="Headers to include in requests to the Azure API.",
103
+ default=None
90
104
  )
@@ -5,8 +5,9 @@ from pydantic import BaseModel, Field, model_validator
5
5
 
6
6
  class AzureAISearchConfig(BaseModel):
7
7
  collection_name: str = Field("mem0", description="Name of the collection")
8
- service_name: str = Field(None, description="Azure AI Search service name")
9
- api_key: str = Field(None, description="API key for the Azure AI Search service")
8
+ service_name: str | None = Field(None, description="Azure AI Search service name")
9
+ api_key: str | None = Field(None, description="API key for the Azure AI Search service")
10
+ azure_ad_token: str | None = Field(None, description="Azure AD token for authentication")
10
11
  embedding_model_dims: int = Field(1536, description="Dimension of the embedding vector")
11
12
  compression_type: Optional[str] = Field(
12
13
  None, description="Type of vector compression to use. Options: 'scalar', 'binary', or None"
@@ -12,21 +12,36 @@ class AzureOpenAIEmbedding(EmbeddingBase):
12
12
  def __init__(self, config: Optional[BaseEmbedderConfig] = None):
13
13
  super().__init__(config)
14
14
 
15
+ api_key = self.config.azure_kwargs.api_key or os.getenv("EMBEDDING_AZURE_API_KEY")
16
+ azure_ad_token = self.config.azure_kwargs.azure_ad_token or os.getenv("EMBEDDING_AZURE_AD_TOKEN")
15
17
  azure_deployment = self.config.azure_kwargs.azure_deployment or os.getenv("EMBEDDING_AZURE_DEPLOYMENT")
16
18
  azure_endpoint = self.config.azure_kwargs.azure_endpoint or os.getenv("EMBEDDING_AZURE_ENDPOINT")
17
19
  api_version = self.config.azure_kwargs.api_version or os.getenv("EMBEDDING_AZURE_API_VERSION")
18
20
  default_headers = self.config.azure_kwargs.default_headers
19
21
 
20
- credential = DefaultAzureCredential()
21
- token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
22
+ if azure_endpoint is None:
23
+ raise ValueError("Azure endpoint must be provided either in config or as an environment variable.")
24
+
25
+ auth_kwargs = {}
26
+ if api_key:
27
+ print("Using API key for Azure OpenAI Embedding authentication.")
28
+ auth_kwargs["api_key"] = api_key
29
+ elif azure_ad_token:
30
+ print("Using Azure AD token for Azure OpenAI Embedding authentication.")
31
+ auth_kwargs["azure_ad_token"] = azure_ad_token
32
+ else:
33
+ print("Using DefaultAzureCredential for Azure OpenAI Embedding authentication.")
34
+ credential = DefaultAzureCredential()
35
+ token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
36
+ auth_kwargs["azure_ad_token_provider"] = token_provider
22
37
 
23
38
  self.client = AzureOpenAI(
24
39
  azure_deployment=azure_deployment,
25
40
  azure_endpoint=azure_endpoint,
26
41
  api_version=api_version,
27
- azure_ad_token_provider=token_provider,
28
42
  http_client=self.config.http_client,
29
43
  default_headers=default_headers,
44
+ **auth_kwargs,
30
45
  )
31
46
 
32
47
  def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
mem0/llms/azure_openai.py CHANGED
@@ -18,20 +18,36 @@ class AzureOpenAILLM(LLMBase):
18
18
  if not self.config.model:
19
19
  self.config.model = "gpt-4o"
20
20
 
21
+ api_key = self.config.azure_kwargs.api_key or os.getenv("LLM_AZURE_API_KEY")
22
+ azure_ad_token = self.config.azure_kwargs.azure_ad_token or os.getenv("LLM_AZURE_AD_TOKEN")
21
23
  azure_deployment = self.config.azure_kwargs.azure_deployment or os.getenv("LLM_AZURE_DEPLOYMENT")
22
24
  azure_endpoint = self.config.azure_kwargs.azure_endpoint or os.getenv("LLM_AZURE_ENDPOINT")
23
25
  api_version = self.config.azure_kwargs.api_version or os.getenv("LLM_AZURE_API_VERSION")
24
26
  default_headers = self.config.azure_kwargs.default_headers
25
27
 
26
- credential = DefaultAzureCredential()
27
- token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
28
+ if azure_endpoint is None:
29
+ raise ValueError("Azure endpoint must be provided either in config or as an environment variable.")
30
+
31
+ auth_kwargs = {}
32
+ if api_key:
33
+ print("Using API key for Azure OpenAI authentication.")
34
+ auth_kwargs["api_key"] = api_key
35
+ elif azure_ad_token:
36
+ print("Using Azure AD token for Azure OpenAI authentication.")
37
+ auth_kwargs["azure_ad_token"] = azure_ad_token
38
+ else:
39
+ print("Using DefaultAzureCredential for Azure OpenAI authentication.")
40
+ credential = DefaultAzureCredential()
41
+ token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
42
+ auth_kwargs["azure_ad_token_provider"] = token_provider
43
+
28
44
  self.client = AzureOpenAI(
29
45
  azure_deployment=azure_deployment,
30
46
  azure_endpoint=azure_endpoint,
31
47
  api_version=api_version,
32
- azure_ad_token_provider=token_provider,
33
48
  http_client=self.config.http_client,
34
49
  default_headers=default_headers,
50
+ **auth_kwargs,
35
51
  )
36
52
 
37
53
  def _parse_response(self, response, tools):
@@ -2,8 +2,11 @@ import json
2
2
  import logging
3
3
  import re
4
4
  from typing import List, Optional
5
+ from datetime import datetime, timedelta, timezone
5
6
 
6
7
  from pydantic import BaseModel
8
+ from azure.core.credentials import AccessToken
9
+ from azure.core.credentials import TokenCredential, AzureKeyCredential
7
10
  from azure.identity import DefaultAzureCredential
8
11
 
9
12
  from mem0.memory.utils import extract_json
@@ -33,6 +36,18 @@ except ImportError:
33
36
  logger = logging.getLogger(__name__)
34
37
 
35
38
 
39
+ class StaticTokenCredential(TokenCredential):
40
+ def __init__(self, token: str, expires_on: int | None = None):
41
+ self._token = token
42
+
43
+ if expires_on is None:
44
+ expires_on = int((datetime.now(timezone.utc) + timedelta(days=2)).timestamp())
45
+ self._expires_on = expires_on
46
+
47
+ def get_token(self, *scopes, **kwargs) -> AccessToken:
48
+ return AccessToken(self._token, self._expires_on)
49
+
50
+
36
51
  class OutputData(BaseModel):
37
52
  id: Optional[str]
38
53
  score: Optional[float]
@@ -50,6 +65,7 @@ class AzureAISearch(VectorStoreBase):
50
65
  use_float16: bool = False,
51
66
  hybrid_search: bool = False,
52
67
  vector_filter_mode: Optional[str] = None,
68
+ azure_ad_token: Optional[str] = None
53
69
  ):
54
70
  """
55
71
  Initialize the Azure AI Search vector store.
@@ -65,6 +81,7 @@ class AzureAISearch(VectorStoreBase):
65
81
  (Note: This flag is preserved from the initial implementation per feedback.)
66
82
  hybrid_search (bool): Whether to use hybrid search. Default is False.
67
83
  vector_filter_mode (Optional[str]): Mode for vector filtering. Default is "preFilter".
84
+ azure_ad_token (Optional[str]): Azure AD token for authentication.
68
85
  """
69
86
  self.service_name = service_name
70
87
  self.api_key = api_key
@@ -76,24 +93,43 @@ class AzureAISearch(VectorStoreBase):
76
93
  self.use_float16 = use_float16
77
94
  self.hybrid_search = hybrid_search
78
95
  self.vector_filter_mode = vector_filter_mode
96
+ self.azure_ad_token = azure_ad_token
97
+
98
+ self.search_client, self.index_client = self.create_clients()
99
+
100
+ collections = self.list_cols()
101
+ if collection_name not in collections:
102
+ self.create_col()
103
+
104
+ def create_clients(self):
105
+ """
106
+ Create the search and index clients.
107
+ This method is called during initialization to set up the clients.
108
+ """
109
+ if self.api_key:
110
+ print("Using API key for Azure AI Search authentication.")
111
+ credential = AzureKeyCredential(self.api_key)
112
+ elif self.azure_ad_token:
113
+ print("Using Azure AD token for Azure AI Search authentication.")
114
+ credential = StaticTokenCredential(self.azure_ad_token)
115
+ else:
116
+ print("Using DefaultAzureCredential for Azure AI Search authentication.")
117
+ credential = DefaultAzureCredential()
79
118
 
80
- credential = DefaultAzureCredential()
81
- self.search_client = SearchClient(
82
- endpoint=f"https://{service_name}.search.windows.net",
119
+ search_client = SearchClient(
120
+ endpoint=f"https://{self.service_name}.search.windows.net",
83
121
  index_name=self.index_name,
84
122
  credential=credential,
85
123
  )
86
- self.index_client = SearchIndexClient(
87
- endpoint=f"https://{service_name}.search.windows.net",
124
+ index_client = SearchIndexClient(
125
+ endpoint=f"https://{self.service_name}.search.windows.net",
88
126
  credential=credential,
89
127
  )
90
128
 
91
- self.search_client._client._config.user_agent_policy.add_user_agent("mem0")
92
- self.index_client._client._config.user_agent_policy.add_user_agent("mem0")
129
+ search_client._client._config.user_agent_policy.add_user_agent("mem0")
130
+ index_client._client._config.user_agent_policy.add_user_agent("mem0")
93
131
 
94
- collections = self.list_cols()
95
- if collection_name not in collections:
96
- self.create_col()
132
+ return search_client, index_client
97
133
 
98
134
  def create_col(self):
99
135
  """Create a new index in Azure AI Search."""
@@ -359,22 +395,8 @@ class AzureAISearch(VectorStoreBase):
359
395
  # Delete the collection
360
396
  self.delete_col()
361
397
 
362
- # Reinitialize the clients
363
- credential = DefaultAzureCredential()
364
- service_endpoint = f"https://{self.service_name}.search.windows.net"
365
- self.search_client = SearchClient(
366
- endpoint=service_endpoint,
367
- index_name=self.index_name,
368
- credential=credential,
369
- )
370
- self.index_client = SearchIndexClient(
371
- endpoint=service_endpoint,
372
- credential=credential,
373
- )
374
-
375
- # Add user agent
376
- self.search_client._client._config.user_agent_policy.add_user_agent("mem0")
377
- self.index_client._client._config.user_agent_policy.add_user_agent("mem0")
398
+ # Recreate the clients
399
+ self.search_client, self.index_client = self.create_clients()
378
400
 
379
401
  # Create the collection
380
402
  self.create_col()
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: mem0ai-azure-mysql
3
- Version: 0.1.115
3
+ Version: 0.1.115.2
4
4
  Summary: Long-term memory for AI Agents with Azure DefaultAzureCredential authentication and MySQL history database support
5
5
  License-File: LICENSE
6
6
  Requires-Python: <4.0,>=3.9
@@ -4,7 +4,7 @@ mem0/client/main.py,sha256=MtYDxoP9Gauv3dblAozoDmDX0uNqOUG8BhN51nzOIrY,55178
4
4
  mem0/client/project.py,sha256=6nrOQbAY-7fIEGgh6kXI90AltvFaI-TnwnTEY_8K8h0,26577
5
5
  mem0/client/utils.py,sha256=YRS-qpofVngRbu9CqGuHMgkt6edCLGWSLWe67sSBKGQ,719
6
6
  mem0/configs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
7
- mem0/configs/base.py,sha256=ux9CeExSJ1SfyDElGNLSq9777ar7_NCjG9ahaCEdM0g,3588
7
+ mem0/configs/base.py,sha256=ZYgQP7I_nRNWjKlKKMOvwyfKLDnUQOlHSvHFqeh462M,3827
8
8
  mem0/configs/enums.py,sha256=5qj-Ptly_-5_hnrQd5ViBR2wgQEVlendYiMCrA_EK9E,151
9
9
  mem0/configs/prompts.py,sha256=FApgNZPdonc16NBPBjvjhWjQQcV_qU0f4UaqSauh1sY,16373
10
10
  mem0/configs/dbs/__init__.py,sha256=hcfsoti8rHf5KukRSahiz3jzWhXCcoYmy_4Uv__yy24,136
@@ -15,7 +15,7 @@ mem0/configs/embeddings/base.py,sha256=4oJi2QCy7i4yUKQh2_9OCSA2zUe4SidE3i_p7-n0B
15
15
  mem0/configs/llms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
16
16
  mem0/configs/llms/base.py,sha256=vI7BQelHRrRDlWsdBf1lutZ8jiXfbQ-2kGmfDvvOnXg,6646
17
17
  mem0/configs/vector_stores/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
18
- mem0/configs/vector_stores/azure_ai_search.py,sha256=0c5oZr3gwJH19MbU3UhLqfPQ29kgThYUMFvkdftiZ7k,2545
18
+ mem0/configs/vector_stores/azure_ai_search.py,sha256=wzQsyr7ADDW7hHpZlMTnpBwQuPSOKaND8yc4SW9f6Pw,2653
19
19
  mem0/configs/vector_stores/baidu.py,sha256=YauoQZ2BgqzJg8hRAP8Ub5ntzYxhMuthX8dOvBCjS2I,1267
20
20
  mem0/configs/vector_stores/chroma.py,sha256=5mEJaNr8-rpwUd_Rx1t7mbaIFwMKezt0Yy72Pfn50Fc,1706
21
21
  mem0/configs/vector_stores/elasticsearch.py,sha256=IuAxw-wsV4f5pYfEQfRlMt5yVi1Qjdf8vhS0MuL3gxI,2364
@@ -38,7 +38,7 @@ mem0/dbs/configs.py,sha256=EoWnTNM_YPR39VgDDv8wS0KfA1nJ8Q0bpGtJdx1wu84,624
38
38
  mem0/dbs/mysql.py,sha256=AdrJQPrzqAN1cDjYgT-tA5n7CLPjJwR8LPlBhIJbT80,12731
39
39
  mem0/embeddings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
40
40
  mem0/embeddings/aws_bedrock.py,sha256=MvlmXfUxua2sIEoxzJKN9RPA5Xk9pN9peeDr3QBjuGw,3531
41
- mem0/embeddings/azure_openai.py,sha256=WU65mN6VQPROwmBguVwF8o8ydLFktZEjLmy9m6dQ660,1843
41
+ mem0/embeddings/azure_openai.py,sha256=Yj-xtQ4wZutfzI97KPQINwQYLpQ_oxjYUWn4OZR7eQ0,2687
42
42
  mem0/embeddings/base.py,sha256=O-oFuizpx9tyL8qjhRu0l42GDkjgORy_UMtMCkRfc4k,973
43
43
  mem0/embeddings/configs.py,sha256=ucqSXa_Wiiy_LQJRpulD_-VFQQc2tWOOzEV3pLrHMow,881
44
44
  mem0/embeddings/gemini.py,sha256=6WPlY1GRYSfX2aD0ywou3uM4XenEQm6oB1c7x7R5oOk,1509
@@ -60,7 +60,7 @@ mem0/graphs/neptune/main.py,sha256=JoyMpFqRe9RNLnh4f_HyfUaxwHtaDC6ZCdr0hqR8Vcw,1
60
60
  mem0/llms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
61
61
  mem0/llms/anthropic.py,sha256=LtKWVTAQxQnpiMrB0jwihn4mOO9RBjUT_vEjeXPBYhE,2250
62
62
  mem0/llms/aws_bedrock.py,sha256=awRKZvsYpab3jYFu0e3_gTRxAToF3-M1nm9Sbr1QwzA,10520
63
- mem0/llms/azure_openai.py,sha256=6l-sTZ6RvFbMyLrxLdDJFoanaM0lsTQn4X9nvpcz5aI,4170
63
+ mem0/llms/azure_openai.py,sha256=dAqgLV5G4N-SHC1rTDUFbWWOJMLqpgC6l6M9JCZ_Zqc,4973
64
64
  mem0/llms/azure_openai_structured.py,sha256=gWFt73wn7okFcJAI8_a7SSxcnoV8_Qrf9axiV5qKNf8,2722
65
65
  mem0/llms/base.py,sha256=Uit80Gs0MNKKH2z9h0Fkmr-4u4QATCpKFlPXj_8DBuo,1061
66
66
  mem0/llms/configs.py,sha256=yvtU5WDBsiygCiMHjhG9ZCb1NIAMv8uXYD_KL8_X54w,997
@@ -90,7 +90,7 @@ mem0/proxy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
90
90
  mem0/proxy/main.py,sha256=DhGH9dvVd3GTS1yyZzTBpyYvgArDRcYIgKLBCcosCLE,7398
91
91
  mem0/utils/factory.py,sha256=FZ-8qKK4K7unnEfcmcYAY77i2IVDK1nbcsB1WeGlR2I,5448
92
92
  mem0/vector_stores/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
93
- mem0/vector_stores/azure_ai_search.py,sha256=pnTzzYGfeU2Q_pm8D_o_uJu-uRDNMgoYERGu7HZZb6M,14804
93
+ mem0/vector_stores/azure_ai_search.py,sha256=rigXGyIp35yDn1RmHuAC0nYit_xQU0xDjxcV0rZNdwE,15686
94
94
  mem0/vector_stores/baidu.py,sha256=Cp7TGtMnZRkuDAnRw0F0OjNg229dH31dXYZL1JmPQes,12683
95
95
  mem0/vector_stores/base.py,sha256=yb9TCcxximzdJd8vfNJRROHBrUyYLbH26-E63fclcDo,1344
96
96
  mem0/vector_stores/chroma.py,sha256=fHGbvecYPvlhHrS4lkSWt4S1EnxicEWuqG6qOqDleQA,7539
@@ -109,8 +109,8 @@ mem0/vector_stores/supabase.py,sha256=RCegadYW4Q8CT_NSeCywUVW01TLJQT4mDzeStNHEJH
109
109
  mem0/vector_stores/upstash_vector.py,sha256=dAIef-bIB6AE_dVR8UpVWZUPQs-wLhqfxEHyyXO5kQU,9073
110
110
  mem0/vector_stores/vertex_ai_vector_search.py,sha256=RaaPLTTbTPftlg7nm9WTROOhurSxqnAyx-g3tYoxv0I,24659
111
111
  mem0/vector_stores/weaviate.py,sha256=v91CA79Z_zJBohwaXlOW6jtudDlGq0kYKQFHMHGqUy4,11833
112
- mem0ai_azure_mysql-0.1.115.data/data/README.md,sha256=HhoTn6hBxTIsFwbm5zHT7aNmQOa5yaTgwdRIF6V7hX4,6525
113
- mem0ai_azure_mysql-0.1.115.dist-info/METADATA,sha256=BTfI2Osc41_2mBpXlJAynVJsWOtUvq1MbPMILl52D5M,8948
114
- mem0ai_azure_mysql-0.1.115.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
115
- mem0ai_azure_mysql-0.1.115.dist-info/licenses/LICENSE,sha256=C7y-kxw1MpOi-vzggyYYHf7qDlaMVmr9TOgzenD14hk,11349
116
- mem0ai_azure_mysql-0.1.115.dist-info/RECORD,,
112
+ mem0ai_azure_mysql-0.1.115.2.data/data/README.md,sha256=HhoTn6hBxTIsFwbm5zHT7aNmQOa5yaTgwdRIF6V7hX4,6525
113
+ mem0ai_azure_mysql-0.1.115.2.dist-info/METADATA,sha256=VVr2yasEfvc9zVfzUDm8JKm6jLkTYouXKNO8DCLRDIY,8950
114
+ mem0ai_azure_mysql-0.1.115.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
115
+ mem0ai_azure_mysql-0.1.115.2.dist-info/licenses/LICENSE,sha256=C7y-kxw1MpOi-vzggyYYHf7qDlaMVmr9TOgzenD14hk,11349
116
+ mem0ai_azure_mysql-0.1.115.2.dist-info/RECORD,,