Skip to content

Commit

Permalink
Adds "NIN" metadata filter for pgvector to all checking for set absen…
Browse files Browse the repository at this point in the history
…ce (#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"
  • Loading branch information
jeffpezzone authored Dec 6, 2023
1 parent 20d2b4a commit 7c2ef06
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 7 deletions.
20 changes: 14 additions & 6 deletions libs/langchain/langchain/vectorstores/pgvector.py
Original file line number Diff line number Diff line change
Expand Up @@ -434,16 +434,24 @@ def __query_collection(

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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
password=os.environ.get("TEST_PGVECTOR_PASSWORD", "postgres"),
)


ADA_TOKEN_COUNT = 1536


Expand Down Expand Up @@ -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"]
Expand Down

0 comments on commit 7c2ef06

Please sign in to comment.