⚡️ Speed up method AgentUpdateRequest.serialize_model by 45%

[codeflash-ai]

📄 45% (0.45x) speedup for AgentUpdateRequest.serialize_model in src/mistralai/models/agentupdaterequest.py

⏱️ Runtime : 137 microseconds → 94.3 microseconds (best of 67 runs)

📝 Explanation and details

The optimization achieves a 45% speedup by converting expensive O(n) list lookups to O(1) set operations and eliminating redundant computations:

Key Optimizations:

1. Set-based lookups: Changed optional_fields and nullable_fields from lists to sets, transforming k in optional_fields from O(n) to O(1) operations
2. Precomputed field membership: Instead of calling self.__pydantic_fields_set__.intersection({n}) for each field (which creates a new set), the code now precomputes fields_set once and uses direct n in fields_set checks
3. Optimized comparison: Changed val != UNSET_SENTINEL to val is not UNSET_SENTINEL, using identity comparison instead of equality comparison for the sentinel object

Performance Impact:
The test results show consistent 27-56% improvements across all scenarios, with particularly strong gains for:

• Sparse field usage (55.8% faster for unset fields)
• Empty collections (53.7-53.9% faster for empty lists/strings)
• Single field operations (47.8% faster for minimal fields)

These optimizations are especially effective when processing objects with many fields but only a few set values, which appears to be the common use case for AgentUpdateRequest serialization.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 26 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 List, Optional, Union

# imports
import pytest
from mistralai.models.agentupdaterequest import AgentUpdateRequest
from typing_extensions import Annotated


# Dummy sentinel and type for testing
class UnsetSentinelType:
    pass

UNSET_SENTINEL = UnsetSentinelType()
UNSET = UNSET_SENTINEL

# Dummy BaseModel for testing
class BaseModel:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
        self.__pydantic_fields_set__ = set(kwargs.keys())
    # mimic Pydantic's model_fields
    @classmethod
    def model_fields(cls):
        return {k: type('Field', (), {'alias': None}) for k in cls.__annotations__}

# Dummy tools and args for testing
class CodeInterpreterTool:
    def __init__(self, type="code_interpreter"):
        self.type = type

class CompletionArgs:
    def __init__(self, arg1=None):
        self.arg1 = arg1

# The function under test, adapted for testability
def serialize_model(self, handler):
    optional_fields = [
        "instructions",
        "tools",
        "completion_args",
        "model",
        "name",
        "description",
        "handoffs",
    ]
    nullable_fields = ["instructions", "model", "name", "description", "handoffs"]
    null_default_fields = []

    serialized = handler(self)

    m = {}

    # simulate model_fields for test class
    model_fields = self.__class__.model_fields()
    for n, f in model_fields.items():
        k = f.alias or n
        val = serialized.get(k)
        serialized.pop(k, None)

        optional_nullable = k in optional_fields and k in nullable_fields
        is_set = (
            getattr(self, "__pydantic_fields_set__", set()).intersection({n})
            or k in null_default_fields
        )

        if val is not None and val != UNSET_SENTINEL:
            m[k] = val
        elif val != UNSET_SENTINEL and (
            not k in optional_fields or (optional_nullable and is_set)
        ):
            m[k] = val

    return m
from mistralai.models.agentupdaterequest import AgentUpdateRequest


# Handler function for serialization (identity for testing)
def handler(obj):
    # Simulate serialization by returning __dict__ with all fields
    d = {}
    for k in obj.__class__.__annotations__:
        v = getattr(obj, k, None)
        d[k] = v
    return d

# ------------------- UNIT TESTS -------------------

# 1. Basic Test Cases








def test_serialize_model_handoffs_none_and_empty():
    # handoffs set to None and to empty list
    req_none = AgentUpdateRequest(handoffs=None)
    req_empty = AgentUpdateRequest(handoffs=[])
    codeflash_output = serialize_model(req_none, handler); result_none = codeflash_output
    codeflash_output = serialize_model(req_empty, handler); result_empty = codeflash_output






#------------------------------------------------
import pytest
from mistralai.models.agentupdaterequest import AgentUpdateRequest

UNSET_SENTINEL = object()
UNSET = UNSET_SENTINEL

# OptionalNullable type
def OptionalNullable(type_):
    return type_

def model_serializer(mode=None):
    def decorator(fn):
        return fn
    return decorator

# Dummy BaseModel with Pydantic-like behavior
class BaseModel:
    model_fields = {}
    __pydantic_fields_set__ = set()
    def __init_subclass__(cls, **kwargs):
        # Build model_fields for all subclasses
        fields = {}
        for k, v in cls.__annotations__.items():
            fields[k] = type('Field', (), {'alias': None})()
        cls.model_fields = fields

    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
        self.__pydantic_fields_set__ = set(kwargs.keys())

# --- Dummy tool classes for tools field ---
class CodeInterpreterTool(BaseModel): type = "code_interpreter"
class DocumentLibraryTool(BaseModel): type = "document_library"
class FunctionTool(BaseModel): type = "function"
class ImageGenerationTool(BaseModel): type = "image_generation"
class WebSearchTool(BaseModel): type = "web_search"
class WebSearchPremiumTool(BaseModel): type = "web_search_premium"

# Dummy CompletionArgs
class CompletionArgs(BaseModel):
    def __init__(self, temperature=None, max_tokens=None):
        super().__init__(temperature=temperature, max_tokens=max_tokens)
    def __eq__(self, other):
        return (
            isinstance(other, CompletionArgs)
            and self.__dict__ == other.__dict__
        )
from mistralai.models.agentupdaterequest import AgentUpdateRequest


# Handler function to simulate Pydantic serialization
def default_handler(obj):
    # Only include fields that are not UNSET_SENTINEL
    d = {}
    for k in obj.model_fields:
        v = getattr(obj, k)
        d[k] = v
    return d

# ------------------ UNIT TESTS ------------------

# Basic Test Cases

def test_serialize_model_minimal_fields():
    """Test with no fields set (all defaults/unset). Should return empty dict."""
    req = AgentUpdateRequest()
    codeflash_output = req.serialize_model(default_handler); result = codeflash_output # 13.0μs -> 8.82μs (47.8% faster)

def test_serialize_model_only_instructions():
    """Test with only instructions set."""
    req = AgentUpdateRequest(instructions="Do something")
    codeflash_output = req.serialize_model(default_handler); result = codeflash_output # 11.1μs -> 7.73μs (43.2% faster)


def test_serialize_model_some_fields_set_none():
    """Test with some fields explicitly set to None."""
    req = AgentUpdateRequest(
        instructions=None,
        model=None,
        name=None,
        description=None,
        handoffs=None
    )
    codeflash_output = req.serialize_model(default_handler); result = codeflash_output # 11.7μs -> 9.19μs (27.5% faster)


def test_serialize_model_instructions_empty_string():
    """Test with instructions set to empty string."""
    req = AgentUpdateRequest(instructions="")
    codeflash_output = req.serialize_model(default_handler); result = codeflash_output # 13.4μs -> 8.70μs (53.9% faster)

def test_serialize_model_handoffs_empty_list():
    """Test with handoffs set to an empty list."""
    req = AgentUpdateRequest(handoffs=[])
    codeflash_output = req.serialize_model(default_handler); result = codeflash_output # 10.9μs -> 7.09μs (53.7% faster)

def test_serialize_model_tools_empty_list():
    """Test with tools set to an empty list."""
    req = AgentUpdateRequest(tools=[])
    codeflash_output = req.serialize_model(default_handler); result = codeflash_output # 10.4μs -> 7.12μs (45.5% faster)

def test_serialize_model_completion_args_none():
    """Test with completion_args explicitly set to None."""
    req = AgentUpdateRequest(completion_args=None)
    codeflash_output = req.serialize_model(default_handler); result = codeflash_output # 10.3μs -> 6.99μs (47.2% faster)


def test_serialize_model_tools_none():
    """Test with tools explicitly set to None."""
    req = AgentUpdateRequest(tools=None)
    codeflash_output = req.serialize_model(default_handler); result = codeflash_output # 13.1μs -> 8.73μs (49.9% faster)

def test_serialize_model_handoffs_none_vs_unset():
    """Test difference between handoffs=None and handoffs=UNSET."""
    req_none = AgentUpdateRequest(handoffs=None)
    req_unset = AgentUpdateRequest()
    codeflash_output = req_none.serialize_model(default_handler); result_none = codeflash_output # 10.7μs -> 7.44μs (43.6% faster)
    codeflash_output = req_unset.serialize_model(default_handler); result_unset = codeflash_output # 6.06μs -> 3.89μs (55.8% faster)




def test_serialize_model_large_sparse_fields():
    """Test with only one large field set, rest unset."""
    handoffs = [f"handoff_{i}" for i in range(999)]
    req = AgentUpdateRequest(handoffs=handoffs)
    codeflash_output = req.serialize_model(default_handler); result = codeflash_output # 13.2μs -> 9.16μs (43.6% faster)
# 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-AgentUpdateRequest.serialize_model-mh4km42u and push.