JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

merge: integrate main into feat/v3-pipeline + port AsyncMemory to v3

Browse Source 
Merge 190 commits from main including:
- limit→top_k rename across all vector store APIs
- custom_fact_extraction_prompt→custom_instructions rename
- enable_graph flag removal (use self.graph truthiness)
- US/Pacific→timezone.utc normalization
- 18 Category A bug fixes (partial updates, graph cleanup on delete,
  telemetry gating, metadata mutation prevention, filter operator
  merging, delete timestamps, actor_id preservation, etc.)

Conflict resolutions:
- openai.py: merged _pass_dimensions_to_api with embed_batch
- milvus.py: merged partial-update logic with BM25 text field
- qdrant.py: took main's enhanced _create_filter, merged BM25
  sparse vectors with partial-update endpoints
- test_main.py: adopted main's test structure, removed obsolete
  test_custom_prompts, kept hybrid search over-fetch assertion

AsyncMemory v3 port (full parity):
- Rewrote AsyncMemory._add_to_vector_store with the same 8-phase
  batch pipeline as sync Memory (single LLM call, batch embed,
  hash dedup, batch persist, batch entity linking, message saving)
- All blocking calls wrapped in await asyncio.to_thread()
- Moved _build_session_scope to module level for shared access

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Soumil Rathi
2026-04-12 16:26:10 -07:00
co-authored by Claude Opus 4.6
parent 90f9e6a550 9d6b79a14e
commit b9e26acf4d
655 changed files with 75693 additions and 8779 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0/client/main.py
+206 -217
View File
@@ -2,13 +2,22 @@ import hashlib
import logging
import os
import warnings
from typing import Any, Dict, List, Optional, Union
from typing import Any, Dict, List, Optional
import httpx
import requests
from mem0.client.project import AsyncProject, Project
from mem0.client.types import (
AddMemoryOptions,
DeleteAllMemoryOptions,
GetAllMemoryOptions,
ProjectUpdateOptions,
SearchMemoryOptions,
UpdateMemoryOptions,
)
from mem0.client.utils import api_error_handler
# Exception classes are referenced in docstrings only
from mem0.memory.setup import get_user_id, setup_config
from mem0.memory.telemetry import capture_client_event
@@ -31,8 +40,6 @@ class MemoryClient:
api_key (str): The API key for authenticating with the Mem0 API.
host (str): The base URL for the Mem0 API.
client (httpx.Client): The HTTP client used for making API requests.
org_id (str, optional): Organization ID.
project_id (str, optional): Project ID.
user_id (str): Unique identifier for the user.
"""
@@ -40,8 +47,6 @@ class MemoryClient:
self,
api_key: Optional[str] = None,
host: Optional[str] = None,
org_id: Optional[str] = None,
project_id: Optional[str] = None,
client: Optional[httpx.Client] = None,
):
"""Initialize the MemoryClient.
@@ -52,8 +57,6 @@ class MemoryClient:
environment variable.
host: The base URL for the Mem0 API. Defaults to
"https://api.mem0.ai".
org_id: The ID of the organization.
project_id: The ID of the project.
client: A custom httpx.Client instance. If provided, it will be
used instead of creating a new one. Note that base_url and
headers will be set/overridden as needed.
@@ -63,8 +66,8 @@ class MemoryClient:
"""
self.api_key = api_key or os.getenv("MEM0_API_KEY")
self.host = host or "https://api.mem0.ai"
self.org_id = org_id
self.project_id = project_id
self.org_id = None
self.project_id = None
self.user_id = get_user_id()
if not self.api_key:
@@ -128,15 +131,16 @@ class MemoryClient:
raise ValueError(f"Error: {error_message}")
@api_error_handler
def add(self, messages, **kwargs) -> Dict[str, Any]:
def add(self, messages, options: Optional[AddMemoryOptions] = None, **kwargs) -> Dict[str, Any]:
"""Add a new memory.
Args:
messages: A list of message dictionaries, a single message dictionary,
or a string. If a string is provided, it will be converted to
a user message.
options: Typed options for the add operation (AddMemoryOptions).
**kwargs: Additional parameters such as user_id, agent_id, app_id,
metadata, filters, async_mode.
metadata, filters.
Returns:
A dictionary containing the API response in v1.1 format.
@@ -149,24 +153,19 @@ class MemoryClient:
NetworkError: If network connectivity issues occur.
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
"""
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
# Handle different message input formats (align with OSS behavior)
if isinstance(messages, str):
messages = [{"role": "user", "content": messages}]
elif isinstance(messages, dict):
messages = [messages]
elif not isinstance(messages, list):
raise ValueError(
f"messages must be str, dict, or list[dict], got {type(messages).__name__}"
)
kwargs = self._prepare_params(kwargs)
# Set async_mode to True by default, but allow user override
if "async_mode" not in kwargs:
kwargs["async_mode"] = True
raise ValueError(f"messages must be str, dict, or list[dict], got {type(messages).__name__}")
# Force v1.1 format for all add operations
kwargs["output_format"] = "v1.1"
kwargs = self._prepare_params(kwargs)
payload = self._prepare_payload(messages, kwargs)
response = self.client.post("/v1/memories/", json=payload)
response.raise_for_status()
@@ -200,10 +199,11 @@ class MemoryClient:
return response.json()
@api_error_handler
def get_all(self, **kwargs) -> Dict[str, Any]:
def get_all(self, options: Optional[GetAllMemoryOptions] = None, **kwargs) -> Dict[str, Any]:
"""Retrieve all memories, with optional filtering.
Args:
options: Typed options for the get_all operation (GetAllMemoryOptions).
**kwargs: Optional parameters for filtering (user_id, agent_id,
app_id, top_k, page, page_size).
@@ -218,8 +218,8 @@ class MemoryClient:
NetworkError: If network connectivity issues occur.
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
"""
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
params = self._prepare_params(kwargs)
params.pop("async_mode", None)
if "page" in params and "page_size" in params:
query_params = {
@@ -236,7 +236,6 @@ class MemoryClient:
"client.get_all",
self,
{
"api_version": "v2",
"keys": list(kwargs.keys()),
"sync_type": "sync",
},
@@ -249,13 +248,13 @@ class MemoryClient:
return result
@api_error_handler
def search(self, query: str, **kwargs) -> Dict[str, Any]:
def search(self, query: str, options: Optional[SearchMemoryOptions] = None, **kwargs) -> Dict[str, Any]:
"""Search memories based on a query.
Args:
query: The search query string.
**kwargs: Additional parameters such as user_id, agent_id, app_id,
top_k, filters.
options: Typed options for the search operation (SearchMemoryOptions).
**kwargs: Additional parameters such as filters, top_k, rerank.
Returns:
A dictionary containing search results in v1.1 format: {"results": [...]}
@@ -268,11 +267,9 @@ class MemoryClient:
NetworkError: If network connectivity issues occur.
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
"""
payload = {"query": query}
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
params = self._prepare_params(kwargs)
params.pop("async_mode", None)
payload.update(params)
payload = {"query": query, **params}
response = self.client.post("/v2/memories/search/", json=payload)
response.raise_for_status()
@@ -282,7 +279,6 @@ class MemoryClient:
"client.search",
self,
{
"api_version": "v2",
"keys": list(kwargs.keys()),
"sync_type": "sync",
},
@@ -298,36 +294,32 @@ class MemoryClient:
def update(
self,
memory_id: str,
text: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
timestamp: Optional[Union[int, float, str]] = None,
options: Optional[UpdateMemoryOptions] = None,
**kwargs,
) -> Dict[str, Any]:
"""
Update a memory by ID.
"""Update a memory by ID.
Args:
memory_id (str): Memory ID.
text (str, optional): New content to update the memory with.
metadata (dict, optional): Metadata to update in the memory.
timestamp (int, float, or str, optional): Unix epoch timestamp or ISO 8601 string.
memory_id: The ID of the memory to update.
options: Typed options (UpdateMemoryOptions) with text, metadata,
and/or timestamp fields.
**kwargs: Alternatively pass text, metadata, timestamp as keyword args.
Returns:
Dict[str, Any]: The response from the server.
Example:
>>> client.update(memory_id="mem_123", text="Likes to play tennis on weekends")
>>> client.update(memory_id="mem_123", timestamp="2025-01-15T12:00:00Z")
"""
if text is None and metadata is None and timestamp is None:
raise ValueError("At least one of text, metadata, or timestamp must be provided for update.")
Raises:
ValueError: If none of text, metadata, or timestamp are provided.
payload = {}
if text is not None:
payload["text"] = text
if metadata is not None:
payload["metadata"] = metadata
if timestamp is not None:
payload["timestamp"] = timestamp
Example:
>>> client.update("mem_123", UpdateMemoryOptions(text="Updated text"))
>>> client.update("mem_123", text="Updated text")
"""
payload = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
payload = {k: v for k, v in payload.items() if v is not None}
if not payload:
raise ValueError("At least one of text, metadata, or timestamp must be provided for update.")
capture_client_event("client.update", self, {"memory_id": memory_id, "sync_type": "sync"})
params = self._prepare_params()
@@ -360,10 +352,11 @@ class MemoryClient:
return response.json()
@api_error_handler
def delete_all(self, **kwargs) -> Dict[str, str]:
def delete_all(self, options: Optional[DeleteAllMemoryOptions] = None, **kwargs) -> Dict[str, str]:
"""Delete all memories, with optional filtering.
Args:
options: Typed options for the delete_all operation (DeleteAllMemoryOptions).
**kwargs: Optional parameters for filtering (user_id, agent_id,
app_id).
@@ -378,6 +371,7 @@ class MemoryClient:
NetworkError: If network connectivity issues occur.
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
"""
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
params = self._prepare_params(kwargs)
response = self.client.delete("/v1/memories/", params=params)
response.raise_for_status()
@@ -667,77 +661,62 @@ class MemoryClient:
@api_error_handler
def update_project(
self,
options: Optional[ProjectUpdateOptions] = None,
custom_instructions: Optional[str] = None,
custom_categories: Optional[List[str]] = None,
retrieval_criteria: Optional[List[Dict[str, Any]]] = None,
enable_graph: Optional[bool] = None,
version: Optional[str] = None,
inclusion_prompt: Optional[str] = None,
exclusion_prompt: Optional[str] = None,
memory_depth: Optional[str] = None,
usecase_setting: Optional[str] = None,
multilingual: Optional[bool] = None,
) -> Dict[str, Any]:
"""Update the project settings.
Args:
custom_instructions: New instructions for the project
custom_categories: New categories for the project
retrieval_criteria: New retrieval criteria for the project
enable_graph: Enable or disable the graph for the project
version: Version of the project
inclusion_prompt: Inclusion prompt for the project
exclusion_prompt: Exclusion prompt for the project
memory_depth: Memory depth for the project
usecase_setting: Usecase setting for the project
options: Typed options for the update operation (ProjectUpdateOptions).
custom_instructions: New instructions for the project.
custom_categories: New categories for the project.
retrieval_criteria: New retrieval criteria for the project.
memory_depth: Memory depth for the project.
usecase_setting: Usecase setting for the project.
multilingual: Whether to use the input language for memory storage and retrieval.
Returns:
Dictionary containing the API response.
Raises:
ValidationError: If the input data is invalid.
AuthenticationError: If authentication fails.
RateLimitError: If rate limits are exceeded.
MemoryQuotaExceededError: If memory quota is exceeded.
NetworkError: If network connectivity issues occur.
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
ValueError: If org_id or project_id are not set.
ValueError: If org_id or project_id are not set, or no update fields provided.
"""
logger.warning(
"update_project() method is going to be deprecated in version v1.0 of the package. Please use the client.project.update() method instead."
"update_project() method is going to be deprecated in version v1.0 of the package. "
"Please use the client.project.update() method instead."
)
if not (self.org_id and self.project_id):
raise ValueError("org_id and project_id must be set to update instructions or categories")
if (
custom_instructions is None
and custom_categories is None
and retrieval_criteria is None
and enable_graph is None
and version is None
and inclusion_prompt is None
and exclusion_prompt is None
and memory_depth is None
and usecase_setting is None
):
kwargs = {
**(options.model_dump(exclude_unset=True) if options else {}),
**{
k: v
for k, v in {
"custom_instructions": custom_instructions,
"custom_categories": custom_categories,
"retrieval_criteria": retrieval_criteria,
"memory_depth": memory_depth,
"usecase_setting": usecase_setting,
"multilingual": multilingual,
}.items()
if v is not None
},
}
if not kwargs:
raise ValueError(
"Currently we only support updating custom_instructions or "
"custom_categories or retrieval_criteria, so you must "
"provide at least one of them"
)
payload = self._prepare_params(
{
"custom_instructions": custom_instructions,
"custom_categories": custom_categories,
"retrieval_criteria": retrieval_criteria,
"enable_graph": enable_graph,
"version": version,
"inclusion_prompt": inclusion_prompt,
"exclusion_prompt": exclusion_prompt,
"memory_depth": memory_depth,
"usecase_setting": usecase_setting,
}
)
payload = self._prepare_params(kwargs)
response = self.client.patch(
f"/api/v1/orgs/organizations/{self.org_id}/projects/{self.project_id}/",
json=payload,
@@ -746,18 +725,7 @@ class MemoryClient:
capture_client_event(
"client.update_project",
self,
{
"custom_instructions": custom_instructions,
"custom_categories": custom_categories,
"retrieval_criteria": retrieval_criteria,
"enable_graph": enable_graph,
"version": version,
"inclusion_prompt": inclusion_prompt,
"exclusion_prompt": exclusion_prompt,
"memory_depth": memory_depth,
"usecase_setting": usecase_setting,
"sync_type": "sync",
},
{**kwargs, "sync_type": "sync"},
)
return response.json()
@@ -932,20 +900,12 @@ class MemoryClient:
Returns:
A dictionary containing the prepared parameters.
Raises:
ValueError: If either org_id or project_id is provided but not both.
"""
if kwargs is None:
kwargs = {}
# Add org_id and project_id if both are available
if self.org_id and self.project_id:
kwargs["org_id"] = self.org_id
kwargs["project_id"] = self.project_id
elif self.org_id or self.project_id:
raise ValueError("Please provide both org_id and project_id")
# org_id and project_id are resolved from API key — not injected into params
return {k: v for k, v in kwargs.items() if v is not None}
@@ -961,8 +921,6 @@ class AsyncMemoryClient:
self,
api_key: Optional[str] = None,
host: Optional[str] = None,
org_id: Optional[str] = None,
project_id: Optional[str] = None,
client: Optional[httpx.AsyncClient] = None,
):
"""Initialize the AsyncMemoryClient.
@@ -973,8 +931,6 @@ class AsyncMemoryClient:
environment variable.
host: The base URL for the Mem0 API. Defaults to
"https://api.mem0.ai".
org_id: The ID of the organization.
project_id: The ID of the project.
client: A custom httpx.AsyncClient instance. If provided, it will
be used instead of creating a new one. Note that base_url
and headers will be set/overridden as needed.
@@ -984,8 +940,8 @@ class AsyncMemoryClient:
"""
self.api_key = api_key or os.getenv("MEM0_API_KEY")
self.host = host or "https://api.mem0.ai"
self.org_id = org_id
self.project_id = project_id
self.org_id = None
self.project_id = None
self.user_id = get_user_id()
if not self.api_key:
@@ -1080,20 +1036,12 @@ class AsyncMemoryClient:
Returns:
A dictionary containing the prepared parameters.
Raises:
ValueError: If either org_id or project_id is provided but not both.
"""
if kwargs is None:
kwargs = {}
# Add org_id and project_id if both are available
if self.org_id and self.project_id:
kwargs["org_id"] = self.org_id
kwargs["project_id"] = self.project_id
elif self.org_id or self.project_id:
raise ValueError("Please provide both org_id and project_id")
# org_id and project_id are resolved from API key — not injected into params
return {k: v for k, v in kwargs.items() if v is not None}
@@ -1104,25 +1052,41 @@ class AsyncMemoryClient:
await self.async_client.aclose()
@api_error_handler
async def add(self, messages, **kwargs) -> Dict[str, Any]:
async def add(self, messages, options: Optional[AddMemoryOptions] = None, **kwargs) -> Dict[str, Any]:
"""Add a new memory.
Args:
messages: A list of message dictionaries, a single message dictionary,
or a string. If a string is provided, it will be converted to
a user message.
options: Typed options for the add operation (AddMemoryOptions).
**kwargs: Additional parameters such as user_id, agent_id, app_id,
metadata, filters.
Returns:
A dictionary containing the API response in v1.1 format.
Raises:
ValidationError: If the input data is invalid.
AuthenticationError: If authentication fails.
RateLimitError: If rate limits are exceeded.
MemoryQuotaExceededError: If memory quota is exceeded.
NetworkError: If network connectivity issues occur.
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
"""
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
# Handle different message input formats (align with OSS behavior)
if isinstance(messages, str):
messages = [{"role": "user", "content": messages}]
elif isinstance(messages, dict):
messages = [messages]
elif not isinstance(messages, list):
raise ValueError(
f"messages must be str, dict, or list[dict], got {type(messages).__name__}"
)
kwargs = self._prepare_params(kwargs)
# Set async_mode to True by default, but allow user override
if "async_mode" not in kwargs:
kwargs["async_mode"] = True
raise ValueError(f"messages must be str, dict, or list[dict], got {type(messages).__name__}")
# Force v1.1 format for all add operations
kwargs["output_format"] = "v1.1"
kwargs = self._prepare_params(kwargs)
payload = self._prepare_payload(messages, kwargs)
response = await self.async_client.post("/v1/memories/", json=payload)
response.raise_for_status()
@@ -1140,9 +1104,26 @@ class AsyncMemoryClient:
return response.json()
@api_error_handler
async def get_all(self, **kwargs) -> Dict[str, Any]:
async def get_all(self, options: Optional[GetAllMemoryOptions] = None, **kwargs) -> Dict[str, Any]:
"""Retrieve all memories, with optional filtering.
Args:
options: Typed options for the get_all operation (GetAllMemoryOptions).
**kwargs: Optional parameters for filtering (filters, page, page_size).
Returns:
A dictionary containing memories in v1.1 format: {"results": [...]}
Raises:
ValidationError: If the input data is invalid.
AuthenticationError: If authentication fails.
RateLimitError: If rate limits are exceeded.
MemoryQuotaExceededError: If memory quota is exceeded.
NetworkError: If network connectivity issues occur.
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
"""
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
params = self._prepare_params(kwargs)
params.pop("async_mode", None)
if "page" in params and "page_size" in params:
query_params = {
@@ -1159,7 +1140,6 @@ class AsyncMemoryClient:
"client.get_all",
self,
{
"api_version": "v2",
"keys": list(kwargs.keys()),
"sync_type": "async",
},
@@ -1172,12 +1152,28 @@ class AsyncMemoryClient:
return result
@api_error_handler
async def search(self, query: str, **kwargs) -> Dict[str, Any]:
payload = {"query": query}
params = self._prepare_params(kwargs)
params.pop("async_mode", None)
async def search(self, query: str, options: Optional[SearchMemoryOptions] = None, **kwargs) -> Dict[str, Any]:
"""Search memories based on a query.
payload.update(params)
Args:
query: The search query string.
options: Typed options for the search operation (SearchMemoryOptions).
**kwargs: Additional parameters such as filters, top_k, rerank.
Returns:
A dictionary containing search results in v1.1 format: {"results": [...]}
Raises:
ValidationError: If the input data is invalid.
AuthenticationError: If authentication fails.
RateLimitError: If rate limits are exceeded.
MemoryQuotaExceededError: If memory quota is exceeded.
NetworkError: If network connectivity issues occur.
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
"""
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
params = self._prepare_params(kwargs)
payload = {"query": query, **params}
response = await self.async_client.post("/v2/memories/search/", json=payload)
response.raise_for_status()
@@ -1187,7 +1183,6 @@ class AsyncMemoryClient:
"client.search",
self,
{
"api_version": "v2",
"keys": list(kwargs.keys()),
"sync_type": "async",
},
@@ -1203,36 +1198,32 @@ class AsyncMemoryClient:
async def update(
self,
memory_id: str,
text: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
timestamp: Optional[Union[int, float, str]] = None,
options: Optional[UpdateMemoryOptions] = None,
**kwargs,
) -> Dict[str, Any]:
"""
Update a memory by ID asynchronously.
"""Update a memory by ID asynchronously.
Args:
memory_id (str): Memory ID.
text (str, optional): New content to update the memory with.
metadata (dict, optional): Metadata to update in the memory.
timestamp (int, float, or str, optional): Unix epoch timestamp or ISO 8601 string.
memory_id: The ID of the memory to update.
options: Typed options (UpdateMemoryOptions) with text, metadata,
and/or timestamp fields.
**kwargs: Alternatively pass text, metadata, timestamp as keyword args.
Returns:
Dict[str, Any]: The response from the server.
Example:
>>> await client.update(memory_id="mem_123", text="Likes to play tennis on weekends")
>>> await client.update(memory_id="mem_123", timestamp="2025-01-15T12:00:00Z")
"""
if text is None and metadata is None and timestamp is None:
raise ValueError("At least one of text, metadata, or timestamp must be provided for update.")
Raises:
ValueError: If none of text, metadata, or timestamp are provided.
payload = {}
if text is not None:
payload["text"] = text
if metadata is not None:
payload["metadata"] = metadata
if timestamp is not None:
payload["timestamp"] = timestamp
Example:
>>> await client.update("mem_123", UpdateMemoryOptions(text="Updated text"))
>>> await client.update("mem_123", text="Updated text")
"""
payload = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
payload = {k: v for k, v in payload.items() if v is not None}
if not payload:
raise ValueError("At least one of text, metadata, or timestamp must be provided for update.")
capture_client_event("client.update", self, {"memory_id": memory_id, "sync_type": "async"})
params = self._prepare_params()
@@ -1265,10 +1256,11 @@ class AsyncMemoryClient:
return response.json()
@api_error_handler
async def delete_all(self, **kwargs) -> Dict[str, str]:
async def delete_all(self, options: Optional[DeleteAllMemoryOptions] = None, **kwargs) -> Dict[str, str]:
"""Delete all memories, with optional filtering.
Args:
options: Typed options for the delete_all operation (DeleteAllMemoryOptions).
**kwargs: Optional parameters for filtering (user_id, agent_id, app_id).
Returns:
@@ -1282,6 +1274,7 @@ class AsyncMemoryClient:
NetworkError: If network connectivity issues occur.
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
"""
kwargs = {**(options.model_dump(exclude_unset=True) if options else {}), **kwargs}
params = self._prepare_params(kwargs)
response = await self.async_client.delete("/v1/memories/", params=params)
response.raise_for_status()
@@ -1549,59 +1542,62 @@ class AsyncMemoryClient:
@api_error_handler
async def update_project(
self,
options: Optional[ProjectUpdateOptions] = None,
custom_instructions: Optional[str] = None,
custom_categories: Optional[List[str]] = None,
retrieval_criteria: Optional[List[Dict[str, Any]]] = None,
enable_graph: Optional[bool] = None,
version: Optional[str] = None,
memory_depth: Optional[str] = None,
usecase_setting: Optional[str] = None,
multilingual: Optional[bool] = None,
) -> Dict[str, Any]:
"""Update the project settings.
Args:
custom_instructions: New instructions for the project
custom_categories: New categories for the project
retrieval_criteria: New retrieval criteria for the project
enable_graph: Enable or disable the graph for the project
version: Version of the project
options: Typed options for the update operation (ProjectUpdateOptions).
custom_instructions: New instructions for the project.
custom_categories: New categories for the project.
retrieval_criteria: New retrieval criteria for the project.
memory_depth: Memory depth for the project.
usecase_setting: Usecase setting for the project.
multilingual: Whether to use the input language for memory storage and retrieval.
Returns:
Dictionary containing the API response.
Raises:
ValidationError: If the input data is invalid.
AuthenticationError: If authentication fails.
RateLimitError: If rate limits are exceeded.
MemoryQuotaExceededError: If memory quota is exceeded.
NetworkError: If network connectivity issues occur.
MemoryNotFoundError: If the memory doesn't exist (for updates/deletes).
ValueError: If org_id or project_id are not set.
ValueError: If org_id or project_id are not set, or no update fields provided.
"""
logger.warning(
"update_project() method is going to be deprecated in version v1.0 of the package. Please use the client.project.update() method instead."
"update_project() method is going to be deprecated in version v1.0 of the package. "
"Please use the client.project.update() method instead."
)
if not (self.org_id and self.project_id):
raise ValueError("org_id and project_id must be set to update instructions or categories")
if (
custom_instructions is None
and custom_categories is None
and retrieval_criteria is None
and enable_graph is None
and version is None
):
kwargs = {
**(options.model_dump(exclude_unset=True) if options else {}),
**{
k: v
for k, v in {
"custom_instructions": custom_instructions,
"custom_categories": custom_categories,
"retrieval_criteria": retrieval_criteria,
"memory_depth": memory_depth,
"usecase_setting": usecase_setting,
"multilingual": multilingual,
}.items()
if v is not None
},
}
if not kwargs:
raise ValueError(
"Currently we only support updating custom_instructions or custom_categories or retrieval_criteria, so you must provide at least one of them"
"Currently we only support updating custom_instructions or "
"custom_categories or retrieval_criteria, so you must "
"provide at least one of them"
)
payload = self._prepare_params(
{
"custom_instructions": custom_instructions,
"custom_categories": custom_categories,
"retrieval_criteria": retrieval_criteria,
"enable_graph": enable_graph,
"version": version,
}
)
payload = self._prepare_params(kwargs)
response = await self.async_client.patch(
f"/api/v1/orgs/organizations/{self.org_id}/projects/{self.project_id}/",
json=payload,
@@ -1610,14 +1606,7 @@ class AsyncMemoryClient:
capture_client_event(
"client.update_project",
self,
{
"custom_instructions": custom_instructions,
"custom_categories": custom_categories,
"retrieval_criteria": retrieval_criteria,
"enable_graph": enable_graph,
"version": version,
"sync_type": "async",
},
{**kwargs, "sync_type": "async"},
)
return response.json()
mem0/client/project.py
+13 -14
View File
@@ -7,6 +7,7 @@ from pydantic import BaseModel, ConfigDict, Field
from mem0.client.utils import api_error_handler
from mem0.memory.telemetry import capture_client_event
# Exception classes are referenced in docstrings only
logger = logging.getLogger(__name__)
@@ -177,7 +178,6 @@ class BaseProject(ABC):
custom_instructions: Optional[str] = None,
custom_categories: Optional[List[str]] = None,
retrieval_criteria: Optional[List[Dict[str, Any]]] = None,
enable_graph: Optional[bool] = None,
) -> Dict[str, Any]:
"""
Update project settings.
@@ -186,7 +186,6 @@ class BaseProject(ABC):
custom_instructions: New instructions for the project
custom_categories: New categories for the project
retrieval_criteria: New retrieval criteria for the project
enable_graph: Enable or disable the graph for the project
Returns:
Dictionary containing the API response.
@@ -398,7 +397,7 @@ class Project(BaseProject):
custom_instructions: Optional[str] = None,
custom_categories: Optional[List[str]] = None,
retrieval_criteria: Optional[List[Dict[str, Any]]] = None,
enable_graph: Optional[bool] = None,
multilingual: Optional[bool] = None,
) -> Dict[str, Any]:
"""
Update project settings.
@@ -407,7 +406,7 @@ class Project(BaseProject):
custom_instructions: New instructions for the project
custom_categories: New categories for the project
retrieval_criteria: New retrieval criteria for the project
enable_graph: Enable or disable the graph for the project
multilingual: Whether to use the input language for memory storage and retrieval
Returns:
Dictionary containing the API response.
@@ -423,12 +422,12 @@ class Project(BaseProject):
custom_instructions is None
and custom_categories is None
and retrieval_criteria is None
and enable_graph is None
and multilingual is None
):
raise ValueError(
"At least one parameter must be provided for update: "
"custom_instructions, custom_categories, retrieval_criteria, "
"enable_graph"
"multilingual"
)
payload = self._prepare_params(
@@ -436,7 +435,7 @@ class Project(BaseProject):
"custom_instructions": custom_instructions,
"custom_categories": custom_categories,
"retrieval_criteria": retrieval_criteria,
"enable_graph": enable_graph,
"multilingual": multilingual,
}
)
response = self._client.patch(
@@ -451,7 +450,7 @@ class Project(BaseProject):
"custom_instructions": custom_instructions,
"custom_categories": custom_categories,
"retrieval_criteria": retrieval_criteria,
"enable_graph": enable_graph,
"multilingual": multilingual,
"sync_type": "sync",
},
)
@@ -715,7 +714,7 @@ class AsyncProject(BaseProject):
custom_instructions: Optional[str] = None,
custom_categories: Optional[List[str]] = None,
retrieval_criteria: Optional[List[Dict[str, Any]]] = None,
enable_graph: Optional[bool] = None,
multilingual: Optional[bool] = None,
) -> Dict[str, Any]:
"""
Update project settings.
@@ -724,7 +723,7 @@ class AsyncProject(BaseProject):
custom_instructions: New instructions for the project
custom_categories: New categories for the project
retrieval_criteria: New retrieval criteria for the project
enable_graph: Enable or disable the graph for the project
multilingual: Whether to use the input language for memory storage and retrieval
Returns:
Dictionary containing the API response.
@@ -740,12 +739,12 @@ class AsyncProject(BaseProject):
custom_instructions is None
and custom_categories is None
and retrieval_criteria is None
and enable_graph is None
and multilingual is None
):
raise ValueError(
"At least one parameter must be provided for update: "
"custom_instructions, custom_categories, retrieval_criteria, "
"enable_graph"
"multilingual"
)
payload = self._prepare_params(
@@ -753,7 +752,7 @@ class AsyncProject(BaseProject):
"custom_instructions": custom_instructions,
"custom_categories": custom_categories,
"retrieval_criteria": retrieval_criteria,
"enable_graph": enable_graph,
"multilingual": multilingual,
}
)
response = await self._client.patch(
@@ -768,7 +767,7 @@ class AsyncProject(BaseProject):
"custom_instructions": custom_instructions,
"custom_categories": custom_categories,
"retrieval_criteria": retrieval_criteria,
"enable_graph": enable_graph,
"multilingual": multilingual,
"sync_type": "async",
},
)
mem0/client/types.py
+99
View File
@@ -0,0 +1,99 @@
"""Pydantic option models for MemoryClient methods.
These models provide IDE autocompletion, runtime validation, and type safety.
Methods accept both typed options and **kwargs for backward compatibility.
"""
from typing import Any, Dict, List, Optional, Union
from pydantic import BaseModel, Field
class EntityOptions(BaseModel):
"""Identity options for add/delete operations (top-level entity IDs)."""
user_id: Optional[str] = Field(default=None, description="The user ID to associate with the memory")
agent_id: Optional[str] = Field(default=None, description="The agent ID to associate with the memory")
app_id: Optional[str] = Field(default=None, description="The app ID to associate with the memory")
run_id: Optional[str] = Field(default=None, description="The run ID to associate with the memory")
class AddMemoryOptions(EntityOptions):
"""Options for the add() method."""
metadata: Optional[Dict[str, Any]] = Field(default=None, description="Additional metadata for the memory")
infer: Optional[bool] = Field(default=None, description="Whether to infer memories from the input")
custom_categories: Optional[List[Dict[str, Any]]] = Field(
default=None, description="Custom categories for memory classification"
)
custom_instructions: Optional[str] = Field(default=None, description="Custom instructions for fact extraction")
timestamp: Optional[int] = Field(default=None, description="Unix timestamp for the memory")
structured_data_schema: Optional[Dict[str, Any]] = Field(
default=None, description="Schema for structured data extraction"
)
class SearchMemoryOptions(BaseModel):
"""Options for the search() method.
Identity fields (user_id, agent_id, etc.) must be passed inside the
``filters`` dict — the v2 API does not accept them at the top level.
"""
filters: Optional[Dict[str, Any]] = Field(
default=None, description="Filters for the search (e.g. {'user_id': '...'})"
)
metadata: Optional[Dict[str, Any]] = Field(default=None, description="Additional metadata for the search")
top_k: Optional[int] = Field(default=None, description="Number of results to return")
rerank: Optional[bool] = Field(default=None, description="Whether to rerank results")
threshold: Optional[float] = Field(default=None, description="Minimum similarity score threshold")
fields: Optional[List[str]] = Field(default=None, description="Fields to include in the response")
categories: Optional[List[str]] = Field(default=None, description="Categories to filter by")
class GetAllMemoryOptions(BaseModel):
"""Options for the get_all() method.
Identity fields (user_id, agent_id, etc.) must be passed inside the
``filters`` dict — the v2 API does not accept them at the top level.
"""
filters: Optional[Dict[str, Any]] = Field(
default=None, description="Filters for retrieval (e.g. {'user_id': '...'})"
)
page: Optional[int] = Field(default=None, description="Page number for pagination")
page_size: Optional[int] = Field(default=None, description="Number of items per page")
start_date: Optional[str] = Field(
default=None, description="Filter memories created on or after this date (ISO 8601)"
)
end_date: Optional[str] = Field(
default=None, description="Filter memories created on or before this date (ISO 8601)"
)
categories: Optional[List[str]] = Field(default=None, description="Categories to filter by")
class DeleteAllMemoryOptions(EntityOptions):
"""Options for the delete_all() method."""
pass
class UpdateMemoryOptions(BaseModel):
"""Options for the update() method."""
text: Optional[str] = Field(default=None, description="New text content for the memory")
metadata: Optional[Dict[str, Any]] = Field(default=None, description="Updated metadata")
timestamp: Optional[Union[int, float, str]] = Field(default=None, description="Updated timestamp")
class ProjectUpdateOptions(BaseModel):
"""Options for project update operations."""
custom_instructions: Optional[str] = Field(default=None, description="Custom instructions for fact extraction")
custom_categories: Optional[List[Dict[str, Any]]] = Field(
default=None, description="Custom categories for classification"
)
memory_depth: Optional[str] = Field(default=None, description="Memory depth configuration")
usecase_setting: Optional[Any] = Field(default=None, description="Use case specific settings")
multilingual: Optional[bool] = Field(default=None, description="Whether to enable multilingual support")
retrieval_criteria: Optional[List[Any]] = Field(default=None, description="Criteria for memory retrieval")
mem0/configs/base.py
+3 -7
View File
@@ -3,11 +3,11 @@ from typing import Any, Dict, Optional
from pydantic import BaseModel, Field
from mem0.configs.rerankers.config import RerankerConfig
from mem0.embeddings.configs import EmbedderConfig
from mem0.graphs.configs import GraphStoreConfig
from mem0.llms.configs import LlmConfig
from mem0.vector_stores.configs import VectorStoreConfig
from mem0.configs.rerankers.config import RerankerConfig
# Set up the directory path
home_dir = os.path.expanduser("~")
@@ -56,18 +56,14 @@ class MemoryConfig(BaseModel):
description="The version of the API",
default="v1.1",
)
custom_fact_extraction_prompt: Optional[str] = Field(
description="Custom prompt for the fact extraction",
custom_instructions: Optional[str] = Field(
description="Custom instructions for fact extraction",
default=None,
)
custom_update_memory_prompt: Optional[str] = Field(
description="Custom prompt for the update memory (deprecated: use custom_instructions)",
default=None,
)
custom_instructions: Optional[str] = Field(
description="Custom instructions injected into the extraction prompt",
default=None,
)
class AzureConfig(BaseModel):
mem0/configs/llms/anthropic.py
+2 -2
View File
@@ -16,7 +16,7 @@ class AnthropicConfig(BaseLlmConfig):
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_p: Optional[float] = None,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
@@ -32,7 +32,7 @@ class AnthropicConfig(BaseLlmConfig):
temperature: Controls randomness, defaults to 0.1
api_key: Anthropic API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_p: Nucleus sampling parameter, defaults to None (omitted to avoid conflict with temperature)
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
mem0/configs/llms/aws_bedrock.py
+8 -4
View File
@@ -16,7 +16,7 @@ class AWSBedrockConfig(BaseLlmConfig):
model: Optional[str] = None,
temperature: float = 0.1,
max_tokens: int = 2000,
top_p: float = 0.9,
top_p: Optional[float] = None,
top_k: int = 1,
aws_access_key_id: Optional[str] = None,
aws_secret_access_key: Optional[str] = None,
@@ -33,7 +33,8 @@ class AWSBedrockConfig(BaseLlmConfig):
model: Bedrock model identifier (e.g., "amazon.nova-3-mini-20241119-v1:0")
temperature: Controls randomness (0.0 to 2.0)
max_tokens: Maximum tokens to generate
top_p: Nucleus sampling parameter (0.0 to 1.0)
top_p: Nucleus sampling (0.0–1.0). Default None omits topP on Converse
(required for Anthropic, which rejects temperature and topP together).
top_k: Top-k sampling parameter (1 to 40)
aws_access_key_id: AWS access key (optional, uses env vars if not provided)
aws_secret_access_key: AWS secret key (optional, uses env vars if not provided)
@@ -75,13 +76,16 @@ class AWSBedrockConfig(BaseLlmConfig):
def get_model_config(self) -> Dict[str, Any]:
"""Get model-specific configuration parameters."""
base_config = {
base_config: Dict[str, Any] = {
"temperature": self.temperature,
"max_tokens": self.max_tokens,
"top_p": self.top_p,
"top_k": self.top_k,
}
# Only include top_p when explicitly set by the user.
if self.top_p is not None:
base_config["top_p"] = self.top_p
# Add custom model kwargs
base_config.update(self.model_kwargs)
mem0/configs/llms/azure.py
+3
View File
@@ -21,6 +21,7 @@ class AzureOpenAIConfig(BaseLlmConfig):
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
reasoning_effort: Optional[str] = None,
http_client_proxies: Optional[dict] = None,
# Azure OpenAI-specific parameters
azure_kwargs: Optional[Dict[str, Any]] = None,
@@ -37,6 +38,7 @@ class AzureOpenAIConfig(BaseLlmConfig):
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
reasoning_effort: Effort level for reasoning models ("low", "medium", "high"), defaults to None
http_client_proxies: HTTP client proxy settings, defaults to None
azure_kwargs: Azure-specific configuration, defaults to None
"""
@@ -50,6 +52,7 @@ class AzureOpenAIConfig(BaseLlmConfig):
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
reasoning_effort=reasoning_effort,
http_client_proxies=http_client_proxies,
)
mem0/configs/llms/base.py
+5
View File
@@ -23,6 +23,7 @@ class BaseLlmConfig(ABC):
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
reasoning_effort: Optional[str] = None,
http_client_proxies: Optional[Union[Dict, str]] = None,
):
"""
@@ -48,6 +49,9 @@ class BaseLlmConfig(ABC):
Only applicable to vision-enabled models. Defaults to False
vision_details: Level of detail for vision processing.
Options: "low", "high", "auto". Defaults to "auto"
reasoning_effort: Effort level for reasoning models (e.g., o1, o3, gpt-5).
Options: "low", "medium", "high". Only applicable to reasoning models.
Defaults to None (uses the model's default reasoning effort)
http_client_proxies: Proxy settings for HTTP client.
Can be a dict or string. Defaults to None
"""
@@ -59,4 +63,5 @@ class BaseLlmConfig(ABC):
self.top_k = top_k
self.enable_vision = enable_vision
self.vision_details = vision_details
self.reasoning_effort = reasoning_effort
self.http_client = httpx.Client(proxies=http_client_proxies) if http_client_proxies else None
mem0/configs/llms/minimax.py
+56
View File
@@ -0,0 +1,56 @@
from typing import Optional
from mem0.configs.llms.base import BaseLlmConfig
class MinimaxConfig(BaseLlmConfig):
"""
Configuration class for MiniMax-specific parameters.
Inherits from BaseLlmConfig and adds MiniMax-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = None,
temperature: float = 0.1,
api_key: Optional[str] = None,
max_tokens: int = 2000,
top_p: float = 0.1,
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
http_client_proxies: Optional[dict] = None,
# MiniMax-specific parameters
minimax_base_url: Optional[str] = None,
):
"""
Initialize MiniMax configuration.
Args:
model: MiniMax model to use, defaults to None
temperature: Controls randomness, defaults to 0.1
api_key: MiniMax API key, defaults to None
max_tokens: Maximum tokens to generate, defaults to 2000
top_p: Nucleus sampling parameter, defaults to 0.1
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
http_client_proxies: HTTP client proxy settings, defaults to None
minimax_base_url: MiniMax API base URL, defaults to None
"""
# Initialize base parameters
super().__init__(
model=model,
temperature=temperature,
api_key=api_key,
max_tokens=max_tokens,
top_p=top_p,
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
http_client_proxies=http_client_proxies,
)
# MiniMax-specific parameters
self.minimax_base_url = minimax_base_url
mem0/configs/llms/openai.py
+3
View File
@@ -20,6 +20,7 @@ class OpenAIConfig(BaseLlmConfig):
top_k: int = 1,
enable_vision: bool = False,
vision_details: Optional[str] = "auto",
reasoning_effort: Optional[str] = None,
http_client_proxies: Optional[dict] = None,
# OpenAI-specific parameters
openai_base_url: Optional[str] = None,
@@ -44,6 +45,7 @@ class OpenAIConfig(BaseLlmConfig):
top_k: Top-k sampling parameter, defaults to 1
enable_vision: Enable vision capabilities, defaults to False
vision_details: Vision detail level, defaults to "auto"
reasoning_effort: Effort level for reasoning models ("low", "medium", "high"), defaults to None
http_client_proxies: HTTP client proxy settings, defaults to None
openai_base_url: OpenAI API base URL, defaults to None
models: List of models for OpenRouter, defaults to None
@@ -63,6 +65,7 @@ class OpenAIConfig(BaseLlmConfig):
top_k=top_k,
enable_vision=enable_vision,
vision_details=vision_details,
reasoning_effort=reasoning_effort,
http_client_proxies=http_client_proxies,
)
mem0/configs/rerankers/llm.py
+7 -1
View File
@@ -1,4 +1,5 @@
from typing import Optional
from typing import Any, Dict, Optional
from pydantic import Field
from mem0.configs.rerankers.base import BaseRerankerConfig
@@ -46,3 +47,8 @@ class LLMRerankerConfig(BaseRerankerConfig):
default=None,
description="Custom prompt template for scoring documents"
)
llm: Optional[Dict[str, Any]] = Field(
default=None,
description="Nested LLM configuration with 'provider' and 'config' keys. "
"Overrides top-level provider/model/api_key when provided.",
)
mem0/configs/vector_stores/azure_mysql.py
+2 -3
View File
@@ -1,6 +1,6 @@
from typing import Any, Dict, Optional
from pydantic import BaseModel, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, model_validator
class AzureMySQLConfig(BaseModel):
@@ -80,5 +80,4 @@ class AzureMySQLConfig(BaseModel):
return values
class Config:
arbitrary_types_allowed = True
model_config = ConfigDict(arbitrary_types_allowed=True)
mem0/configs/vector_stores/cassandra.py
+2 -3
View File
@@ -1,6 +1,6 @@
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, model_validator
class CassandraConfig(BaseModel):
@@ -72,6 +72,5 @@ class CassandraConfig(BaseModel):
return values
class Config:
arbitrary_types_allowed = True
model_config = ConfigDict(arbitrary_types_allowed=True)
mem0/configs/vector_stores/chroma.py
+2 -2
View File
@@ -32,8 +32,8 @@ class ChromaDbConfig(BaseModel):
values.pop("path", None)
return values
# Check if local/server configuration is provided (excluding default tmp path for cloud config)
local_config = bool(path and path != "/tmp/chroma") or bool(host and port)
# Check if local/server configuration is provided
local_config = bool(path) or bool(host and port)
if not cloud_config and not local_config:
raise ValueError("Either ChromaDB Cloud configuration (api_key, tenant) or local configuration (path or host/port) must be provided.")
mem0/configs/vector_stores/elasticsearch.py
+4 -2
View File
@@ -1,7 +1,7 @@
from collections.abc import Callable
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, model_validator
class ElasticsearchConfig(BaseModel):
@@ -17,7 +17,7 @@ class ElasticsearchConfig(BaseModel):
use_ssl: bool = Field(True, description="Use SSL for connection")
auto_create_index: bool = Field(True, description="Automatically create index during initialization")
custom_search_query: Optional[Callable[[List[float], int, Optional[Dict]], Dict]] = Field(
None, description="Custom search query function. Parameters: (query, limit, filters) -> Dict"
None, description="Custom search query function. Parameters: (query, top_k, filters) -> Dict"
)
headers: Optional[Dict[str, str]] = Field(None, description="Custom headers to include in requests")
@@ -63,3 +63,5 @@ class ElasticsearchConfig(BaseModel):
f"Please input only the following fields: {', '.join(allowed_fields)}"
)
return values
model_config = ConfigDict(arbitrary_types_allowed=True)
mem0/configs/vector_stores/mongodb.py
+3 -1
View File
@@ -1,6 +1,6 @@
from typing import Any, Dict, Optional
from pydantic import BaseModel, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, model_validator
class MongoDBConfig(BaseModel):
@@ -23,3 +23,5 @@ class MongoDBConfig(BaseModel):
f"Please provide only the following fields: {', '.join(allowed_fields)}."
)
return values
model_config = ConfigDict(arbitrary_types_allowed=False)
mem0/configs/vector_stores/neptune.py
+2 -4
View File
@@ -5,7 +5,7 @@ This module provides configuration settings for integrating with Amazon Neptune
as a vector store backend for Mem0's memory layer.
"""
from pydantic import BaseModel, Field
from pydantic import BaseModel, ConfigDict, Field
class NeptuneAnalyticsConfig(BaseModel):
@@ -22,6 +22,4 @@ class NeptuneAnalyticsConfig(BaseModel):
collection_name: str = Field("mem0", description="Default name for the collection")
endpoint: str = Field("endpoint", description="Graph ID for the runtime")
model_config = {
"arbitrary_types_allowed": False,
}
model_config = ConfigDict(arbitrary_types_allowed=False)
mem0/configs/vector_stores/opensearch.py
+3 -1
View File
@@ -1,6 +1,6 @@
from typing import Any, Dict, Optional, Type, Union
from pydantic import BaseModel, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, model_validator
class OpenSearchConfig(BaseModel):
@@ -39,3 +39,5 @@ class OpenSearchConfig(BaseModel):
f"Extra fields not allowed: {', '.join(extra_fields)}. Allowed fields: {', '.join(allowed_fields)}"
)
return values
model_config = ConfigDict(arbitrary_types_allowed=True)
mem0/configs/vector_stores/pgvector.py
+3 -1
View File
@@ -1,6 +1,6 @@
from typing import Any, Dict, Optional
from pydantic import BaseModel, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, model_validator
class PGVectorConfig(BaseModel):
@@ -50,3 +50,5 @@ class PGVectorConfig(BaseModel):
f"Extra fields not allowed: {', '.join(extra_fields)}. Please input only the following fields: {', '.join(allowed_fields)}"
)
return values
model_config = ConfigDict(arbitrary_types_allowed=True)
mem0/configs/vector_stores/qdrant.py
+1 -1
View File
@@ -16,7 +16,7 @@ class QdrantConfig(BaseModel):
path: Optional[str] = Field("/tmp/qdrant", description="Path for local Qdrant database")
url: Optional[str] = Field(None, description="Full URL for Qdrant server")
api_key: Optional[str] = Field(None, description="API key for Qdrant server")
on_disk: Optional[bool] = Field(False, description="Enables persistent storage")
on_disk: Optional[bool] = Field(False,description="Enables persistent storage. Vectors are kept on disk (True) or in memory (False). Does not delete the local database path.")
@model_validator(mode="before")
@classmethod
mem0/configs/vector_stores/supabase.py
+3 -1
View File
@@ -1,7 +1,7 @@
from enum import Enum
from typing import Any, Dict, Optional
from pydantic import BaseModel, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, model_validator
class IndexMethod(str, Enum):
@@ -42,3 +42,5 @@ class SupabaseConfig(BaseModel):
f"Extra fields not allowed: {', '.join(extra_fields)}. Please input only the following fields: {', '.join(allowed_fields)}"
)
return values
model_config = ConfigDict(arbitrary_types_allowed=False)
mem0/configs/vector_stores/turbopuffer.py
+45
View File
@@ -0,0 +1,45 @@
import os
from typing import Any, Dict, Optional
from pydantic import BaseModel, ConfigDict, Field, model_validator
class TurbopufferConfig(BaseModel):
collection_name: str = Field("mem0", description="Name of the namespace/collection")
embedding_model_dims: int = Field(1536, description="Dimensions of the embedding model")
api_key: Optional[str] = Field(None, description="API key for Turbopuffer")
region: str = Field("gcp-us-central1", description="Turbopuffer region (e.g., 'gcp-us-central1', 'aws-us-west-2')")
distance_metric: str = Field(
"cosine_distance",
description="Distance metric for vector similarity ('cosine_distance' or 'euclidean_squared')",
)
batch_size: int = Field(100, description="Batch size for bulk operations")
extra_params: Optional[Dict[str, Any]] = Field(
None,
description="Additional parameters for Turbopuffer client",
)
@model_validator(mode="before")
@classmethod
def check_api_key(cls, values: Dict[str, Any]) -> Dict[str, Any]:
api_key = values.get("api_key")
if not api_key and "TURBOPUFFER_API_KEY" not in os.environ:
raise ValueError(
"Either 'api_key' must be provided or TURBOPUFFER_API_KEY environment variable must be set."
)
return values
@model_validator(mode="before")
@classmethod
def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:
allowed_fields = set(cls.model_fields.keys())
input_fields = set(values.keys())
extra_fields = input_fields - allowed_fields
if extra_fields:
raise ValueError(
f"Extra fields not allowed: {', '.join(extra_fields)}. "
f"Please input only the following fields: {', '.join(allowed_fields)}"
)
return values
model_config = ConfigDict(arbitrary_types_allowed=True)
mem0/configs/vector_stores/valkey.py
+26 -10
View File
@@ -1,15 +1,31 @@
from pydantic import BaseModel
from typing import Any, Dict
from pydantic import BaseModel, ConfigDict, Field, model_validator
class ValkeyConfig(BaseModel):
"""Configuration for Valkey vector store."""
valkey_url: str
collection_name: str
embedding_model_dims: int
timezone: str = "UTC"
index_type: str = "hnsw" # Default to HNSW, can be 'hnsw' or 'flat'
# HNSW specific parameters with recommended defaults
hnsw_m: int = 16 # Number of connections per layer (default from Valkey docs)
hnsw_ef_construction: int = 200 # Search width during construction
hnsw_ef_runtime: int = 10 # Search width during queries
valkey_url: str = Field(..., description="Valkey server URL (e.g., redis://localhost:6379)")
collection_name: str = Field(..., description="Name of the index / collection")
embedding_model_dims: int = Field(..., description="Dimensions of the embedding model")
timezone: str = Field("UTC", description="Timezone for timestamp handling")
index_type: str = Field("hnsw", description="Index type: 'hnsw' (default) or 'flat'")
hnsw_m: int = Field(16, description="HNSW: number of connections per layer")
hnsw_ef_construction: int = Field(200, description="HNSW: search width during index construction")
hnsw_ef_runtime: int = Field(10, description="HNSW: search width during queries")
@model_validator(mode="before")
@classmethod
def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:
allowed_fields = set(cls.model_fields.keys())
input_fields = set(values.keys())
extra_fields = input_fields - allowed_fields
if extra_fields:
raise ValueError(
f"Extra fields not allowed: {', '.join(extra_fields)}. "
f"Please input only the following fields: {', '.join(allowed_fields)}"
)
return values
model_config = ConfigDict(arbitrary_types_allowed=False)
mem0/embeddings/gemini.py
+1 -1
View File
@@ -12,7 +12,7 @@ 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.model = self.config.model or "models/gemini-embedding-001"
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")
mem0/embeddings/ollama.py
+15 -3
View File
@@ -31,12 +31,21 @@ class OllamaEmbedding(EmbeddingBase):
self.client = Client(host=self.config.ollama_base_url)
self._ensure_model_exists()
@staticmethod
def _normalize_model_name(name: str) -> str:
return name if ":" in name else f"{name}:latest"
def _ensure_model_exists(self):
"""
Ensure the specified model exists locally. If not, pull it from Ollama.
"""
local_models = self.client.list()["models"]
if not any(model.get("name") == self.config.model or model.get("model") == self.config.model for model in local_models):
target = self._normalize_model_name(self.config.model)
if not any(
self._normalize_model_name(model.get("name", "")) == target
or self._normalize_model_name(model.get("model", "")) == target
for model in local_models
):
self.client.pull(self.config.model)
def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
@@ -49,5 +58,8 @@ class OllamaEmbedding(EmbeddingBase):
Returns:
list: The embedding vector.
"""
response = self.client.embeddings(model=self.config.model, prompt=text)
return response["embedding"]
response = self.client.embed(model=self.config.model, input=text)
embeddings = response.get("embeddings") or []
if not embeddings:
raise ValueError(f"Ollama embed() returned no embeddings for model '{self.config.model}'")
return embeddings[0]
mem0/embeddings/openai.py
+19 -10
View File
@@ -13,6 +13,9 @@ class OpenAIEmbedding(EmbeddingBase):
super().__init__(config)
self.config.model = self.config.model or "text-embedding-3-small"
# Only pass `dimensions` to the API when the user set embedding_dims; non-matryoshka
# OpenAI-compatible backends (vLLM, Voyage, etc.) reject the parameter
self._pass_dimensions_to_api = self.config.embedding_dims is not None
self.config.embedding_dims = self.config.embedding_dims or 1536
api_key = self.config.api_key or os.getenv("OPENAI_API_KEY")
@@ -42,18 +45,24 @@ class OpenAIEmbedding(EmbeddingBase):
list: The embedding vector.
"""
text = text.replace("\n", " ")
return (
self.client.embeddings.create(input=[text], model=self.config.model, dimensions=self.config.embedding_dims)
.data[0]
.embedding
)
kwargs = {
"input": [text],
"model": self.config.model,
"encoding_format": "float",
}
if self._pass_dimensions_to_api:
kwargs["dimensions"] = self.config.embedding_dims
return self.client.embeddings.create(**kwargs).data[0].embedding
def embed_batch(self, texts, memory_action="add"):
"""Embed multiple texts in a single OpenAI API call."""
texts = [text.replace("\n", " ") for text in texts]
response = self.client.embeddings.create(
input=texts,
model=self.config.model,
dimensions=self.config.embedding_dims,
)
kwargs = {
"input": texts,
"model": self.config.model,
"encoding_format": "float",
}
if self._pass_dimensions_to_api:
kwargs["dimensions"] = self.config.embedding_dims
response = self.client.embeddings.create(**kwargs)
return [item.embedding for item in sorted(response.data, key=lambda x: x.index)]
mem0/embeddings/vertexai.py
+1 -1
View File
@@ -12,7 +12,7 @@ class VertexAIEmbedding(EmbeddingBase):
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
super().__init__(config)
self.config.model = self.config.model or "text-embedding-004"
self.config.model = self.config.model or "gemini-embedding-001"
self.config.embedding_dims = self.config.embedding_dims or 256
self.embedding_types = {
mem0/graphs/configs.py
+24 -2
View File
@@ -77,12 +77,32 @@ class KuzuConfig(BaseModel):
db: Optional[str] = Field(":memory:", description="Path to a Kuzu database file")
class ApacheAgeConfig(BaseModel):
host: Optional[str] = Field("localhost", description="PostgreSQL server hostname")
port: Optional[int] = Field(5432, description="PostgreSQL server port")
database: Optional[str] = Field(None, description="PostgreSQL database name")
username: Optional[str] = Field(None, description="PostgreSQL username")
password: Optional[str] = Field(None, description="PostgreSQL password")
graph_name: Optional[str] = Field("mem0_graph", description="Name of the Apache AGE graph")
@model_validator(mode="before")
def check_required_fields(cls, values):
database, username, password = (
values.get("database"),
values.get("username"),
values.get("password"),
)
if not database or not username or not password:
raise ValueError("Please provide 'database', 'username' and 'password'.")
return values
class GraphStoreConfig(BaseModel):
provider: str = Field(
description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu')",
description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu', 'apache_age')",
default="neo4j",
)
config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig] = Field(
config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig, ApacheAgeConfig] = Field(
description="Configuration for the specific data store", default=None
)
llm: Optional[LlmConfig] = Field(description="LLM configuration for querying the graph store", default=None)
@@ -110,5 +130,7 @@ class GraphStoreConfig(BaseModel):
return NeptuneConfig(**v.model_dump())
elif provider == "kuzu":
return KuzuConfig(**v.model_dump())
elif provider == "apache_age":
return ApacheAgeConfig(**v.model_dump())
else:
raise ValueError(f"Unsupported graph store provider: {provider}")
mem0/graphs/neptune/base.py
+35 -17
View File
@@ -1,7 +1,7 @@
import logging
from abc import ABC, abstractmethod
from mem0.memory.utils import format_entities
from mem0.memory.utils import format_entities, remove_spaces_from_entities
try:
from rank_bm25 import BM25Okapi
@@ -97,7 +97,7 @@ class NeptuneBase(ABC):
for tool_call in search_results["tool_calls"]:
if tool_call["name"] != "extract_entities":
continue
for item in tool_call["arguments"]["entities"]:
for item in tool_call.get("arguments", {}).get("entities", []):
entity_type_map[item["entity"]] = item["entity_type"]
except Exception as e:
logger.exception(
@@ -144,18 +144,14 @@ class NeptuneBase(ABC):
entities = []
if extracted_entities["tool_calls"]:
entities = extracted_entities["tool_calls"][0]["arguments"]["entities"]
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
entities = self._remove_spaces_from_entities(entities)
logger.debug(f"Extracted entities: {entities}")
return entities
def _remove_spaces_from_entities(self, entity_list):
for item in entity_list:
item["source"] = item["source"].lower().replace(" ", "_")
item["relationship"] = item["relationship"].lower().replace(" ", "_")
item["destination"] = item["destination"].lower().replace(" ", "_")
return entity_list
return remove_spaces_from_entities(entity_list, sanitize_relationship=False)
def _get_delete_entities_from_search_output(self, search_output, data, filters):
"""
@@ -350,14 +346,14 @@ class NeptuneBase(ABC):
):
pass
def search(self, query, filters, limit=100):
def search(self, query, filters, top_k=100):
"""
Search for memories and related graph data.
Args:
query (str): Query to search for.
filters (dict): A dictionary containing filters to be applied during the search.
limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
dict: A dictionary containing:
@@ -409,6 +405,28 @@ class NeptuneBase(ABC):
"""
pass
def delete(self, data, filters):
"""
Delete graph entities associated with the given memory text.
Extracts entities and relationships from the memory text using the same
pipeline as add(), then deletes the matching relationships in the graph.
Args:
data (str): The memory text whose graph entities should be removed.
filters (dict): Scope filters (user_id, agent_id, run_id).
"""
try:
entity_type_map = self._retrieve_nodes_from_data(data, filters)
if not entity_type_map:
logger.debug("No entities found in memory text, skipping graph cleanup")
return
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
if to_be_deleted:
self._delete_entities(to_be_deleted, filters["user_id"])
except Exception as e:
logger.error(f"Error during graph cleanup for memory delete: {e}")
def delete_all(self, filters):
cypher, params = self._delete_all_cypher(filters)
self.graph.query(cypher, params=params)
@@ -420,13 +438,13 @@ class NeptuneBase(ABC):
"""
pass
def get_all(self, filters, limit=100):
def get_all(self, filters, top_k=100):
"""
Retrieves all nodes and relationships from the graph database based on filtering criteria.
Args:
filters (dict): A dictionary containing filters to be applied during the retrieval.
limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
list: A list of dictionaries, each containing:
- 'contexts': The base data store response for each memory.
@@ -434,7 +452,7 @@ class NeptuneBase(ABC):
"""
# return all nodes and relationships
query, params = self._get_all_cypher(filters, limit)
query, params = self._get_all_cypher(filters, top_k)
results = self.graph.query(query, params=params)
final_results = []
@@ -452,13 +470,13 @@ class NeptuneBase(ABC):
return final_results
@abstractmethod
def _get_all_cypher(self, filters, limit):
def _get_all_cypher(self, filters, top_k):
"""
Returns the OpenCypher query and parameters to get all edges/nodes in the memory store
"""
pass
def _search_graph_db(self, node_list, filters, limit=100):
def _search_graph_db(self, node_list, filters, top_k=100):
"""
Search similar nodes among and their respective incoming and outgoing relations.
"""
@@ -466,14 +484,14 @@ class NeptuneBase(ABC):
for node in node_list:
n_embedding = self.embedding_model.embed(node)
cypher_query, params = self._search_graph_db_cypher(n_embedding, filters, limit)
cypher_query, params = self._search_graph_db_cypher(n_embedding, filters, top_k)
ans = self.graph.query(cypher_query, params=params)
result_relations.extend(ans)
return result_relations
@abstractmethod
def _search_graph_db_cypher(self, n_embedding, filters, limit):
def _search_graph_db_cypher(self, n_embedding, filters, top_k):
"""
Returns the OpenCypher query and parameters to search for similar nodes in the memory store
"""
mem0/graphs/neptune/neptunedb.py
+14 -15
View File
@@ -1,7 +1,6 @@
import logging
import uuid
from datetime import datetime
import pytz
from datetime import datetime, timezone
from .base import NeptuneBase
@@ -114,7 +113,7 @@ class MemoryGraph(NeptuneBase):
"name": destination,
"type": destination_type,
"user_id": user_id,
"created_at": datetime.now(pytz.timezone("US/Pacific")).isoformat(),
"created_at": datetime.now(timezone.utc).isoformat(),
}
self.vector_store.insert(
vectors=[dest_embedding],
@@ -189,7 +188,7 @@ class MemoryGraph(NeptuneBase):
"name": source,
"type": source_type,
"user_id": user_id,
"created_at": datetime.now(pytz.timezone("US/Pacific")).isoformat(),
"created_at": datetime.now(timezone.utc).isoformat(),
}
self.vector_store.insert(
vectors=[source_embedding],
@@ -316,14 +315,14 @@ class MemoryGraph(NeptuneBase):
"name": source,
"type": source_type,
"user_id": user_id,
"created_at": datetime.now(pytz.timezone("US/Pacific")).isoformat(),
"created_at": datetime.now(timezone.utc).isoformat(),
}
destination_id = str(uuid.uuid4())
destination_payload = {
"name": destination,
"type": destination_type,
"user_id": user_id,
"created_at": datetime.now(pytz.timezone("US/Pacific")).isoformat(),
"created_at": datetime.now(timezone.utc).isoformat(),
}
self.vector_store.insert(
vectors=[source_embedding, dest_embedding],
@@ -381,7 +380,7 @@ class MemoryGraph(NeptuneBase):
source_nodes = self.vector_store.search(
query="",
vectors=source_embedding,
limit=self.vector_store_limit,
top_k=self.vector_store_limit,
filters={"user_id": user_id},
)
@@ -414,7 +413,7 @@ class MemoryGraph(NeptuneBase):
destination_nodes = self.vector_store.search(
query="",
vectors=destination_embedding,
limit=self.vector_store_limit,
top_k=self.vector_store_limit,
filters={"user_id": user_id},
)
@@ -456,12 +455,12 @@ class MemoryGraph(NeptuneBase):
logger.debug(f"delete_all query={cypher}")
return cypher, params
def _get_all_cypher(self, filters, limit):
def _get_all_cypher(self, filters, top_k):
"""
Returns the OpenCypher query and parameters to get all edges/nodes in the memory store
:param filters: search filters
:param limit: return limit
:param top_k: return limit
:return: str, dict
"""
@@ -470,16 +469,16 @@ class MemoryGraph(NeptuneBase):
RETURN n.name AS source, type(r) AS relationship, m.name AS target
LIMIT $limit
"""
params = {"user_id": filters["user_id"], "limit": limit}
params = {"user_id": filters["user_id"], "limit": top_k}
return cypher, params
def _search_graph_db_cypher(self, n_embedding, filters, limit):
def _search_graph_db_cypher(self, n_embedding, filters, top_k):
"""
Returns the OpenCypher query and parameters to search for similar nodes in the memory store
:param n_embedding: node vector
:param filters: search filters
:param limit: return limit
:param top_k: return limit
:return: str, dict
"""
@@ -487,7 +486,7 @@ class MemoryGraph(NeptuneBase):
search_nodes = self.vector_store.search(
query="",
vectors=n_embedding,
limit=self.vector_store_limit,
top_k=self.vector_store_limit,
filters=filters,
)
@@ -505,7 +504,7 @@ class MemoryGraph(NeptuneBase):
params = {
"n_ids": ids,
"user_id": filters["user_id"],
"limit": limit,
"limit": top_k,
}
logger.debug(f"_search_graph_db\n query={cypher_query}")
mem0/graphs/neptune/neptunegraph.py
+7 -7
View File
@@ -3,8 +3,8 @@ import logging
from .base import NeptuneBase
try:
from langchain_aws import NeptuneAnalyticsGraph
from botocore.config import Config
from langchain_aws import NeptuneAnalyticsGraph
except ImportError:
raise ImportError("langchain_aws is not installed. Please install it using 'make install_all'.")
@@ -412,12 +412,12 @@ class MemoryGraph(NeptuneBase):
logger.debug(f"delete_all query={cypher}")
return cypher, params
def _get_all_cypher(self, filters, limit):
def _get_all_cypher(self, filters, top_k):
"""
Returns the OpenCypher query and parameters to get all edges/nodes in the memory store
:param filters: search filters
:param limit: return limit
:param top_k: return limit
:return: str, dict
"""
@@ -426,16 +426,16 @@ class MemoryGraph(NeptuneBase):
RETURN n.name AS source, type(r) AS relationship, m.name AS target
LIMIT $limit
"""
params = {"user_id": filters["user_id"], "limit": limit}
params = {"user_id": filters["user_id"], "limit": top_k}
return cypher, params
def _search_graph_db_cypher(self, n_embedding, filters, limit):
def _search_graph_db_cypher(self, n_embedding, filters, top_k):
"""
Returns the OpenCypher query and parameters to search for similar nodes in the memory store
:param n_embedding: node vector
:param filters: search filters
:param limit: return limit
:param top_k: return limit
:return: str, dict
"""
@@ -468,7 +468,7 @@ class MemoryGraph(NeptuneBase):
"n_embedding": n_embedding,
"threshold": self.threshold,
"user_id": filters["user_id"],
"limit": limit,
"limit": top_k,
}
logger.debug(f"_search_graph_db\n query={cypher_query}")
mem0/llms/anthropic.py
+25
View File
@@ -40,6 +40,31 @@ class AnthropicLLM(LLMBase):
api_key = self.config.api_key or os.getenv("ANTHROPIC_API_KEY")
self.client = anthropic.Anthropic(api_key=api_key)
def _get_common_params(self, **kwargs) -> Dict:
"""Get common parameters, avoiding sending both temperature and top_p together.
Anthropic rejects requests that include both temperature and top_p.
When both are set, we keep temperature and drop top_p.
"""
params = {}
if self.config.max_tokens is not None:
params["max_tokens"] = self.config.max_tokens
has_temperature = self.config.temperature is not None
has_top_p = self.config.top_p is not None
if has_temperature and has_top_p:
# Anthropic forbids both; prefer temperature
params["temperature"] = self.config.temperature
elif has_temperature:
params["temperature"] = self.config.temperature
elif has_top_p:
params["top_p"] = self.config.top_p
params.update(kwargs)
return params
def generate_response(
self,
messages: List[Dict[str, str]],
mem0/llms/aws_bedrock.py
+92 -44
View File
@@ -17,8 +17,9 @@ from mem0.memory.utils import extract_json
logger = logging.getLogger(__name__)
PROVIDERS = [
"ai21", "amazon", "anthropic", "cohere", "meta", "mistral", "stability", "writer",
"deepseek", "gpt-oss", "perplexity", "snowflake", "titan", "command", "j2", "llama"
"ai21", "amazon", "anthropic", "cohere", "meta", "mistral", "stability", "writer",
"deepseek", "gpt-oss", "perplexity", "snowflake", "titan", "command", "j2", "llama",
"minimax",
]
@@ -120,6 +121,10 @@ class AWSBedrockLLM(LLMBase):
"""Initialize provider-specific settings and capabilities."""
# Determine capabilities based on provider and model
self.supports_tools = self.provider in ["anthropic", "cohere", "amazon"]
# MiniMax M2.x is intentionally excluded from supports_tools: tool use for MiniMax
# on Amazon Bedrock is only available via the bedrock-mantle (OpenAI-compatible)
# endpoint, not via the bedrock-runtime Converse API used by this class.
# See: https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-minimax-minimax-m2-5.html
self.supports_vision = self.provider in ["anthropic", "amazon", "meta", "mistral"]
self.supports_streaming = self.provider in ["anthropic", "cohere", "mistral", "amazon", "meta"]
@@ -228,6 +233,12 @@ class AWSBedrockLLM(LLMBase):
return "\n\nHuman: " + "".join(formatted_messages) + "\n\nAssistant:"
def _merge_optional_top_p(self, target: Dict[str, Any], *, key: str = "top_p") -> None:
"""Add nucleus sampling to ``target`` only when ``model_config`` has ``top_p`` set."""
top_p = self.model_config.get("top_p")
if top_p is not None:
target[key] = top_p
def _prepare_input(self, prompt: str) -> Dict[str, Any]:
"""
Prepare input for the current provider's model.
@@ -268,44 +279,38 @@ class AWSBedrockLLM(LLMBase):
"messages": [{"role": "user", "content": prompt}],
"max_tokens": self.model_config.get("max_tokens", 5000),
"temperature": self.model_config.get("temperature", 0.1),
"top_p": self.model_config.get("top_p", 0.9),
}
self._merge_optional_top_p(input_body)
else:
# Legacy Amazon models
input_body = {
"inputText": prompt,
"textGenerationConfig": {
"maxTokenCount": self.model_config.get("max_tokens", 5000),
"topP": self.model_config.get("top_p", 0.9),
"temperature": self.model_config.get("temperature", 0.1),
},
}
# Remove None values
input_body["textGenerationConfig"] = {
k: v for k, v in input_body["textGenerationConfig"].items() if v is not None
text_gen_config: Dict[str, Any] = {
"maxTokenCount": self.model_config.get("max_tokens", 5000),
"temperature": self.model_config.get("temperature", 0.1),
}
self._merge_optional_top_p(text_gen_config, key="topP")
input_body = {"inputText": prompt, "textGenerationConfig": text_gen_config}
elif self.provider == "anthropic":
input_body = {
"messages": [{"role": "user", "content": [{"type": "text", "text": prompt}]}],
"max_tokens": self.model_config.get("max_tokens", 2000),
"temperature": self.model_config.get("temperature", 0.1),
"top_p": self.model_config.get("top_p", 0.9),
"anthropic_version": "bedrock-2023-05-31",
}
self._merge_optional_top_p(input_body)
elif self.provider == "meta":
input_body = {
"prompt": prompt,
"max_gen_len": self.model_config.get("max_tokens", 5000),
"temperature": self.model_config.get("temperature", 0.1),
"top_p": self.model_config.get("top_p", 0.9),
}
self._merge_optional_top_p(input_body)
elif self.provider == "mistral":
input_body = {
"prompt": prompt,
"max_tokens": self.model_config.get("max_tokens", 5000),
"temperature": self.model_config.get("temperature", 0.1),
"top_p": self.model_config.get("top_p", 0.9),
}
self._merge_optional_top_p(input_body)
else:
# Generic case - add all model config parameters
input_body.update(self.model_config)
@@ -479,6 +484,37 @@ class AWSBedrockLLM(LLMBase):
return converse_tools
def _default_max_tokens_for_converse(self) -> int:
"""Default maxTokens if ``max_tokens`` is missing (Nova: 5000, else 2000)."""
model_id = (self.config.model or "").lower()
if self.provider == "amazon" and "nova" in model_id:
return 5000
return 2000
def _build_inference_config(self) -> Dict[str, Any]:
"""Build Converse ``inferenceConfig``.
Anthropic and MiniMax reasoning models reject requests that include both
``temperature`` and ``topP`` simultaneously, so ``topP`` is omitted for
those providers even when the user has configured it.
"""
inference_config: Dict[str, Any] = {
"maxTokens": self.model_config.get("max_tokens", self._default_max_tokens_for_converse()),
"temperature": self.model_config.get("temperature", 0.1),
}
top_p = self.model_config.get("top_p")
if top_p is not None:
if self.provider in ("anthropic", "minimax"):
# Both Anthropic and MiniMax M2.x (reasoning models) raise a
# ValidationException when temperature and topP are both present
# in inferenceConfig. Omit topP and rely on temperature only.
logger.debug("Omitting topP for %s Converse (using temperature); top_p=%s", self.provider, top_p)
else:
inference_config["topP"] = top_p
return inference_config
def _generate_with_tools(self, messages: List[Dict[str, str]], tools: List[Dict], stream: bool = False) -> Dict[str, Any]:
"""Generate response with tool calling support using correct message format."""
# Format messages for tool-enabled models
@@ -501,11 +537,7 @@ class AWSBedrockLLM(LLMBase):
converse_params = {
"modelId": self.config.model,
"messages": formatted_messages,
"inferenceConfig": {
"maxTokens": self.model_config.get("max_tokens", 2000),
"temperature": self.model_config.get("temperature", 0.1),
"topP": self.model_config.get("top_p", 0.9),
}
"inferenceConfig": self._build_inference_config(),
}
# Add system message if present (for Anthropic)
@@ -531,11 +563,7 @@ class AWSBedrockLLM(LLMBase):
converse_params = {
"modelId": self.config.model,
"messages": formatted_messages,
"inferenceConfig": {
"maxTokens": self.model_config.get("max_tokens", 2000),
"temperature": self.model_config.get("temperature", 0.1),
"topP": self.model_config.get("top_p", 0.9),
}
"inferenceConfig": self._build_inference_config(),
}
# Add system message if present
@@ -553,27 +581,47 @@ class AWSBedrockLLM(LLMBase):
else:
return str(response)
elif self.provider == "amazon" and "nova" in self.config.model.lower():
# Nova models use converse API even without tools
formatted_messages = self._format_messages_amazon(messages)
input_body = {
"messages": formatted_messages,
"max_tokens": self.model_config.get("max_tokens", 5000),
"temperature": self.model_config.get("temperature", 0.1),
"top_p": self.model_config.get("top_p", 0.9),
elif self.provider == "minimax":
# MiniMax models (e.g. minimax.minimax-m2.5) use the Bedrock Converse API.
# M2.5 is a reasoning model whose response content array may include a
# `reasoningContent` block before the actual `text` block, so we iterate
# to find the first block that contains a "text" key.
# System messages must be passed via the top-level `system` parameter
# (not as a message with role="system") per the Converse API spec.
system_parts = []
converse_messages = []
for msg in messages:
role = msg.get("role", "user")
content = msg.get("content", "")
if not isinstance(content, str):
content = str(content)
if role == "system":
system_parts.append(content)
else:
converse_messages.append({"role": role, "content": [{"text": content}]})
if not converse_messages:
converse_messages = [{"role": "user", "content": [{"text": ""}]}]
converse_params = {
"modelId": self.config.model,
"messages": converse_messages,
"inferenceConfig": self._build_inference_config(),
}
# Use converse API for Nova models
if system_parts:
converse_params["system"] = [{"text": "\n".join(system_parts)}]
response = self.client.converse(**converse_params)
for block in response["output"]["message"]["content"]:
if "text" in block:
return block["text"]
return ""
elif self.provider == "amazon" and "nova" in self.config.model.lower():
# Nova models use the Converse API even without tools
formatted_messages = self._format_messages_amazon(messages)
response = self.client.converse(
modelId=self.config.model,
messages=input_body["messages"],
inferenceConfig={
"maxTokens": input_body["max_tokens"],
"temperature": input_body["temperature"],
"topP": input_body["top_p"],
}
messages=formatted_messages,
inferenceConfig=self._build_inference_config(),
)
return self._parse_response(response)
else:
# For other providers and legacy Amazon models (like Titan)
mem0/llms/azure_openai.py
+3
View File
@@ -31,6 +31,7 @@ class AzureOpenAILLM(LLMBase):
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
reasoning_effort=getattr(config, 'reasoning_effort', None),
http_client_proxies=config.http_client,
)
@@ -133,6 +134,8 @@ class AzureOpenAILLM(LLMBase):
"messages": messages,
})
if response_format:
params["response_format"] = response_format
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
mem0/llms/azure_openai_structured.py
+32 -4
View File
@@ -1,3 +1,4 @@
import json
import os
from typing import Dict, List, Optional
@@ -6,6 +7,7 @@ from openai import AzureOpenAI
from mem0.configs.llms.base import BaseLlmConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
SCOPE = "https://cognitiveservices.azure.com/.default"
@@ -83,9 +85,35 @@ class AzureOpenAIStructuredLLM(LLMBase):
params["tools"] = tools
params["tool_choice"] = tool_choice
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
response = self.client.chat.completions.create(**params)
return self._parse_response(response, tools)
def _parse_response(self, response, tools):
"""
Process the response based on whether tools are used or not.
Args:
response: The raw response from API.
tools: The list of tools provided in the request.
Returns:
str or dict: The processed response.
"""
if tools:
processed_response = {
"content": response.choices[0].message.content,
"tool_calls": [],
}
if response.choices[0].message.tool_calls:
for tool_call in response.choices[0].message.tool_calls:
processed_response["tool_calls"].append(
{
"name": tool_call.function.name,
"arguments": json.loads(extract_json(tool_call.function.arguments)),
}
)
return processed_response
else:
return response.choices[0].message.content
mem0/llms/base.py
+6 -1
View File
@@ -88,7 +88,12 @@ class LLMBase(ABC):
supported_params["tools"] = kwargs["tools"]
if "tool_choice" in kwargs:
supported_params["tool_choice"] = kwargs["tool_choice"]
# Add reasoning_effort if configured
reasoning_effort = getattr(self.config, 'reasoning_effort', None)
if reasoning_effort:
supported_params["reasoning_effort"] = reasoning_effort
return supported_params
else:
# For regular models, include all common parameters
mem0/llms/configs.py
+1
View File
@@ -23,6 +23,7 @@ class LlmConfig(BaseModel):
"azure_openai_structured",
"gemini",
"deepseek",
"minimax",
"xai",
"sarvam",
"lmstudio",
mem0/llms/deepseek.py
+2
View File
@@ -99,6 +99,8 @@ class DeepSeekLLM(LLMBase):
}
)
if response_format:
params["response_format"] = response_format
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
mem0/llms/gemini.py
+11 -8
View File
@@ -32,22 +32,25 @@ class GeminiLLM(LLMBase):
Returns:
str or dict: The processed response.
"""
# Get parts safely — content can be None when Gemini blocks the response
candidate = response.candidates[0] if response.candidates else None
parts = candidate.content.parts if candidate and candidate.content else None
if tools:
processed_response = {
"content": None,
"tool_calls": [],
}
# Extract content from the first candidate
if response.candidates and response.candidates[0].content.parts:
for part in response.candidates[0].content.parts:
if parts:
# Extract content from the first candidate
for part in parts:
if hasattr(part, "text") and part.text:
processed_response["content"] = part.text
break
# Extract function calls
if response.candidates and response.candidates[0].content.parts:
for part in response.candidates[0].content.parts:
# Extract function calls
for part in parts:
if hasattr(part, "function_call") and part.function_call:
fn = part.function_call
processed_response["tool_calls"].append(
@@ -59,8 +62,8 @@ class GeminiLLM(LLMBase):
return processed_response
else:
if response.candidates and response.candidates[0].content.parts:
for part in response.candidates[0].content.parts:
if parts:
for part in parts:
if hasattr(part, "text") and part.text:
return part.text
return ""
mem0/llms/groq.py
+1 -1
View File
@@ -17,7 +17,7 @@ class GroqLLM(LLMBase):
super().__init__(config)
if not self.config.model:
self.config.model = "llama3-70b-8192"
self.config.model = "llama-3.3-70b-versatile"
api_key = self.config.api_key or os.getenv("GROQ_API_KEY")
self.client = Groq(api_key=api_key)
mem0/llms/minimax.py
+114
View File
@@ -0,0 +1,114 @@
import json
import os
from typing import Dict, List, Optional, Union
from openai import OpenAI
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.minimax import MinimaxConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
class MiniMaxLLM(LLMBase):
def __init__(self, config: Optional[Union[BaseLlmConfig, MinimaxConfig, Dict]] = None):
# Convert to MinimaxConfig if needed
if config is None:
config = MinimaxConfig()
elif isinstance(config, dict):
config = MinimaxConfig(**config)
elif isinstance(config, BaseLlmConfig) and not isinstance(config, MinimaxConfig):
# Convert BaseLlmConfig to MinimaxConfig
config = MinimaxConfig(
model=config.model,
temperature=config.temperature,
api_key=config.api_key,
max_tokens=config.max_tokens,
top_p=config.top_p,
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
http_client_proxies=config.http_client,
)
super().__init__(config)
if not self.config.model:
self.config.model = "MiniMax-M2.7"
api_key = self.config.api_key or os.getenv("MINIMAX_API_KEY")
base_url = (
self.config.minimax_base_url
or os.getenv("MINIMAX_API_BASE")
or "https://api.minimax.io/v1"
)
self.client = OpenAI(api_key=api_key, base_url=base_url)
def _parse_response(self, response, tools):
"""
Process the response based on whether tools are used or not.
Args:
response: The raw response from API.
tools: The list of tools provided in the request.
Returns:
str or dict: The processed response.
"""
if tools:
processed_response = {
"content": response.choices[0].message.content,
"tool_calls": [],
}
if response.choices[0].message.tool_calls:
for tool_call in response.choices[0].message.tool_calls:
processed_response["tool_calls"].append(
{
"name": tool_call.function.name,
"arguments": json.loads(extract_json(tool_call.function.arguments)),
}
)
return processed_response
else:
return response.choices[0].message.content
def generate_response(
self,
messages: List[Dict[str, str]],
response_format=None,
tools: Optional[List[Dict]] = None,
tool_choice: str = "auto",
**kwargs,
):
"""
Generate a response based on the given messages using MiniMax.
Args:
messages (list): List of message dicts containing 'role' and 'content'.
response_format (str or object, optional): Format of the response. Defaults to None.
tools (list, optional): List of tools that the model can call. Defaults to None.
tool_choice (str, optional): Tool choice method. Defaults to "auto".
**kwargs: Additional MiniMax-specific parameters.
Returns:
str: The generated response.
"""
params = self._get_supported_params(messages=messages, **kwargs)
params.update(
{
"model": self.config.model,
"messages": messages,
}
)
if response_format:
params["response_format"] = response_format
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
response = self.client.chat.completions.create(**params)
return self._parse_response(response, tools)
mem0/llms/ollama.py
+27 -1
View File
@@ -1,3 +1,4 @@
import json
from typing import Dict, List, Optional, Union
try:
@@ -8,6 +9,7 @@ except ImportError:
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.ollama import OllamaConfig
from mem0.llms.base import LLMBase
from mem0.memory.utils import extract_json
class OllamaLLM(LLMBase):
@@ -61,7 +63,28 @@ class OllamaLLM(LLMBase):
"tool_calls": [],
}
# Ollama doesn't support tool calls in the same way, so we return the content
if isinstance(response, dict):
raw_calls = response.get("message", {}).get("tool_calls") or []
else:
raw_calls = getattr(response.message, "tool_calls", None) or []
for tool_call in raw_calls:
if isinstance(tool_call, dict):
fn = tool_call.get("function", {})
name = fn.get("name", "")
arguments = fn.get("arguments", {})
else:
fn = getattr(tool_call, "function", None)
name = getattr(fn, "name", "") if fn else ""
arguments = getattr(fn, "arguments", {}) if fn else {}
if isinstance(arguments, str):
arguments = json.loads(extract_json(arguments))
processed_response["tool_calls"].append(
{"name": name, "arguments": arguments}
)
return processed_response
else:
return content
@@ -113,5 +136,8 @@ class OllamaLLM(LLMBase):
# Remove OpenAI-specific parameters that Ollama doesn't support
params.pop("max_tokens", None) # Ollama uses different parameter names
if tools:
params["tools"] = tools
response = self.client.chat(**params)
return self._parse_response(response, tools)
mem0/llms/openai.py
+1
View File
@@ -29,6 +29,7 @@ class OpenAILLM(LLMBase):
top_k=config.top_k,
enable_vision=config.enable_vision,
vision_details=config.vision_details,
reasoning_effort=getattr(config, 'reasoning_effort', None),
http_client_proxies=config.http_client,
)
mem0/llms/vllm.py
+2
View File
@@ -99,6 +99,8 @@ class VllmLLM(LLMBase):
}
)
if response_format:
params["response_format"] = response_format
if tools:
params["tools"] = tools
params["tool_choice"] = tool_choice
mem0/memory/apache_age_memory.py
+595
View File
@@ -0,0 +1,595 @@
import json
import logging
import time
from mem0.memory.utils import format_entities, sanitize_relationship_for_cypher
try:
import age
except ImportError:
raise ImportError("apache-age-python is not installed. Please install it using pip install apache-age-python")
try:
from rank_bm25 import BM25Okapi
except ImportError:
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
from mem0.graphs.tools import (
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
DELETE_MEMORY_TOOL_GRAPH,
EXTRACT_ENTITIES_STRUCT_TOOL,
EXTRACT_ENTITIES_TOOL,
RELATIONS_STRUCT_TOOL,
RELATIONS_TOOL,
)
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
from mem0.utils.factory import EmbedderFactory, LlmFactory
logger = logging.getLogger(__name__)
def _cosine_similarity(vec1, vec2):
"""Compute cosine similarity between two vectors without numpy."""
dot = sum(a * b for a, b in zip(vec1, vec2))
norm1 = sum(a * a for a in vec1) ** 0.5
norm2 = sum(b * b for b in vec2) ** 0.5
if norm1 == 0 or norm2 == 0:
return 0.0
return dot / (norm1 * norm2)
def _get_similar_nodes(nodes, query_embedding, filters, threshold):
"""Find nodes above the similarity threshold from a fetched node list.
Shared by ``_find_similar_node`` and ``_search_graph_db`` to avoid
duplicating the client-side cosine similarity logic.
"""
matches = []
for node in nodes:
props = node if isinstance(node, dict) else {}
stored_emb = props.get("embedding")
if not stored_emb:
continue
if isinstance(stored_emb, str):
stored_emb = json.loads(stored_emb)
if filters.get("agent_id") and props.get("agent_id") != filters["agent_id"]:
continue
if filters.get("run_id") and props.get("run_id") != filters["run_id"]:
continue
sim = _cosine_similarity(query_embedding, stored_emb)
if sim >= threshold:
matches.append({"name": props.get("name"), "similarity": sim, "props": props})
matches.sort(key=lambda x: x["similarity"], reverse=True)
return matches
class MemoryGraph:
def __init__(self, config):
self.config = config
graph_cfg = self.config.graph_store.config
self.graph_name = graph_cfg.graph_name
# Connect using the Apache AGE Python driver (psycopg2-based)
self.ag = age.connect(
graph=self.graph_name,
host=graph_cfg.host,
port=graph_cfg.port,
dbname=graph_cfg.database,
user=graph_cfg.username,
password=graph_cfg.password,
)
self.embedding_model = EmbedderFactory.create(
self.config.embedder.provider, self.config.embedder.config, self.config.vector_store.config
)
# Default to openai if no specific provider is configured
self.llm_provider = "openai"
if self.config.llm and self.config.llm.provider:
self.llm_provider = self.config.llm.provider
if self.config.graph_store and self.config.graph_store.llm and self.config.graph_store.llm.provider:
self.llm_provider = self.config.graph_store.llm.provider
# Get LLM config with proper null checks
llm_config = None
if self.config.graph_store and self.config.graph_store.llm and hasattr(self.config.graph_store.llm, "config"):
llm_config = self.config.graph_store.llm.config
elif hasattr(self.config.llm, "config"):
llm_config = self.config.llm.config
self.llm = LlmFactory.create(self.llm_provider, llm_config)
self.user_id = None
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, "threshold") else 0.7
# -- helpers ---------------------------------------------------------------
def _exec_cypher(self, cypher_stmt, cols=None, params=None):
"""Execute a Cypher query via the AGE driver and return results.
Uses ``ag.execCypher`` which delegates to psycopg2's safe parameter
substitution (``%s`` placeholders). The *cols* argument specifies the
column names in the ``AS (…)`` clause — when ``None`` the driver
defaults to a single ``v agtype`` column.
When *cols* are provided, returns a list of dicts keyed by column name.
When *cols* is ``None``, returns vertex/edge property dicts or raw values.
"""
cursor = self.ag.execCypher(cypher_stmt, cols=cols, params=params)
rows = cursor.fetchall()
if not rows:
return []
col_names = [desc[0] for desc in cursor.description] if cursor.description else None
results = []
for row in rows:
if col_names and len(col_names) > 1:
record = {}
for i, col_name in enumerate(col_names):
val = row[i]
if hasattr(val, "properties"):
record[col_name] = val.properties
else:
record[col_name] = val
results.append(record)
else:
val = row[0] if len(row) == 1 else row
if hasattr(val, "properties"):
results.append(val.properties)
else:
results.append(val)
return results
def _fetch_user_nodes_with_embeddings(self, user_id):
"""Fetch all nodes with embeddings for a given user_id."""
return self._exec_cypher(
"MATCH (n {user_id: %s}) WHERE n.embedding IS NOT NULL RETURN n",
params=(user_id,),
)
def _find_similar_node(self, embedding, filters, threshold=0.9):
"""Find the most similar existing node by cosine similarity.
Apache AGE does not have a built-in vector index, so we fetch all
node embeddings matching the filters and compute cosine similarity
on the client side. This is adequate for moderate graph sizes; for
very large graphs consider pairing AGE with pgvector.
"""
nodes = self._fetch_user_nodes_with_embeddings(filters["user_id"])
matches = _get_similar_nodes(nodes, embedding, filters, threshold)
return matches[0]["props"] if matches else None
def _merge_node(self, user_id, name, embedding, agent_id=None, run_id=None):
"""Create a node if it doesn't exist, or update mentions if it does.
Apache AGE does not support ``ON CREATE SET`` / ``ON MATCH SET``, so
we use ``MERGE … SET`` which always applies the SET clause. Embeddings
and optional filter properties are set in a single query.
"""
set_parts = [
"n.embedding = %s",
"n.mentions = coalesce(n.mentions, 0) + 1",
"n.created = coalesce(n.created, %s)",
]
params = [user_id, name, json.dumps(embedding), int(time.time() * 1000)]
if agent_id:
set_parts.append("n.agent_id = %s")
params.append(agent_id)
if run_id:
set_parts.append("n.run_id = %s")
params.append(run_id)
set_clause = ", ".join(set_parts)
self._exec_cypher(
f"MERGE (n {{user_id: %s, name: %s}}) SET {set_clause}",
params=tuple(params),
)
def close(self):
"""Close the underlying database connection."""
if self.ag:
self.ag.close()
# -- public API ------------------------------------------------------------
def add(self, data, filters):
"""
Adds data to the graph.
Args:
data (str): The data to add to the graph.
filters (dict): A dictionary containing filters to be applied during the addition.
"""
entity_type_map = self._retrieve_nodes_from_data(data, filters)
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
deleted_entities = self._delete_entities(to_be_deleted, filters)
added_entities = self._add_entities(to_be_added, filters, entity_type_map)
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
def search(self, query, filters, top_k=100):
"""
Search for memories and related graph data.
Args:
query (str): Query to search for.
filters (dict): A dictionary containing filters to be applied during the search.
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
list: A list of dicts with keys "source", "relationship", "destination".
"""
entity_type_map = self._retrieve_nodes_from_data(query, filters)
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
if not search_output:
return []
search_outputs_sequence = [
[item["source"], item["relationship"], item["destination"]] for item in search_output
]
bm25 = BM25Okapi(search_outputs_sequence)
tokenized_query = query.split(" ")
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
search_results = []
for item in reranked_results:
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
logger.info(f"Returned {len(search_results)} search results")
return search_results
def delete(self, data, filters):
"""
Delete graph entities associated with the given memory text.
Extracts entities and relationships from the memory text using the same
pipeline as add(), then deletes the matching relationships in the graph.
Args:
data (str): The memory text whose graph entities should be removed.
filters (dict): Scope filters (user_id, agent_id, run_id).
"""
try:
entity_type_map = self._retrieve_nodes_from_data(data, filters)
if not entity_type_map:
logger.debug("No entities found in memory text, skipping graph cleanup")
return
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
if to_be_deleted:
self._delete_entities(to_be_deleted, filters)
except Exception as e:
logger.error(f"Error during graph cleanup for memory delete: {e}")
def delete_all(self, filters):
"""Delete all nodes and relationships for a user or specific agent."""
where_parts = ["n.user_id = %s"]
params = [filters["user_id"]]
if filters.get("agent_id"):
where_parts.append("n.agent_id = %s")
params.append(filters["agent_id"])
if filters.get("run_id"):
where_parts.append("n.run_id = %s")
params.append(filters["run_id"])
where_clause = " AND ".join(where_parts)
self._exec_cypher(
f"MATCH (n) WHERE {where_clause} DETACH DELETE n",
params=tuple(params),
)
self.ag.commit()
def get_all(self, filters, top_k=100):
"""
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
Args:
filters (dict): A dictionary containing filters to be applied during the retrieval.
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
list: A list of dictionaries, each containing:
- 'source': The source node name.
- 'relationship': The relationship type.
- 'target': The target node name.
"""
where_parts = ["n.user_id = %s", "m.user_id = %s"]
params = [filters["user_id"], filters["user_id"]]
if filters.get("agent_id"):
where_parts.extend(["n.agent_id = %s", "m.agent_id = %s"])
params.extend([filters["agent_id"], filters["agent_id"]])
if filters.get("run_id"):
where_parts.extend(["n.run_id = %s", "m.run_id = %s"])
params.extend([filters["run_id"], filters["run_id"]])
where_clause = " AND ".join(where_parts)
params.append(top_k)
results = self._exec_cypher(
f"MATCH (n)-[r]->(m) WHERE {where_clause} "
f"RETURN n.name, type(r), m.name LIMIT %s",
cols=["source", "relationship", "target"],
params=tuple(params),
)
final_results = []
for result in results:
final_results.append(
{
"source": result["source"],
"relationship": result["relationship"],
"target": result["target"],
}
)
logger.info(f"Retrieved {len(final_results)} relationships")
return final_results
# -- LLM-driven extraction -------------------------------------------------
def _retrieve_nodes_from_data(self, data, filters):
"""Extracts all the entities mentioned in the query."""
_tools = [EXTRACT_ENTITIES_TOOL]
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
search_results = self.llm.generate_response(
messages=[
{
"role": "system",
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
},
{"role": "user", "content": data},
],
tools=_tools,
)
entity_type_map = {}
try:
for tool_call in search_results["tool_calls"]:
if tool_call["name"] != "extract_entities":
continue
for item in tool_call.get("arguments", {}).get("entities", []):
if "entity" in item and "entity_type" in item:
entity_type_map[item["entity"]] = item["entity_type"]
except Exception as e:
logger.exception(
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
)
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
logger.debug(f"Entity type map: {entity_type_map}\n search_results={search_results}")
return entity_type_map
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
"""Establish relations among the extracted nodes."""
user_identity = f"user_id: {filters['user_id']}"
if filters.get("agent_id"):
user_identity += f", agent_id: {filters['agent_id']}"
if filters.get("run_id"):
user_identity += f", run_id: {filters['run_id']}"
if self.config.graph_store.custom_prompt:
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
system_content = system_content.replace("CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}")
messages = [
{"role": "system", "content": system_content},
{"role": "user", "content": data},
]
else:
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
messages = [
{"role": "system", "content": system_content},
{"role": "user", "content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}"},
]
_tools = [RELATIONS_TOOL]
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
_tools = [RELATIONS_STRUCT_TOOL]
extracted_entities = self.llm.generate_response(
messages=messages,
tools=_tools,
)
entities = []
if extracted_entities and extracted_entities.get("tool_calls"):
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
entities = self._remove_spaces_from_entities(entities)
logger.debug(f"Extracted entities: {entities}")
return entities
# -- graph DB operations ---------------------------------------------------
def _search_graph_db(self, node_list, filters, top_k=100):
"""Search similar nodes and their respective incoming and outgoing relations."""
result_relations = []
for node in node_list:
n_embedding = self.embedding_model.embed(node)
nodes = self._fetch_user_nodes_with_embeddings(filters["user_id"])
similar_nodes = _get_similar_nodes(nodes, n_embedding, filters, self.threshold)
# Build WHERE clause for relationship target filtering
rel_where_parts = ["m.user_id = %s"]
rel_params_suffix = [filters["user_id"]]
if filters.get("agent_id"):
rel_where_parts.append("m.agent_id = %s")
rel_params_suffix.append(filters["agent_id"])
if filters.get("run_id"):
rel_where_parts.append("m.run_id = %s")
rel_params_suffix.append(filters["run_id"])
rel_where = " AND ".join(rel_where_parts)
# For each similar node, fetch its relationships
for sn in similar_nodes[:top_k]:
node_name = sn["name"]
similarity = sn["similarity"]
out_params = (filters["user_id"], node_name) + tuple(rel_params_suffix)
out_results = self._exec_cypher(
f"MATCH (n {{user_id: %s, name: %s}})-[r]->(m) "
f"WHERE {rel_where} "
f"RETURN n.name, type(r), m.name",
cols=["source", "relationship", "destination"],
params=out_params,
)
in_params = (filters["user_id"], node_name) + tuple(rel_params_suffix)
in_results = self._exec_cypher(
f"MATCH (n {{user_id: %s, name: %s}})<-[r]-(m) "
f"WHERE {rel_where} "
f"RETURN m.name, type(r), n.name",
cols=["source", "relationship", "destination"],
params=in_params,
)
for rel in out_results + in_results:
rel["similarity"] = similarity
result_relations.append(rel)
return result_relations
def _get_delete_entities_from_search_output(self, search_output, data, filters):
"""Get the entities to be deleted from the search output."""
search_output_string = format_entities(search_output)
user_identity = f"user_id: {filters['user_id']}"
if filters.get("agent_id"):
user_identity += f", agent_id: {filters['agent_id']}"
if filters.get("run_id"):
user_identity += f", run_id: {filters['run_id']}"
system_prompt, user_prompt = get_delete_messages(search_output_string, data, user_identity)
_tools = [DELETE_MEMORY_TOOL_GRAPH]
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
_tools = [DELETE_MEMORY_STRUCT_TOOL_GRAPH]
memory_updates = self.llm.generate_response(
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
tools=_tools,
)
to_be_deleted = []
for item in memory_updates.get("tool_calls", []):
if item.get("name") == "delete_graph_memory":
to_be_deleted.append(item.get("arguments"))
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
logger.debug(f"Deleted relationships: {to_be_deleted}")
return to_be_deleted
def _delete_entities(self, to_be_deleted, filters):
"""Delete the entities from the graph."""
user_id = filters["user_id"]
agent_id = filters.get("agent_id")
run_id = filters.get("run_id")
results = []
try:
for item in to_be_deleted:
source = item["source"]
destination = item["destination"]
relationship = item["relationship"]
where_parts = [
"n.user_id = %s", "n.name = %s",
"m.user_id = %s", "m.name = %s",
]
params = [user_id, source, user_id, destination]
if agent_id:
where_parts.extend(["n.agent_id = %s", "m.agent_id = %s"])
params.extend([agent_id, agent_id])
if run_id:
where_parts.extend(["n.run_id = %s", "m.run_id = %s"])
params.extend([run_id, run_id])
where_clause = " AND ".join(where_parts)
result = self._exec_cypher(
f"MATCH (n)-[r:{relationship}]->(m) "
f"WHERE {where_clause} "
f"DELETE r "
f"RETURN n.name, type(r), m.name",
cols=["source", "relationship", "target"],
params=tuple(params),
)
results.append(result)
self.ag.commit()
except Exception:
self.ag.rollback()
raise
return results
def _add_entities(self, to_be_added, filters, entity_type_map):
"""Add new entities to the graph. Merge nodes if they already exist.
Apache AGE does not support ``ON CREATE SET`` / ``ON MATCH SET``, so we
use ``MERGE … SET`` which always applies. The ``coalesce`` pattern
ensures ``created`` is only set on the first merge.
"""
user_id = filters["user_id"]
agent_id = filters.get("agent_id")
run_id = filters.get("run_id")
results = []
try:
for item in to_be_added:
source = item["source"]
destination = item["destination"]
relationship = item["relationship"]
source_embedding = self.embedding_model.embed(source)
dest_embedding = self.embedding_model.embed(destination)
source_match = self._find_similar_node(source_embedding, filters, threshold=self.threshold)
dest_match = self._find_similar_node(dest_embedding, filters, threshold=self.threshold)
effective_source = source_match["name"] if source_match else source
effective_dest = dest_match["name"] if dest_match else destination
# Merge source and destination nodes
self._merge_node(user_id, effective_source, source_embedding, agent_id, run_id)
self._merge_node(user_id, effective_dest, dest_embedding, agent_id, run_id)
# Merge relationship
result = self._exec_cypher(
f"MATCH (s {{user_id: %s, name: %s}}), (d {{user_id: %s, name: %s}}) "
f"MERGE (s)-[r:{relationship}]->(d) "
f"RETURN s.name, type(r), d.name",
cols=["source", "relationship", "target"],
params=(user_id, effective_source, user_id, effective_dest),
)
results.append(result)
self.ag.commit()
except Exception:
self.ag.rollback()
raise
return results
def _remove_spaces_from_entities(self, entity_list):
for item in entity_list:
item["source"] = item["source"].lower().replace(" ", "_")
item["relationship"] = sanitize_relationship_for_cypher(item["relationship"].lower().replace(" ", "_"))
item["destination"] = item["destination"].lower().replace(" ", "_")
return entity_list
def reset(self):
"""Reset the graph by clearing all nodes and relationships."""
logger.warning("Clearing graph...")
self._exec_cypher("MATCH (n) DETACH DELETE n")
self.ag.commit()
mem0/memory/graph_memory.py
+82 -37
View File
@@ -1,6 +1,6 @@
import logging
from mem0.memory.utils import format_entities, sanitize_relationship_for_cypher
from mem0.memory.utils import format_entities, remove_spaces_from_entities
try:
from langchain_neo4j import Neo4jGraph
@@ -30,10 +30,10 @@ class MemoryGraph:
def __init__(self, config):
self.config = config
self.graph = Neo4jGraph(
self.config.graph_store.config.url,
self.config.graph_store.config.username,
self.config.graph_store.config.password,
self.config.graph_store.config.database,
url=self.config.graph_store.config.url,
username=self.config.graph_store.config.username,
password=self.config.graph_store.config.password,
database=self.config.graph_store.config.database,
refresh_schema=False,
driver_config={"notifications_min_severity": "OFF"},
)
@@ -93,14 +93,14 @@ class MemoryGraph:
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
def search(self, query, filters, limit=100):
def search(self, query, filters, top_k=100):
"""
Search for memories and related graph data.
Args:
query (str): Query to search for.
filters (dict): A dictionary containing filters to be applied during the search.
limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
dict: A dictionary containing:
@@ -129,6 +129,28 @@ class MemoryGraph:
return search_results
def delete(self, data, filters):
"""
Delete graph entities associated with the given memory text.
Extracts entities and relationships from the memory text using the same
pipeline as add(), then soft-deletes the matching relationships in the graph.
Args:
data (str): The memory text whose graph entities should be removed.
filters (dict): Scope filters (user_id, agent_id, run_id).
"""
try:
entity_type_map = self._retrieve_nodes_from_data(data, filters)
if not entity_type_map:
logger.debug("No entities found in memory text, skipping graph cleanup")
return
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
if to_be_deleted:
self._delete_entities(to_be_deleted, filters)
except Exception as e:
logger.error(f"Error during graph cleanup for memory delete: {e}")
def delete_all(self, filters):
# Build node properties for filtering
node_props = ["user_id: $user_id"]
@@ -149,18 +171,18 @@ class MemoryGraph:
params["run_id"] = filters["run_id"]
self.graph.query(cypher, params=params)
def get_all(self, filters, limit=100):
def get_all(self, filters, top_k=100):
"""
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
Args:
filters (dict): A dictionary containing filters to be applied during the retrieval.
limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
list: A list of dictionaries, each containing:
- 'contexts': The base data store response for each memory.
- 'entities': A list of strings representing the nodes and relationships
"""
params = {"user_id": filters["user_id"], "limit": limit}
params = {"user_id": filters["user_id"], "limit": top_k}
# Build node properties based on filters
node_props = ["user_id: $user_id"]
@@ -174,6 +196,7 @@ class MemoryGraph:
query = f"""
MATCH (n {self.node_label} {{{node_props_str}}})-[r]->(m {self.node_label} {{{node_props_str}}})
WHERE r.valid IS NULL OR r.valid = true
RETURN n.name AS source, type(r) AS relationship, m.name AS target
LIMIT $limit
"""
@@ -215,7 +238,7 @@ class MemoryGraph:
for tool_call in search_results["tool_calls"]:
if tool_call["name"] != "extract_entities":
continue
for item in tool_call["arguments"]["entities"]:
for item in tool_call.get("arguments", {}).get("entities", []):
entity_type_map[item["entity"]] = item["entity_type"]
except Exception as e:
logger.exception(
@@ -268,7 +291,7 @@ class MemoryGraph:
logger.debug(f"Extracted entities: {entities}")
return entities
def _search_graph_db(self, node_list, filters, limit=100):
def _search_graph_db(self, node_list, filters, top_k=100):
"""Search similar nodes among and their respective incoming and outgoing relations."""
result_relations = []
@@ -291,10 +314,12 @@ class MemoryGraph:
CALL {{
WITH n
MATCH (n)-[r]->(m {self.node_label} {{{node_props_str}}})
WHERE r.valid IS NULL OR r.valid = true
RETURN n.name AS source, elementId(n) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, m.name AS destination, elementId(m) AS destination_id
UNION
WITH n
MATCH (n)<-[r]-(m {self.node_label} {{{node_props_str}}})
WHERE r.valid IS NULL OR r.valid = true
RETURN m.name AS source, elementId(m) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, n.name AS destination, elementId(n) AS destination_id
}}
WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
@@ -307,7 +332,7 @@ class MemoryGraph:
"n_embedding": n_embedding,
"threshold": self.threshold,
"user_id": filters["user_id"],
"limit": limit,
"limit": top_k,
}
if filters.get("agent_id"):
params["agent_id"] = filters["agent_id"]
@@ -392,13 +417,15 @@ class MemoryGraph:
source_props_str = ", ".join(source_props)
dest_props_str = ", ".join(dest_props)
# Delete the specific relationship between nodes
# Soft-delete: mark relationship as invalid instead of removing it,
# enabling temporal reasoning over historical graph state.
# See: https://github.com/mem0ai/mem0/issues/4187
cypher = f"""
MATCH (n {self.node_label} {{{source_props_str}}})
-[r:{relationship}]->
(m {self.node_label} {{{dest_props_str}}})
DELETE r
WHERE r.valid IS NULL OR r.valid = true
SET r.valid = false, r.invalidated_at = datetime()
RETURN
n.name AS source,
m.name AS target,
@@ -464,11 +491,16 @@ class MemoryGraph:
CALL db.create.setNodeVectorProperty(destination, 'embedding', $destination_embedding)
WITH source, destination
MERGE (source)-[r:{relationship}]->(destination)
ON CREATE SET
r.created = timestamp(),
r.mentions = 1
ON CREATE SET
r.created_at = timestamp(),
r.updated_at = timestamp(),
r.mentions = 1,
r.valid = true
ON MATCH SET
r.mentions = coalesce(r.mentions, 0) + 1
r.mentions = coalesce(r.mentions, 0) + 1,
r.valid = true,
r.updated_at = timestamp(),
r.invalidated_at = null
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
"""
@@ -508,11 +540,16 @@ class MemoryGraph:
CALL db.create.setNodeVectorProperty(source, 'embedding', $source_embedding)
WITH source, destination
MERGE (source)-[r:{relationship}]->(destination)
ON CREATE SET
r.created = timestamp(),
r.mentions = 1
ON CREATE SET
r.created_at = timestamp(),
r.updated_at = timestamp(),
r.mentions = 1,
r.valid = true
ON MATCH SET
r.mentions = coalesce(r.mentions, 0) + 1
r.mentions = coalesce(r.mentions, 0) + 1,
r.valid = true,
r.updated_at = timestamp(),
r.invalidated_at = null
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
"""
@@ -537,11 +574,16 @@ class MemoryGraph:
WHERE elementId(destination) = $destination_id
SET destination.mentions = coalesce(destination.mentions, 0) + 1
MERGE (source)-[r:{relationship}]->(destination)
ON CREATE SET
ON CREATE SET
r.created_at = timestamp(),
r.updated_at = timestamp(),
r.mentions = 1
ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
r.mentions = 1,
r.valid = true
ON MATCH SET
r.mentions = coalesce(r.mentions, 0) + 1,
r.valid = true,
r.updated_at = timestamp(),
r.invalidated_at = null
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
"""
@@ -585,10 +627,18 @@ class MemoryGraph:
WITH source, destination
CALL db.create.setNodeVectorProperty(destination, 'embedding', $dest_embedding)
WITH source, destination
MERGE (source)-[rel:{relationship}]->(destination)
ON CREATE SET rel.created = timestamp(), rel.mentions = 1
ON MATCH SET rel.mentions = coalesce(rel.mentions, 0) + 1
RETURN source.name AS source, type(rel) AS relationship, destination.name AS target
MERGE (source)-[r:{relationship}]->(destination)
ON CREATE SET
r.created_at = timestamp(),
r.updated_at = timestamp(),
r.mentions = 1,
r.valid = true
ON MATCH SET
r.mentions = coalesce(r.mentions, 0) + 1,
r.valid = true,
r.updated_at = timestamp(),
r.invalidated_at = null
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
"""
params = {
@@ -607,12 +657,7 @@ class MemoryGraph:
return results
def _remove_spaces_from_entities(self, entity_list):
for item in entity_list:
item["source"] = item["source"].lower().replace(" ", "_")
# Use the sanitization function for relationships to handle special characters
item["relationship"] = sanitize_relationship_for_cypher(item["relationship"].lower().replace(" ", "_"))
item["destination"] = item["destination"].lower().replace(" ", "_")
return entity_list
return remove_spaces_from_entities(entity_list, sanitize_relationship=True)
def _search_source_node(self, source_embedding, filters, threshold=0.9):
# Build WHERE conditions
mem0/memory/kuzu_memory.py
+34 -16
View File
@@ -1,6 +1,6 @@
import logging
from mem0.memory.utils import format_entities
from mem0.memory.utils import format_entities, remove_spaces_from_entities
try:
import kuzu
@@ -113,14 +113,14 @@ class MemoryGraph:
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
def search(self, query, filters, limit=5):
def search(self, query, filters, top_k=5):
"""
Search for memories and related graph data.
Args:
query (str): Query to search for.
filters (dict): A dictionary containing filters to be applied during the search.
limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
dict: A dictionary containing:
@@ -139,7 +139,7 @@ class MemoryGraph:
bm25 = BM25Okapi(search_outputs_sequence)
tokenized_query = query.split(" ")
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=limit)
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=top_k)
search_results = []
for item in reranked_results:
@@ -149,6 +149,28 @@ class MemoryGraph:
return search_results
def delete(self, data, filters):
"""
Delete graph entities associated with the given memory text.
Extracts entities and relationships from the memory text using the same
pipeline as add(), then deletes the matching relationships in the graph.
Args:
data (str): The memory text whose graph entities should be removed.
filters (dict): Scope filters (user_id, agent_id, run_id).
"""
try:
entity_type_map = self._retrieve_nodes_from_data(data, filters)
if not entity_type_map:
logger.debug("No entities found in memory text, skipping graph cleanup")
return
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
if to_be_deleted:
self._delete_entities(to_be_deleted, filters)
except Exception as e:
logger.error(f"Error during graph cleanup for memory delete: {e}")
def delete_all(self, filters):
# Build node properties for filtering
node_props = ["user_id: $user_id"]
@@ -169,12 +191,12 @@ class MemoryGraph:
params["run_id"] = filters["run_id"]
self.kuzu_execute(cypher, parameters=params)
def get_all(self, filters, limit=100):
def get_all(self, filters, top_k=100):
"""
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
Args:
filters (dict): A dictionary containing filters to be applied during the retrieval.
limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
list: A list of dictionaries, each containing:
- 'contexts': The base data store response for each memory.
@@ -183,7 +205,7 @@ class MemoryGraph:
params = {
"user_id": filters["user_id"],
"limit": limit,
"limit": top_k,
}
# Build node properties based on filters
node_props = ["user_id: $user_id"]
@@ -241,7 +263,7 @@ class MemoryGraph:
for tool_call in search_results["tool_calls"]:
if tool_call["name"] != "extract_entities":
continue
for item in tool_call["arguments"]["entities"]:
for item in tool_call.get("arguments", {}).get("entities", []):
entity_type_map[item["entity"]] = item["entity_type"]
except Exception as e:
logger.exception(
@@ -294,14 +316,14 @@ class MemoryGraph:
logger.debug(f"Extracted entities: {entities}")
return entities
def _search_graph_db(self, node_list, filters, limit=100, threshold=None):
def _search_graph_db(self, node_list, filters, top_k=100, threshold=None):
"""Search similar nodes among and their respective incoming and outgoing relations."""
result_relations = []
params = {
"threshold": threshold if threshold else self.threshold,
"user_id": filters["user_id"],
"limit": limit,
"limit": top_k,
}
# Build node properties for filtering
node_props = ["user_id: $user_id"]
@@ -342,7 +364,7 @@ class MemoryGraph:
parameters=params))
# Kuzu does not support sort/limit over unions. Do it manually for now.
result_relations.extend(sorted(results, key=lambda x: x["similarity"], reverse=True)[:limit])
result_relations.extend(sorted(results, key=lambda x: x["similarity"], reverse=True)[:top_k])
return result_relations
@@ -632,11 +654,7 @@ class MemoryGraph:
return results
def _remove_spaces_from_entities(self, entity_list):
for item in entity_list:
item["source"] = item["source"].lower().replace(" ", "_")
item["relationship"] = item["relationship"].lower().replace(" ", "_")
item["destination"] = item["destination"].lower().replace(" ", "_")
return entity_list
return remove_spaces_from_entities(entity_list, sanitize_relationship=False)
def _search_source_node(self, source_embedding, filters, threshold=0.9):
params = {
mem0/memory/main.py
+486 -314
View File
File diff suppressed because it is too large Load Diff
mem0/memory/memgraph_memory.py
+34 -17
View File
@@ -1,6 +1,6 @@
import logging
from mem0.memory.utils import format_entities, sanitize_relationship_for_cypher
from mem0.memory.utils import format_entities, remove_spaces_from_entities
try:
from langchain_memgraph.graphs.memgraph import Memgraph
@@ -98,14 +98,14 @@ class MemoryGraph:
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
def search(self, query, filters, limit=100):
def search(self, query, filters, top_k=100):
"""
Search for memories and related graph data.
Args:
query (str): Query to search for.
filters (dict): A dictionary containing filters to be applied during the search.
limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
dict: A dictionary containing:
@@ -134,6 +134,28 @@ class MemoryGraph:
return search_results
def delete(self, data, filters):
"""
Delete graph entities associated with the given memory text.
Extracts entities and relationships from the memory text using the same
pipeline as add(), then deletes the matching relationships in the graph.
Args:
data (str): The memory text whose graph entities should be removed.
filters (dict): Scope filters (user_id, agent_id).
"""
try:
entity_type_map = self._retrieve_nodes_from_data(data, filters)
if not entity_type_map:
logger.debug("No entities found in memory text, skipping graph cleanup")
return
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
if to_be_deleted:
self._delete_entities(to_be_deleted, filters)
except Exception as e:
logger.error(f"Error during graph cleanup for memory delete: {e}")
def delete_all(self, filters):
"""Delete all nodes and relationships for a user or specific agent."""
if filters.get("agent_id"):
@@ -150,14 +172,14 @@ class MemoryGraph:
params = {"user_id": filters["user_id"]}
self.graph.query(cypher, params=params)
def get_all(self, filters, limit=100):
def get_all(self, filters, top_k=100):
"""
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
Args:
filters (dict): A dictionary containing filters to be applied during the retrieval.
Supports 'user_id' (required) and 'agent_id' (optional).
limit (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
Returns:
list: A list of dictionaries, each containing:
- 'source': The source node name.
@@ -171,14 +193,14 @@ class MemoryGraph:
RETURN n.name AS source, type(r) AS relationship, m.name AS target
LIMIT $limit
"""
params = {"user_id": filters["user_id"], "agent_id": filters["agent_id"], "limit": limit}
params = {"user_id": filters["user_id"], "agent_id": filters["agent_id"], "limit": top_k}
else:
query = """
MATCH (n:Entity {user_id: $user_id})-[r]->(m:Entity {user_id: $user_id})
RETURN n.name AS source, type(r) AS relationship, m.name AS target
LIMIT $limit
"""
params = {"user_id": filters["user_id"], "limit": limit}
params = {"user_id": filters["user_id"], "limit": top_k}
results = self.graph.query(query, params=params)
@@ -218,7 +240,7 @@ class MemoryGraph:
for tool_call in search_results["tool_calls"]:
if tool_call["name"] != "extract_entities":
continue
for item in tool_call["arguments"]["entities"]:
for item in tool_call.get("arguments", {}).get("entities", []):
if "entity" in item and "entity_type" in item:
entity_type_map[item["entity"]] = item["entity_type"]
except Exception as e:
@@ -271,7 +293,7 @@ class MemoryGraph:
logger.debug(f"Extracted entities: {entities}")
return entities
def _search_graph_db(self, node_list, filters, limit=100):
def _search_graph_db(self, node_list, filters, top_k=100):
"""Search similar nodes among and their respective incoming and outgoing relations."""
result_relations = []
@@ -302,7 +324,7 @@ class MemoryGraph:
"threshold": self.threshold,
"user_id": filters["user_id"],
"agent_id": filters["agent_id"],
"limit": limit,
"limit": top_k,
}
else:
cypher_query = """
@@ -326,7 +348,7 @@ class MemoryGraph:
"n_embedding": n_embedding,
"threshold": self.threshold,
"user_id": filters["user_id"],
"limit": limit,
"limit": top_k,
}
ans = self.graph.query(cypher_query, params=params)
@@ -528,12 +550,7 @@ class MemoryGraph:
return results
def _remove_spaces_from_entities(self, entity_list):
for item in entity_list:
item["source"] = item["source"].lower().replace(" ", "_")
# Use the sanitization function for relationships to handle special characters
item["relationship"] = sanitize_relationship_for_cypher(item["relationship"].lower().replace(" ", "_"))
item["destination"] = item["destination"].lower().replace(" ", "_")
return entity_list
return remove_spaces_from_entities(entity_list, sanitize_relationship=True)
def _search_source_node(self, source_embedding, filters, threshold=0.9):
"""Search for source nodes with similar embeddings."""
mem0/memory/telemetry.py
+102 -7
View File
@@ -1,7 +1,10 @@
import atexit
import logging
import os
import platform
import random
import sys
import threading
from posthog import Posthog
@@ -20,16 +23,65 @@ if not isinstance(MEM0_TELEMETRY, bool):
logging.getLogger("posthog").setLevel(logging.CRITICAL + 1)
logging.getLogger("urllib3").setLevel(logging.CRITICAL + 1)
_logger = logging.getLogger(__name__)
# Default sampling rate for hot-path OSS events. Lifecycle events always fire at 100%.
# Override via MEM0_TELEMETRY_SAMPLE_RATE env var.
_DEFAULT_SAMPLE_RATE = 0.1
def _parse_sample_rate(raw):
"""Parse MEM0_TELEMETRY_SAMPLE_RATE env var. Never raises."""
try:
value = float(raw)
except (TypeError, ValueError):
_logger.debug("MEM0_TELEMETRY_SAMPLE_RATE %r is not a number, defaulting to %s", raw, _DEFAULT_SAMPLE_RATE)
return _DEFAULT_SAMPLE_RATE
if not 0.0 <= value <= 1.0:
_logger.debug("MEM0_TELEMETRY_SAMPLE_RATE %s out of [0.0, 1.0], defaulting to %s", value, _DEFAULT_SAMPLE_RATE)
return _DEFAULT_SAMPLE_RATE
return value
MEM0_TELEMETRY_SAMPLE_RATE = _parse_sample_rate(os.environ.get("MEM0_TELEMETRY_SAMPLE_RATE", str(_DEFAULT_SAMPLE_RATE)))
# Events that bypass sampling and always fire. Keep this set in sync with the
# event names passed to capture_event() in mem0/memory/main.py.
_LIFECYCLE_EVENTS = frozenset({"mem0.init", "mem0.reset", "mem0._create_procedural_memory"})
def _sampling_before_send(msg):
"""PostHog before_send hook: drop sampled hot-path events, annotate survivors with sample_rate."""
if not isinstance(msg, dict):
return None
event_name = msg.get("event", "")
is_lifecycle = event_name in _LIFECYCLE_EVENTS
# >= so that rate=0 drops everything and rate=1 keeps everything (random ∈ [0, 1)).
if not is_lifecycle and random.random() >= MEM0_TELEMETRY_SAMPLE_RATE:
return None
# Annotate so PostHog dashboards can extrapolate true counts via 1/sample_rate.
properties = msg.setdefault("properties", {})
properties["sample_rate"] = 1.0 if is_lifecycle else MEM0_TELEMETRY_SAMPLE_RATE
return msg
class AnonymousTelemetry:
def __init__(self, vector_store=None):
def __init__(self, vector_store=None, before_send=None):
if not MEM0_TELEMETRY:
self.posthog = None
self.user_id = None
return
self.posthog = Posthog(project_api_key=PROJECT_API_KEY, host=HOST)
try:
self.posthog = Posthog(project_api_key=PROJECT_API_KEY, host=HOST, before_send=before_send)
except TypeError:
# posthog <4.5.0 does not accept before_send; fall back without sampling.
_logger.debug("posthog.Posthog does not accept before_send; upgrade to >=4.5.0 for sampling")
self.posthog = Posthog(project_api_key=PROJECT_API_KEY, host=HOST)
self.user_id = get_or_create_user_id(vector_store)
def capture_event(self, event_name, properties=None, user_email=None):
@@ -55,20 +107,63 @@ class AnonymousTelemetry:
def close(self):
if self.posthog is not None:
self.posthog.shutdown()
self.posthog = None
# Thread-safe lazy singleton for OSS telemetry.
# A single AnonymousTelemetry instance (and its underlying PostHog client /
# background thread) is reused for the lifetime of the process instead of
# creating a new one on every capture_event() call. The singleton is shut down
# once at process exit via an atexit handler.
_oss_telemetry_instance = None
_oss_telemetry_lock = threading.Lock()
_oss_telemetry_shutting_down = False
def _get_oss_telemetry():
"""Return the process-wide AnonymousTelemetry singleton, creating it on first call.
Returns None after _shutdown_oss_telemetry() has run (interpreter exit).
"""
global _oss_telemetry_instance
if _oss_telemetry_shutting_down:
return None
if _oss_telemetry_instance is not None:
return _oss_telemetry_instance
with _oss_telemetry_lock:
if _oss_telemetry_shutting_down:
return None
# Double-checked locking
if _oss_telemetry_instance is not None:
return _oss_telemetry_instance
_oss_telemetry_instance = AnonymousTelemetry(before_send=_sampling_before_send)
atexit.register(_shutdown_oss_telemetry)
return _oss_telemetry_instance
def _shutdown_oss_telemetry():
global _oss_telemetry_instance, _oss_telemetry_shutting_down
with _oss_telemetry_lock:
_oss_telemetry_shutting_down = True
if _oss_telemetry_instance is not None:
_oss_telemetry_instance.close()
_oss_telemetry_instance = None
# Module-level client telemetry singleton (used by capture_client_event).
# No before_send — hosted MemoryClient traffic must never be sampled.
client_telemetry = AnonymousTelemetry()
atexit.register(client_telemetry.close)
def capture_event(event_name, memory_instance, additional_data=None):
if not MEM0_TELEMETRY:
return
oss_telemetry = AnonymousTelemetry(
vector_store=memory_instance._telemetry_vector_store
if hasattr(memory_instance, "_telemetry_vector_store")
else None,
)
oss_telemetry = _get_oss_telemetry()
if oss_telemetry is None:
return
event_data = {
"collection": memory_instance.collection_name,
mem0/memory/utils.py
+90 -3
View File
@@ -1,12 +1,16 @@
import hashlib
import logging
import re
from typing import Any, Dict, List
from mem0.configs.prompts import (
AGENT_MEMORY_EXTRACTION_PROMPT,
FACT_RETRIEVAL_PROMPT,
USER_MEMORY_EXTRACTION_PROMPT,
AGENT_MEMORY_EXTRACTION_PROMPT,
)
logger = logging.getLogger(__name__)
def get_fact_retrieval_messages(message, is_agent_memory=False):
"""Get fact retrieval messages based on the memory type.
@@ -29,6 +33,31 @@ def get_fact_retrieval_messages_legacy(message):
return FACT_RETRIEVAL_PROMPT, f"Input:\n{message}"
def ensure_json_instruction(system_prompt, user_prompt):
"""Ensure the word 'json' appears in the prompts when using json_object response format.
OpenAI's API requires the word 'json' to appear in the messages when
response_format is set to {"type": "json_object"}. When users provide a
custom_instructions that doesn't include 'json', this causes a
400 error. This function appends a JSON format instruction to the system
prompt if 'json' is not already present in either prompt.
Args:
system_prompt: The system prompt string
user_prompt: The user prompt string
Returns:
tuple: (system_prompt, user_prompt) with JSON instruction added if needed
"""
combined = (system_prompt + user_prompt).lower()
if "json" not in combined:
system_prompt += (
"\n\nYou must return your response in valid JSON format "
"with a 'facts' key containing an array of strings."
)
return system_prompt, user_prompt
def parse_messages(messages):
response = ""
for msg in messages:
@@ -52,6 +81,30 @@ def format_entities(entities):
return "\n".join(formatted_lines)
def normalize_facts(raw_facts):
"""Normalize LLM-extracted facts to a list of strings.
Smaller LLMs (e.g. llama3.1:8b) sometimes return facts as objects
like {"fact": "..."} or {"text": "..."} instead of plain strings.
This mirrors the TypeScript FactRetrievalSchema validation.
"""
if not raw_facts:
return []
normalized = []
for item in raw_facts:
if isinstance(item, str):
fact = item
elif isinstance(item, dict):
fact = item.get("fact") or item.get("text")
if fact is None:
logger.warning("Unexpected fact shape from LLM, skipping: %s", item)
continue
else:
fact = str(item)
if fact:
normalized.append(fact)
return normalized
def remove_code_blocks(content: str) -> str:
"""
@@ -72,14 +125,20 @@ def remove_code_blocks(content: str) -> str:
def extract_json(text):
"""
Extracts JSON content from a string, removing enclosing triple backticks and optional 'json' tag if present.
If no code block is found, returns the text as-is.
If no code block is found, attempts to locate JSON by finding the first '{' and last '}'.
If that also fails, returns the text as-is.
"""
text = text.strip()
match = re.search(r"```(?:json)?\s*(.*?)\s*```", text, re.DOTALL)
if match:
json_str = match.group(1)
else:
json_str = text # assume it's raw JSON
start_idx = text.find("{")
end_idx = text.rfind("}")
if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
json_str = text[start_idx : end_idx + 1]
else:
json_str = text
return json_str
@@ -197,6 +256,7 @@ def sanitize_relationship_for_cypher(relationship) -> str:
"}": "_rbrace_",
"<": "_langle_",
">": "_rangle_",
"-": "_",
}
# Apply replacements and clean up
@@ -206,3 +266,30 @@ def sanitize_relationship_for_cypher(relationship) -> str:
return re.sub(r"_+", "_", sanitized).strip("_")
def remove_spaces_from_entities(
entity_list: List[Any],
*,
sanitize_relationship: bool = True,
) -> List[Dict[str, Any]]:
"""
Normalize entity relation dicts from LLM/tool output: lowercase, spaces to underscores.
Skips entries that are not non-empty dicts or that lack any of
``source``, ``relationship``, or ``destination`` (avoids KeyError on ``[{}]``
or partial dicts).
"""
required = ("source", "relationship", "destination")
cleaned: List[Dict[str, Any]] = []
for item in entity_list:
if not isinstance(item, dict) or not item:
continue
if not all(key in item for key in required):
continue
item["source"] = item["source"].lower().replace(" ", "_")
rel = item["relationship"].lower().replace(" ", "_")
item["relationship"] = sanitize_relationship_for_cypher(rel) if sanitize_relationship else rel
item["destination"] = item["destination"].lower().replace(" ", "_")
cleaned.append(item)
return cleaned
mem0/proxy/main.py
+4 -4
View File
@@ -59,7 +59,7 @@ class Completions:
run_id: Optional[str] = None,
metadata: Optional[dict] = None,
filters: Optional[dict] = None,
limit: Optional[int] = 10,
top_k: Optional[int] = 10,
# LLM arguments
timeout: Optional[Union[float, str, httpx.Timeout]] = None,
temperature: Optional[float] = None,
@@ -103,7 +103,7 @@ class Completions:
prepared_messages = self._prepare_messages(messages)
if prepared_messages[-1]["role"] == "user":
self._async_add_to_memory(messages, user_id, agent_id, run_id, metadata, filters)
relevant_memories = self._fetch_relevant_memories(messages, user_id, agent_id, run_id, filters, limit)
relevant_memories = self._fetch_relevant_memories(messages, user_id, agent_id, run_id, filters, top_k)
logger.debug(f"Retrieved {len(relevant_memories)} relevant memories")
prepared_messages[-1]["content"] = self._format_query_with_memories(messages, relevant_memories)
@@ -163,7 +163,7 @@ class Completions:
threading.Thread(target=add_task, daemon=True).start()
def _fetch_relevant_memories(self, messages, user_id, agent_id, run_id, filters, limit):
def _fetch_relevant_memories(self, messages, user_id, agent_id, run_id, filters, top_k):
# Currently, only pass the last 6 messages to the search API to prevent long query
message_input = [f"{message['role']}: {message['content']}" for message in messages][-6:]
# TODO: Make it better by summarizing the past conversation
@@ -173,7 +173,7 @@ class Completions:
agent_id=agent_id,
run_id=run_id,
filters=filters,
limit=limit,
top_k=top_k,
)
def _format_query_with_memories(self, messages, relevant_memories):
mem0/reranker/llm_reranker.py
+25 -14
View File
@@ -1,10 +1,10 @@
import re
from typing import List, Dict, Any, Union
from typing import Any, Dict, List, Union
from mem0.reranker.base import BaseReranker
from mem0.utils.factory import LlmFactory
from mem0.configs.rerankers.base import BaseRerankerConfig
from mem0.configs.rerankers.llm import LLMRerankerConfig
from mem0.reranker.base import BaseReranker
from mem0.utils.factory import LlmFactory
class LLMReranker(BaseReranker):
@@ -33,19 +33,30 @@ class LLMReranker(BaseReranker):
self.config = config
# Create LLM configuration for the factory
llm_config = {
"model": self.config.model,
"temperature": self.config.temperature,
"max_tokens": self.config.max_tokens,
}
# Add API key if provided
if self.config.api_key:
llm_config["api_key"] = self.config.api_key
# If a nested ``llm`` dict is provided (e.g. for non-OpenAI providers
# like Ollama that need provider-specific fields such as
# ``ollama_base_url``), use it to configure the LLM factory.
if self.config.llm:
nested = self.config.llm
llm_provider = nested.get("provider", self.config.provider)
llm_config: dict = dict(nested.get("config") or {})
llm_config.setdefault("model", self.config.model)
llm_config.setdefault("temperature", self.config.temperature)
llm_config.setdefault("max_tokens", self.config.max_tokens)
if self.config.api_key:
llm_config.setdefault("api_key", self.config.api_key)
else:
llm_provider = self.config.provider
llm_config = {
"model": self.config.model,
"temperature": self.config.temperature,
"max_tokens": self.config.max_tokens,
}
if self.config.api_key:
llm_config["api_key"] = self.config.api_key
# Initialize LLM using the factory
self.llm = LlmFactory.create(self.config.provider, llm_config)
self.llm = LlmFactory.create(llm_provider, llm_config)
# Default scoring prompt
self.scoring_prompt = getattr(self.config, 'scoring_prompt', None) or self._get_default_prompt()
mem0/utils/factory.py
+6 -1
View File
@@ -3,9 +3,11 @@ from typing import Dict, Optional, Union
from mem0.configs.embeddings.base import BaseEmbedderConfig
from mem0.configs.llms.anthropic import AnthropicConfig
from mem0.configs.llms.aws_bedrock import AWSBedrockConfig
from mem0.configs.llms.azure import AzureOpenAIConfig
from mem0.configs.llms.base import BaseLlmConfig
from mem0.configs.llms.deepseek import DeepSeekConfig
from mem0.configs.llms.minimax import MinimaxConfig
from mem0.configs.llms.lmstudio import LMStudioConfig
from mem0.configs.llms.ollama import OllamaConfig
from mem0.configs.llms.openai import OpenAIConfig
@@ -37,7 +39,7 @@ class LlmFactory:
"openai": ("mem0.llms.openai.OpenAILLM", OpenAIConfig),
"groq": ("mem0.llms.groq.GroqLLM", BaseLlmConfig),
"together": ("mem0.llms.together.TogetherLLM", BaseLlmConfig),
"aws_bedrock": ("mem0.llms.aws_bedrock.AWSBedrockLLM", BaseLlmConfig),
"aws_bedrock": ("mem0.llms.aws_bedrock.AWSBedrockLLM", AWSBedrockConfig),
"litellm": ("mem0.llms.litellm.LiteLLM", BaseLlmConfig),
"azure_openai": ("mem0.llms.azure_openai.AzureOpenAILLM", AzureOpenAIConfig),
"openai_structured": ("mem0.llms.openai_structured.OpenAIStructuredLLM", OpenAIConfig),
@@ -45,6 +47,7 @@ class LlmFactory:
"azure_openai_structured": ("mem0.llms.azure_openai_structured.AzureOpenAIStructuredLLM", AzureOpenAIConfig),
"gemini": ("mem0.llms.gemini.GeminiLLM", BaseLlmConfig),
"deepseek": ("mem0.llms.deepseek.DeepSeekLLM", DeepSeekConfig),
"minimax": ("mem0.llms.minimax.MiniMaxLLM", MinimaxConfig),
"xai": ("mem0.llms.xai.XAILLM", BaseLlmConfig),
"sarvam": ("mem0.llms.sarvam.SarvamLLM", BaseLlmConfig),
"lmstudio": ("mem0.llms.lmstudio.LMStudioLLM", LMStudioConfig),
@@ -186,6 +189,7 @@ class VectorStoreFactory:
"baidu": "mem0.vector_stores.baidu.BaiduDB",
"cassandra": "mem0.vector_stores.cassandra.CassandraDB",
"neptune": "mem0.vector_stores.neptune_analytics.NeptuneAnalyticsVector",
"turbopuffer": "mem0.vector_stores.turbopuffer.TurbopufferDB",
}
@classmethod
@@ -216,6 +220,7 @@ class GraphStoreFactory:
"neptune": "mem0.graphs.neptune.neptunegraph.MemoryGraph",
"neptunedb": "mem0.graphs.neptune.neptunedb.MemoryGraph",
"kuzu": "mem0.memory.kuzu_memory.MemoryGraph",
"apache_age": "mem0.memory.apache_age_memory.MemoryGraph",
"default": "mem0.memory.graph_memory.MemoryGraph",
}
mem0/vector_stores/azure_ai_search.py
+11 -11
View File
@@ -205,14 +205,14 @@ class AzureAISearch(VectorStoreBase):
filter_expression = " and ".join(filter_conditions)
return filter_expression
def search(self, query, vectors, limit=5, filters=None):
def search(self, query, vectors, top_k=5, filters=None):
"""
Search for similar vectors.
Args:
query (str): Query.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -222,13 +222,13 @@ class AzureAISearch(VectorStoreBase):
if filters:
filter_expression = self._build_filter_expression(filters)
vector_query = VectorizedQuery(vector=vectors, k_nearest_neighbors=limit, fields="vector")
vector_query = VectorizedQuery(vector=vectors, k_nearest_neighbors=top_k, fields="vector")
if self.hybrid_search:
search_results = self.search_client.search(
search_text=query,
vector_queries=[vector_query],
filter=filter_expression,
top=limit,
top=top_k,
vector_filter_mode=self.vector_filter_mode,
search_fields=["payload"],
)
@@ -236,7 +236,7 @@ class AzureAISearch(VectorStoreBase):
search_results = self.search_client.search(
vector_queries=[vector_query],
filter=filter_expression,
top=limit,
top=top_k,
vector_filter_mode=self.vector_filter_mode,
)
@@ -246,12 +246,12 @@ class AzureAISearch(VectorStoreBase):
results.append(OutputData(id=result["id"], score=result["@search.score"], payload=payload))
return results
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""Search for memories using keyword/BM25 text matching (no vector queries).
Args:
query (str): The text query to search for.
limit (int): Maximum number of results to return. Defaults to 5.
top_k (int): Maximum number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search.
Returns:
@@ -264,7 +264,7 @@ class AzureAISearch(VectorStoreBase):
search_results = self.search_client.search(
search_text=query,
filter=filter_expression,
top=limit,
top=top_k,
search_fields=["payload"],
)
@@ -355,13 +355,13 @@ class AzureAISearch(VectorStoreBase):
index = self.index_client.get_index(self.index_name)
return {"name": index.name, "fields": index.fields}
def list(self, filters=None, limit=100):
def list(self, filters=None, top_k=100):
"""
List all vectors in the index.
Args:
filters (dict, optional): Filters to apply to the list.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
@@ -370,7 +370,7 @@ class AzureAISearch(VectorStoreBase):
if filters:
filter_expression = self._build_filter_expression(filters)
search_results = self.search_client.search(search_text="*", filter=filter_expression, top=limit)
search_results = self.search_client.search(search_text="*", filter=filter_expression, top=top_k)
results = []
for result in search_results:
payload = json.loads(extract_json(result["payload"]))
mem0/vector_stores/azure_mysql.py
+11 -11
View File
@@ -7,8 +7,8 @@ from pydantic import BaseModel
try:
import pymysql
from pymysql.cursors import DictCursor
from dbutils.pooled_db import PooledDB
from pymysql.cursors import DictCursor
except ImportError:
raise ImportError(
"Azure MySQL vector store requires PyMySQL and DBUtils. "
@@ -258,7 +258,7 @@ class AzureMySQL(VectorStoreBase):
self,
query: str,
vectors: List[float],
limit: int = 5,
top_k: int = 5,
filters: Optional[Dict] = None,
) -> List[OutputData]:
"""
@@ -267,7 +267,7 @@ class AzureMySQL(VectorStoreBase):
Args:
query (str): Query string (not used in vector search)
vectors (List[float]): Query vector
limit (int): Number of results to return
top_k (int): Number of results to return
filters (Dict, optional): Filters to apply to the search
Returns:
@@ -306,16 +306,16 @@ class AzureMySQL(VectorStoreBase):
distance = 1 - similarity
scored_results.append((row['id'], distance, row['payload']))
# Sort by distance and limit
# Sort by distance and apply limit
scored_results.sort(key=lambda x: x[1])
scored_results = scored_results[:limit]
scored_results = scored_results[:top_k]
return [
OutputData(id=r[0], score=float(r[1]), payload=json.loads(r[2]) if isinstance(r[2], str) else r[2])
for r in scored_results
]
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""
Search for memories using MySQL FULLTEXT search via MATCH() AGAINST().
@@ -324,7 +324,7 @@ class AzureMySQL(VectorStoreBase):
Args:
query (str): The text query for keyword-based search.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -354,7 +354,7 @@ class AzureMySQL(VectorStoreBase):
ORDER BY score DESC
LIMIT %s
"""
params = [query, query] + filter_params + [limit]
params = [query, query] + filter_params + [top_k]
cur.execute(query_sql, params)
results = cur.fetchall()
@@ -476,14 +476,14 @@ class AzureMySQL(VectorStoreBase):
def list(
self,
filters: Optional[Dict] = None,
limit: int = 100
top_k: int = 100
) -> List[List[OutputData]]:
"""
List all vectors in the collection.
Args:
filters (Dict, optional): Filters to apply
limit (int): Number of vectors to return
top_k (int): Number of vectors to return
Returns:
List[List[OutputData]]: List of vectors
@@ -506,7 +506,7 @@ class AzureMySQL(VectorStoreBase):
{filter_clause}
LIMIT %s
""",
(*filter_params, limit)
(*filter_params, top_k)
)
results = cur.fetchall()
mem0/vector_stores/baidu.py
+9 -9
View File
@@ -186,14 +186,14 @@ class BaiduDB(VectorStoreBase):
row = Row(id=idx, vector=vector, metadata=metadata)
self._table.upsert(rows=[row])
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None) -> list:
"""
Search for similar vectors.
Args:
query (str): Query string.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -208,7 +208,7 @@ class BaiduDB(VectorStoreBase):
request = VectorTopkSearchRequest(
vector_field="vector",
vector=FloatVector(vectors),
limit=limit,
limit=top_k,
filter=search_filter,
config=VectorSearchConfig(ef=200),
)
@@ -228,13 +228,13 @@ class BaiduDB(VectorStoreBase):
return output
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""
Perform keyword-based search using Baidu Mochow's BM25 search.
Args:
query (str): The text query to search for.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search.
Returns:
@@ -248,7 +248,7 @@ class BaiduDB(VectorStoreBase):
request = BM25SearchRequest(
index_name="data_bm25_idx",
search_text=query,
limit=limit,
limit=top_k,
filter=search_filter,
)
@@ -356,20 +356,20 @@ class BaiduDB(VectorStoreBase):
"""
return self._table.stats()
def list(self, filters: dict = None, limit: int = 100) -> list:
def list(self, filters: dict = None, top_k: int = 100) -> list:
"""
List all vectors in the table.
Args:
filters (Dict, optional): Filters to apply to the list.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
"""
projections = ["id", "metadata"]
list_filter = self._create_filter(filters) if filters else None
result = self._table.select(filter=list_filter, projections=projections, limit=limit)
result = self._table.select(filter=list_filter, projections=projections, limit=top_k)
memories = []
for row in result.rows:
mem0/vector_stores/base.py
+7 -7
View File
@@ -13,7 +13,7 @@ class VectorStoreBase(ABC):
pass
@abstractmethod
def search(self, query, vectors, limit=5, filters=None):
def search(self, query, vectors, top_k=5, filters=None):
"""Search for similar vectors."""
pass
@@ -48,7 +48,7 @@ class VectorStoreBase(ABC):
pass
@abstractmethod
def list(self, filters=None, limit=None):
def list(self, filters=None, top_k=None):
"""List all memories."""
pass
@@ -57,7 +57,7 @@ class VectorStoreBase(ABC):
"""Reset by delete the collection and recreate it."""
pass
def keyword_search(self, query: str, limit: int = 5, filters: dict = None):
def keyword_search(self, query: str, top_k: int = 5, filters: dict = None):
"""Keyword/BM25 full-text search. Returns None if not supported by this store.
Override in subclasses that support native keyword/BM25 search.
@@ -66,7 +66,7 @@ class VectorStoreBase(ABC):
Args:
query: The search query text (should be lemmatized for best results).
limit: Maximum number of results to return.
top_k: Maximum number of results to return.
filters: Optional metadata filters (same format as search filters).
Returns:
@@ -74,7 +74,7 @@ class VectorStoreBase(ABC):
"""
return None
def search_batch(self, queries: list, vectors_list: list, limit: int = 1, filters: dict = None):
def search_batch(self, queries: list, vectors_list: list, top_k: int = 1, filters: dict = None):
"""Batch search for multiple queries at once.
Default implementation calls search() sequentially. Override in subclasses
@@ -83,10 +83,10 @@ class VectorStoreBase(ABC):
Args:
queries: List of query texts.
vectors_list: List of query vectors (one per query).
limit: Maximum results per query.
top_k: Maximum results per query.
filters: Optional metadata filters applied to all queries.
Returns:
List of result lists, one per query.
"""
return [self.search(q, v, limit=limit, filters=filters) for q, v in zip(queries, vectors_list)]
return [self.search(q, v, top_k=top_k, filters=filters) for q, v in zip(queries, vectors_list)]
mem0/vector_stores/cassandra.py
+8 -8
View File
@@ -7,8 +7,8 @@ import numpy as np
from pydantic import BaseModel
try:
from cassandra.cluster import Cluster
from cassandra.auth import PlainTextAuthProvider
from cassandra.cluster import Cluster
except ImportError:
raise ImportError(
"Apache Cassandra vector store requires cassandra-driver. "
@@ -214,7 +214,7 @@ class CassandraDB(VectorStoreBase):
self,
query: str,
vectors: List[float],
limit: int = 5,
top_k: int = 5,
filters: Optional[Dict] = None,
) -> List[OutputData]:
"""
@@ -223,7 +223,7 @@ class CassandraDB(VectorStoreBase):
Args:
query (str): Query string (not used in vector search)
vectors (List[float]): Query vector
limit (int): Number of results to return
top_k (int): Number of results to return
filters (Dict, optional): Filters to apply to the search
Returns:
@@ -263,9 +263,9 @@ class CassandraDB(VectorStoreBase):
scored_results.append((row.id, distance, row.payload))
# Sort by distance and limit
# Sort by distance and apply limit
scored_results.sort(key=lambda x: x[1])
scored_results = scored_results[:limit]
scored_results = scored_results[:top_k]
return [
OutputData(
@@ -427,14 +427,14 @@ class CassandraDB(VectorStoreBase):
def list(
self,
filters: Optional[Dict] = None,
limit: int = 100
top_k: int = 100
) -> List[List[OutputData]]:
"""
List all vectors in the collection.
Args:
filters (Dict, optional): Filters to apply
limit (int): Number of vectors to return
top_k (int): Number of vectors to return
Returns:
List[List[OutputData]]: List of vectors
@@ -443,7 +443,7 @@ class CassandraDB(VectorStoreBase):
query = f"""
SELECT id, vector, payload
FROM {self.keyspace}.{self.collection_name}
LIMIT {limit}
LIMIT {top_k}
"""
rows = self.session.execute(query)
mem0/vector_stores/chroma.py
+6 -6
View File
@@ -141,7 +141,7 @@ class ChromaDB(VectorStoreBase):
self.collection.add(ids=ids, embeddings=vectors, metadatas=payloads)
def search(
self, query: str, vectors: List[list], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[list], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
@@ -149,14 +149,14 @@ class ChromaDB(VectorStoreBase):
Args:
query (str): Query.
vectors (List[list]): List of vectors to search.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
Returns:
List[OutputData]: Search results.
"""
where_clause = self._generate_where_clause(filters) if filters else None
results = self.collection.query(query_embeddings=vectors, where=where_clause, n_results=limit)
results = self.collection.query(query_embeddings=vectors, where=where_clause, n_results=top_k)
final_results = self._parse_output(results)
return final_results
@@ -222,19 +222,19 @@ class ChromaDB(VectorStoreBase):
"""
return self.client.get_collection(name=self.collection_name)
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List all vectors in a collection.
Args:
filters (Optional[Dict], optional): Filters to apply to the list. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
"""
where_clause = self._generate_where_clause(filters) if filters else None
results = self.collection.get(where=where_clause, limit=limit)
results = self.collection.get(where=where_clause, limit=top_k)
return [self._parse_output(results)]
def reset(self):
mem0/vector_stores/configs.py
+1
View File
@@ -34,6 +34,7 @@ class VectorStoreConfig(BaseModel):
"faiss": "FAISSConfig",
"langchain": "LangchainConfig",
"s3_vectors": "S3VectorsConfig",
"turbopuffer": "TurbopufferConfig",
}
@model_validator(mode="after")
mem0/vector_stores/databricks.py
+132 -67
View File
@@ -1,19 +1,29 @@
import json
import logging
import re
import uuid
from typing import Optional, List
from datetime import datetime, date
from databricks.sdk.service.catalog import ColumnInfo, ColumnTypeName, TableType, DataSourceFormat
from databricks.sdk.service.catalog import TableConstraint, PrimaryKeyConstraint
from datetime import date, datetime
from typing import List, Optional
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.catalog import (
ColumnInfo,
ColumnTypeName,
DataSourceFormat,
PrimaryKeyConstraint,
TableConstraint,
TableType,
)
from databricks.sdk.service.sql import StatementParameterListItem
from databricks.sdk.service.vectorsearch import (
VectorIndexType,
DeltaSyncVectorIndexSpecRequest,
DirectAccessVectorIndexSpec,
EmbeddingSourceColumn,
EmbeddingVectorColumn,
VectorIndexType,
)
from pydantic import BaseModel
from mem0.memory.utils import extract_json
from mem0.vector_stores.base import VectorStoreBase
@@ -28,6 +38,9 @@ class MemoryResult(BaseModel):
excluded_keys = {"user_id", "agent_id", "run_id", "hash", "data", "created_at", "updated_at"}
# Pattern for valid SQL identifiers to prevent column name injection
_VALID_SQL_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
class Databricks(VectorStoreBase):
def __init__(
@@ -65,7 +78,7 @@ class Databricks(VectorStoreBase):
catalog (str): Unity Catalog catalog name.
schema (str): Unity Catalog schema name.
table_name (str): Source Delta table name.
index_name (str, optional): Vector search index name (default: "mem0").
collection_name (str, optional): Vector search index name (default: "mem0").
index_type (str, optional): Index type, either "DELTA_SYNC" or "DIRECT_ACCESS" (default: "DELTA_SYNC").
embedding_model_endpoint_name (str, optional): Embedding model endpoint for Databricks-computed embeddings.
embedding_dimension (int, optional): Vector embedding dimensions (default: 1536).
@@ -85,7 +98,7 @@ class Databricks(VectorStoreBase):
self.fully_qualified_index_name = f"{self.catalog}.{self.schema}.{self.index_name}"
# Configuration
self.index_type = index_type
self.index_type = VectorIndexType(index_type) if isinstance(index_type, str) else index_type
self.embedding_model_endpoint_name = embedding_model_endpoint_name
self.embedding_dimension = embedding_dimension
self.endpoint_type = endpoint_type
@@ -261,11 +274,11 @@ class Databricks(VectorStoreBase):
)
logger.info(f"Successfully created source table '{self.fully_qualified_table_name}'")
self.client.table_constraints.create(
full_name_arg="logistics_dev.ai.dev_memory",
full_name_arg=self.fully_qualified_table_name,
constraint=TableConstraint(
primary_key_constraint=PrimaryKeyConstraint(
name="pk_dev_memory", # Name of the primary key constraint
child_columns=["memory_id"], # Columns that make up the primary key
name=f"pk_{self.table_name}",
child_columns=["memory_id"],
)
),
)
@@ -388,28 +401,46 @@ class Databricks(VectorStoreBase):
# Determine the number of items to process
num_items = len(payloads) if payloads else len(vectors) if vectors else 0
params = []
value_tuples = []
for i in range(num_items):
values = []
placeholders = []
for col in self.columns:
param_name = f"{col.name}_{i}"
if col.name == "memory_id":
val = ids[i] if ids and i < len(ids) else str(uuid.uuid4())
elif col.name == "embedding":
val = vectors[i] if vectors and i < len(vectors) else []
# Vectors are numeric arrays — ARRAY type not supported by StatementParameterListItem,
# so we inline using _format_sql_value (values are floats from the embedding model).
placeholders.append(self._format_sql_value(val))
continue
elif col.name == "memory":
val = payloads[i].get("data") if payloads and i < len(payloads) else None
else:
val = payloads[i].get(col.name) if payloads and i < len(payloads) else None
values.append(val)
formatted = [self._format_sql_value(v) for v in values]
value_tuples.append(f"({', '.join(formatted)})")
if val is None:
placeholders.append("NULL")
else:
placeholders.append(f":{param_name}")
if isinstance(val, dict):
val = json.dumps(val)
# Use explicit type for TIMESTAMP columns so Databricks doesn't
# rely on implicit STRING→TIMESTAMP casting.
param_type = "TIMESTAMP" if col.type_name == ColumnTypeName.TIMESTAMP else None
params.append(StatementParameterListItem(name=param_name, value=str(val), type=param_type))
value_tuples.append(f"({', '.join(placeholders)})")
insert_sql = f"INSERT INTO {self.fully_qualified_table_name} ({', '.join(self.column_names)}) VALUES {', '.join(value_tuples)}"
# Execute the insert
try:
response = self.client.statement_execution.execute_statement(
statement=insert_sql, warehouse_id=self.warehouse_id, wait_timeout="30s"
statement=insert_sql,
warehouse_id=self.warehouse_id,
wait_timeout="30s",
parameters=params,
)
if response.status.state.value == "SUCCEEDED":
logger.info(
@@ -423,14 +454,14 @@ class Databricks(VectorStoreBase):
logger.error(f"Insert operation failed: {e}")
raise
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> List[MemoryResult]:
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None) -> List[MemoryResult]:
"""
Search for similar vectors or text using the Databricks Vector Search index.
Args:
query (str): Search query text (for text-based search).
vectors (list): Query vector (for vector-based search).
limit (int): Maximum number of results.
top_k (int): Maximum number of results.
filters (dict): Filters to apply.
Returns:
@@ -439,29 +470,29 @@ class Databricks(VectorStoreBase):
try:
filters_json = json.dumps(filters) if filters else None
# Choose query type
if self.index_type == VectorIndexType.DELTA_SYNC and query:
# Text-based search
sdk_results = self.client.vector_search_indexes.query_index(
index_name=self.fully_qualified_index_name,
columns=self.column_names,
query_text=query,
num_results=limit,
query_type=self.query_type,
filters_json=filters_json,
)
elif self.index_type == VectorIndexType.DIRECT_ACCESS and vectors:
# Vector-based search
sdk_results = self.client.vector_search_indexes.query_index(
index_name=self.fully_qualified_index_name,
columns=self.column_names,
query_vector=vectors,
num_results=limit,
query_type=self.query_type,
filters_json=filters_json,
)
# Choose query mode per Databricks SDK contract:
# - query_text: for Delta Sync Index with model endpoint
# - query_vector: for Direct Access Index and Delta Sync Index with self-managed vectors
query_kwargs = {
"index_name": self.fully_qualified_index_name,
"columns": self.column_names,
"num_results": top_k,
"query_type": self.query_type,
"filters_json": filters_json,
}
uses_model_endpoint = (
self.index_type == VectorIndexType.DELTA_SYNC and self.embedding_model_endpoint_name
)
if uses_model_endpoint:
if not query:
raise ValueError("Query text is required for Delta Sync Index with model endpoint.")
query_kwargs["query_text"] = query
elif vectors:
query_kwargs["query_vector"] = vectors
else:
raise ValueError("Must provide query text for DELTA_SYNC or vectors for DIRECT_ACCESS.")
raise ValueError("Must provide vectors for search.")
sdk_results = self.client.vector_search_indexes.query_index(**query_kwargs)
# Parse results
result_data = sdk_results.result if hasattr(sdk_results, "result") else sdk_results
@@ -484,7 +515,7 @@ class Databricks(VectorStoreBase):
logger.error(f"Search failed: {e}")
raise
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""
Search for memories using full-text keyword search.
@@ -492,7 +523,7 @@ class Databricks(VectorStoreBase):
Args:
query (str): Search query text.
limit (int): Maximum number of results. Defaults to 5.
top_k (int): Maximum number of results. Defaults to 5.
filters (dict, optional): Filters to apply.
Returns:
@@ -509,7 +540,7 @@ class Databricks(VectorStoreBase):
index_name=self.fully_qualified_index_name,
columns=self.column_names,
query_text=query,
num_results=limit,
num_results=top_k,
query_type="FULL_TEXT",
filters_json=filters_json,
)
@@ -543,10 +574,13 @@ class Databricks(VectorStoreBase):
try:
logger.info(f"Deleting vector with ID {vector_id} from Delta table {self.fully_qualified_table_name}")
delete_sql = f"DELETE FROM {self.fully_qualified_table_name} WHERE memory_id = '{vector_id}'"
delete_sql = f"DELETE FROM {self.fully_qualified_table_name} WHERE memory_id = :vector_id"
response = self.client.statement_execution.execute_statement(
statement=delete_sql, warehouse_id=self.warehouse_id, wait_timeout="30s"
statement=delete_sql,
warehouse_id=self.warehouse_id,
wait_timeout="30s",
parameters=[StatementParameterListItem(name="vector_id", value=str(vector_id))],
)
if response.status.state.value == "SUCCEEDED":
@@ -568,8 +602,8 @@ class Databricks(VectorStoreBase):
payload (dict, optional): New payload data.
"""
update_sql = f"UPDATE {self.fully_qualified_table_name} SET "
set_clauses = []
params = []
if not vector_id:
logger.error("vector_id is required for update operation")
return
@@ -577,25 +611,38 @@ class Databricks(VectorStoreBase):
if not isinstance(vector, list):
logger.error("vector must be a list of float values")
return
set_clauses.append(f"embedding = {vector}")
# Vectors are numeric arrays — safe to inline since StatementParameterListItem
# doesn't support ARRAY types, and values are validated as list of floats above.
# Use array() SQL syntax, not Python list repr which is invalid Databricks SQL.
set_clauses.append(f"embedding = {self._format_sql_value(vector)}")
if payload:
if not isinstance(payload, dict):
logger.error("payload must be a dictionary")
return
for key, value in payload.items():
if key not in excluded_keys:
set_clauses.append(f"{key} = '{value}'")
if not _VALID_SQL_IDENTIFIER.match(key):
logger.warning(f"Skipping invalid column name in payload: {key!r}")
continue
param_name = f"payload_{key}"
set_clauses.append(f"{key} = :{param_name}")
params.append(StatementParameterListItem(name=param_name, value=str(value)))
if not set_clauses:
logger.error("No fields to update")
return
update_sql = f"UPDATE {self.fully_qualified_table_name} SET "
update_sql += ", ".join(set_clauses)
update_sql += f" WHERE memory_id = '{vector_id}'"
update_sql += " WHERE memory_id = :vector_id"
params.append(StatementParameterListItem(name="vector_id", value=str(vector_id)))
try:
logger.info(f"Updating vector with ID {vector_id} in Delta table {self.fully_qualified_table_name}")
response = self.client.statement_execution.execute_statement(
statement=update_sql, warehouse_id=self.warehouse_id, wait_timeout="30s"
statement=update_sql,
warehouse_id=self.warehouse_id,
wait_timeout="30s",
parameters=params,
)
if response.status.state.value == "SUCCEEDED":
@@ -621,14 +668,23 @@ class Databricks(VectorStoreBase):
filters = {"memory_id": vector_id}
filters_json = json.dumps(filters)
results = self.client.vector_search_indexes.query_index(
index_name=self.fully_qualified_index_name,
columns=self.column_names,
query_text=" ", # Empty query, rely on filters
num_results=1,
query_type=self.query_type,
filters_json=filters_json,
# Use query_text for Delta Sync with model endpoint, query_vector otherwise
query_kwargs = {
"index_name": self.fully_qualified_index_name,
"columns": self.column_names,
"num_results": 1,
"query_type": self.query_type,
"filters_json": filters_json,
}
uses_model_endpoint = (
self.index_type == VectorIndexType.DELTA_SYNC and self.embedding_model_endpoint_name
)
if uses_model_endpoint:
query_kwargs["query_text"] = " "
else:
query_kwargs["query_vector"] = [0.0] * self.embedding_dimension
results = self.client.vector_search_indexes.query_index(**query_kwargs)
# Process results
result_data = results.result if hasattr(results, "result") else results
@@ -638,7 +694,7 @@ class Databricks(VectorStoreBase):
raise KeyError(f"Vector with ID {vector_id} not found")
result = data_array[0]
columns = columns = [col.name for col in results.manifest.columns] if results.manifest and results.manifest.columns else []
columns = [col.name for col in results.manifest.columns] if results.manifest and results.manifest.columns else []
row_data = dict(zip(columns, result))
# Build payload following the standard schema
@@ -720,29 +776,38 @@ class Databricks(VectorStoreBase):
logger.error(f"Failed to get info for index '{name or self.index_name}': {e}")
raise
def list(self, filters: dict = None, limit: int = None) -> list[MemoryResult]:
def list(self, filters: dict = None, top_k: int = None) -> list[MemoryResult]:
"""
List all recent created memories from the vector store.
Args:
filters (dict, optional): Filters to apply.
limit (int, optional): Maximum number of results.
top_k (int, optional): Maximum number of results.
Returns:
List containing list of MemoryResult objects.
"""
try:
filters_json = json.dumps(filters) if filters else None
num_results = limit or 100
num_results = top_k or 100
columns = self.column_names
sdk_results = self.client.vector_search_indexes.query_index(
index_name=self.fully_qualified_index_name,
columns=columns,
query_text=" ",
num_results=num_results,
query_type=self.query_type,
filters_json=filters_json,
# Use query_text for Delta Sync with model endpoint, query_vector otherwise
query_kwargs = {
"index_name": self.fully_qualified_index_name,
"columns": columns,
"num_results": num_results,
"query_type": self.query_type,
"filters_json": filters_json,
}
uses_model_endpoint = (
self.index_type == VectorIndexType.DELTA_SYNC and self.embedding_model_endpoint_name
)
if uses_model_endpoint:
query_kwargs["query_text"] = " "
else:
query_kwargs["query_vector"] = [0.0] * self.embedding_dimension
sdk_results = self.client.vector_search_indexes.query_index(**query_kwargs)
result_data = sdk_results.result if hasattr(sdk_results, "result") else sdk_results
data_array = result_data.data_array if hasattr(result_data, "data_array") else []
mem0/vector_stores/elasticsearch.py
+9 -9
View File
@@ -129,7 +129,7 @@ class ElasticsearchDB(VectorStoreBase):
return results
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search with two options:
@@ -137,10 +137,10 @@ class ElasticsearchDB(VectorStoreBase):
2. Use KNN search on vectors with pre-filtering if no custom search query is provided
"""
if self.custom_search_query:
search_query = self.custom_search_query(vectors, limit, filters)
search_query = self.custom_search_query(vectors, top_k, filters)
else:
search_query = {
"knn": {"field": "vector", "query_vector": vectors, "k": limit, "num_candidates": limit * 2}
"knn": {"field": "vector", "query_vector": vectors, "k": top_k, "num_candidates": top_k * 2}
}
if filters:
filter_conditions = []
@@ -158,12 +158,12 @@ class ElasticsearchDB(VectorStoreBase):
return results
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""Search for memories using BM25 keyword matching.
Args:
query (str): The text query to search for.
limit (int): Maximum number of results to return. Defaults to 5.
top_k (int): Maximum number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search.
Returns:
@@ -187,7 +187,7 @@ class ElasticsearchDB(VectorStoreBase):
bool_query["filter"] = filter_conditions
search_query = {
"size": limit,
"size": top_k,
"query": {"bool": bool_query},
}
@@ -246,7 +246,7 @@ class ElasticsearchDB(VectorStoreBase):
"""Get information about a collection (index)."""
return self.client.indices.get(index=name)
def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None) -> List[List[OutputData]]:
def list(self, filters: Optional[Dict] = None, top_k: Optional[int] = None) -> List[List[OutputData]]:
"""List all memories."""
query: Dict[str, Any] = {"query": {"match_all": {}}}
@@ -256,8 +256,8 @@ class ElasticsearchDB(VectorStoreBase):
filter_conditions.append({"term": {f"metadata.{key}": value}})
query["query"] = {"bool": {"must": filter_conditions}}
if limit:
query["size"] = limit
if top_k:
query["size"] = top_k
response = self.client.search(index=self.collection_name, body=query)
mem0/vector_stores/faiss.py
+31 -17
View File
@@ -2,14 +2,13 @@ import logging
import os
import pickle
import uuid
import warnings
from pathlib import Path
from typing import Dict, List, Optional
import numpy as np
from pydantic import BaseModel
import warnings
try:
# Suppress SWIG deprecation warnings from FAISS
warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*SwigPy.*")
@@ -115,23 +114,23 @@ class FAISS(VectorStoreBase):
except Exception as e:
logger.warning(f"Failed to save FAISS index: {e}")
def _parse_output(self, scores, ids, limit=None) -> List[OutputData]:
def _parse_output(self, scores, ids, top_k=None) -> List[OutputData]:
"""
Parse the output data.
Args:
scores: Similarity scores from FAISS.
ids: Indices from FAISS.
limit: Maximum number of results to return.
top_k: Maximum number of results to return.
Returns:
List[OutputData]: Parsed output data.
"""
if limit is None:
limit = len(ids)
if top_k is None:
top_k = len(ids)
results = []
for i in range(min(len(ids), limit)):
for i in range(min(len(ids), top_k)):
if ids[i] == -1: # FAISS returns -1 for empty results
continue
@@ -225,7 +224,7 @@ class FAISS(VectorStoreBase):
logger.info(f"Inserted {len(vectors)} vectors into collection {self.collection_name}")
def search(
self, query: str, vectors: List[list], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[list], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
@@ -233,7 +232,7 @@ class FAISS(VectorStoreBase):
Args:
query (str): Query (not used, kept for API compatibility).
vectors (List[list]): List of vectors to search.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -250,19 +249,19 @@ class FAISS(VectorStoreBase):
if self.normalize_L2 and self.distance_strategy.lower() == "euclidean":
faiss.normalize_L2(query_vectors)
fetch_k = limit * 2 if filters else limit
fetch_k = top_k * 2 if filters else top_k
scores, indices = self.index.search(query_vectors, fetch_k)
results = self._parse_output(scores[0], indices[0], limit)
results = self._parse_output(scores[0], indices[0], top_k)
if filters:
filtered_results = []
for result in results:
if self._apply_filters(result.payload, filters):
filtered_results.append(result)
if len(filtered_results) >= limit:
if len(filtered_results) >= top_k:
break
results = filtered_results[:limit]
results = filtered_results[:top_k]
return results
@@ -309,8 +308,23 @@ class FAISS(VectorStoreBase):
break
if index_to_delete is not None:
# Reconstruct remaining vectors and rebuild the FAISS index
remaining_vectors = []
new_index_to_id = {}
new_idx = 0
for old_idx in sorted(self.index_to_id.keys()):
if old_idx == index_to_delete:
continue
remaining_vectors.append(self.index.reconstruct(int(old_idx)))
new_index_to_id[new_idx] = self.index_to_id[old_idx]
new_idx += 1
self.index.reset()
if remaining_vectors:
self.index.add(np.array(remaining_vectors, dtype=np.float32))
self.docstore.pop(vector_id, None)
self.index_to_id.pop(index_to_delete, None)
self.index_to_id = new_index_to_id
self._save()
@@ -435,13 +449,13 @@ class FAISS(VectorStoreBase):
"distance": self.distance_strategy,
}
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List all vectors in a collection.
Args:
filters (Optional[Dict], optional): Filters to apply to the list. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
@@ -467,7 +481,7 @@ class FAISS(VectorStoreBase):
)
count += 1
if count >= limit:
if count >= top_k:
break
return [results]
mem0/vector_stores/langchain.py
+6 -6
View File
@@ -91,15 +91,15 @@ class Langchain(VectorStoreBase):
texts = [payload.get("data", "") for payload in payloads] if payloads else [""] * len(vectors)
self.client.add_texts(texts=texts, metadatas=payloads, ids=ids)
def search(self, query: str, vectors: List[List[float]], limit: int = 5, filters: Optional[Dict] = None):
def search(self, query: str, vectors: List[List[float]], top_k: int = 5, filters: Optional[Dict] = None):
"""
Search for similar vectors in LangChain.
"""
# For each vector, perform a similarity search
if filters:
results = self.client.similarity_search_by_vector(embedding=vectors, k=limit, filter=filters)
results = self.client.similarity_search_by_vector(embedding=vectors, k=top_k, filter=filters)
else:
results = self.client.similarity_search_by_vector(embedding=vectors, k=limit)
results = self.client.similarity_search_by_vector(embedding=vectors, k=top_k)
final_results = self._parse_output(results)
return final_results
@@ -115,7 +115,7 @@ class Langchain(VectorStoreBase):
Update a vector and its payload.
"""
self.delete(vector_id)
self.insert(vector, payload, [vector_id])
self.insert([vector], [payload], [vector_id])
def get(self, vector_id):
"""
@@ -152,7 +152,7 @@ class Langchain(VectorStoreBase):
"""
return {"name": self.collection_name}
def list(self, filters=None, limit=None):
def list(self, filters=None, top_k=None):
"""
List all vectors in a collection.
"""
@@ -164,7 +164,7 @@ class Langchain(VectorStoreBase):
# Handle all filters, not just user_id
where_clause = filters
result = self.client._collection.get(where=where_clause, limit=limit)
result = self.client._collection.get(where=where_clause, limit=top_k)
# Convert the result to the expected format
if result and isinstance(result, dict):
mem0/vector_stores/milvus.py
+20 -9
View File
@@ -165,14 +165,14 @@ class MilvusDB(VectorStoreBase):
return memory
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None) -> list:
"""
Search for similar vectors.
Args:
query (str): Query.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -182,14 +182,14 @@ class MilvusDB(VectorStoreBase):
hits = self.client.search(
collection_name=self.collection_name,
data=[vectors],
limit=limit,
limit=top_k,
filter=query_filter,
output_fields=["*"],
)
result = self._parse_output(data=hits[0])
return result
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""
Search for memories using BM25-based full-text search via Milvus sparse vector support.
@@ -199,7 +199,7 @@ class MilvusDB(VectorStoreBase):
Args:
query (str): The text query for keyword-based search.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -212,7 +212,7 @@ class MilvusDB(VectorStoreBase):
collection_name=self.collection_name,
data=[query],
anns_field="sparse",
limit=limit,
limit=top_k,
filter=query_filter,
output_fields=["*"],
)
@@ -240,6 +240,17 @@ class MilvusDB(VectorStoreBase):
vector (List[float], optional): Updated vector.
payload (Dict, optional): Updated payload.
"""
if vector is None or payload is None:
existing = self.client.get(collection_name=self.collection_name, ids=vector_id)
if not existing:
raise ValueError(f"Vector with id {vector_id} not found in collection {self.collection_name}")
if vector is None:
vector = existing[0].get("vectors")
if vector is None:
raise ValueError(f"Existing record {vector_id} has no vector data")
if payload is None:
payload = existing[0].get("metadata")
text = ""
if payload:
text = (payload.get("text_lemmatized") or payload.get("data", ""))[:65535]
@@ -286,19 +297,19 @@ class MilvusDB(VectorStoreBase):
"""
return self.client.get_collection_stats(collection_name=self.collection_name)
def list(self, filters: dict = None, limit: int = 100) -> list:
def list(self, filters: dict = None, top_k: int = 100) -> list:
"""
List all vectors in a collection.
Args:
filters (Dict, optional): Filters to apply to the list.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
"""
query_filter = self._create_filter(filters) if filters else None
result = self.client.query(collection_name=self.collection_name, filter=query_filter, limit=limit)
result = self.client.query(collection_name=self.collection_name, filter=query_filter, limit=top_k)
memories = []
for data in result:
obj = OutputData(id=data.get("id"), score=None, payload=data.get("metadata"))
mem0/vector_stores/mongodb.py
+24 -24
View File
@@ -26,7 +26,7 @@ class OutputData(BaseModel):
class MongoDB(VectorStoreBase):
VECTOR_TYPE = "knnVector"
VECTOR_TYPE = "vector"
SIMILARITY_METRIC = "cosine"
def __init__(self, db_name: str, collection_name: str, embedding_model_dims: int, mongo_uri: str):
@@ -51,7 +51,7 @@ class MongoDB(VectorStoreBase):
"""Create new collection with vector search index."""
try:
database = self.client[self.db_name]
collection_names = database.list_collection_names()
collection_names = database.list_collection_names(authorizedCollections=True)
if self.collection_name not in collection_names:
logger.info(f"Collection '{self.collection_name}' does not exist. Creating it now.")
collection = database[self.collection_name]
@@ -69,17 +69,16 @@ class MongoDB(VectorStoreBase):
else:
search_index_model = SearchIndexModel(
name=self.index_name,
type="vectorSearch",
definition={
"mappings": {
"dynamic": False,
"fields": {
"embedding": {
"type": self.VECTOR_TYPE,
"dimensions": self.embedding_model_dims,
"similarity": self.SIMILARITY_METRIC,
}
},
}
"fields": [
{
"type": self.VECTOR_TYPE,
"path": "embedding",
"numDimensions": self.embedding_model_dims,
"similarity": self.SIMILARITY_METRIC,
}
]
},
)
collection.create_search_index(search_index_model)
@@ -149,14 +148,14 @@ class MongoDB(VectorStoreBase):
except PyMongoError as e:
logger.error(f"Error inserting data: {e}")
def search(self, query: str, vectors: List[float], limit=5, filters: Optional[Dict] = None) -> List[OutputData]:
def search(self, query: str, vectors: List[float], top_k=5, filters: Optional[Dict] = None) -> List[OutputData]:
"""
Search for similar vectors using the vector search index.
Args:
query (str): Query string
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search.
Returns:
@@ -175,8 +174,8 @@ class MongoDB(VectorStoreBase):
{
"$vectorSearch": {
"index": self.index_name,
"limit": limit,
"numCandidates": limit,
"limit": top_k,
"numCandidates": min(top_k * 20, 10000),
"queryVector": vectors,
"path": "embedding",
}
@@ -204,13 +203,13 @@ class MongoDB(VectorStoreBase):
output = [OutputData(id=str(doc["_id"]), score=doc.get("score"), payload=doc.get("payload")) for doc in results]
return output
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""
Perform keyword-based search using MongoDB Atlas Search.
Args:
query (str): The text query to search for.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search.
Returns:
@@ -242,7 +241,7 @@ class MongoDB(VectorStoreBase):
if filter_conditions:
pipeline.insert(1, {"$match": {"$and": filter_conditions}})
pipeline.append({"$limit": limit})
pipeline.append({"$limit": top_k})
results = list(collection.aggregate(pipeline))
logger.info(f"Keyword search completed. Found {len(results)} documents.")
@@ -285,7 +284,8 @@ class MongoDB(VectorStoreBase):
if vector is not None:
update_fields["embedding"] = vector
if payload is not None:
update_fields["payload"] = payload
for key, value in payload.items():
update_fields[f"payload.{key}"] = value
if update_fields:
try:
@@ -327,7 +327,7 @@ class MongoDB(VectorStoreBase):
List[str]: List of collection names.
"""
try:
collections = self.db.list_collection_names()
collections = self.db.list_collection_names(authorizedCollections=True)
logger.info(f"Listing collections in database '{self.db_name}': {collections}")
return collections
except PyMongoError as e:
@@ -358,13 +358,13 @@ class MongoDB(VectorStoreBase):
logger.error(f"Error getting collection info: {e}")
return {}
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List vectors in the collection.
Args:
filters (Dict, optional): Filters to apply to the list.
limit (int, optional): Number of vectors to return.
top_k (int, optional): Number of vectors to return.
Returns:
List[OutputData]: List of vectors.
@@ -379,7 +379,7 @@ class MongoDB(VectorStoreBase):
if filter_conditions:
query = {"$and": filter_conditions}
cursor = self.collection.find(query).limit(limit)
cursor = self.collection.find(query).limit(top_k)
results = [OutputData(id=str(doc["_id"]), score=None, payload=doc.get("payload")) for doc in cursor]
logger.info(f"Retrieved {len(results)} documents from collection '{self.collection_name}'.")
return results
mem0/vector_stores/neptune_analytics.py
+6 -6
View File
@@ -133,7 +133,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors using embedding similarity.
@@ -144,7 +144,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
Args:
query (str): Search query text (unused in vector search).
vectors (List[float]): Query embedding vector.
limit (int, optional): Maximum number of results to return. Defaults to 5.
top_k (int, optional): Maximum number of results to return. Defaults to 5.
filters (Optional[Dict]): Optional filters to apply to search results.
Returns:
@@ -159,7 +159,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
query_string = f"""
CALL neptune.algo.vectors.topKByEmbeddingWithFiltering({{
topK: {limit},
topK: {top_k},
embedding: {vectors}
{filter_clause}
}}
@@ -309,7 +309,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
pass
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List all vectors in the collection with optional filtering.
@@ -317,7 +317,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
Args:
filters (Optional[Dict]): Optional filters to apply based on metadata.
limit (int, optional): Maximum number of vectors to return. Defaults to 100.
top_k (int, optional): Maximum number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors with their metadata.
@@ -325,7 +325,7 @@ class NeptuneAnalyticsVector(VectorStoreBase):
where_clause = self._get_where_clause(filters) if filters else ""
para = {
"limit": limit,
"limit": top_k,
}
query_string = f"""
MATCH (n :{self.collection_name})
mem0/vector_stores/opensearch.py
+54 -24
View File
@@ -113,6 +113,25 @@ class OpenSearchDB(VectorStoreBase):
if payloads is None:
payloads = [{} for _ in range(len(vectors))]
for idx, vec in enumerate(vectors):
if vec is None:
raise ValueError(
f"Vector at index {idx} is null. "
f"This usually means the embedding model failed to generate an embedding. "
f"Check that your embedding model is configured correctly and returning valid vectors."
)
if len(vec) == 0:
raise ValueError(
f"Vector at index {idx} is empty. "
f"Expected a vector of dimension {self.embedding_model_dims}, got an empty vector."
)
if len(vec) != self.embedding_model_dims:
raise ValueError(
f"Vector at index {idx} has dimension {len(vec)}, "
f"but the index '{self.collection_name}' expects dimension {self.embedding_model_dims}. "
f"Ensure your embedding model's output dimensions match the vector store configuration."
)
results = []
for i, (vec, id_) in enumerate(zip(vectors, ids)):
body = {
@@ -124,20 +143,22 @@ class OpenSearchDB(VectorStoreBase):
self.client.index(index=self.collection_name, body=body)
# Force refresh to make documents immediately searchable for tests
self.client.indices.refresh(index=self.collection_name)
results.append(OutputData(
id=id_,
score=1.0, # No score for inserts
payload=payloads[i]
))
results.append(
OutputData(
id=id_,
score=1.0, # No score for inserts
payload=payloads[i],
)
)
except Exception as e:
logger.error(f"Error inserting vector {id_}: {e}")
logger.error(f"Error inserting vector {id_}: {e}", exc_info=True)
raise
return results
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""Search for similar vectors using OpenSearch k-NN search with optional filters."""
@@ -146,13 +167,13 @@ class OpenSearchDB(VectorStoreBase):
"knn": {
"vector_field": {
"vector": vectors,
"k": limit * 2,
"k": top_k * 2,
}
}
}
# Start building the full query
query_body = {"size": limit * 2, "query": None}
query_body = {"size": top_k * 2, "query": None}
# Prepare filter conditions if applicable
filter_clauses = []
@@ -175,19 +196,19 @@ class OpenSearchDB(VectorStoreBase):
hits = response["hits"]["hits"]
results = [
OutputData(id=hit["_source"].get("id"), score=hit["_score"], payload=hit["_source"].get("payload", {}))
for hit in hits[:limit] # Ensure we don't exceed limit
for hit in hits[:top_k] # Ensure we don't exceed top_k
]
return results
except Exception as e:
logger.error(f"Error during search: {e}")
logger.error(f"Error during search: {e}", exc_info=True)
return []
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""Search for memories using BM25 keyword matching.
Args:
query (str): The text query to search for.
limit (int): Maximum number of results to return. Defaults to 5.
top_k (int): Maximum number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search.
Returns:
@@ -216,7 +237,7 @@ class OpenSearchDB(VectorStoreBase):
bool_query["filter"] = filter_clauses
query_body = {
"size": limit,
"size": top_k,
"query": {"bool": bool_query},
}
@@ -226,7 +247,7 @@ class OpenSearchDB(VectorStoreBase):
hits = response["hits"]["hits"]
results = [
OutputData(id=hit["_source"].get("id"), score=hit["_score"], payload=hit["_source"].get("payload", {}))
for hit in hits[:limit]
for hit in hits[:top_k]
]
return results
except Exception as e:
@@ -251,6 +272,15 @@ class OpenSearchDB(VectorStoreBase):
def update(self, vector_id: str, vector: Optional[List[float]] = None, payload: Optional[Dict] = None) -> None:
"""Update a vector and its payload using the custom 'id' field."""
if vector is not None:
if len(vector) == 0:
raise ValueError("Cannot update with an empty vector.")
if len(vector) != self.embedding_model_dims:
raise ValueError(
f"Update vector has dimension {len(vector)}, "
f"but the index '{self.collection_name}' expects dimension {self.embedding_model_dims}. "
f"Ensure your embedding model's output dimensions match the vector store configuration."
)
# First, find the document by custom ID
search_query = {"query": {"term": {"id": vector_id}}}
@@ -273,8 +303,9 @@ class OpenSearchDB(VectorStoreBase):
if doc:
try:
response = self.client.update(index=self.collection_name, id=opensearch_id, body={"doc": doc})
except Exception:
pass
except Exception as e:
logger.error(f"Error updating vector {vector_id}: {e}", exc_info=True)
raise
def get(self, vector_id: str) -> Optional[OutputData]:
"""Retrieve a vector by ID."""
@@ -289,7 +320,7 @@ class OpenSearchDB(VectorStoreBase):
return OutputData(id=hits[0]["_source"].get("id"), score=1.0, payload=hits[0]["_source"].get("payload", {}))
except Exception as e:
logger.error(f"Error retrieving vector {vector_id}: {str(e)}")
logger.error(f"Error retrieving vector {vector_id}: {str(e)}", exc_info=True)
return None
def list_cols(self) -> List[str]:
@@ -304,7 +335,7 @@ class OpenSearchDB(VectorStoreBase):
"""Get information about a collection (index)."""
return self.client.indices.get(index=name)
def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: Optional[int] = None) -> List[OutputData]:
try:
"""List all memories with optional filters."""
query: Dict = {"query": {"match_all": {}}}
@@ -319,8 +350,8 @@ class OpenSearchDB(VectorStoreBase):
if filter_clauses:
query["query"] = {"bool": {"filter": filter_clauses}}
if limit:
query["size"] = limit
if top_k:
query["size"] = top_k
response = self.client.search(index=self.collection_name, body=query)
hits = response["hits"]["hits"]
@@ -332,9 +363,8 @@ class OpenSearchDB(VectorStoreBase):
]
return [results] # VectorStore expects tuple/list format
except Exception as e:
logger.error(f"Error listing vectors: {e}")
logger.error(f"Error listing vectors: {e}", exc_info=True)
return []
def reset(self):
"""Reset the index by deleting and recreating it."""
mem0/vector_stores/pgvector.py
+9 -9
View File
@@ -210,7 +210,7 @@ class PGVector(VectorStoreBase):
self,
query: str,
vectors: list[float],
limit: Optional[int] = 5,
top_k: Optional[int] = 5,
filters: Optional[dict] = None,
) -> List[OutputData]:
"""
@@ -219,7 +219,7 @@ class PGVector(VectorStoreBase):
Args:
query (str): Query.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -244,19 +244,19 @@ class PGVector(VectorStoreBase):
ORDER BY distance
LIMIT %s
""",
(vectors, *filter_params, limit),
(vectors, *filter_params, top_k),
)
results = cur.fetchall()
return [OutputData(id=str(r[0]), score=float(r[1]), payload=r[2]) for r in results]
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""
Search using PostgreSQL full-text search on lemmatized text.
Args:
query (str): The search query text.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -285,7 +285,7 @@ class PGVector(VectorStoreBase):
ORDER BY score DESC
LIMIT %s
""",
(query, query, *filter_params, limit),
(query, query, *filter_params, top_k),
)
results = cur.fetchall()
@@ -401,14 +401,14 @@ class PGVector(VectorStoreBase):
def list(
self,
filters: Optional[dict] = None,
limit: Optional[int] = 100
top_k: Optional[int] = 100
) -> List[OutputData]:
"""
List all vectors in a collection.
Args:
filters (Dict, optional): Filters to apply to the list.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors.
@@ -431,7 +431,7 @@ class PGVector(VectorStoreBase):
"""
with self._get_cursor() as cur:
cur.execute(query, (*filter_params, limit))
cur.execute(query, (*filter_params, top_k))
results = cur.fetchall()
return [[OutputData(id=str(r[0]), score=None, payload=r[2]) for r in results]]
mem0/vector_stores/pinecone.py
+9 -9
View File
@@ -204,7 +204,7 @@ class PineconeDB(VectorStoreBase):
return pinecone_filter
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
@@ -212,7 +212,7 @@ class PineconeDB(VectorStoreBase):
Args:
query (str): Query.
vectors (list): List of vectors to search.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -222,7 +222,7 @@ class PineconeDB(VectorStoreBase):
query_params = {
"vector": vectors,
"top_k": limit,
"top_k": top_k,
"include_metadata": True,
"include_values": False,
}
@@ -241,13 +241,13 @@ class PineconeDB(VectorStoreBase):
results = self._parse_output(response.matches)
return results
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""
Search using BM25 sparse vectors for keyword-based retrieval.
Args:
query (str): The search query text.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -263,7 +263,7 @@ class PineconeDB(VectorStoreBase):
query_params = {
"sparse_vector": sparse_vector,
"top_k": limit,
"top_k": top_k,
"include_metadata": True,
"include_values": False,
}
@@ -356,13 +356,13 @@ class PineconeDB(VectorStoreBase):
"""
return self.client.describe_index(self.collection_name)
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List vectors in an index with optional filtering.
Args:
filters (dict, optional): Filters to apply to the list. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
dict: List of vectors with their metadata.
@@ -376,7 +376,7 @@ class PineconeDB(VectorStoreBase):
query_params = {
"vector": zero_vector,
"top_k": limit,
"top_k": top_k,
"include_metadata": True,
"include_values": True,
}
mem0/vector_stores/qdrant.py
+196 -32
View File
@@ -1,15 +1,20 @@
import logging
import os
import shutil
import re
from typing import Optional
from qdrant_client import QdrantClient, models
from qdrant_client.models import (
DatetimeRange,
Distance,
FieldCondition,
Filter,
MatchAny,
MatchExcept,
MatchText,
MatchValue,
PointIdsList,
PointStruct,
PointVectors,
Range,
SparseVector,
SparseVectorParams,
@@ -46,7 +51,8 @@ class Qdrant(VectorStoreBase):
path (str, optional): Path for local Qdrant database. Defaults to None.
url (str, optional): Full URL for Qdrant server. Defaults to None.
api_key (str, optional): API key for Qdrant server. Defaults to None.
on_disk (bool, optional): Enables persistent storage. Defaults to False.
on_disk (bool, optional): Enables persistent storage. Vectors are stored on disk (True) or in memory (False).
Does not delete the local database path. Defaults to False.
"""
if client:
self.client = client
@@ -64,9 +70,6 @@ class Qdrant(VectorStoreBase):
if not params:
params["path"] = path
self.is_local = True
if not on_disk:
if os.path.exists(path) and os.path.isdir(path):
shutil.rmtree(path)
else:
self.is_local = False
@@ -186,35 +189,181 @@ class Qdrant(VectorStoreBase):
self.client.upsert(collection_name=self.collection_name, points=points)
def _create_filter(self, filters: dict) -> Filter:
# ISO 8601 datetime pattern for detecting datetime strings in range filters
_ISO_DATETIME_RE = re.compile(
r"^\d{4}-\d{2}-\d{2}" # date part
r"([T ]\d{2}:\d{2}(:\d{2})?" # optional time part
r"(\.\d+)?" # optional fractional seconds
r"(Z|[+-]\d{2}:?\d{2})?" # optional timezone
r")?$"
)
@staticmethod
def _is_datetime_range(range_kwargs: dict) -> bool:
"""Check if all values in range kwargs are ISO datetime strings."""
return all(
isinstance(v, str) and Qdrant._ISO_DATETIME_RE.match(v)
for v in range_kwargs.values()
)
def _build_field_condition(self, key: str, value) -> Optional[FieldCondition]:
"""
Build a single FieldCondition from a key-value filter pair.
Supports the enhanced filter syntax documented at
https://docs.mem0.ai/open-source/features/metadata-filtering
Args:
key (str): The payload field name.
value: A scalar for simple equality, or a dict with one operator key.
Returns:
Optional[FieldCondition]: The Qdrant field condition, or None if the
value is the wildcard '*' (match any / field exists — skip filter).
"""
if not isinstance(value, dict):
if value == "*":
# Wildcard: match any value. Qdrant has no direct "field exists"
# condition via FieldCondition, so we skip this filter (match all).
return None
if isinstance(value, list):
# List shorthand: {"field": ["a", "b"]} treated as in-operator.
return FieldCondition(key=key, match=MatchAny(any=value))
# Simple equality: {"field": "value"}
return FieldCondition(key=key, match=MatchValue(value=value))
ops = set(value.keys())
range_ops = {"gt", "gte", "lt", "lte"}
non_range_ops = ops - range_ops
if ops & range_ops:
if non_range_ops:
raise ValueError(
f"Cannot mix range operators ({ops & range_ops}) with "
f"non-range operators ({non_range_ops}) for field '{key}'. "
f"Use AND to combine them as separate conditions."
)
range_kwargs = {op: value[op] for op in range_ops if op in value}
if self._is_datetime_range(range_kwargs):
try:
return FieldCondition(key=key, range=DatetimeRange(**range_kwargs))
except (ValueError, TypeError) as e:
raise ValueError(
f"Invalid datetime value in range filter for field '{key}': {e}"
) from e
return FieldCondition(key=key, range=Range(**range_kwargs))
elif "eq" in value:
return FieldCondition(key=key, match=MatchValue(value=value["eq"]))
elif "ne" in value:
return FieldCondition(key=key, match=MatchExcept(**{"except": [value["ne"]]}))
elif "in" in value:
return FieldCondition(key=key, match=MatchAny(any=value["in"]))
elif "nin" in value:
return FieldCondition(key=key, match=MatchExcept(**{"except": value["nin"]}))
elif "contains" in value or "icontains" in value:
# MatchText: with a full-text index, tokenized matching (all words must appear).
# Without a full-text index, exact substring match.
op = "icontains" if "icontains" in value else "contains"
text = value[op]
if op == "icontains":
logger.debug(
"icontains on field '%s': Qdrant MatchText case sensitivity depends on "
"full-text index configuration. Without a full-text index this behaves "
"as a case-sensitive substring match (same as 'contains').",
key,
)
return FieldCondition(key=key, match=MatchText(text=text))
else:
supported = {"eq", "ne", "gt", "gte", "lt", "lte", "in", "nin", "contains", "icontains"}
raise ValueError(
f"Unsupported filter operator(s) for field '{key}': {ops}. "
f"Supported operators: {supported}"
)
def _create_filter(self, filters: dict) -> Optional[Filter]:
"""
Create a Filter object from the provided filters.
Supports the enhanced filter syntax with comparison operators (eq, ne,
gt, gte, lt, lte), list operators (in, nin), string operators (contains,
icontains), and logical operators (AND, OR, NOT).
Args:
filters (dict): Filters to apply.
Returns:
Filter: The created Filter object.
Filter: The created Filter object, or None if filters is empty.
"""
if not filters:
return None
conditions = []
# Normalize $or/$not/$and → OR/NOT/AND and deduplicate.
# Memory._process_metadata_filters() renames OR→$or and NOT→$not,
# but effective_filters retains the original OR/NOT keys from
# deepcopy(input_filters). Without dedup the same sub-conditions
# would be evaluated twice.
key_map = {"$or": "OR", "$not": "NOT", "$and": "AND"}
normalized = {}
for key, value in filters.items():
if isinstance(value, dict) and "gte" in value and "lte" in value:
conditions.append(FieldCondition(key=key, range=Range(gte=value["gte"], lte=value["lte"])))
else:
conditions.append(FieldCondition(key=key, match=MatchValue(value=value)))
return Filter(must=conditions) if conditions else None
norm_key = key_map.get(key, key)
if norm_key not in normalized:
normalized[norm_key] = value
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
must = []
should = []
must_not = []
for key, value in normalized.items():
if key in ("AND", "OR", "NOT"):
if not isinstance(value, list):
raise ValueError(
f"{key} filter value must be a list of filter dicts, "
f"got {type(value).__name__}"
)
for i, item in enumerate(value):
if not isinstance(item, dict):
raise ValueError(
f"{key} filter list item at index {i} must be a dict, "
f"got {type(item).__name__}: {item!r}"
)
if key == "AND":
for sub in value:
built = self._create_filter(sub)
if built:
must.append(built)
elif key == "OR":
for sub in value:
built = self._create_filter(sub)
if built:
should.append(built)
elif key == "NOT":
for sub in value:
built = self._create_filter(sub)
if built:
must_not.append(built)
else:
condition = self._build_field_condition(key, value)
if condition is not None:
must.append(condition)
if not any([must, should, must_not]):
return None
return Filter(
must=must or None,
should=should or None,
must_not=must_not or None,
)
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None) -> list:
"""
Search for similar vectors.
Args:
query (str): Query.
vectors (list): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -225,15 +374,15 @@ class Qdrant(VectorStoreBase):
collection_name=self.collection_name,
query=vectors,
query_filter=query_filter,
limit=limit,
limit=top_k,
)
return hits.points
def search_batch(self, queries: list, vectors_list: list, limit: int = 1, filters: dict = None):
def search_batch(self, queries: list, vectors_list: list, top_k: int = 1, filters: dict = None):
"""Batch search using Qdrant's query_batch_points for efficiency."""
query_filter = self._create_filter(filters) if filters else None
requests = [
models.QueryRequest(query=vec, filter=query_filter, limit=limit)
models.QueryRequest(query=vec, filter=query_filter, limit=top_k)
for vec in vectors_list
]
try:
@@ -244,15 +393,15 @@ class Qdrant(VectorStoreBase):
return [r.points for r in results]
except Exception as e:
logger.warning(f"Batch search failed, falling back to sequential: {e}")
return [self.search(q, v, limit=limit, filters=filters) for q, v in zip(queries, vectors_list)]
return [self.search(q, v, top_k=top_k, filters=filters) for q, v in zip(queries, vectors_list)]
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""
Search using BM25 sparse vectors for keyword-based retrieval.
Args:
query (str): The search query text.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -269,7 +418,7 @@ class Qdrant(VectorStoreBase):
query=sparse_query,
using="bm25",
query_filter=query_filter,
limit=limit,
limit=top_k,
)
return hits.points
except Exception as e:
@@ -299,17 +448,32 @@ class Qdrant(VectorStoreBase):
vector (list, optional): Updated vector. Defaults to None.
payload (dict, optional): Updated payload. Defaults to None.
"""
# Build named vectors with BM25 if payload has text
named_vectors = {"": vector} if vector else vector
if payload and vector:
if vector is not None and payload is not None:
# Full update: attach BM25 sparse vector alongside dense vector
named_vectors = {"": vector}
text_for_bm25 = payload.get("text_lemmatized") or payload.get("data", "")
if text_for_bm25:
sparse = self._encode_bm25(text_for_bm25)
if sparse is not None:
named_vectors["bm25"] = sparse
point = PointStruct(id=vector_id, vector=named_vectors, payload=payload)
self.client.upsert(collection_name=self.collection_name, points=[point])
point = PointStruct(id=vector_id, vector=named_vectors, payload=payload)
self.client.upsert(collection_name=self.collection_name, points=[point])
else:
# Partial update: use Qdrant's dedicated endpoints.
# Note: BM25 sparse vector cannot be refreshed via set_payload alone;
# payload-only updates will leave any existing BM25 vector stale. In
# practice v3 re-embeds on memory text change, so this is acceptable.
if payload is not None:
self.client.set_payload(
collection_name=self.collection_name,
payload=payload,
points=[vector_id],
)
if vector is not None:
self.client.update_vectors(
collection_name=self.collection_name,
points=[PointVectors(id=vector_id, vector=vector)],
)
def get(self, vector_id: int) -> dict:
"""
@@ -346,13 +510,13 @@ class Qdrant(VectorStoreBase):
"""
return self.client.get_collection(collection_name=self.collection_name)
def list(self, filters: dict = None, limit: int = 100) -> list:
def list(self, filters: dict = None, top_k: int = 100) -> list:
"""
List all vectors in a collection.
Args:
filters (dict, optional): Filters to apply to the list. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
list: List of vectors.
@@ -361,7 +525,7 @@ class Qdrant(VectorStoreBase):
result = self.client.scroll(
collection_name=self.collection_name,
scroll_filter=query_filter,
limit=limit,
limit=top_k,
with_payload=True,
with_vectors=False,
)
mem0/vector_stores/redis.py
+20 -18
View File
@@ -1,10 +1,9 @@
import json
import logging
from datetime import datetime
from datetime import datetime, timezone
from functools import reduce
import numpy as np
import pytz
import redis
from redis.commands.search.query import Query
from redisvl.index import SearchIndex
@@ -142,7 +141,7 @@ class RedisDB(VectorStoreBase):
data.append(entry)
self.index.load(data, id_field="memory_id")
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None):
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None):
conditions = [Tag(key) == value for key, value in filters.items() if value is not None]
filter = reduce(lambda x, y: x & y, conditions)
@@ -151,7 +150,7 @@ class RedisDB(VectorStoreBase):
vector_field_name="embedding",
return_fields=["memory_id", "hash", "agent_id", "run_id", "user_id", "memory", "metadata", "created_at"],
filter_expression=filter,
num_results=limit,
num_results=top_k,
)
results = self.index.query(v)
@@ -159,17 +158,17 @@ class RedisDB(VectorStoreBase):
return [
MemoryResult(
id=result["memory_id"],
score=result["vector_distance"],
score=float(result["vector_distance"]),
payload={
"hash": result["hash"],
"data": result["memory"],
"created_at": datetime.fromtimestamp(
int(result["created_at"]), tz=pytz.timezone("US/Pacific")
int(result["created_at"]), tz=timezone.utc
).isoformat(timespec="microseconds"),
**(
{
"updated_at": datetime.fromtimestamp(
int(result["updated_at"]), tz=pytz.timezone("US/Pacific")
int(result["updated_at"]), tz=timezone.utc
).isoformat(timespec="microseconds")
}
if "updated_at" in result
@@ -182,13 +181,13 @@ class RedisDB(VectorStoreBase):
for result in results
]
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""
Search for memories using BM25 keyword search on the memory field.
Args:
query (str): Search query text.
limit (int): Maximum number of results. Defaults to 5.
top_k (int): Maximum number of results. Defaults to 5.
filters (dict, optional): Filters to apply (user_id, agent_id, run_id).
Returns:
@@ -205,7 +204,7 @@ class RedisDB(VectorStoreBase):
text_field_name="memory",
return_fields=["memory_id", "hash", "agent_id", "run_id", "user_id", "memory", "metadata", "created_at"],
filter_expression=filter_expression,
num_results=limit,
num_results=top_k,
)
results = self.index.query(t)
@@ -246,9 +245,12 @@ class RedisDB(VectorStoreBase):
"memory": payload["data"],
"created_at": int(datetime.fromisoformat(payload["created_at"]).timestamp()),
"updated_at": int(datetime.fromisoformat(payload["updated_at"]).timestamp()),
"embedding": np.array(vector, dtype=np.float32).tobytes(),
}
# Only update embedding if vector is provided
if vector is not None:
data["embedding"] = np.array(vector, dtype=np.float32).tobytes()
for field in ["agent_id", "run_id", "user_id"]:
if field in payload:
data[field] = payload[field]
@@ -261,13 +263,13 @@ class RedisDB(VectorStoreBase):
payload = {
"hash": result["hash"],
"data": result["memory"],
"created_at": datetime.fromtimestamp(int(result["created_at"]), tz=pytz.timezone("US/Pacific")).isoformat(
"created_at": datetime.fromtimestamp(int(result["created_at"]), tz=timezone.utc).isoformat(
timespec="microseconds"
),
**(
{
"updated_at": datetime.fromtimestamp(
int(result["updated_at"]), tz=pytz.timezone("US/Pacific")
int(result["updated_at"]), tz=timezone.utc
).isoformat(timespec="microseconds")
}
if "updated_at" in result
@@ -306,15 +308,15 @@ class RedisDB(VectorStoreBase):
# Recreate the index with the same parameters
self.create_col(collection_name, self.embedding_model_dims)
def list(self, filters: dict = None, limit: int = None) -> list:
def list(self, filters: dict = None, top_k: int = None) -> list:
"""
List all recent created memories from the vector store.
"""
conditions = [Tag(key) == value for key, value in filters.items() if value is not None]
filter = reduce(lambda x, y: x & y, conditions)
query = Query(str(filter)).sort_by("created_at", asc=False)
if limit is not None:
query = Query(str(filter)).sort_by("created_at", asc=False).paging(0, limit)
if top_k is not None:
query = Query(str(filter)).sort_by("created_at", asc=False).paging(0, top_k)
results = self.index.search(query)
return [
@@ -325,12 +327,12 @@ class RedisDB(VectorStoreBase):
"hash": result["hash"],
"data": result["memory"],
"created_at": datetime.fromtimestamp(
int(result["created_at"]), tz=pytz.timezone("US/Pacific")
int(result["created_at"]), tz=timezone.utc
).isoformat(timespec="microseconds"),
**(
{
"updated_at": datetime.fromtimestamp(
int(result["updated_at"]), tz=pytz.timezone("US/Pacific")
int(result["updated_at"]), tz=timezone.utc
).isoformat(timespec="microseconds")
}
if result.__dict__.get("updated_at")
mem0/vector_stores/s3_vectors.py
+5 -5
View File
@@ -99,12 +99,12 @@ class S3Vectors(VectorStoreBase):
vectors=vectors_to_put,
)
def search(self, query, vectors, limit=5, filters=None):
def search(self, query, vectors, top_k=5, filters=None):
params = {
"vectorBucketName": self.vector_bucket_name,
"indexName": self.collection_name,
"queryVector": {"float32": vectors},
"topK": limit,
"topK": top_k,
"returnMetadata": True,
"returnDistance": True,
}
@@ -149,7 +149,7 @@ class S3Vectors(VectorStoreBase):
response = self.client.get_index(vectorBucketName=self.vector_bucket_name, indexName=self.collection_name)
return response.get("index", {})
def list(self, filters=None, limit=None):
def list(self, filters=None, top_k=None):
# Note: list_vectors does not support metadata filtering.
if filters:
logger.warning("S3 Vectors `list` does not support metadata filtering. Ignoring filters.")
@@ -160,8 +160,8 @@ class S3Vectors(VectorStoreBase):
"returnData": False,
"returnMetadata": True,
}
if limit:
params["maxResults"] = limit
if top_k:
params["maxResults"] = top_k
paginator = self.client.get_paginator("list_vectors")
pages = paginator.paginate(**params)
mem0/vector_stores/supabase.py
+6 -6
View File
@@ -116,7 +116,7 @@ class Supabase(VectorStoreBase):
self.collection.upsert(records)
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
@@ -124,7 +124,7 @@ class Supabase(VectorStoreBase):
Args:
query (str): Query.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search. Defaults to None.
Returns:
@@ -132,7 +132,7 @@ class Supabase(VectorStoreBase):
"""
filters = self._preprocess_filters(filters)
results = self.collection.query(
data=vectors, limit=limit, filters=filters, include_metadata=True, include_value=True
data=vectors, limit=top_k, filters=filters, include_metadata=True, include_value=True
)
return [OutputData(id=str(result[0]), score=float(result[1]), payload=result[2]) for result in results]
@@ -209,13 +209,13 @@ class Supabase(VectorStoreBase):
"index": {"method": info.index_method, "metric": info.distance_metric},
}
def list(self, filters: Optional[dict] = None, limit: int = 100) -> List[OutputData]:
def list(self, filters: Optional[dict] = None, top_k: int = 100) -> List[OutputData]:
"""
List vectors in the collection.
Args:
filters (Dict, optional): Filters to apply
limit (int, optional): Maximum number of results to return. Defaults to 100.
top_k (int, optional): Maximum number of results to return. Defaults to 100.
Returns:
List[OutputData]: List of vectors
@@ -223,7 +223,7 @@ class Supabase(VectorStoreBase):
filters = self._preprocess_filters(filters)
query = [0] * self.embedding_model_dims
ids = self.collection.query(
data=query, limit=limit, filters=filters, include_metadata=True, include_value=False
data=query, limit=top_k, filters=filters, include_metadata=True, include_value=False
)
ids = [id[0] for id in ids]
records = self.collection.fetch(ids=ids)
mem0/vector_stores/turbopuffer.py
+337
View File
@@ -0,0 +1,337 @@
import logging
import os
from typing import Any, Dict, List, Optional, Union
try:
from turbopuffer import Turbopuffer as TurbopufferClient
except ImportError:
raise ImportError(
"Turbopuffer requires extra dependencies. Install with `pip install turbopuffer`"
) from None
from pydantic import BaseModel
from mem0.vector_stores.base import VectorStoreBase
logger = logging.getLogger(__name__)
class OutputData(BaseModel):
id: Optional[str]
score: Optional[float]
payload: Optional[Dict]
class TurbopufferDB(VectorStoreBase):
def __init__(
self,
collection_name: str,
embedding_model_dims: int,
api_key: Optional[str] = None,
region: str = "gcp-us-central1",
distance_metric: str = "cosine_distance",
batch_size: int = 100,
extra_params: Optional[Dict[str, Any]] = None,
):
"""
Initialize the Turbopuffer vector store.
Args:
collection_name (str): Name of the namespace/collection.
embedding_model_dims (int): Dimensions of the embedding model.
api_key (str, optional): API key for Turbopuffer. Defaults to None.
region (str, optional): Turbopuffer region. Defaults to "gcp-us-central1".
distance_metric (str, optional): Distance metric for vector similarity.
Options: "cosine_distance" or "euclidean_squared". Defaults to "cosine_distance".
batch_size (int, optional): Batch size for operations. Defaults to 100.
extra_params (Dict, optional): Additional parameters for Turbopuffer client. Defaults to None.
"""
api_key = api_key or os.environ.get("TURBOPUFFER_API_KEY")
if not api_key:
raise ValueError(
"Turbopuffer API key must be provided either as a parameter or via TURBOPUFFER_API_KEY environment variable"
)
params = extra_params or {}
params["region"] = region
self.client = TurbopufferClient(api_key=api_key, **params)
self.collection_name = collection_name
self.embedding_model_dims = embedding_model_dims
self.distance_metric = distance_metric
self.batch_size = batch_size
self.namespace = self.client.namespace(self.collection_name)
def create_col(self, name=None, vector_size=None, distance=None):
"""
Create a new namespace in Turbopuffer.
Namespaces are created implicitly on first upsert, so this is a no-op.
"""
pass
def insert(
self,
vectors: List[List[float]],
payloads: Optional[List[Dict]] = None,
ids: Optional[List[Union[str, int]]] = None,
):
"""
Insert vectors into the namespace.
Args:
vectors (list): List of vectors to insert.
payloads (list, optional): List of payloads corresponding to vectors. Defaults to None.
ids (list, optional): List of IDs corresponding to vectors. Defaults to None.
"""
logger.info(f"Inserting {len(vectors)} vectors into namespace {self.collection_name}")
if ids is None:
ids = [str(i) for i in range(len(vectors))]
for i in range(0, len(vectors), self.batch_size):
batch_end = i + self.batch_size
rows = []
for j in range(i, min(batch_end, len(vectors))):
row = {}
if payloads and payloads[j]:
row.update(payloads[j])
row["id"] = str(ids[j])
row["vector"] = vectors[j]
rows.append(row)
self.namespace.write(
upsert_rows=rows,
distance_metric=self.distance_metric,
)
def _parse_output(self, rows) -> List[OutputData]:
"""
Parse the output data from Turbopuffer query results.
Args:
rows: List of Row objects from Turbopuffer query.
Returns:
List[OutputData]: Parsed output data.
"""
results = []
for row in rows:
row_dict = row.model_dump()
row_id = str(row_dict.pop("id"))
dist = row_dict.pop("$dist", None)
row_dict.pop("vector", None)
score = 1 - dist if dist is not None else None
results.append(OutputData(
id=row_id,
score=score,
payload=row_dict,
))
return results
def _convert_filters(self, filters: Optional[Dict]):
"""
Convert mem0 filters to Turbopuffer filter format.
Turbopuffer filters use tuple format: ("And", (("field", "Op", value), ...))
"""
if not filters:
return None
conditions = []
for key, value in filters.items():
if isinstance(value, dict):
if "gte" in value:
conditions.append((key, "Gte", value["gte"]))
if "lte" in value:
conditions.append((key, "Lte", value["lte"]))
else:
conditions.append((key, "Eq", value))
if not conditions:
return None
if len(conditions) == 1:
return conditions[0]
return ("And", tuple(conditions))
def search(
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
Args:
query (str): Query text (unused in vector search, kept for interface consistency).
vectors (list): Query vector to search with.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
Returns:
list: Search results.
"""
query_params = {
"rank_by": ("vector", "ANN", vectors),
"top_k": top_k,
"include_attributes": True,
}
tpuf_filters = self._convert_filters(filters)
if tpuf_filters is not None:
query_params["filters"] = tpuf_filters
response = self.namespace.query(**query_params)
return self._parse_output(response.rows or [])
def delete(self, vector_id: Union[str, int]):
"""
Delete a vector by ID.
Args:
vector_id (Union[str, int]): ID of the vector to delete.
"""
self.namespace.write(deletes=[str(vector_id)])
def update(
self,
vector_id: Union[str, int],
vector: Optional[List[float]] = None,
payload: Optional[Dict] = None,
):
"""
Update a vector and its payload.
Args:
vector_id (Union[str, int]): ID of the vector to update.
vector (list, optional): Updated vector. Defaults to None.
payload (dict, optional): Updated payload. Defaults to None.
"""
if vector is not None:
row = {}
if payload:
row.update(payload)
row["id"] = str(vector_id)
row["vector"] = vector
self.namespace.write(
upsert_rows=[row],
distance_metric=self.distance_metric,
)
elif payload is not None:
row = dict(payload)
row["id"] = str(vector_id)
self.namespace.write(patch_rows=[row])
def get(self, vector_id: Union[str, int]) -> Optional[OutputData]:
"""
Retrieve a vector by ID.
Args:
vector_id (Union[str, int]): ID of the vector to retrieve.
Returns:
OutputData: Retrieved vector data, or None if not found.
"""
try:
response = self.namespace.query(
top_k=1,
rank_by=("vector", "ANN", [0.0] * self.embedding_model_dims),
filters=("id", "Eq", str(vector_id)),
include_attributes=True,
)
rows = response.rows or []
if rows:
return self._parse_output(rows)[0]
return None
except Exception as e:
logger.error(f"Error retrieving vector {vector_id}: {e}")
return None
def list_cols(self) -> list:
"""
List all namespaces.
Returns:
list: List of namespace summaries.
"""
try:
result = []
for ns in self.client.namespaces():
result.append(ns)
return result
except Exception as e:
logger.error(f"Error listing namespaces: {e}")
return []
def delete_col(self):
"""Delete the entire namespace."""
try:
self.namespace.delete_all()
logger.info(f"Namespace {self.collection_name} deleted successfully")
except Exception as e:
logger.error(f"Error deleting namespace {self.collection_name}: {e}")
def col_info(self) -> Dict:
"""
Get information about the namespace.
Returns:
dict: Namespace metadata.
"""
try:
metadata = self.namespace.metadata()
return {
"name": self.collection_name,
"approx_row_count": metadata.approx_row_count,
"approx_logical_bytes": metadata.approx_logical_bytes,
"created_at": str(metadata.created_at),
"updated_at": str(metadata.updated_at),
}
except Exception:
return {"name": self.collection_name}
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> list:
"""
List vectors in the namespace with optional filtering.
Args:
filters (dict, optional): Filters to apply. Defaults to None.
top_k (int, optional): Number of vectors to return. Defaults to 100.
Returns:
list: Wrapped list of OutputData objects ([[results]]).
"""
query_params = {
"rank_by": ("vector", "ANN", [0.0] * self.embedding_model_dims),
"top_k": top_k,
"include_attributes": True,
}
tpuf_filters = self._convert_filters(filters)
if tpuf_filters is not None:
query_params["filters"] = tpuf_filters
try:
response = self.namespace.query(**query_params)
results = self._parse_output(response.rows or [])
except Exception as e:
logger.error(f"Error listing vectors: {e}")
results = []
return [results]
def count(self) -> int:
"""
Get approximate count of vectors in the namespace.
Returns:
int: Approximate number of vectors.
"""
try:
metadata = self.namespace.metadata()
return metadata.approx_row_count
except Exception:
return 0
def reset(self):
"""Reset the namespace by deleting all vectors."""
self.delete_col()
mem0/vector_stores/upstash_vector.py
+10 -10
View File
@@ -98,7 +98,7 @@ class UpstashVector(VectorStoreBase):
self,
query: str,
vectors: List[list],
limit: int = 5,
top_k: int = 5,
filters: Optional[Dict] = None,
) -> List[OutputData]:
"""
@@ -106,7 +106,7 @@ class UpstashVector(VectorStoreBase):
Args:
query (list): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search.
Returns:
@@ -120,7 +120,7 @@ class UpstashVector(VectorStoreBase):
if self.enable_embeddings:
response = self.client.query(
data=query,
top_k=limit,
top_k=top_k,
filter=filters_str or "",
include_metadata=True,
namespace=self.collection_name,
@@ -129,7 +129,7 @@ class UpstashVector(VectorStoreBase):
queries = [
{
"vector": v,
"top_k": limit,
"top_k": top_k,
"filter": filters_str or "",
"include_metadata": True,
"namespace": self.collection_name,
@@ -149,13 +149,13 @@ class UpstashVector(VectorStoreBase):
for res in response
]
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""
Perform keyword-based search using Upstash's BM25 sparse search.
Args:
query (str): The text query to search for.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Dict, optional): Filters to apply to the search.
Returns:
@@ -170,7 +170,7 @@ class UpstashVector(VectorStoreBase):
response = self.client.query(
data=query,
top_k=limit,
top_k=top_k,
filter=filters_str or "",
include_metadata=True,
namespace=self.collection_name,
@@ -244,12 +244,12 @@ class UpstashVector(VectorStoreBase):
return None
return OutputData(id=vector.id, score=None, payload=vector.metadata)
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[List[OutputData]]:
def list(self, filters: Optional[Dict] = None, top_k: int = 100) -> List[List[OutputData]]:
"""
List all memories.
Args:
filters (Dict, optional): Filters to apply to the search. Defaults to None.
limit (int, optional): Number of results to return. Defaults to 100.
top_k (int, optional): Number of results to return. Defaults to 100.
Returns:
List[OutputData]: Search results.
"""
@@ -272,7 +272,7 @@ class UpstashVector(VectorStoreBase):
)
with query:
while True:
if len(results) >= limit:
if len(results) >= top_k:
break
res = query.fetch_next(100)
if not res:
mem0/vector_stores/valkey.py
+12 -9
View File
@@ -410,14 +410,14 @@ class ValkeyDB(VectorStoreBase):
return memory_results
def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None, ef_runtime: int = None):
def search(self, query: str, vectors: list, top_k: int = 5, filters: dict = None, ef_runtime: int = None):
"""
Search for similar vectors in the index.
Args:
query (str): The search query.
vectors (list): The vector to search for.
limit (int, optional): Maximum number of results to return. Defaults to 5.
top_k (int, optional): Maximum number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None.
ef_runtime (int, optional): HNSW ef_runtime parameter for this query. Only used with HNSW index. Defaults to None.
@@ -429,10 +429,10 @@ class ValkeyDB(VectorStoreBase):
# Build the KNN part with optional EF_RUNTIME for HNSW
if self.index_type == "hnsw" and ef_runtime is not None:
knn_part = f"[KNN {limit} @embedding $vec_param EF_RUNTIME {ef_runtime} AS vector_score]"
knn_part = f"[KNN {top_k} @embedding $vec_param EF_RUNTIME {ef_runtime} AS vector_score]"
else:
# For FLAT indexes or when ef_runtime is None, use basic KNN
knn_part = f"[KNN {limit} @embedding $vec_param AS vector_score]"
knn_part = f"[KNN {top_k} @embedding $vec_param AS vector_score]"
# Build the complete query
q = self._build_search_query(knn_part, filters)
@@ -491,9 +491,12 @@ class ValkeyDB(VectorStoreBase):
"hash": payload.get("hash", f"hash_{vector_id}"), # Use a default hash if not provided
"memory": payload.get("data", f"data_{vector_id}"), # Use a default data if not provided
"created_at": int(datetime.fromisoformat(payload["created_at"]).timestamp()),
"embedding": np.array(vector, dtype=np.float32).tobytes(),
}
# Only update embedding if vector is provided
if vector is not None:
hash_data["embedding"] = np.array(vector, dtype=np.float32).tobytes()
# Add updated_at if available
if "updated_at" in payload:
hash_data["updated_at"] = int(datetime.fromisoformat(payload["updated_at"]).timestamp())
@@ -765,7 +768,7 @@ class ValkeyDB(VectorStoreBase):
return q
def list(self, filters: dict = None, limit: int = None) -> list:
def list(self, filters: dict = None, top_k: int = None) -> list:
"""
List all recent created memories from the vector store.
@@ -775,7 +778,7 @@ class ValkeyDB(VectorStoreBase):
Values are used as-is without validation - wildcards, special characters,
lists, etc. are passed through literally to Valkey search.
Multiple filters are combined with AND logic.
limit (int, optional): Maximum number of results to return. Defaults to 1000
top_k (int, optional): Maximum number of results to return. Defaults to 1000
if not specified.
Returns:
@@ -786,10 +789,10 @@ class ValkeyDB(VectorStoreBase):
# Since Valkey search requires vector format, use a dummy vector search
# that returns all documents by using a zero vector and large K
dummy_vector = [0.0] * self.embedding_model_dims
search_limit = limit if limit is not None else 1000 # Large default
search_limit = top_k if top_k is not None else 1000 # Large default
# Use the existing search method which handles filters properly
search_results = self.search("", dummy_vector, limit=search_limit, filters=filters)
search_results = self.search("", dummy_vector, top_k=search_limit, filters=filters)
# Convert search results to list format (match Redis format)
class MemoryResult:
mem0/vector_stores/vertex_ai_vector_search.py
+15 -13
View File
@@ -5,7 +5,9 @@ from typing import Any, Dict, List, Optional, Tuple
import google.api_core.exceptions
from google.cloud import aiplatform, aiplatform_v1
from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import Namespace
from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import (
Namespace,
)
from google.oauth2 import service_account
from pydantic import BaseModel
@@ -206,20 +208,20 @@ class GoogleMatchingEngine(VectorStoreBase):
raise
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
Args:
query (str): Query.
vectors (List[float]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5.
top_k (int, optional): Number of results to return. Defaults to 5.
filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
Returns:
List[OutputData]: Search results (unwrapped)
"""
logger.debug("Starting search")
logger.debug("Limit: %d, Filters: %s", limit, filters)
logger.debug("Limit: %d, Filters: %s", top_k, filters)
try:
filter_namespaces = []
@@ -241,7 +243,7 @@ class GoogleMatchingEngine(VectorStoreBase):
response = self.index_endpoint.find_neighbors(
deployed_index_id=self.deployment_index_id,
queries=[vectors],
num_neighbors=limit,
num_neighbors=top_k,
filter=filter_namespaces if filter_namespaces else None,
return_full_datapoint=True,
)
@@ -274,7 +276,7 @@ class GoogleMatchingEngine(VectorStoreBase):
logger.error("Stack trace: %s", traceback.format_exc())
raise
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
# Vertex AI hybrid search requires sparse embeddings configuration.
# Not yet supported - requires HybridQuery with sparse encoder setup.
return None
@@ -458,12 +460,12 @@ class GoogleMatchingEngine(VectorStoreBase):
"region": self.region,
}
def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None) -> List[List[OutputData]]:
def list(self, filters: Optional[Dict] = None, top_k: Optional[int] = None) -> List[List[OutputData]]:
"""List vectors matching the given filters.
Args:
filters: Optional filters to apply
limit: Optional maximum number of results to return
top_k: Optional maximum number of results to return
Returns:
List[List[OutputData]]: List of matching vectors wrapped in an extra array
@@ -471,17 +473,17 @@ class GoogleMatchingEngine(VectorStoreBase):
"""
logger.debug("Starting list operation")
logger.debug("Filters: %s", filters)
logger.debug("Limit: %s", limit)
logger.debug("Limit: %s", top_k)
try:
# Use a zero vector for the search
dimension = 768 # This should be configurable based on the model
zero_vector = [0.0] * dimension
# Use a large limit if none specified
search_limit = limit if limit is not None else 10000
# Use a large top_k if none specified
search_limit = top_k if top_k is not None else 10000
results = self.search(query=zero_vector, limit=search_limit, filters=filters)
results = self.search(query=zero_vector, top_k=search_limit, filters=filters)
logger.debug("Found %d results", len(results))
return [results] # Wrap in extra array to match interface
@@ -614,7 +616,7 @@ class GoogleMatchingEngine(VectorStoreBase):
logger.debug("Filter: %s", filter)
embedding = self.embedder.embed_query(query)
results = self.search(query=embedding, limit=k, filters=filter)
results = self.search(query=embedding, top_k=k, filters=filter)
docs_and_scores = [
(Document(page_content=result.payload.get("text", ""), metadata=result.payload), result.score)
mem0/vector_stores/weaviate.py
+7 -7
View File
@@ -179,7 +179,7 @@ class Weaviate(VectorStoreBase):
batch.add_object(collection=self.collection_name, properties=data_object, uuid=object_id, vector=vector)
def search(
self, query: str, vectors: List[float], limit: int = 5, filters: Optional[Dict] = None
self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[Dict] = None
) -> List[OutputData]:
"""
Search for similar vectors.
@@ -194,7 +194,7 @@ class Weaviate(VectorStoreBase):
response = collection.query.hybrid(
query="",
vector=vectors,
limit=limit,
limit=top_k,
filters=combined_filter,
return_properties=["hash", "created_at", "updated_at", "user_id", "agent_id", "run_id", "data", "category"],
return_metadata=MetadataQuery(score=True),
@@ -223,13 +223,13 @@ class Weaviate(VectorStoreBase):
)
return results
def keyword_search(self, query, limit=5, filters=None):
def keyword_search(self, query, top_k=5, filters=None):
"""
Search for memories using BM25 keyword search.
Args:
query (str): Search query text.
limit (int): Maximum number of results. Defaults to 5.
top_k (int): Maximum number of results. Defaults to 5.
filters (dict, optional): Filters to apply (user_id, agent_id, run_id).
Returns:
@@ -245,7 +245,7 @@ class Weaviate(VectorStoreBase):
response = collection.query.bm25(
query=query,
query_properties=["data"],
limit=limit,
limit=top_k,
filters=combined_filter,
return_properties=["hash", "created_at", "updated_at", "user_id", "agent_id", "run_id", "data", "category"],
return_metadata=MetadataQuery(score=True),
@@ -362,7 +362,7 @@ class Weaviate(VectorStoreBase):
return schema
return None
def list(self, filters=None, limit=100) -> List[OutputData]:
def list(self, filters=None, top_k=100) -> List[OutputData]:
"""
List all vectors in a collection.
"""
@@ -374,7 +374,7 @@ class Weaviate(VectorStoreBase):
filter_conditions.append(Filter.by_property(key).equal(value))
combined_filter = Filter.all_of(filter_conditions) if filter_conditions else None
response = collection.query.fetch_objects(
limit=limit,
limit=top_k,
filters=combined_filter,
return_properties=["hash", "created_at", "updated_at", "user_id", "agent_id", "run_id", "data", "category"],
)