Adds "NIN" metadata filter for pgvector to all checking for set absence (#14205)

This PR adds support for metadata filters of the form:

`{"filter": {"key": { "NIN" : ["list", "of", "values"]}}}`

"IN" is already supported, so this is a quick & related update to add
"NIN"
pull/13138/head^2
jeffpezzone 6 months ago committed by GitHub
parent 20d2b4a6ba
commit 7c2ef06136
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -434,16 +434,24 @@ class PGVector(VectorStore):
if filter is not None:
filter_clauses = []
IN, NIN = "in", "nin"
for key, value in filter.items():
IN = "in"
if isinstance(value, dict) and IN in map(str.lower, value):
if isinstance(value, dict):
value_case_insensitive = {
k.lower(): v for k, v in value.items()
}
filter_by_metadata = self.EmbeddingStore.cmetadata[
key
].astext.in_(value_case_insensitive[IN])
filter_clauses.append(filter_by_metadata)
if IN in map(str.lower, value):
filter_by_metadata = self.EmbeddingStore.cmetadata[
key
].astext.in_(value_case_insensitive[IN])
elif NIN in map(str.lower, value):
filter_by_metadata = self.EmbeddingStore.cmetadata[
key
].astext.not_in(value_case_insensitive[NIN])
else:
filter_by_metadata = None
if filter_by_metadata is not None:
filter_clauses.append(filter_by_metadata)
else:
filter_by_metadata = self.EmbeddingStore.cmetadata[
key

@ -17,7 +17,6 @@ CONNECTION_STRING = PGVector.connection_string_from_db_params(
password=os.environ.get("TEST_PGVECTOR_PASSWORD", "postgres"),
)
ADA_TOKEN_COUNT = 1536
@ -186,6 +185,27 @@ def test_pgvector_with_filter_in_set() -> None:
]
def test_pgvector_with_filter_nin_set() -> None:
"""Test end to end construction and search."""
texts = ["foo", "bar", "baz"]
metadatas = [{"page": str(i)} for i in range(len(texts))]
docsearch = PGVector.from_texts(
texts=texts,
collection_name="test_collection_filter",
embedding=FakeEmbeddingsWithAdaDimension(),
metadatas=metadatas,
connection_string=CONNECTION_STRING,
pre_delete_collection=True,
)
output = docsearch.similarity_search_with_score(
"foo", k=2, filter={"page": {"NIN": ["1"]}}
)
assert output == [
(Document(page_content="foo", metadata={"page": "0"}), 0.0),
(Document(page_content="baz", metadata={"page": "2"}), 0.0013003906671379406),
]
def test_pgvector_delete_docs() -> None:
"""Add and delete documents."""
texts = ["foo", "bar", "baz"]

Loading…
Cancel
Save