diff --git a/MIGRATION_GUIDE_v1.0.md b/MIGRATION_GUIDE_v1.0.md index 209d86c9d..db921169f 100644 --- a/MIGRATION_GUIDE_v1.0.md +++ b/MIGRATION_GUIDE_v1.0.md @@ -29,14 +29,27 @@ client.search(query) # 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 -- **Get/Search operations**: Always use v2 endpoints for enhanced filtering capabilities -- **Removed parameters**: `version` parameter from `get_all()` and `search()` methods -- **Async mode**: Now default behavior (no parameter needed) +The API versioning strategy in mem0ai 1.0.0 has been unified and simplified: + +#### **Endpoint vs Format Distinction** +- **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 diff --git a/mem0/client/main.py b/mem0/client/main.py index 8d6bbb68c..c2545420e 100644 --- a/mem0/client/main.py +++ b/mem0/client/main.py @@ -133,21 +133,29 @@ class MemoryClient: Args: messages: A list of message dictionaries. **kwargs: Additional parameters such as user_id, agent_id, app_id, - metadata, filters, output_format. + metadata, filters. Returns: - A dictionary containing the API response. + A dictionary containing the API response in v1.1 format. Raises: APIError: If the API request fails. """ kwargs = self._prepare_params(kwargs) - # Handle output_format parameter for add operations (default to v1.1) - if "output_format" not in kwargs: - kwargs["output_format"] = "v1.1" + + # Remove deprecated parameters + 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 - # 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) response = self.client.post("/v1/memories/", json=payload) response.raise_for_status() @@ -979,12 +987,20 @@ class AsyncMemoryClient: @api_error_handler async def add(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]: kwargs = self._prepare_params(kwargs) - # Handle output_format parameter for add operations (default to v1.1) - if "output_format" not in kwargs: - kwargs["output_format"] = "v1.1" + + # Remove deprecated parameters + 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 - # 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) response = await self.async_client.post("/v1/memories/", json=payload) response.raise_for_status() diff --git a/mem0/configs/base.py b/mem0/configs/base.py index 147d2593a..dd0dd9df4 100644 --- a/mem0/configs/base.py +++ b/mem0/configs/base.py @@ -7,6 +7,7 @@ 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("~") @@ -47,6 +48,10 @@ class MemoryConfig(BaseModel): description="Configuration for the graph", default_factory=GraphStoreConfig, ) + reranker: Optional[RerankerConfig] = Field( + description="Configuration for the reranker", + default=None, + ) version: str = Field( description="The version of the API", default="v1.1", diff --git a/mem0/configs/rerankers/__init__.py b/mem0/configs/rerankers/__init__.py new file mode 100644 index 000000000..8544b7c5b --- /dev/null +++ b/mem0/configs/rerankers/__init__.py @@ -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"] \ No newline at end of file diff --git a/mem0/configs/rerankers/base.py b/mem0/configs/rerankers/base.py new file mode 100644 index 000000000..cc2be05a6 --- /dev/null +++ b/mem0/configs/rerankers/base.py @@ -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 \ No newline at end of file diff --git a/mem0/configs/rerankers/cohere.py b/mem0/configs/rerankers/cohere.py new file mode 100644 index 000000000..52230aedc --- /dev/null +++ b/mem0/configs/rerankers/cohere.py @@ -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 \ No newline at end of file diff --git a/mem0/configs/rerankers/config.py b/mem0/configs/rerankers/config.py new file mode 100644 index 000000000..46367327d --- /dev/null +++ b/mem0/configs/rerankers/config.py @@ -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"} \ No newline at end of file diff --git a/mem0/configs/rerankers/sentence_transformer.py b/mem0/configs/rerankers/sentence_transformer.py new file mode 100644 index 000000000..6e320be94 --- /dev/null +++ b/mem0/configs/rerankers/sentence_transformer.py @@ -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 \ No newline at end of file diff --git a/mem0/memory/main.py b/mem0/memory/main.py index d466751f7..c7a7c0da4 100644 --- a/mem0/memory/main.py +++ b/mem0/memory/main.py @@ -36,6 +36,7 @@ from mem0.utils.factory import ( GraphStoreFactory, LlmFactory, VectorStoreFactory, + RerankerFactory, ) @@ -137,6 +138,14 @@ class Memory(MemoryBase): self.db = SQLiteManager(self.config.history_db_path) self.collection_name = self.config.vector_store.config.collection_name 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 @@ -265,14 +274,11 @@ class Memory(MemoryBase): graph_result = future2.result() if self.api_version == "v1.0": - warnings.warn( - "The v1.0 API format is deprecated and will be removed in mem0ai 2.0.0. " - "Please upgrade to v1.1 format which returns a dict with 'results' key. " - "Set version='v1.1' in your MemoryConfig.", - DeprecationWarning, - stacklevel=2, + 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 vector_store_result if self.enable_graph: return { @@ -562,14 +568,11 @@ class Memory(MemoryBase): return {"results": all_memories_result, "relations": graph_entities_result} if self.api_version == "v1.0": - warnings.warn( - "The v1.0 API format is deprecated and will be removed in mem0ai 2.0.0. " - "Please upgrade to v1.1 format which returns a dict with 'results' key. " - "Set version='v1.1' in your MemoryConfig.", - DeprecationWarning, - stacklevel=2, + 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 all_memories_result else: return {"results": all_memories_result} @@ -622,6 +625,7 @@ class Memory(MemoryBase): limit: int = 100, filters: Optional[Dict[str, Any]] = None, threshold: Optional[float] = None, + rerank: bool = True, ): """ 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. 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. - 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. + 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: 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")): 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) capture_event( "mem0.search", @@ -657,6 +681,7 @@ class Memory(MemoryBase): "encoded_ids": encoded_ids, "sync_type": "sync", "threshold": threshold, + "advanced_filters": bool(filters and self._has_advanced_operators(filters)), }, ) @@ -673,21 +698,97 @@ class Memory(MemoryBase): original_memories = future_memories.result() 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: return {"results": original_memories, "relations": graph_entities} if self.api_version == "v1.0": - warnings.warn( - "The v1.0 API format is deprecated and will be removed in mem0ai 2.0.0. " - "Please upgrade to v1.1 format which returns a dict with 'results' key. " - "Set version='v1.1' in your MemoryConfig.", - DeprecationWarning, - stacklevel=2, + 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": original_memories} else: 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): embeddings = self.embedding_model.embed(query, "search") memories = self.vector_store.search(query=query, vectors=embeddings, limit=limit, filters=filters) @@ -990,6 +1091,14 @@ class AsyncMemory(MemoryBase): self.db = SQLiteManager(self.config.history_db_path) self.collection_name = self.config.vector_store.config.collection_name 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 @@ -1103,14 +1212,11 @@ class AsyncMemory(MemoryBase): vector_store_result, graph_result = await asyncio.gather(vector_store_task, graph_task) if self.api_version == "v1.0": - warnings.warn( - "The v1.0 API format is deprecated and will be removed in mem0ai 2.0.0. " - "Please upgrade to v1.1 format which returns a dict with 'results' key. " - "Set version='v1.1' in your MemoryConfig.", - DeprecationWarning, - stacklevel=2, + 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 vector_store_result if self.enable_graph: return { @@ -1427,7 +1533,11 @@ class AsyncMemory(MemoryBase): DeprecationWarning, 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 @@ -1480,6 +1590,8 @@ class AsyncMemory(MemoryBase): limit: int = 100, filters: Optional[Dict[str, Any]] = None, threshold: Optional[float] = None, + metadata_filters: Optional[Dict[str, Any]] = None, + rerank: bool = True, ): """ 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. 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. - 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. + 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: 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")): 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) capture_event( "mem0.search", @@ -1516,6 +1648,7 @@ class AsyncMemory(MemoryBase): "encoded_ids": encoded_ids, "sync_type": "async", "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 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: return {"results": original_memories, "relations": graph_entities} if self.api_version == "v1.0": - warnings.warn( - "The v1.0 API format is deprecated and will be removed in mem0ai 2.0.0. " - "Please upgrade to v1.1 format which returns a dict with 'results' key. " - "Set version='v1.1' in your MemoryConfig.", - DeprecationWarning, - stacklevel=2, + 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": original_memories} else: return {"results": original_memories} diff --git a/mem0/memory/memgraph_memory.py b/mem0/memory/memgraph_memory.py index 2414746c7..e1d8a8e98 100644 --- a/mem0/memory/memgraph_memory.py +++ b/mem0/memory/memgraph_memory.py @@ -281,23 +281,17 @@ class MemoryGraph: # Build query based on whether agent_id is provided if filters.get("agent_id"): cypher_query = """ - MATCH (n:Entity {user_id: $user_id, agent_id: $agent_id}) - WHERE n.embedding IS NOT NULL - WITH n, $n_embedding as n_embedding - CALL node_similarity.cosine_pairwise("embedding", [n_embedding], [n.embedding]) - YIELD node1, node2, similarity - WITH n, similarity - WHERE similarity >= $threshold + CALL vector_search.search("memzero", $limit, $n_embedding) + YIELD distance, node, similarity + WITH node AS n, similarity + WHERE n:Entity AND n.user_id = $user_id AND n.agent_id = $agent_id AND n.embedding IS NOT NULL AND similarity >= $threshold 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 UNION - MATCH (n:Entity {user_id: $user_id, agent_id: $agent_id}) - WHERE n.embedding IS NOT NULL - WITH n, $n_embedding as n_embedding - CALL node_similarity.cosine_pairwise("embedding", [n_embedding], [n.embedding]) - YIELD node1, node2, similarity - WITH n, similarity - WHERE similarity >= $threshold + CALL vector_search.search("memzero", $limit, $n_embedding) + YIELD distance, node, similarity + WITH node AS n, similarity + WHERE n:Entity AND n.user_id = $user_id AND n.agent_id = $agent_id AND n.embedding IS NOT NULL AND similarity >= $threshold 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 ORDER BY similarity DESC @@ -312,23 +306,17 @@ class MemoryGraph: } else: cypher_query = """ - MATCH (n:Entity {user_id: $user_id}) - WHERE n.embedding IS NOT NULL - WITH n, $n_embedding as n_embedding - CALL node_similarity.cosine_pairwise("embedding", [n_embedding], [n.embedding]) - YIELD node1, node2, similarity - WITH n, similarity - WHERE similarity >= $threshold + CALL vector_search.search("memzero", $limit, $n_embedding) + YIELD distance, node, similarity + WITH node AS n, similarity + WHERE n:Entity AND n.user_id = $user_id AND n.embedding IS NOT NULL AND similarity >= $threshold 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 UNION - MATCH (n:Entity {user_id: $user_id}) - WHERE n.embedding IS NOT NULL - WITH n, $n_embedding as n_embedding - CALL node_similarity.cosine_pairwise("embedding", [n_embedding], [n.embedding]) - YIELD node1, node2, similarity - WITH n, similarity - WHERE similarity >= $threshold + CALL vector_search.search("memzero", $limit, $n_embedding) + YIELD distance, node, similarity + WITH node AS n, similarity + WHERE n:Entity AND n.user_id = $user_id AND n.embedding IS NOT NULL AND similarity >= $threshold 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 ORDER BY similarity DESC diff --git a/mem0/reranker/__init__.py b/mem0/reranker/__init__.py new file mode 100644 index 000000000..33919b788 --- /dev/null +++ b/mem0/reranker/__init__.py @@ -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"] \ No newline at end of file diff --git a/mem0/reranker/base.py b/mem0/reranker/base.py new file mode 100644 index 000000000..2d91ab964 --- /dev/null +++ b/mem0/reranker/base.py @@ -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 \ No newline at end of file diff --git a/mem0/reranker/cohere_reranker.py b/mem0/reranker/cohere_reranker.py new file mode 100644 index 000000000..55501b81a --- /dev/null +++ b/mem0/reranker/cohere_reranker.py @@ -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 \ No newline at end of file diff --git a/mem0/reranker/sentence_transformer_reranker.py b/mem0/reranker/sentence_transformer_reranker.py new file mode 100644 index 000000000..efe9ca8fa --- /dev/null +++ b/mem0/reranker/sentence_transformer_reranker.py @@ -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 \ No newline at end of file diff --git a/mem0/utils/factory.py b/mem0/utils/factory.py index e7eced5d2..e3a453622 100644 --- a/mem0/utils/factory.py +++ b/mem0/utils/factory.py @@ -10,6 +10,9 @@ from mem0.configs.llms.lmstudio import LMStudioConfig from mem0.configs.llms.ollama import OllamaConfig from mem0.configs.llms.openai import OpenAIConfig 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 @@ -214,3 +217,54 @@ class GraphStoreFactory: except (ImportError, AttributeError) as e: raise ImportError(f"Could not import MemoryGraph for provider '{provider_name}': {e}") 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)