Added metadata filtering and reranker for the OSS
This commit is contained in:
+19
-6
@@ -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
@@ -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()
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
@@ -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"]
|
||||||
@@ -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
|
||||||
@@ -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
|
||||||
@@ -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
@@ -36,6 +36,7 @@ from mem0.utils.factory import (
|
|||||||
GraphStoreFactory,
|
GraphStoreFactory,
|
||||||
LlmFactory,
|
LlmFactory,
|
||||||
VectorStoreFactory,
|
VectorStoreFactory,
|
||||||
|
RerankerFactory,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -137,6 +138,14 @@ class Memory(MemoryBase):
|
|||||||
self.db = SQLiteManager(self.config.history_db_path)
|
self.db = SQLiteManager(self.config.history_db_path)
|
||||||
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
|
||||||
|
|
||||||
@@ -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)
|
||||||
@@ -990,6 +1091,14 @@ class AsyncMemory(MemoryBase):
|
|||||||
self.db = SQLiteManager(self.config.history_db_path)
|
self.db = SQLiteManager(self.config.history_db_path)
|
||||||
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
|
||||||
|
|
||||||
@@ -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}
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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"]
|
||||||
@@ -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
|
||||||
@@ -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
|
||||||
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user