⚡️ Speed up method AstraDBVectorStoreComponent.reset_database_list by 36% in PR #6236 (LFOSS-492)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #6236

If you approve this dependent PR, these changes will be merged into the original PR branch LFOSS-492.

This PR will be automatically closed if the original PR is merged.

---

📄 36% (0.36x) speedup for AstraDBVectorStoreComponent.reset_database_list in src/backend/base/langflow/components/vectorstores/astradb.py

⏱️ Runtime : 1.87 millisecond → 1.38 millisecond (best of 131 runs)

📝 Explanation and details

To optimize the given Python program for better runtime performance, I'll focus on restructuring the code to reduce redundant work and aim for more efficient data access patterns where applicable. While the initial implementation is not inherently slow, there are some improvements that can be made.

Here's the optimized version of your code.

Changes and Optimizations.

1. Database List Fetching.

   • Saved the result of self.get_database_list() to database_list to avoid multiple calls and improve readability.

2. Optimized Metadata List Comprehension.

   • Used specific key extraction within the dictionary comprehension for options_metadata to make it clearer and possibly slightly more efficient.

3. Variable Assignment.

   • Separated the list of database names and metadata dictionary creation to improve readability and simplify checking conditions on the list of database names.

These changes slightly improve the clarity and performance of the code by reducing redundant operations and making the code more maintainable.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 18 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	undefined

🌀 Generated Regression Tests Details

import pytest  # used for our unit tests
# function to test
from langflow.base.vectorstores.model import LCVectorStoreComponent
from langflow.components.vectorstores.astradb import \
    AstraDBVectorStoreComponent

# unit tests

@pytest.fixture
def astra_db_component():
    class MockAstraDBVectorStoreComponent(AstraDBVectorStoreComponent):
        def get_database_list(self):
            return {
                "db1": {"status": "active", "collections": 3, "api_endpoint": "http://api1"},
                "db2": {"status": "inactive", "collections": 5, "api_endpoint": "http://api2"},
            }
    return MockAstraDBVectorStoreComponent()

def test_valid_token_and_database_list(astra_db_component):
    build_config = {
        "token": {"value": "valid_token"},
        "database_name": {"value": "db1", "options": [], "options_metadata": [], "advanced": True},
        "api_endpoint": {"value": ""},
        "collection_name": {"advanced": True}
    }
    codeflash_output = astra_db_component.reset_database_list(build_config)

def test_empty_database_list(astra_db_component):
    class MockAstraDBVectorStoreComponentEmpty(AstraDBVectorStoreComponent):
        def get_database_list(self):
            return {}
    empty_component = MockAstraDBVectorStoreComponentEmpty()
    build_config = {
        "token": {"value": "valid_token"},
        "database_name": {"value": "db1", "options": [], "options_metadata": [], "advanced": True},
        "api_endpoint": {"value": ""},
        "collection_name": {"advanced": True}
    }
    codeflash_output = empty_component.reset_database_list(build_config)

def test_no_token_provided(astra_db_component):
    build_config = {
        "token": {"value": ""},
        "database_name": {"value": "db1", "options": [], "options_metadata": [], "advanced": True},
        "api_endpoint": {"value": ""},
        "collection_name": {"advanced": True}
    }
    codeflash_output = astra_db_component.reset_database_list(build_config)

def test_invalid_selected_database(astra_db_component):
    build_config = {
        "token": {"value": "valid_token"},
        "database_name": {"value": "invalid_db", "options": [], "options_metadata": [], "advanced": True},
        "api_endpoint": {"value": ""},
        "collection_name": {"advanced": True}
    }
    codeflash_output = astra_db_component.reset_database_list(build_config)

def test_exception_in_initialize_database_options():
    class MockAstraDBVectorStoreComponentError(AstraDBVectorStoreComponent):
        def get_database_list(self):
            raise Exception("Test Exception")
    error_component = MockAstraDBVectorStoreComponentError()
    build_config = {
        "token": {"value": "valid_token"},
        "database_name": {"value": "db1", "options": [], "options_metadata": [], "advanced": True},
        "api_endpoint": {"value": ""},
        "collection_name": {"advanced": True}
    }
    with pytest.raises(ValueError, match="Error fetching database options: Test Exception"):
        error_component.reset_database_list(build_config)

def test_large_number_of_databases():
    class MockAstraDBVectorStoreComponentLarge(AstraDBVectorStoreComponent):
        def get_database_list(self):
            return {f"db{i}": {"status": "active", "collections": i, "api_endpoint": f"http://api{i}"} for i in range(1000)}
    large_component = MockAstraDBVectorStoreComponentLarge()
    build_config = {
        "token": {"value": "valid_token"},
        "database_name": {"value": "db1", "options": [], "options_metadata": [], "advanced": True},
        "api_endpoint": {"value": ""},
        "collection_name": {"advanced": True}
    }
    codeflash_output = large_component.reset_database_list(build_config)

def test_correct_metadata_assignment(astra_db_component):
    build_config = {
        "token": {"value": "valid_token"},
        "database_name": {"value": "db1", "options": [], "options_metadata": [], "advanced": True},
        "api_endpoint": {"value": ""},
        "collection_name": {"advanced": True}
    }
    codeflash_output = astra_db_component.reset_database_list(build_config)

def test_consistent_output(astra_db_component):
    build_config = {
        "token": {"value": "valid_token"},
        "database_name": {"value": "db1", "options": [], "options_metadata": [], "advanced": True},
        "api_endpoint": {"value": ""},
        "collection_name": {"advanced": True}
    }
    codeflash_output = astra_db_component.reset_database_list(build_config)
    codeflash_output = astra_db_component.reset_database_list(build_config)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests
# function to test
from langflow.base.vectorstores.model import LCVectorStoreComponent
from langflow.components.vectorstores.astradb import \
    AstraDBVectorStoreComponent


# unit tests
class TestAstraDBVectorStoreComponent:
    @pytest.fixture
    def component(self, mocker):
        component = AstraDBVectorStoreComponent()
        mocker.patch.object(component, 'get_database_list', return_value={
            "db1": {"status": "active", "collections": 5, "api_endpoint": "http://example.com/db1"},
            "db2": {"status": "inactive", "collections": 3, "api_endpoint": "http://example.com/db2"}
        })
        return component

    def test_basic_functionality_with_token(self, component):
        build_config = {
            "database_name": {"value": "db1", "options": [], "options_metadata": [], "advanced": False},
            "api_endpoint": {"value": ""},
            "collection_name": {"advanced": False},
            "token": {"value": True}
        }
        codeflash_output = component.reset_database_list(build_config)

    def test_basic_functionality_without_token(self, component):
        build_config = {
            "database_name": {"value": "db1", "options": [], "options_metadata": [], "advanced": False},
            "api_endpoint": {"value": ""},
            "collection_name": {"advanced": False},
            "token": {"value": False}
        }
        codeflash_output = component.reset_database_list(build_config)

    def test_empty_build_config(self, component):
        build_config = {}
        with pytest.raises(KeyError):
            component.reset_database_list(build_config)