⚡️ Speed up method AmazonMistralConfig.map_openai_params by 36%

[codeflash-ai]

📄 36% (0.36x) speedup for AmazonMistralConfig.map_openai_params in litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py

⏱️ Runtime : 327 microseconds → 241 microseconds (best of 316 runs)

📝 Explanation and details

The optimization delivers a 35% speedup by addressing two key performance bottlenecks:

1. Replaced multiple sequential if statements with a single set membership check
The original code used 5 separate if statements to check each parameter (max_tokens, temperature, top_p, stop, stream), executing all checks for every key-value pair. The optimized version uses if k in {"max_tokens", "temperature", "top_p", "stop", "stream"}:, which performs a single O(1) hash table lookup instead of up to 5 string comparisons.

2. Eliminated redundant parameter assignment logic
The original code repeated optional_params[param_name] = v for each parameter type. The optimized version uses optional_params[k] = v, directly assigning the key without additional branching.

Performance impact analysis from line profiler:

• Original: 3.46ms total time with 9,840 individual parameter checks (1,968 checks × 5 parameters)
• Optimized: 1.57ms total time with just 1,968 set membership operations
• The set membership check (k in {set}) takes 551μs vs 2.4ms for sequential string comparisons

Test case performance patterns:

• Single parameter cases: 3-6% improvement (small overhead reduction)
• Multiple parameter cases: 10-22% improvement (compound benefits)
• Large-scale cases with many parameters: Up to 160% improvement when processing mixed valid/invalid parameters, as the set membership efficiently filters out unsupported keys

This optimization is particularly effective for workloads with multiple parameter mappings or when processing configuration dictionaries with many keys, both common patterns in ML model parameter handling.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import List, Optional

# imports
import pytest  # used for our unit tests
from litellm.llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import \
    AmazonMistralConfig


# function to test (AmazonMistralConfig and map_openai_params)
class BaseConfig:
    pass

class BaseAWSLLM:
    pass

class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
    def __init__(self, **kwargs):
        BaseConfig.__init__(self, **kwargs)
        BaseAWSLLM.__init__(self, **kwargs)
from litellm.llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import \
    AmazonMistralConfig

# ---------- UNIT TESTS BEGIN HERE ----------

@pytest.fixture
def config():
    # Fixture to create a config instance for reuse
    return AmazonMistralConfig()

# 1. Basic Test Cases

def test_basic_max_tokens(config):
    # Test mapping of max_tokens param
    codeflash_output = config.map_openai_params(
        non_default_params={"max_tokens": 42},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.15μs -> 1.31μs (12.1% slower)

def test_basic_temperature(config):
    # Test mapping of temperature param
    codeflash_output = config.map_openai_params(
        non_default_params={"temperature": 0.5},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.24μs -> 1.19μs (3.52% faster)

def test_basic_top_p(config):
    # Test mapping of top_p param
    codeflash_output = config.map_openai_params(
        non_default_params={"top_p": 0.9},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.19μs -> 1.18μs (0.508% faster)

def test_basic_stop(config):
    # Test mapping of stop param
    stop_list = ["\n", "END"]
    codeflash_output = config.map_openai_params(
        non_default_params={"stop": stop_list},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.12μs -> 1.11μs (1.26% faster)

def test_basic_stream(config):
    # Test mapping of stream param
    codeflash_output = config.map_openai_params(
        non_default_params={"stream": True},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.23μs -> 1.16μs (6.22% faster)

def test_basic_multiple_params(config):
    # Test mapping multiple params at once
    codeflash_output = config.map_openai_params(
        non_default_params={
            "max_tokens": 123,
            "temperature": 0.7,
            "top_p": 0.8,
            "stop": ["STOP"],
            "stream": False,
        },
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.87μs -> 1.53μs (21.6% faster)

def test_basic_optional_params_preserved(config):
    # Test that optional_params are preserved and updated
    codeflash_output = config.map_openai_params(
        non_default_params={"temperature": 0.3},
        optional_params={"max_tokens": 99, "top_p": 0.5},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.23μs -> 1.14μs (7.93% faster)

# 2. Edge Test Cases

def test_edge_empty_non_default_params(config):
    # Test with empty non_default_params
    codeflash_output = config.map_openai_params(
        non_default_params={},
        optional_params={"max_tokens": 10},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 922ns -> 938ns (1.71% slower)

def test_edge_empty_optional_params(config):
    # Test with empty optional_params
    codeflash_output = config.map_openai_params(
        non_default_params={"max_tokens": 7},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.19μs -> 1.13μs (5.04% faster)

def test_edge_unknown_param_ignored(config):
    # Test that unknown params are ignored
    codeflash_output = config.map_openai_params(
        non_default_params={"unknown_param": "value"},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.08μs -> 1.07μs (1.50% faster)

def test_edge_none_param_value(config):
    # Test that None values are mapped
    codeflash_output = config.map_openai_params(
        non_default_params={"max_tokens": None},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.16μs -> 1.16μs (0.432% faster)

def test_edge_stop_empty_list(config):
    # Test stop param as empty list
    codeflash_output = config.map_openai_params(
        non_default_params={"stop": []},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.14μs -> 1.12μs (1.42% faster)

def test_edge_stop_large_list(config):
    # Test stop param as large list of strings
    stop_list = [str(i) for i in range(100)]
    codeflash_output = config.map_openai_params(
        non_default_params={"stop": stop_list},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.12μs -> 1.12μs (0.537% faster)

def test_edge_overwrite_optional_params(config):
    # Test that non_default_params overwrite optional_params
    codeflash_output = config.map_openai_params(
        non_default_params={"temperature": 0.9},
        optional_params={"temperature": 0.1, "max_tokens": 5},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.22μs -> 1.12μs (8.77% faster)

def test_edge_type_variations(config):
    # Test with type variations (int vs float)
    codeflash_output = config.map_openai_params(
        non_default_params={"max_tokens": 50, "temperature": 1, "top_p": 1.0},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.51μs -> 1.38μs (9.60% faster)

def test_edge_stream_false(config):
    # Test stream param with False value
    codeflash_output = config.map_openai_params(
        non_default_params={"stream": False},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.11μs -> 1.02μs (8.81% faster)

def test_edge_drop_params_flag(config):
    # Test drop_params flag has no effect (since not used)
    codeflash_output = config.map_openai_params(
        non_default_params={"max_tokens": 1},
        optional_params={},
        model="mistral",
        drop_params=True,
    ); result = codeflash_output # 1.17μs -> 1.11μs (4.86% faster)

# 3. Large Scale Test Cases

def test_large_scale_many_params(config):
    # Test with large number of params in optional_params
    optional_params = {f"param_{i}": i for i in range(900)}
    non_default_params = {"max_tokens": 1000, "temperature": 0.2}
    codeflash_output = config.map_openai_params(
        non_default_params=non_default_params,
        optional_params=optional_params.copy(),
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.37μs -> 1.34μs (1.79% faster)
    # All original params should be preserved
    for i in range(900):
        pass

def test_large_scale_stop_list(config):
    # Test with a large stop list (close to 1000 elements)
    stop_list = [f"stop_{i}" for i in range(999)]
    codeflash_output = config.map_openai_params(
        non_default_params={"stop": stop_list},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.22μs -> 1.17μs (4.64% faster)

def test_large_scale_non_default_params(config):
    # Test with many non_default_params, but only supported keys should be mapped
    non_default_params = {f"param_{i}": i for i in range(900)}
    # Add supported params
    non_default_params.update({
        "max_tokens": 500,
        "temperature": 0.99,
        "top_p": 0.77,
        "stop": ["a", "b"],
        "stream": True,
    })
    codeflash_output = config.map_openai_params(
        non_default_params=non_default_params,
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 57.0μs -> 22.0μs (160% faster)
    # Unsupported keys should not be present
    for i in range(900):
        pass

def test_large_scale_update_existing(config):
    # Test updating existing optional_params with large scale
    optional_params = {f"param_{i}": i for i in range(500)}
    optional_params["max_tokens"] = 100
    optional_params["temperature"] = 0.1
    non_default_params = {"max_tokens": 200, "temperature": 0.2}
    codeflash_output = config.map_openai_params(
        non_default_params=non_default_params,
        optional_params=optional_params.copy(),
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 1.40μs -> 1.39μs (0.719% faster)
    # Other keys remain unchanged
    for i in range(500):
        pass

def test_large_scale_empty(config):
    # Test with both dicts empty (should return empty dict)
    codeflash_output = config.map_openai_params(
        non_default_params={},
        optional_params={},
        model="mistral",
        drop_params=False,
    ); result = codeflash_output # 947ns -> 920ns (2.93% 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 List, Optional

# imports
import pytest
from litellm.llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import \
    AmazonMistralConfig


class BaseConfig:
    def __init__(self, **kwargs):
        pass

class BaseAWSLLM:
    def __init__(self, **kwargs):
        pass

class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
    def __init__(self, **kwargs):
        BaseConfig.__init__(self, **kwargs)
        BaseAWSLLM.__init__(self, **kwargs)
from litellm.llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import \
    AmazonMistralConfig

# unit tests

@pytest.fixture
def config():
    # Fixture for AmazonMistralConfig instance
    return AmazonMistralConfig()

# 1. Basic Test Cases

def test_basic_single_param_max_tokens(config):
    # Test mapping a single 'max_tokens' parameter
    codeflash_output = config.map_openai_params(
        non_default_params={"max_tokens": 123},
        optional_params={},
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.19μs -> 1.16μs (3.29% faster)

def test_basic_multiple_params(config):
    # Test mapping several parameters at once
    params = {
        "max_tokens": 50,
        "temperature": 0.5,
        "top_p": 0.9,
        "stop": ["\n", "<END>"]
    }
    codeflash_output = config.map_openai_params(
        non_default_params=params,
        optional_params={},
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.67μs -> 1.41μs (17.8% faster)

def test_basic_optional_params_preserved_and_updated(config):
    # Test that optional_params are updated but unrelated keys are preserved
    non_default = {"max_tokens": 42, "temperature": 0.2}
    optional = {"foo": "bar", "max_tokens": 1}
    codeflash_output = config.map_openai_params(
        non_default_params=non_default,
        optional_params=optional.copy(),
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.32μs -> 1.19μs (10.9% faster)
    expected = {"foo": "bar", "max_tokens": 42, "temperature": 0.2}

def test_basic_stream_param(config):
    # Test that 'stream' param is mapped
    codeflash_output = config.map_openai_params(
        non_default_params={"stream": True},
        optional_params={},
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.15μs -> 1.10μs (4.55% faster)

# 2. Edge Test Cases

def test_edge_empty_non_default_params(config):
    # Test with empty non_default_params (should not modify optional_params)
    codeflash_output = config.map_openai_params(
        non_default_params={},
        optional_params={"existing": 1},
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 904ns -> 853ns (5.98% faster)

def test_edge_non_default_params_with_unrecognized_key(config):
    # Test that keys not in the mapping are ignored
    codeflash_output = config.map_openai_params(
        non_default_params={"unknown": 999, "max_tokens": 10},
        optional_params={},
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.29μs -> 1.24μs (3.78% faster)

def test_edge_none_values_in_non_default_params(config):
    # Test that None values are mapped as-is
    codeflash_output = config.map_openai_params(
        non_default_params={"max_tokens": None, "temperature": None},
        optional_params={},
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.39μs -> 1.29μs (7.82% faster)

def test_edge_stop_param_empty_list(config):
    # Test that an empty list for 'stop' is mapped correctly
    codeflash_output = config.map_openai_params(
        non_default_params={"stop": []},
        optional_params={},
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.18μs -> 1.10μs (8.03% faster)

def test_edge_stop_param_single_string(config):
    # Test that a single string for 'stop' is mapped (even if not a list)
    codeflash_output = config.map_openai_params(
        non_default_params={"stop": "<END>"},
        optional_params={},
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.16μs -> 1.09μs (6.71% faster)

def test_edge_optional_params_is_mutated(config):
    # Test that optional_params is mutated in-place
    optional = {}
    config.map_openai_params(
        non_default_params={"max_tokens": 1},
        optional_params=optional,
        model="mistral",
        drop_params=False
    ) # 1.13μs -> 1.10μs (2.82% faster)

def test_edge_optional_params_with_existing_keys(config):
    # Test that existing keys in optional_params are overwritten if present in non_default_params
    optional = {"temperature": 0.1, "max_tokens": 5}
    non_default = {"temperature": 0.9, "max_tokens": 99}
    codeflash_output = config.map_openai_params(
        non_default_params=non_default,
        optional_params=optional.copy(),
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.29μs -> 1.26μs (2.22% faster)

def test_edge_drop_params_flag_is_ignored(config):
    # Test that drop_params flag does not affect the output
    codeflash_output = config.map_openai_params(
        non_default_params={"max_tokens": 1},
        optional_params={},
        model="mistral",
        drop_params=True
    ); result1 = codeflash_output # 1.06μs -> 1.03μs (2.82% faster)
    codeflash_output = config.map_openai_params(
        non_default_params={"max_tokens": 1},
        optional_params={},
        model="mistral",
        drop_params=False
    ); result2 = codeflash_output # 479ns -> 400ns (19.8% faster)

def test_edge_model_is_ignored(config):
    # Test that the model argument does not affect the output
    codeflash_output = config.map_openai_params(
        non_default_params={"max_tokens": 1},
        optional_params={},
        model="mistral",
        drop_params=False
    ); result1 = codeflash_output # 1.04μs -> 1.03μs (0.774% faster)
    codeflash_output = config.map_openai_params(
        non_default_params={"max_tokens": 1},
        optional_params={},
        model="other-model",
        drop_params=False
    ); result2 = codeflash_output # 492ns -> 382ns (28.8% faster)

def test_edge_top_k_param(config):
    # Test that 'top_k' is not mapped (since not in mapping)
    codeflash_output = config.map_openai_params(
        non_default_params={"top_k": 5},
        optional_params={},
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.00μs -> 956ns (4.92% faster)

def test_edge_top_p_param(config):
    # Test that 'top_p' is mapped
    codeflash_output = config.map_openai_params(
        non_default_params={"top_p": 0.8},
        optional_params={},
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.08μs -> 1.03μs (5.26% faster)

# 3. Large Scale Test Cases

def test_large_scale_many_params(config):
    # Test with many recognized and unrecognized keys
    non_default = {f"max_tokens_{i}": i for i in range(500)}
    # Add one valid key
    non_default["max_tokens"] = 999
    optional = {f"foo_{i}": i for i in range(500)}
    codeflash_output = config.map_openai_params(
        non_default_params=non_default,
        optional_params=optional.copy(),
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 32.0μs -> 12.8μs (150% faster)
    # Only 'max_tokens' should be mapped, rest preserved
    expected = optional.copy()
    expected["max_tokens"] = 999

def test_large_scale_all_supported_params(config):
    # Test with all supported keys and large values
    non_default = {
        "max_tokens": 1000,
        "temperature": 1.0,
        "top_p": 0.99,
        "stop": ["<END>", "\n", "STOP"],
        "stream": False
    }
    optional = {f"bar_{i}": i for i in range(100)}
    codeflash_output = config.map_openai_params(
        non_default_params=non_default,
        optional_params=optional.copy(),
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.74μs -> 1.61μs (8.33% faster)
    expected = optional.copy()
    expected.update(non_default)

def test_large_scale_stop_param_large_list(config):
    # Test with a large list for 'stop'
    stop_list = [f"<STOP_{i}>" for i in range(999)]
    codeflash_output = config.map_openai_params(
        non_default_params={"stop": stop_list},
        optional_params={},
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.10μs -> 1.16μs (5.51% slower)

def test_large_scale_stream_param_many_times(config):
    # Test mapping 'stream' param many times
    for i in range(500):
        codeflash_output = config.map_openai_params(
            non_default_params={"stream": bool(i % 2)},
            optional_params={},
            model="mistral",
            drop_params=False
        ); result = codeflash_output # 187μs -> 158μs (18.5% faster)

def test_large_scale_overwrite_existing_keys(config):
    # Test that many keys in optional_params are preserved except those overwritten
    optional = {f"key_{i}": i for i in range(1000)}
    non_default = {"max_tokens": 123, "temperature": 0.1}
    codeflash_output = config.map_openai_params(
        non_default_params=non_default,
        optional_params=optional.copy(),
        model="mistral",
        drop_params=False
    ); result = codeflash_output # 1.35μs -> 1.34μs (0.745% faster)
    expected = optional.copy()
    expected["max_tokens"] = 123
    expected["temperature"] = 0.1
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from litellm.llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import AmazonMistralConfig

def test_AmazonMistralConfig_map_openai_params():
    AmazonMistralConfig.map_openai_params(AmazonMistralConfig(max_tokens=0, temperature=None, top_p=0, top_k=0.0, stop=None), {'max_tokens': ''}, {}, '', False)

def test_AmazonMistralConfig_map_openai_params_2():
    AmazonMistralConfig.map_openai_params(AmazonMistralConfig(max_tokens=None, temperature=0.0, top_p=None, top_k=None, stop=None), {'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00': '', 0: '', 'temperature': 0}, {'temperature': 0}, '', True)

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_zbim32de/tmpuyi6m9dz/test_concolic_coverage.py::test_AmazonMistralConfig_map_openai_params	788ns	828ns	-4.83%⚠️
codeflash_concolic_zbim32de/tmpuyi6m9dz/test_concolic_coverage.py::test_AmazonMistralConfig_map_openai_params_2	1.10μs	1.04μs	6.26%✅

To edit these changes git checkout codeflash/optimize-AmazonMistralConfig.map_openai_params-mhdnlzut and push.