⚡️ Speed up function required_tool by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for required_tool in src/xai_sdk/chat.py

⏱️ Runtime : 471 microseconds → 428 microseconds (best of 10 runs)

📝 Explanation and details

The optimization caches the chat_pb2.ToolChoice class reference in a module-level variable _ToolChoice instead of performing attribute lookup on chat_pb2 every time the function is called.

Key change: Added _ToolChoice = chat_pb2.ToolChoice at module level and replaced chat_pb2.ToolChoice(function_name=name) with _ToolChoice(function_name=name).

Why it's faster: In Python, attribute lookups like chat_pb2.ToolChoice involve dictionary lookups in the module's namespace on every function call. By caching the class reference once at import time, we eliminate this repeated attribute lookup overhead. The line profiler shows the per-call time decreased from 679.6ns to 636.1ns (6.4% improvement per call).

Performance characteristics: This optimization is most effective for frequently called functions where the same class is instantiated repeatedly. The test results show consistent 2-16% speedups across different input scenarios, with the most significant improvement (15.8%) when handling type errors with dictionary inputs. The optimization provides steady gains regardless of input type since the class lookup elimination occurs before any input validation.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
aio/chat_test.py::test_chat_create_with_required_tool	1.76μs	1.46μs	20.5%✅
sync/chat_test.py::test_chat_create_with_required_tool	1.21μs	1.14μs	5.69%✅

🌀 Generated Regression Tests and Runtime

import pytest
from xai_sdk.chat import required_tool

# --- Mocking the external dependency for a self-contained test suite ---
# Since we don't have access to xai_sdk.proto.chat_pb2, we'll mock ToolChoice.
# This mock should behave like the real ToolChoice for the purpose of these tests.

class MockToolChoice:
    def __init__(self, function_name):
        self.function_name = function_name

    def __eq__(self, other):
        # Equality based on function_name and type
        return isinstance(other, MockToolChoice) and self.function_name == other.function_name

    def __repr__(self):
        return f"MockToolChoice(function_name={self.function_name!r})"
from xai_sdk.chat import required_tool

# --- Unit tests ---

# 1. Basic Test Cases
















#------------------------------------------------
import pytest  # used for our unit tests
from xai_sdk.chat import required_tool

# --- Begin: Minimal mock of xai_sdk.proto.chat_pb2.ToolChoice for testing purposes ---
# Since we cannot import xai_sdk.proto.chat_pb2, we create a minimal mock class for ToolChoice.
# This mock matches the interface required for testing the required_tool function.

class ToolChoice:
    def __init__(self, function_name):
        self.function_name = function_name

    def __eq__(self, other):
        # Equality for test assertions
        if not isinstance(other, ToolChoice):
            return False
        return self.function_name == other.function_name

    def __repr__(self):
        return f"ToolChoice(function_name={self.function_name!r})"
from xai_sdk.chat import required_tool

# unit tests

# 1. Basic Test Cases






def test_required_tool_none_input():
    # Test with None as input (should raise TypeError)
    with pytest.raises(TypeError):
        required_tool(None)

def test_required_tool_integer_input():
    # Test with integer as input (should raise TypeError)
    with pytest.raises(TypeError):
        required_tool(123) # 5.35μs -> 5.04μs (6.13% faster)

def test_required_tool_list_input():
    # Test with list as input (should raise TypeError)
    with pytest.raises(TypeError):
        required_tool(["summarize"]) # 3.62μs -> 3.53μs (2.38% faster)

def test_required_tool_dict_input():
    # Test with dict as input (should raise TypeError)
    with pytest.raises(TypeError):
        required_tool({"function_name": "summarize"}) # 3.44μs -> 2.97μs (15.8% faster)







#------------------------------------------------
from xai_sdk.chat import required_tool
import pytest

def test_required_tool():
    with pytest.raises(TypeError, match='bad\\ argument\\ type\\ for\\ built\\-in\\ operation'):
        required_tool('')

To edit these changes git checkout codeflash/optimize-required_tool-mgu3fd6t and push.