⚡️ Speed up method ConversationUsageInfo.serialize_model by 25%

[codeflash-ai]

📄 25% (0.25x) speedup for ConversationUsageInfo.serialize_model in src/mistralai/models/conversationusageinfo.py

⏱️ Runtime : 110 microseconds → 88.1 microseconds (best of 187 runs)

📝 Explanation and details

The optimized code achieves a 25% speedup through several targeted micro-optimizations that reduce overhead in the serialization loop:

Key Changes:

1. Data structure optimization: Changed nullable_fields from a list to a set ({"connector_tokens", "connectors"}) for O(1) membership testing instead of O(n) list scanning.

2. Loop variable hoisting: Moved self.__pydantic_fields_set__ and type(self).model_fields outside the loop to avoid repeated attribute lookups on each iteration.

3. Set intersection elimination: Replaced self.__pydantic_fields_set__.intersection({n}) with direct n in fields_set, avoiding the overhead of creating a temporary set and calling the intersection method on every loop iteration.

Why it's faster:

• The set lookup for nullable_fields provides constant-time membership testing, which is particularly beneficial when checking multiple fields
• Hoisting attribute lookups eliminates repeated dot notation access to self.__pydantic_fields_set__ and type(self).model_fields
• Direct set membership (n in fields_set) is much faster than set intersection operations, especially since we're only checking a single element

Performance characteristics:
Based on the test results, the optimization shows consistent 19-32% improvements across all test cases, with larger gains on simpler serialization scenarios (up to 31.7% on large-scale tests). The optimizations are particularly effective for models with multiple fields since the improvements compound with each loop iteration.

The changes maintain identical behavior while reducing computational overhead in the hot path of model serialization.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

from typing import Dict, Optional

# imports
import pytest
from mistralai.models.conversationusageinfo import ConversationUsageInfo


# Minimal stand-ins for external dependencies
class UNSETType:
    pass

UNSET_SENTINEL = object()
UNSET = UNSETType()

class OptionalNullable:
    def __init__(self, value):
        self.value = value

    def __eq__(self, other):
        if isinstance(other, OptionalNullable):
            return self.value == other.value
        return self.value == other

    def __repr__(self):
        return f"OptionalNullable({self.value!r})"

# Minimal BaseModel and model_serializer decorator
class BaseModel:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
        self.__pydantic_fields_set__ = set(kwargs.keys())

    @classmethod
    def model_fields(cls):
        # Return a dict of field names to dummy objects with 'alias' attribute
        fields = {}
        for k in cls.__annotations__:
            # For this test, no aliases
            fields[k] = type('F', (), {'alias': None})()
        return fields

def model_serializer(mode=None):
    def decorator(fn):
        return fn
    return decorator
from mistralai.models.conversationusageinfo import ConversationUsageInfo


# Helper handler function for serialization (simulates Pydantic's behavior)
def handler(obj):
    # Return a dict of all attributes, except private ones
    result = {}
    for k in obj.model_fields:
        v = getattr(obj, k, None)
        # Unwrap OptionalNullable for serialization
        if isinstance(v, OptionalNullable):
            result[k] = v.value
        elif v is UNSET:
            result[k] = UNSET_SENTINEL
        else:
            result[k] = v
    return result

# -------------------------------
# Unit tests for serialize_model
# -------------------------------

# 1. Basic Test Cases



def test_serialize_only_required_fields():
    # Only required fields set (simulate minimal input)
    obj = ConversationUsageInfo(prompt_tokens=1)
    codeflash_output = obj.serialize_model(handler); result = codeflash_output # 11.2μs -> 9.26μs (21.3% faster)

# 2. Edge Test Cases


def test_serialize_unset_fields():
    # Fields left unset (simulate not provided)
    obj = ConversationUsageInfo()
    codeflash_output = obj.serialize_model(handler); result = codeflash_output # 10.7μs -> 8.98μs (19.1% faster)








#------------------------------------------------
from __future__ import annotations

from typing import Dict, Optional

# imports
import pytest
from mistralai.models.conversationusageinfo import ConversationUsageInfo


# Simulated UNSET/UNSET_SENTINEL and OptionalNullable for test purposes
class _UnsetType:
    pass

UNSET = _UnsetType()
UNSET_SENTINEL = UNSET
OptionalNullable = Optional

# Simulated BaseModel for test purposes
class BaseModel:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
        # Track explicitly set fields
        self.__pydantic_fields_set__ = set(kwargs.keys())

    @classmethod
    def model_fields(cls):
        # Simulate pydantic field info
        return {k: type('Field', (), {'alias': None}) for k in cls.__annotations__}

# Simulated model_serializer decorator
def model_serializer(mode=None):
    def decorator(fn):
        return fn
    return decorator
from mistralai.models.conversationusageinfo import ConversationUsageInfo


# Helper handler function for serialization
def default_handler(obj):
    # Simulate default serialization as dict
    result = {}
    for k in obj.__annotations__:
        val = getattr(obj, k, None)
        result[k] = val
    return result

# unit tests

# Basic Test Cases

def test_basic_all_defaults():
    """Test serialization with all default values."""
    obj = ConversationUsageInfo()
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 9.22μs -> 7.74μs (19.1% faster)

def test_basic_explicit_values():
    """Test serialization with all fields explicitly set."""
    obj = ConversationUsageInfo(
        prompt_tokens=10,
        completion_tokens=20,
        total_tokens=30,
        connector_tokens=40,
        connectors={"a": 1, "b": 2}
    )
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 6.62μs -> 5.15μs (28.5% faster)

def test_basic_partial_fields():
    """Test serialization with only some fields set."""
    obj = ConversationUsageInfo(prompt_tokens=5)
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 7.58μs -> 5.98μs (26.7% faster)

# Edge Test Cases

def test_edge_none_values():
    """Test serialization with None for nullable fields."""
    obj = ConversationUsageInfo(
        connector_tokens=None,
        connectors=None
    )
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 6.30μs -> 4.94μs (27.6% faster)


def test_edge_missing_fields():
    """Test serialization with missing fields (simulate by not setting them)."""
    obj = ConversationUsageInfo()
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 9.16μs -> 7.59μs (20.8% faster)
    # All fields should have their default values
    for field in obj.__annotations__:
        pass

def test_edge_zero_and_empty_dict():
    """Test serialization with zero and empty dict for connectors."""
    obj = ConversationUsageInfo(
        connector_tokens=0,
        connectors={}
    )
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 6.94μs -> 5.37μs (29.3% faster)

def test_edge_negative_and_large_int():
    """Test serialization with negative and large integer values."""
    obj = ConversationUsageInfo(
        prompt_tokens=-1,
        completion_tokens=999999999,
        total_tokens=0
    )
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 7.64μs -> 6.01μs (27.2% faster)

def test_edge_connectors_with_non_str_keys():
    """Test serialization with connectors dict having non-str keys."""
    obj = ConversationUsageInfo(
        connectors={1: 10, 2: 20}
    )
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 7.27μs -> 5.63μs (29.1% faster)

def test_edge_connectors_with_none_value():
    """Test serialization with connectors dict having None value."""
    obj = ConversationUsageInfo(
        connectors={"x": None}
    )
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 7.11μs -> 5.48μs (29.7% faster)

# Large Scale Test Cases

def test_large_connectors_dict():
    """Test serialization with large connectors dict."""
    large_dict = {str(i): i for i in range(1000)}
    obj = ConversationUsageInfo(connectors=large_dict)
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 7.22μs -> 5.61μs (28.8% faster)

def test_large_tokens_values():
    """Test serialization with large integer values for tokens."""
    obj = ConversationUsageInfo(
        prompt_tokens=999999,
        completion_tokens=888888,
        total_tokens=777777,
        connector_tokens=666666
    )
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 6.82μs -> 5.44μs (25.3% faster)

def test_large_scale_all_fields():
    """Test serialization with all fields set to large values."""
    large_dict = {str(i): i for i in range(1000)}
    obj = ConversationUsageInfo(
        prompt_tokens=1000000,
        completion_tokens=2000000,
        total_tokens=3000000,
        connector_tokens=4000000,
        connectors=large_dict
    )
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 6.49μs -> 4.93μs (31.7% faster)

To edit these changes git checkout codeflash/optimize-ConversationUsageInfo.serialize_model-mh4e82p7 and push.