diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/tests/test_vector_stores_deeplake.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/tests/test_vector_stores_deeplake.py index 088cb73af5c6f08da644b48d769b63421d09fdd5..8d404bbf1a072e9eff61737da4e99e6975b08d7f 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/tests/test_vector_stores_deeplake.py +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/tests/test_vector_stores_deeplake.py @@ -1,4 +1,5 @@ import jwt # noqa +import pytest from llama_index.core import Document from llama_index.core.vector_stores.types import ( @@ -16,7 +17,8 @@ def test_class(): assert BasePydanticVectorStore.__name__ in names_of_base_classes -def test_e2e(): +@pytest.fixture() +def vs_ids(): vs = DeepLakeVectorStore(dataset_path="mem://test", overwrite=True) ids = vs.add( nodes=[ @@ -25,6 +27,12 @@ def test_e2e(): Document(text="Doc 3", embedding=[1, 2, 3], metadata={"a": "3", "b": 12}), ] ) + yield (vs, ids) + vs.clear() + + +def test_filters(vs_ids): + vs, ids = vs_ids nodes = vs.get_nodes(node_ids=[ids[0], ids[2]]) assert [x.text for x in nodes] == ["Doc 1", "Doc 3"] @@ -32,95 +40,75 @@ def test_e2e(): nodes = vs.get_nodes(node_ids=["a"]) assert len(nodes) == 0 - assert [ - x.text - for x in vs.get_nodes( - filters=MetadataFilters( - filters=[ - MetadataFilter(key="a", value="2"), - ] - ) - ) - ] == ["Doc 2"] - - assert [ - x.text - for x in vs.get_nodes( - filters=MetadataFilters( - filters=[ - MetadataFilter(key="a", value="2"), - MetadataFilter(key="a", value="3"), - ] - ) + nodes = vs.get_nodes( + filters=MetadataFilters(filters=[MetadataFilter(key="a", value="2")]) + ) + assert [x.text for x in nodes] == ["Doc 2"] + + nodes = vs.get_nodes( + filters=MetadataFilters( + filters=[ + MetadataFilter(key="a", value="2"), + MetadataFilter(key="a", value="3"), + ] ) - ] == [] - - assert [ - x.text - for x in vs.get_nodes( - filters=MetadataFilters( - condition=FilterCondition.OR, - filters=[ - MetadataFilter(key="a", value="2"), - MetadataFilter(key="a", value="3"), - ], - ) + ) + assert [x.text for x in nodes] == [] + + nodes = vs.get_nodes( + filters=MetadataFilters( + condition=FilterCondition.OR, + filters=[ + MetadataFilter(key="a", value="2"), + MetadataFilter(key="a", value="3"), + ], ) - ] == ["Doc 2", "Doc 3"] - - assert [ - x.text - for x in vs.get_nodes( - filters=MetadataFilters( - condition=FilterCondition.OR, - filters=[ - MetadataFilter(key="a", value="2"), - MetadataFilter(key="a", value="3"), - ], - ) + ) + assert [x.text for x in nodes] == ["Doc 2", "Doc 3"] + + nodes = vs.get_nodes( + filters=MetadataFilters( + condition=FilterCondition.OR, + filters=[ + MetadataFilter(key="a", value="2"), + MetadataFilter(key="a", value="3"), + ], ) - ] == ["Doc 2", "Doc 3"] - - assert [ - x.text - for x in vs.get_nodes( - filters=MetadataFilters( - filters=[ - MetadataFilter(key="b", value=10, operator=FilterOperator.GT), - ] - ) + ) + assert [x.text for x in nodes] == ["Doc 2", "Doc 3"] + + nodes = vs.get_nodes( + filters=MetadataFilters( + filters=[ + MetadataFilter(key="b", value=10, operator=FilterOperator.GT), + ] ) - ] == ["Doc 2", "Doc 3"] - - assert [ - x.text - for x in vs.get_nodes( - filters=MetadataFilters( - filters=[ - MetadataFilter(key="b", value=11, operator=FilterOperator.LTE), - ] - ) + ) + assert [x.text for x in nodes] == ["Doc 2", "Doc 3"] + + nodes = vs.get_nodes( + filters=MetadataFilters( + filters=[ + MetadataFilter(key="b", value=11, operator=FilterOperator.LTE), + ] ) - ] == ["Doc 1", "Doc 2"] + ) + assert [x.text for x in nodes] == ["Doc 1", "Doc 2"] + +def test_delete_id(vs_ids): + vs, ids = vs_ids vs.delete_nodes(node_ids=[ids[0], ids[2]]) assert [x.text for x in vs.get_nodes()] == ["Doc 2"] - vs.add( - nodes=[ - Document(text="Doc 4", embedding=[1, 2, 1], metadata={"a": "4", "b": 14}), - Document(text="Doc 5", embedding=[1, 2, 2], metadata={"a": "5", "b": 15}), - Document(text="Doc 6", embedding=[1, 2, 3], metadata={"a": "6", "b": 16}), - ] - ) +def test_delete_filter(vs_ids): + vs, ids = vs_ids vs.delete_nodes( filters=MetadataFilters( filters=[ - MetadataFilter(key="b", value=14, operator=FilterOperator.GT), + MetadataFilter(key="b", value=10, operator=FilterOperator.GT), ] ) ) - assert [x.text for x in vs.get_nodes()] == ["Doc 2", "Doc 4"] - - vs.clear() + assert [x.text for x in vs.get_nodes()] == ["Doc 1"]