⚡️ Speed up method VertexAITextEmbeddingConfig.get_mapped_special_auth_params by 33%

[codeflash-ai]

📄 33% (0.33x) speedup for VertexAITextEmbeddingConfig.get_mapped_special_auth_params in litellm/llms/vertex_ai/vertex_embeddings/transformation.py

⏱️ Runtime : 573 microseconds → 429 microseconds (best of 147 runs)

📝 Explanation and details

The optimization replaces dictionary creation on every method call with a pre-allocated module-level constant. Instead of creating a new dictionary {"project": "vertex_project", "region_name": "vertex_location"} each time get_mapped_special_auth_params() is called, the optimized version defines _MAPPED_SPECIAL_AUTH_PARAMS as a constant and returns it directly.

This eliminates the overhead of:

• Dictionary object allocation (heap allocation)
• String key/value pair creation and insertion
• Dictionary initialization overhead

The line profiler shows the per-call time dropping from 327.4ns to 235.1ns (28% improvement per call), which compounds significantly under load. Test results demonstrate consistent 25-47% speedup across all scenarios, with the optimization being particularly effective for:

• High-frequency method calls (large scale tests show 30-35% improvements)
• Repeated calls on the same instance (determinism test shows 52-62% improvement on subsequent calls)
• Any usage pattern since the method behavior is independent of instance state

This is a classic constant hoisting optimization that's especially valuable for methods that return static data and are called frequently in production systems.

✅ Correctness verification report:

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

# 1. Basic Test Cases

def test_basic_output_type_and_keys():
    # Test that the function returns a dict with expected keys and values
    config = VertexAITextEmbeddingConfig()
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 620ns -> 436ns (42.2% faster)

def test_basic_with_all_fields_set():
    # Test with all config fields set
    config = VertexAITextEmbeddingConfig(
        auto_truncate=True,
        task_type="RETRIEVAL_DOCUMENT",
        title="Sample Title"
    )
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 438ns -> 325ns (34.8% faster)

def test_basic_with_partial_fields_set():
    # Test with some config fields set
    config = VertexAITextEmbeddingConfig(
        auto_truncate=False
    )
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 405ns -> 306ns (32.4% faster)

def test_basic_with_none_fields():
    # Test with all fields as None
    config = VertexAITextEmbeddingConfig(
        auto_truncate=None,
        task_type=None,
        title=None
    )
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 444ns -> 312ns (42.3% faster)

# 2. Edge Test Cases

def test_edge_task_type_invalid():
    # Test with an invalid task_type (should not affect output)
    config = VertexAITextEmbeddingConfig(
        task_type="INVALID_TASK_TYPE"
    )
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 401ns -> 321ns (24.9% faster)

def test_edge_title_empty_string():
    # Test with title as empty string
    config = VertexAITextEmbeddingConfig(
        title=""
    )
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 374ns -> 296ns (26.4% faster)

def test_edge_title_long_string():
    # Test with a very long title string
    long_title = "A" * 1000
    config = VertexAITextEmbeddingConfig(
        title=long_title
    )
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 394ns -> 274ns (43.8% faster)

def test_edge_auto_truncate_non_bool():
    # Test with auto_truncate as a non-bool value (should not affect output)
    config = VertexAITextEmbeddingConfig(
        auto_truncate="not_a_bool"
    )
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 359ns -> 268ns (34.0% faster)

def test_edge_task_type_none_and_title_none():
    # Test with both task_type and title as None
    config = VertexAITextEmbeddingConfig(
        task_type=None,
        title=None
    )
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 385ns -> 326ns (18.1% faster)

def test_edge_task_type_literal_variants():
    # Test all valid Literal values for task_type
    valid_types = [
        "RETRIEVAL_QUERY",
        "RETRIEVAL_DOCUMENT",
        "SEMANTIC_SIMILARITY",
        "CLASSIFICATION",
        "CLUSTERING",
        "QUESTION_ANSWERING",
        "FACT_VERIFICATION",
    ]
    for vtype in valid_types:
        config = VertexAITextEmbeddingConfig(task_type=vtype)
        codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 1.54μs -> 1.16μs (32.6% faster)

# 3. Large Scale Test Cases

def test_large_scale_multiple_instances():
    # Create 1000 instances and verify output for each
    for i in range(1000):
        config = VertexAITextEmbeddingConfig(
            auto_truncate=bool(i % 2),
            task_type="RETRIEVAL_QUERY" if i % 3 == 0 else "RETRIEVAL_DOCUMENT",
            title=f"title_{i}"
        )
        codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 182μs -> 138μs (32.0% faster)

def test_large_scale_with_varied_inputs():
    # Test with a variety of input combinations for 500 instances
    for i in range(500):
        auto_truncate = None if i % 5 == 0 else bool(i % 2)
        task_type = None if i % 7 == 0 else (
            "RETRIEVAL_QUERY" if i % 3 == 0 else "FACT_VERIFICATION"
        )
        title = None if i % 11 == 0 else f"title_{i}"
        config = VertexAITextEmbeddingConfig(
            auto_truncate=auto_truncate,
            task_type=task_type,
            title=title
        )
        codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 91.5μs -> 69.4μs (31.9% faster)

def test_large_scale_edge_values():
    # Test with extreme values for title and all valid task_types
    long_title = "Z" * 999
    for vtype in [
        "RETRIEVAL_QUERY",
        "RETRIEVAL_DOCUMENT",
        "SEMANTIC_SIMILARITY",
        "CLASSIFICATION",
        "CLUSTERING",
        "QUESTION_ANSWERING",
        "FACT_VERIFICATION",
    ]:
        config = VertexAITextEmbeddingConfig(
            auto_truncate=True,
            task_type=vtype,
            title=long_title
        )
        codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 1.54μs -> 1.13μs (36.2% faster)
# 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
from pydantic import BaseModel

# unit tests

# 1. Basic Test Cases

def test_get_mapped_special_auth_params_basic():
    # Test with default config (no params)
    config = VertexAITextEmbeddingConfig()
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 379ns -> 279ns (35.8% faster)

def test_get_mapped_special_auth_params_with_auto_truncate_true():
    # Test with auto_truncate True
    config = VertexAITextEmbeddingConfig(auto_truncate=True)
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 385ns -> 287ns (34.1% faster)

def test_get_mapped_special_auth_params_with_task_type_and_title():
    # Test with all fields set
    config = VertexAITextEmbeddingConfig(
        auto_truncate=False,
        task_type="RETRIEVAL_DOCUMENT",
        title="Example Title"
    )
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 377ns -> 268ns (40.7% faster)

def test_get_mapped_special_auth_params_with_semantic_similarity_task():
    # Test with a different valid task_type
    config = VertexAITextEmbeddingConfig(
        task_type="SEMANTIC_SIMILARITY"
    )
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 355ns -> 253ns (40.3% faster)

# 2. Edge Test Cases

def test_get_mapped_special_auth_params_with_none_values():
    # All fields explicitly set to None
    config = VertexAITextEmbeddingConfig(
        auto_truncate=None,
        task_type=None,
        title=None
    )
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 397ns -> 294ns (35.0% faster)

def test_get_mapped_special_auth_params_with_empty_title():
    # Title is empty string
    config = VertexAITextEmbeddingConfig(title="")
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 380ns -> 284ns (33.8% faster)

def test_get_mapped_special_auth_params_with_unusual_title():
    # Title contains special characters
    config = VertexAITextEmbeddingConfig(title="!@#$%^&*()_+-=[]{};':\",.<>/?\\|`~")
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 427ns -> 298ns (43.3% faster)


def test_get_mapped_special_auth_params_with_long_title():
    # Title is very long
    long_title = "A" * 1000
    config = VertexAITextEmbeddingConfig(title=long_title)
    codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 600ns -> 406ns (47.8% faster)

def test_get_mapped_special_auth_params_with_boolean_auto_truncate():
    # Edge boolean values
    config_true = VertexAITextEmbeddingConfig(auto_truncate=True)
    config_false = VertexAITextEmbeddingConfig(auto_truncate=False)
    codeflash_output = config_true.get_mapped_special_auth_params() # 435ns -> 336ns (29.5% faster)
    codeflash_output = config_false.get_mapped_special_auth_params() # 210ns -> 159ns (32.1% faster)

# 3. Large Scale Test Cases

def test_get_mapped_special_auth_params_many_instances():
    # Create many instances and call the method to check for consistency and performance
    instances = [
        VertexAITextEmbeddingConfig(
            auto_truncate=(i % 2 == 0),
            task_type="RETRIEVAL_QUERY" if i % 3 == 0 else "CLASSIFICATION",
            title=f"Title {i}"
        )
        for i in range(500)
    ]
    for config in instances:
        codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 89.5μs -> 66.1μs (35.3% faster)

def test_get_mapped_special_auth_params_all_task_types():
    # Test every valid task_type
    valid_task_types = [
        "RETRIEVAL_QUERY",
        "RETRIEVAL_DOCUMENT",
        "SEMANTIC_SIMILARITY",
        "CLASSIFICATION",
        "CLUSTERING",
        "QUESTION_ANSWERING",
        "FACT_VERIFICATION",
    ]
    for task_type in valid_task_types:
        config = VertexAITextEmbeddingConfig(task_type=task_type)
        codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 1.53μs -> 1.14μs (33.6% faster)

def test_get_mapped_special_auth_params_large_titles():
    # Test with many configs with large titles
    for i in range(100):
        config = VertexAITextEmbeddingConfig(title="X" * (i + 1) * 10)
        codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 18.8μs -> 14.4μs (30.9% faster)

def test_get_mapped_special_auth_params_stress():
    # Stress test with maximum allowed elements (under 1000)
    configs = [
        VertexAITextEmbeddingConfig(
            auto_truncate=True,
            task_type="RETRIEVAL_QUERY",
            title=f"StressTestTitle{i}"
        )
        for i in range(999)
    ]
    for config in configs:
        codeflash_output = config.get_mapped_special_auth_params(); result = codeflash_output # 177μs -> 131μs (34.8% faster)

# 4. Determinism Test

def test_get_mapped_special_auth_params_determinism():
    # Multiple calls on the same object should return the same result
    config = VertexAITextEmbeddingConfig(
        auto_truncate=True,
        task_type="RETRIEVAL_DOCUMENT",
        title="Determinism"
    )
    codeflash_output = config.get_mapped_special_auth_params(); result1 = codeflash_output # 467ns -> 288ns (62.2% faster)
    codeflash_output = config.get_mapped_special_auth_params(); result2 = codeflash_output # 247ns -> 162ns (52.5% faster)

# 5. Mutation Safety Test

def test_get_mapped_special_auth_params_mutation_safety():
    # Changing config attributes should not affect the output of get_mapped_special_auth_params
    config = VertexAITextEmbeddingConfig(
        auto_truncate=False,
        task_type="RETRIEVAL_QUERY",
        title="MutationTest"
    )
    codeflash_output = config.get_mapped_special_auth_params(); result_before = codeflash_output
    config.auto_truncate = True
    config.task_type = "FACT_VERIFICATION"
    config.title = "ChangedTitle"
    codeflash_output = config.get_mapped_special_auth_params(); result_after = codeflash_output
# 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.get_mapped_special_auth_params-mhc6bz0p and push.