This is an automated email from the ASF dual-hosted git repository.
joewitt pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi-python-extensions.git
The following commit(s) were added to refs/heads/main by this push:
new 26adfa6 NIFI-13448 Added initial Python Processors This closes #3
26adfa6 is described below
commit 26adfa6767416be7296f3c33b90f44561b6f257f
Author: exceptionfactory <exceptionfact...@apache.org>
AuthorDate: Wed Jun 26 13:01:11 2024 -0500
NIFI-13448 Added initial Python Processors
This closes #3
- Added pyproject.toml with source bundles directory
- Added Check Formatting step with Hatch
- Added Build Distribution step with Hatch
Signed-off-by: David Handermann <exceptionfact...@apache.org>
NIFI-13448 Excluded GitHub and Check Licenses
---
.github/workflows/build.yml | 4 +
.ratignore | 1 +
README.md | 20 ++
pyproject.toml | 74 +++++
src/__about__.py | 3 +
src/__init__.py | 1 +
src/extensions/chunking/ChunkDocument.py | 277 ++++++++++++++++++
src/extensions/chunking/ParseDocument.py | 314 +++++++++++++++++++++
src/extensions/chunking/__init__.py | 1 +
src/extensions/openai/PromptChatGPT.py | 219 ++++++++++++++
src/extensions/openai/__init__.py | 1 +
src/extensions/vectorstores/ChromaUtils.py | 138 +++++++++
src/extensions/vectorstores/EmbeddingUtils.py | 165 +++++++++++
.../vectorstores/OpenSearchVectorUtils.py | 130 +++++++++
src/extensions/vectorstores/PutChroma.py | 128 +++++++++
src/extensions/vectorstores/PutOpenSearchVector.py | 267 ++++++++++++++++++
src/extensions/vectorstores/PutPinecone.py | 220 +++++++++++++++
src/extensions/vectorstores/PutQdrant.py | 151 ++++++++++
src/extensions/vectorstores/QdrantUtils.py | 112 ++++++++
src/extensions/vectorstores/QueryChroma.py | 181 ++++++++++++
.../vectorstores/QueryOpenSearchVector.py | 242 ++++++++++++++++
src/extensions/vectorstores/QueryPinecone.py | 205 ++++++++++++++
src/extensions/vectorstores/QueryQdrant.py | 161 +++++++++++
src/extensions/vectorstores/QueryUtils.py | 171 +++++++++++
src/extensions/vectorstores/__init__.py | 1 +
src/extensions/vectorstores/requirements.txt | 24 ++
26 files changed, 3211 insertions(+)
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 15b5e4a..3ba97bb 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -43,3 +43,7 @@ jobs:
run: |
python -m pip install --upgrade pip
pip install hatch
+ - name: Check Formatting
+ run: hatch fmt --check
+ - name: Build Distribution
+ run: hatch build
diff --git a/.ratignore b/.ratignore
index af3daed..3ca979f 100644
--- a/.ratignore
+++ b/.ratignore
@@ -14,3 +14,4 @@ share/python-wheels/*
.idea/*
.git/*
.cache/*
+.ruff_cache/*
diff --git a/README.md b/README.md
index 9ff87aa..85e459c 100644
--- a/README.md
+++ b/README.md
@@ -2,15 +2,35 @@
[](https://github.com/apache/nifi-python-extensions/blob/main/LICENSE)
[](https://github.com/apache/nifi-python-extensions/actions/workflows/build.yml)
+[](https://github.com/pypa/hatch)
+[](https://github.com/astral-sh/ruff)
The [Apache NiFi](https://nifi.apache.org) Python Extensions repository
contains Processors implemented in [Python](https://www.python.org/)
for deployment in Apache NiFi 2.
+## Building
+
+This project uses [Hatch](https://hatch.pypa.io) to build distribution
packages.
+
+```
+hatch build
+```
+
+The build command creates a source distribution in the `dist` directory.
+
+The source distribution contains an `extensions` directory can be copied into
Apache NiFi to use the packaged Processors.
+
## Developing
The Apache NiFi [Python Developer's
Guide](https://nifi.apache.org/documentation/nifi-2.0.0-M3/html/python-developer-guide.html)
provides the API and implementation guidelines for Python Processors.
+The Hatch format command supports evaluating Python Processors against
configured rules.
+
+```
+hatch fmt --check
+```
+
## Documentation
The Apache NiFi [Documentation](https://nifi.apache.org/documentation/)
includes reference information for project capabilities.
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..f47f4e3
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "nifi-python-extensions"
+dynamic = ["version"]
+description = "Apache NiFi Processors implemented in Python"
+requires-python = ">=3.11"
+keywords = ["apache", "nifi", "extensions", "processors"]
+readme = "README.md"
+authors = [
+ { name = "Apache NiFi Developers", email = "d...@nifi.apache.org" },
+]
+maintainers = [
+ { name = "Apache NiFi Developers", email = "d...@nifi.apache.org" },
+]
+classifiers = [
+ "Development Status :: 5 - Production/Stable",
+ "License :: OSI Approved :: Apache Software License",
+ "Intended Audience :: Developers",
+ "Programming Language :: Python",
+ "Programming Language :: Python :: 3.11",
+ "Programming Language :: Python :: 3.12",
+ "Framework :: Hatch",
+]
+
+[project.urls]
+Homepage = "https://nifi.apache.org";
+Issues = "https://issues.apache.org/jira/projects/NIFI/issues";
+Source = "https://github.com/apache/nifi-python-extensions";
+
+[tool.hatch.version]
+path = "src/__about__.py"
+
+[[tool.hatch.envs.all.matrix]]
+python = ["3.11", "3.12"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/extensions"]
+
+[tool.hatch.build.targets.sdist]
+exclude = [
+ ".asf.yaml",
+ ".github",
+ ".ratignore",
+ "check-licenses.sh",
+]
+
+[tool.ruff]
+preview = true
+lint.pep8-naming.extend-ignore-names = [
+ "flowFile",
+ "getPropertyDescriptors",
+ "onScheduled",
+]
+lint.flake8-self.extend-ignore-names = [
+ "_standard_validators"
+]
+lint.extend-select = [
+ "CPY001"
+]
+lint.ignore = [
+ "G004", # Allow f-string for logging
+ "N999", # Allow Processor module names that do not follow pep8-naming
+ "PERF401", # Allow manual list comprehension
+ "RUF012", # Allow mutable class attributes without typing.ClassVar
+ "S105", # Avoid checking for hardcoded-password-string values
+]
+
+[tool.ruff.lint.flake8-copyright]
+notice-rgx = "# SPDX-License-Identifier: Apache-2.0\n"
diff --git a/src/__about__.py b/src/__about__.py
new file mode 100644
index 0000000..c442147
--- /dev/null
+++ b/src/__about__.py
@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+
+__version__ = "2.0.0.dev0"
diff --git a/src/__init__.py b/src/__init__.py
new file mode 100644
index 0000000..9881313
--- /dev/null
+++ b/src/__init__.py
@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0
diff --git a/src/extensions/chunking/ChunkDocument.py
b/src/extensions/chunking/ChunkDocument.py
new file mode 100644
index 0000000..576e914
--- /dev/null
+++ b/src/extensions/chunking/ChunkDocument.py
@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+from langchain.text_splitter import Language
+from nifiapi.documentation import ProcessorConfiguration,
multi_processor_use_case, use_case
+from nifiapi.flowfiletransform import FlowFileTransform,
FlowFileTransformResult
+from nifiapi.properties import ExpressionLanguageScope, PropertyDependency,
PropertyDescriptor, StandardValidators
+
+SPLIT_BY_CHARACTER = "Split by Character"
+SPLIT_CODE = "Split Code"
+RECURSIVELY_SPLIT_BY_CHARACTER = "Recursively Split by Character"
+
+TEXT_KEY = "text"
+METADATA_KEY = "metadata"
+
+
+@use_case(
+ description="Create chunks of text from a single larger chunk.",
+ notes="The input for this use case is expected to be a FlowFile whose
content is a JSON Lines document, with each line having a 'text' and a
'metadata' element.",
+ keywords=["embedding", "vector", "text", "rag", "retrieval augmented
generation"],
+ configuration="""
+ Set "Input Format" to "Plain Text"
+ Set "Element Strategy" to "Single Document"
+ """,
+)
+@multi_processor_use_case(
+ description="""
+ Chunk Plaintext data in order to prepare it for storage in a vector
store. The output is in "json-lines" format,
+ containing the chunked data as text, as well as metadata pertaining to
the chunk.""",
+ notes="The input for this use case is expected to be a FlowFile whose
content is a plaintext document.",
+ keywords=["embedding", "vector", "text", "rag", "retrieval augmented
generation"],
+ configurations=[
+ ProcessorConfiguration(
+ processor_type="ParseDocument",
+ configuration="""
+ Set "Input Format" to "Plain Text"
+ Set "Element Strategy" to "Single Document"
+
+ Connect the 'success' Relationship to ChunkDocument.
+ """,
+ ),
+ ProcessorConfiguration(
+ processor_type="ChunkDocument",
+ configuration="""
+ Set the following properties:
+ "Chunking Strategy" = "Recursively Split by Character"
+ "Separator" = "\\n\\n,\\n, ,"
+ "Separator Format" = "Plain Text"
+ "Chunk Size" = "4000"
+ "Chunk Overlap" = "200"
+ "Keep Separator" = "false"
+
+ Connect the 'success' Relationship to the appropriate
destination to store data in the desired vector store.
+ """,
+ ),
+ ],
+)
+@multi_processor_use_case(
+ description="""
+ Parse and chunk the textual contents of a PDF document in order to
prepare it for storage in a vector store. The output is in "json-lines" format,
+ containing the chunked data as text, as well as metadata pertaining to
the chunk.""",
+ notes="The input for this use case is expected to be a FlowFile whose
content is a PDF document.",
+ keywords=["pdf", "embedding", "vector", "text", "rag", "retrieval
augmented generation"],
+ configurations=[
+ ProcessorConfiguration(
+ processor_type="ParseDocument",
+ configuration="""
+ Set "Input Format" to "PDF"
+ Set "Element Strategy" to "Single Document"
+ Set "Include Extracted Metadata" to "false"
+
+ Connect the 'success' Relationship to ChunkDocument.
+ """,
+ ),
+ ProcessorConfiguration(
+ processor_type="ChunkDocument",
+ configuration="""
+ Set the following properties:
+ "Chunking Strategy" = "Recursively Split by Character"
+ "Separator" = "\\n\\n,\\n, ,"
+ "Separator Format" = "Plain Text"
+ "Chunk Size" = "4000"
+ "Chunk Overlap" = "200"
+ "Keep Separator" = "false"
+
+ Connect the 'success' Relationship to the appropriate
destination to store data in the desired vector store.
+ """,
+ ),
+ ],
+)
+class ChunkDocument(FlowFileTransform):
+ class Java:
+ implements = ["org.apache.nifi.python.processor.FlowFileTransform"]
+
+ class ProcessorDetails:
+ version = "2.0.0.dev0"
+ description = """Chunks incoming documents that are formatted as JSON
Lines into chunks that are appropriately sized for creating Text Embeddings.
+ The input is expected to be in "json-lines" format, with each line
having a 'text' and a 'metadata' element.
+ Each line will then be split into one or more lines in the
output."""
+ tags = [
+ "text",
+ "split",
+ "chunk",
+ "langchain",
+ "embeddings",
+ "vector",
+ "machine learning",
+ "ML",
+ "artificial intelligence",
+ "ai",
+ "document",
+ ]
+ dependencies = ["langchain"]
+
+ CHUNK_STRATEGY = PropertyDescriptor(
+ name="Chunking Strategy",
+ description="Specifies which splitter should be used to split the
text",
+ allowable_values=[RECURSIVELY_SPLIT_BY_CHARACTER, SPLIT_BY_CHARACTER,
SPLIT_CODE],
+ required=True,
+ default_value=RECURSIVELY_SPLIT_BY_CHARACTER,
+ )
+ SEPARATOR = PropertyDescriptor(
+ name="Separator",
+ description="""Specifies the character sequence to use for splitting
apart the text. If using a Chunking Strategy of Recursively Split by Character,
+ it is a comma-separated list of character sequences.
Meta-characters \\n, \\r and \\t are automatically un-escaped.""",
+ required=True,
+ default_value="\\n\\n,\\n, ,",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ dependencies=[PropertyDependency(CHUNK_STRATEGY, SPLIT_BY_CHARACTER,
RECURSIVELY_SPLIT_BY_CHARACTER)],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ SEPARATOR_FORMAT = PropertyDescriptor(
+ name="Separator Format",
+ description="Specifies how to interpret the value of the <Separator>
property",
+ required=True,
+ default_value="Plain Text",
+ allowable_values=["Plain Text", "Regular Expression"],
+ dependencies=[PropertyDependency(CHUNK_STRATEGY, SPLIT_BY_CHARACTER,
RECURSIVELY_SPLIT_BY_CHARACTER)],
+ )
+ CHUNK_SIZE = PropertyDescriptor(
+ name="Chunk Size",
+ description="The maximum size of a chunk that should be returned",
+ required=True,
+ default_value="4000",
+ validators=[StandardValidators.POSITIVE_INTEGER_VALIDATOR],
+ )
+ CHUNK_OVERLAP = PropertyDescriptor(
+ name="Chunk Overlap",
+ description="The number of characters that should be overlapped
between each chunk of text",
+ required=True,
+ default_value="200",
+ validators=[StandardValidators.NON_NEGATIVE_INTEGER_VALIDATOR],
+ )
+ KEEP_SEPARATOR = PropertyDescriptor(
+ name="Keep Separator",
+ description="Whether or not to keep the text separator in each chunk
of data",
+ required=True,
+ default_value="false",
+ allowable_values=["true", "false"],
+ dependencies=[PropertyDependency(CHUNK_STRATEGY, SPLIT_BY_CHARACTER,
RECURSIVELY_SPLIT_BY_CHARACTER)],
+ )
+ STRIP_WHITESPACE = PropertyDescriptor(
+ name="Strip Whitespace",
+ description="Whether or not to strip the whitespace at the beginning
and end of each chunk",
+ required=True,
+ default_value="true",
+ allowable_values=["true", "false"],
+ dependencies=[PropertyDependency(CHUNK_STRATEGY, SPLIT_BY_CHARACTER,
RECURSIVELY_SPLIT_BY_CHARACTER)],
+ )
+ LANGUAGE = PropertyDescriptor(
+ name="Language",
+ description="The language to use for the Code's syntax",
+ required=True,
+ default_value="python",
+ allowable_values=[e.value for e in Language],
+ dependencies=[PropertyDependency(CHUNK_STRATEGY, SPLIT_CODE)],
+ )
+
+ property_descriptors = [
+ CHUNK_STRATEGY,
+ SEPARATOR,
+ SEPARATOR_FORMAT,
+ CHUNK_SIZE,
+ CHUNK_OVERLAP,
+ KEEP_SEPARATOR,
+ STRIP_WHITESPACE,
+ LANGUAGE,
+ ]
+
+ def __init__(self, **kwargs):
+ pass
+
+ def getPropertyDescriptors(self):
+ return self.property_descriptors
+
+ def split_docs(self, context, flowfile, documents):
+ from langchain.text_splitter import CharacterTextSplitter,
RecursiveCharacterTextSplitter
+
+ strategy = context.getProperty(self.CHUNK_STRATEGY).getValue()
+ if strategy == SPLIT_BY_CHARACTER:
+ text_splitter = CharacterTextSplitter(
+
separator=context.getProperty(self.SEPARATOR).evaluateAttributeExpressions(flowfile).getValue(),
+
keep_separator=context.getProperty(self.KEEP_SEPARATOR).asBoolean(),
+
is_separator_regex=context.getProperty(self.SEPARATOR_FORMAT).getValue() ==
"Regular Expression",
+ chunk_size=context.getProperty(self.CHUNK_SIZE).asInteger(),
+
chunk_overlap=context.getProperty(self.CHUNK_OVERLAP).asInteger(),
+ length_function=len,
+
strip_whitespace=context.getProperty(self.STRIP_WHITESPACE).asBoolean(),
+ )
+ elif strategy == SPLIT_CODE:
+ text_splitter = RecursiveCharacterTextSplitter.from_language(
+ language=context.getProperty(self.LANGUAGE).getValue(),
+ chunk_size=context.getProperty(self.CHUNK_SIZE).asInteger(),
+
chunk_overlap=context.getProperty(self.CHUNK_OVERLAP).asInteger(),
+ )
+ else:
+ separator_text =
context.getProperty(self.SEPARATOR).evaluateAttributeExpressions(flowfile).getValue()
+ splits = separator_text.split(",")
+ unescaped = []
+ for split in splits:
+ unescaped.append(split.replace("\\n", "\n").replace("\\r",
"\r").replace("\\t", "\t"))
+ text_splitter = RecursiveCharacterTextSplitter(
+ separators=unescaped,
+
keep_separator=context.getProperty(self.KEEP_SEPARATOR).asBoolean(),
+
is_separator_regex=context.getProperty(self.SEPARATOR_FORMAT).getValue() ==
"Regular Expression",
+ chunk_size=context.getProperty(self.CHUNK_SIZE).asInteger(),
+
chunk_overlap=context.getProperty(self.CHUNK_OVERLAP).asInteger(),
+ length_function=len,
+
strip_whitespace=context.getProperty(self.STRIP_WHITESPACE).asBoolean(),
+ )
+
+ return text_splitter.split_documents(documents)
+
+ def to_json(self, docs) -> str:
+ json_docs = []
+
+ for i, doc in enumerate(docs):
+ doc.metadata["chunk_index"] = i
+ doc.metadata["chunk_count"] = len(docs)
+
+ json_doc = json.dumps({TEXT_KEY: doc.page_content, METADATA_KEY:
doc.metadata})
+ json_docs.append(json_doc)
+
+ return "\n".join(json_docs)
+
+ def load_docs(self, flowfile):
+ from langchain.schema import Document
+
+ flowfile_contents = flowfile.getContentsAsBytes().decode()
+ docs = []
+ for line in flowfile_contents.split("\n"):
+ stripped = line.strip()
+ if stripped == "":
+ continue
+
+ json_element = json.loads(stripped)
+ page_content = json_element.get(TEXT_KEY)
+ if page_content is None:
+ continue
+
+ metadata = json_element.get(METADATA_KEY)
+ if metadata is None:
+ metadata = {}
+
+ doc = Document(page_content=page_content, metadata=metadata)
+ docs.append(doc)
+
+ return docs
+
+ def transform(self, context, flowfile):
+ documents = self.load_docs(flowfile)
+ split_docs = self.split_docs(context, flowfile, documents)
+
+ output_json = self.to_json(split_docs)
+ attributes = {"document.count": str(len(split_docs))}
+ return FlowFileTransformResult("success", contents=output_json,
attributes=attributes)
diff --git a/src/extensions/chunking/ParseDocument.py
b/src/extensions/chunking/ParseDocument.py
new file mode 100644
index 0000000..2a2d05e
--- /dev/null
+++ b/src/extensions/chunking/ParseDocument.py
@@ -0,0 +1,314 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import io
+import json
+
+from nifiapi.flowfiletransform import FlowFileTransform,
FlowFileTransformResult
+from nifiapi.properties import PropertyDependency, PropertyDescriptor,
StandardValidators
+
+PLAIN_TEXT = "Plain Text"
+HTML = "HTML"
+MARKDOWN = "Markdown"
+PDF = "PDF"
+EXCEL = "Microsoft Excel"
+POWERPOINT = "Microsoft PowerPoint"
+WORD = "Microsoft Word"
+
+PARSING_STRATEGY_AUTO = "Automatic"
+PARSING_STRATEGY_HIGH_RES = "High Resolution"
+PARSING_STRATEGY_OCR_ONLY = "OCR Only"
+PARSING_STRATEGY_FAST = "Fast"
+
+SINGLE_DOCUMENT = "Single Document"
+DOCUMENT_PER_ELEMENT = "Document Per Element"
+
+TEXT_KEY = "text"
+METADATA_KEY = "metadata"
+
+
+class ParseDocument(FlowFileTransform):
+ class Java:
+ implements = ["org.apache.nifi.python.processor.FlowFileTransform"]
+
+ class ProcessorDetails:
+ version = "2.0.0.dev0"
+ description = """Parses incoming unstructured text documents and
performs optical character recognition (OCR) in order to extract text from PDF
and image files.
+ The output is formatted as "json-lines" with two keys: 'text' and
'metadata'.
+ Note that use of this Processor may require significant storage
space and RAM utilization due to third-party dependencies necessary for
processing PDF and image files.
+ Also note that in order to process PDF or Images, Tesseract and
Poppler must be installed on the system."""
+ tags = [
+ "text",
+ "embeddings",
+ "vector",
+ "machine learning",
+ "ML",
+ "artificial intelligence",
+ "ai",
+ "document",
+ "langchain",
+ "pdf",
+ "html",
+ "markdown",
+ "word",
+ "excel",
+ "powerpoint",
+ ]
+ dependencies = [
+ "pikepdf==8.12.0",
+ "pypdf==4.0.1",
+ "langchain==0.1.7",
+ "unstructured==0.14.8",
+ "unstructured-inference==0.7.36",
+ "unstructured_pytesseract==0.3.12",
+ "pillow-heif==0.15.0",
+ "numpy==1.26.4",
+ "opencv-python==4.9.0.80",
+ "pdf2image==1.17.0",
+ "pdfminer.six==20221105",
+ "python-docx==1.1.0",
+ "openpyxl==3.1.2",
+ "python-pptx==0.6.23",
+ ]
+
+ INPUT_FORMAT = PropertyDescriptor(
+ name="Input Format",
+ description="""The format of the input FlowFile. This dictates which
TextLoader will be used to parse the input.
+ Note that in order to process images or extract tables from PDF
files,you must have both 'poppler' and 'tesseract' installed on your system.""",
+ allowable_values=[PLAIN_TEXT, HTML, MARKDOWN, PDF, WORD, EXCEL,
POWERPOINT],
+ required=True,
+ default_value=PLAIN_TEXT,
+ )
+ PDF_PARSING_STRATEGY = PropertyDescriptor(
+ name="PDF Parsing Strategy",
+ display_name="Parsing Strategy",
+ description="Specifies the strategy to use when parsing a PDF",
+ allowable_values=[
+ PARSING_STRATEGY_AUTO,
+ PARSING_STRATEGY_HIGH_RES,
+ PARSING_STRATEGY_OCR_ONLY,
+ PARSING_STRATEGY_FAST,
+ ],
+ required=True,
+ default_value=PARSING_STRATEGY_AUTO,
+ dependencies=[PropertyDependency(INPUT_FORMAT, PDF)],
+ )
+ PDF_MODEL_NAME = PropertyDescriptor(
+ name="PDF Parsing Model",
+ description="The model to use for parsing. Different models will have
their own strengths and weaknesses.",
+ allowable_values=["yolox", "detectron2_onnx", "chipper"],
+ required=True,
+ default_value="yolox",
+ dependencies=[PropertyDependency(INPUT_FORMAT, PDF)],
+ )
+ ELEMENT_STRATEGY = PropertyDescriptor(
+ name="Element Strategy",
+ description="Specifies whether the input should be loaded as a single
Document, or if each element in the input should be separated out into its own
Document",
+ allowable_values=[SINGLE_DOCUMENT, DOCUMENT_PER_ELEMENT],
+ required=True,
+ default_value=DOCUMENT_PER_ELEMENT,
+ dependencies=[PropertyDependency(INPUT_FORMAT, HTML, MARKDOWN)],
+ )
+ INCLUDE_PAGE_BREAKS = PropertyDescriptor(
+ name="Include Page Breaks",
+ description="Specifies whether or not page breaks should be considered
when creating Documents from the input",
+ allowable_values=["true", "false"],
+ required=True,
+ default_value="false",
+ dependencies=[
+ PropertyDependency(INPUT_FORMAT, HTML, MARKDOWN),
+ PropertyDependency(ELEMENT_STRATEGY, DOCUMENT_PER_ELEMENT),
+ ],
+ )
+ PDF_INFER_TABLE_STRUCTURE = PropertyDescriptor(
+ name="Infer Table Structure",
+ description="If true, any table that is identified in the PDF will be
parsed and translated into an HTML structure. The HTML of that table will then
be added to the \
+ Document's metadata in a key named 'text_as_html'.
Regardless of the value of this property, the textual contents of the table
will be written to the contents \
+ without the structure.",
+ allowable_values=["true", "false"],
+ default_value="false",
+ required=True,
+ dependencies=[PropertyDependency(PDF_PARSING_STRATEGY,
PARSING_STRATEGY_HIGH_RES)],
+ )
+ LANGUAGES = PropertyDescriptor(
+ name="Languages",
+ description="A comma-separated list of language codes that should be
used when using OCR to determine the text.",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ default_value="Eng",
+ required=True,
+ dependencies=[PropertyDependency(INPUT_FORMAT, PDF)],
+ )
+ METADATA_FIELDS = PropertyDescriptor(
+ name="Metadata Fields",
+ description="A comma-separated list of FlowFile attributes that will
be added to the Documents' Metadata",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ default_value="filename, uuid",
+ required=True,
+ )
+ EXTRACT_METADATA = PropertyDescriptor(
+ name="Include Extracted Metadata",
+ description="Whether or not to include the metadata that is extracted
from the input in each of the Documents",
+ allowable_values=["true", "false"],
+ default_value="true",
+ required=True,
+ )
+
+ property_descriptors = [
+ INPUT_FORMAT,
+ PDF_PARSING_STRATEGY,
+ PDF_MODEL_NAME,
+ ELEMENT_STRATEGY,
+ INCLUDE_PAGE_BREAKS,
+ PDF_INFER_TABLE_STRUCTURE,
+ LANGUAGES,
+ METADATA_FIELDS,
+ EXTRACT_METADATA,
+ ]
+
+ def __init__(self, **kwargs):
+ pass
+
+ def getPropertyDescriptors(self):
+ return self.property_descriptors
+
+ def get_parsing_strategy(self, nifi_value: str, default_value: str) -> str:
+ if nifi_value == PARSING_STRATEGY_OCR_ONLY:
+ return "ocr_only"
+ if nifi_value == PARSING_STRATEGY_HIGH_RES:
+ return "hi_res"
+ if nifi_value == PARSING_STRATEGY_FAST:
+ return "fast"
+ if nifi_value == PARSING_STRATEGY_AUTO:
+ return "auto"
+ return default_value
+
+ def get_languages(self, nifi_value: str) -> list[str]:
+ return [lang.strip() for lang in nifi_value.split(",")]
+
+ def create_docs(self, context, flowFile):
+ from langchain.schema import Document
+
+ metadata = {}
+
+ for attribute_name in
context.getProperty(self.METADATA_FIELDS).getValue().split(","):
+ trimmed = attribute_name.strip()
+ value = flowFile.getAttribute(trimmed)
+ metadata[trimmed] = value
+
+ input_format =
context.getProperty(self.INPUT_FORMAT).evaluateAttributeExpressions(flowFile).getValue()
+ if input_format == PLAIN_TEXT:
+ return
[Document(page_content=flowFile.getContentsAsBytes().decode("utf-8"),
metadata=metadata)]
+
+ element_strategy =
context.getProperty(self.ELEMENT_STRATEGY).getValue()
+ mode = "single" if element_strategy == SINGLE_DOCUMENT else "elements"
+
+ include_page_breaks =
context.getProperty(self.INCLUDE_PAGE_BREAKS).asBoolean()
+ include_metadata =
context.getProperty(self.EXTRACT_METADATA).asBoolean()
+
+ if input_format == HTML:
+ from langchain.document_loaders import UnstructuredHTMLLoader
+
+ loader = UnstructuredHTMLLoader(
+ None,
+ file=io.BytesIO(flowFile.getContentsAsBytes()),
+ mode=mode,
+ include_page_breaks=include_page_breaks,
+ include_metadata=include_metadata,
+ )
+
+ elif input_format == PDF:
+ from langchain.document_loaders import UnstructuredPDFLoader
+
+ infer_table_structure =
context.getProperty(self.PDF_INFER_TABLE_STRUCTURE).asBoolean()
+ strategy = self.get_parsing_strategy(
+ context.getProperty(self.PDF_PARSING_STRATEGY).getValue(),
PARSING_STRATEGY_AUTO
+ )
+ languages =
self.get_languages(context.getProperty(self.LANGUAGES).getValue())
+ model_name = context.getProperty(self.PDF_MODEL_NAME).getValue()
+
+ loader = UnstructuredPDFLoader(
+ None,
+ file=io.BytesIO(flowFile.getContentsAsBytes()),
+ mode=mode,
+ infer_table_structure=infer_table_structure,
+ include_page_breaks=include_page_breaks,
+ languages=languages,
+ strategy=strategy,
+ include_metadata=include_metadata,
+ model_name=model_name,
+ )
+
+ elif input_format == MARKDOWN:
+ from langchain.document_loaders import UnstructuredMarkdownLoader
+
+ loader = UnstructuredMarkdownLoader(
+ None,
+ file=io.BytesIO(flowFile.getContentsAsBytes()),
+ mode=mode,
+ include_page_breaks=include_page_breaks,
+ include_metadata=include_metadata,
+ )
+
+ elif input_format == WORD:
+ from langchain.document_loaders import
UnstructuredWordDocumentLoader
+
+ loader = UnstructuredWordDocumentLoader(
+ None,
+ file=io.BytesIO(flowFile.getContentsAsBytes()),
+ mode=mode,
+ include_page_breaks=include_page_breaks,
+ include_metadata=include_metadata,
+ )
+
+ elif input_format == EXCEL:
+ from langchain.document_loaders import UnstructuredExcelLoader
+
+ loader = UnstructuredExcelLoader(
+ None,
+ file=io.BytesIO(flowFile.getContentsAsBytes()),
+ mode=mode,
+ include_page_breaks=include_page_breaks,
+ include_metadata=include_metadata,
+ )
+
+ elif input_format == POWERPOINT:
+ from langchain.document_loaders import UnstructuredPowerPointLoader
+
+ loader = UnstructuredPowerPointLoader(
+ None,
+ file=io.BytesIO(flowFile.getContentsAsBytes()),
+ mode=mode,
+ include_page_breaks=include_page_breaks,
+ include_metadata=include_metadata,
+ )
+
+ else:
+ raise ValueError("Configured Input Format is invalid: " +
input_format)
+
+ documents = loader.load()
+
+ if len(metadata) > 0:
+ for doc in documents:
+ if doc.metadata is None:
+ doc.metadata = metadata
+ else:
+ doc.metadata.update(metadata)
+
+ return documents
+
+ def to_json(self, docs) -> str:
+ json_docs = []
+
+ for i, doc in enumerate(docs):
+ doc.metadata["chunk_index"] = i
+ doc.metadata["chunk_count"] = len(docs)
+
+ json_doc = json.dumps({"text": doc.page_content, "metadata":
doc.metadata})
+ json_docs.append(json_doc)
+
+ return "\n".join(json_docs)
+
+ def transform(self, context, flowFile):
+ documents = self.create_docs(context, flowFile)
+ output_json = self.to_json(documents)
+
+ return FlowFileTransformResult("success", contents=output_json,
attributes={"mime.type": "application/json"})
diff --git a/src/extensions/chunking/__init__.py
b/src/extensions/chunking/__init__.py
new file mode 100644
index 0000000..9881313
--- /dev/null
+++ b/src/extensions/chunking/__init__.py
@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0
diff --git a/src/extensions/openai/PromptChatGPT.py
b/src/extensions/openai/PromptChatGPT.py
new file mode 100644
index 0000000..a8470f4
--- /dev/null
+++ b/src/extensions/openai/PromptChatGPT.py
@@ -0,0 +1,219 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import re
+
+from nifiapi.flowfiletransform import FlowFileTransform,
FlowFileTransformResult
+from nifiapi.properties import ExpressionLanguageScope, PropertyDescriptor,
StandardValidators, TimeUnit
+
+FLOWFILE_CONTENT = "flowfile_content"
+FLOWFILE_CONTENT_REFERENCE = "{" + FLOWFILE_CONTENT + "}"
+# Regex to match { followed by any number of characters other than { or },
followed by }. But do not match if it starts with {{
+VAR_NAME_REGEX = r"(?<!{)\{([^{]*?)\}"
+
+
+class PromptChatGPT(FlowFileTransform):
+ class Java:
+ implements = ["org.apache.nifi.python.processor.FlowFileTransform"]
+
+ class ProcessorDetails:
+ version = "2.0.0.dev0"
+ description = "Submits a prompt to ChatGPT, writing the results either
to a FlowFile attribute or to the contents of the FlowFile"
+ tags = [
+ "text",
+ "chatgpt",
+ "gpt",
+ "machine learning",
+ "ML",
+ "artificial intelligence",
+ "ai",
+ "document",
+ "langchain",
+ ]
+ dependencies = ["langchain==0.1.2", "openai==1.9.0", "jsonpath-ng"]
+
+ MODEL = PropertyDescriptor(
+ name="OpenAI Model Name",
+ description="The name of the OpenAI Model to use in order to answer
the prompt",
+ default_value="gpt-3.5-turbo",
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ )
+ PROMPT = PropertyDescriptor(
+ name="Prompt",
+ description="""The prompt to issue to ChatGPT. This may use FlowFile
attributes via Expression Language and may also reference the FlowFile content
by using the literal
+ {flowfile_content} (including braces) in the prompt. If
the FlowFile's content is JSON formatted, a reference may also include JSONPath
Expressions
+ to reference specific fields in the FlowFile content, such
as {$.page_content}""",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ required=True,
+ )
+ TEMPERATURE = PropertyDescriptor(
+ name="Temperature",
+ description="""The Temperature parameter to submit to OpenAI. A lower
value will result in more consistent answers while a higher value will result
in a more creative answer.
+ "The value must be between 0 and 2, inclusive.""",
+
validators=[StandardValidators._standard_validators.createNonNegativeFloatingPointValidator(2.0)],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ required=True,
+ default_value="1.0",
+ )
+ RESULT_ATTRIBUTE = PropertyDescriptor(
+ name="Result Attribute",
+ description="If specified, the result will be added to the attribute
whose name is given. If not specified, the result will be written to the
FlowFile's content",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=False,
+ )
+ API_KEY = PropertyDescriptor(
+ name="API Key",
+ description="The OpenAI API Key to use",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ sensitive=True,
+ )
+ TIMEOUT = PropertyDescriptor(
+ name="Request Timeout",
+ description="The amount of time to wait before timing out the request",
+ validators=[StandardValidators.TIME_PERIOD_VALIDATOR],
+ default_value="60 secs",
+ required=True,
+ )
+ MAX_TOKENS = PropertyDescriptor(
+ name="Max Tokens to Generate",
+ description="The maximum number of tokens that ChatGPT should
generate",
+ validators=[StandardValidators.POSITIVE_INTEGER_VALIDATOR],
+ required=False,
+ )
+ ORGANIZATION = PropertyDescriptor(
+ name="OpenAI Organization ID",
+ description="The OpenAI Organization ID",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=False,
+ )
+ API_BASE = PropertyDescriptor(
+ name="API Base URL Path",
+ description="The API Base URL to use for interacting with OpenAI. This
should be populated only if using a proxy or an emulator.",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=False,
+ )
+
+ property_descriptors = [
+ MODEL,
+ PROMPT,
+ TEMPERATURE,
+ RESULT_ATTRIBUTE,
+ API_KEY,
+ TIMEOUT,
+ MAX_TOKENS,
+ ORGANIZATION,
+ API_BASE,
+ ]
+
+ def __init__(self, **kwargs):
+ pass
+
+ def getPropertyDescriptors(self):
+ return self.property_descriptors
+
+ def transform(self, context, flowFile):
+ from langchain import PromptTemplate
+ from langchain.chains.llm import LLMChain
+ from langchain.chat_models import ChatOpenAI
+
+ prompt =
context.getProperty(self.PROMPT).evaluateAttributeExpressions(flowFile).getValue()
+
+ # We want to allow referencing FlowFile content using JSONPath
Expressions.
+ # To do that, we allow the same {variable} syntax as Langchain. But
Langchain does not allow '$' characters
+ # to exist in the variable names. So we need to replace those
variables in the prompt with new variables, such as
+ # jsonpath_var_0, jsonpath_var_1, etc. To do this, we will use a Regex
to detect any variables that are referenced
+ # and if it starts with a $ we will replace it with
jsonpath_var_<index> and we will keep a mapping from that name to
+ # the substituted variable name so that we can later determine what
the JSONPath expression was.
+ variable_references = list(set(re.findall(VAR_NAME_REGEX, prompt)))
+
+ input_variables = []
+ jsonpath_to_var_mapping = {}
+ index = 0
+ for ref in variable_references:
+ if ref.startswith("$"):
+ var_name = "jsonpath_var_" + str(index)
+ index += 1
+ input_variables.append(var_name)
+ jsonpath_to_var_mapping[ref] = var_name
+ prompt = prompt.replace("{" + ref + "}", "{" + var_name + "}")
+ elif ref == FLOWFILE_CONTENT:
+ input_variables.append(ref)
+ else:
+ raise ValueError(
+ "Prompt contained an invalid variable reference: {"
+ + ref
+ + "}. Valid references are flowfile_content or any
JSONPath expression."
+ )
+
+ temperature =
context.getProperty(self.TEMPERATURE).evaluateAttributeExpressions(flowFile).asFloat()
+ model_name =
context.getProperty(self.MODEL).evaluateAttributeExpressions(flowFile).getValue()
+ api_key = context.getProperty(self.API_KEY).getValue()
+ timeout =
context.getProperty(self.TIMEOUT).asTimePeriod(TimeUnit.SECONDS)
+ max_tokens = context.getProperty(self.MAX_TOKENS).asInteger()
+ organization = context.getProperty(self.ORGANIZATION).getValue()
+ api_base = context.getProperty(self.API_BASE).getValue()
+
+ # Build out our LLMChain
+ llm = ChatOpenAI(
+ model_name=model_name,
+ temperature=temperature,
+ openai_api_key=api_key,
+ request_timeout=timeout,
+ max_retries=0,
+ max_tokens=max_tokens,
+ openai_organization=organization,
+ openai_api_base=api_base,
+ )
+
+ prompt_template = PromptTemplate(template=prompt,
input_variables=input_variables)
+
+ llm_chain = LLMChain(llm=llm, prompt=prompt_template)
+
+ # Substitute in any JSON Path Expressions or references to
{flowfile_content}.
+ llm_args = {}
+ json_content = None
+ for var_name in variable_references:
+ # If variable references {flowfile_content} substitute the content
+ if var_name == FLOWFILE_CONTENT:
+ llm_args[FLOWFILE_CONTENT] =
flowFile.getContentsAsBytes().decode()
+ if var_name.startswith("$"):
+ # Load the FlowFile's contents into the json_content variable
only once
+ if json_content is None:
+ json_content =
json.loads(flowFile.getContentsAsBytes().decode())
+
+ # Import jsonpath_ng so that we can evaluate JSONPath against
the FlowFile content.
+ from jsonpath_ng import parse
+
+ try:
+ jsonpath_expression = parse(var_name)
+ matches = jsonpath_expression.find(json_content)
+ variable_value = "\n".join([match.value for match in
matches])
+ except:
+ self.logger.exception(f"Invalid JSONPath reference in
prompt: {var_name}")
+ raise
+
+ # Insert the resolved value into llm_args
+ resolved_var_name = jsonpath_to_var_mapping.get(var_name)
+ llm_args[resolved_var_name] = variable_value
+
+ self.logger.debug(f"Evaluating prompt\nPrompt: {prompt}\nArgs:
#{llm_args}")
+
+ # Run the LLM Chain in order to prompt ChatGPT
+ results = llm_chain(llm_args)
+
+ # Create the output content or FLowFile attribute
+ text = results["text"]
+ attribute_name = context.getProperty(self.RESULT_ATTRIBUTE).getValue()
+ if attribute_name is None:
+ output_content = text
+ output_attributes = None
+ else:
+ output_content = None
+ output_attributes = {attribute_name: text}
+
+ # Return the results
+ return FlowFileTransformResult("success", contents=output_content,
attributes=output_attributes)
diff --git a/src/extensions/openai/__init__.py
b/src/extensions/openai/__init__.py
new file mode 100644
index 0000000..9881313
--- /dev/null
+++ b/src/extensions/openai/__init__.py
@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0
diff --git a/src/extensions/vectorstores/ChromaUtils.py
b/src/extensions/vectorstores/ChromaUtils.py
new file mode 100644
index 0000000..4a5ada0
--- /dev/null
+++ b/src/extensions/vectorstores/ChromaUtils.py
@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from nifiapi.properties import ExpressionLanguageScope, PropertyDependency,
PropertyDescriptor, StandardValidators
+
+# Connection Strategies
+LOCAL_DISK = "Local Disk"
+REMOTE_SERVER = "Remote Chroma Server"
+
+# Authentication Strategies
+TOKEN = "Token Authentication"
+BASIC_AUTH = "Basic Authentication"
+NONE = "None"
+
+# Transport Protocols
+HTTP = "http"
+HTTPS = "https"
+
+CONNECTION_STRATEGY = PropertyDescriptor(
+ name="Connection Strategy",
+ description="Specifies how to connect to the Chroma server",
+ allowable_values=[LOCAL_DISK, REMOTE_SERVER],
+ default_value=REMOTE_SERVER,
+ required=True,
+)
+DIRECTORY = PropertyDescriptor(
+ name="Directory",
+ description="The Directory that Chroma should use to persist data",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ default_value="./chroma",
+ dependencies=[PropertyDependency(CONNECTION_STRATEGY, LOCAL_DISK)],
+)
+HOSTNAME = PropertyDescriptor(
+ name="Hostname",
+ description="The hostname to connect to in order to communicate with
Chroma",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ default_value="localhost",
+ required=True,
+ dependencies=[PropertyDependency(CONNECTION_STRATEGY, REMOTE_SERVER)],
+)
+PORT = PropertyDescriptor(
+ name="Port",
+ description="The port that the Chroma server is listening on",
+ validators=[StandardValidators.PORT_VALIDATOR],
+ default_value="8000",
+ required=True,
+ dependencies=[PropertyDependency(CONNECTION_STRATEGY, REMOTE_SERVER)],
+)
+TRANSPORT_PROTOCOL = PropertyDescriptor(
+ name="Transport Protocol",
+ description="Specifies whether connections should be made over http or
https",
+ allowable_values=[HTTP, HTTPS],
+ default_value=HTTPS,
+ required=True,
+ dependencies=[PropertyDependency(CONNECTION_STRATEGY, REMOTE_SERVER)],
+)
+AUTH_STRATEGY = PropertyDescriptor(
+ name="Authentication Strategy",
+ description="Specifies how to authenticate to Chroma server",
+ allowable_values=[TOKEN, BASIC_AUTH, NONE],
+ default_value=TOKEN,
+ required=True,
+ dependencies=[PropertyDependency(CONNECTION_STRATEGY, REMOTE_SERVER)],
+)
+AUTH_TOKEN = PropertyDescriptor(
+ name="Authentication Token",
+ description="The token to use for authenticating to Chroma server",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ sensitive=True,
+ dependencies=[PropertyDependency(AUTH_STRATEGY, TOKEN)],
+)
+USERNAME = PropertyDescriptor(
+ name="Username",
+ description="The username to use for authenticating to Chroma server",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ dependencies=[PropertyDependency(AUTH_STRATEGY, BASIC_AUTH)],
+)
+PASSWORD = PropertyDescriptor(
+ name="Password",
+ description="The password to use for authenticating to Chroma server",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ sensitive=True,
+ dependencies=[PropertyDependency(AUTH_STRATEGY, BASIC_AUTH)],
+)
+COLLECTION_NAME = PropertyDescriptor(
+ name="Collection Name",
+ description="The name of the Chroma Collection",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ default_value="nifi",
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+)
+
+PROPERTIES = [
+ CONNECTION_STRATEGY,
+ DIRECTORY,
+ HOSTNAME,
+ PORT,
+ TRANSPORT_PROTOCOL,
+ AUTH_STRATEGY,
+ AUTH_TOKEN,
+ USERNAME,
+ PASSWORD,
+ COLLECTION_NAME,
+]
+
+
+def create_client(context):
+ import chromadb
+ from chromadb import Settings
+
+ connection_strategy = context.getProperty(CONNECTION_STRATEGY).getValue()
+ if connection_strategy == LOCAL_DISK:
+ directory = context.getProperty(DIRECTORY).getValue()
+ return chromadb.PersistentClient(directory)
+ hostname = context.getProperty(HOSTNAME).getValue()
+ port = context.getProperty(PORT).asInteger()
+ headers = {}
+ ssl = context.getProperty(TRANSPORT_PROTOCOL).getValue() == HTTPS
+
+ auth_strategy = context.getProperty(AUTH_STRATEGY).getValue()
+ if auth_strategy == TOKEN:
+ auth_provider = "chromadb.auth.token.TokenAuthClientProvider"
+ credentials = context.getProperty(AUTH_TOKEN).getValue()
+ elif auth_strategy == BASIC_AUTH:
+ auth_provider = "chromadb.auth.basic.BasicAuthClientProvider"
+ username = context.getProperty(USERNAME).getValue()
+ password = context.getProperty(PASSWORD).getValue()
+ credentials = username + ":" + password
+ else:
+ auth_provider = None
+ credentials = None
+
+ settings = Settings(chroma_client_auth_provider=auth_provider,
chroma_client_auth_credentials=credentials)
+ return chromadb.HttpClient(hostname, port, ssl, headers, settings)
diff --git a/src/extensions/vectorstores/EmbeddingUtils.py
b/src/extensions/vectorstores/EmbeddingUtils.py
new file mode 100644
index 0000000..df45086
--- /dev/null
+++ b/src/extensions/vectorstores/EmbeddingUtils.py
@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from langchain.embeddings.huggingface import HuggingFaceInferenceAPIEmbeddings
+from langchain.embeddings.openai import OpenAIEmbeddings
+from nifiapi.properties import PropertyDependency, PropertyDescriptor,
StandardValidators
+
+# Embedding Functions
+ONNX_ALL_MINI_LM_L6_V2 = "ONNX all-MiniLM-L6-v2 Model"
+HUGGING_FACE = "Hugging Face Model"
+OPENAI = "OpenAI Model"
+SENTENCE_TRANSFORMERS = "Sentence Transformers"
+
+
+EMBEDDING_FUNCTION = PropertyDescriptor(
+ name="Embedding Function",
+ description="Specifies which embedding function should be used in order to
create embeddings from incoming Documents",
+ allowable_values=[ONNX_ALL_MINI_LM_L6_V2, HUGGING_FACE, OPENAI,
SENTENCE_TRANSFORMERS],
+ default_value=ONNX_ALL_MINI_LM_L6_V2,
+ required=True,
+)
+HUGGING_FACE_MODEL_NAME = PropertyDescriptor(
+ name="HuggingFace Model Name",
+ description="The name of the HuggingFace model to use",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ default_value="sentence-transformers/all-MiniLM-L6-v2",
+ required=True,
+ dependencies=[PropertyDependency(EMBEDDING_FUNCTION, HUGGING_FACE)],
+)
+HUGGING_FACE_API_KEY = PropertyDescriptor(
+ name="HuggingFace API Key",
+ description="The API Key for interacting with HuggingFace",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ sensitive=True,
+ dependencies=[PropertyDependency(EMBEDDING_FUNCTION, HUGGING_FACE)],
+)
+OPENAI_API_KEY = PropertyDescriptor(
+ name="OpenAI API Key",
+ description="The API Key for interacting with OpenAI",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ sensitive=True,
+ dependencies=[PropertyDependency(EMBEDDING_FUNCTION, OPENAI)],
+)
+OPENAI_MODEL_NAME = PropertyDescriptor(
+ name="OpenAI Model Name",
+ description="The name of the OpenAI model to use",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ default_value="text-embedding-ada-002",
+ required=True,
+ dependencies=[PropertyDependency(EMBEDDING_FUNCTION, OPENAI)],
+)
+OPENAI_ORGANIZATION = PropertyDescriptor(
+ name="OpenAI Organization ID",
+ description="The OpenAI Organization ID",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=False,
+ dependencies=[PropertyDependency(EMBEDDING_FUNCTION, OPENAI)],
+)
+OPENAI_API_BASE = PropertyDescriptor(
+ name="OpenAI API Base Path",
+ description="The API Base to use for interacting with OpenAI. This is used
for interacting with different deployments, such as an Azure deployment.",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=False,
+ dependencies=[PropertyDependency(EMBEDDING_FUNCTION, OPENAI)],
+)
+OPENAI_API_TYPE = PropertyDescriptor(
+ name="OpenAI API Deployment Type",
+ description="The type of the OpenAI API Deployment. This is used for
interacting with different deployments, such as an Azure deployment.",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=False,
+ dependencies=[PropertyDependency(EMBEDDING_FUNCTION, OPENAI)],
+)
+OPENAI_API_VERSION = PropertyDescriptor(
+ name="OpenAI API Version",
+ description="The OpenAI API Version. This is used for interacting with
different deployments, such as an Azure deployment.",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=False,
+ dependencies=[PropertyDependency(EMBEDDING_FUNCTION, OPENAI)],
+)
+SENTENCE_TRANSFORMER_MODEL_NAME = PropertyDescriptor(
+ name="Sentence Transformer Model Name",
+ description="The name of the Sentence Transformer model to use",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ default_value="all-MiniLM-L6-v2",
+ required=True,
+ dependencies=[PropertyDependency(EMBEDDING_FUNCTION,
SENTENCE_TRANSFORMERS)],
+)
+SENTENCE_TRANSFORMER_DEVICE = PropertyDescriptor(
+ name="Sentence Transformer Device Type",
+ description="""The type of device to use for performing the embeddings
using the Sentence Transformer, such as 'cpu', 'cuda', 'mps', 'cuda:0', etc.
+ If not specified, a GPU will be used if possible, otherwise
a CPU.""",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=False,
+ dependencies=[PropertyDependency(EMBEDDING_FUNCTION,
SENTENCE_TRANSFORMERS)],
+)
+EMBEDDING_MODEL = PropertyDescriptor(
+ name="Embedding Model",
+ description="Specifies which embedding model should be used in order to
create embeddings from incoming Documents. Default model is OpenAI.",
+ allowable_values=[HUGGING_FACE, OPENAI],
+ default_value=OPENAI,
+ required=True,
+)
+PROPERTIES = [
+ EMBEDDING_FUNCTION,
+ HUGGING_FACE_MODEL_NAME,
+ HUGGING_FACE_API_KEY,
+ OPENAI_MODEL_NAME,
+ OPENAI_API_KEY,
+ OPENAI_ORGANIZATION,
+ OPENAI_API_BASE,
+ OPENAI_API_TYPE,
+ OPENAI_API_VERSION,
+ SENTENCE_TRANSFORMER_MODEL_NAME,
+ SENTENCE_TRANSFORMER_DEVICE,
+ EMBEDDING_MODEL,
+]
+
+
+def create_embedding_function(context):
+ from chromadb.utils.embedding_functions import (
+ HuggingFaceEmbeddingFunction,
+ ONNXMiniLM_L6_V2,
+ OpenAIEmbeddingFunction,
+ SentenceTransformerEmbeddingFunction,
+ )
+
+ function_name = context.getProperty(EMBEDDING_FUNCTION).getValue()
+ if function_name == ONNX_ALL_MINI_LM_L6_V2:
+ return ONNXMiniLM_L6_V2()
+
+ if function_name == OPENAI:
+ api_key = context.getProperty(OPENAI_API_KEY).getValue()
+ model_name = context.getProperty(OPENAI_MODEL_NAME).getValue()
+ organization_id = context.getProperty(OPENAI_ORGANIZATION).getValue()
+ api_base = context.getProperty(OPENAI_API_BASE).getValue()
+ api_type = context.getProperty(OPENAI_API_TYPE).getValue()
+ api_version = context.getProperty(OPENAI_API_VERSION).getValue()
+ return OpenAIEmbeddingFunction(
+ api_key=api_key,
+ model_name=model_name,
+ organization_id=organization_id,
+ api_base=api_base,
+ api_type=api_type,
+ api_version=api_version,
+ )
+
+ if function_name == HUGGING_FACE:
+ api_key = context.getProperty(HUGGING_FACE_API_KEY).getValue()
+ model_name = context.getProperty(HUGGING_FACE_MODEL_NAME).getValue()
+ return HuggingFaceEmbeddingFunction(api_key=api_key,
model_name=model_name)
+
+ model_name =
context.getProperty(SENTENCE_TRANSFORMER_MODEL_NAME).getValue()
+ device = context.getProperty(SENTENCE_TRANSFORMER_DEVICE).getValue()
+ return SentenceTransformerEmbeddingFunction(model_name=model_name,
device=device)
+
+
+def create_embedding_service(context):
+ embedding_service = context.getProperty(EMBEDDING_MODEL).getValue()
+
+ if embedding_service == OPENAI:
+ openai_api_key = context.getProperty(OPENAI_API_KEY).getValue()
+ return OpenAIEmbeddings(openai_api_key=openai_api_key)
+ huggingface_api_key = context.getProperty(HUGGING_FACE_API_KEY).getValue()
+ return HuggingFaceInferenceAPIEmbeddings(api_key=huggingface_api_key)
diff --git a/src/extensions/vectorstores/OpenSearchVectorUtils.py
b/src/extensions/vectorstores/OpenSearchVectorUtils.py
new file mode 100644
index 0000000..4f527e0
--- /dev/null
+++ b/src/extensions/vectorstores/OpenSearchVectorUtils.py
@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from EmbeddingUtils import EMBEDDING_MODEL, HUGGING_FACE, OPENAI
+from nifiapi.properties import ExpressionLanguageScope, PropertyDependency,
PropertyDescriptor, StandardValidators
+
+# Space types
+L2 = ("L2 (Euclidean distance)", "l2")
+L1 = ("L1 (Manhattan distance)", "l1")
+LINF = ("L-infinity (chessboard) distance", "linf")
+COSINESIMIL = ("Cosine similarity", "cosinesimil")
+
+HUGGING_FACE_API_KEY = PropertyDescriptor(
+ name="HuggingFace API Key",
+ description="The API Key for interacting with HuggingFace",
+ required=True,
+ sensitive=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, HUGGING_FACE)],
+)
+HUGGING_FACE_MODEL = PropertyDescriptor(
+ name="HuggingFace Model",
+ description="The name of the HuggingFace model to use",
+ default_value="sentence-transformers/all-MiniLM-L6-v2",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, HUGGING_FACE)],
+)
+OPENAI_API_KEY = PropertyDescriptor(
+ name="OpenAI API Key",
+ description="The API Key for OpenAI in order to create embeddings",
+ required=True,
+ sensitive=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, OPENAI)],
+)
+OPENAI_API_MODEL = PropertyDescriptor(
+ name="OpenAI Model",
+ description="The API Key for OpenAI in order to create embeddings",
+ default_value="text-embedding-ada-002",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, OPENAI)],
+)
+HTTP_HOST = PropertyDescriptor(
+ name="HTTP Host",
+ description="URL where OpenSearch is hosted.",
+ default_value="http://localhost:9200";,
+ required=True,
+ validators=[StandardValidators.URL_VALIDATOR],
+)
+USERNAME = PropertyDescriptor(
+ name="Username",
+ description="The username to use for authenticating to OpenSearch server",
+ required=False,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+)
+PASSWORD = PropertyDescriptor(
+ name="Password",
+ description="The password to use for authenticating to OpenSearch server",
+ required=False,
+ sensitive=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+)
+INDEX_NAME = PropertyDescriptor(
+ name="Index Name",
+ description="The name of the OpenSearch index.",
+ sensitive=False,
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+)
+VECTOR_FIELD = PropertyDescriptor(
+ name="Vector Field Name",
+ description="The name of field in the document where the embeddings are
stored. This field need to be a 'knn_vector' typed field.",
+ default_value="vector_field",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+)
+TEXT_FIELD = PropertyDescriptor(
+ name="Text Field Name",
+ description="The name of field in the document where the text is stored.",
+ default_value="text",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+)
+
+
+def create_authentication_params(context):
+ username = context.getProperty(USERNAME).getValue()
+ password = context.getProperty(PASSWORD).getValue()
+
+ params = {"verify_certs": "true"}
+
+ if username is not None and password is not None:
+ params["http_auth"] = (username, password)
+
+ return params
+
+
+def parse_documents(json_lines, id_field_name, file_name):
+ import json
+
+ texts = []
+ metadatas = []
+ ids = []
+ for i, line in enumerate(json_lines.split("\n"), start=1):
+ try:
+ doc = json.loads(line)
+ except Exception as e:
+ message = f"Could not parse line {i} as JSON"
+ raise ValueError(message) from e
+
+ text = doc.get("text")
+ metadata = doc.get("metadata")
+ texts.append(text)
+
+ # Remove any null values, or it will cause the embedding to fail
+ filtered_metadata = {key: value for key, value in metadata.items() if
value is not None}
+ metadatas.append(filtered_metadata)
+
+ doc_id = None
+ if id_field_name is not None:
+ doc_id = metadata.get(id_field_name)
+ if doc_id is None:
+ doc_id = file_name + "-" + str(i)
+ ids.append(doc_id)
+
+ return {"texts": texts, "metadatas": metadatas, "ids": ids}
diff --git a/src/extensions/vectorstores/PutChroma.py
b/src/extensions/vectorstores/PutChroma.py
new file mode 100644
index 0000000..f9aad49
--- /dev/null
+++ b/src/extensions/vectorstores/PutChroma.py
@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+import ChromaUtils
+import EmbeddingUtils
+from nifiapi.flowfiletransform import FlowFileTransform,
FlowFileTransformResult
+from nifiapi.properties import ExpressionLanguageScope, PropertyDescriptor,
StandardValidators
+
+
+class PutChroma(FlowFileTransform):
+ class Java:
+ implements = ["org.apache.nifi.python.processor.FlowFileTransform"]
+
+ class ProcessorDetails:
+ version = "2.0.0.dev0"
+ description = """Publishes JSON data to a Chroma VectorDB. The
Incoming data must be in single JSON per Line format, each with two keys:
'text' and 'metadata'.
+ The text must be a string, while metadata must be a map
with strings for values. Any additional fields will be ignored. If the
collection name specified
+ does not exist, the Processor will automatically create
the collection."""
+ tags = [
+ "chroma",
+ "vector",
+ "vectordb",
+ "embeddings",
+ "ai",
+ "artificial intelligence",
+ "ml",
+ "machine learning",
+ "text",
+ "LLM",
+ ]
+
+ STORE_TEXT = PropertyDescriptor(
+ name="Store Document Text",
+ description="""Specifies whether or not the text of the document
should be stored in Chroma. If so, both the document's text and its embedding
will be stored. If not,
+ only the vector/embedding will be stored.""",
+ allowable_values=["true", "false"],
+ required=True,
+ default_value="true",
+ )
+ DISTANCE_METHOD = PropertyDescriptor(
+ name="Distance Method",
+ description="If the specified collection does not exist, it will be
created using this Distance Method. If the collection exists, this property
will be ignored.",
+ allowable_values=["cosine", "l2", "ip"],
+ default_value="cosine",
+ required=True,
+ )
+ DOC_ID_FIELD_NAME = PropertyDescriptor(
+ name="Document ID Field Name",
+ description="""Specifies the name of the field in the 'metadata'
element of each document where the document's ID can be found.
+ If not specified, an ID will be generated based on the
FlowFile's filename and a one-up number.""",
+ required=False,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+
+ client = None
+ embedding_function = None
+
+ def __init__(self):
+ self.property_descriptors = list(ChromaUtils.PROPERTIES) + [
+ prop for prop in EmbeddingUtils.PROPERTIES if prop !=
EmbeddingUtils.EMBEDDING_MODEL
+ ]
+ self.property_descriptors.append(self.STORE_TEXT)
+ self.property_descriptors.append(self.DISTANCE_METHOD)
+ self.property_descriptors.append(self.DOC_ID_FIELD_NAME)
+
+ def getPropertyDescriptors(self):
+ return self.property_descriptors
+
+ def onScheduled(self, context):
+ self.client = ChromaUtils.create_client(context)
+ self.embedding_function =
EmbeddingUtils.create_embedding_function(context)
+
+ def transform(self, context, flowfile):
+ client = self.client
+ embedding_function = self.embedding_function
+ collection_name = (
+
context.getProperty(ChromaUtils.COLLECTION_NAME).evaluateAttributeExpressions(flowfile).getValue()
+ )
+ distance_method = context.getProperty(self.DISTANCE_METHOD).getValue()
+ id_field_name =
context.getProperty(self.DOC_ID_FIELD_NAME).evaluateAttributeExpressions(flowfile).getValue()
+
+ collection = client.get_or_create_collection(
+ name=collection_name, embedding_function=embedding_function,
metadata={"hnsw:space": distance_method}
+ )
+
+ json_lines = flowfile.getContentsAsBytes().decode()
+ i = 0
+ texts = []
+ metadatas = []
+ ids = []
+ for line in json_lines.split("\n"):
+ doc = json.loads(line)
+ text = doc.get("text")
+ metadata = doc.get("metadata")
+ texts.append(text)
+
+ # Remove any null values, or it will cause the embedding to fail
+ filtered_metadata = {}
+ for key, value in metadata.items():
+ if value is not None:
+ if isinstance(value, list):
+ for i, element in enumerate(value):
+ element_count = i + 1
+ indexed_key = f"{key}_{element_count}"
+ filtered_metadata[indexed_key] = element
+ else:
+ filtered_metadata[key] = value
+
+ metadatas.append(filtered_metadata)
+
+ doc_id = None
+ if id_field_name is not None:
+ doc_id = metadata.get(id_field_name)
+ if doc_id is None:
+ doc_id = flowfile.getAttribute("filename") + "-" + str(i)
+ ids.append(doc_id)
+
+ i += 1
+
+ embeddings = embedding_function(texts)
+ if not context.getProperty(self.STORE_TEXT).asBoolean():
+ texts = None
+
+ collection.upsert(ids, embeddings, metadatas, texts)
+
+ return FlowFileTransformResult(relationship="success")
diff --git a/src/extensions/vectorstores/PutOpenSearchVector.py
b/src/extensions/vectorstores/PutOpenSearchVector.py
new file mode 100644
index 0000000..92821db
--- /dev/null
+++ b/src/extensions/vectorstores/PutOpenSearchVector.py
@@ -0,0 +1,267 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from EmbeddingUtils import EMBEDDING_MODEL, create_embedding_service
+from langchain.vectorstores import OpenSearchVectorSearch
+from nifiapi.documentation import use_case
+from nifiapi.flowfiletransform import FlowFileTransform,
FlowFileTransformResult
+from nifiapi.properties import ExpressionLanguageScope, PropertyDependency,
PropertyDescriptor, StandardValidators
+from OpenSearchVectorUtils import (
+ COSINESIMIL,
+ HTTP_HOST,
+ HUGGING_FACE_API_KEY,
+ HUGGING_FACE_MODEL,
+ INDEX_NAME,
+ L1,
+ L2,
+ LINF,
+ OPENAI_API_KEY,
+ OPENAI_API_MODEL,
+ PASSWORD,
+ TEXT_FIELD,
+ USERNAME,
+ VECTOR_FIELD,
+ create_authentication_params,
+ parse_documents,
+)
+
+
+@use_case(
+ description="Create vectors/embeddings that represent text content and
send the vectors to OpenSearch",
+ notes="This use case assumes that the data has already been formatted in
JSONL format with the text to store in OpenSearch provided in the 'text'
field.",
+ keywords=["opensearch", "embedding", "vector", "text", "vectorstore",
"insert"],
+ configuration="""
+ Configure the 'HTTP Host' to an appropriate URL where
OpenSearch is accessible.
+ Configure 'Embedding Model' to indicate whether OpenAI
embeddings should be used or a HuggingFace embedding model should be used:
'Hugging Face Model' or 'OpenAI Model'
+ Configure the 'OpenAI API Key' or 'HuggingFace API Key',
depending on the chosen Embedding Model.
+ Set 'Index Name' to the name of your OpenSearch Index.
+ Set 'Vector Field Name' to the name of the field in the
document which will store the vector data.
+ Set 'Text Field Name' to the name of the field in the document
which will store the text data.
+
+ If the documents to send to OpenSearch contain a unique
identifier, set the 'Document ID Field Name' property to the name of the field
that contains the document ID.
+ This property can be left blank, in which case a unique ID
will be generated based on the FlowFile's filename.
+
+ If the provided index does not exists in OpenSearch then the
processor is capable to create it. The 'New Index Strategy' property defines
+ that the index needs to be created from the default template
or it should be configured with custom values.
+ """,
+)
+@use_case(
+ description="Update vectors/embeddings in OpenSearch",
+ notes="This use case assumes that the data has already been formatted in
JSONL format with the text to store in OpenSearch provided in the 'text'
field.",
+ keywords=["opensearch", "embedding", "vector", "text", "vectorstore",
"update", "upsert"],
+ configuration="""
+ Configure the 'HTTP Host' to an appropriate URL where
OpenSearch is accessible.
+ Configure 'Embedding Model' to indicate whether OpenAI
embeddings should be used or a HuggingFace embedding model should be used:
'Hugging Face Model' or 'OpenAI Model'
+ Configure the 'OpenAI API Key' or 'HuggingFace API Key',
depending on the chosen Embedding Model.
+ Set 'Index Name' to the name of your OpenSearch Index.
+ Set 'Vector Field Name' to the name of the field in the
document which will store the vector data.
+ Set 'Text Field Name' to the name of the field in the document
which will store the text data.
+ Set the 'Document ID Field Name' property to the name of the
field that contains the identifier of the document in OpenSearch to update.
+ """,
+)
+class PutOpenSearchVector(FlowFileTransform):
+ class Java:
+ implements = ["org.apache.nifi.python.processor.FlowFileTransform"]
+
+ class ProcessorDetails:
+ version = "2.0.0.dev0"
+ description = """Publishes JSON data to OpenSearch. The Incoming data
must be in single JSON per Line format, each with two keys: 'text' and
'metadata'.
+ The text must be a string, while metadata must be a map
with strings for values. Any additional fields will be ignored."""
+ tags = [
+ "opensearch",
+ "vector",
+ "vectordb",
+ "vectorstore",
+ "embeddings",
+ "ai",
+ "artificial intelligence",
+ "ml",
+ "machine learning",
+ "text",
+ "LLM",
+ ]
+
+ # Engine types
+ NMSLIB = ("nmslib (Non-Metric Space Library)", "nmslib")
+ FAISS = ("faiss (Facebook AI Similarity Search)", "faiss")
+ LUCENE = ("lucene", "lucene")
+
+ ENGINE_VALUES = dict([NMSLIB, FAISS, LUCENE])
+
+ # Space types
+ INNERPRODUCT = ("Inner product", "innerproduct")
+
+ NMSLIB_SPACE_TYPE_VALUES = dict([L2, L1, LINF, COSINESIMIL, INNERPRODUCT])
+ FAISS_SPACE_TYPE_VALUES = dict([L2, INNERPRODUCT])
+ LUCENE_SPACE_TYPE_VALUES = dict([L2, COSINESIMIL])
+
+ # New Index Mapping Strategy
+ DEFAULT_INDEX_MAPPING = "Default index mapping"
+ CUSTOM_INDEX_MAPPING = "Custom index mapping"
+
+ DOC_ID_FIELD_NAME = PropertyDescriptor(
+ name="Document ID Field Name",
+ description="""Specifies the name of the field in the 'metadata'
element of each document where the document's ID can be found.
+ If not specified, an ID will be generated based on the
FlowFile's filename and a one-up number.""",
+ required=False,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ NEW_INDEX_STRATEGY = PropertyDescriptor(
+ name="New Index Strategy",
+ description="""Specifies the Mapping strategy to use for new index
creation. The default template values are the following:
+ {engine: nmslib, space_type: l2, ef_search: 512,
ef_construction: 512, m: 16}""",
+ allowable_values=[DEFAULT_INDEX_MAPPING, CUSTOM_INDEX_MAPPING],
+ default_value=DEFAULT_INDEX_MAPPING,
+ required=False,
+ )
+ ENGINE = PropertyDescriptor(
+ name="Engine",
+ description="The approximate k-NN library to use for indexing and
search.",
+ allowable_values=ENGINE_VALUES.keys(),
+ default_value=NMSLIB[0],
+ required=False,
+ dependencies=[PropertyDependency(NEW_INDEX_STRATEGY,
CUSTOM_INDEX_MAPPING)],
+ )
+ NMSLIB_SPACE_TYPE = PropertyDescriptor(
+ name="NMSLIB Space Type",
+ description="The vector space used to calculate the distance between
vectors.",
+ allowable_values=NMSLIB_SPACE_TYPE_VALUES.keys(),
+ default_value=L2[0],
+ required=False,
+ dependencies=[
+ PropertyDependency(NEW_INDEX_STRATEGY, CUSTOM_INDEX_MAPPING),
+ PropertyDependency(ENGINE, NMSLIB[0]),
+ ],
+ )
+ FAISS_SPACE_TYPE = PropertyDescriptor(
+ name="FAISS Space Type",
+ description="The vector space used to calculate the distance between
vectors.",
+ allowable_values=FAISS_SPACE_TYPE_VALUES.keys(),
+ default_value=L2[0],
+ required=False,
+ dependencies=[
+ PropertyDependency(NEW_INDEX_STRATEGY, CUSTOM_INDEX_MAPPING),
+ PropertyDependency(ENGINE, FAISS[0]),
+ ],
+ )
+ LUCENE_SPACE_TYPE = PropertyDescriptor(
+ name="Lucene Space Type",
+ description="The vector space used to calculate the distance between
vectors.",
+ allowable_values=LUCENE_SPACE_TYPE_VALUES.keys(),
+ default_value=L2[0],
+ required=False,
+ dependencies=[
+ PropertyDependency(NEW_INDEX_STRATEGY, CUSTOM_INDEX_MAPPING),
+ PropertyDependency(ENGINE, LUCENE[0]),
+ ],
+ )
+ EF_SEARCH = PropertyDescriptor(
+ name="EF Search",
+ description="The size of the dynamic list used during k-NN searches.
Higher values lead to more accurate but slower searches.",
+ default_value="512",
+ required=False,
+ validators=[StandardValidators.NON_NEGATIVE_INTEGER_VALIDATOR],
+ dependencies=[PropertyDependency(NEW_INDEX_STRATEGY,
CUSTOM_INDEX_MAPPING)],
+ )
+ EF_CONSTRUCTION = PropertyDescriptor(
+ name="EF Construction",
+ description="The size of the dynamic list used during k-NN graph
creation. Higher values lead to a more accurate graph but slower indexing
speed.",
+ default_value="512",
+ required=False,
+ validators=[StandardValidators.NON_NEGATIVE_INTEGER_VALIDATOR],
+ dependencies=[PropertyDependency(NEW_INDEX_STRATEGY,
CUSTOM_INDEX_MAPPING)],
+ )
+ M = PropertyDescriptor(
+ name="M",
+ description="The number of bidirectional links that the plugin creates
for each new element. Increasing and "
+ "decreasing this value can have a large impact on memory consumption.
Keep this value between 2 and 100.",
+ default_value="16",
+ required=False,
+
validators=[StandardValidators._standard_validators.createLongValidator(2, 100,
True)],
+ dependencies=[PropertyDependency(NEW_INDEX_STRATEGY,
CUSTOM_INDEX_MAPPING)],
+ )
+
+ properties = [
+ EMBEDDING_MODEL,
+ OPENAI_API_KEY,
+ OPENAI_API_MODEL,
+ HUGGING_FACE_API_KEY,
+ HUGGING_FACE_MODEL,
+ HTTP_HOST,
+ USERNAME,
+ PASSWORD,
+ INDEX_NAME,
+ DOC_ID_FIELD_NAME,
+ VECTOR_FIELD,
+ TEXT_FIELD,
+ NEW_INDEX_STRATEGY,
+ ENGINE,
+ NMSLIB_SPACE_TYPE,
+ FAISS_SPACE_TYPE,
+ LUCENE_SPACE_TYPE,
+ EF_SEARCH,
+ EF_CONSTRUCTION,
+ M,
+ ]
+
+ embeddings = None
+
+ def __init__(self, **kwargs):
+ pass
+
+ def getPropertyDescriptors(self):
+ return self.properties
+
+ def onScheduled(self, context):
+ self.embeddings = create_embedding_service(context)
+
+ def transform(self, context, flowfile):
+ file_name = flowfile.getAttribute("filename")
+ http_host =
context.getProperty(HTTP_HOST).evaluateAttributeExpressions(flowfile).getValue()
+ index_name =
context.getProperty(INDEX_NAME).evaluateAttributeExpressions(flowfile).getValue()
+ id_field_name =
context.getProperty(self.DOC_ID_FIELD_NAME).evaluateAttributeExpressions(flowfile).getValue()
+ vector_field =
context.getProperty(VECTOR_FIELD).evaluateAttributeExpressions(flowfile).getValue()
+ text_field =
context.getProperty(TEXT_FIELD).evaluateAttributeExpressions(flowfile).getValue()
+ new_index_strategy =
context.getProperty(self.NEW_INDEX_STRATEGY).evaluateAttributeExpressions().getValue()
+
+ params = {"vector_field": vector_field, "text_field": text_field}
+ params.update(create_authentication_params(context))
+
+ if new_index_strategy == self.CUSTOM_INDEX_MAPPING:
+ engine =
context.getProperty(self.ENGINE).evaluateAttributeExpressions().getValue()
+ params["engine"] = self.ENGINE_VALUES.get(engine)
+
+ if engine == self.NMSLIB[0]:
+ space_type =
context.getProperty(self.NMSLIB_SPACE_TYPE).evaluateAttributeExpressions().getValue()
+ params["space_type"] =
self.NMSLIB_SPACE_TYPE_VALUES.get(space_type)
+ if engine == self.FAISS[0]:
+ space_type =
context.getProperty(self.FAISS_SPACE_TYPE).evaluateAttributeExpressions().getValue()
+ params["space_type"] =
self.FAISS_SPACE_TYPE_VALUES.get(space_type)
+ if engine == self.LUCENE[0]:
+ space_type =
context.getProperty(self.LUCENE_SPACE_TYPE).evaluateAttributeExpressions().getValue()
+ params["space_type"] =
self.LUCENE_SPACE_TYPE_VALUES.get(space_type)
+
+ ef_search =
context.getProperty(self.EF_SEARCH).evaluateAttributeExpressions().asInteger()
+ params["ef_search"] = ef_search
+
+ ef_construction =
context.getProperty(self.EF_CONSTRUCTION).evaluateAttributeExpressions().asInteger()
+ params["ef_construction"] = ef_construction
+
+ m =
context.getProperty(self.M).evaluateAttributeExpressions().asInteger()
+ params["m"] = m
+
+ # Read the FlowFile content as "json-lines".
+ json_lines = flowfile.getContentsAsBytes().decode()
+ parsed_documents = parse_documents(json_lines, id_field_name,
file_name)
+
+ vectorstore = OpenSearchVectorSearch(
+ opensearch_url=http_host, index_name=index_name,
embedding_function=self.embeddings, **params
+ )
+ vectorstore.add_texts(
+ texts=parsed_documents["texts"],
+ metadatas=parsed_documents["metadatas"],
+ ids=parsed_documents["ids"],
+ **params,
+ )
+
+ return FlowFileTransformResult(relationship="success")
diff --git a/src/extensions/vectorstores/PutPinecone.py
b/src/extensions/vectorstores/PutPinecone.py
new file mode 100644
index 0000000..8625082
--- /dev/null
+++ b/src/extensions/vectorstores/PutPinecone.py
@@ -0,0 +1,220 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+import langchain.vectorstores
+from EmbeddingUtils import EMBEDDING_MODEL, HUGGING_FACE, OPENAI,
create_embedding_service
+from nifiapi.documentation import use_case
+from nifiapi.flowfiletransform import FlowFileTransform,
FlowFileTransformResult
+from nifiapi.properties import ExpressionLanguageScope, PropertyDependency,
PropertyDescriptor, StandardValidators
+from pinecone import Pinecone
+
+
+@use_case(
+ description="Create vectors/embeddings that represent text content and
send the vectors to Pinecone",
+ notes="This use case assumes that the data has already been formatted in
JSONL format with the text to store in Pinecone provided in the 'text' field.",
+ keywords=["pinecone", "embedding", "vector", "text", "vectorstore",
"insert"],
+ configuration="""
+ Configure the 'Pinecone API Key' to the appropriate
authentication token for interacting with Pinecone.
+ Configure 'Embedding Model' to indicate whether OpenAI
embeddings should be used or a HuggingFace embedding model should be used:
'Hugging Face Model' or 'OpenAI Model'
+ Configure the 'OpenAI API Key' or 'HuggingFace API Key',
depending on the chosen Embedding Model.
+ Set 'Pinecone Environment' to the name of your Pinecone
environment
+ Set 'Index Name' to the name of your Pinecone Index.
+ Set 'Namespace' to appropriate namespace, or leave it empty to
use the default Namespace.
+
+ If the documents to send to Pinecone contain a unique
identifier, set the 'Document ID Field Name' property to the name of the field
that contains the document ID.
+ This property can be left blank, in which case a unique ID
will be generated based on the FlowFile's filename.
+ """,
+)
+@use_case(
+ description="Update vectors/embeddings in Pinecone",
+ notes="This use case assumes that the data has already been formatted in
JSONL format with the text to store in Pinecone provided in the 'text' field.",
+ keywords=["pinecone", "embedding", "vector", "text", "vectorstore",
"update", "upsert"],
+ configuration="""
+ Configure the 'Pinecone API Key' to the appropriate
authentication token for interacting with Pinecone.
+ Configure 'Embedding Model' to indicate whether OpenAI
embeddings should be used or a HuggingFace embedding model should be used:
'Hugging Face Model' or 'OpenAI Model'
+ Configure the 'OpenAI API Key' or 'HuggingFace API Key',
depending on the chosen Embedding Model.
+ Set 'Pinecone Environment' to the name of your Pinecone
environment
+ Set 'Index Name' to the name of your Pinecone Index.
+ Set 'Namespace' to appropriate namespace, or leave it empty to
use the default Namespace.
+ Set the 'Document ID Field Name' property to the name of the
field that contains the identifier of the document in Pinecone to update.
+ """,
+)
+class PutPinecone(FlowFileTransform):
+ class Java:
+ implements = ["org.apache.nifi.python.processor.FlowFileTransform"]
+
+ class ProcessorDetails:
+ version = "2.0.0.dev0"
+ description = """Publishes JSON data to Pinecone. The Incoming data
must be in single JSON per Line format, each with two keys: 'text' and
'metadata'.
+ The text must be a string, while metadata must be a map
with strings for values. Any additional fields will be ignored."""
+ tags = [
+ "pinecone",
+ "vector",
+ "vectordb",
+ "vectorstore",
+ "embeddings",
+ "ai",
+ "artificial intelligence",
+ "ml",
+ "machine learning",
+ "text",
+ "LLM",
+ ]
+
+ PINECONE_API_KEY = PropertyDescriptor(
+ name="Pinecone API Key",
+ description="The API Key to use in order to authentication with
Pinecone",
+ sensitive=True,
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ )
+ HUGGING_FACE_API_KEY = PropertyDescriptor(
+ name="HuggingFace API Key",
+ description="The API Key for interacting with HuggingFace",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ sensitive=True,
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, HUGGING_FACE)],
+ )
+ HUGGING_FACE_MODEL = PropertyDescriptor(
+ name="HuggingFace Model",
+ description="The name of the HuggingFace model to use",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ default_value="sentence-transformers/all-MiniLM-L6-v2",
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, HUGGING_FACE)],
+ )
+ OPENAI_API_KEY = PropertyDescriptor(
+ name="OpenAI API Key",
+ description="The API Key for OpenAI in order to create embeddings",
+ sensitive=True,
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, OPENAI)],
+ )
+ OPENAI_API_MODEL = PropertyDescriptor(
+ name="OpenAI Model",
+ description="The API Key for OpenAI in order to create embeddings",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ default_value="text-embedding-ada-002",
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, OPENAI)],
+ )
+ PINECONE_ENV = PropertyDescriptor(
+ name="Pinecone Environment",
+ description="The name of the Pinecone Environment. This can be found
in the Pinecone console next to the API Key.",
+ sensitive=False,
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ )
+ INDEX_NAME = PropertyDescriptor(
+ name="Index Name",
+ description="The name of the Pinecone index.",
+ sensitive=False,
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ TEXT_KEY = PropertyDescriptor(
+ name="Text Key",
+ description="The key in the document that contains the text to create
embeddings for.",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ default_value="text",
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ NAMESPACE = PropertyDescriptor(
+ name="Namespace",
+ description="The name of the Pinecone Namespace to put the documents
to.",
+ required=False,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ DOC_ID_FIELD_NAME = PropertyDescriptor(
+ name="Document ID Field Name",
+ description="""Specifies the name of the field in the 'metadata'
element of each document where the document's ID can be found.
+ If not specified, an ID will be generated based on the
FlowFile's filename and a one-up number.""",
+ required=False,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+
+ properties = [
+ PINECONE_API_KEY,
+ EMBEDDING_MODEL,
+ OPENAI_API_KEY,
+ OPENAI_API_MODEL,
+ HUGGING_FACE_API_KEY,
+ HUGGING_FACE_MODEL,
+ PINECONE_ENV,
+ INDEX_NAME,
+ TEXT_KEY,
+ NAMESPACE,
+ DOC_ID_FIELD_NAME,
+ ]
+
+ embeddings = None
+ pc = None
+
+ def __init__(self, **kwargs):
+ pass
+
+ def getPropertyDescriptors(self):
+ return self.properties
+
+ def onScheduled(self, context):
+ # initialize pinecone
+ self.pc = Pinecone(
+ api_key=context.getProperty(self.PINECONE_API_KEY).getValue(),
+ environment=context.getProperty(self.PINECONE_ENV).getValue(),
+ )
+ # initialize embedding service
+ self.embeddings = create_embedding_service(context)
+
+ def transform(self, context, flowfile):
+ # First, check if our index already exists. If it doesn't, we create it
+ index_name =
context.getProperty(self.INDEX_NAME).evaluateAttributeExpressions(flowfile).getValue()
+ namespace =
context.getProperty(self.NAMESPACE).evaluateAttributeExpressions(flowfile).getValue()
+ id_field_name =
context.getProperty(self.DOC_ID_FIELD_NAME).evaluateAttributeExpressions(flowfile).getValue()
+
+ index = self.pc.Index(index_name)
+
+ # Read the FlowFile content as "json-lines".
+ json_lines = flowfile.getContentsAsBytes().decode()
+ i = 1
+ texts = []
+ metadatas = []
+ ids = []
+ for line in json_lines.split("\n"):
+ try:
+ doc = json.loads(line)
+ except Exception as e:
+ message = f"Could not parse line {i} as JSON"
+ raise ValueError(message) from e
+
+ text = doc.get("text")
+ metadata = doc.get("metadata")
+ texts.append(text)
+
+ # Remove any null values, or it will cause the embedding to fail
+ filtered_metadata = {}
+ for key, value in metadata.items():
+ if value is not None:
+ filtered_metadata[key] = value
+
+ metadatas.append(filtered_metadata)
+
+ doc_id = None
+ if id_field_name is not None:
+ doc_id = metadata.get(id_field_name)
+ if doc_id is None:
+ doc_id = flowfile.getAttribute("filename") + "-" + str(i)
+ ids.append(doc_id)
+
+ i += 1
+
+ text_key =
context.getProperty(self.TEXT_KEY).evaluateAttributeExpressions().getValue()
+ vectorstore = langchain.vectorstores.Pinecone(index,
self.embeddings.embed_query, text_key)
+ vectorstore.add_texts(texts=texts, metadatas=metadatas, ids=ids,
namespace=namespace)
+ return FlowFileTransformResult(relationship="success")
diff --git a/src/extensions/vectorstores/PutQdrant.py
b/src/extensions/vectorstores/PutQdrant.py
new file mode 100644
index 0000000..42e7d7c
--- /dev/null
+++ b/src/extensions/vectorstores/PutQdrant.py
@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+import QdrantUtils
+from EmbeddingUtils import (
+ create_embedding_service,
+)
+from langchain.vectorstores.qdrant import Qdrant
+from nifiapi.documentation import use_case
+from nifiapi.flowfiletransform import FlowFileTransform,
FlowFileTransformResult
+from nifiapi.properties import (
+ ExpressionLanguageScope,
+ PropertyDescriptor,
+ StandardValidators,
+)
+from qdrant_client.models import Distance
+
+
+@use_case(
+ description="Create embeddings that semantically represent text content
and upload to Qdrant - https://qdrant.tech/";,
+ notes="This processor assumes that the data has already been formatted in
JSONL format with the text to store in Qdrant provided in the 'text' field.",
+ keywords=["qdrant", "embedding", "vector", "text", "vectorstore",
"insert"],
+ configuration="""
+ Configure 'Collection Name' to the name of the Qdrant
collection to use.
+ Configure 'Qdrant URL' to the fully qualified URL of the
Qdrant instance.
+ Configure 'Qdrant API Key' to the API Key to use in order to
authenticate with Qdrant.
+ Configure 'Prefer gRPC' to True if you want to use gRPC for
interfacing with Qdrant.
+ Configure 'Use HTTPS' to True if you want to use TLS(HTTPS)
while interfacing with Qdrant.
+ Configure 'Embedding Model' to indicate whether OpenAI
embeddings should be used or a HuggingFace embedding model should be used:
'Hugging Face Model' or 'OpenAI Model'
+ Configure 'HuggingFace API Key' or 'OpenAI API Key', depending
on the chosen Embedding Model.
+ Configure 'HuggingFace Model' or 'OpenAI Model' to the name of
the model to use.
+ Configure 'Force Recreate Collection' to True if you want to
recreate the collection if it already exists.
+ Configure 'Similarity Metric' to the similarity metric to use
when querying Qdrant.
+
+ If the documents to send to Qdrant contain a unique
identifier(UUID), set the 'Document ID Field Name' property to the name of the
field that contains the document ID.
+ This property can be left blank, in which case a UUID will be
generated based on the FlowFile's filename.
+ """,
+)
+class PutQdrant(FlowFileTransform):
+ class Java:
+ implements = ["org.apache.nifi.python.processor.FlowFileTransform"]
+
+ class ProcessorDetails:
+ version = "2.0.0.dev0"
+ description = """Publishes JSON data to Qdrant. The Incoming data must
be in single JSON per Line format, each with two keys: 'text' and 'metadata'.
+ The text must be a string, while metadata must be a map
with strings for values. Any additional fields will be ignored."""
+ tags = [
+ "qdrant",
+ "vector",
+ "vectordb",
+ "vectorstore",
+ "embeddings",
+ "ai",
+ "artificial intelligence",
+ "ml",
+ "machine learning",
+ "text",
+ "LLM",
+ ]
+
+ DOC_ID_FIELD_NAME = PropertyDescriptor(
+ name="Document ID Field Name",
+ description="""Specifies the name of the field in the 'metadata'
element of each document where the document's ID can be found.
+ If not specified, a UUID will be generated based on the
FlowFile's filename and an incremental number.""",
+ required=False,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ FORCE_RECREATE_COLLECTION = PropertyDescriptor(
+ name="Force Recreate Collection",
+ description="Specifies whether to recreate the collection if it
already exists. Essentially clearing the existing data.",
+ required=True,
+ default_value="False",
+ allowable_values=["True", "False"],
+ validators=[StandardValidators.BOOLEAN_VALIDATOR],
+ )
+ SIMILARITY_METRIC = PropertyDescriptor(
+ name="Similarity Metric",
+ description="Specifies the similarity metric when creating the
collection.",
+ required=True,
+ default_value=Distance.COSINE,
+ allowable_values=[
+ Distance.COSINE,
+ Distance.EUCLID,
+ Distance.DOT,
+ Distance.MANHATTAN,
+ ],
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ )
+
+ properties = (
+ QdrantUtils.QDRANT_PROPERTIES
+ + QdrantUtils.EMBEDDING_MODEL_PROPERTIES
+ + [
+ FORCE_RECREATE_COLLECTION,
+ SIMILARITY_METRIC,
+ DOC_ID_FIELD_NAME,
+ ]
+ )
+
+ def __init__(self, **kwargs):
+ pass
+
+ def getPropertyDescriptors(self):
+ return self.properties
+
+ def onScheduled(self, context):
+ # The Qdrant#construct_instance() internally checks if the collection
exists
+ # and creates it if it doesn't with the appropriate dimesions and
configurations.
+ self.vector_store = Qdrant.construct_instance(
+ texts=["Some text to obtain the embeddings dimension when creating
the collection"],
+ embedding=create_embedding_service(context),
+
collection_name=context.getProperty(QdrantUtils.COLLECTION_NAME).getValue(),
+ url=context.getProperty(QdrantUtils.QDRANT_URL).getValue(),
+ api_key=context.getProperty(QdrantUtils.QDRANT_API_KEY).getValue(),
+
prefer_grpc=context.getProperty(QdrantUtils.PREFER_GRPC).asBoolean(),
+ https=context.getProperty(QdrantUtils.HTTPS).asBoolean(),
+
force_recreate=context.getProperty(self.FORCE_RECREATE_COLLECTION).asBoolean(),
+
distance_func=context.getProperty(self.SIMILARITY_METRIC).getValue(),
+ )
+
+ def transform(self, context, flowfile):
+ id_field_name =
context.getProperty(self.DOC_ID_FIELD_NAME).evaluateAttributeExpressions(flowfile).getValue()
+
+ # Read the FlowFile content as "json-lines".
+ json_lines = flowfile.getContentsAsBytes().decode()
+ i = 1
+ texts, metadatas, ids = [], [], []
+ for line in json_lines.split("\n"):
+ try:
+ doc = json.loads(line)
+ except Exception as e:
+ message = f"Could not parse line {i} as JSON"
+ raise ValueError(message) from e
+
+ metadata = doc.get("metadata")
+ texts.append(doc.get("text"))
+ metadatas.append(metadata)
+
+ doc_id = None
+ if id_field_name is not None:
+ doc_id = metadata.get(id_field_name)
+ if doc_id is None:
+ doc_id =
QdrantUtils.convert_id(flowfile.getAttribute("filename") + "-" + str(i))
+ ids.append(doc_id)
+
+ i += 1
+
+ self.vector_store.add_texts(texts=texts, metadatas=metadatas, ids=ids)
+ return FlowFileTransformResult(relationship="success")
diff --git a/src/extensions/vectorstores/QdrantUtils.py
b/src/extensions/vectorstores/QdrantUtils.py
new file mode 100644
index 0000000..7435124
--- /dev/null
+++ b/src/extensions/vectorstores/QdrantUtils.py
@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import uuid
+
+from EmbeddingUtils import (
+ EMBEDDING_MODEL,
+ HUGGING_FACE,
+ OPENAI,
+)
+from nifiapi.properties import (
+ ExpressionLanguageScope,
+ PropertyDependency,
+ PropertyDescriptor,
+ StandardValidators,
+)
+
+DEFAULT_COLLECTION_NAME = "apache-nifi"
+
+
+COLLECTION_NAME = PropertyDescriptor(
+ name="Collection Name",
+ description="The name of the Qdrant collection to use.",
+ sensitive=False,
+ required=True,
+ default_value=DEFAULT_COLLECTION_NAME,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+)
+QDRANT_URL = PropertyDescriptor(
+ name="Qdrant URL",
+ description="The fully qualified URL to the Qdrant instance.",
+ sensitive=False,
+ required=True,
+ default_value="http://localhost:6333";,
+ validators=[StandardValidators.URL_VALIDATOR],
+)
+QDRANT_API_KEY = PropertyDescriptor(
+ name="Qdrant API Key",
+ description="The API Key to use in order to authentication with Qdrant.
Can be empty.",
+ sensitive=True,
+ required=True,
+)
+
+PREFER_GRPC = PropertyDescriptor(
+ name="Prefer gRPC",
+ description="Specifies whether to use gRPC for interfacing with Qdrant.",
+ required=True,
+ default_value=False,
+ allowable_values=["True", "False"],
+ validators=[StandardValidators.BOOLEAN_VALIDATOR],
+)
+HTTPS = PropertyDescriptor(
+ name="Use HTTPS",
+ description="Specifies whether to TLS(HTTPS) while interfacing with
Qdrant.",
+ required=True,
+ default_value=False,
+ allowable_values=["True", "False"],
+ validators=[StandardValidators.BOOLEAN_VALIDATOR],
+)
+
+QDRANT_PROPERTIES = [COLLECTION_NAME, QDRANT_URL, QDRANT_API_KEY, PREFER_GRPC,
HTTPS]
+
+HUGGING_FACE_API_KEY = PropertyDescriptor(
+ name="HuggingFace API Key",
+ description="The API Key for interacting with HuggingFace",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ sensitive=True,
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, HUGGING_FACE)],
+)
+HUGGING_FACE_MODEL = PropertyDescriptor(
+ name="HuggingFace Model",
+ description="The name of the HuggingFace model to use.",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ default_value="sentence-transformers/all-MiniLM-L6-v2",
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, HUGGING_FACE)],
+)
+OPENAI_API_KEY = PropertyDescriptor(
+ name="OpenAI API Key",
+ description="The API Key for OpenAI in order to create embeddings.",
+ sensitive=True,
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, OPENAI)],
+)
+OPENAI_API_MODEL = PropertyDescriptor(
+ name="OpenAI Model",
+ description="The name of the OpenAI model to use.",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ default_value="text-embedding-ada-002",
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, OPENAI)],
+)
+
+EMBEDDING_MODEL_PROPERTIES = [
+ EMBEDDING_MODEL,
+ HUGGING_FACE_API_KEY,
+ HUGGING_FACE_MODEL,
+ OPENAI_API_KEY,
+ OPENAI_API_MODEL,
+]
+
+
+def convert_id(_id: str) -> str:
+ """
+ Converts any string into a UUID string deterministically.
+
+ Qdrant accepts UUID strings and unsigned integers as point ID.
+ This allows us to overwrite the same point with the original ID.
+ """
+ return str(uuid.uuid5(uuid.NAMESPACE_DNS, _id))
diff --git a/src/extensions/vectorstores/QueryChroma.py
b/src/extensions/vectorstores/QueryChroma.py
new file mode 100644
index 0000000..3d28ccf
--- /dev/null
+++ b/src/extensions/vectorstores/QueryChroma.py
@@ -0,0 +1,181 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+import ChromaUtils
+import EmbeddingUtils
+import QueryUtils
+from nifiapi.flowfiletransform import FlowFileTransform,
FlowFileTransformResult
+from nifiapi.properties import ExpressionLanguageScope, PropertyDescriptor,
StandardValidators
+
+
+class QueryChroma(FlowFileTransform):
+ class Java:
+ implements = ["org.apache.nifi.python.processor.FlowFileTransform"]
+
+ class ProcessorDetails:
+ version = "2.0.0.dev0"
+ description = "Queries a Chroma Vector Database in order to gather a
specified number of documents that are most closely related to the given query."
+ tags = [
+ "chroma",
+ "vector",
+ "vectordb",
+ "embeddings",
+ "enrich",
+ "enrichment",
+ "ai",
+ "artificial intelligence",
+ "ml",
+ "machine learning",
+ "text",
+ "LLM",
+ ]
+
+ QUERY = PropertyDescriptor(
+ name="Query",
+ description="""The query to issue to the Chroma VectorDB. The query is
always converted into embeddings using the configured embedding function, and
the embedding is
+ then sent to Chroma. The text itself is not sent to
Chroma.""",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ NUMBER_OF_RESULTS = PropertyDescriptor(
+ name="Number of Results",
+ description="The number of results to return from Chroma",
+ required=True,
+ validators=[StandardValidators.POSITIVE_INTEGER_VALIDATOR],
+ default_value="10",
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ METADATA_FILTER = PropertyDescriptor(
+ name="Metadata Filter",
+ description="""A JSON representation of a Metadata Filter that can be
applied against the Chroma documents in order to narrow down the documents that
can be returned.
+ For example: { "metadata_field": "some_value" }""",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ required=False,
+ )
+ DOCUMENT_FILTER = PropertyDescriptor(
+ name="Document Filter",
+ description="""A JSON representation of a Document Filter that can be
applied against the Chroma documents' text in order to narrow down the
documents that can be returned.
+ For example: { "$contains": "search_string" }""",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ required=False,
+ )
+
+ client = None
+ embedding_function = None
+ include_ids = None
+ include_metadatas = None
+ include_documents = None
+ include_distances = None
+ include_embeddings = None
+ results_field = None
+
+ property_descriptors = (
+ list(ChromaUtils.PROPERTIES)
+ + [prop for prop in EmbeddingUtils.PROPERTIES if prop !=
EmbeddingUtils.EMBEDDING_MODEL]
+ + [
+ QUERY,
+ NUMBER_OF_RESULTS,
+ QueryUtils.OUTPUT_STRATEGY,
+ QueryUtils.RESULTS_FIELD,
+ METADATA_FILTER,
+ DOCUMENT_FILTER,
+ QueryUtils.INCLUDE_IDS,
+ QueryUtils.INCLUDE_METADATAS,
+ QueryUtils.INCLUDE_DOCUMENTS,
+ QueryUtils.INCLUDE_DISTANCES,
+ QueryUtils.INCLUDE_EMBEDDINGS,
+ ]
+ )
+
+ def __init__(self, **kwargs):
+ pass
+
+ def getPropertyDescriptors(self):
+ return self.property_descriptors
+
+ def onScheduled(self, context):
+ self.client = ChromaUtils.create_client(context)
+ self.embedding_function =
EmbeddingUtils.create_embedding_function(context)
+ self.include_ids =
context.getProperty(QueryUtils.INCLUDE_IDS).asBoolean()
+ self.include_metadatas =
context.getProperty(QueryUtils.INCLUDE_METADATAS).asBoolean()
+ self.include_documents =
context.getProperty(QueryUtils.INCLUDE_DOCUMENTS).asBoolean()
+ self.include_distances =
context.getProperty(QueryUtils.INCLUDE_DISTANCES).asBoolean()
+ self.include_embeddings =
context.getProperty(QueryUtils.INCLUDE_EMBEDDINGS).asBoolean()
+ self.results_field =
context.getProperty(QueryUtils.RESULTS_FIELD).getValue()
+ self.query_utils = QueryUtils.QueryUtils(context)
+
+ def transform(self, context, flowfile):
+ client = self.client
+ embedding_function = self.embedding_function
+ collection_name = (
+
context.getProperty(ChromaUtils.COLLECTION_NAME).evaluateAttributeExpressions(flowfile).getValue()
+ )
+
+ collection = client.get_collection(name=collection_name,
embedding_function=embedding_function)
+
+ query_text =
context.getProperty(self.QUERY).evaluateAttributeExpressions(flowfile).getValue()
+ embeddings = embedding_function([query_text])
+
+ included_fields = []
+ if self.include_distances:
+ included_fields.append("distances")
+ if self.include_documents:
+ included_fields.append("documents")
+ if self.include_embeddings:
+ included_fields.append("embeddings")
+ if self.include_metadatas:
+ included_fields.append("metadatas")
+
+ where = None
+ where_clause =
context.getProperty(self.METADATA_FILTER).evaluateAttributeExpressions(flowfile).getValue()
+ if where_clause is not None:
+ where = json.loads(where_clause)
+
+ where_document = None
+ where_document_clause = (
+
context.getProperty(self.DOCUMENT_FILTER).evaluateAttributeExpressions(flowfile).getValue()
+ )
+ if where_document_clause is not None:
+ where_document = json.loads(where_document_clause)
+
+ query_results = collection.query(
+ query_embeddings=embeddings,
+
n_results=context.getProperty(self.NUMBER_OF_RESULTS).evaluateAttributeExpressions(flowfile).asInteger(),
+ include=included_fields,
+ where_document=where_document,
+ where=where,
+ )
+
+ ids = query_results["ids"][0]
+ distances = (
+ None
+ if (not self.include_distances or query_results["distances"] is
None)
+ else query_results["distances"][0]
+ )
+ metadatas = (
+ None
+ if (not self.include_metadatas or query_results["metadatas"] is
None)
+ else query_results["metadatas"][0]
+ )
+ documents = (
+ None
+ if (not self.include_documents or query_results["documents"] is
None)
+ else query_results["documents"][0]
+ )
+ embeddings = (
+ None
+ if (not self.include_embeddings or query_results["embeddings"] is
None)
+ else query_results["embeddings"][0]
+ )
+
+ (output_contents, mime_type) = self.query_utils.create_json(
+ flowfile, documents, metadatas, embeddings, distances, ids
+ )
+
+ # Return the results
+ attributes = {"mime.type": mime_type}
+ return FlowFileTransformResult(relationship="success",
contents=output_contents, attributes=attributes)
diff --git a/src/extensions/vectorstores/QueryOpenSearchVector.py
b/src/extensions/vectorstores/QueryOpenSearchVector.py
new file mode 100644
index 0000000..365f824
--- /dev/null
+++ b/src/extensions/vectorstores/QueryOpenSearchVector.py
@@ -0,0 +1,242 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+from EmbeddingUtils import EMBEDDING_MODEL, create_embedding_service
+from langchain.vectorstores import OpenSearchVectorSearch
+from nifiapi.flowfiletransform import FlowFileTransform,
FlowFileTransformResult
+from nifiapi.properties import ExpressionLanguageScope, PropertyDependency,
PropertyDescriptor, StandardValidators
+from OpenSearchVectorUtils import (
+ COSINESIMIL,
+ HTTP_HOST,
+ HUGGING_FACE_API_KEY,
+ HUGGING_FACE_MODEL,
+ INDEX_NAME,
+ L1,
+ L2,
+ LINF,
+ OPENAI_API_KEY,
+ OPENAI_API_MODEL,
+ PASSWORD,
+ TEXT_FIELD,
+ USERNAME,
+ VECTOR_FIELD,
+ create_authentication_params,
+)
+from QueryUtils import INCLUDE_DISTANCES, INCLUDE_METADATAS, OUTPUT_STRATEGY,
RESULTS_FIELD, QueryUtils
+
+
+class QueryOpenSearchVector(FlowFileTransform):
+ class Java:
+ implements = ["org.apache.nifi.python.processor.FlowFileTransform"]
+
+ class ProcessorDetails:
+ version = "2.0.0.dev0"
+ description = "Queries OpenSearch in order to gather a specified
number of documents that are most closely related to the given query."
+ tags = [
+ "opensearch",
+ "vector",
+ "vectordb",
+ "vectorstore",
+ "embeddings",
+ "ai",
+ "artificial intelligence",
+ "ml",
+ "machine learning",
+ "text",
+ "LLM",
+ ]
+
+ # Search types
+ APPROXIMATE_SEARCH = ("Approximate Search", "approximate_search")
+ SCRIPT_SCORING_SEARCH = ("Script Scoring Search", "script_scoring")
+ PAINLESS_SCRIPTING_SEARCH = ("Painless Scripting Search",
"painless_scripting")
+
+ SEARCH_TYPE_VALUES = dict([APPROXIMATE_SEARCH, SCRIPT_SCORING_SEARCH,
PAINLESS_SCRIPTING_SEARCH])
+
+ # Script Scoring Search space types
+ HAMMINGBIT = ("Hamming distance", "hammingbit")
+
+ SCRIPT_SCORING_SPACE_TYPE_VALUES = dict([L2, L1, LINF, COSINESIMIL,
HAMMINGBIT])
+
+ # Painless Scripting Search space types
+ L2_SQUARED = ("L2 (Euclidean distance)", "l2Squared")
+ L1_NORM = ("L1 (Manhattan distance)", "l1Norm")
+ COSINE_SIMILARITY = ("Cosine similarity", "cosineSimilarity")
+
+ PAINLESS_SCRIPTING_SPACE_TYPE_VALUES = dict([L2_SQUARED, L1_NORM,
COSINE_SIMILARITY])
+
+ QUERY = PropertyDescriptor(
+ name="Query",
+ description="The text of the query to send to OpenSearch.",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ NUMBER_OF_RESULTS = PropertyDescriptor(
+ name="Number of Results",
+ description="The number of results to return from OpenSearch",
+ default_value="10",
+ required=True,
+ validators=[StandardValidators.POSITIVE_INTEGER_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ SEARCH_TYPE = PropertyDescriptor(
+ name="Search Type",
+ description="Specifies the type of the search to be performed.",
+ allowable_values=SEARCH_TYPE_VALUES.keys(),
+ default_value=APPROXIMATE_SEARCH[0],
+ required=True,
+ )
+ SCRIPT_SCORING_SPACE_TYPE = PropertyDescriptor(
+ name="Script Scoring Space Type",
+ description="Used to measure the distance between two points in order
to determine the k-nearest neighbors.",
+ allowable_values=SCRIPT_SCORING_SPACE_TYPE_VALUES.keys(),
+ default_value=L2[0],
+ required=False,
+ dependencies=[PropertyDependency(SEARCH_TYPE,
SCRIPT_SCORING_SEARCH[0])],
+ )
+ PAINLESS_SCRIPTING_SPACE_TYPE = PropertyDescriptor(
+ name="Painless Scripting Space Type",
+ description="Used to measure the distance between two points in order
to determine the k-nearest neighbors.",
+ allowable_values=PAINLESS_SCRIPTING_SPACE_TYPE_VALUES.keys(),
+ default_value=L2_SQUARED[0],
+ required=False,
+ dependencies=[PropertyDependency(SEARCH_TYPE,
PAINLESS_SCRIPTING_SEARCH[0])],
+ )
+ BOOLEAN_FILTER = PropertyDescriptor(
+ name="Boolean Filter",
+ description="A Boolean filter is a post filter consists of a Boolean
query that contains a k-NN query and a filter. "
+ "The value of the field must be a JSON representation of the filter.",
+ required=False,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ dependencies=[PropertyDependency(SEARCH_TYPE, APPROXIMATE_SEARCH[0])],
+ )
+ EFFICIENT_FILTER = PropertyDescriptor(
+ name="Efficient Filter",
+ description="The Lucene Engine or Faiss Engine decides whether to
perform an exact k-NN search with "
+ "pre-filtering or an approximate search with modified post-filtering.
The value of the field must "
+ "be a JSON representation of the filter.",
+ required=False,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ dependencies=[PropertyDependency(SEARCH_TYPE, APPROXIMATE_SEARCH[0])],
+ )
+ PRE_FILTER = PropertyDescriptor(
+ name="Pre Filter",
+ description="Script Score query to pre-filter documents before
identifying nearest neighbors. The value of "
+ "the field must be a JSON representation of the filter.",
+ default_value='{"match_all": {}}',
+ required=False,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ dependencies=[PropertyDependency(SEARCH_TYPE,
SCRIPT_SCORING_SEARCH[0], PAINLESS_SCRIPTING_SEARCH[0])],
+ )
+
+ properties = [
+ EMBEDDING_MODEL,
+ OPENAI_API_KEY,
+ OPENAI_API_MODEL,
+ HUGGING_FACE_API_KEY,
+ HUGGING_FACE_MODEL,
+ HTTP_HOST,
+ USERNAME,
+ PASSWORD,
+ INDEX_NAME,
+ QUERY,
+ VECTOR_FIELD,
+ TEXT_FIELD,
+ NUMBER_OF_RESULTS,
+ SEARCH_TYPE,
+ SCRIPT_SCORING_SPACE_TYPE,
+ PAINLESS_SCRIPTING_SPACE_TYPE,
+ BOOLEAN_FILTER,
+ EFFICIENT_FILTER,
+ PRE_FILTER,
+ OUTPUT_STRATEGY,
+ RESULTS_FIELD,
+ INCLUDE_METADATAS,
+ INCLUDE_DISTANCES,
+ ]
+
+ embeddings = None
+ query_utils = None
+
+ def __init__(self, **kwargs):
+ pass
+
+ def getPropertyDescriptors(self):
+ return self.properties
+
+ def onScheduled(self, context):
+ # initialize embedding service
+ self.embeddings = create_embedding_service(context)
+ self.query_utils = QueryUtils(context)
+
+ def transform(self, context, flowfile):
+ http_host =
context.getProperty(HTTP_HOST).evaluateAttributeExpressions(flowfile).getValue()
+ index_name =
context.getProperty(INDEX_NAME).evaluateAttributeExpressions(flowfile).getValue()
+ query =
context.getProperty(self.QUERY).evaluateAttributeExpressions(flowfile).getValue()
+ num_results =
context.getProperty(self.NUMBER_OF_RESULTS).evaluateAttributeExpressions(flowfile).asInteger()
+ vector_field =
context.getProperty(VECTOR_FIELD).evaluateAttributeExpressions(flowfile).getValue()
+ text_field =
context.getProperty(TEXT_FIELD).evaluateAttributeExpressions(flowfile).getValue()
+ search_type =
context.getProperty(self.SEARCH_TYPE).evaluateAttributeExpressions().getValue()
+
+ params = {
+ "vector_field": vector_field,
+ "text_field": text_field,
+ "search_type": self.SEARCH_TYPE_VALUES.get(search_type),
+ }
+ params.update(create_authentication_params(context))
+
+ if search_type == self.APPROXIMATE_SEARCH[0]:
+ boolean_filter =
context.getProperty(self.BOOLEAN_FILTER).evaluateAttributeExpressions().getValue()
+ if boolean_filter is not None:
+ params["boolean_filter"] = json.loads(boolean_filter)
+
+ efficient_filter =
context.getProperty(self.EFFICIENT_FILTER).evaluateAttributeExpressions().getValue()
+ if efficient_filter is not None:
+ params["efficient_filter"] = json.loads(efficient_filter)
+ else:
+ pre_filter =
context.getProperty(self.PRE_FILTER).evaluateAttributeExpressions().getValue()
+ if pre_filter is not None:
+ params["pre_filter"] = json.loads(pre_filter)
+ if search_type == self.SCRIPT_SCORING_SEARCH[0]:
+ space_type = (
+
context.getProperty(self.SCRIPT_SCORING_SPACE_TYPE).evaluateAttributeExpressions().getValue()
+ )
+ params["space_type"] =
self.SCRIPT_SCORING_SPACE_TYPE_VALUES.get(space_type)
+ elif search_type == self.PAINLESS_SCRIPTING_SEARCH[0]:
+ space_type = (
+
context.getProperty(self.PAINLESS_SCRIPTING_SPACE_TYPE).evaluateAttributeExpressions().getValue()
+ )
+ params["space_type"] =
self.PAINLESS_SCRIPTING_SPACE_TYPE_VALUES.get(space_type)
+
+ vectorstore = OpenSearchVectorSearch(
+ index_name=index_name, embedding_function=self.embeddings,
opensearch_url=http_host, **params
+ )
+
+ results = vectorstore.similarity_search_with_score(query=query,
k=num_results, **params)
+
+ documents = []
+ for result in results:
+ documents.append(result[0].page_content)
+
+ if context.getProperty(INCLUDE_METADATAS):
+ metadatas = []
+ for result in results:
+ metadatas.append(result[0].metadata)
+ else:
+ metadatas = None
+
+ if context.getProperty(INCLUDE_DISTANCES):
+ distances = []
+ for result in results:
+ distances.append(result[1])
+ else:
+ distances = None
+
+ (output_contents, mime_type) = self.query_utils.create_json(
+ flowfile, documents, metadatas, None, distances, None
+ )
+ attributes = {"mime.type": mime_type}
+
+ return FlowFileTransformResult(relationship="success",
contents=output_contents, attributes=attributes)
diff --git a/src/extensions/vectorstores/QueryPinecone.py
b/src/extensions/vectorstores/QueryPinecone.py
new file mode 100644
index 0000000..aeb416a
--- /dev/null
+++ b/src/extensions/vectorstores/QueryPinecone.py
@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+import langchain.vectorstores
+import QueryUtils
+from EmbeddingUtils import EMBEDDING_MODEL, HUGGING_FACE, OPENAI,
create_embedding_service
+from nifiapi.flowfiletransform import FlowFileTransform,
FlowFileTransformResult
+from nifiapi.properties import ExpressionLanguageScope, PropertyDependency,
PropertyDescriptor, StandardValidators
+from pinecone import Pinecone
+
+
+class QueryPinecone(FlowFileTransform):
+ class Java:
+ implements = ["org.apache.nifi.python.processor.FlowFileTransform"]
+
+ class ProcessorDetails:
+ version = "2.0.0.dev0"
+ description = "Queries Pinecone in order to gather a specified number
of documents that are most closely related to the given query."
+ tags = [
+ "pinecone",
+ "vector",
+ "vectordb",
+ "vectorstore",
+ "embeddings",
+ "ai",
+ "artificial intelligence",
+ "ml",
+ "machine learning",
+ "text",
+ "LLM",
+ ]
+
+ PINECONE_API_KEY = PropertyDescriptor(
+ name="Pinecone API Key",
+ description="The API Key to use in order to authentication with
Pinecone",
+ sensitive=True,
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ )
+ OPENAI_API_KEY = PropertyDescriptor(
+ name="OpenAI API Key",
+ description="The API Key for OpenAI in order to create embeddings",
+ sensitive=True,
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, OPENAI)],
+ )
+ HUGGING_FACE_API_KEY = PropertyDescriptor(
+ name="HuggingFace API Key",
+ description="The API Key for interacting with HuggingFace",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ sensitive=True,
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, HUGGING_FACE)],
+ )
+ OPENAI_MODEL = PropertyDescriptor(
+ name="OpenAI Model",
+ description="The API Key for OpenAI in order to create embeddings",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ default_value="text-embedding-ada-002",
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, OPENAI)],
+ )
+ HUGGING_FACE_MODEL = PropertyDescriptor(
+ name="HuggingFace Model",
+ description="The name of the HuggingFace model to use",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=True,
+ default_value="sentence-transformers/all-MiniLM-L6-v2",
+ dependencies=[PropertyDependency(EMBEDDING_MODEL, HUGGING_FACE)],
+ )
+ PINECONE_ENV = PropertyDescriptor(
+ name="Pinecone Environment",
+ description="The name of the Pinecone Environment. This can be found
in the Pinecone console next to the API Key.",
+ sensitive=False,
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ )
+ INDEX_NAME = PropertyDescriptor(
+ name="Index Name",
+ description="The name of the Pinecone index.",
+ sensitive=False,
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ QUERY = PropertyDescriptor(
+ name="Query",
+ description="The text of the query to send to Pinecone.",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ NUMBER_OF_RESULTS = PropertyDescriptor(
+ name="Number of Results",
+ description="The number of results to return from Pinecone",
+ required=True,
+ validators=[StandardValidators.POSITIVE_INTEGER_VALIDATOR],
+ default_value="10",
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ TEXT_KEY = PropertyDescriptor(
+ name="Text Key",
+ description="The key in the document that contains the text to create
embeddings for.",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ default_value="text",
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ NAMESPACE = PropertyDescriptor(
+ name="Namespace",
+ 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,
+ )
+
+ properties = [
+ PINECONE_API_KEY,
+ EMBEDDING_MODEL,
+ OPENAI_API_KEY,
+ OPENAI_MODEL,
+ HUGGING_FACE_API_KEY,
+ HUGGING_FACE_MODEL,
+ PINECONE_ENV,
+ INDEX_NAME,
+ QUERY,
+ FILTER,
+ NUMBER_OF_RESULTS,
+ NAMESPACE,
+ TEXT_KEY,
+ QueryUtils.OUTPUT_STRATEGY,
+ QueryUtils.RESULTS_FIELD,
+ QueryUtils.INCLUDE_METADATAS,
+ QueryUtils.INCLUDE_DISTANCES,
+ ]
+
+ embeddings = None
+ query_utils = None
+ pc = None
+
+ def __init__(self, **kwargs):
+ pass
+
+ def getPropertyDescriptors(self):
+ return self.properties
+
+ def onScheduled(self, context):
+ # initialize pinecone
+ self.pc = Pinecone(
+ api_key=context.getProperty(self.PINECONE_API_KEY).getValue(),
+ environment=context.getProperty(self.PINECONE_ENV).getValue(),
+ )
+ # initialize embedding service
+ self.embeddings = create_embedding_service(context)
+ self.query_utils = QueryUtils.QueryUtils(context)
+
+ def transform(self, context, flowfile):
+ # First, check if our index already exists. If it doesn't, we create it
+ index_name =
context.getProperty(self.INDEX_NAME).evaluateAttributeExpressions(flowfile).getValue()
+ query =
context.getProperty(self.QUERY).evaluateAttributeExpressions(flowfile).getValue()
+ namespace =
context.getProperty(self.NAMESPACE).evaluateAttributeExpressions(flowfile).getValue()
+ num_results =
context.getProperty(self.NUMBER_OF_RESULTS).evaluateAttributeExpressions(flowfile).asInteger()
+
+ index = self.pc.Index(index_name)
+
+ text_key =
context.getProperty(self.TEXT_KEY).evaluateAttributeExpressions().getValue()
+ filter_definition =
context.getProperty(self.FILTER).evaluateAttributeExpressions(flowfile).getValue()
+ vectorstore = langchain.vectorstores.Pinecone(index,
self.embeddings.embed_query, text_key, namespace=namespace)
+ results = vectorstore.similarity_search_with_score(
+ query, num_results, filter=None if filter_definition is None else
json.loads(filter_definition)
+ )
+
+ documents = []
+ for result in results:
+ documents.append(result[0].page_content)
+
+ if context.getProperty(QueryUtils.INCLUDE_METADATAS):
+ metadatas = []
+ for result in results:
+ metadatas.append(result[0].metadata)
+ else:
+ metadatas = None
+
+ if context.getProperty(QueryUtils.INCLUDE_DISTANCES):
+ distances = []
+ for result in results:
+ distances.append(result[1])
+ else:
+ distances = None
+
+ (output_contents, mime_type) = self.query_utils.create_json(
+ flowfile, documents, metadatas, None, distances, None
+ )
+ attributes = {"mime.type": mime_type}
+
+ return FlowFileTransformResult(relationship="success",
contents=output_contents, attributes=attributes)
diff --git a/src/extensions/vectorstores/QueryQdrant.py
b/src/extensions/vectorstores/QueryQdrant.py
new file mode 100644
index 0000000..a7a63e3
--- /dev/null
+++ b/src/extensions/vectorstores/QueryQdrant.py
@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+import QdrantUtils
+import QueryUtils
+from EmbeddingUtils import (
+ create_embedding_service,
+)
+from langchain.vectorstores.qdrant import Qdrant
+from nifiapi.documentation import use_case
+from nifiapi.flowfiletransform import FlowFileTransform,
FlowFileTransformResult
+from nifiapi.properties import (
+ ExpressionLanguageScope,
+ PropertyDescriptor,
+ StandardValidators,
+)
+from qdrant_client import QdrantClient
+
+
+@use_case(
+ description="Semantically search for documents stored in Qdrant -
https://qdrant.tech/";,
+ keywords=["qdrant", "embedding", "vector", "text", "vectorstore",
"search"],
+ configuration="""
+ Configure 'Collection Name' to the name of the Qdrant
collection to use.
+ Configure 'Qdrant URL' to the fully qualified URL of the
Qdrant instance.
+ Configure 'Qdrant API Key' to the API Key to use in order to
authenticate with Qdrant.
+ Configure 'Prefer gRPC' to True if you want to use gRPC for
interfacing with Qdrant.
+ Configure 'Use HTTPS' to True if you want to use TLS(HTTPS)
while interfacing with Qdrant.
+ Configure 'Embedding Model' to indicate whether OpenAI
embeddings should be used or a HuggingFace embedding model should be used:
'Hugging Face Model' or 'OpenAI Model'
+ Configure 'HuggingFace API Key' or 'OpenAI API Key', depending
on the chosen Embedding Model.
+ Configure 'HuggingFace Model' or 'OpenAI Model' to the name of
the model to use.
+ Configure 'Query' to the text of the query to send to Qdrant.
+ Configure 'Number of Results' to the number of results to
return from Qdrant.
+ Configure 'Metadata Filter' to apply an optional metadata
filter with the query. For example: { "author": "john.doe" }
+ Configure 'Output Strategy' to indicate how the output should
be formatted: 'Row-Oriented', 'Text', or 'Column-Oriented'.
+ Configure 'Results Field' to the name of the field to insert
the results, if the input FlowFile is JSON Formatted,.
+ Configure 'Include Metadatas' to True if metadata should be
included in the output.
+ Configure 'Include Distances' to True if distances should be
included in the output.
+ """,
+)
+class QueryQdrant(FlowFileTransform):
+ class Java:
+ implements = ["org.apache.nifi.python.processor.FlowFileTransform"]
+
+ class ProcessorDetails:
+ version = "2.0.0.dev0"
+ description = "Queries Qdrant in order to gather a specified number of
documents that are most closely related to the given query."
+ tags = [
+ "qdrant",
+ "vector",
+ "vectordb",
+ "vectorstore",
+ "embeddings",
+ "ai",
+ "artificial intelligence",
+ "ml",
+ "machine learning",
+ "text",
+ "LLM",
+ ]
+
+ QUERY = PropertyDescriptor(
+ name="Query",
+ description="The text of the query to send to Qdrant.",
+ required=True,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ NUMBER_OF_RESULTS = PropertyDescriptor(
+ name="Number of Results",
+ description="The number of results to return from Qdrant.",
+ required=True,
+ validators=[StandardValidators.POSITIVE_INTEGER_VALIDATOR],
+ default_value="10",
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+ FILTER = PropertyDescriptor(
+ name="Metadata Filter",
+ description='Optional metadata filter to apply with the query. For
example: { "author": "john.doe" }',
+ required=False,
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ expression_language_scope=ExpressionLanguageScope.FLOWFILE_ATTRIBUTES,
+ )
+
+ properties = (
+ QdrantUtils.QDRANT_PROPERTIES
+ + QdrantUtils.EMBEDDING_MODEL_PROPERTIES
+ + [
+ QUERY,
+ FILTER,
+ NUMBER_OF_RESULTS,
+ QueryUtils.OUTPUT_STRATEGY,
+ QueryUtils.RESULTS_FIELD,
+ QueryUtils.INCLUDE_METADATAS,
+ QueryUtils.INCLUDE_DISTANCES,
+ ]
+ )
+
+ embeddings = None
+ query_utils = None
+ client = None
+
+ def __init__(self, **kwargs):
+ pass
+
+ def getPropertyDescriptors(self):
+ return self.properties
+
+ def onScheduled(self, context):
+ self.client = QdrantClient(
+ url=context.getProperty(QdrantUtils.QDRANT_URL).getValue(),
+ api_key=context.getProperty(QdrantUtils.QDRANT_API_KEY).getValue(),
+
prefer_grpc=context.getProperty(QdrantUtils.PREFER_GRPC).asBoolean(),
+ https=context.getProperty(QdrantUtils.HTTPS).asBoolean(),
+ )
+ self.embeddings = create_embedding_service(context)
+ self.query_utils = QueryUtils.QueryUtils(context)
+
+ def transform(self, context, flowfile):
+ collection_name = (
+
context.getProperty(QdrantUtils.COLLECTION_NAME).evaluateAttributeExpressions(flowfile).getValue()
+ )
+ query =
context.getProperty(self.QUERY).evaluateAttributeExpressions(flowfile).getValue()
+ num_results =
context.getProperty(self.NUMBER_OF_RESULTS).evaluateAttributeExpressions(flowfile).asInteger()
+ filter_definition =
context.getProperty(self.FILTER).evaluateAttributeExpressions(flowfile).getValue()
+ vector_store = Qdrant(
+ client=self.client,
+ collection_name=collection_name,
+ embeddings=self.embeddings,
+ )
+ results = vector_store.similarity_search_with_score(
+ query=query,
+ k=num_results,
+ filter=None if filter is None else json.loads(filter_definition),
+ )
+
+ documents = []
+ for result in results:
+ documents.append(result[0].page_content)
+
+ if context.getProperty(QueryUtils.INCLUDE_METADATAS).asBoolean():
+ metadatas = []
+ for result in results:
+ metadatas.append(result[0].metadata)
+ else:
+ metadatas = None
+
+ if context.getProperty(QueryUtils.INCLUDE_DISTANCES).asBoolean():
+ distances = []
+ for result in results:
+ distances.append(result[1])
+ else:
+ distances = None
+
+ (output_contents, mime_type) = self.query_utils.create_json(
+ flowfile, documents, metadatas, None, distances, None
+ )
+ attributes = {"mime.type": mime_type}
+
+ return FlowFileTransformResult(relationship="success",
contents=output_contents, attributes=attributes)
diff --git a/src/extensions/vectorstores/QueryUtils.py
b/src/extensions/vectorstores/QueryUtils.py
new file mode 100644
index 0000000..c9d9a73
--- /dev/null
+++ b/src/extensions/vectorstores/QueryUtils.py
@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+from nifiapi.properties import PropertyDependency, PropertyDescriptor,
StandardValidators
+
+ROW_ORIENTED = "Row-Oriented"
+TEXT = "Text"
+COLUMN_ORIENTED = "Column-Oriented"
+
+
+OUTPUT_STRATEGY = PropertyDescriptor(
+ name="Output Strategy",
+ description="""Specifies whether the output should contain only the text
of the documents (each document separated by \\n\\n), or if it
+ should be formatted as either single column-oriented JSON
object,
+ consisting of a keys 'ids', 'embeddings', 'documents',
'distances', and 'metadatas'; or if the results should be row-oriented,
+ a JSON per line, each consisting of a single id, document,
metadata, embedding, and distance.""",
+ allowable_values=[ROW_ORIENTED, TEXT, COLUMN_ORIENTED],
+ default_value=ROW_ORIENTED,
+ required=True,
+)
+RESULTS_FIELD = PropertyDescriptor(
+ name="Results Field",
+ description="""If the input FlowFile is JSON Formatted, this represents
the name of the field to insert the results. This allows the results to be
inserted into
+ "an existing input in order to enrich it. If this property is
unset, the results will be written to the FlowFile contents, overwriting any
pre-existing content.""",
+ validators=[StandardValidators.NON_EMPTY_VALIDATOR],
+ required=False,
+)
+
+INCLUDE_IDS = PropertyDescriptor(
+ name="Include Document IDs",
+ description="Whether or not to include the Documents' IDs in the response",
+ allowable_values=["true", "false"],
+ default_value="true",
+ required=False,
+ dependencies=[PropertyDependency(OUTPUT_STRATEGY, ROW_ORIENTED,
COLUMN_ORIENTED)],
+)
+INCLUDE_METADATAS = PropertyDescriptor(
+ name="Include Metadata",
+ description="Whether or not to include the Documents' Metadata in the
response",
+ allowable_values=["true", "false"],
+ default_value="true",
+ required=False,
+ dependencies=[PropertyDependency(OUTPUT_STRATEGY, ROW_ORIENTED,
COLUMN_ORIENTED)],
+)
+INCLUDE_DOCUMENTS = PropertyDescriptor(
+ name="Include Document",
+ description="Whether or not to include the Documents' Text in the
response",
+ allowable_values=["true", "false"],
+ default_value="true",
+ required=False,
+ dependencies=[PropertyDependency(OUTPUT_STRATEGY, ROW_ORIENTED,
COLUMN_ORIENTED)],
+)
+INCLUDE_DISTANCES = PropertyDescriptor(
+ name="Include Distances",
+ description="Whether or not to include the Documents' Distances (i.e., how
far the Document was away from the query) in the response",
+ allowable_values=["true", "false"],
+ default_value="true",
+ required=False,
+ dependencies=[PropertyDependency(OUTPUT_STRATEGY, ROW_ORIENTED,
COLUMN_ORIENTED)],
+)
+INCLUDE_EMBEDDINGS = PropertyDescriptor(
+ name="Include Embeddings",
+ description="Whether or not to include the Documents' Embeddings in the
response",
+ allowable_values=["true", "false"],
+ default_value="false",
+ required=False,
+ dependencies=[PropertyDependency(OUTPUT_STRATEGY, ROW_ORIENTED,
COLUMN_ORIENTED)],
+)
+
+
+class QueryUtils:
+ context = None
+
+ def __init__(self, context):
+ self.context = context
+ self.results_field = context.getProperty(RESULTS_FIELD).getValue()
+ self.output_strategy = context.getProperty(OUTPUT_STRATEGY).getValue()
+
+ ids_property = context.getProperty(INCLUDE_IDS)
+ self.include_ids = ids_property.asBoolean() if ids_property else False
+
+ embeddings_property = context.getProperty(INCLUDE_EMBEDDINGS)
+ self.include_embeddings = embeddings_property.asBoolean() if
embeddings_property else False
+
+ self.include_distances =
context.getProperty(INCLUDE_DISTANCES).asBoolean()
+
+ documents_property = context.getProperty(INCLUDE_DOCUMENTS)
+ self.include_documents = documents_property.asBoolean() if
documents_property else True
+ self.include_metadatas =
context.getProperty(INCLUDE_METADATAS).asBoolean()
+
+ def create_json(self, flowfile, documents, metadatas, embeddings,
distances, ids) -> tuple[str, str]:
+ input_json = None if self.results_field is None else
json.loads(flowfile.getContentsAsBytes().decode())
+
+ if self.output_strategy == TEXT:
+ # Delete any document that is None or an empty-string
+ documents = [doc for doc in documents if doc is not None and doc
!= ""]
+
+ # Join the documents with two newlines
+ text = "\n\n".join(documents)
+
+ # Create either JSON or text output, based on whether or not an
results field was specified
+ if input_json is None:
+ mime_type = "text/plain"
+ output_contents = text
+ else:
+ input_json[self.results_field] = text
+ output_contents = json.dumps(input_json)
+ mime_type = "application/json"
+ elif self.output_strategy == COLUMN_ORIENTED:
+ doc = {}
+ if self.include_ids:
+ doc["ids"] = ids
+ if self.include_distances:
+ doc["distances"] = distances
+ if self.include_documents:
+ doc["documents"] = documents
+ if self.include_metadatas:
+ doc["metadatas"] = metadatas
+ if self.include_embeddings:
+ doc["embeddings"] = embeddings
+
+ # Create the JSON from the Document
+ if input_json is None:
+ output_contents = json.dumps(doc)
+ else:
+ input_json[self.results_field] = doc
+ output_contents = json.dumps(input_json)
+
+ mime_type = "application/json"
+ else:
+ # Build the Documents
+ docs = []
+
+ count = len(ids) if ids else len(documents)
+ for i in range(count):
+ doc_id = None if ids is None else ids[i]
+ distance = None if distances is None else distances[i]
+ metadata = None if metadatas is None else metadatas[i]
+ document = None if documents is None else documents[i]
+ embedding = None if embeddings is None else embeddings[i]
+
+ # Create the document but do not include any key that we don't
want to include in the output.
+ doc = {}
+ if self.include_ids:
+ doc["id"] = doc_id
+ if self.include_distances:
+ doc["distance"] = distance
+ if self.include_documents:
+ doc["document"] = document
+ if self.include_metadatas:
+ doc["metadata"] = metadata
+ if self.include_embeddings:
+ doc["embedding"] = embedding
+
+ docs.append(doc)
+
+ # If input_json is None, we just create JSON based on the
Documents.
+ # If input_json is populated, we insert the documents into the
input JSON using the specified key.
+ if input_json is None:
+ jsons = []
+ for doc in docs:
+ jsons.append(json.dumps(doc))
+ output_contents = "\n".join(jsons)
+ else:
+ input_json[self.results_field] = docs
+ output_contents = json.dumps(input_json)
+
+ mime_type = "application/json"
+
+ return output_contents, mime_type
diff --git a/src/extensions/vectorstores/__init__.py
b/src/extensions/vectorstores/__init__.py
new file mode 100644
index 0000000..9881313
--- /dev/null
+++ b/src/extensions/vectorstores/__init__.py
@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0
diff --git a/src/extensions/vectorstores/requirements.txt
b/src/extensions/vectorstores/requirements.txt
new file mode 100644
index 0000000..31e215a
--- /dev/null
+++ b/src/extensions/vectorstores/requirements.txt
@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Shared requirements
+openai==1.9.0
+tiktoken
+langchain==0.1.11
+
+# Chroma requirements
+chromadb==0.4.22
+onnxruntime
+tokenizers
+tqdm
+requests
+
+# Pinecone requirements
+pinecone-client==3.0.1
+tiktoken
+langchain==0.1.11
+
+# OpenSearch requirements
+opensearch-py==2.5.0
+
+# Qdrant requirements
+qdrant-client==1.9.1
