From 24c52758990244bd2ee8ac6bb41657cfeacbae46 Mon Sep 17 00:00:00 2001
From: Rodrigo <rodrigosilvanader@gmail.com>
Date: Tue, 18 Jun 2024 12:55:36 -0300
Subject: [PATCH] refactor: Add SelectivePassThroughComponent to experimental
 components and update ChromaVectorStoreComponent

---
 .../base/langflow/components/experimental/__init__.py    | 3 +++
 .../base/langflow/components/vectorstores/Chroma.py      | 9 +++++----
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/backend/base/langflow/components/experimental/__init__.py b/src/backend/base/langflow/components/experimental/__init__.py
index 550d4d6dd..15c0806db 100644
--- a/src/backend/base/langflow/components/experimental/__init__.py
+++ b/src/backend/base/langflow/components/experimental/__init__.py
@@ -13,6 +13,8 @@ from .SplitText import SplitTextComponent
 from .SQLExecutor import SQLExecutorComponent
 from .SubFlow import SubFlowComponent
 from .ConditionalRouter import ConditionalRouterComponent
+from .SelectivePassThrough import SelectivePassThroughComponent
+
 
 __all__ = [
     "AgentComponent",
@@ -30,4 +32,5 @@ __all__ = [
     "SplitTextComponent",
     "SQLExecutorComponent",
     "SubFlowComponent",
+    "SelectivePassThroughComponent",
 ]
diff --git a/src/backend/base/langflow/components/vectorstores/Chroma.py b/src/backend/base/langflow/components/vectorstores/Chroma.py
index 0fc167dbb..ae9b08ede 100644
--- a/src/backend/base/langflow/components/vectorstores/Chroma.py
+++ b/src/backend/base/langflow/components/vectorstores/Chroma.py
@@ -7,7 +7,7 @@ from loguru import logger
 
 from langflow.base.vectorstores.model import LCVectorStoreComponent
 from langflow.base.vectorstores.utils import chroma_collection_to_data
-from langflow.inputs import BoolInput, DataInput, DropdownInput, HandleInput, IntInput, MessageInput, StrInput
+from langflow.inputs import BoolInput, DataInput, DropdownInput, HandleInput, IntInput, MessageInput, StrInput, TextInput
 from langflow.schema import Data
 
 if TYPE_CHECKING:
@@ -39,7 +39,7 @@ class ChromaVectorStoreComponent(LCVectorStoreComponent):
             display_name="Code",
             advanced=True,
         ),
-        MessageInput(
+        TextInput(
             name="search_query",
             display_name="Search Query",
             is_list=True,
@@ -139,6 +139,7 @@ class ChromaVectorStoreComponent(LCVectorStoreComponent):
             collection_name=self.collection_name,
         )
 
+        self._add_documents_to_vector_store(chroma)
         self.status = chroma_collection_to_data(chroma.get(self.limit))
         return chroma
 
@@ -177,7 +178,7 @@ class ChromaVectorStoreComponent(LCVectorStoreComponent):
         """
         Search for documents in the Chroma vector store.
         """
-        if not self.search_query.text:
+        if not self.search_query:
             self.status = ""
             return
 
@@ -194,6 +195,6 @@ class ChromaVectorStoreComponent(LCVectorStoreComponent):
                 self.search_query = self.search_query[0]
 
         search_results = self.search_with_vector_store(
-            self.search_query.text, self.search_type, vector_store, k=self.number_of_results
+            self.search_query, self.search_type, vector_store, k=self.number_of_results
         )
         return search_results