Apache
/
nifi
mirror of https://github.com/apache/nifi.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

NIFI-12629 - adding metadata filtering to QueryPinecone (#8264) 

* NIFI-12629 - adding metadata filtering to QueryPinecone
This commit is contained in:
Pierre Villard 2024-01-19 22:03:38 +04:00 committed by GitHub
parent 6649115a34
commit f402970132
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 11 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
nifi-python-extensions/nifi-text-embeddings-module/src/main/python/vectorstores/QueryPinecone.py
View File

@ -18,6 +18,7 @@ from nifiapi.flowfiletransform import FlowFileTransform, FlowFileTransformResult
from nifiapi.properties import PropertyDescriptor, StandardValidators, ExpressionLanguageScope, PropertyDependency
import QueryUtils
import pinecone
import json
from EmbeddingUtils import OPENAI, HUGGING_FACE, EMBEDDING_MODEL, create_embedding_service
@ -109,7 +110,14 @@ class QueryPinecone(FlowFileTransform):
)
NAMESPACE = PropertyDescriptor(
name="Namespace",
description="The name of the Pinecone Namespace to put the documents to.",
description="The name of the Pinecone Namespace to query into.",
required=False,
validators=[StandardValidators.NON_EMPTY_VALIDATOR],
expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES
)
FILTER = PropertyDescriptor(
name="Metadata Filter",
description="Optional metadata filter to apply with the query. For example: { \"author\": {\"$eq\": \"john.doe\"} }",
required=False,
validators=[StandardValidators.NON_EMPTY_VALIDATOR],
expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES
@ -164,8 +172,9 @@ class QueryPinecone(FlowFileTransform):
index = pinecone.Index(index_name)
text_key = context.getProperty(self.TEXT_KEY).evaluateAttributeExpressions().getValue()
filter = context.getProperty(self.FILTER).evaluateAttributeExpressions(flowfile).getValue()
vectorstore = Pinecone(index, self.embeddings.embed_query, text_key, namespace=namespace)
results = vectorstore.similarity_search_with_score(query, num_results)
results = vectorstore.similarity_search_with_score(query, num_results, filter=None if filter is None else json.loads(filter))
documents = []
for result in results: