⚡️ Speed up function reconstitue_entries by 100%

[codeflash-ai]

📄 100% (1.00x) speedup for reconstitue_entries in src/mistralai/extra/run/result.py

⏱️ Runtime : 85.5 microseconds → 42.7 microseconds (best of 104 runs)

📝 Explanation and details

The optimized code achieves a 100% speedup (from 85.5μs to 42.7μs) through two key optimizations:

1. Pre-sorting the dictionary items outside the loop

sorted_items = sorted(received_event_tracker.items())
for idx, events in sorted_items:

Instead of sorting on every iteration with sorted(received_event_tracker.items(), key=lambda x: x[0]), the sorting is done once upfront. This eliminates the lambda function overhead and repeated sorting calls that were consuming 26.5% of the original runtime.

2. Replacing isinstance() with type() for exact type checking

if type(first_event) is MessageOutputEvent:
elif type(first_event) is FunctionCallEvent:

Using type() instead of isinstance() avoids Method Resolution Order (MRO) traversal for subclass checking. Since the code expects exact type matches, this direct comparison is both faster and semantically correct, reducing the type checking overhead from 62.3% to 50.9% of total runtime.

Performance impact by test case:

• Small datasets (single events): 50-74% faster, showing the type checking optimization's impact
• Large mixed datasets (100+ entries): Up to 185% faster, demonstrating how the pre-sorting optimization scales with dataset size
• Empty inputs: Still 16-22% faster due to eliminated lambda overhead

These optimizations are particularly effective for workloads with frequent event processing and larger event dictionaries, which appear to be the common use case based on the test patterns.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from datetime import datetime

# imports
import pytest
from mistralai.extra.run.result import reconstitue_entries

# --- Minimal class definitions to support tests (mocking mistralai.models) ---

# Since we cannot import mistralai, we define minimal versions of the required classes and types.

class TextChunk:
    def __init__(self, text):
        self.text = text
    def __eq__(self, other):
        return isinstance(other, TextChunk) and self.text == other.text
    def __repr__(self):
        return f"TextChunk({self.text!r})"

class MessageOutputEvent:
    def __init__(self, content, created_at, id, agent_id, model, role):
        self.content = content
        self.created_at = created_at
        self.id = id
        self.agent_id = agent_id
        self.model = model
        self.role = role

class MessageOutputEntry:
    def __init__(self, content, created_at, id, agent_id, model, role):
        self.content = content
        self.created_at = created_at
        self.id = id
        self.agent_id = agent_id
        self.model = model
        self.role = role
    def __eq__(self, other):
        return (
            isinstance(other, MessageOutputEntry)
            and self.content == other.content
            and self.created_at == other.created_at
            and self.id == other.id
            and self.agent_id == other.agent_id
            and self.model == other.model
            and self.role == other.role
        )
    def __repr__(self):
        return (
            f"MessageOutputEntry(content={self.content!r}, created_at={self.created_at!r}, id={self.id!r}, "
            f"agent_id={self.agent_id!r}, model={self.model!r}, role={self.role!r})"
        )

class FunctionCallEvent:
    def __init__(self, name, arguments, created_at, id, tool_call_id):
        self.name = name
        self.arguments = arguments
        self.created_at = created_at
        self.id = id
        self.tool_call_id = tool_call_id

class FunctionCallEntry:
    def __init__(self, name, arguments, created_at, id, tool_call_id):
        self.name = name
        self.arguments = arguments
        self.created_at = created_at
        self.id = id
        self.tool_call_id = tool_call_id
    def __eq__(self, other):
        return (
            isinstance(other, FunctionCallEntry)
            and self.name == other.name
            and self.arguments == other.arguments
            and self.created_at == other.created_at
            and self.id == other.id
            and self.tool_call_id == other.tool_call_id
        )
    def __repr__(self):
        return (
            f"FunctionCallEntry(name={self.name!r}, arguments={self.arguments!r}, created_at={self.created_at!r}, "
            f"id={self.id!r}, tool_call_id={self.tool_call_id!r})"
        )
from mistralai.extra.run.result import reconstitue_entries

# --- Unit tests ---

# Basic Test Cases

def test_single_message_event_with_string_content():
    # Test with one MessageOutputEvent, content is a string
    event = MessageOutputEvent(
        content="Hello, world!",
        created_at=123,
        id="msg1",
        agent_id="agentA",
        model="modelX",
        role="user"
    )
    tracker = {0: [event]}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 2.73μs -> 1.57μs (74.0% faster)
    expected = [
        MessageOutputEntry(
            content="Hello, world!",
            created_at=123,
            id="msg1",
            agent_id="agentA",
            model="modelX",
            role="user"
        )
    ]

def test_single_message_event_with_textchunk_content():
    # Test with one MessageOutputEvent, content is a TextChunk
    event = MessageOutputEvent(
        content=TextChunk("Hi!"),
        created_at=456,
        id="msg2",
        agent_id="agentB",
        model="modelY",
        role="assistant"
    )
    tracker = {0: [event]}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 2.28μs -> 1.32μs (73.2% faster)
    expected = [
        MessageOutputEntry(
            content=[TextChunk("Hi!")],
            created_at=456,
            id="msg2",
            agent_id="agentB",
            model="modelY",
            role="assistant"
        )
    ]

def test_multiple_message_events_string_content():
    # Multiple MessageOutputEvent, content are strings, should concatenate
    events = [
        MessageOutputEvent(
            content="Hello, ",
            created_at=100,
            id="msg3",
            agent_id="agentC",
            model="modelZ",
            role="user"
        ),
        MessageOutputEvent(
            content="world!",
            created_at=100,
            id="msg3",
            agent_id="agentC",
            model="modelZ",
            role="user"
        )
    ]
    tracker = {0: events}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 1.99μs -> 1.30μs (52.6% faster)
    expected = [
        MessageOutputEntry(
            content="Hello, world!",
            created_at=100,
            id="msg3",
            agent_id="agentC",
            model="modelZ",
            role="user"
        )
    ]

def test_multiple_message_events_textchunk_content():
    # Multiple MessageOutputEvent, content are TextChunk, should merge text
    events = [
        MessageOutputEvent(
            content=TextChunk("Hi "),
            created_at=200,
            id="msg4",
            agent_id="agentD",
            model="modelW",
            role="assistant"
        ),
        MessageOutputEvent(
            content=TextChunk("there!"),
            created_at=200,
            id="msg4",
            agent_id="agentD",
            model="modelW",
            role="assistant"
        )
    ]
    tracker = {0: events}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 2.13μs -> 1.36μs (56.1% faster)
    expected = [
        MessageOutputEntry(
            content=[TextChunk("Hi there!")],
            created_at=200,
            id="msg4",
            agent_id="agentD",
            model="modelW",
            role="assistant"
        )
    ]

def test_single_function_call_event():
    # Test with one FunctionCallEvent
    event = FunctionCallEvent(
        name="do_something",
        arguments="arg1",
        created_at=111,
        id="func1",
        tool_call_id="toolA"
    )
    tracker = {0: [event]}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 1.96μs -> 1.30μs (50.9% faster)
    expected = [
        FunctionCallEntry(
            name="do_something",
            arguments="arg1",
            created_at=111,
            id="func1",
            tool_call_id="toolA"
        )
    ]

def test_multiple_function_call_events():
    # Multiple FunctionCallEvent, arguments are strings, should concatenate
    events = [
        FunctionCallEvent(
            name="do_something",
            arguments="foo",
            created_at=222,
            id="func2",
            tool_call_id="toolB"
        ),
        FunctionCallEvent(
            name="do_something",
            arguments="bar",
            created_at=222,
            id="func2",
            tool_call_id="toolB"
        )
    ]
    tracker = {0: events}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 2.01μs -> 1.24μs (61.6% faster)
    expected = [
        FunctionCallEntry(
            name="do_something",
            arguments="foobar",
            created_at=222,
            id="func2",
            tool_call_id="toolB"
        )
    ]

# Edge Test Cases

def test_empty_tracker():
    # Empty input dict should return empty list
    tracker = {}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 1.26μs -> 1.07μs (16.9% faster)

def test_message_event_mixed_chunks():
    # MessageOutputEvent with mixed str and TextChunk
    events = [
        MessageOutputEvent(
            content="Hello, ",
            created_at=333,
            id="msg5",
            agent_id="agentE",
            model="modelQ",
            role="user"
        ),
        MessageOutputEvent(
            content=TextChunk("world!"),
            created_at=333,
            id="msg5",
            agent_id="agentE",
            model="modelQ",
            role="user"
        )
    ]
    tracker = {0: events}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 2.12μs -> 1.28μs (66.4% faster)
    # Should merge into a single TextChunk
    expected = [
        MessageOutputEntry(
            content=[TextChunk("Hello, world!")],
            created_at=333,
            id="msg5",
            agent_id="agentE",
            model="modelQ",
            role="user"
        )
    ]

def test_message_event_chunks_with_empty_string():
    # MessageOutputEvent with empty string chunks
    events = [
        MessageOutputEvent(
            content="",
            created_at=444,
            id="msg6",
            agent_id="agentF",
            model="modelR",
            role="assistant"
        ),
        MessageOutputEvent(
            content="Hi!",
            created_at=444,
            id="msg6",
            agent_id="agentF",
            model="modelR",
            role="assistant"
        )
    ]
    tracker = {0: events}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 1.85μs -> 1.25μs (48.1% faster)
    expected = [
        MessageOutputEntry(
            content="Hi!",
            created_at=444,
            id="msg6",
            agent_id="agentF",
            model="modelR",
            role="assistant"
        )
    ]

def test_function_call_event_empty_arguments():
    # FunctionCallEvent with empty arguments
    events = [
        FunctionCallEvent(
            name="do_nothing",
            arguments="",
            created_at=555,
            id="func3",
            tool_call_id="toolC"
        )
    ]
    tracker = {0: events}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 1.99μs -> 1.26μs (57.9% faster)
    expected = [
        FunctionCallEntry(
            name="do_nothing",
            arguments="",
            created_at=555,
            id="func3",
            tool_call_id="toolC"
        )
    ]

def test_multiple_entries_different_types():
    # Tracker with both MessageOutputEvent and FunctionCallEvent
    tracker = {
        0: [
            MessageOutputEvent(
                content="Hi!",
                created_at=666,
                id="msg7",
                agent_id="agentG",
                model="modelS",
                role="user"
            )
        ],
        1: [
            FunctionCallEvent(
                name="greet",
                arguments="hello",
                created_at=777,
                id="func4",
                tool_call_id="toolD"
            )
        ]
    }
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 2.62μs -> 1.92μs (36.5% faster)
    expected = [
        MessageOutputEntry(
            content="Hi!",
            created_at=666,
            id="msg7",
            agent_id="agentG",
            model="modelS",
            role="user"
        ),
        FunctionCallEntry(
            name="greet",
            arguments="hello",
            created_at=777,
            id="func4",
            tool_call_id="toolD"
        )
    ]

def test_tracker_unsorted_keys():
    # Tracker keys are out of order, should return sorted
    tracker = {
        2: [
            FunctionCallEvent(
                name="sum",
                arguments="1",
                created_at=888,
                id="func5",
                tool_call_id="toolE"
            )
        ],
        0: [
            MessageOutputEvent(
                content="First",
                created_at=999,
                id="msg8",
                agent_id="agentH",
                model="modelT",
                role="user"
            )
        ],
        1: [
            MessageOutputEvent(
                content="Second",
                created_at=1000,
                id="msg9",
                agent_id="agentI",
                model="modelU",
                role="assistant"
            )
        ]
    }
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 3.10μs -> 1.91μs (62.0% faster)
    expected = [
        MessageOutputEntry(
            content="First",
            created_at=999,
            id="msg8",
            agent_id="agentH",
            model="modelT",
            role="user"
        ),
        MessageOutputEntry(
            content="Second",
            created_at=1000,
            id="msg9",
            agent_id="agentI",
            model="modelU",
            role="assistant"
        ),
        FunctionCallEntry(
            name="sum",
            arguments="1",
            created_at=888,
            id="func5",
            tool_call_id="toolE"
        )
    ]

# Large Scale Test Cases

def test_large_number_of_message_events():
    # Test with 500 MessageOutputEvent, all with string content, should concatenate
    n = 500
    events = [
        MessageOutputEvent(
            content=str(i),
            created_at=12345,
            id="msg_large",
            agent_id="agentLarge",
            model="modelLarge",
            role="user"
        )
        for i in range(n)
    ]
    tracker = {0: events}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 2.34μs -> 1.39μs (68.1% faster)
    expected_content = "".join(str(i) for i in range(n))
    expected = [
        MessageOutputEntry(
            content=expected_content,
            created_at=12345,
            id="msg_large",
            agent_id="agentLarge",
            model="modelLarge",
            role="user"
        )
    ]

def test_large_number_of_function_call_events():
    # Test with 500 FunctionCallEvent, arguments are single characters, should concatenate
    n = 500
    events = [
        FunctionCallEvent(
            name="bigfunc",
            arguments=chr(65 + (i % 26)),  # A-Z cycling
            created_at=54321,
            id="func_large",
            tool_call_id="toolLarge"
        )
        for i in range(n)
    ]
    tracker = {0: events}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 2.25μs -> 1.47μs (53.7% faster)
    expected_arguments = "".join(chr(65 + (i % 26)) for i in range(n))
    expected = [
        FunctionCallEntry(
            name="bigfunc",
            arguments=expected_arguments,
            created_at=54321,
            id="func_large",
            tool_call_id="toolLarge"
        )
    ]

def test_large_mixed_entries():
    # 100 MessageOutputEvent and 100 FunctionCallEvent, interleaved keys
    n = 100
    tracker = {}
    for i in range(n):
        tracker[i*2] = [
            MessageOutputEvent(
                content=f"msg{i}",
                created_at=10000+i,
                id=f"msg{i}",
                agent_id=f"agent{i}",
                model=f"model{i}",
                role="user"
            )
        ]
        tracker[i*2+1] = [
            FunctionCallEvent(
                name=f"func{i}",
                arguments=f"arg{i}",
                created_at=20000+i,
                id=f"func{i}",
                tool_call_id=f"tool{i}"
            )
        ]
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 43.7μs -> 15.3μs (185% faster)
    # Should be sorted by keys
    expected = []
    for i in range(n):
        expected.append(
            MessageOutputEntry(
                content=f"msg{i}",
                created_at=10000+i,
                id=f"msg{i}",
                agent_id=f"agent{i}",
                model=f"model{i}",
                role="user"
            )
        )
        expected.append(
            FunctionCallEntry(
                name=f"func{i}",
                arguments=f"arg{i}",
                created_at=20000+i,
                id=f"func{i}",
                tool_call_id=f"tool{i}"
            )
        )

def test_large_message_event_with_mixed_chunks():
    # 100 MessageOutputEvent, alternating str and TextChunk, should merge into one TextChunk
    n = 100
    events = []
    for i in range(n):
        if i % 2 == 0:
            events.append(MessageOutputEvent(
                content=str(i),
                created_at=30000,
                id="msgmix",
                agent_id="agentMix",
                model="modelMix",
                role="user"
            ))
        else:
            events.append(MessageOutputEvent(
                content=TextChunk(str(i)),
                created_at=30000,
                id="msgmix",
                agent_id="agentMix",
                model="modelMix",
                role="user"
            ))
    tracker = {0: events}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 2.15μs -> 1.40μs (53.6% faster)
    expected_content = "".join(str(i) for i in range(n))
    expected = [
        MessageOutputEntry(
            content=[TextChunk(expected_content)],
            created_at=30000,
            id="msgmix",
            agent_id="agentMix",
            model="modelMix",
            role="user"
        )
    ]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import typing
from dataclasses import dataclass

# imports
import pytest
from mistralai.extra.run.result import reconstitue_entries

# --- Minimal stubs for the required classes/types ---

@dataclass
class TextChunk:
    text: str

MessageOutputContentChunks = typing.Union[str, TextChunk]
MessageOutputEntryContent = typing.Union[str, list[MessageOutputContentChunks]]
MessageOutputEventContent = typing.Union[str, TextChunk]

@dataclass
class MessageOutputEvent:
    content: MessageOutputEventContent
    created_at: int
    id: str
    agent_id: str
    model: str
    role: str

@dataclass
class MessageOutputEntry:
    content: MessageOutputEntryContent
    created_at: int
    id: str
    agent_id: str
    model: str
    role: str

@dataclass
class FunctionCallEvent:
    name: str
    arguments: str
    created_at: int
    id: str
    tool_call_id: str

@dataclass
class FunctionCallEntry:
    name: str
    arguments: str
    created_at: int
    id: str
    tool_call_id: str

# Stubs for unused types in the function
class AgentHandoffEntry: pass
class FunctionResultEntry: pass
class ToolExecutionEntry: pass
ConversationEventsData = typing.Union[MessageOutputEvent, FunctionCallEvent]
RunOutputEntries = typing.Union[
    MessageOutputEntry,
    FunctionCallEntry,
    FunctionResultEntry,
    AgentHandoffEntry,
    ToolExecutionEntry,
]
from mistralai.extra.run.result import reconstitue_entries

# --- Unit tests for reconstitue_entries ---

# 1. Basic Test Cases







def test_mixed_message_and_function_entries():
    """Test multiple keys with both MessageOutputEvent and FunctionCallEvent."""
    tracker = {
        0: [
            MessageOutputEvent(
                content="Hi.",
                created_at=7,
                id="msg5",
                agent_id="agentE",
                model="modelB",
                role="user"
            )
        ],
        1: [
            FunctionCallEvent(
                name="sum",
                arguments="1,2",
                created_at=8,
                id="fn3",
                tool_call_id="toolC"
            )
        ]
    }
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 3.69μs -> 2.16μs (70.8% faster)

# 2. Edge Test Cases

def test_empty_tracker_returns_empty_list():
    """Test that an empty tracker returns an empty list."""
    tracker = {}
    codeflash_output = reconstitue_entries(tracker); result = codeflash_output # 1.34μs -> 1.09μs (22.7% faster)











#------------------------------------------------
from mistralai.extra.run.result import reconstitue_entries
import pytest

def test_reconstitue_entries():
    with pytest.raises(IndexError):
        reconstitue_entries({2: []})

def test_reconstitue_entries_2():
    reconstitue_entries({})

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_dwma23ol/tmpb2wpc4dc/test_concolic_coverage.py::test_reconstitue_entries	2.51μs	1.83μs	37.0%✅
codeflash_concolic_dwma23ol/tmpb2wpc4dc/test_concolic_coverage.py::test_reconstitue_entries_2	1.46μs	1.24μs	17.7%✅

To edit these changes git checkout codeflash/optimize-reconstitue_entries-mh2yvfc2 and push.