⚡️ Speed up method Memory.set_messages by 88% in PR #1059 (feat/agentic-codeflash)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1059

If you approve this dependent PR, these changes will be merged into the original PR branch feat/agentic-codeflash.

This PR will be automatically closed if the original PR is merged.

---

📄 88% (0.88x) speedup for Memory.set_messages in codeflash/agent/memory.py

⏱️ Runtime : 659 microseconds → 350 microseconds (best of 189 runs)

📝 Explanation and details

The optimization replaces the call to encoded_tokens_len(message["content"]) with an inlined approximation len(message["content"]) // 4 in the get_total_tokens() method.

What changed:

• Instead of calling encoded_tokens_len() for each message (which multiplies len(s) by 0.25 and converts to int), the code now directly computes len(message["content"]) // 4 inline.

Why it's faster:

1. Eliminates function call overhead: Each call to encoded_tokens_len() incurs Python function call overhead (argument passing, stack frame creation, return). With hundreds of messages, these microseconds accumulate significantly. The line profiler shows get_total_tokens() dropping from 5.06ms to 2.32ms—a 54% reduction.

2. Avoids floating-point arithmetic: The original code multiplies by 0.25 (float) then converts to int. The optimized version uses integer division (// 4), which is faster as it stays in integer arithmetic throughout.

3. Better CPU cache locality: Inlining keeps the hot loop tighter, improving instruction cache utilization during the list comprehension.

Impact on workloads:
The annotated tests show consistent speedups across all scenarios:

• Small workloads (single message): ~25-30% faster
• Medium workloads (100 messages): ~85-93% faster
• Large workloads (1000 messages): ~106-110% faster

The speedup scales linearly with message count because the optimization eliminates per-message overhead. Functions that process many messages (batch operations, conversation histories) benefit most. Since set_messages() calls get_total_tokens() to check against max_tokens, any code path that validates message lists sees this improvement.

Test case performance:

• Best for large-scale scenarios (500-1000 messages): 85-110% speedup
• Good for typical workloads (10-100 messages): 30-93% speedup
• Minimal but positive impact on edge cases (empty lists, single messages): 2-34% speedup

The optimization maintains identical behavior—both int(len(s) * 0.25) and len(s) // 4 produce the same token approximation for all string lengths.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from typing import Dict, List

# imports
import pytest  # used for our unit tests

from codeflash.agent.memory import Memory
from codeflash.code_utils.code_utils import encoded_tokens_len

# Unit tests for Memory.set_messages
#
# Each test is self-contained and deterministic. Tests are grouped from basic to edge
# to large-scale scenarios. Comments explain the intent and the assertions made.


def test_basic_set_and_get_messages():
    # Basic functionality: setting a typical list of message dicts should return the same list
    mem = Memory()  # create a Memory instance
    messages: List[Dict[str, str]] = [
        {"role": "user", "content": "Hello"},
        {"role": "assistant", "content": "Hi there!"},
    ]
    # set_messages should store and return the same object (assignment -> same identity)
    codeflash_output = mem.set_messages(messages)
    returned = codeflash_output  # 2.38μs -> 1.77μs (34.5% faster)
    # get_total_tokens should sum encoded token lengths of each message content
    expected_tokens = sum(encoded_tokens_len(m["content"]) for m in messages)


def test_empty_messages():
    # Edge case: empty list of messages
    mem = Memory()
    messages: List[Dict[str, str]] = []
    codeflash_output = mem.set_messages(messages)
    returned = codeflash_output  # 1.29μs -> 1.32μs (2.34% slower)


def test_non_string_content_raises_type_error():
    # Edge case: message content that is not a string should cause a TypeError when tokenizing
    mem = Memory()
    # Use an integer as content which will break len(s) inside encoded_tokens_len
    messages = [{"role": "user", "content": 123}]
    # set_messages calls get_total_tokens internally; this should raise a TypeError
    with pytest.raises(TypeError):
        mem.set_messages(messages)  # 4.14μs -> 3.83μs (8.10% faster)


def test_missing_content_key_raises_key_error():
    # Edge case: a message dict missing "content" key should raise KeyError when computing tokens
    mem = Memory()
    messages = [{"role": "user"}]  # missing "content"
    with pytest.raises(KeyError):
        mem.set_messages(messages)  # 3.12μs -> 3.13μs (0.320% slower)


def test_non_dict_message_item_raises_type_error():
    # Edge case: list contains an element that is not a dict (e.g., a string)
    mem = Memory()
    messages = ["not a dict"]
    # Attempting to index into the list element as message["content"] should raise TypeError
    with pytest.raises(TypeError):
        mem.set_messages(messages)  # 3.47μs -> 3.44μs (0.902% faster)


def test_mutable_input_reference_preserved():
    # Behavior test: set_messages assigns the list by reference, so subsequent mutations to the
    # original list should be visible inside Memory._messages
    mem = Memory()
    messages: List[Dict[str, str]] = [{"role": "user", "content": "start"}]
    codeflash_output = mem.set_messages(messages)
    returned = codeflash_output  # 2.21μs -> 1.76μs (25.1% faster)
    # Mutate the original list: append a new message
    messages.append({"role": "assistant", "content": "response"})


def test_max_tokens_threshold_behavior():
    # Edge: when total tokens exceed max_tokens, current implementation does nothing special
    mem = Memory()
    # Set a very small max_tokens to trigger the "exceeds" branch
    mem.max_tokens = 1
    messages = [{"role": "user", "content": "This is a long message that will exceed the tiny max tokens."}]
    # Although get_total_tokens() will be > max_tokens, implementation only has a 'pass' in that branch,
    # so set_messages should still return the messages and not raise an exception.
    codeflash_output = mem.set_messages(messages)
    returned = codeflash_output  # 2.08μs -> 1.65μs (26.1% faster)


def test_large_scale_messages_1000():
    # Large-scale test: use 1000 messages (within the constraint of <= 1000) to validate scalability.
    # Each message has 100 characters -> encoded_tokens_len should return int(100 * 0.25) == 25 per message.
    mem = Memory()
    single_content = "a" * 100  # 100 chars
    count = 1000  # maximum allowed by the guidance
    messages = [{"role": "user", "content": single_content} for _ in range(count)]
    codeflash_output = mem.set_messages(messages)
    returned = codeflash_output  # 136μs -> 64.9μs (110% faster)
    # Validate total token accounting
    expected_per_message = encoded_tokens_len(single_content)


def test_unicode_content_token_count():
    # Edge: content containing multi-byte/unicode characters (emojis) should be counted by Python's len()
    mem = Memory()
    # 40 emoji characters; len counts code points, so encoded tokens should be int(40 * 0.25) == 10
    emoji_content = "🙂" * 40
    messages = [{"role": "user", "content": emoji_content}]
    mem.set_messages(messages)  # 2.11μs -> 1.64μs (28.7% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from codeflash.agent.memory import Memory
from codeflash.code_utils.code_utils import encoded_tokens_len

# ============================================================================
# BASIC TEST CASES
# ============================================================================


def test_set_messages_empty_list():
    """Test setting an empty list of messages."""
    memory = Memory()
    codeflash_output = memory.set_messages([])
    result = codeflash_output  # 1.40μs -> 1.37μs (2.11% faster)


def test_set_messages_single_message():
    """Test setting a single message."""
    memory = Memory()
    messages = [{"content": "Hello, world!"}]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 2.20μs -> 1.69μs (30.2% faster)


def test_set_messages_multiple_messages():
    """Test setting multiple messages."""
    memory = Memory()
    messages = [{"content": "First message"}, {"content": "Second message"}, {"content": "Third message"}]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 2.45μs -> 1.88μs (30.3% faster)


def test_set_messages_returns_same_reference():
    """Test that the function returns the same list reference."""
    memory = Memory()
    messages = [{"content": "Test"}]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 2.02μs -> 1.59μs (27.1% faster)


def test_set_messages_overwrites_previous():
    """Test that calling set_messages overwrites previous messages."""
    memory = Memory()
    first_messages = [{"content": "First"}]
    second_messages = [{"content": "Second"}]

    memory.set_messages(first_messages)  # 2.01μs -> 1.58μs (27.2% faster)

    memory.set_messages(second_messages)  # 961ns -> 731ns (31.5% faster)


def test_set_messages_with_empty_content():
    """Test setting messages with empty content strings."""
    memory = Memory()
    messages = [{"content": ""}]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 1.93μs -> 1.64μs (17.7% faster)


def test_set_messages_with_multiple_keys():
    """Test that messages with multiple keys are preserved."""
    memory = Memory()
    messages = [
        {"content": "Hello", "role": "user", "id": "123"},
        {"content": "Hi there", "role": "assistant", "id": "456"},
    ]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 2.28μs -> 1.70μs (34.1% faster)


def test_set_messages_token_counting():
    """Test that total tokens are calculated correctly after setting messages."""
    memory = Memory()
    message_content = "a" * 100  # 100 characters = 25 tokens (100 * 0.25)
    messages = [{"content": message_content}]
    memory.set_messages(messages)  # 1.98μs -> 1.57μs (26.1% faster)
    expected_tokens = encoded_tokens_len(message_content)


def test_set_messages_multiple_token_counting():
    """Test token counting with multiple messages."""
    memory = Memory()
    messages = [
        {"content": "a" * 100},  # 25 tokens
        {"content": "b" * 200},  # 50 tokens
        {"content": "c" * 80},  # 20 tokens
    ]
    memory.set_messages(messages)  # 2.48μs -> 1.87μs (32.1% faster)
    expected_tokens = 25 + 50 + 20  # 95 tokens total


# ============================================================================
# EDGE TEST CASES
# ============================================================================


def test_set_messages_with_special_characters():
    """Test messages containing special characters."""
    memory = Memory()
    messages = [
        {"content": "Hello! @#$%^&*() world"},
        {"content": "Line 1\nLine 2\tTabbed"},
        {"content": "Unicode: 你好 مرحبا"},
    ]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 2.25μs -> 1.82μs (23.6% faster)


def test_set_messages_with_very_long_content():
    """Test message with very long content string."""
    memory = Memory()
    long_content = "x" * 10000
    messages = [{"content": long_content}]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 2.20μs -> 1.83μs (20.2% faster)


def test_set_messages_with_only_whitespace():
    """Test message containing only whitespace characters."""
    memory = Memory()
    messages = [{"content": "   "}, {"content": "\n\n\n"}, {"content": "\t\t"}]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 2.46μs -> 1.79μs (36.9% faster)


def test_set_messages_near_token_limit():
    """Test messages that approach but don't exceed token limit."""
    memory = Memory()
    # Create content that results in tokens just under the limit
    # max_tokens = 16000, so we need content that gives ~16000 tokens
    # 1 token ≈ 4 characters, so 16000 * 4 = 64000 characters max
    content_size = int(memory.max_tokens * 4 * 0.99)  # 99% of limit
    messages = [{"content": "a" * content_size}]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 1.99μs -> 1.74μs (14.4% faster)
    total_tokens = memory.get_total_tokens()


def test_set_messages_exceeds_token_limit():
    """Test that function executes when messages exceed token limit."""
    memory = Memory()
    # Create content that significantly exceeds the limit
    content_size = int(memory.max_tokens * 4 * 2)  # 200% of limit
    messages = [{"content": "a" * content_size}]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 2.20μs -> 1.79μs (22.9% faster)


def test_set_messages_with_zero_length_key():
    """Test messages with various key names."""
    memory = Memory()
    messages = [{"content": "test", "": "empty_key"}, {"content": "test2"}]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 2.29μs -> 1.78μs (28.2% faster)


def test_set_messages_preserves_order():
    """Test that message order is preserved."""
    memory = Memory()
    messages = []
    for i in range(100):
        messages.append({"content": f"Message {i}"})

    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 15.6μs -> 8.09μs (93.1% faster)
    for i, msg in enumerate(result):
        pass


def test_set_messages_does_not_modify_input():
    """Test that the input list is not modified (shallow reference is OK)."""
    memory = Memory()
    messages = [{"content": "Original"}]
    original_content = messages[0]["content"]
    memory.set_messages(messages)  # 1.96μs -> 1.57μs (24.9% faster)


def test_set_messages_with_numeric_string_content():
    """Test messages with numeric content."""
    memory = Memory()
    messages = [{"content": "123456"}, {"content": "3.14159"}, {"content": "-999"}]
    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 2.30μs -> 1.77μs (29.9% faster)


def test_set_messages_consecutive_calls_independent():
    """Test that consecutive calls don't interfere with each other."""
    memory = Memory()

    codeflash_output = memory.set_messages([{"content": "First"}])
    result1 = codeflash_output  # 1.91μs -> 1.53μs (24.9% faster)
    tokens1 = memory.get_total_tokens()

    codeflash_output = memory.set_messages([{"content": "Second call"}])
    result2 = codeflash_output  # 782ns -> 661ns (18.3% faster)
    tokens2 = memory.get_total_tokens()


# ============================================================================
# LARGE SCALE TEST CASES
# ============================================================================


def test_set_messages_large_number_of_messages():
    """Test setting a large number of messages (500 messages)."""
    memory = Memory()
    messages = []
    for i in range(500):
        messages.append({"content": f"Message number {i}"})

    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 69.6μs -> 34.1μs (104% faster)


def test_set_messages_large_content_size():
    """Test setting messages with large individual content sizes."""
    memory = Memory()
    # Create 50 messages each with 5000 characters
    messages = []
    for i in range(50):
        messages.append({"content": "x" * 5000})

    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 10.3μs -> 6.17μs (66.4% faster)
    total_tokens = memory.get_total_tokens()


def test_set_messages_many_keys_per_message():
    """Test messages with many additional keys."""
    memory = Memory()
    messages = []
    for i in range(100):
        msg = {"content": f"Message {i}"}
        for j in range(20):  # Add 20 additional keys per message
            msg[f"key_{j}"] = f"value_{j}"
        messages.append(msg)

    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 16.8μs -> 9.09μs (85.1% faster)


def test_set_messages_varying_content_sizes():
    """Test messages with varying content sizes."""
    memory = Memory()
    messages = []
    for i in range(200):
        # Vary content size from 10 to 5000 characters
        size = 10 + (i * 25) % 5000
        messages.append({"content": "a" * size})

    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 35.4μs -> 20.9μs (69.3% faster)
    total_tokens = memory.get_total_tokens()


def test_set_messages_unicode_at_scale():
    """Test handling of unicode content at scale."""
    memory = Memory()
    unicode_strings = ["Hello 世界", "مرحبا بالعالم", "Привет мир", "Γεια σου κόσμε", "שלום עולם"]

    messages = []
    for i in range(100):
        messages.append({"content": unicode_strings[i % 5]})

    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 15.7μs -> 8.20μs (91.8% faster)
    total_tokens = memory.get_total_tokens()


def test_set_messages_performance_with_large_messages():
    """Test that function completes in reasonable time with large data."""
    memory = Memory()
    # Create a realistic large scenario: 300 messages with moderate content
    messages = []
    for i in range(300):
        messages.append(
            {"content": f"Message {i}: " + ("x" * 500), "timestamp": "2024-01-01T00:00:00Z", "sender": f"user_{i % 10}"}
        )

    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 46.6μs -> 25.2μs (85.2% faster)


def test_set_messages_empty_strings_at_scale():
    """Test handling of empty content strings at large scale."""
    memory = Memory()
    messages = []
    for i in range(1000):
        if i % 10 == 0:
            messages.append({"content": ""})
        else:
            messages.append({"content": f"Content {i}"})

    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 138μs -> 67.3μs (106% faster)
    # Count empty messages
    empty_count = sum(1 for msg in result if msg["content"] == "")


def test_set_messages_maximum_realistic_load():
    """Test with maximum realistic message load within limits."""
    memory = Memory()
    # Create 800 messages with minimal content to stay within reasonable bounds
    messages = []
    for i in range(800):
        messages.append({"content": f"M{i}"})

    codeflash_output = memory.set_messages(messages)
    result = codeflash_output  # 111μs -> 53.9μs (107% faster)
    # Verify the internal state is correctly updated
    total_tokens = memory.get_total_tokens()


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from codeflash.agent.memory import Memory


def test_Memory_set_messages():
    Memory.set_messages(Memory(), [])

🔎 Click to see Concolic Coverage Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_ik_zub_8/tmpj95ssauz/test_concolic_coverage.py::test_Memory_set_messages	1.35μs	1.34μs	0.745%✅

To edit these changes git checkout codeflash/optimize-pr1059-2026-01-15T14.17.06 and push.