⚡️ Speed up method VertexAITextEmbeddingConfig.map_special_auth_params by 113%

[codeflash-ai]

📄 113% (1.13x) speedup for VertexAITextEmbeddingConfig.map_special_auth_params in litellm/llms/vertex_ai/vertex_embeddings/transformation.py

⏱️ Runtime : 141 microseconds → 66.3 microseconds (best of 250 runs)

📝 Explanation and details

The optimization replaces the inefficient loop-and-check pattern in map_special_auth_params with Python's set intersection operation (&).

Key optimization:

• Original approach: Iterates through ALL parameters in non_default_params.items() and checks if param in mapped_params for each one
• Optimized approach: Uses non_default_params.keys() & mapped_params.keys() to find only the intersection of keys, then iterates only over matching parameters

Why this is faster:
The original code performs O(n×m) operations where n is the number of parameters and m is the size of the mapping dict. For each parameter, it does a dictionary lookup to check membership. The optimized version uses set intersection which is O(n+m) and only processes relevant keys.

Performance impact from test results:

• Small inputs (1-5 params): 20-30% slower due to set operation overhead
• Large inputs (500+ params): 500-700% faster when most parameters don't need mapping
• The optimization shines when non_default_params contains many keys that aren't in the small mapped_params dict (only 2 keys: "project" and "region_name")

When this optimization matters:
This is particularly beneficial for configurations with many authentication parameters where only a few need special mapping - a common pattern in cloud service integrations where you might have dozens of config options but only specific ones need transformation.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Literal, Optional

# imports
import pytest
from litellm.llms.vertex_ai.vertex_embeddings.transformation import \
    VertexAITextEmbeddingConfig
from pydantic import BaseModel

# unit tests

# BASIC TEST CASES

def test_basic_single_mapping_project():
    # Test mapping of 'project' param
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"project": "my_proj"}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.32μs -> 1.84μs (28.0% slower)

def test_basic_single_mapping_region_name():
    # Test mapping of 'region_name' param
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"region_name": "us-central1"}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.25μs -> 1.78μs (30.0% slower)

def test_basic_multiple_mappings():
    # Test mapping both 'project' and 'region_name'
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"project": "my_proj", "region_name": "us-central1"}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.35μs -> 1.92μs (29.4% slower)

def test_basic_no_mappable_params():
    # Test when non_default_params has no mappable keys
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"foo": "bar", "baz": 42}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.21μs -> 1.52μs (20.6% slower)

def test_basic_optional_params_prepopulated():
    # Test mapping when optional_params is prepopulated
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"project": "my_proj"}
    optional_params = {"existing_key": "existing_value"}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.19μs -> 1.65μs (28.2% slower)

# EDGE TEST CASES

def test_edge_empty_non_default_params():
    # Test with empty non_default_params
    config = VertexAITextEmbeddingConfig()
    non_default_params = {}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.04μs -> 1.39μs (24.6% slower)

def test_edge_empty_optional_params():
    # Test with empty optional_params but mappable non_default_params
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"project": "my_proj"}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.20μs -> 1.76μs (31.7% slower)

def test_edge_none_values_in_non_default_params():
    # Test with None value in non_default_params
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"project": None, "region_name": "us-central1"}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.33μs -> 1.88μs (29.3% slower)

def test_edge_non_string_values():
    # Test with non-string values (e.g., int, bool)
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"project": 123, "region_name": True}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.25μs -> 1.78μs (29.7% slower)

def test_edge_overwrite_existing_optional_params():
    # Test that mapped params overwrite existing keys in optional_params
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"project": "new_proj"}
    optional_params = {"vertex_project": "old_proj"}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.23μs -> 1.69μs (27.3% slower)

def test_edge_non_default_params_with_extra_keys():
    # Test that extra keys in non_default_params are ignored
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"project": "my_proj", "region_name": "us-central1", "extra": "value"}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.39μs -> 1.75μs (21.0% slower)

def test_edge_case_sensitive_keys():
    # Test that keys are case sensitive
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"Project": "should_not_map", "region_name": "us-central1"}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.22μs -> 1.76μs (30.8% slower)

def test_edge_mutable_optional_params():
    # Test that optional_params is mutated in place
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"project": "my_proj"}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.21μs -> 1.62μs (25.5% slower)

# LARGE SCALE TEST CASES

def test_large_scale_many_non_mappable_keys():
    # Test with many non-mappable keys
    config = VertexAITextEmbeddingConfig()
    non_default_params = {f"key_{i}": i for i in range(1000)}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 28.0μs -> 1.75μs (1499% faster)

def test_large_scale_many_mappable_keys():
    # Test with many mappable keys, but only two are mapped
    config = VertexAITextEmbeddingConfig()
    # Only 'project' and 'region_name' should be mapped
    non_default_params = {f"key_{i}": i for i in range(998)}
    non_default_params["project"] = "large_proj"
    non_default_params["region_name"] = "large_region"
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 27.8μs -> 2.21μs (1162% faster)

def test_large_scale_prepopulated_optional_params():
    # Test with large prepopulated optional_params and mappable keys
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"project": "big_proj", "region_name": "big_region"}
    optional_params = {f"existing_{i}": i for i in range(998)}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.50μs -> 1.90μs (21.4% slower)
    # All original keys should still be present
    for i in range(998):
        pass

def test_large_scale_overwrite():
    # Test overwriting mapped keys in large optional_params
    config = VertexAITextEmbeddingConfig()
    non_default_params = {"project": "new_proj", "region_name": "new_region"}
    optional_params = {f"existing_{i}": i for i in range(998)}
    optional_params["vertex_project"] = "old_proj"
    optional_params["vertex_location"] = "old_region"
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.54μs -> 1.94μs (20.5% slower)
    # All original keys should still be present
    for i in range(998):
        pass

def test_large_scale_only_mappable_keys():
    # Test with only mappable keys in large input
    config = VertexAITextEmbeddingConfig()
    non_default_params = {}
    for i in range(500):
        non_default_params["project"] = f"proj_{i}"
        non_default_params["region_name"] = f"region_{i}"
    optional_params = {}
    # Only the last values should be present, since dict keys are overwritten
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.43μs -> 1.89μs (24.5% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from typing import Literal, Optional

# imports
import pytest
from litellm.llms.vertex_ai.vertex_embeddings.transformation import \
    VertexAITextEmbeddingConfig

# unit tests

@pytest.fixture
def config():
    # Provide a fresh config instance for each test
    return VertexAITextEmbeddingConfig()

# ---------------------- BASIC TEST CASES ----------------------

def test_map_single_project_param(config):
    # Only 'project' param should map to 'vertex_project'
    non_default_params = {'project': 'my_project'}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.45μs -> 1.92μs (24.3% slower)

def test_map_single_region_param(config):
    # Only 'region_name' param should map to 'vertex_location'
    non_default_params = {'region_name': 'us-central1'}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.36μs -> 1.84μs (26.0% slower)

def test_map_both_params(config):
    # Both 'project' and 'region_name' should map correctly
    non_default_params = {'project': 'p1', 'region_name': 'r1'}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.48μs -> 1.99μs (25.5% slower)

def test_map_with_unrelated_param(config):
    # Unrecognized params should be ignored
    non_default_params = {'project': 'p2', 'foo': 'bar'}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.38μs -> 1.81μs (23.9% slower)

def test_map_with_empty_non_default_params(config):
    # Empty input should result in no changes
    non_default_params = {}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.24μs -> 1.59μs (21.9% slower)

def test_map_with_existing_optional_params(config):
    # Should update optional_params, not overwrite unrelated keys
    non_default_params = {'project': 'p3'}
    optional_params = {'existing': 123}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.40μs -> 1.82μs (23.0% slower)

def test_map_with_overwrite_existing_key(config):
    # Should overwrite the mapped key if it exists
    non_default_params = {'project': 'p4'}
    optional_params = {'vertex_project': 'old_project'}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.36μs -> 1.74μs (21.5% slower)

# ---------------------- EDGE TEST CASES ----------------------

def test_map_with_none_values(config):
    # None values should be mapped as-is
    non_default_params = {'project': None, 'region_name': None}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.50μs -> 1.92μs (21.8% slower)

def test_map_with_non_string_values(config):
    # Non-string values (e.g. int, bool) should be mapped as-is
    non_default_params = {'project': 123, 'region_name': True}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.44μs -> 1.78μs (19.1% slower)

def test_map_with_extra_params_and_overlap(config):
    # Should only map recognized keys, even with overlap in optional_params
    non_default_params = {'project': 'p5', 'region_name': 'r5', 'extra': 'x'}
    optional_params = {'vertex_project': 'old', 'vertex_location': 'old', 'other': 1}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.48μs -> 1.83μs (19.2% slower)

def test_map_with_case_sensitivity(config):
    # Should be case-sensitive, so 'Project' is ignored
    non_default_params = {'Project': 'upper', 'project': 'lower'}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.38μs -> 1.80μs (23.4% slower)

def test_map_with_mutable_optional_params(config):
    # Should mutate the passed optional_params dict in-place
    non_default_params = {'project': 'p6'}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.29μs -> 1.69μs (23.4% slower)

def test_map_with_non_dict_inputs(config):
    # Should raise exception if non-dict is passed (TypeError)
    with pytest.raises(AttributeError):
        config.map_special_auth_params(['project', 'p7'], {})

    with pytest.raises(AttributeError):
        config.map_special_auth_params({'project': 'p7'}, [])

def test_map_with_empty_strings(config):
    # Should map empty string values as-is
    non_default_params = {'project': '', 'region_name': ''}
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.76μs -> 2.35μs (25.1% slower)

# ---------------------- LARGE SCALE TEST CASES ----------------------

def test_map_large_number_of_unrelated_params(config):
    # Only the two mapped keys should be present, rest ignored
    non_default_params = {f'foo{i}': i for i in range(500)}
    non_default_params['project'] = 'big_project'
    non_default_params['region_name'] = 'big_region'
    optional_params = {}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 14.7μs -> 2.29μs (542% faster)

def test_map_large_optional_params(config):
    # Should add mapped keys to a large optional_params dict
    optional_params = {f'bar{i}': i for i in range(500)}
    non_default_params = {'project': 'p8', 'region_name': 'r8'}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.56μs -> 2.06μs (24.3% slower)
    # All original keys must still exist
    for i in range(500):
        pass

def test_map_large_both_large(config):
    # Large unrelated keys in both, only mapped keys should be updated/added
    non_default_params = {f'foo{i}': i for i in range(500)}
    non_default_params['project'] = 'large_proj'
    optional_params = {f'bar{i}': i for i in range(500)}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 14.3μs -> 2.03μs (603% faster)
    for i in range(500):
        pass

def test_map_large_scale_overwrite(config):
    # Overwrite mapped keys in a large optional_params dict
    optional_params = {f'bar{i}': i for i in range(500)}
    optional_params['vertex_project'] = 'old_proj'
    optional_params['vertex_location'] = 'old_loc'
    non_default_params = {'project': 'new_proj', 'region_name': 'new_loc'}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 1.63μs -> 2.07μs (21.1% slower)

def test_map_large_scale_no_mapped_keys(config):
    # No mapped keys present, optional_params should be unchanged
    non_default_params = {f'foo{i}': i for i in range(500)}
    optional_params = {f'bar{i}': i for i in range(500)}
    codeflash_output = config.map_special_auth_params(non_default_params, optional_params); result = codeflash_output # 14.2μs -> 1.70μs (735% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-VertexAITextEmbeddingConfig.map_special_auth_params-mhoaj1wa and push.