Added metadata filtering and reranker for the OSS

This commit is contained in:
parshvadaftari
2025-09-02 01:15:17 +05:30
parent a9a1d57dfa
commit 61e8668584
15 changed files with 624 additions and 82 deletions
+19 -6
View File
@@ -29,14 +29,27 @@ client.search(query) # Uses v2 API endpoint
client.get_all() # Uses v2 API endpoint client.get_all() # Uses v2 API endpoint
``` ```
### 2. API Versioning Behavior ### 2. API Versioning Strategy Clarification
**New API versioning behavior:** **IMPORTANT: Understanding the New Versioning Strategy**
- **Add operations**: Use v1 endpoint by default, v1.1 when `output_format="v1.1"` is specified The API versioning strategy in mem0ai 1.0.0 has been unified and simplified:
- **Get/Search operations**: Always use v2 endpoints for enhanced filtering capabilities
- **Removed parameters**: `version` parameter from `get_all()` and `search()` methods #### **Endpoint vs Format Distinction**
- **Async mode**: Now default behavior (no parameter needed) - **API Endpoints** (`/v1/`, `/v2/`): Control which REST API version to use
- **Response Formats** (v1.0, v1.1): Control the structure of the returned data
#### **New Unified Strategy:**
- **Add operations**: Always use `/v1/` endpoint with v1.1 response format (no more output_format parameter)
- **Search operations**: Always use `/v2/` endpoint with v1.1 response format
- **Get_all operations**: Always use `/v2/` endpoint with v1.1 response format
- **Response format**: All operations now return v1.1 format (`{"results": [...]}`)
#### **What Changed:**
- ✅ **Consistent response format**: Everything returns v1.1 format
- ✅ **Simplified API**: No more `output_format` or `version` parameters to manage
- ✅ **Endpoint optimization**: Add uses v1, Search/Get use v2 for best performance
- ❌ **Removed v1.0 support**: v1.0 response format is no longer supported
### 3. Response Format Standardization ### 3. Response Format Standardization
+26 -10
View File
@@ -133,21 +133,29 @@ class MemoryClient:
Args: Args:
messages: A list of message dictionaries. messages: A list of message dictionaries.
**kwargs: Additional parameters such as user_id, agent_id, app_id, **kwargs: Additional parameters such as user_id, agent_id, app_id,
metadata, filters, output_format. metadata, filters.
Returns: Returns:
A dictionary containing the API response. A dictionary containing the API response in v1.1 format.
Raises: Raises:
APIError: If the API request fails. APIError: If the API request fails.
""" """
kwargs = self._prepare_params(kwargs) kwargs = self._prepare_params(kwargs)
# Handle output_format parameter for add operations (default to v1.1)
if "output_format" not in kwargs: # Remove deprecated parameters
kwargs["output_format"] = "v1.1" if "output_format" in kwargs:
warnings.warn(
"output_format parameter is deprecated and ignored. All responses now use v1.1 format.",
DeprecationWarning,
stacklevel=2,
)
kwargs.pop("output_format")
kwargs.pop("async_mode", None) # Remove async_mode as it's now default kwargs.pop("async_mode", None) # Remove async_mode as it's now default
# For add operations, use output_format parameter (keep it in kwargs) # Force v1.1 format for all add operations
kwargs["output_format"] = "v1.1"
payload = self._prepare_payload(messages, kwargs) payload = self._prepare_payload(messages, kwargs)
response = self.client.post("/v1/memories/", json=payload) response = self.client.post("/v1/memories/", json=payload)
response.raise_for_status() response.raise_for_status()
@@ -979,12 +987,20 @@ class AsyncMemoryClient:
@api_error_handler @api_error_handler
async def add(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]: async def add(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
kwargs = self._prepare_params(kwargs) kwargs = self._prepare_params(kwargs)
# Handle output_format parameter for add operations (default to v1.1)
if "output_format" not in kwargs: # Remove deprecated parameters
kwargs["output_format"] = "v1.1" if "output_format" in kwargs:
warnings.warn(
"output_format parameter is deprecated and ignored. All responses now use v1.1 format.",
DeprecationWarning,
stacklevel=2,
)
kwargs.pop("output_format")
kwargs.pop("async_mode", None) # Remove async_mode as it's now default kwargs.pop("async_mode", None) # Remove async_mode as it's now default
# For add operations, use output_format parameter (keep it in kwargs) # Force v1.1 format for all add operations
kwargs["output_format"] = "v1.1"
payload = self._prepare_payload(messages, kwargs) payload = self._prepare_payload(messages, kwargs)
response = await self.async_client.post("/v1/memories/", json=payload) response = await self.async_client.post("/v1/memories/", json=payload)
response.raise_for_status() response.raise_for_status()
+5
View File
@@ -7,6 +7,7 @@ from mem0.embeddings.configs import EmbedderConfig
from mem0.graphs.configs import GraphStoreConfig from mem0.graphs.configs import GraphStoreConfig
from mem0.llms.configs import LlmConfig from mem0.llms.configs import LlmConfig
from mem0.vector_stores.configs import VectorStoreConfig from mem0.vector_stores.configs import VectorStoreConfig
from mem0.configs.rerankers.config import RerankerConfig
# Set up the directory path # Set up the directory path
home_dir = os.path.expanduser("~") home_dir = os.path.expanduser("~")
@@ -47,6 +48,10 @@ class MemoryConfig(BaseModel):
description="Configuration for the graph", description="Configuration for the graph",
default_factory=GraphStoreConfig, default_factory=GraphStoreConfig,
) )
reranker: Optional[RerankerConfig] = Field(
description="Configuration for the reranker",
default=None,
)
version: str = Field( version: str = Field(
description="The version of the API", description="The version of the API",
default="v1.1", default="v1.1",
+6
View File
@@ -0,0 +1,6 @@
from .base import BaseRerankerConfig
from .cohere import CohereRerankerConfig
from .sentence_transformer import SentenceTransformerRerankerConfig
from .config import RerankerConfig
__all__ = ["BaseRerankerConfig", "CohereRerankerConfig", "SentenceTransformerRerankerConfig", "RerankerConfig"]
+30
View File
@@ -0,0 +1,30 @@
from abc import ABC
from typing import Optional
class BaseRerankerConfig(ABC):
"""
Base configuration for rerankers with only common parameters.
Provider-specific configurations should be handled by separate config classes.
This class contains only the parameters that are common across all reranker providers.
For provider-specific parameters, use the appropriate provider config class.
"""
def __init__(
self,
model: Optional[str] = None,
api_key: Optional[str] = None,
top_k: Optional[int] = None,
):
"""
Initialize a base configuration class instance for the reranker.
Args:
model (str, optional): The reranker model to use.
api_key (str, optional): The API key for the reranker service.
top_k (int, optional): Maximum number of documents to return after reranking.
"""
self.model = model
self.api_key = api_key
self.top_k = top_k
+34
View File
@@ -0,0 +1,34 @@
from typing import Optional
from mem0.configs.rerankers.base import BaseRerankerConfig
class CohereRerankerConfig(BaseRerankerConfig):
"""
Configuration class for Cohere reranker-specific parameters.
Inherits from BaseRerankerConfig and adds Cohere-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = "rerank-english-v3.0",
api_key: Optional[str] = None,
top_k: Optional[int] = None,
# Cohere-specific parameters
return_documents: bool = False,
max_chunks_per_doc: Optional[int] = None,
):
"""
Initialize Cohere reranker configuration.
Args:
model (str): The Cohere rerank model to use.
api_key (str, optional): The Cohere API key.
top_k (int, optional): Maximum number of documents to return after reranking.
return_documents (bool): Whether to return the document texts in the response.
max_chunks_per_doc (int, optional): Maximum number of chunks per document.
"""
super().__init__(model=model, api_key=api_key, top_k=top_k)
self.return_documents = return_documents
self.max_chunks_per_doc = max_chunks_per_doc
+14
View File
@@ -0,0 +1,14 @@
from typing import Optional
from pydantic import BaseModel, Field
from mem0.configs.rerankers.base import BaseRerankerConfig
class RerankerConfig(BaseModel):
"""Configuration for rerankers."""
provider: str = Field(description="Reranker provider (e.g., 'cohere', 'sentence_transformer')", default="cohere")
config: Optional[BaseRerankerConfig] = Field(description="Provider-specific reranker configuration", default=None)
model_config = {"extra": "forbid"}
@@ -0,0 +1,37 @@
from typing import Optional
from mem0.configs.rerankers.base import BaseRerankerConfig
class SentenceTransformerRerankerConfig(BaseRerankerConfig):
"""
Configuration class for Sentence Transformer reranker-specific parameters.
Inherits from BaseRerankerConfig and adds Sentence Transformer-specific settings.
"""
def __init__(
self,
# Base parameters
model: Optional[str] = "cross-encoder/ms-marco-MiniLM-L-6-v2",
api_key: Optional[str] = None, # Not used for sentence transformers
top_k: Optional[int] = None,
# Sentence Transformer-specific parameters
device: Optional[str] = None,
batch_size: int = 32,
show_progress_bar: bool = False,
):
"""
Initialize Sentence Transformer reranker configuration.
Args:
model (str): The cross-encoder model name to use.
api_key (str, optional): Not used for sentence transformers.
top_k (int, optional): Maximum number of documents to return after reranking.
device (str, optional): Device to run the model on ('cpu', 'cuda', etc.).
batch_size (int): Batch size for processing documents.
show_progress_bar (bool): Whether to show progress bar during processing.
"""
super().__init__(model=model, api_key=api_key, top_k=top_k)
self.device = device
self.batch_size = batch_size
self.show_progress_bar = show_progress_bar
+179 -38
View File
@@ -36,6 +36,7 @@ from mem0.utils.factory import (
GraphStoreFactory, GraphStoreFactory,
LlmFactory, LlmFactory,
VectorStoreFactory, VectorStoreFactory,
RerankerFactory,
) )
@@ -138,6 +139,14 @@ class Memory(MemoryBase):
self.collection_name = self.config.vector_store.config.collection_name self.collection_name = self.config.vector_store.config.collection_name
self.api_version = self.config.version self.api_version = self.config.version
# Initialize reranker if configured
self.reranker = None
if config.reranker:
self.reranker = RerankerFactory.create(
config.reranker.provider,
config.reranker.config
)
self.enable_graph = False self.enable_graph = False
if self.config.graph_store.config: if self.config.graph_store.config:
@@ -265,14 +274,11 @@ class Memory(MemoryBase):
graph_result = future2.result() graph_result = future2.result()
if self.api_version == "v1.0": if self.api_version == "v1.0":
warnings.warn( raise ValueError(
"The v1.0 API format is deprecated and will be removed in mem0ai 2.0.0. " "The v1.0 API format is no longer supported in mem0ai 1.0.0+. "
"Please upgrade to v1.1 format which returns a dict with 'results' key. " "Please use v1.1 format which returns a dict with 'results' key. "
"Set version='v1.1' in your MemoryConfig.", "Remove version='v1.0' from your MemoryConfig or set it to version='v1.1'."
DeprecationWarning,
stacklevel=2,
) )
return vector_store_result
if self.enable_graph: if self.enable_graph:
return { return {
@@ -562,14 +568,11 @@ class Memory(MemoryBase):
return {"results": all_memories_result, "relations": graph_entities_result} return {"results": all_memories_result, "relations": graph_entities_result}
if self.api_version == "v1.0": if self.api_version == "v1.0":
warnings.warn( raise ValueError(
"The v1.0 API format is deprecated and will be removed in mem0ai 2.0.0. " "The v1.0 API format is no longer supported in mem0ai 1.0.0+. "
"Please upgrade to v1.1 format which returns a dict with 'results' key. " "Please use v1.1 format which returns a dict with 'results' key. "
"Set version='v1.1' in your MemoryConfig.", "Remove version='v1.0' from your MemoryConfig or set it to version='v1.1'."
DeprecationWarning,
stacklevel=2,
) )
return all_memories_result
else: else:
return {"results": all_memories_result} return {"results": all_memories_result}
@@ -622,6 +625,7 @@ class Memory(MemoryBase):
limit: int = 100, limit: int = 100,
filters: Optional[Dict[str, Any]] = None, filters: Optional[Dict[str, Any]] = None,
threshold: Optional[float] = None, threshold: Optional[float] = None,
rerank: bool = True,
): ):
""" """
Searches for memories based on a query Searches for memories based on a query
@@ -631,8 +635,20 @@ class Memory(MemoryBase):
agent_id (str, optional): ID of the agent 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. 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. limit (int, optional): Limit the number of results. Defaults to 100.
filters (dict, optional): Filters to apply to the search. Defaults to None.. 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. threshold (float, optional): Minimum score for a memory to be included in the results. Defaults to None.
metadata_filters (dict, optional): Enhanced metadata filtering with operators:
- {"key": "value"} - exact match
- {"key": {"$eq": "value"}} - equals
- {"key": {"$ne": "value"}} - not equals
- {"key": {"$in": ["val1", "val2"]}} - in list
- {"key": {"$nin": ["val1", "val2"]}} - not in list
- {"key": {"$gt": 10}} - greater than
- {"key": {"$gte": 10}} - greater than or equal
- {"key": {"$lt": 10}} - less than
- {"key": {"$lte": 10}} - less than or equal
- {"$and": [filter1, filter2]} - logical AND
- {"$or": [filter1, filter2]} - logical OR
Returns: Returns:
dict: A dictionary containing the search results, typically under a "results" key, dict: A dictionary containing the search results, typically under a "results" key,
@@ -646,6 +662,14 @@ class Memory(MemoryBase):
if not any(key in effective_filters for key in ("user_id", "agent_id", "run_id")): if not any(key in effective_filters for key in ("user_id", "agent_id", "run_id")):
raise ValueError("At least one of 'user_id', 'agent_id', or 'run_id' must be specified.") raise ValueError("At least one of 'user_id', 'agent_id', or 'run_id' must be specified.")
# Apply enhanced metadata filtering if advanced operators are detected
if filters and self._has_advanced_operators(filters):
processed_filters = self._process_metadata_filters(filters)
effective_filters.update(processed_filters)
elif filters:
# Simple filters, merge directly
effective_filters.update(filters)
keys, encoded_ids = process_telemetry_filters(effective_filters) keys, encoded_ids = process_telemetry_filters(effective_filters)
capture_event( capture_event(
"mem0.search", "mem0.search",
@@ -657,6 +681,7 @@ class Memory(MemoryBase):
"encoded_ids": encoded_ids, "encoded_ids": encoded_ids,
"sync_type": "sync", "sync_type": "sync",
"threshold": threshold, "threshold": threshold,
"advanced_filters": bool(filters and self._has_advanced_operators(filters)),
}, },
) )
@@ -673,21 +698,97 @@ class Memory(MemoryBase):
original_memories = future_memories.result() original_memories = future_memories.result()
graph_entities = future_graph_entities.result() if future_graph_entities else None graph_entities = future_graph_entities.result() if future_graph_entities else None
# 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)
original_memories = reranked_memories
except Exception as e:
logger.warning(f"Reranking failed, using original results: {e}")
if self.enable_graph: if self.enable_graph:
return {"results": original_memories, "relations": graph_entities} return {"results": original_memories, "relations": graph_entities}
if self.api_version == "v1.0": if self.api_version == "v1.0":
warnings.warn( raise ValueError(
"The v1.0 API format is deprecated and will be removed in mem0ai 2.0.0. " "The v1.0 API format is no longer supported in mem0ai 1.0.0+. "
"Please upgrade to v1.1 format which returns a dict with 'results' key. " "Please use v1.1 format which returns a dict with 'results' key. "
"Set version='v1.1' in your MemoryConfig.", "Remove version='v1.0' from your MemoryConfig or set it to version='v1.1'."
DeprecationWarning,
stacklevel=2,
) )
return {"results": original_memories}
else: else:
return {"results": original_memories} return {"results": original_memories}
def _process_metadata_filters(self, metadata_filters: Dict[str, Any]) -> Dict[str, Any]:
"""
Process enhanced metadata filters and convert them to vector store compatible format.
Args:
metadata_filters: Enhanced metadata filters with operators
Returns:
Dict of processed filters compatible with vector store
"""
processed_filters = {}
def process_condition(key: str, condition: Any) -> Dict[str, Any]:
if not isinstance(condition, dict):
# Simple equality: {"key": "value"}
return {key: {"$eq": condition}}
result = {}
for operator, value in condition.items():
if operator in ["$eq", "$ne", "$gt", "$gte", "$lt", "$lte", "$in", "$nin"]:
result[key] = {operator: value}
else:
raise ValueError(f"Unsupported metadata filter operator: {operator}")
return result
for key, value in metadata_filters.items():
if key == "$and":
# Logical AND: combine multiple conditions
if not isinstance(value, list):
raise ValueError("$and operator requires a list of conditions")
for condition in value:
for sub_key, sub_value in condition.items():
processed_filters.update(process_condition(sub_key, sub_value))
elif key == "$or":
# Logical OR: for now, we'll apply the first condition
# Note: Full OR support would require vector store level implementation
if not isinstance(value, list) or not value:
raise ValueError("$or operator requires a non-empty list of conditions")
# Apply first condition as fallback
first_condition = value[0]
for sub_key, sub_value in first_condition.items():
processed_filters.update(process_condition(sub_key, sub_value))
else:
processed_filters.update(process_condition(key, value))
return processed_filters
def _has_advanced_operators(self, filters: Dict[str, Any]) -> bool:
"""
Check if filters contain advanced operators that need special processing.
Args:
filters: Dictionary of filters to check
Returns:
bool: True if advanced operators are detected
"""
if not isinstance(filters, dict):
return False
for key, value in filters.items():
# Check for logical operators
if key in ["$and", "$or"]:
return True
# Check for advanced comparison operators
if isinstance(value, dict):
for op in value.keys():
if op in ["$eq", "$ne", "$gt", "$gte", "$lt", "$lte", "$in", "$nin"]:
return True
return False
def _search_vector_store(self, query, filters, limit, threshold: Optional[float] = None): def _search_vector_store(self, query, filters, limit, threshold: Optional[float] = None):
embeddings = self.embedding_model.embed(query, "search") 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, limit=limit, filters=filters)
@@ -991,6 +1092,14 @@ class AsyncMemory(MemoryBase):
self.collection_name = self.config.vector_store.config.collection_name self.collection_name = self.config.vector_store.config.collection_name
self.api_version = self.config.version self.api_version = self.config.version
# Initialize reranker if configured
self.reranker = None
if config.reranker:
self.reranker = RerankerFactory.create(
config.reranker.provider,
config.reranker.config
)
self.enable_graph = False self.enable_graph = False
if self.config.graph_store.config: if self.config.graph_store.config:
@@ -1103,14 +1212,11 @@ class AsyncMemory(MemoryBase):
vector_store_result, graph_result = await asyncio.gather(vector_store_task, graph_task) vector_store_result, graph_result = await asyncio.gather(vector_store_task, graph_task)
if self.api_version == "v1.0": if self.api_version == "v1.0":
warnings.warn( raise ValueError(
"The v1.0 API format is deprecated and will be removed in mem0ai 2.0.0. " "The v1.0 API format is no longer supported in mem0ai 1.0.0+. "
"Please upgrade to v1.1 format which returns a dict with 'results' key. " "Please use v1.1 format which returns a dict with 'results' key. "
"Set version='v1.1' in your MemoryConfig.", "Remove version='v1.0' from your MemoryConfig or set it to version='v1.1'."
DeprecationWarning,
stacklevel=2,
) )
return vector_store_result
if self.enable_graph: if self.enable_graph:
return { return {
@@ -1427,7 +1533,11 @@ class AsyncMemory(MemoryBase):
DeprecationWarning, DeprecationWarning,
stacklevel=2, stacklevel=2,
) )
return results_dict["results"] raise ValueError(
"The v1.0 API format is no longer supported in mem0ai 1.0.0+. "
"Please use v1.1 format which returns a dict with 'results' key. "
"Remove version='v1.0' from your MemoryConfig or set it to version='v1.1'."
)
return results_dict return results_dict
@@ -1480,6 +1590,8 @@ class AsyncMemory(MemoryBase):
limit: int = 100, limit: int = 100,
filters: Optional[Dict[str, Any]] = None, filters: Optional[Dict[str, Any]] = None,
threshold: Optional[float] = None, threshold: Optional[float] = None,
metadata_filters: Optional[Dict[str, Any]] = None,
rerank: bool = True,
): ):
""" """
Searches for memories based on a query Searches for memories based on a query
@@ -1489,8 +1601,20 @@ class AsyncMemory(MemoryBase):
agent_id (str, optional): ID of the agent 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. 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. limit (int, optional): Limit the number of results. Defaults to 100.
filters (dict, optional): Filters to apply to the search. Defaults to None. 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. threshold (float, optional): Minimum score for a memory to be included in the results. Defaults to None.
metadata_filters (dict, optional): Enhanced metadata filtering with operators:
- {"key": "value"} - exact match
- {"key": {"$eq": "value"}} - equals
- {"key": {"$ne": "value"}} - not equals
- {"key": {"$in": ["val1", "val2"]}} - in list
- {"key": {"$nin": ["val1", "val2"]}} - not in list
- {"key": {"$gt": 10}} - greater than
- {"key": {"$gte": 10}} - greater than or equal
- {"key": {"$lt": 10}} - less than
- {"key": {"$lte": 10}} - less than or equal
- {"$and": [filter1, filter2]} - logical AND
- {"$or": [filter1, filter2]} - logical OR
Returns: Returns:
dict: A dictionary containing the search results, typically under a "results" key, dict: A dictionary containing the search results, typically under a "results" key,
@@ -1505,6 +1629,14 @@ class AsyncMemory(MemoryBase):
if not any(key in effective_filters for key in ("user_id", "agent_id", "run_id")): if not any(key in effective_filters for key in ("user_id", "agent_id", "run_id")):
raise ValueError("at least one of 'user_id', 'agent_id', or 'run_id' must be specified ") raise ValueError("at least one of 'user_id', 'agent_id', or 'run_id' must be specified ")
# Apply enhanced metadata filtering if advanced operators are detected
if filters and self._has_advanced_operators(filters):
processed_filters = self._process_metadata_filters(filters)
effective_filters.update(processed_filters)
elif filters:
# Simple filters, merge directly
effective_filters.update(filters)
keys, encoded_ids = process_telemetry_filters(effective_filters) keys, encoded_ids = process_telemetry_filters(effective_filters)
capture_event( capture_event(
"mem0.search", "mem0.search",
@@ -1516,6 +1648,7 @@ class AsyncMemory(MemoryBase):
"encoded_ids": encoded_ids, "encoded_ids": encoded_ids,
"sync_type": "async", "sync_type": "async",
"threshold": threshold, "threshold": threshold,
"advanced_filters": bool(filters and self._has_advanced_operators(filters)),
}, },
) )
@@ -1534,18 +1667,26 @@ class AsyncMemory(MemoryBase):
original_memories = await vector_store_task original_memories = await vector_store_task
graph_entities = None graph_entities = None
# Apply reranking if enabled and reranker is available
if rerank and self.reranker and original_memories:
try:
# Run reranking in thread pool to avoid blocking async loop
reranked_memories = await asyncio.to_thread(
self.reranker.rerank, query, original_memories, limit
)
original_memories = reranked_memories
except Exception as e:
logger.warning(f"Reranking failed, using original results: {e}")
if self.enable_graph: if self.enable_graph:
return {"results": original_memories, "relations": graph_entities} return {"results": original_memories, "relations": graph_entities}
if self.api_version == "v1.0": if self.api_version == "v1.0":
warnings.warn( raise ValueError(
"The v1.0 API format is deprecated and will be removed in mem0ai 2.0.0. " "The v1.0 API format is no longer supported in mem0ai 1.0.0+. "
"Please upgrade to v1.1 format which returns a dict with 'results' key. " "Please use v1.1 format which returns a dict with 'results' key. "
"Set version='v1.1' in your MemoryConfig.", "Remove version='v1.0' from your MemoryConfig or set it to version='v1.1'."
DeprecationWarning,
stacklevel=2,
) )
return {"results": original_memories}
else: else:
return {"results": original_memories} return {"results": original_memories}
+16 -28
View File
@@ -281,23 +281,17 @@ class MemoryGraph:
# Build query based on whether agent_id is provided # Build query based on whether agent_id is provided
if filters.get("agent_id"): if filters.get("agent_id"):
cypher_query = """ cypher_query = """
MATCH (n:Entity {user_id: $user_id, agent_id: $agent_id}) CALL vector_search.search("memzero", $limit, $n_embedding)
WHERE n.embedding IS NOT NULL YIELD distance, node, similarity
WITH n, $n_embedding as n_embedding WITH node AS n, similarity
CALL node_similarity.cosine_pairwise("embedding", [n_embedding], [n.embedding]) WHERE n:Entity AND n.user_id = $user_id AND n.agent_id = $agent_id AND n.embedding IS NOT NULL AND similarity >= $threshold
YIELD node1, node2, similarity
WITH n, similarity
WHERE similarity >= $threshold
MATCH (n)-[r]->(m:Entity) MATCH (n)-[r]->(m:Entity)
RETURN n.name AS source, id(n) AS source_id, type(r) AS relationship, id(r) AS relation_id, m.name AS destination, id(m) AS destination_id, similarity RETURN n.name AS source, id(n) AS source_id, type(r) AS relationship, id(r) AS relation_id, m.name AS destination, id(m) AS destination_id, similarity
UNION UNION
MATCH (n:Entity {user_id: $user_id, agent_id: $agent_id}) CALL vector_search.search("memzero", $limit, $n_embedding)
WHERE n.embedding IS NOT NULL YIELD distance, node, similarity
WITH n, $n_embedding as n_embedding WITH node AS n, similarity
CALL node_similarity.cosine_pairwise("embedding", [n_embedding], [n.embedding]) WHERE n:Entity AND n.user_id = $user_id AND n.agent_id = $agent_id AND n.embedding IS NOT NULL AND similarity >= $threshold
YIELD node1, node2, similarity
WITH n, similarity
WHERE similarity >= $threshold
MATCH (m:Entity)-[r]->(n) MATCH (m:Entity)-[r]->(n)
RETURN m.name AS source, id(m) AS source_id, type(r) AS relationship, id(r) AS relation_id, n.name AS destination, id(n) AS destination_id, similarity RETURN m.name AS source, id(m) AS source_id, type(r) AS relationship, id(r) AS relation_id, n.name AS destination, id(n) AS destination_id, similarity
ORDER BY similarity DESC ORDER BY similarity DESC
@@ -312,23 +306,17 @@ class MemoryGraph:
} }
else: else:
cypher_query = """ cypher_query = """
MATCH (n:Entity {user_id: $user_id}) CALL vector_search.search("memzero", $limit, $n_embedding)
WHERE n.embedding IS NOT NULL YIELD distance, node, similarity
WITH n, $n_embedding as n_embedding WITH node AS n, similarity
CALL node_similarity.cosine_pairwise("embedding", [n_embedding], [n.embedding]) WHERE n:Entity AND n.user_id = $user_id AND n.embedding IS NOT NULL AND similarity >= $threshold
YIELD node1, node2, similarity
WITH n, similarity
WHERE similarity >= $threshold
MATCH (n)-[r]->(m:Entity) MATCH (n)-[r]->(m:Entity)
RETURN n.name AS source, id(n) AS source_id, type(r) AS relationship, id(r) AS relation_id, m.name AS destination, id(m) AS destination_id, similarity RETURN n.name AS source, id(n) AS source_id, type(r) AS relationship, id(r) AS relation_id, m.name AS destination, id(m) AS destination_id, similarity
UNION UNION
MATCH (n:Entity {user_id: $user_id}) CALL vector_search.search("memzero", $limit, $n_embedding)
WHERE n.embedding IS NOT NULL YIELD distance, node, similarity
WITH n, $n_embedding as n_embedding WITH node AS n, similarity
CALL node_similarity.cosine_pairwise("embedding", [n_embedding], [n.embedding]) WHERE n:Entity AND n.user_id = $user_id AND n.embedding IS NOT NULL AND similarity >= $threshold
YIELD node1, node2, similarity
WITH n, similarity
WHERE similarity >= $threshold
MATCH (m:Entity)-[r]->(n) MATCH (m:Entity)-[r]->(n)
RETURN m.name AS source, id(m) AS source_id, type(r) AS relationship, id(r) AS relation_id, n.name AS destination, id(n) AS destination_id, similarity RETURN m.name AS source, id(m) AS source_id, type(r) AS relationship, id(r) AS relation_id, n.name AS destination, id(n) AS destination_id, similarity
ORDER BY similarity DESC ORDER BY similarity DESC
+9
View File
@@ -0,0 +1,9 @@
"""
Reranker implementations for mem0 search functionality.
"""
from .base import BaseReranker
from .cohere_reranker import CohereReranker
from .sentence_transformer_reranker import SentenceTransformerReranker
__all__ = ["BaseReranker", "CohereReranker", "SentenceTransformerReranker"]
+20
View File
@@ -0,0 +1,20 @@
from abc import ABC, abstractmethod
from typing import List, Dict, Any
class BaseReranker(ABC):
"""Abstract base class for all rerankers."""
@abstractmethod
def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
"""
Rerank documents based on relevance to the query.
Args:
query: The search query
documents: List of documents to rerank, each with 'memory' field
top_k: Number of top documents to return (None = return all)
Returns:
List of reranked documents with added 'rerank_score' field
"""
pass
+85
View File
@@ -0,0 +1,85 @@
import os
from typing import List, Dict, Any, Optional
from mem0.reranker.base import BaseReranker
try:
import cohere
COHERE_AVAILABLE = True
except ImportError:
COHERE_AVAILABLE = False
class CohereReranker(BaseReranker):
"""Cohere-based reranker implementation."""
def __init__(self, config):
"""
Initialize Cohere reranker.
Args:
config: CohereRerankerConfig object with configuration parameters
"""
if not COHERE_AVAILABLE:
raise ImportError("cohere package is required for CohereReranker. Install with: pip install cohere")
self.config = config
self.api_key = config.api_key or os.getenv("COHERE_API_KEY")
if not self.api_key:
raise ValueError("Cohere API key is required. Set COHERE_API_KEY environment variable or pass api_key in config.")
self.model = config.model
self.client = cohere.Client(self.api_key)
def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
"""
Rerank documents using Cohere's rerank API.
Args:
query: The search query
documents: List of documents to rerank
top_k: Number of top documents to return
Returns:
List of reranked documents with rerank_score
"""
if not documents:
return documents
# Extract text content for reranking
doc_texts = []
for doc in documents:
if 'memory' in doc:
doc_texts.append(doc['memory'])
elif 'text' in doc:
doc_texts.append(doc['text'])
elif 'content' in doc:
doc_texts.append(doc['content'])
else:
doc_texts.append(str(doc))
try:
# Call Cohere rerank API
response = self.client.rerank(
model=self.model,
query=query,
documents=doc_texts,
top_k=top_k or self.config.top_k or len(documents),
return_documents=self.config.return_documents,
max_chunks_per_doc=self.config.max_chunks_per_doc,
)
# Create reranked results
reranked_docs = []
for result in response.results:
original_doc = documents[result.index].copy()
original_doc['rerank_score'] = result.relevance_score
reranked_docs.append(original_doc)
return reranked_docs
except Exception as e:
# Fallback to original order if reranking fails
for doc in documents:
doc['rerank_score'] = 0.0
return documents[:top_k] if top_k else documents
@@ -0,0 +1,90 @@
from typing import List, Dict, Any, Optional
import numpy as np
from mem0.reranker.base import BaseReranker
try:
from sentence_transformers import SentenceTransformer, util
SENTENCE_TRANSFORMERS_AVAILABLE = True
except ImportError:
SENTENCE_TRANSFORMERS_AVAILABLE = False
class SentenceTransformerReranker(BaseReranker):
"""Sentence Transformer based reranker implementation."""
def __init__(self, config):
"""
Initialize Sentence Transformer reranker.
Args:
config: SentenceTransformerRerankerConfig object with configuration parameters
"""
if not SENTENCE_TRANSFORMERS_AVAILABLE:
raise ImportError("sentence-transformers package is required for SentenceTransformerReranker. Install with: pip install sentence-transformers")
self.config = config
self.model = SentenceTransformer(config.model, device=config.device)
def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
"""
Rerank documents using sentence transformer cross-encoder.
Args:
query: The search query
documents: List of documents to rerank
top_k: Number of top documents to return
Returns:
List of reranked documents with rerank_score
"""
if not documents:
return documents
# Extract text content for reranking
doc_texts = []
for doc in documents:
if 'memory' in doc:
doc_texts.append(doc['memory'])
elif 'text' in doc:
doc_texts.append(doc['text'])
elif 'content' in doc:
doc_texts.append(doc['content'])
else:
doc_texts.append(str(doc))
try:
# Create query-document pairs
pairs = [[query, doc_text] for doc_text in doc_texts]
# Get similarity scores
scores = self.model.predict(pairs)
if isinstance(scores, np.ndarray):
scores = scores.tolist()
# Combine documents with scores
doc_score_pairs = list(zip(documents, scores))
# Sort by score (descending)
doc_score_pairs.sort(key=lambda x: x[1], reverse=True)
# Apply top_k limit
final_top_k = top_k or self.config.top_k
if final_top_k:
doc_score_pairs = doc_score_pairs[:final_top_k]
# Create reranked results
reranked_docs = []
for doc, score in doc_score_pairs:
reranked_doc = doc.copy()
reranked_doc['rerank_score'] = float(score)
reranked_docs.append(reranked_doc)
return reranked_docs
except Exception as e:
# Fallback to original order if reranking fails
for doc in documents:
doc['rerank_score'] = 0.0
final_top_k = top_k or self.config.top_k
return documents[:final_top_k] if final_top_k else documents
+54
View File
@@ -10,6 +10,9 @@ from mem0.configs.llms.lmstudio import LMStudioConfig
from mem0.configs.llms.ollama import OllamaConfig from mem0.configs.llms.ollama import OllamaConfig
from mem0.configs.llms.openai import OpenAIConfig from mem0.configs.llms.openai import OpenAIConfig
from mem0.configs.llms.vllm import VllmConfig from mem0.configs.llms.vllm import VllmConfig
from mem0.configs.rerankers.base import BaseRerankerConfig
from mem0.configs.rerankers.cohere import CohereRerankerConfig
from mem0.configs.rerankers.sentence_transformer import SentenceTransformerRerankerConfig
from mem0.embeddings.mock import MockEmbeddings from mem0.embeddings.mock import MockEmbeddings
@@ -214,3 +217,54 @@ class GraphStoreFactory:
except (ImportError, AttributeError) as e: except (ImportError, AttributeError) as e:
raise ImportError(f"Could not import MemoryGraph for provider '{provider_name}': {e}") raise ImportError(f"Could not import MemoryGraph for provider '{provider_name}': {e}")
return GraphClass(config) return GraphClass(config)
class RerankerFactory:
"""
Factory for creating reranker instances with appropriate configurations.
Supports provider-specific configs following the same pattern as other factories.
"""
# Provider mappings with their config classes
provider_to_class = {
"cohere": ("mem0.reranker.cohere_reranker.CohereReranker", CohereRerankerConfig),
"sentence_transformer": ("mem0.reranker.sentence_transformer_reranker.SentenceTransformerReranker", SentenceTransformerRerankerConfig),
}
@classmethod
def create(cls, provider_name: str, config: Optional[Union[BaseRerankerConfig, Dict]] = None, **kwargs):
"""
Create a reranker instance based on the provider and configuration.
Args:
provider_name: The reranker provider (e.g., 'cohere', 'sentence_transformer')
config: Configuration object or dictionary
**kwargs: Additional configuration parameters
Returns:
Reranker instance configured for the specified provider
Raises:
ImportError: If the provider class cannot be imported
ValueError: If the provider is not supported
"""
if provider_name not in cls.provider_to_class:
raise ValueError(f"Unsupported reranker provider: {provider_name}")
class_path, config_class = cls.provider_to_class[provider_name]
# Handle configuration
if config is None:
config = config_class(**kwargs)
elif isinstance(config, dict):
config = config_class(**config, **kwargs)
elif not isinstance(config, BaseRerankerConfig):
raise ValueError(f"Config must be a {config_class.__name__} instance or dict")
# Import and create the reranker class
try:
reranker_class = load_class(class_path)
except (ImportError, AttributeError) as e:
raise ImportError(f"Could not import reranker for provider '{provider_name}': {e}")
return reranker_class(config)