From ef4a195e7a729e978a2d731c713086c6c69a8ab3 Mon Sep 17 00:00:00 2001
From: anovazzi1 <otavio2204@gmail.com>
Date: Tue, 16 Jan 2024 15:50:19 -0300
Subject: [PATCH] Update text splitter and vector stores

---
 .../textsplitters/CharacterTextSplitter.py    |  9 ++++----
 .../components/vectorstores/Chroma.py         |  2 +-
 .../langflow/components/vectorstores/FAISS.py | 23 +++++--------------
 3 files changed, 11 insertions(+), 23 deletions(-)

diff --git a/src/backend/langflow/components/textsplitters/CharacterTextSplitter.py b/src/backend/langflow/components/textsplitters/CharacterTextSplitter.py
index ba64b0f32..aa8fc77f2 100644
--- a/src/backend/langflow/components/textsplitters/CharacterTextSplitter.py
+++ b/src/backend/langflow/components/textsplitters/CharacterTextSplitter.py
@@ -1,6 +1,6 @@
 
 from langflow import CustomComponent
-from langchain.text_splitter import TextSplitter
+from langchain.text_splitter import CharacterTextSplitter
 from langchain_core.documents.base import Document
 from typing import List
 
@@ -23,10 +23,9 @@ class CharacterTextSplitterComponent(CustomComponent):
         chunk_overlap: int = 200,
         chunk_size: int = 1000,
         separator: str = "\n",
-    ) -> TextSplitter:
-        return TextSplitter(
-            documents=documents,
+    ) -> List[Document]:
+        return CharacterTextSplitter(
             chunk_overlap=chunk_overlap,
             chunk_size=chunk_size,
             separator=separator,
-        )
+        ).split_documents(documents)
diff --git a/src/backend/langflow/components/vectorstores/Chroma.py b/src/backend/langflow/components/vectorstores/Chroma.py
index f546eb361..adc3ed554 100644
--- a/src/backend/langflow/components/vectorstores/Chroma.py
+++ b/src/backend/langflow/components/vectorstores/Chroma.py
@@ -3,7 +3,7 @@ from typing import List, Optional, Union
 import chromadb  # type: ignore
 from langchain.embeddings.base import Embeddings
 from langchain.schema import BaseRetriever, Document
-from langchain.vectorstores import Chroma
+from langchain.vectorstores.chroma import Chroma
 from langchain.vectorstores.base import VectorStore
 
 from langflow import CustomComponent
diff --git a/src/backend/langflow/components/vectorstores/FAISS.py b/src/backend/langflow/components/vectorstores/FAISS.py
index 4d4864e73..2e516e7bc 100644
--- a/src/backend/langflow/components/vectorstores/FAISS.py
+++ b/src/backend/langflow/components/vectorstores/FAISS.py
@@ -1,11 +1,12 @@
 
 from langflow import CustomComponent
-from langchain.vectorstores import FAISS
-from typing import Optional, List
+from langchain_community.vectorstores.faiss import FAISS
+from typing import Optional, List, Union
+from langchain.schema import BaseRetriever
+from langchain.vectorstores.base import VectorStore
 from langflow.field_typing import (
     Document,
     Embeddings,
-    NestedDict,
 )
 
 class FAISSComponent(CustomComponent):
@@ -17,23 +18,11 @@ class FAISSComponent(CustomComponent):
         return {
             "documents": {"display_name": "Documents"},
             "embedding": {"display_name": "Embedding"},
-            "folder_path": {"display_name": "Local Path"},
-            "index_name": {"display_name": "Index Name"},
-            "search_kwargs": {"display_name": "Search Kwargs", "advanced": True},
         }
 
     def build(
         self,
         embedding: Embeddings,
         documents: Optional[List[Document]] = None,
-        folder_path: str = "",
-        index_name: str = "",
-        search_kwargs: Optional[NestedDict] = None,
-    ) -> FAISS:
-        return FAISS(
-            embedding=embedding,
-            documents=documents,
-            folder_path=folder_path,
-            index_name=index_name,
-            search_kwargs=search_kwargs or {},
-        )
+    ) -> Union[VectorStore,FAISS,BaseRetriever]:
+        return FAISS.from_documents(documents=documents,embedding=embedding)