From ff7c3c26ef480637f8656180609de104c21ec4b6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Di=C3=B3genes=20Betat=20Roos?= <diogenes@diogenie.com.br>
Date: Thu, 17 Aug 2023 19:12:21 +0000
Subject: [PATCH] Migrate Chroma to Custom Component

---
 .../components/vectorstores/Chroma.py         | 69 +++++++++++++++++++
 1 file changed, 69 insertions(+)
 create mode 100644 src/backend/langflow/components/vectorstores/Chroma.py

diff --git a/src/backend/langflow/components/vectorstores/Chroma.py b/src/backend/langflow/components/vectorstores/Chroma.py
new file mode 100644
index 000000000..d064af80e
--- /dev/null
+++ b/src/backend/langflow/components/vectorstores/Chroma.py
@@ -0,0 +1,69 @@
+from typing import Optional, Union
+from langflow import CustomComponent
+
+from langchain.vectorstores import Chroma
+from langchain.schema import Document
+from langchain.vectorstores.base import VectorStore
+from langchain.schema import BaseRetriever
+from langchain.embeddings.base import Embeddings
+
+
+class ChromaComponent(CustomComponent):
+    """
+    A custom component for implementing a Vector Store using Chroma.
+    """
+
+    display_name: str = "Chroma (Custom Component)"
+    description: str = "Implementation of Vector Store using Chroma"
+    documentation = "https://python.langchain.com/docs/integrations/vectorstores/chroma"
+    beta = True
+
+    def build_config(self):
+        """
+        Builds the configuration for the component.
+
+        Returns:
+        - dict: A dictionary containing the configuration options for the component.
+        """
+        return {
+            "collection_name": {"display_name": "Collection Name", "value": "langflow"},
+            "persist": {"display_name": "Persist"},
+            "persist_directory": {"display_name": "Persist Directory"},
+            "code": {"show": False, "display_name": "Code"},
+            "documents": {"display_name": "Documents", "is_list": True},
+            "embedding": {"display_name": "Embedding"},
+        }
+
+    def build(
+        self,
+        collection_name: str,
+        persist: bool,
+        persist_directory: Optional[str] = None,
+        embedding: Optional[Embeddings] = None,
+        documents: Optional[Document] = None,
+    ) -> Union[VectorStore, BaseRetriever]:
+        """
+        Builds the Vector Store or BaseRetriever object.
+
+        Args:
+        - collection_name (str): The name of the collection.
+        - persist_directory (Optional[str]): The directory to persist the Vector Store to.
+        - persist (bool): Whether to persist the Vector Store or not.
+        - embedding (Optional[Embeddings]): The embeddings to use for the Vector Store.
+        - documents (Optional[Document]): The documents to use for the Vector Store.
+
+        Returns:
+        - Union[VectorStore, BaseRetriever]: The Vector Store or BaseRetriever object.
+        """
+        # If documents, then we need to create a Chroma instance using .from_documents
+        if documents is not None and embedding is not None:
+            return Chroma.from_documents(
+                documents=documents,  # type: ignore
+                persist_directory=persist_directory if persist else None,
+                collection_name=collection_name,
+                embedding=embedding,
+            )
+
+        return Chroma(
+            persist_directory=persist_directory,
+        )