diff --git a/src/backend/langflow/components/retrievers/VectorStoreRetriever.py b/src/backend/langflow/components/retrievers/VectorStoreRetriever.py new file mode 100644 index 000000000..576b02e69 --- /dev/null +++ b/src/backend/langflow/components/retrievers/VectorStoreRetriever.py @@ -0,0 +1,17 @@ +from langchain_core.vectorstores import VectorStoreRetriever + +from langflow import CustomComponent +from langflow.field_typing import VectorStore + + +class VectoStoreRetrieverComponent(CustomComponent): + display_name = "VectorStore Retriever" + description = "A vector store retriever" + + def build_config(self): + return { + "vectorstore": {"display_name": "Vector Store", "type": VectorStore}, + } + + def build(self, vectorstore: VectorStore) -> VectorStoreRetriever: + return vectorstore.as_retriever() diff --git a/src/backend/langflow/components/tools/RetrieverTool.py b/src/backend/langflow/components/tools/RetrieverTool.py new file mode 100644 index 000000000..5a598dbe2 --- /dev/null +++ b/src/backend/langflow/components/tools/RetrieverTool.py @@ -0,0 +1,32 @@ +from langchain.tools.retriever import create_retriever_tool + +from langflow import CustomComponent +from langflow.field_typing import BaseRetriever, Tool + + +class RetrieverToolComponent(CustomComponent): + display_name = "RetrieverTool" + description = "Tool for interacting with retriever" + + def build_config(self): + return { + "retriever": { + "display_name": "Retriever", + "info": "Retriever to interact with", + "type": BaseRetriever, + }, + "name": {"display_name": "Name", "info": "Name of the tool"}, + "description": {"display_name": "Description", "info": "Description of the tool"}, + } + + def build( + self, + retriever: BaseRetriever, + name: str, + description: str, + ) -> Tool: + return create_retriever_tool( + retriever=retriever, + name=name, + description=description, + ) diff --git a/src/backend/langflow/components/tools/__init__.py b/src/backend/langflow/components/tools/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/backend/langflow/components/vectorstores/MongoDBAtlasVectorSearch.py b/src/backend/langflow/components/vectorstores/MongoDBAtlasVectorSearch.py index d2d215f2b..13b62b708 100644 --- a/src/backend/langflow/components/vectorstores/MongoDBAtlasVectorSearch.py +++ b/src/backend/langflow/components/vectorstores/MongoDBAtlasVectorSearch.py @@ -1,22 +1,18 @@ from typing import List, Optional -from langchain_community.vectorstores import MongoDBAtlasVectorSearch +from langchain_community.vectorstores.mongodb_atlas import MongoDBAtlasVectorSearch from langflow import CustomComponent -from langflow.field_typing import ( - Document, - Embeddings, - NestedDict, -) +from langflow.field_typing import Document, Embeddings, NestedDict class MongoDBAtlasComponent(CustomComponent): display_name = "MongoDB Atlas" - description = "Construct a `MongoDB Atlas Vector Search` vector store from raw documents." + description = "a `MongoDB Atlas Vector Search` vector store from raw documents." def build_config(self): return { - "documents": {"display_name": "Documents"}, + "documents": {"display_name": "Documents", "is_list": True}, "embedding": {"display_name": "Embedding"}, "collection_name": {"display_name": "Collection Name"}, "db_name": {"display_name": "Database Name"}, @@ -27,8 +23,8 @@ class MongoDBAtlasComponent(CustomComponent): def build( self, - documents: List[Document], embedding: Embeddings, + documents: Optional[List[Document]] = None, collection_name: str = "", db_name: str = "", index_name: str = "", @@ -36,12 +32,17 @@ class MongoDBAtlasComponent(CustomComponent): search_kwargs: Optional[NestedDict] = None, ) -> MongoDBAtlasVectorSearch: search_kwargs = search_kwargs or {} - return MongoDBAtlasVectorSearch( - documents=documents, + vector_store = MongoDBAtlasVectorSearch.from_connection_string( + connection_string=mongodb_atlas_cluster_uri, + namespace=f"{db_name}.{collection_name}", embedding=embedding, - collection_name=collection_name, - db_name=db_name, index_name=index_name, - mongodb_atlas_cluster_uri=mongodb_atlas_cluster_uri, - search_kwargs=search_kwargs, ) + + if documents is not None: + if len(documents) == 0: + raise ValueError("If documents are provided, there must be at least one document.") + + vector_store.add_documents(documents) + + return vector_store