Support Custom Search Query for Elasticsearch (#2372)

This commit is contained in:
Wonbin Kim
2025-03-18 14:04:34 +09:00
committed by GitHub
parent 00a2ea9ff0
commit b8f40f728f
4 changed files with 86 additions and 20 deletions

View File

@@ -54,6 +54,7 @@ Let's see the available parameters for the `elasticsearch` config:
| `password` | Password for basic authentication | `None` |
| `verify_certs` | Whether to verify SSL certificates | `True` |
| `auto_create_index` | Whether to automatically create the index | `True` |
| `custom_search_query` | Function returning a custom search query | `None` |
### Features
@@ -62,3 +63,46 @@ Let's see the available parameters for the `elasticsearch` config:
- Multiple authentication methods (Basic Auth, API Key)
- Automatic index creation with optimized mappings for vector search
- Memory isolation through payload filtering
- Custom search query function to customize the search query
### Custom Search Query
The `custom_search_query` parameter allows you to customize the search query when `Memory.search` is called.
__Example__
```python
import os
from typing import List, Optional, Dict
from mem0 import Memory
def custom_search_query(query: List[float], limit: int, filters: Optional[Dict]) -> Dict:
return {
"knn": {
"field": "vector",
"query_vector": query,
"k": limit,
"num_candidates": limit * 2
}
}
os.environ["OPENAI_API_KEY"] = "sk-xx"
config = {
"vector_store": {
"provider": "elasticsearch",
"config": {
"collection_name": "mem0",
"host": "localhost",
"port": 9200,
"embedding_model_dims": 1536,
"custom_search_query": custom_search_query
}
}
}
```
It should be a function that takes the following parameters:
- `query`: a query vector used in `Memory.search`
- `limit`: a number of results used in `Memory.search`
- `filters`: a dictionary of key-value pairs used in `Memory.search`. You can add custom pairs for the custom search query.
The function should return a query body for the Elasticsearch search API.