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

fix(sdk): removing deprecating param from our sdk and docs changes with it (#4740)

Browse Source
This commit is contained in:
Kartik
2026-04-12 00:34:58 +05:30
committed by GitHub
parent 3882af7450
commit e44b46ef2e
138 changed files with 1186 additions and 2840 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0/client/main.py
+208 -223
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,13 +661,11 @@ 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,
@@ -681,67 +673,53 @@ class MemoryClient:
"""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
multilingual: Whether to use the input language for memory storage and retrieval
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.
enable_graph: Enable or disable the graph 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
and multilingual 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,
"enable_graph": enable_graph,
"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,
"multilingual": multilingual,
}
)
payload = self._prepare_params(kwargs)
response = self.client.patch(
f"/api/v1/orgs/organizations/{self.org_id}/projects/{self.project_id}/",
json=payload,
@@ -750,19 +728,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,
"multilingual": multilingual,
"sync_type": "sync",
},
{**kwargs, "sync_type": "sync"},
)
return response.json()
@@ -937,20 +903,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}
@@ -966,8 +924,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.
@@ -978,8 +934,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.
@@ -989,8 +943,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:
@@ -1085,20 +1039,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}
@@ -1109,25 +1055,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()
@@ -1145,9 +1107,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 = {
@@ -1164,7 +1143,6 @@ class AsyncMemoryClient:
"client.get_all",
self,
{
"api_version": "v2",
"keys": list(kwargs.keys()),
"sync_type": "async",
},
@@ -1177,12 +1155,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()
@@ -1192,7 +1186,6 @@ class AsyncMemoryClient:
"client.search",
self,
{
"api_version": "v2",
"keys": list(kwargs.keys()),
"sync_type": "async",
},
@@ -1208,36 +1201,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()
@@ -1270,10 +1259,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:
@@ -1287,6 +1277,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()
@@ -1554,63 +1545,65 @@ 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
multilingual: Whether to use the input language for memory storage and retrieval
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.
enable_graph: Enable or disable the graph 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 multilingual 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,
"enable_graph": enable_graph,
"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,
"multilingual": multilingual,
}
)
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,
@@ -1619,15 +1612,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,
"multilingual": multilingual,
"sync_type": "async",
},
{**kwargs, "sync_type": "async"},
)
return response.json()
mem0/client/types.py
+103
View File
@@ -0,0 +1,103 @@
"""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"
)
enable_graph: Optional[bool] = Field(default=None, description="Whether to enable graph memory for this operation")
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")
enable_graph: Optional[bool] = Field(default=None, description="Whether to enable graph memory for this search")
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")
enable_graph: Optional[bool] = Field(default=None, description="Whether to enable graph memory for retrieval")
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"
)
enable_graph: Optional[bool] = Field(default=None, description="Whether to enable graph memory")
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 -3
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,8 +56,8 @@ 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(
mem0/configs/vector_stores/elasticsearch.py
+1 -1
View File
@@ -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")
mem0/graphs/neptune/base.py
+9 -9
View File
@@ -346,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:
@@ -438,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.
@@ -452,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 = []
@@ -470,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.
"""
@@ -484,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
+9 -9
View File
@@ -380,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},
)
@@ -413,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},
)
@@ -455,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
"""
@@ -469,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
"""
@@ -486,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,
)
@@ -504,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/memory/apache_age_memory.py
+7 -7
View File
@@ -213,14 +213,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:
list: A list of dicts with keys "source", "relationship", "destination".
@@ -286,13 +286,13 @@ class MemoryGraph:
)
self.ag.commit()
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:
- 'source': The source node name.
@@ -308,7 +308,7 @@ class MemoryGraph:
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(limit)
params.append(top_k)
results = self._exec_cypher(
f"MATCH (n)-[r]->(m) WHERE {where_clause} "
@@ -408,7 +408,7 @@ class MemoryGraph:
# -- graph DB operations ---------------------------------------------------
def _search_graph_db(self, node_list, filters, limit=100):
def _search_graph_db(self, node_list, filters, top_k=100):
"""Search similar nodes and their respective incoming and outgoing relations."""
result_relations = []
@@ -430,7 +430,7 @@ class MemoryGraph:
rel_where = " AND ".join(rel_where_parts)
# For each similar node, fetch its relationships
for sn in similar_nodes[:limit]:
for sn in similar_nodes[:top_k]:
node_name = sn["name"]
similarity = sn["similarity"]
mem0/memory/graph_memory.py
+7 -7
View File
@@ -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:
@@ -171,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"]
@@ -291,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 = []
@@ -332,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"]
mem0/memory/kuzu_memory.py
+9 -9
View File
@@ -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:
@@ -191,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.
@@ -205,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"]
@@ -316,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"]
@@ -364,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
mem0/memory/main.py
+39 -39
View File
@@ -258,7 +258,7 @@ class Memory(MemoryBase):
def __init__(self, config: MemoryConfig = MemoryConfig()):
self.config = config
self.custom_fact_extraction_prompt = self.config.custom_fact_extraction_prompt
self.custom_instructions = self.config.custom_instructions
self.custom_update_memory_prompt = self.config.custom_update_memory_prompt
self.embedding_model = EmbedderFactory.create(
self.config.embedder.provider,
@@ -525,8 +525,8 @@ class Memory(MemoryBase):
parsed_messages = parse_messages(messages)
if self.config.custom_fact_extraction_prompt:
system_prompt = self.config.custom_fact_extraction_prompt
if self.config.custom_instructions:
system_prompt = self.config.custom_instructions
user_prompt = f"Input:\n{parsed_messages}"
else:
# Determine if this should use agent memory extraction based on agent_id presence
@@ -582,7 +582,7 @@ class Memory(MemoryBase):
existing_memories = self.vector_store.search(
query=new_mem,
vectors=messages_embeddings,
limit=5,
top_k=5,
filters=search_filters,
)
for mem in existing_memories:
@@ -785,7 +785,7 @@ class Memory(MemoryBase):
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
filters: Optional[Dict[str, Any]] = None,
limit: int = 100,
top_k: int = 100,
):
"""
List all memories.
@@ -797,7 +797,7 @@ class Memory(MemoryBase):
filters (dict, optional): Additional custom key-value filters to apply to the search.
These are merged with the ID-based scoping filters. For example,
`filters={"actor_id": "some_user"}`.
limit (int, optional): The maximum number of memories to return. Defaults to 100.
top_k (int, optional): The maximum number of memories to return. Defaults to 100.
Returns:
dict: A dictionary containing a list of memories under the "results" key,
@@ -815,13 +815,13 @@ class Memory(MemoryBase):
keys, encoded_ids = process_telemetry_filters(effective_filters)
capture_event(
"mem0.get_all", self, {"limit": limit, "keys": keys, "encoded_ids": encoded_ids, "sync_type": "sync"}
"mem0.get_all", self, {"top_k": top_k, "keys": keys, "encoded_ids": encoded_ids, "sync_type": "sync"}
)
with concurrent.futures.ThreadPoolExecutor() as executor:
future_memories = executor.submit(self._get_all_from_vector_store, effective_filters, limit)
future_memories = executor.submit(self._get_all_from_vector_store, effective_filters, top_k)
future_graph_entities = (
executor.submit(self.graph.get_all, effective_filters, limit) if self.enable_graph else None
executor.submit(self.graph.get_all, effective_filters, top_k) if self.enable_graph else None
)
concurrent.futures.wait(
@@ -836,8 +836,8 @@ class Memory(MemoryBase):
return {"results": all_memories_result}
def _get_all_from_vector_store(self, filters, limit):
memories_result = self.vector_store.list(filters=filters, limit=limit)
def _get_all_from_vector_store(self, filters, top_k):
memories_result = self.vector_store.list(filters=filters, top_k=top_k)
# Handle different vector store return formats by inspecting first element
if isinstance(memories_result, (tuple, list)) and len(memories_result) > 0:
@@ -890,7 +890,7 @@ class Memory(MemoryBase):
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
limit: int = 100,
top_k: int = 100,
filters: Optional[Dict[str, Any]] = None,
threshold: Optional[float] = None,
rerank: bool = True,
@@ -902,7 +902,7 @@ class Memory(MemoryBase):
user_id (str, optional): ID of the user to search for. Defaults to None.
agent_id (str, optional): ID of the agent to search for. Defaults to None.
run_id (str, optional): ID of the run to search for. Defaults to None.
limit (int, optional): Limit the number of results. Defaults to 100.
top_k (int, optional): Maximum number of results to return. Defaults to 100.
filters (dict, optional): Legacy filters to apply to the search. Defaults to None.
threshold (float, optional): Minimum score for a memory to be included in the results. Defaults to None.
filters (dict, optional): Enhanced metadata filtering with operators:
@@ -947,7 +947,7 @@ class Memory(MemoryBase):
"mem0.search",
self,
{
"limit": limit,
"top_k": top_k,
"version": self.api_version,
"keys": keys,
"encoded_ids": encoded_ids,
@@ -958,9 +958,9 @@ class Memory(MemoryBase):
)
with concurrent.futures.ThreadPoolExecutor() as executor:
future_memories = executor.submit(self._search_vector_store, query, effective_filters, limit, threshold)
future_memories = executor.submit(self._search_vector_store, query, effective_filters, top_k, threshold)
future_graph_entities = (
executor.submit(self.graph.search, query, effective_filters, limit) if self.enable_graph else None
executor.submit(self.graph.search, query, effective_filters, top_k) if self.enable_graph else None
)
concurrent.futures.wait(
@@ -973,7 +973,7 @@ class Memory(MemoryBase):
# Apply reranking if enabled and reranker is available
if rerank and self.reranker and original_memories:
try:
reranked_memories = self.reranker.rerank(query, original_memories, limit)
reranked_memories = self.reranker.rerank(query, original_memories, top_k)
original_memories = reranked_memories
except Exception as e:
logger.warning(f"Reranking failed, using original results: {e}")
@@ -1079,9 +1079,9 @@ class Memory(MemoryBase):
return True
return False
def _search_vector_store(self, query, filters, limit, threshold: Optional[float] = None):
def _search_vector_store(self, query, filters, top_k, threshold: Optional[float] = None):
embeddings = self.embedding_model.embed(query, "search")
memories = self.vector_store.search(query=query, vectors=embeddings, limit=limit, filters=filters)
memories = self.vector_store.search(query=query, vectors=embeddings, top_k=top_k, filters=filters)
promoted_payload_keys = [
"user_id",
@@ -1643,8 +1643,8 @@ class AsyncMemory(MemoryBase):
return returned_memories
parsed_messages = parse_messages(messages)
if self.config.custom_fact_extraction_prompt:
system_prompt = self.config.custom_fact_extraction_prompt
if self.config.custom_instructions:
system_prompt = self.config.custom_instructions
user_prompt = f"Input:\n{parsed_messages}"
else:
# Determine if this should use agent memory extraction based on agent_id presence
@@ -1699,7 +1699,7 @@ class AsyncMemory(MemoryBase):
self.vector_store.search,
query=new_mem_content,
vectors=embeddings,
limit=5,
top_k=5,
filters=search_filters,
)
return [{"id": mem.id, "text": mem.payload.get("data", "")} for mem in existing_mems]
@@ -1922,7 +1922,7 @@ class AsyncMemory(MemoryBase):
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
filters: Optional[Dict[str, Any]] = None,
limit: int = 100,
top_k: int = 100,
):
"""
List all memories.
@@ -1934,7 +1934,7 @@ class AsyncMemory(MemoryBase):
filters (dict, optional): Additional custom key-value filters to apply to the search.
These are merged with the ID-based scoping filters. For example,
`filters={"actor_id": "some_user"}`.
limit (int, optional): The maximum number of memories to return. Defaults to 100.
top_k (int, optional): The maximum number of memories to return. Defaults to 100.
Returns:
dict: A dictionary containing a list of memories under the "results" key,
@@ -1955,19 +1955,19 @@ class AsyncMemory(MemoryBase):
keys, encoded_ids = process_telemetry_filters(effective_filters)
capture_event(
"mem0.get_all", self, {"limit": limit, "keys": keys, "encoded_ids": encoded_ids, "sync_type": "async"}
"mem0.get_all", self, {"top_k": top_k, "keys": keys, "encoded_ids": encoded_ids, "sync_type": "async"}
)
vector_store_task = asyncio.create_task(self._get_all_from_vector_store(effective_filters, limit))
vector_store_task = asyncio.create_task(self._get_all_from_vector_store(effective_filters, top_k))
graph_task = None
if self.enable_graph:
graph_get_all = getattr(self.graph, "get_all", None)
if callable(graph_get_all):
if asyncio.iscoroutinefunction(graph_get_all):
graph_task = asyncio.create_task(graph_get_all(effective_filters, limit))
graph_task = asyncio.create_task(graph_get_all(effective_filters, top_k))
else:
graph_task = asyncio.create_task(asyncio.to_thread(graph_get_all, effective_filters, limit))
graph_task = asyncio.create_task(asyncio.to_thread(graph_get_all, effective_filters, top_k))
results_dict = {}
if graph_task:
@@ -1978,8 +1978,8 @@ class AsyncMemory(MemoryBase):
return results_dict
async def _get_all_from_vector_store(self, filters, limit):
memories_result = await asyncio.to_thread(self.vector_store.list, filters=filters, limit=limit)
async def _get_all_from_vector_store(self, filters, top_k):
memories_result = await asyncio.to_thread(self.vector_store.list, filters=filters, top_k=top_k)
# Handle different vector store return formats by inspecting first element
if isinstance(memories_result, (tuple, list)) and len(memories_result) > 0:
@@ -2032,7 +2032,7 @@ class AsyncMemory(MemoryBase):
user_id: Optional[str] = None,
agent_id: Optional[str] = None,
run_id: Optional[str] = None,
limit: int = 100,
top_k: int = 100,
filters: Optional[Dict[str, Any]] = None,
threshold: Optional[float] = None,
metadata_filters: Optional[Dict[str, Any]] = None,
@@ -2045,7 +2045,7 @@ class AsyncMemory(MemoryBase):
user_id (str, optional): ID of the user to search for. Defaults to None.
agent_id (str, optional): ID of the agent to search for. Defaults to None.
run_id (str, optional): ID of the run to search for. Defaults to None.
limit (int, optional): Limit the number of results. Defaults to 100.
top_k (int, optional): Maximum number of results to return. Defaults to 100.
filters (dict, optional): Legacy filters to apply to the search. Defaults to None.
threshold (float, optional): Minimum score for a memory to be included in the results. Defaults to None.
filters (dict, optional): Enhanced metadata filtering with operators:
@@ -2091,7 +2091,7 @@ class AsyncMemory(MemoryBase):
"mem0.search",
self,
{
"limit": limit,
"top_k": top_k,
"version": self.api_version,
"keys": keys,
"encoded_ids": encoded_ids,
@@ -2101,14 +2101,14 @@ class AsyncMemory(MemoryBase):
},
)
vector_store_task = asyncio.create_task(self._search_vector_store(query, effective_filters, limit, threshold))
vector_store_task = asyncio.create_task(self._search_vector_store(query, effective_filters, top_k, threshold))
graph_task = None
if self.enable_graph:
if hasattr(self.graph.search, "__await__"): # Check if graph search is async
graph_task = asyncio.create_task(self.graph.search(query, effective_filters, limit))
graph_task = asyncio.create_task(self.graph.search(query, effective_filters, top_k))
else:
graph_task = asyncio.create_task(asyncio.to_thread(self.graph.search, query, effective_filters, limit))
graph_task = asyncio.create_task(asyncio.to_thread(self.graph.search, query, effective_filters, top_k))
if graph_task:
original_memories, graph_entities = await asyncio.gather(vector_store_task, graph_task)
@@ -2121,7 +2121,7 @@ class AsyncMemory(MemoryBase):
try:
# Run reranking in thread pool to avoid blocking async loop
reranked_memories = await asyncio.to_thread(
self.reranker.rerank, query, original_memories, limit
self.reranker.rerank, query, original_memories, top_k
)
original_memories = reranked_memories
except Exception as e:
@@ -2228,10 +2228,10 @@ class AsyncMemory(MemoryBase):
return True
return False
async def _search_vector_store(self, query, filters, limit, threshold: Optional[float] = None):
async def _search_vector_store(self, query, filters, top_k, threshold: Optional[float] = None):
embeddings = await asyncio.to_thread(self.embedding_model.embed, query, "search")
memories = await asyncio.to_thread(
self.vector_store.search, query=query, vectors=embeddings, limit=limit, filters=filters
self.vector_store.search, query=query, vectors=embeddings, top_k=top_k, filters=filters
)
promoted_payload_keys = [
mem0/memory/memgraph_memory.py
+9 -9
View File
@@ -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:
@@ -172,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.
@@ -193,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)
@@ -293,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 = []
@@ -324,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 = """
@@ -348,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)
mem0/memory/utils.py
+1 -1
View File
@@ -38,7 +38,7 @@ def ensure_json_instruction(system_prompt, user_prompt):
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_fact_extraction_prompt that doesn't include 'json', this causes 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.
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/vector_stores/azure_ai_search.py
+8 -8
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,
)
@@ -327,13 +327,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.
@@ -342,7 +342,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
+8 -8
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. "
@@ -243,7 +243,7 @@ class AzureMySQL(VectorStoreBase):
self,
query: str,
vectors: List[float],
limit: int = 5,
top_k: int = 5,
filters: Optional[Dict] = None,
) -> List[OutputData]:
"""
@@ -252,7 +252,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:
@@ -291,9 +291,9 @@ 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])
@@ -406,14 +406,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
@@ -436,7 +436,7 @@ class AzureMySQL(VectorStoreBase):
{filter_clause}
LIMIT %s
""",
(*filter_params, limit)
(*filter_params, top_k)
)
results = cur.fetchall()
mem0/vector_stores/baidu.py
+6 -6
View File
@@ -185,14 +185,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:
@@ -207,7 +207,7 @@ class BaiduDB(VectorStoreBase):
request = VectorTopkSearchRequest(
vector_field="vector",
vector=FloatVector(vectors),
limit=limit,
limit=top_k,
filter=search_filter,
config=VectorSearchConfig(ef=200),
)
@@ -313,20 +313,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
+2 -2
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
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/databricks.py
+19 -11
View File
@@ -2,20 +2,28 @@ 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
@@ -446,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:
@@ -468,7 +476,7 @@ class Databricks(VectorStoreBase):
query_kwargs = {
"index_name": self.fully_qualified_index_name,
"columns": self.column_names,
"num_results": limit,
"num_results": top_k,
"query_type": self.query_type,
"filters_json": filters_json,
}
@@ -719,20 +727,20 @@ 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
# Use query_text for Delta Sync with model endpoint, query_vector otherwise
query_kwargs = {
mem0/vector_stores/elasticsearch.py
+6 -6
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 = []
@@ -203,7 +203,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": {}}}
@@ -213,8 +213,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
+15 -16
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
@@ -450,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.
@@ -482,7 +481,7 @@ class FAISS(VectorStoreBase):
)
count += 1
if count >= limit:
if count >= top_k:
break
return [results]
mem0/vector_stores/langchain.py
+5 -5
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
@@ -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
+6 -6
View File
@@ -139,14 +139,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:
@@ -156,7 +156,7 @@ class MilvusDB(VectorStoreBase):
hits = self.client.search(
collection_name=self.collection_name,
data=[vectors],
limit=limit,
limit=top_k,
filter=query_filter,
output_fields=["*"],
)
@@ -235,19 +235,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
+7 -7
View File
@@ -113,14 +113,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:
@@ -139,8 +139,8 @@ class MongoDB(VectorStoreBase):
{
"$vectorSearch": {
"index": self.index_name,
"limit": limit,
"numCandidates": min(limit * 20, 10000),
"limit": top_k,
"numCandidates": min(top_k * 20, 10000),
"queryVector": vectors,
"path": "embedding",
}
@@ -271,13 +271,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.
@@ -292,7 +292,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
+7 -7
View File
@@ -158,7 +158,7 @@ class OpenSearchDB(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 for similar vectors using OpenSearch k-NN search with optional filters."""
@@ -167,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 = []
@@ -196,7 +196,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] # 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:
@@ -284,7 +284,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": {}}}
@@ -299,8 +299,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"]
mem0/vector_stores/pgvector.py
+6 -6
View File
@@ -203,7 +203,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]:
"""
@@ -212,7 +212,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:
@@ -237,7 +237,7 @@ class PGVector(VectorStoreBase):
ORDER BY distance
LIMIT %s
""",
(vectors, *filter_params, limit),
(vectors, *filter_params, top_k),
)
results = cur.fetchall()
@@ -350,14 +350,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.
@@ -380,7 +380,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
+6 -6
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,
}
@@ -320,13 +320,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.
@@ -340,7 +340,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
+6 -6
View File
@@ -308,14 +308,14 @@ class Qdrant(VectorStoreBase):
must_not=must_not or None,
)
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): 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:
@@ -326,7 +326,7 @@ class Qdrant(VectorStoreBase):
collection_name=self.collection_name,
query=vectors,
query_filter=query_filter,
limit=limit,
limit=top_k,
)
return hits.points
@@ -404,13 +404,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.
@@ -419,7 +419,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
+5 -5
View File
@@ -141,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)
@@ -150,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)
@@ -254,15 +254,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 [
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
+6 -6
View File
@@ -157,7 +157,7 @@ class TurbopufferDB(VectorStoreBase):
return ("And", tuple(conditions))
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.
@@ -165,7 +165,7 @@ class TurbopufferDB(VectorStoreBase):
Args:
query (str): Query text (unused in vector search, kept for interface consistency).
vectors (list): Query vector to search with.
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:
@@ -173,7 +173,7 @@ class TurbopufferDB(VectorStoreBase):
"""
query_params = {
"rank_by": ("vector", "ANN", vectors),
"top_k": limit,
"top_k": top_k,
"include_attributes": True,
}
@@ -290,20 +290,20 @@ class TurbopufferDB(VectorStoreBase):
except Exception:
return {"name": self.collection_name}
def list(self, filters: Optional[Dict] = None, limit: int = 100) -> list:
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.
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: Wrapped list of OutputData objects ([[results]]).
"""
query_params = {
"rank_by": ("vector", "ANN", [0.0] * self.embedding_model_dims),
"top_k": limit,
"top_k": top_k,
"include_attributes": True,
}
mem0/vector_stores/upstash_vector.py
+7 -7
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,
@@ -205,12 +205,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.
"""
@@ -233,7 +233,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
+8 -8
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)
@@ -768,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.
@@ -778,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:
@@ -789,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
+14 -12
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,
)
@@ -453,12 +455,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
@@ -466,17 +468,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
@@ -609,7 +611,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
+4 -4
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),
@@ -313,7 +313,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.
"""
@@ -325,7 +325,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"],
)