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: filters in chroma integration #1072

Merged
merged 18 commits into from
Sep 20, 2024
Merged

fix: filters in chroma integration #1072

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
3ae4118
Fix bugs in logic
Amnah199 Sep 6, 2024
6611846
Fix bug in filters and add tests
Amnah199 Sep 9, 2024
4b45f8e
PR comments and new tests
Amnah199 Sep 9, 2024
92f81f7
Merge branch 'main' into fix-chroma-filters
Amnah199 Sep 9, 2024
ac1fd63
Add filter conversion from haystack to chroma
Amnah199 Sep 13, 2024
1fc48d9
Fix linting
Amnah199 Sep 13, 2024
88a6f09
add new tests and fix docstrings
Amnah199 Sep 16, 2024
1550501
Update docstrings
Amnah199 Sep 16, 2024
3c13680
Merge branch 'main' into fix-chroma-filters
Amnah199 Sep 16, 2024
06edc97
Update filters.py
Amnah199 Sep 16, 2024
323a4d4
Improved code
Amnah199 Sep 16, 2024
b98e472
Merge branch 'main' into fix-chroma-filters
Amnah199 Sep 16, 2024
fabe3a1
Add a clear error message
Amnah199 Sep 16, 2024
af9121b
Updates based on PR review
Amnah199 Sep 18, 2024
d0ddb5e
Small fix
Amnah199 Sep 18, 2024
734024a
Small fixes
Amnah199 Sep 18, 2024
78860cc
Revised docstrings
Amnah199 Sep 20, 2024
58c3ff7
Merge branch 'main' into fix-chroma-filters
Amnah199 Sep 20, 2024
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
118 changes: 51 additions & 67 deletions integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,35 +121,26 @@ def filter_documents(self, filters: Optional[Dict[str, Any]] = None) -> List[Doc
`"$or"`, `"$not"`), a comparison operator (`"$eq"`, `$ne`, `"$in"`, `$nin`, `"$gt"`, `"$gte"`, `"$lt"`,
`"$lte"`) or a metadata field name.

Logical operator keys take a dictionary of metadata field names and/or logical operators as value. Metadata
field names take a dictionary of comparison operators as value. Comparison operator keys take a single value or
(in case of `"$in"`) a list of values as value. If no logical operator is provided, `"$and"` is used as default
operation. If no comparison operator is provided, `"$eq"` (or `"$in"` if the comparison value is a list) is used
as default operation.
Logical operator keys take a list of dictionaries of metadata field names and/or logical operators as value.
Metadata field names take a dictionary of comparison operators as value. Comparison operator keys take a single
value or (in case of `"$in"`) a list of values as value. If no comparison operator is provided, `"$eq"`
(or `"$in"` if the comparison value is a list) is used
as default operation. The logical operators `$and` and `$or` can be used to combine multiple filters.

Example:

```python
filters = {
"$and": {
"type": {"$eq": "article"},
"date": {"$gte": "2015-01-01", "$lt": "2021-01-01"},
"rating": {"$gte": 3},
"$or": {
"genre": {"$in": ["economy", "politics"]},
"publisher": {"$eq": "nytimes"}
"$and": [
{"type": {"$eq": "article"}},
{"date": {"$eq": "2015-01-01"}},
{"rating": {"$gte": 3}},
{"$or": [
{"genre": {"$in": ["economy", "politics"]}},
{"publisher": {"$eq": "nytimes"}}
]
}
}
}
# or simpler using default operators
filters = {
"type": "article",
"date": {"$gte": "2015-01-01", "$lt": "2021-01-01"},
"rating": {"$gte": 3},
"$or": {
"genre": ["economy", "politics"],
"publisher": "nytimes"
}
]
}
```

Expand All @@ -162,21 +153,17 @@ def filter_documents(self, filters: Optional[Dict[str, Any]] = None) -> List[Doc
filters = {
"$or": [
{
"$and": {
"Type": "News Paper",
"Date": {
"$lt": "2019-01-01"
}
}
"$and": [
{"type": "news paper"},
{"date": {"$eq": "2019-01-01"}}
]
},
{
"$and": {
"Type": "Blog Post",
"Date": {
"$gte": "2019-01-01"
}
}
}
"$and": [
{"type": "blog post"},
{"date": {"$eq": "2019-01-01"}}
]
},
]
}
```
Expand Down Expand Up @@ -363,50 +350,47 @@ def _normalize_filters(filters: Dict[str, Any]) -> Tuple[List[str], Dict[str, An
passed to `ids`, `where` and `where_document` respectively.
"""
if not isinstance(filters, dict):
msg = "'filters' parameter must be a dictionary"
msg = f"'{filters}' parameter must be a dictionary"
raise ChromaDocumentStoreFilterError(msg)

ids = []
where = defaultdict(list)
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
where_document = defaultdict(list)
keys_to_remove = []
where_document = {}
document_flag = False
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved

for field, value in filters.items():
# Process filters
for field, value in list(filters.items()):
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
if field == "content":
# Schedule for removal the original key, we're going to change it
keys_to_remove.append(field)
where_document["$contains"] = value
document_flag = True
del filters[field]
elif field == "id":
# Schedule for removal the original key, we're going to change it
keys_to_remove.append(field)
ids.append(value)
elif isinstance(value, (list, tuple)):
# Schedule for removal the original key, we're going to change it
keys_to_remove.append(field)

# if the list is empty the filter is invalid, let's just remove it
if len(value) == 0:
del filters[field]
elif isinstance(value, (list, tuple)) and field not in ["$and", "$or"]:
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
if not value: # Skip empty lists
del filters[field]
continue

# if the list has a single item, just make it a regular key:value filter pair
if len(value) == 1:
where[field] = value[0]
continue

# if the list contains multiple items, we need an $or chain
for v in value:
where["$or"].append({field: v})

for k in keys_to_remove:
del filters[k]

final_where = dict(filters)
final_where.update(dict(where))
where[field] = value[0] # Simplify single-item lists
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
# Check if list contains special operators that apply to documents
elif any(isinstance(v, dict) and any(k in v for k in ["$contains", "$not_contains"]) for v in value):
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
where_document = filters.copy()
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
document_flag = True
break
else:
where[field] = {"$in": value} # Create $in chain for multiple items
del filters[field]

final_where = {}
try:
if final_where:
validate_where(final_where)
if where_document:
if document_flag:
validate_where_document(where_document)
elif filters or where:
final_where.update(filters)
final_where.update(where)
validate_where(final_where)
except ValueError as e:
raise ChromaDocumentStoreFilterError(e) from e
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved

Expand Down Expand Up @@ -453,7 +437,7 @@ def _query_result_to_documents(result: QueryResult) -> List[List[Document]]:
for j in range(len(answers)):
document_dict: Dict[str, Any] = {
"id": result["ids"][i][j],
"content": documents[i][j],
"content": answers[j],
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
}

# prepare metadata
Expand Down
77 changes: 67 additions & 10 deletions integrations/chroma/tests/test_document_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,73 @@ def test_metadata_initialization(self, caplog):
assert store._collection.metadata["hnsw:space"] == "ip"
assert new_store._collection.metadata["hnsw:space"] == "ip"

# Override tests from LegacyFilterDocumentsTest
def test_filter_nested_or_and(self, document_store: ChromaDocumentStore, filterable_docs: List[Document]):
filters = {
"$and": [
{"page": {"$eq": "123"}},
{"$or": [{"name": {"$in": ["name_0", "name_1"]}}, {"number": {"$lt": 1.0}}]},
]
}

document_store.write_documents(filterable_docs)
result = document_store.filter_documents(filters)
self.assert_documents_are_equal(
result,
[
doc
for doc in filterable_docs
if (
doc.meta.get("page") in ["123"]
and (
doc.meta.get("name") in ["name_0", "name_1"]
or ("number" in doc.meta and doc.meta["number"] < 1)
)
)
],
)

def test_filter_nested_multiple_identical_operators_same_level(
self, document_store: ChromaDocumentStore, filterable_docs: List[Document]
):
filters = {
"$and": [
{"page": {"$eq": "123"}},
{"$or": [{"name": {"$in": ["name_0", "name_1"]}}, {"number": {"$lt": 1.0}}]},
{"$or": [{"name": {"$in": ["name_0", "name_1"]}}, {"number": {"$lt": 1.0}}]},
]
}

document_store.write_documents(filterable_docs)
filters = {
"$or": [
{
"$and": [
{"name": {"$in": ["name_0", "name_1"]}},
{"page": {"$eq": "100"}},
]
},
{
"$and": [
{"chapter": {"$in": ["intro", "abstract"]}},
{"page": "123"},
]
},
]
}
result = document_store.filter_documents(filters=filters)
self.assert_documents_are_equal(
result,
[
doc
for doc in filterable_docs
if (
(doc.meta.get("name") in ["name_0", "name_1"] and doc.meta.get("page") == "100")
or (doc.meta.get("chapter") in ["intro", "abstract"] and doc.meta.get("page") == "123")
)
],
)

@pytest.mark.skip(reason="Filter on dataframe contents is not supported.")
def test_filter_document_dataframe(self, document_store: ChromaDocumentStore, filterable_docs: List[Document]):
pass
Expand Down Expand Up @@ -304,13 +371,3 @@ def test_filter_nested_and_or_explicit(self, document_store: ChromaDocumentStore
@pytest.mark.skip(reason="Filter syntax not supported.")
def test_filter_nested_and_or_implicit(self, document_store: ChromaDocumentStore, filterable_docs: List[Document]):
pass

@pytest.mark.skip(reason="Filter syntax not supported.")
def test_filter_nested_or_and(self, document_store: ChromaDocumentStore, filterable_docs: List[Document]):
pass

@pytest.mark.skip(reason="Filter syntax not supported.")
def test_filter_nested_multiple_identical_operators_same_level(
self, document_store: ChromaDocumentStore, filterable_docs: List[Document]
):
pass
Loading