Added mulit id filters support for all vectorstores (#3269)
This commit is contained in:
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user