⚡️ Speed up method CohereChatConfig._construct_cohere_tool by 9%

[codeflash-ai]

📄 9% (0.09x) speedup for CohereChatConfig._construct_cohere_tool in litellm/llms/cohere/chat/transformation.py

⏱️ Runtime : 4.39 milliseconds → 4.04 milliseconds (best of 234 runs)

📝 Explanation and details

The optimized code achieves an 8% speedup through three key optimizations:

1. Eliminated locals() overhead in __init__
The original code used locals().copy() and iterated through all local variables, which creates an unnecessary dictionary copy and performs attribute lookup on the class. The optimized version directly checks each parameter individually and sets instance attributes only when values are not None, avoiding the costly dictionary operations.

2. Replaced loop with list comprehension in _construct_cohere_tool
Changed from manually creating an empty list and appending items in a loop to using a list comprehension. This eliminates the overhead of repeated append() calls and leverages Python's optimized list comprehension implementation.

3. Reduced repeated dictionary lookups in _translate_openai_tool_to_cohere
The original code repeatedly accessed nested dictionary paths like openai_tool["function"]["parameters"] multiple times. The optimized version extracts these into local variables (function_section, properties, required_params, pd) and reuses them, significantly reducing dictionary lookup overhead.

Performance characteristics by test case:

• Large-scale operations benefit most: 47% faster for 100 tools with 10 parameters each, and 9.7% faster for single tools with 500 parameters
• Small operations show slight regression (30-40% slower for empty/single tool cases) due to the overhead of list comprehension setup, but this is negligible in absolute terms (microseconds)
• Medium-scale operations show modest improvements of 3-7% faster

The optimizations are particularly effective for workloads with many tools or complex parameter structures, where the reduced dictionary lookups and more efficient list creation compound the performance gains.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Any, Optional

# imports
import pytest  # used for our unit tests
from litellm.llms.cohere.chat.transformation import CohereChatConfig

# unit tests

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

# 1. BASIC TEST CASES

def test_empty_tools_returns_empty_list(config):
    # Test that passing None returns an empty list
    codeflash_output = config._construct_cohere_tool(None) # 642ns -> 1.09μs (41.3% slower)
    # Test that passing empty list returns an empty list
    codeflash_output = config._construct_cohere_tool([]) # 310ns -> 446ns (30.5% slower)

def test_single_tool_basic_translation(config):
    # Test translation of a single simple OpenAI tool
    openai_tool = {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get weather info",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "City name",
                    },
                },
                "required": ["location"],
            },
        },
    }
    expected = [{
        "name": "get_weather",
        "description": "Get weather info",
        "parameter_definitions": {
            "location": {
                "description": "City name",
                "type": "string",
                "required": True,
            }
        }
    }]
    codeflash_output = config._construct_cohere_tool([openai_tool]) # 2.38μs -> 2.48μs (3.76% slower)

def test_multiple_tools_basic_translation(config):
    # Test translation of multiple tools
    openai_tools = [
        {
            "type": "function",
            "function": {
                "name": "func1",
                "description": "desc1",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "param1": {"type": "string", "description": "desc param1"},
                    },
                    "required": ["param1"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "func2",
                "description": "desc2",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "param2": {"type": "int", "description": "desc param2"},
                    },
                    "required": [],
                },
            },
        }
    ]
    expected = [
        {
            "name": "func1",
            "description": "desc1",
            "parameter_definitions": {
                "param1": {
                    "description": "desc param1",
                    "type": "string",
                    "required": True,
                }
            }
        },
        {
            "name": "func2",
            "description": "desc2",
            "parameter_definitions": {
                "param2": {
                    "description": "desc param2",
                    "type": "int",
                    "required": False,
                }
            }
        }
    ]
    codeflash_output = config._construct_cohere_tool(openai_tools) # 3.10μs -> 3.16μs (1.93% slower)

# 2. EDGE TEST CASES

def test_tool_with_no_required_parameters(config):
    # Test tool with no required parameters
    openai_tool = {
        "type": "function",
        "function": {
            "name": "no_required",
            "description": "No required params",
            "parameters": {
                "type": "object",
                "properties": {
                    "foo": {"type": "string", "description": "foo desc"},
                    "bar": {"type": "int", "description": "bar desc"},
                },
                "required": [],
            },
        },
    }
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.62μs -> 2.55μs (2.78% faster)

def test_tool_with_multiple_required_parameters(config):
    # Test tool with multiple required parameters
    openai_tool = {
        "type": "function",
        "function": {
            "name": "multi_required",
            "description": "Multiple required params",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "string", "description": "desc a"},
                    "b": {"type": "int", "description": "desc b"},
                    "c": {"type": "float", "description": "desc c"},
                },
                "required": ["a", "c"],
            },
        },
    }
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 3.05μs -> 3.02μs (0.993% faster)

def test_tool_with_missing_description_and_type(config):
    # Test tool with missing description and type for a parameter
    openai_tool = {
        "type": "function",
        "function": {
            "name": "missing_fields",
            "description": "Missing fields",
            "parameters": {
                "type": "object",
                "properties": {
                    "x": {},  # missing both description and type
                },
                "required": ["x"],
            },
        },
    }
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.31μs -> 2.37μs (2.62% slower)

def test_tool_with_extra_parameter_fields(config):
    # Test that extra fields in parameter definitions are ignored
    openai_tool = {
        "type": "function",
        "function": {
            "name": "extra_fields",
            "description": "Extra fields test",
            "parameters": {
                "type": "object",
                "properties": {
                    "foo": {
                        "type": "string",
                        "description": "desc foo",
                        "extra": "should be ignored",
                    },
                },
                "required": ["foo"],
            },
        },
    }
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.13μs -> 2.37μs (9.85% slower)

def test_tool_with_empty_properties(config):
    # Test tool with no parameters (empty properties)
    openai_tool = {
        "type": "function",
        "function": {
            "name": "no_params",
            "description": "No params",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": [],
            },
        },
    }
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 1.49μs -> 1.93μs (22.9% slower)

def test_tool_with_non_string_param_types(config):
    # Test tool with param types other than string
    openai_tool = {
        "type": "function",
        "function": {
            "name": "non_string_types",
            "description": "Non-string param types",
            "parameters": {
                "type": "object",
                "properties": {
                    "count": {"type": "integer", "description": "Number of items"},
                    "flag": {"type": "boolean", "description": "A boolean flag"},
                },
                "required": ["count"],
            },
        },
    }
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.66μs -> 2.65μs (0.529% faster)

def test_tool_with_enum_field_in_param(config):
    # Test that enum field in param is ignored
    openai_tool = {
        "type": "function",
        "function": {
            "name": "enum_test",
            "description": "Enum param test",
            "parameters": {
                "type": "object",
                "properties": {
                    "unit": {
                        "type": "string",
                        "description": "Unit type",
                        "enum": ["celsius", "fahrenheit"],
                    },
                },
                "required": ["unit"],
            },
        },
    }
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.21μs -> 2.33μs (5.03% slower)
    param_def = result[0]["parameter_definitions"]["unit"]

def test_tool_with_required_param_not_in_properties(config):
    # Test required parameter not present in properties
    openai_tool = {
        "type": "function",
        "function": {
            "name": "missing_param",
            "description": "Required param not in properties",
            "parameters": {
                "type": "object",
                "properties": {
                    "foo": {"type": "string", "description": "desc foo"},
                },
                "required": ["bar"],  # 'bar' not in properties
            },
        },
    }
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.24μs -> 2.40μs (6.55% slower)

# 3. LARGE SCALE TEST CASES

def test_large_number_of_tools_and_parameters(config):
    # Test with 100 tools, each with 10 parameters
    num_tools = 100
    num_params = 10
    openai_tools = []
    for i in range(num_tools):
        properties = {}
        required = []
        for j in range(num_params):
            param_name = f"param{j}"
            properties[param_name] = {
                "type": "string",
                "description": f"desc {param_name}",
            }
            if j % 2 == 0:
                required.append(param_name)
        openai_tools.append({
            "type": "function",
            "function": {
                "name": f"tool{i}",
                "description": f"description{i}",
                "parameters": {
                    "type": "object",
                    "properties": properties,
                    "required": required,
                },
            },
        })
    codeflash_output = config._construct_cohere_tool(openai_tools); result = codeflash_output # 292μs -> 199μs (47.0% faster)
    for i, tool in enumerate(result):
        for j in range(num_params):
            param = tool["parameter_definitions"][f"param{j}"]

def test_large_number_of_parameters_in_single_tool(config):
    # Test with a single tool with 500 parameters
    num_params = 500
    properties = {}
    required = []
    for i in range(num_params):
        param_name = f"p{i}"
        properties[param_name] = {
            "type": "string",
            "description": f"desc {param_name}",
        }
        if i % 3 == 0:
            required.append(param_name)
    openai_tool = {
        "type": "function",
        "function": {
            "name": "bigtool",
            "description": "A big tool",
            "parameters": {
                "type": "object",
                "properties": properties,
                "required": required,
            },
        },
    }
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 520μs -> 474μs (9.68% faster)
    for i in range(num_params):
        param = result[0]["parameter_definitions"][f"p{i}"]

def test_large_empty_tools_list(config):
    # Test with a large empty tools list
    tools = []
    codeflash_output = config._construct_cohere_tool(tools) # 572ns -> 928ns (38.4% slower)

def test_large_tools_list_with_empty_properties(config):
    # Test with 100 tools, each with empty properties
    openai_tools = []
    for i in range(100):
        openai_tools.append({
            "type": "function",
            "function": {
                "name": f"tool{i}",
                "description": f"description{i}",
                "parameters": {
                    "type": "object",
                    "properties": {},
                    "required": [],
                },
            },
        })
    codeflash_output = config._construct_cohere_tool(openai_tools); result = codeflash_output # 27.5μs -> 31.5μs (12.8% slower)
    for tool in result:
        pass
# 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 Any, Optional

# imports
import pytest
from litellm.llms.cohere.chat.transformation import CohereChatConfig

# unit tests

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

# ------------------------
# Basic Test Cases
# ------------------------

def test_empty_tools_list_returns_empty(config):
    # Test: empty tools list returns empty list
    codeflash_output = config._construct_cohere_tool([]); result = codeflash_output # 593ns -> 1.02μs (41.9% slower)

def test_none_tools_returns_empty(config):
    # Test: None as tools returns empty list
    codeflash_output = config._construct_cohere_tool(None); result = codeflash_output # 621ns -> 1.01μs (38.8% slower)

def test_single_simple_tool(config):
    # Test: single OpenAI tool with one required parameter
    openai_tool = {
        "type": "function",
        "function": {
            "name": "get_current_weather",
            "description": "Get the current weather in a given location",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The city and state, e.g. San Francisco, CA",
                    }
                },
                "required": ["location"],
            },
        },
    }
    expected = [{
        "name": "get_current_weather",
        "description": "Get the current weather in a given location",
        "parameter_definitions": {
            "location": {
                "description": "The city and state, e.g. San Francisco, CA",
                "type": "string",
                "required": True,
            }
        }
    }]
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.67μs -> 2.68μs (0.411% slower)

def test_multiple_tools(config):
    # Test: multiple OpenAI tools
    openai_tool1 = {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get weather",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string",
                        "description": "City name"
                    }
                },
                "required": ["city"],
            },
        },
    }
    openai_tool2 = {
        "type": "function",
        "function": {
            "name": "get_time",
            "description": "Get time",
            "parameters": {
                "type": "object",
                "properties": {
                    "timezone": {
                        "type": "string",
                        "description": "Timezone"
                    }
                },
                "required": ["timezone"],
            },
        },
    }
    expected = [
        {
            "name": "get_weather",
            "description": "Get weather",
            "parameter_definitions": {
                "city": {
                    "description": "City name",
                    "type": "string",
                    "required": True,
                }
            }
        },
        {
            "name": "get_time",
            "description": "Get time",
            "parameter_definitions": {
                "timezone": {
                    "description": "Timezone",
                    "type": "string",
                    "required": True,
                }
            }
        }
    ]
    codeflash_output = config._construct_cohere_tool([openai_tool1, openai_tool2]); result = codeflash_output # 3.32μs -> 3.17μs (4.96% faster)

# ------------------------
# Edge Test Cases
# ------------------------

def test_tool_with_no_required_parameters(config):
    # Test: OpenAI tool with no required parameters
    openai_tool = {
        "type": "function",
        "function": {
            "name": "say_hello",
            "description": "Say hello",
            "parameters": {
                "type": "object",
                "properties": {
                    "greeting": {
                        "type": "string",
                        "description": "Greeting message",
                    }
                },
                "required": [],
            },
        },
    }
    expected = [{
        "name": "say_hello",
        "description": "Say hello",
        "parameter_definitions": {
            "greeting": {
                "description": "Greeting message",
                "type": "string",
                "required": False,
            }
        }
    }]
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.27μs -> 2.47μs (8.18% slower)

def test_tool_with_multiple_parameters_and_types(config):
    # Test: OpenAI tool with multiple parameters, types, and required subset
    openai_tool = {
        "type": "function",
        "function": {
            "name": "complex_func",
            "description": "A complex function",
            "parameters": {
                "type": "object",
                "properties": {
                    "foo": {
                        "type": "string",
                        "description": "Foo parameter"
                    },
                    "bar": {
                        "type": "integer",
                        "description": "Bar parameter"
                    },
                    "baz": {
                        "type": "boolean",
                        "description": "Baz parameter"
                    }
                },
                "required": ["foo", "baz"],
            },
        },
    }
    expected = [{
        "name": "complex_func",
        "description": "A complex function",
        "parameter_definitions": {
            "foo": {
                "description": "Foo parameter",
                "type": "string",
                "required": True,
            },
            "bar": {
                "description": "Bar parameter",
                "type": "integer",
                "required": False,
            },
            "baz": {
                "description": "Baz parameter",
                "type": "boolean",
                "required": True,
            }
        }
    }]
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 3.32μs -> 3.09μs (7.37% faster)

def test_tool_with_missing_description(config):
    # Test: parameter missing description should default to empty string
    openai_tool = {
        "type": "function",
        "function": {
            "name": "missing_desc_func",
            "description": "Function with missing param description",
            "parameters": {
                "type": "object",
                "properties": {
                    "foo": {
                        "type": "string"
                    }
                },
                "required": ["foo"],
            },
        },
    }
    expected = [{
        "name": "missing_desc_func",
        "description": "Function with missing param description",
        "parameter_definitions": {
            "foo": {
                "description": "",
                "type": "string",
                "required": True,
            }
        }
    }]
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.27μs -> 2.36μs (4.10% slower)

def test_tool_with_no_parameters(config):
    # Test: OpenAI tool with no parameters
    openai_tool = {
        "type": "function",
        "function": {
            "name": "no_params_func",
            "description": "Function with no parameters",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": [],
            },
        },
    }
    expected = [{
        "name": "no_params_func",
        "description": "Function with no parameters",
        "parameter_definitions": {}
    }]
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 1.47μs -> 1.93μs (23.6% slower)

def test_tool_with_additional_properties(config):
    # Test: OpenAI tool with extra fields in parameter properties
    openai_tool = {
        "type": "function",
        "function": {
            "name": "extra_func",
            "description": "Function with extra param fields",
            "parameters": {
                "type": "object",
                "properties": {
                    "foo": {
                        "type": "string",
                        "description": "Foo",
                        "extra_field": "should be ignored"
                    }
                },
                "required": ["foo"],
            },
        },
    }
    expected = [{
        "name": "extra_func",
        "description": "Function with extra param fields",
        "parameter_definitions": {
            "foo": {
                "description": "Foo",
                "type": "string",
                "required": True,
            }
        }
    }]
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.10μs -> 2.27μs (7.11% slower)

def test_tool_with_enum_type(config):
    # Test: OpenAI tool with enum type (should just copy type and description)
    openai_tool = {
        "type": "function",
        "function": {
            "name": "enum_func",
            "description": "Function with enum param",
            "parameters": {
                "type": "object",
                "properties": {
                    "unit": {
                        "type": "string",
                        "description": "Unit type",
                        "enum": ["celsius", "fahrenheit"]
                    }
                },
                "required": ["unit"],
            },
        },
    }
    expected = [{
        "name": "enum_func",
        "description": "Function with enum param",
        "parameter_definitions": {
            "unit": {
                "description": "Unit type",
                "type": "string",
                "required": True,
            }
        }
    }]
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.22μs -> 2.32μs (4.10% slower)

def test_tool_with_non_string_type(config):
    # Test: OpenAI tool with non-string parameter type
    openai_tool = {
        "type": "function",
        "function": {
            "name": "int_func",
            "description": "Function with int param",
            "parameters": {
                "type": "object",
                "properties": {
                    "count": {
                        "type": "integer",
                        "description": "Count value"
                    }
                },
                "required": ["count"],
            },
        },
    }
    expected = [{
        "name": "int_func",
        "description": "Function with int param",
        "parameter_definitions": {
            "count": {
                "description": "Count value",
                "type": "integer",
                "required": True,
            }
        }
    }]
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.22μs -> 2.35μs (5.33% slower)

# ------------------------
# Large Scale Test Cases
# ------------------------

def test_large_number_of_tools(config):
    # Test: many tools (up to 1000)
    tools = []
    expected = []
    for i in range(1000):
        tool = {
            "type": "function",
            "function": {
                "name": f"func_{i}",
                "description": f"Function number {i}",
                "parameters": {
                    "type": "object",
                    "properties": {
                        f"param_{i}": {
                            "type": "string",
                            "description": f"Parameter {i}"
                        }
                    },
                    "required": [f"param_{i}"],
                },
            },
        }
        tools.append(tool)
        expected.append({
            "name": f"func_{i}",
            "description": f"Function number {i}",
            "parameter_definitions": {
                f"param_{i}": {
                    "description": f"Parameter {i}",
                    "type": "string",
                    "required": True,
                }
            }
        })
    codeflash_output = config._construct_cohere_tool(tools); result = codeflash_output # 651μs -> 608μs (6.95% faster)

def test_tool_with_many_parameters(config):
    # Test: single tool with many parameters (up to 1000)
    properties = {}
    required = []
    for i in range(1000):
        properties[f"param_{i}"] = {
            "type": "string",
            "description": f"Parameter {i}"
        }
        if i % 2 == 0:
            required.append(f"param_{i}")
    openai_tool = {
        "type": "function",
        "function": {
            "name": "many_params_func",
            "description": "Function with many parameters",
            "parameters": {
                "type": "object",
                "properties": properties,
                "required": required,
            },
        },
    }
    expected_param_defs = {}
    for i in range(1000):
        expected_param_defs[f"param_{i}"] = {
            "description": f"Parameter {i}",
            "type": "string",
            "required": (i % 2 == 0),
        }
    expected = [{
        "name": "many_params_func",
        "description": "Function with many parameters",
        "parameter_definitions": expected_param_defs
    }]
    codeflash_output = config._construct_cohere_tool([openai_tool]); result = codeflash_output # 2.56ms -> 2.48ms (3.53% faster)

def test_large_tools_and_parameters(config):
    # Test: 100 tools, each with 10 parameters
    tools = []
    expected = []
    for i in range(100):
        properties = {}
        required = []
        for j in range(10):
            properties[f"param_{j}"] = {
                "type": "string",
                "description": f"Parameter {j} of tool {i}"
            }
            if j % 3 == 0:
                required.append(f"param_{j}")
        tool = {
            "type": "function",
            "function": {
                "name": f"tool_{i}",
                "description": f"Tool number {i}",
                "parameters": {
                    "type": "object",
                    "properties": properties,
                    "required": required,
                },
            },
        }
        param_defs = {}
        for j in range(10):
            param_defs[f"param_{j}"] = {
                "description": f"Parameter {j} of tool {i}",
                "type": "string",
                "required": (j % 3 == 0),
            }
        expected.append({
            "name": f"tool_{i}",
            "description": f"Tool number {i}",
            "parameter_definitions": param_defs
        })
        tools.append(tool)
    codeflash_output = config._construct_cohere_tool(tools); result = codeflash_output # 287μs -> 199μs (44.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 litellm.llms.cohere.chat.transformation import CohereChatConfig
import pytest

def test_CohereChatConfig__construct_cohere_tool():
    with pytest.raises(TypeError, match="<class\\ 'str'>"):
        CohereChatConfig._construct_cohere_tool(CohereChatConfig(preamble='', chat_history=None, generation_id='', response_id='', conversation_id='', prompt_truncation='', connectors=[], search_queries_only=False, documents=[], temperature=0, max_tokens=0, max_completion_tokens=None, k=None, p=0, frequency_penalty=0, presence_penalty=0, tools=None, tool_results=[], seed=0), tools=[''])

def test_CohereChatConfig__construct_cohere_tool_2():
    CohereChatConfig._construct_cohere_tool(CohereChatConfig(preamble=None, chat_history=[], generation_id='', response_id='', conversation_id='', prompt_truncation='', connectors=[], search_queries_only=None, documents=[], temperature=0, max_tokens=None, max_completion_tokens=0, k=0, p=None, frequency_penalty=0, presence_penalty=0, tools=[], tool_results=[], seed=0), tools=None)

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_zbim32de/tmp4c6jznoz/test_concolic_coverage.py::test_CohereChatConfig__construct_cohere_tool_2	880ns	1.41μs	-37.4%⚠️

To edit these changes git checkout codeflash/optimize-CohereChatConfig._construct_cohere_tool-mhdonf6b and push.