⚡️ Speed up method AstraDBVectorStoreComponent._initialize_collection_options by 12% in PR #6028 (PlaygroundPage)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #6028

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

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

---

📄 12% (0.12x) speedup for AstraDBVectorStoreComponent._initialize_collection_options in src/backend/base/langflow/components/vectorstores/astradb.py

⏱️ Runtime : 509 milliseconds → 453 milliseconds (best of 5 runs)

📝 Explanation and details

Certainly! To optimize this code for better performance, we'll minimize repeated operations, reduce the number of API calls by caching results where appropriate, and employ efficient error handling and logging to ensure minimal impact on run time.

Here's the optimized code.

Explanation of Optimization.

1. Client Initialization.

   • The client is initialized once and reused across multiple calls to avoid repeated object creation.

2. Database Object Caching.

   • Database object is retrieved and cached on first access. Subsequent accesses reuse the cached database object. This avoids making multiple calls to fetch the same database object.

3. Error Logging.

   • Error logging has been kept minimal and occurs only when an exception truly affects the execution.

4. Loop Optimizations.

   • The collection options initialization function now directly builds a results list minimizing the overhead of list comprehensions.

5. Keyspace Retrieval.

   • A quick keyspace check is done once per function call, stored in a local variable, and reused within that scope.

By caching objects and reducing the overhead of repeated function calls and object creations, the optimized code will run more efficiently in terms of both runtime and resource usage.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from unittest.mock import MagicMock, patch

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

# unit tests

@pytest.fixture
def mock_component():
    component = AstraDBVectorStoreComponent()
    component.token = "fake_token"
    component.environment = "fake_environment"
    component.keyspace = "fake_keyspace"
    component.log = MagicMock()
    return component



def test_empty_database(mock_component):
    # Mock the database with no collections
    mock_database = MagicMock()
    mock_database.list_collections.return_value = []
    mock_component.get_database_object = MagicMock(return_value=mock_database)
    
    codeflash_output = mock_component._initialize_collection_options()


def test_invalid_api_endpoint(mock_component):
    # Simulate an invalid API endpoint
    mock_component.get_database_object = MagicMock(side_effect=ValueError("Invalid API endpoint"))
    
    with pytest.raises(ValueError, match="Invalid API endpoint"):
        mock_component._initialize_collection_options(api_endpoint="invalid_endpoint")







from unittest.mock import MagicMock

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


# unit tests
class TestAstraDBVectorStoreComponent:
    @pytest.fixture
    def mock_component(self):
        component = AstraDBVectorStoreComponent()
        component.token = "mock_token"
        component.environment = "mock_environment"
        component.keyspace = "mock_keyspace"
        return component

    @pytest.fixture
    def mock_database(self):
        database = MagicMock()
        collection1 = MagicMock()
        collection1.name = "col1"
        collection1.options.vector.service.provider = "provider1"
        collection1.options.vector.service.model_name = "model1"
        collection2 = MagicMock()
        collection2.name = "col2"
        collection2.options.vector.service.provider = "provider2"
        collection2.options.vector.service.model_name = "model2"
        database.list_collections.return_value = [collection1, collection2]
        database.get_collection.return_value.estimated_document_count.return_value = 10
        return database

    def test_basic_functionality_single_collection(self, mock_component, mock_database):
        # Mock methods
        mock_component.get_database_object = MagicMock(return_value=mock_database)
        mock_component.collection_data = MagicMock(return_value=10)

        codeflash_output = mock_component._initialize_collection_options()

    def test_edge_case_no_collections(self, mock_component):
        # Mock methods
        mock_database = MagicMock()
        mock_database.list_collections.return_value = []
        mock_component.get_database_object = MagicMock(return_value=mock_database)

        codeflash_output = mock_component._initialize_collection_options()

    def test_edge_case_empty_keyspace(self, mock_component):
        # Mock methods
        mock_component.get_keyspace = MagicMock(return_value="")
        mock_database = MagicMock()
        mock_database.list_collections.return_value = []
        mock_component.get_database_object = MagicMock(return_value=mock_database)

        codeflash_output = mock_component._initialize_collection_options()

    def test_invalid_api_endpoint(self, mock_component):
        # Mock methods
        mock_component.get_database_object = MagicMock(side_effect=ValueError("Invalid API endpoint"))

        with pytest.raises(ValueError):
            mock_component._initialize_collection_options(api_endpoint="invalid_endpoint")

    def test_missing_vector_options(self, mock_component, mock_database):
        # Mock methods
        collection_without_vector = MagicMock()
        collection_without_vector.name = "col3"
        collection_without_vector.options.vector = None
        mock_database.list_collections.return_value.append(collection_without_vector)
        mock_component.get_database_object = MagicMock(return_value=mock_database)
        mock_component.collection_data = MagicMock(return_value=10)

        codeflash_output = mock_component._initialize_collection_options()

[codeflash-ai]

Suggested change

	return self.keyspace.strip() if self.keyspace else None
	keyspace = self.keyspace
	return keyspace.strip() if keyspace else None

[codeflash-ai]

⚡️ Codeflash found optimizations for this PR

📄 28% (0.28x) speedup for AstraDBVectorStoreComponent.get_keyspace in src/backend/base/langflow/components/vectorstores/astradb.py

⏱️ Runtime : 58.6 microseconds → 45.8 microseconds (best of 108 runs)

📝 Explanation and details

To optimize this Python program, there are a few adjustments we can make for potentially faster performance and better memory usage.

1. Directly Access Attributes: Ensure that the attributes keyspace are being accessed directly to avoid any potential unnecessary lookups.
2. Avoid Stripping if Not Necessary: Only call the strip method on keyspace if it is set and has a value. This avoids calling an additional method unnecessarily.
3. Conditional Return Optimization: Optimize how the keyspace value is checked and returned.

Here's an optimized version of the get_keyspace method.

In this optimized version.

• keyspace is retrieved once and assigned to a local variable.
• The strip() method is called only once if keyspace has a value, avoiding multiple attribute accesses and making the code more efficient.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 42 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

# Basic Functionality
def test_get_keyspace_basic():
    component = AstraDBVectorStoreComponent()
    component.keyspace = "test_keyspace"
    codeflash_output = component.get_keyspace()
    
    component.keyspace = "example"
    codeflash_output = component.get_keyspace()

# Handling Whitespace
def test_get_keyspace_whitespace():
    component = AstraDBVectorStoreComponent()
    component.keyspace = "  test_keyspace  "
    codeflash_output = component.get_keyspace()
    
    component.keyspace = "\ttest_keyspace\n"
    codeflash_output = component.get_keyspace()

    component.keyspace = "   "
    codeflash_output = component.get_keyspace()
    
    component.keyspace = "\n\t "
    codeflash_output = component.get_keyspace()

# Empty and None Keyspace
def test_get_keyspace_empty_none():
    component = AstraDBVectorStoreComponent()
    component.keyspace = ""
    codeflash_output = component.get_keyspace()
    
    component.keyspace = None
    codeflash_output = component.get_keyspace()

# Special Characters in Keyspace
def test_get_keyspace_special_characters():
    component = AstraDBVectorStoreComponent()
    component.keyspace = "keyspace_with_underscores"
    codeflash_output = component.get_keyspace()
    
    component.keyspace = "keyspace-with-dashes"
    codeflash_output = component.get_keyspace()
    
    component.keyspace = "keyspace.with.dots"
    codeflash_output = component.get_keyspace()
    
    component.keyspace = "keyspace@with#special$chars"
    codeflash_output = component.get_keyspace()

# Unicode and Internationalization
def test_get_keyspace_unicode():
    component = AstraDBVectorStoreComponent()
    component.keyspace = "ключевое_пространство"  # Russian for "keyspace"
    codeflash_output = component.get_keyspace()
    
    component.keyspace = "空間"  # Chinese for "space"
    codeflash_output = component.get_keyspace()
    
    component.keyspace = "スペース"  # Japanese for "space"
    codeflash_output = component.get_keyspace()

# Large Scale Test Cases
def test_get_keyspace_large_scale():
    component = AstraDBVectorStoreComponent()
    component.keyspace = "a" * 1000  # 1000 characters long
    codeflash_output = component.get_keyspace()

# Edge Cases
def test_get_keyspace_edge_cases():
    component = AstraDBVectorStoreComponent()
    component.keyspace = "  valid_keyspace  "
    codeflash_output = component.get_keyspace()
    
    component.keyspace = "\tvalid_keyspace\n"
    codeflash_output = component.get_keyspace()
    
    component.keyspace = "\n\n\n"
    codeflash_output = component.get_keyspace()

# Keyspace with Embedded Whitespace
def test_get_keyspace_embedded_whitespace():
    component = AstraDBVectorStoreComponent()
    component.keyspace = "key space"
    codeflash_output = component.get_keyspace()
    
    component.keyspace = "key\tspace"
    codeflash_output = component.get_keyspace()

# Run the tests
if __name__ == "__main__":
    pytest.main()
# 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
@pytest.mark.parametrize("keyspace, expected", [
    # Basic Functionality
    ("test_keyspace", "test_keyspace"),
    ("  test_keyspace  ", "test_keyspace"),
    ("\ttest_keyspace\t", "test_keyspace"),
    ("\ntest_keyspace\n", "test_keyspace"),
    # Empty and None Keyspace
    ("", None),
    (None, None),
    # Whitespace-Only Keyspace
    ("   ", None),
    ("\t\t\t", None),
    ("\n\n\n", None),
    # Special Characters in Keyspace
    ("test_keyspace!@#", "test_keyspace!@#"),
    ("keyspace_with_underscores", "keyspace_with_underscores"),
    ("keyspace-with-dashes", "keyspace-with-dashes"),
    # Unicode and Non-ASCII Characters
    ("ключевое_пространство", "ключевое_пространство"),
    ("espacio_de_clave", "espacio_de_clave"),
    ("空間", "空間"),
    # Long Keyspace Strings
    ("a" * 1000, "a" * 1000),
    (" " * 1000 + "keyspace" + " " * 1000, "keyspace"),
    # Mixed Content Keyspace
    ("  keyspace123!@#  ", "keyspace123!@#"),
    ("\tkeyspace_with_tabs\t", "keyspace_with_tabs"),
    # Edge Cases
    ("key\nspace", "key\nspace"),
    ("\n\tkeyspace\t\n", "keyspace"),
])
def test_get_keyspace(keyspace, expected):
    # Create an instance of AstraDBVectorStoreComponent with the provided keyspace
    component = AstraDBVectorStoreComponent(keyspace=keyspace)
    # Assert that the get_keyspace method returns the expected value
    codeflash_output = component.get_keyspace()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

[codeflash-ai]

⚡️ Codeflash found optimizations for this PR

📄 38% (0.38x) speedup for AstraDBVectorStoreComponent._initialize_collection_options in src/backend/base/langflow/components/vectorstores/astradb.py

⏱️ Runtime : 519 milliseconds → 376 milliseconds (best of 5 runs)

I created a new dependent PR with the suggested changes. Please review:

• ⚡️ Speed up method AstraDBVectorStoreComponent._initialize_collection_options by 38% in PR #6202 (codeflash/optimize-pr6028-2025-02-07T20.23.30) #6204

If you approve, it will be merged into this PR (branch codeflash/optimize-pr6028-2025-02-07T20.23.30).