⚡️ Speed up method AstraDBVectorStoreComponent._initialize_collection_options by 21% in PR #6085 (codeflash/optimize-pr6048-2025-02-03T13.54.58)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #6085

If you approve this dependent PR, these changes will be merged into the original PR branch codeflash/optimize-pr6048-2025-02-03T13.54.58.

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

---

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

⏱️ Runtime : 408 milliseconds → 336 milliseconds (best of 5 runs)

📝 Explanation and details

Here's the optimized version.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 17 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
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 = "test_token"
    component.environment = "test_env"
    component.keyspace = "test_keyspace"
    return component

def test_basic_functionality(mock_component):
    # Mock the database object and collections
    mock_database = MagicMock()
    mock_collection1 = MagicMock()
    mock_collection1.name = "collection1"
    mock_collection1.options.vector.service.provider = "provider1"
    mock_collection1.options.vector.service.model_name = "model1"
    
    mock_collection2 = MagicMock()
    mock_collection2.name = "collection2"
    mock_collection2.options.vector = None
    
    mock_database.list_collections.return_value = [mock_collection1, mock_collection2]
    mock_component.get_database_object = MagicMock(return_value=mock_database)
    mock_component.collection_data = MagicMock(side_effect=lambda collection_name, database: f"records_{collection_name}")

    codeflash_output = mock_component._initialize_collection_options()
    
    expected = [
        {
            "name": "collection1",
            "records": "records_collection1",
            "provider": "provider1",
            "icon": "",
            "model": "model1",
        },
        {
            "name": "collection2",
            "records": "records_collection2",
            "provider": None,
            "icon": "",
            "model": None,
        },
    ]

def test_empty_keyspace(mock_component):
    mock_component.keyspace = ""
    mock_component.get_keyspace = MagicMock(return_value=None)
    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_no_collections_in_keyspace(mock_component):
    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_collections_with_and_without_vector_options(mock_component):
    mock_database = MagicMock()
    mock_collection1 = MagicMock()
    mock_collection1.name = "collection1"
    mock_collection1.options.vector.service.provider = "provider1"
    mock_collection1.options.vector.service.model_name = "model1"
    
    mock_collection2 = MagicMock()
    mock_collection2.name = "collection2"
    mock_collection2.options.vector = None
    
    mock_database.list_collections.return_value = [mock_collection1, mock_collection2]
    mock_component.get_database_object = MagicMock(return_value=mock_database)
    mock_component.collection_data = MagicMock(side_effect=lambda collection_name, database: f"records_{collection_name}")

    codeflash_output = mock_component._initialize_collection_options()
    
    expected = [
        {
            "name": "collection1",
            "records": "records_collection1",
            "provider": "provider1",
            "icon": "",
            "model": "model1",
        },
        {
            "name": "collection2",
            "records": "records_collection2",
            "provider": None,
            "icon": "",
            "model": None,
        },
    ]

def test_invalid_api_endpoint(mock_component):
    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_url")


def test_special_characters_in_collection_names(mock_component):
    mock_database = MagicMock()
    mock_collection1 = MagicMock()
    mock_collection1.name = "collection_1"
    mock_collection1.options.vector = None
    
    mock_collection2 = MagicMock()
    mock_collection2.name = "collection-2"
    mock_collection2.options.vector = None
    
    mock_collection3 = MagicMock()
    mock_collection3.name = "collection 3"
    mock_collection3.options.vector = None
    
    mock_collection4 = MagicMock()
    mock_collection4.name = "collection_ñ"
    mock_collection4.options.vector = None
    
    mock_database.list_collections.return_value = [mock_collection1, mock_collection2, mock_collection3, mock_collection4]
    mock_component.get_database_object = MagicMock(return_value=mock_database)
    mock_component.collection_data = MagicMock(side_effect=lambda collection_name, database: f"records_{collection_name}")

    codeflash_output = mock_component._initialize_collection_options()
    
    expected = [
        {
            "name": "collection_1",
            "records": "records_collection_1",
            "provider": None,
            "icon": "",
            "model": None,
        },
        {
            "name": "collection-2",
            "records": "records_collection-2",
            "provider": None,
            "icon": "",
            "model": None,
        },
        {
            "name": "collection 3",
            "records": "records_collection 3",
            "provider": None,
            "icon": "",
            "model": None,
        },
        {
            "name": "collection_ñ",
            "records": "records_collection_ñ",
            "provider": None,
            "icon": "",
            "model": None,
        },
    ]

def test_database_object_returns_none(mock_component):
    mock_component.get_database_object = MagicMock(return_value=None)
    
    with pytest.raises(AttributeError):
        mock_component._initialize_collection_options()

def test_keyspace_with_similar_names(mock_component):
    mock_database = MagicMock()
    mock_collection1 = MagicMock()
    mock_collection1.name = "collection1"
    mock_collection1.options.vector = None
    
    mock_collection2 = MagicMock()
    mock_collection2.name = "collection1_"
    mock_collection2.options.vector = None
    
    mock_database.list_collections.return_value = [mock_collection1, mock_collection2]
    mock_component.get_database_object = MagicMock(return_value=mock_database)
    mock_component.collection_data = MagicMock(side_effect=lambda collection_name, database: f"records_{collection_name}")

    codeflash_output = mock_component._initialize_collection_options()
    
    expected = [
        {
            "name": "collection1",
            "records": "records_collection1",
            "provider": None,
            "icon": "",
            "model": None,
        },
        {
            "name": "collection1_",
            "records": "records_collection1_",
            "provider": None,
            "icon": "",
            "model": None,
        },
    ]


def test_collections_with_circular_references(mock_component):
    mock_database = MagicMock()
    mock_collection1 = MagicMock()
    mock_collection1.name = "collection1"
    mock_collection1.options.vector.service.provider = "provider1"
    mock_collection1.options.vector.service.model_name = "model1"
    
    mock_collection2 = MagicMock()
    mock_collection2.name = "collection2"
    mock_collection2.options.vector.service.provider = "provider2"
    mock_collection2.options.vector.service.model_name = "model2"
    
    mock_database.list_collections.return_value = [mock_collection1, mock_collection2]
    mock_component.get_database_object = MagicMock(return_value=mock_database)
    mock_component.collection_data = MagicMock(side_effect=lambda collection_name, database: f"records_{collection_name}")

    codeflash_output = mock_component._initialize_collection_options()
    
    expected = [
        {
            "name": "collection1",
            "records": "records_collection1",
            "provider": "provider1",
            "icon": "",
            "model": "model1",
        },
        {
            "name": "collection2",
            "records": "records_collection2",
            "provider": "provider2",
            "icon": "",
            "model": "model2",
        },
    ]


def test_collections_with_missing_attributes(mock_component):
    mock_database = MagicMock()
    mock_collection = MagicMock()
    del mock_collection.name  # Missing name attribute
    mock_collection.options.vector.service.provider = "provider1"
    
    mock_database.list_collections.return_value = [mock_collection]
    mock_component.get_database_object = MagicMock(return_value=mock_database)
    mock_component.collection_data = MagicMock(side_effect=lambda collection_name, database: f"records_{collection_name}")

    with pytest.raises(AttributeError):
        mock_component._initialize_collection_options()






from unittest.mock import MagicMock, patch

# imports
import pytest  # used for our unit tests
# function to test
from astrapy import DataAPIClient
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 = "test_token"
    component.environment = "test_env"
    component.keyspace = "test_keyspace"
    return component

def test_single_collection_with_vector_options(mock_component):
    # Mock the database object and its methods
    mock_database = MagicMock()
    mock_collection = MagicMock()
    mock_collection.name = "col1"
    mock_collection.options.vector.service.provider = "provider1"
    mock_collection.options.vector.service.model_name = "model1"
    mock_database.list_collections.return_value = [mock_collection]
    
    with patch.object(mock_component, 'get_database_object', return_value=mock_database):
        with patch.object(mock_component, 'collection_data', return_value={"data": "sample"}):
            codeflash_output = mock_component._initialize_collection_options()

def test_single_collection_without_vector_options(mock_component):
    mock_database = MagicMock()
    mock_collection = MagicMock()
    mock_collection.name = "col1"
    mock_collection.options.vector = None
    mock_database.list_collections.return_value = [mock_collection]
    
    with patch.object(mock_component, 'get_database_object', return_value=mock_database):
        with patch.object(mock_component, 'collection_data', return_value={"data": "sample"}):
            codeflash_output = mock_component._initialize_collection_options()

def test_multiple_collections(mock_component):
    mock_database = MagicMock()
    mock_collection1 = MagicMock()
    mock_collection1.name = "col1"
    mock_collection1.options.vector.service.provider = "provider1"
    mock_collection1.options.vector.service.model_name = "model1"
    
    mock_collection2 = MagicMock()
    mock_collection2.name = "col2"
    mock_collection2.options.vector = None
    
    mock_database.list_collections.return_value = [mock_collection1, mock_collection2]
    
    with patch.object(mock_component, 'get_database_object', return_value=mock_database):
        with patch.object(mock_component, 'collection_data', return_value={"data": "sample"}):
            codeflash_output = mock_component._initialize_collection_options()

def test_empty_keyspace(mock_component):
    mock_component.keyspace = None
    mock_database = MagicMock()
    mock_database.list_collections.return_value = []
    
    with patch.object(mock_component, 'get_database_object', return_value=mock_database):
        codeflash_output = mock_component._initialize_collection_options()

def test_invalid_api_endpoint(mock_component):
    with patch.object(mock_component, 'get_database_object', side_effect=ValueError("Invalid endpoint")):
        with pytest.raises(ValueError, match="Invalid endpoint"):
            mock_component._initialize_collection_options(api_endpoint="invalid_endpoint")


def test_special_characters_in_keyspace(mock_component):
    mock_component.keyspace = "keyspace_with_special_!@#"
    mock_database = MagicMock()
    mock_collection = MagicMock()
    mock_collection.name = "col1"
    mock_collection.options.vector.service.provider = "provider1"
    mock_collection.options.vector.service.model_name = "model1"
    mock_database.list_collections.return_value = [mock_collection]
    
    with patch.object(mock_component, 'get_database_object', return_value=mock_database):
        with patch.object(mock_component, 'collection_data', return_value={"data": "sample"}):
            codeflash_output = mock_component._initialize_collection_options()

def test_special_characters_in_collection_names(mock_component):
    mock_database = MagicMock()
    mock_collection = MagicMock()
    mock_collection.name = "col_!@#"
    mock_collection.options.vector.service.provider = "provider1"
    mock_collection.options.vector.service.model_name = "model1"
    mock_database.list_collections.return_value = [mock_collection]
    
    with patch.object(mock_component, 'get_database_object', return_value=mock_database):
        with patch.object(mock_component, 'collection_data', return_value={"data": "sample"}):
            codeflash_output = mock_component._initialize_collection_options()
# 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

📄 1,126% (11.26x) speedup for AstraDBVectorStoreComponent.get_database_object in src/backend/base/langflow/components/vectorstores/astradb.py

⏱️ Runtime : 10.2 milliseconds → 834 microseconds (best of 22 runs)

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

• ⚡️ Speed up method AstraDBVectorStoreComponent.get_database_object by 1,126% in PR #6087 (codeflash/optimize-pr6085-2025-02-03T14.05.09) #6088

If you approve, it will be merged into this PR (branch codeflash/optimize-pr6085-2025-02-03T14.05.09).