Feat: add a support for OpenSearch and AstraDB components to yield the langchain vector_store connection object (#6998)

* Added decorator, decorator test, and modified supported vector stores

* Renamed module file name to reflect that this is for generic use, not use for graph rag

* Updated docsstring

* Improved documentation and modification to UT to support graph rag

* Remove extra file from PR

* rollback vector store template

* [autofix.ci] apply automated fixes

---------

Co-authored-by: Nadir J <31660040+NadirJ@users.noreply.github.com>
Co-authored-by: cristhianzl <cristhian.lousa@gmail.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

src/backend/base/langflow/base/vectorstores/vector_store_connection_decorator.py

@@ -0,0 +1,51 @@
+from langchain_core.vectorstores import VectorStore
+
+from langflow.io import Output
+
+
+def vector_store_connection(cls):
+    """A decorator that adds vector store connection functionality to a class.
+
+    This decorator performs the following actions:
+    1. Adds a `decorated` attribute to the class and sets it to True.
+    2. Checks if the class has an `outputs` attribute:
+        - If it does, it makes a copy of the `outputs` attribute to avoid modifying the base class attribute.
+        - It then checks if "vectorstoreconnection" is not already in the output names.
+        - If not, it adds a new `Output` entry for "Vector Store Connection" to the `outputs` attribute.
+    3. Adds an `as_vector_store` method to the class, which returns the result of the `build_vector_store` method.
+
+    Args:
+        cls (type): The class to be decorated.
+
+    Returns:
+        type: The decorated class with added vector store connection functionality.
+    """
+    cls.decorated = True  # Add an attribute to the class
+
+    if hasattr(cls, "outputs"):
+        cls.outputs = cls.outputs.copy()  # Make a copy to avoid modifying the base class attribute
+        output_names = [output.name for output in cls.outputs]
+
+        if "vectorstoreconnection" not in output_names:
+            cls.outputs.extend(
+                [
+                    Output(
+                        display_name="Vector Store Connection",
+                        hidden=True,
+                        name="vectorstoreconnection",
+                        method="as_vector_store",
+                    )
+                ]
+            )
+
+    def as_vector_store(self) -> VectorStore:
+        """Converts the current instance to a VectorStore object.
+
+        Returns:
+            VectorStore: The resulting VectorStore object.
+        """
+        return self.build_vector_store()
+
+    cls.as_vector_store = as_vector_store  # Ensures that the method is added to the class
+
+    return cls

src/backend/base/langflow/components/vectorstores/astradb.py

@@ -6,6 +6,7 @@ from astrapy.info import CollectionDescriptor
 from langchain_astradb import AstraDBVectorStore, CollectionVectorServiceOptions
 
 from langflow.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store
+from langflow.base.vectorstores.vector_store_connection_decorator import vector_store_connection
 from langflow.helpers import docs_to_data
 from langflow.inputs import FloatInput, NestedDictInput
 from langflow.io import (
@@ -20,6 +21,7 @@ from langflow.schema import Data
 from langflow.utils.version import get_version_info
 
 
+@vector_store_connection
 class AstraDBVectorStoreComponent(LCVectorStoreComponent):
     display_name: str = "Astra DB"
     description: str = "Ingest and search documents in Astra DB"

src/backend/base/langflow/components/vectorstores/opensearch.py

@@ -4,6 +4,7 @@ from typing import Any
 from langchain_community.vectorstores import OpenSearchVectorSearch
 
 from langflow.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store
+from langflow.base.vectorstores.vector_store_connection_decorator import vector_store_connection
 from langflow.io import (
     BoolInput,
     DropdownInput,
@@ -17,6 +18,7 @@ from langflow.io import (
 from langflow.schema import Data
 
 
+@vector_store_connection
 class OpenSearchVectorStoreComponent(LCVectorStoreComponent):
     """OpenSearch Vector Store with advanced, customizable search capabilities."""
 

src/backend/tests/unit/base/tools/test_vector_store_decorator.py

@@ -0,0 +1,36 @@
+from typing import Any
+
+import pytest
+from langflow.components.vectorstores import AstraDBVectorStoreComponent
+
+from tests.base import ComponentTestBaseWithoutClient, VersionComponentMapping
+
+
+class TestVectorStoreDecorator(ComponentTestBaseWithoutClient):
+    """Unit tests for the AstraDBVectorStoreComponent decorator.
+
+    This test class inherits from ComponentTestBaseWithoutClient and includes
+    the following tests and fixtures:
+
+    Fixtures:
+        - component_class: Returns the AstraDBVectorStoreComponent class to be tested.
+        - file_names_mapping: Returns an empty list representing the file names mapping for different versions.
+    Tests:
+        - test_decorator_applied: Verifies that the AstraDBVectorStoreComponent has the 'decorated' attribute and that
+        it is set to True.
+    """
+
+    @pytest.fixture
+    def component_class(self) -> type[Any]:
+        """Return the component class to test."""
+        return AstraDBVectorStoreComponent
+
+    @pytest.fixture
+    def file_names_mapping(self) -> list[VersionComponentMapping]:
+        """Return the file names mapping for different versions."""
+        return []
+
+    def test_decorator_applied(self, component_class: type[Any]):
+        component: AstraDBVectorStoreComponent = component_class()
+        assert hasattr(component, "decorated")
+        assert component.decorated

src/backend/tests/unit/initial_setup/starter_projects/test_vector_store_rag.py

@@ -30,7 +30,13 @@ def ingestion_graph():
     openai_embeddings.set(
         openai_api_key="sk-123", openai_api_base="https://api.openai.com/v1", openai_api_type="openai"
     )
+
     vector_store = AstraDBVectorStoreComponent(_id="ingestion-vector-store-123")
+
+    # Mock search_documents by changing the value otherwise set by the vector_store_connection_decorator
+    vector_store.set_on_output(name="vectorstoreconnection", value=[Data(text="This is a test file.")], cache=True)
+
+    vector_store.set_on_output(name="vectorstoreconnection", value=[Data(text="This is a test file.")], cache=True)
     vector_store.set_on_output(name="search_results", value=[Data(text="This is a test file.")], cache=True)
     vector_store.set_on_output(name="dataframe", value=DataFrame(data=[Data(text="This is a test file.")]), cache=True)
     vector_store.set(