Added mulit id filters support for all vectorstores (#3269)

This commit is contained in:
Parshva Daftari
2025-08-04 14:52:57 -07:00
committed by GitHub
parent 57a16aeb4b
commit 42e60d6724
8 changed files with 728 additions and 45 deletions
+27 -2
View File
@@ -142,7 +142,8 @@ class ChromaDB(VectorStoreBase):
Returns:
List[OutputData]: Search results.
"""
results = self.collection.query(query_embeddings=vectors, where=filters, n_results=limit)
where_clause = self._generate_where_clause(filters) if filters else None
results = self.collection.query(query_embeddings=vectors, where=where_clause, n_results=limit)
final_results = self._parse_output(results)
return final_results
@@ -219,7 +220,8 @@ class ChromaDB(VectorStoreBase):
Returns:
List[OutputData]: List of vectors.
"""
results = self.collection.get(where=filters, limit=limit)
where_clause = self._generate_where_clause(filters) if filters else None
results = self.collection.get(where=where_clause, limit=limit)
return [self._parse_output(results)]
def reset(self):
@@ -227,3 +229,26 @@ class ChromaDB(VectorStoreBase):
logger.warning(f"Resetting index {self.collection_name}...")
self.delete_col()
self.collection = self.create_col(self.collection_name)
@staticmethod
def _generate_where_clause(where: dict[str, any]) -> dict[str, any]:
"""
Generate a properly formatted where clause for ChromaDB.
Args:
where (dict[str, any]): The filter conditions.
Returns:
dict[str, any]: Properly formatted where clause for ChromaDB.
"""
# If only one filter is supplied, return it as is
# (no need to wrap in $and based on chroma docs)
if where is None:
return {}
if len(where.keys()) <= 1:
return where
where_filters = []
for k, v in where.items():
if isinstance(v, str):
where_filters.append({k: v})
return {"$and": where_filters}