From ce17bd4e71f69ddc422a68d156fc12ae34d7600d Mon Sep 17 00:00:00 2001
From: ytkimirti <yusuftaha9@gmail.com>
Date: Wed, 14 Aug 2024 17:18:45 +0300
Subject: [PATCH] feat: Add support for metadata filtering and namespaces for
 the Upstash Vector component (#3254)

* feat: add metadata filtering and namespace support for the upstash vector component

* docs: add upstash vector to the vectorstores doc

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 .../Components/components-vector-stores.md    | 21 +++++++++++
 .../components/vectorstores/Upstash.py        | 35 +++++++++++++++++--
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/docs/docs/Components/components-vector-stores.md b/docs/docs/Components/components-vector-stores.md
index ea000342d..02a219f42 100644
--- a/docs/docs/Components/components-vector-stores.md
+++ b/docs/docs/Components/components-vector-stores.md
@@ -472,6 +472,27 @@ Ensure the Supabase service key, URL, and table name are properly configured.
 ---
 
 
+### Upstash Vector
+
+
+`UpstashVector` searches a Upstash Vector Store for documents similar to the input. It has it's own embedding
+model which can be used to search documents without needing an external embedding model.
+
+
+**Parameters:**
+
+- **Index URL:** The URL of the Upstash index.
+- **Index Token:** The token for the Upstash index.
+- **Text Key:** The key in the record to use as text.
+- **Namespace:** The namespace name. A new namespace is created if not found. Leave empty for default namespace.
+- **Search Query:** The search query.
+- **Metadata Filter:** The metadata filter. Filters documents by metadata. Look at the [docs](https://upstash.com/docs/vector/features/filtering) for more information.
+- **Embedding:** The embedding model used. To use Upstash's embeddings, don't provide an embedding.
+- **Number of Results:** The number of results to return.
+
+---
+
+
 ### Vectara {#b4e05230b62a47c792a89c5511af97ac}
 
 
diff --git a/src/backend/base/langflow/components/vectorstores/Upstash.py b/src/backend/base/langflow/components/vectorstores/Upstash.py
index bc39bde83..341e309df 100644
--- a/src/backend/base/langflow/components/vectorstores/Upstash.py
+++ b/src/backend/base/langflow/components/vectorstores/Upstash.py
@@ -4,7 +4,14 @@ from langchain_community.vectorstores import UpstashVectorStore
 
 from langflow.base.vectorstores.model import LCVectorStoreComponent
 from langflow.helpers.data import docs_to_data
-from langflow.io import HandleInput, IntInput, StrInput, SecretStrInput, DataInput, MultilineInput
+from langflow.io import (
+    HandleInput,
+    IntInput,
+    StrInput,
+    SecretStrInput,
+    DataInput,
+    MultilineInput,
+)
 from langflow.schema import Data
 
 
@@ -16,9 +23,17 @@ class UpstashVectorStoreComponent(LCVectorStoreComponent):
     icon = "Upstash"
 
     inputs = [
-        StrInput(name="index_url", display_name="Index URL", info="The URL of the Upstash index.", required=True),
+        StrInput(
+            name="index_url",
+            display_name="Index URL",
+            info="The URL of the Upstash index.",
+            required=True,
+        ),
         SecretStrInput(
-            name="index_token", display_name="Index Token", info="The token for the Upstash index.", required=True
+            name="index_token",
+            display_name="Index Token",
+            info="The token for the Upstash index.",
+            required=True,
         ),
         StrInput(
             name="text_key",
@@ -27,7 +42,17 @@ class UpstashVectorStoreComponent(LCVectorStoreComponent):
             value="text",
             advanced=True,
         ),
+        StrInput(
+            name="namespace",
+            display_name="Namespace",
+            info="Leave empty for default namespace.",
+        ),
         MultilineInput(name="search_query", display_name="Search Query"),
+        MultilineInput(
+            name="metadata_filter",
+            display_name="Metadata Filter",
+            info="Filters documents by metadata. Look at the documentation for more information.",
+        ),
         DataInput(
             name="ingest_data",
             display_name="Ingest Data",
@@ -68,6 +93,7 @@ class UpstashVectorStoreComponent(LCVectorStoreComponent):
                     text_key=self.text_key,
                     index_url=self.index_url,
                     index_token=self.index_token,
+                    namespace=self.namespace,
                 )
                 upstash_vs.add_documents(documents)
             else:
@@ -77,6 +103,7 @@ class UpstashVectorStoreComponent(LCVectorStoreComponent):
                     text_key=self.text_key,
                     index_url=self.index_url,
                     index_token=self.index_token,
+                    namespace=self.namespace,
                 )
         else:
             upstash_vs = UpstashVectorStore(
@@ -84,6 +111,7 @@ class UpstashVectorStoreComponent(LCVectorStoreComponent):
                 text_key=self.text_key,
                 index_url=self.index_url,
                 index_token=self.index_token,
+                namespace=self.namespace,
             )
 
         return upstash_vs
@@ -95,6 +123,7 @@ class UpstashVectorStoreComponent(LCVectorStoreComponent):
             docs = vector_store.similarity_search(
                 query=self.search_query,
                 k=self.number_of_results,
+                filter=self.metadata_filter,
             )
 
             data = docs_to_data(docs)