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 4 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
115 changes: 52 additions & 63 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 @@ -120,35 +120,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 @@ -161,21 +152,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 @@ -362,50 +349,52 @@ 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)
where_document = {}
# Flag to check if the filters contain a document filter i.e. "$contains" or "$not_contains"
document_flag = False
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
keys_to_remove = []

for field, value in filters.items():
if field == "content":
# Schedule for removal the original key, we're going to change it
keys_to_remove.append(field)
where_document["$contains"] = value
elif field == "id":
# Schedule for removal the original key, we're going to change it
document_flag = True
keys_to_remove.append(field)
elif field == "id":
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:
elif isinstance(value, (list)) and field not in ["$and", "$or"]:
if not value: # Skip empty lists
keys_to_remove.append(field)
continue

# if the list has a single item, just make it a regular key:value filter pair
# if the list has a single item, make it a key:value filter pair
# e.g. filter = {"name": ["Alice"]}
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))
# Check if list contains special operators that apply to documents
if any(isinstance(v, dict) and any(k in v for k in ["$contains", "$not_contains"]) for v in value):
where_document = filters
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
document_flag = True
break
else:
where[field] = {"$in": value}
keys_to_remove.append(field)

# Remove items from filters after iteration
for field in keys_to_remove:
del filters[field]
Amnah199 marked this conversation as resolved.
Show resolved Hide resolved
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
118 changes: 108 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 @@ -234,6 +234,114 @@ 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_in_implicit(self, document_store: ChromaDocumentStore, filterable_docs: List[Document]):
filters = {"page": {"$in": ["100", "123"]}}

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", "100"]],
)

def test_filter_nested_multiple_identical_operators_same_level(
self, document_store: ChromaDocumentStore, filterable_docs: List[Document]
):

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")
)
],
)

def test_document_filter_contains(self, document_store: ChromaDocumentStore, filterable_docs: List[Document]):
document_store.write_documents(filterable_docs)
filters = {
"$or": [
{
"$and": [
{"name": {"$in": ["name_0", "name_1"]}},
{"page": {"$eq": "100"}},
]
},
{"$contains": "FOO"},
]
}
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.content
and "FOO" in doc.content
)
],
)

def test_logical_and_document_filter_combination(
self, document_store: ChromaDocumentStore, filterable_docs: List[Document]
):
document_store.write_documents(filterable_docs)

result = document_store.filter_documents(filters={"$contains": "FOO"})
self.assert_documents_are_equal(
result,
[doc for doc in filterable_docs if doc.content and "FOO" in doc.content],
)

@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 @@ -309,13 +417,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