diff --git a/src/backend/base/langflow/components/vectorstores/__init__.py b/src/backend/base/langflow/components/vectorstores/__init__.py index 2a548c876..a459a99a8 100644 --- a/src/backend/base/langflow/components/vectorstores/__init__.py +++ b/src/backend/base/langflow/components/vectorstores/__init__.py @@ -9,6 +9,7 @@ from .elasticsearch import ElasticsearchVectorStoreComponent from .faiss import FaissVectorStoreComponent from .graph_rag import GraphRAGComponent from .hcd import HCDVectorStoreComponent +from .local_db import LocalDBComponent from .milvus import MilvusVectorStoreComponent from .mongodb_atlas import MongoVectorStoreComponent from .opensearch import OpenSearchVectorStoreComponent @@ -35,6 +36,7 @@ __all__ = [ "FaissVectorStoreComponent", "GraphRAGComponent", "HCDVectorStoreComponent", + "LocalDBComponent", "MilvusVectorStoreComponent", "MongoVectorStoreComponent", "OpenSearchVectorStoreComponent", diff --git a/src/backend/base/langflow/components/vectorstores/local_db.py b/src/backend/base/langflow/components/vectorstores/local_db.py new file mode 100644 index 000000000..07a5d6773 --- /dev/null +++ b/src/backend/base/langflow/components/vectorstores/local_db.py @@ -0,0 +1,255 @@ +from copy import deepcopy +from pathlib import Path + +from langchain_chroma import Chroma +from loguru import logger +from typing_extensions import override + +from langflow.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store +from langflow.base.vectorstores.utils import chroma_collection_to_data +from langflow.inputs.inputs import MultilineInput +from langflow.io import BoolInput, DropdownInput, HandleInput, IntInput, MessageTextInput, TabInput +from langflow.schema import Data, DataFrame +from langflow.template.field.base import Output + + +class LocalDBComponent(LCVectorStoreComponent): + """Chroma Vector Store with search capabilities.""" + + display_name: str = "Local DB" + description: str = "Local Vector Store with search capabilities" + name = "LocalDB" + icon = "database" + + inputs = [ + TabInput( + name="mode", + display_name="Mode", + options=["Ingest", "Retrieve"], + info="Select the operation mode", + value="Ingest", + real_time_refresh=True, + show=True, + ), + MessageTextInput( + name="collection_name", + display_name="Collection Name", + value="langflow", + ), + MessageTextInput( + name="persist_directory", + display_name="Persist Directory", + info=( + "Custom base directory to save the vector store. " + "Collections will be stored under '{directory}/vector_stores/{collection_name}'. " + "If not specified, it will use your system's cache folder." + ), + advanced=True, + ), + DropdownInput( + name="existing_collections", + display_name="Existing Collections", + options=[], # Will be populated dynamically + info="Select a previously created collection to search through its stored data.", + show=False, + combobox=True, + ), + HandleInput(name="embedding", display_name="Embedding", input_types=["Embeddings"]), + BoolInput( + name="allow_duplicates", + display_name="Allow Duplicates", + advanced=True, + info="If false, will not add documents that are already in the Vector Store.", + ), + DropdownInput( + name="search_type", + display_name="Search Type", + options=["Similarity", "MMR"], + value="Similarity", + advanced=True, + ), + HandleInput( + name="ingest_data", + display_name="Ingest Data", + input_types=["Data", "DataFrame"], + is_list=True, + info="Data to store. It will be embedded and indexed for semantic search.", + show=True, + ), + MultilineInput( + name="search_query", + display_name="Search Query", + tool_mode=True, + info="Enter text to search for similar content in the selected collection.", + show=False, + ), + IntInput( + name="number_of_results", + display_name="Number of Results", + info="Number of results to return.", + advanced=True, + value=10, + ), + IntInput( + name="limit", + display_name="Limit", + advanced=True, + info="Limit the number of records to compare when Allow Duplicates is False.", + ), + ] + outputs = [ + Output(display_name="DataFrame", name="dataframe", method="as_dataframe"), + ] + + def get_vector_store_directory(self, base_dir: str | Path) -> Path: + """Get the full directory path for a collection.""" + # Ensure base_dir is a Path object + base_dir = Path(base_dir) + # Create the full path: base_dir/vector_stores/collection_name + full_path = base_dir / "vector_stores" / self.collection_name + # Create the directory if it doesn't exist + full_path.mkdir(parents=True, exist_ok=True) + return full_path + + def get_default_persist_dir(self) -> str: + """Get the default persist directory from cache.""" + from langflow.services.cache.utils import CACHE_DIR + + return str(self.get_vector_store_directory(CACHE_DIR)) + + def list_existing_collections(self) -> list[str]: + """List existing vector store collections from the persist directory.""" + from langflow.services.cache.utils import CACHE_DIR + + # Get the base directory (either custom or cache) + base_dir = Path(self.persist_directory) if self.persist_directory else Path(CACHE_DIR) + # Get the vector_stores subdirectory + vector_stores_dir = base_dir / "vector_stores" + if not vector_stores_dir.exists(): + return [] + + return [d.name for d in vector_stores_dir.iterdir() if d.is_dir()] + + def update_build_config(self, build_config: dict, field_value: str, field_name: str | None = None) -> dict: + """Update the build configuration when the mode changes.""" + if field_name == "mode": + # Hide all dynamic fields by default + dynamic_fields = [ + "ingest_data", + "search_query", + "search_type", + "number_of_results", + "existing_collections", + "collection_name", + "embedding", + "allow_duplicates", + "limit", + ] + for field in dynamic_fields: + if field in build_config: + build_config[field]["show"] = False + + # Show/hide fields based on selected mode + if field_value == "Ingest": + if "ingest_data" in build_config: + build_config["ingest_data"]["show"] = True + if "collection_name" in build_config: + build_config["collection_name"]["show"] = True + build_config["collection_name"]["display_name"] = "Name Your Collection" + if "persist" in build_config: + build_config["persist"]["show"] = True + if "persist_directory" in build_config: + build_config["persist_directory"]["show"] = True + if "embedding" in build_config: + build_config["embedding"]["show"] = True + if "allow_duplicates" in build_config: + build_config["allow_duplicates"]["show"] = True + if "limit" in build_config: + build_config["limit"]["show"] = True + elif field_value == "Retrieve": + if "persist" in build_config: + build_config["persist"]["show"] = False + build_config["search_query"]["show"] = True + build_config["search_type"]["show"] = True + build_config["number_of_results"]["show"] = True + build_config["embedding"]["show"] = True + build_config["collection_name"]["show"] = False + # Show existing collections dropdown and update its options + if "existing_collections" in build_config: + build_config["existing_collections"]["show"] = True + build_config["existing_collections"]["options"] = self.list_existing_collections() + # Hide collection_name in Retrieve mode since we use existing_collections + elif field_name == "existing_collections": + # Update collection_name when an existing collection is selected + if "collection_name" in build_config: + build_config["collection_name"]["value"] = field_value + + return build_config + + @override + @check_cached_vector_store + def build_vector_store(self) -> Chroma: + """Builds the Chroma object.""" + try: + from langchain_chroma import Chroma + except ImportError as e: + msg = "Could not import Chroma integration package. Please install it with `pip install langchain-chroma`." + raise ImportError(msg) from e + # Chroma settings + # chroma_settings = None + if self.existing_collections: + self.collection_name = self.existing_collections + + # Use user-provided directory or default cache directory + if self.persist_directory: + base_dir = self.resolve_path(self.persist_directory) + persist_directory = str(self.get_vector_store_directory(base_dir)) + logger.debug(f"Using custom persist directory: {persist_directory}") + else: + persist_directory = self.get_default_persist_dir() + logger.debug(f"Using default persist directory: {persist_directory}") + + chroma = Chroma( + persist_directory=persist_directory, + client=None, + embedding_function=self.embedding, + collection_name=self.collection_name, + ) + + self._add_documents_to_vector_store(chroma) + self.status = chroma_collection_to_data(chroma.get(limit=self.limit)) + return chroma + + def _add_documents_to_vector_store(self, vector_store: "Chroma") -> None: + """Adds documents to the Vector Store.""" + ingest_data: list | Data | DataFrame = self.ingest_data + if not ingest_data: + self.status = "" + return + + # Convert DataFrame to Data if needed using parent's method + ingest_data = self._prepare_ingest_data() + + stored_documents_without_id = [] + if self.allow_duplicates: + stored_data = [] + else: + stored_data = chroma_collection_to_data(vector_store.get(limit=self.limit)) + for value in deepcopy(stored_data): + del value.id + stored_documents_without_id.append(value) + + documents = [] + for _input in ingest_data or []: + if isinstance(_input, Data): + if _input not in stored_documents_without_id: + documents.append(_input.to_lc_document()) + else: + msg = "Vector Store Inputs must be Data objects." + raise TypeError(msg) + + if documents and self.embedding is not None: + self.log(f"Adding {len(documents)} documents to the Vector Store.") + vector_store.add_documents(documents) + else: + self.log("No documents to add to the Vector Store.") diff --git a/src/backend/tests/unit/components/vectorstores/test_local_db_component.py b/src/backend/tests/unit/components/vectorstores/test_local_db_component.py new file mode 100644 index 000000000..428934e21 --- /dev/null +++ b/src/backend/tests/unit/components/vectorstores/test_local_db_component.py @@ -0,0 +1,382 @@ +import os +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from langflow.components.vectorstores.local_db import LocalDBComponent +from langflow.schema.data import Data +from langflow.services.cache.utils import CACHE_DIR + +from tests.base import ComponentTestBaseWithoutClient, VersionComponentMapping + + +@pytest.mark.api_key_required +class TestLocalDBComponent(ComponentTestBaseWithoutClient): + @pytest.fixture + def component_class(self) -> type[Any]: + """Return the component class to test.""" + return LocalDBComponent + + @pytest.fixture + def default_kwargs(self, tmp_path: Path) -> dict[str, Any]: + """Return the default kwargs for the component.""" + from langflow.components.embeddings.openai import OpenAIEmbeddingsComponent + + if os.getenv("OPENAI_API_KEY") is None: + pytest.skip("OPENAI_API_KEY is not set") + + api_key = os.getenv("OPENAI_API_KEY") + + return { + "embedding": OpenAIEmbeddingsComponent(openai_api_key=api_key).build_embeddings(), + "collection_name": "test_collection", + "persist": True, + "persist_directory": str(tmp_path), # Convert Path to string + "mode": "Ingest", + } + + @pytest.fixture + def file_names_mapping(self) -> list[VersionComponentMapping]: + """Return the file names mapping for different versions.""" + # Return an empty list since this is a new component + return [] + + def test_create_db(self, component_class: type[LocalDBComponent], default_kwargs: dict[str, Any]) -> None: + """Test creating a vector store.""" + component: LocalDBComponent = component_class().set(**default_kwargs) + component.build_vector_store() + persist_directory = Path(default_kwargs["persist_directory"]) + assert persist_directory.exists() + assert persist_directory.is_dir() + # Assert it isn't empty + assert len(list(persist_directory.iterdir())) > 0 + # Assert there's a chroma.sqlite3 file (since LocalDB uses Chroma underneath) + assert (persist_directory / "chroma.sqlite3").exists() + assert (persist_directory / "chroma.sqlite3").is_file() + + @patch("langchain_chroma.Chroma._collection") + def test_create_db_with_data( + self, + mock_collection, + component_class: type[LocalDBComponent], + default_kwargs: dict[str, Any], + ) -> None: + """Test creating a vector store with data.""" + # Set ingest_data in default_kwargs to a list of Data objects + test_texts = ["test data 1", "test data 2", "something completely different"] + default_kwargs["ingest_data"] = [Data(text=text) for text in test_texts] + + # Mock the collection count to return the expected number + mock_collection.count.return_value = len(test_texts) + mock_collection.name = default_kwargs["collection_name"] + + # Mock the _add_documents_to_vector_store method to ensure add_documents is called + with patch.object(LocalDBComponent, "_add_documents_to_vector_store") as mock_add_docs_method: + component: LocalDBComponent = component_class().set(**default_kwargs) + vector_store = component.build_vector_store() + + # Verify the method was called + mock_add_docs_method.assert_called_once() + + # Verify collection exists and has the correct data + assert vector_store._collection.name == default_kwargs["collection_name"] + assert vector_store._collection.count() == len(test_texts) + + def test_default_persist_dir(self, component_class: type[LocalDBComponent], default_kwargs: dict[str, Any]) -> None: + """Test the default persist directory functionality.""" + # Remove persist_directory from default_kwargs to test default directory + default_kwargs.pop("persist_directory") + + component: LocalDBComponent = component_class().set(**default_kwargs) + + # Call get_default_persist_dir and check the result + default_dir = component.get_default_persist_dir() + expected_dir = Path(CACHE_DIR) / "vector_stores" / default_kwargs["collection_name"] + + assert Path(default_dir) == expected_dir + assert Path(default_dir).exists() + + @patch("langchain_chroma.Chroma.similarity_search") + def test_similarity_search( + self, + mock_similarity_search, + component_class: type[LocalDBComponent], + default_kwargs: dict[str, Any], + ) -> None: + """Test the similarity search functionality.""" + # Create test data with distinct topics + test_data = [ + "The quick brown fox jumps over the lazy dog", + "Python is a popular programming language", + "Machine learning models process data", + "The lazy dog sleeps all day long", + ] + default_kwargs["ingest_data"] = [Data(text=text) for text in test_data] + default_kwargs["search_type"] = "Similarity" + default_kwargs["number_of_results"] = 2 + + # Mock the similarity_search to return documents + from langchain_core.documents import Document + + mock_docs = [ + Document(page_content="The lazy dog sleeps all day long"), + Document(page_content="The quick brown fox jumps over the lazy dog"), + ] + mock_similarity_search.return_value = mock_docs + + component: LocalDBComponent = component_class().set(**default_kwargs) + component.build_vector_store() + + # Switch to Retrieve mode + component.set(mode="Retrieve", search_query="dog sleeping") + results = component.search_documents() + + assert len(results) == 2 + # The most relevant results should be about dogs + assert any("dog" in result.text.lower() for result in results) + mock_similarity_search.assert_called_once_with(query="dog sleeping", k=2) + + # Test with different number of results + component.set(number_of_results=3) + another_doc = Document(page_content="Another document") + mock_similarity_search.return_value = [*mock_docs, another_doc] # Use unpacking instead of concatenation + results = component.search_documents() + assert len(results) == 3 + + @patch("langchain_chroma.Chroma.max_marginal_relevance_search") + def test_mmr_search( + self, + mock_mmr_search, + component_class: type[LocalDBComponent], + default_kwargs: dict[str, Any], + ) -> None: + """Test the MMR search functionality.""" + # Create test data with some similar documents + test_data = [ + "The quick brown fox jumps", + "The quick brown fox leaps", + "The quick brown fox hops", + "Something completely different about cats", + ] + default_kwargs["ingest_data"] = [Data(text=text) for text in test_data] + default_kwargs["search_type"] = "MMR" + default_kwargs["number_of_results"] = 3 + + # Mock the MMR search to return documents + from langchain_core.documents import Document + + mock_docs = [ + Document(page_content="The quick brown fox jumps"), + Document(page_content="The quick brown fox leaps"), + Document(page_content="Something completely different about cats"), + ] + mock_mmr_search.return_value = mock_docs + + component: LocalDBComponent = component_class().set(**default_kwargs) + component.build_vector_store() + + # Switch to Retrieve mode + component.set(mode="Retrieve", search_query="quick fox") + results = component.search_documents() + + assert len(results) == 3 + # Results should be diverse but relevant + assert any("fox" in result.text.lower() for result in results) + mock_mmr_search.assert_called_once_with(query="quick fox", k=3) + + # Test with different settings + component.set(number_of_results=2) + mock_mmr_search.return_value = mock_docs[:2] + diverse_results = component.search_documents() + assert len(diverse_results) == 2 + + @patch("langchain_chroma.Chroma.similarity_search") + @patch("langchain_chroma.Chroma.max_marginal_relevance_search") + def test_search_with_different_types( + self, + mock_mmr_search, + mock_similarity_search, + component_class: type[LocalDBComponent], + default_kwargs: dict[str, Any], + ) -> None: + """Test search with different search types.""" + test_data = [ + "The quick brown fox jumps over the lazy dog", + "Python is a popular programming language", + "Machine learning models process data", + ] + default_kwargs["ingest_data"] = [Data(text=text) for text in test_data] + default_kwargs["number_of_results"] = 2 + + # Mock the search methods to return documents + from langchain_core.documents import Document + + mock_similarity_docs = [ + Document(page_content="Python is a popular programming language"), + Document(page_content="Machine learning models process data"), + ] + mock_similarity_search.return_value = mock_similarity_docs + + mock_mmr_docs = [ + Document(page_content="Python is a popular programming language"), + Document(page_content="The quick brown fox jumps over the lazy dog"), + ] + mock_mmr_search.return_value = mock_mmr_docs + + component: LocalDBComponent = component_class().set(**default_kwargs) + component.build_vector_store() + + # Switch to Retrieve mode and test similarity search + component.set(mode="Retrieve", search_type="Similarity", search_query="programming languages") + similarity_results = component.search_documents() + assert len(similarity_results) == 2 + assert any("python" in result.text.lower() for result in similarity_results) + mock_similarity_search.assert_called_once_with(query="programming languages", k=2) + + # Test MMR search + component.set(search_type="MMR", search_query="programming languages") + mmr_results = component.search_documents() + assert len(mmr_results) == 2 + mock_mmr_search.assert_called_once_with(query="programming languages", k=2) + + # Test with empty query + component.set(search_query="") + empty_results = component.search_documents() + assert len(empty_results) == 0 + + @patch("langchain_chroma.Chroma.get") + @patch("langchain_chroma.Chroma._collection") + def test_duplicate_handling( + self, + mock_collection, + mock_get, + component_class: type[LocalDBComponent], + default_kwargs: dict[str, Any], + ) -> None: + """Test handling of duplicate documents.""" + # Create test data with duplicates + test_data = [ + Data(text_key="text", data={"text": "This is a test document"}), + Data(text_key="text", data={"text": "This is a test document"}), # Duplicate with exact same data + Data(text_key="text", data={"text": "This is another document"}), + ] + default_kwargs["ingest_data"] = test_data + default_kwargs["allow_duplicates"] = False + default_kwargs["limit"] = 100 # Set a high enough limit to get all documents + + # Mock the get method to return documents + mock_get.return_value = { + "documents": ["This is a test document", "This is a test document", "This is another document"], + "metadatas": [{}, {}, {}], + "ids": ["1", "2", "3"], + } + + # Mock collection count + mock_collection.count.return_value = 3 + + component: LocalDBComponent = component_class().set(**default_kwargs) + vector_store = component.build_vector_store() + + # Get all documents + results = vector_store.get(limit=100) + documents = results["documents"] + + # The documents are returned in a list structure + assert len(documents) == 3 # All documents are added, even duplicates + + # Count unique texts + unique_texts = set(documents) + assert len(unique_texts) == 2 # Should have 2 unique texts + + # Test with allow_duplicates=True + test_data = [ + Data(text_key="text", data={"text": "This is a test document"}), + Data(text_key="text", data={"text": "This is a test document"}), # Duplicate + ] + default_kwargs["ingest_data"] = test_data + default_kwargs["allow_duplicates"] = True + default_kwargs["collection_name"] = "test_collection_2" # Use a different collection name + + # Mock for the second test + mock_get.return_value = { + "documents": ["This is a test document", "This is a test document"], + "metadatas": [{}, {}], + "ids": ["1", "2"], + } + mock_collection.count.return_value = 2 + + component = component_class().set(**default_kwargs) + vector_store = component.build_vector_store() + + # Get all documents + results = vector_store.get(limit=100) + documents = results["documents"] + + # With allow_duplicates=True, we should have both documents + assert len(documents) == 2 + assert all("test document" in doc for doc in documents) + + # Verify that we have the expected number of documents + assert vector_store._collection.count() == 2 + + def test_build_config_update(self, component_class: type[LocalDBComponent]) -> None: + """Test the update_build_config method.""" + component = component_class() + + # Test mode=Ingest + build_config = { + "ingest_data": {"show": False}, + "collection_name": {"show": False}, + "persist": {"show": False}, + "persist_directory": {"show": False}, + "embedding": {"show": False}, + "allow_duplicates": {"show": False}, + "limit": {"show": False}, + "search_query": {"show": False}, + "search_type": {"show": False}, + "number_of_results": {"show": False}, + "existing_collections": {"show": False}, + } + + updated_config = component.update_build_config(build_config, "Ingest", "mode") + + assert updated_config["ingest_data"]["show"] is True + assert updated_config["collection_name"]["show"] is True + assert updated_config["persist"]["show"] is True + assert updated_config["search_query"]["show"] is False + + # Test mode=Retrieve + updated_config = component.update_build_config(build_config, "Retrieve", "mode") + + assert updated_config["search_query"]["show"] is True + assert updated_config["search_type"]["show"] is True + assert updated_config["number_of_results"]["show"] is True + assert updated_config["existing_collections"]["show"] is True + assert updated_config["collection_name"]["show"] is False + + # Test persist=True/False + build_config = {"persist_directory": {"show": False}} + # Use keyword arguments to fix FBT003 + updated_config = component.update_build_config(build_config, field_value=True, field_name="persist") + assert updated_config["persist_directory"]["show"] is True + + updated_config = component.update_build_config(build_config, field_value=False, field_name="persist") + assert updated_config["persist_directory"]["show"] is False + + # Test existing_collections update + # Fix the dict entry type issue + build_config = {"collection_name": {"value": "old_name", "show": False}} + updated_config = component.update_build_config(build_config, "new_collection", "existing_collections") + assert updated_config["collection_name"]["value"] == "new_collection" + + @patch("langflow.components.vectorstores.local_db.LocalDBComponent.list_existing_collections") + def test_list_existing_collections(self, mock_list: MagicMock, component_class: type[LocalDBComponent]) -> None: + """Test the list_existing_collections method.""" + mock_list.return_value = ["collection1", "collection2", "collection3"] + + component = component_class() + collections = component.list_existing_collections() + + assert collections == ["collection1", "collection2", "collection3"] + mock_list.assert_called_once()