Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix/es query filter (🚨 URGENT) #2162

Merged
merged 3 commits into from
Jan 21, 2025
Merged

Fix/es query filter (🚨 URGENT) #2162

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7abc8a7
fix: elasticsearch query knn with filter option
alohays Jan 20, 2025
fe6f57b
fix: elasticsearch index setting more standard way
alohays Jan 20, 2025
8766b53
fix: fix test elasticsearch relative to new es query & filter
alohays Jan 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 77 additions & 23 deletions mem0/vector_stores/elasticsearch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,18 +49,28 @@ def __init__(self, **kwargs):
def create_index(self) -> None:
"""Create Elasticsearch index with proper mappings if it doesn't exist"""
index_settings = {
"settings": {
"index": {
"number_of_replicas": 1,
"number_of_shards": 5,
"refresh_interval": "1s"
}
},
"mappings": {
"properties": {
"text": {"type": "text"},
"embedding": {
"vector": {
"type": "dense_vector",
"dims": self.vector_dim,
"index": True,
"similarity": "cosine",
"similarity": "cosine"
},
"metadata": {"type": "object"},
"user_id": {"type": "keyword"},
"hash": {"type": "keyword"},
"metadata": {
"type": "object",
"properties": {
"user_id": {"type": "keyword"}
}
}
}
}
}
Expand Down Expand Up @@ -99,43 +109,76 @@ def insert(

actions = []
for i, (vec, id_) in enumerate(zip(vectors, ids)):
action = {"_index": self.collection_name, "_id": id_, "vector": vec, "payload": payloads[i]}
action = {
"_index": self.collection_name,
"_id": id_,
"_source": {
"vector": vec,
"metadata": payloads[i] # Store all metadata in the metadata field
}
}
actions.append(action)

bulk(self.client, actions)

# Return OutputData objects for inserted documents
results = []
for i, id_ in enumerate(ids):
results.append(
OutputData(
id=id_,
score=1.0, # Default score for inserts
payload=payloads[i],
payload=payloads[i]
)
)
return results

def search(self, query: List[float], limit: int = 5, filters: Optional[Dict] = None) -> List[OutputData]:
"""Search for similar vectors using KNN search with pre-filtering."""
search_query = {
"query": {
"bool": {
"must": [
# Exact match filters for memory isolation
*({"term": {f"payload.{k}": v}} for k, v in (filters or {}).items()),
# KNN vector search
{"knn": {"vector": {"vector": query, "k": limit}}},
]
if not filters:
# If no filters, just do KNN search
search_query = {
"knn": {
"field": "vector",
"query_vector": query,
"k": limit,
"num_candidates": limit * 2
}
}
else:
# If filters exist, apply them with KNN search
filter_conditions = []
for key, value in filters.items():
filter_conditions.append({
"term": {
f"metadata.{key}": value
}
})

search_query = {
"knn": {
"field": "vector",
"query_vector": query,
"k": limit,
"num_candidates": limit * 2,
"filter": {
"bool": {
"must": filter_conditions
}
}
}
}
}

response = self.client.search(index=self.collection_name, body=search_query)

results = []
for hit in response["hits"]["hits"]:
results.append(OutputData(id=hit["_id"], score=hit["_score"], payload=hit["_source"].get("payload", {})))
results.append(
OutputData(
id=hit["_id"],
score=hit["_score"],
payload=hit.get("_source", {}).get("metadata", {})
)
)

return results

Expand All @@ -149,7 +192,7 @@ def update(self, vector_id: str, vector: Optional[List[float]] = None, payload:
if vector is not None:
doc["vector"] = vector
if payload is not None:
doc["payload"] = payload
doc["metadata"] = payload

self.client.update(index=self.collection_name, id=vector_id, body={"doc": doc})

Expand All @@ -160,7 +203,7 @@ def get(self, vector_id: str) -> Optional[OutputData]:
return OutputData(
id=response["_id"],
score=1.0, # Default score for direct get
payload=response["_source"].get("payload", {}),
payload=response["_source"].get("metadata", {})
)
except KeyError as e:
logger.warning(f"Missing key in Elasticsearch response: {e}")
Expand Down Expand Up @@ -189,7 +232,18 @@ def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None) -> L
query: Dict[str, Any] = {"query": {"match_all": {}}}

if filters:
query["query"] = {"bool": {"must": [{"match": {f"payload.{k}": v}} for k, v in filters.items()]}}
filter_conditions = []
for key, value in filters.items():
filter_conditions.append({
"term": {
f"metadata.{key}": value
}
})
query["query"] = {
"bool": {
"must": filter_conditions
}
}

if limit:
query["size"] = limit
Expand All @@ -202,7 +256,7 @@ def list(self, filters: Optional[Dict] = None, limit: Optional[int] = None) -> L
OutputData(
id=hit["_id"],
score=1.0, # Default score for list operation
payload=hit["_source"].get("payload", {}),
payload=hit.get("_source", {}).get("metadata", {})
)
)

Expand Down
53 changes: 20 additions & 33 deletions tests/vector_stores/test_elasticsearch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,13 +92,11 @@ def test_create_index(self):
# Verify field mappings
mappings = create_args["body"]["mappings"]["properties"]
self.assertEqual(mappings["text"]["type"], "text")
self.assertEqual(mappings["embedding"]["type"], "dense_vector")
self.assertEqual(mappings["embedding"]["dims"], 1536)
self.assertEqual(mappings["embedding"]["index"], True)
self.assertEqual(mappings["embedding"]["similarity"], "cosine")
self.assertEqual(mappings["vector"]["type"], "dense_vector")
self.assertEqual(mappings["vector"]["dims"], 1536)
self.assertEqual(mappings["vector"]["index"], True)
self.assertEqual(mappings["vector"]["similarity"], "cosine")
self.assertEqual(mappings["metadata"]["type"], "object")
self.assertEqual(mappings["user_id"]["type"], "keyword")
self.assertEqual(mappings["hash"]["type"], "keyword")

# Reset mocks for next test
self.client_mock.reset_mock()
Expand Down Expand Up @@ -170,8 +168,8 @@ def test_insert(self):
self.assertEqual(len(actions), 2)
self.assertEqual(actions[0]["_index"], "test_collection")
self.assertEqual(actions[0]["_id"], "id1")
self.assertEqual(actions[0]["vector"], vectors[0])
self.assertEqual(actions[0]["payload"], payloads[0])
self.assertEqual(actions[0]["_source"]["vector"], vectors[0])
self.assertEqual(actions[0]["_source"]["metadata"], payloads[0])

# Verify returned objects
self.assertEqual(len(results), 2)
Expand All @@ -189,7 +187,7 @@ def test_search(self):
"_score": 0.8,
"_source": {
"vector": [0.1] * 1536,
"payload": {"key1": "value1"}
"metadata": {"key1": "value1"}
}
}
]
Expand All @@ -210,14 +208,11 @@ def test_search(self):
body = search_args["body"]

# Verify KNN query structure
self.assertIn("query", body)
self.assertIn("bool", body["query"])
self.assertIn("must", body["query"]["bool"])

# Verify KNN parameters
knn_query = body["query"]["bool"]["must"][-1]["knn"]["vector"]
self.assertEqual(knn_query["vector"], query_vector)
self.assertEqual(knn_query["k"], 5)
self.assertIn("knn", body)
self.assertEqual(body["knn"]["field"], "vector")
self.assertEqual(body["knn"]["query_vector"], query_vector)
self.assertEqual(body["knn"]["k"], 5)
self.assertEqual(body["knn"]["num_candidates"], 10)

# Verify results
self.assertEqual(len(results), 1)
Expand All @@ -231,10 +226,8 @@ def test_get(self):
"_id": "id1",
"_source": {
"vector": [0.1] * 1536,
"payload": {"key": "value"},
"text": "sample text",
"user_id": "test_user",
"hash": "sample_hash"
"metadata": {"key": "value"},
"text": "sample text"
}
}
self.client_mock.get.return_value = mock_response
Expand All @@ -248,17 +241,11 @@ def test_get(self):
id="id1"
)

# Basic assertions that should pass if OutputData is created correctly
# Verify result
self.assertIsNotNone(result)
self.assertTrue(hasattr(result, 'id'))
self.assertTrue(hasattr(result, 'score'))
self.assertTrue(hasattr(result, 'payload'))

# If the above assertions pass, we can safely check the values
if result is not None: # This satisfies the linter
self.assertEqual(result.id, "id1")
self.assertEqual(result.score, 1.0)
self.assertEqual(result.payload, {"key": "value"})
self.assertEqual(result.id, "id1")
self.assertEqual(result.score, 1.0)
self.assertEqual(result.payload, {"key": "value"})

def test_get_not_found(self):
# Mock get raising exception
Expand All @@ -277,15 +264,15 @@ def test_list(self):
"_id": "id1",
"_source": {
"vector": [0.1] * 1536,
"payload": {"key1": "value1"}
"metadata": {"key1": "value1"}
},
"_score": 1.0
},
{
"_id": "id2",
"_source": {
"vector": [0.2] * 1536,
"payload": {"key2": "value2"}
"metadata": {"key2": "value2"}
},
"_score": 0.8
}
Expand Down
Loading