⚡️ Speed up method Messages.format by 527%

[codeflash-ai]

📄 527% (5.27x) speedup for Messages.format in guardrails/prompt/messages.py

⏱️ Runtime : 57.2 milliseconds → 9.11 milliseconds (best of 58 runs)

📝 Explanation and details

The optimization achieves a 527% speedup by addressing two major performance bottlenecks identified in the line profiler results:

1. Template Variable Extraction (32% faster)

• Original: Used collections.defaultdict with Template.safe_substitute() to extract variables, creating unnecessary dictionary objects and performing full template substitution
• Optimized: Replaced with direct regex pattern matching using _TEMPLATE_VAR_PATTERN = re.compile(r"\${([_a-zA-Z][_a-zA-Z0-9]*)")
• Why faster: Avoids Template object instantiation and substitution overhead, directly parsing variable names from the string

2. Kwargs Filtering (88% faster)

• Original: {k: v for k, v in kwargs.items() if k in vars} performed O(n×m) operations, iterating through all kwargs for each message
• Optimized: {k: kwargs[k] for k in vars if k in kwargs_keys} with pre-computed kwargs_keys = set(kwargs.keys()) reduces to O(n) set lookups
• Why faster: Set membership testing is O(1) vs list membership O(m), and avoids repeated kwargs.items() iteration

Performance Impact by Test Scale:

• Small messages (1-2 variables): 1-6% improvement - overhead reduction minimal
• Medium scale (100+ messages): 24-40% improvement - reduced per-message overhead compounds
• Large scale (1000+ variables/messages): 400-1000%+ improvement - algorithmic improvements dominate at scale

The optimizations are particularly effective for applications processing many messages with multiple template variables, where the O(n×m) complexity reduction and template parsing efficiency gains provide exponential benefits.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import collections
import re
from string import Template
from typing import Dict, List, Optional, Union

# imports
import pytest  # used for our unit tests
from guardrails.prompt.messages import Messages

# --- Begin: Minimal stubs for dependencies ---

# Minimal stub for constants
class _Constants(dict):
    def __getitem__(self, key):
        # Provide a few example constants
        d = {
            "FOO": "foo_value",
            "BAR": "bar_value",
            "LONG": "x" * 500,
        }
        return d[key]

constants = _Constants()

# Minimal stub for NamespaceTemplate
class NamespaceTemplate(Template):
    pass

# Minimal stub for Instructions and Prompt
class Instructions:
    pass

class Prompt:
    def __init__(self, source):
        self._source = source
from guardrails.prompt.messages import Messages

# --- End: Messages class ---

# --- Begin: Unit tests for Messages.format ---

# Basic Test Cases

def test_format_single_variable():
    # Test with one message, one variable
    messages = Messages([
        {"role": "user", "content": "Hello ${name}!"}
    ])
    codeflash_output = messages.format(name="Alice"); formatted = codeflash_output # 15.4μs -> 14.7μs (4.85% faster)

def test_format_multiple_variables():
    # Test with multiple variables in one message
    messages = Messages([
        {"role": "user", "content": "Hi ${first} ${last}."}
    ])
    codeflash_output = messages.format(first="John", last="Doe"); formatted = codeflash_output # 16.3μs -> 15.6μs (4.28% faster)

def test_format_multiple_messages():
    # Test with multiple messages, each with variables
    messages = Messages([
        {"role": "user", "content": "Hello ${name}!"},
        {"role": "assistant", "content": "Goodbye ${name}."}
    ])
    codeflash_output = messages.format(name="Bob"); formatted = codeflash_output # 18.9μs -> 19.0μs (0.332% slower)

def test_format_no_variables():
    # Test with message containing no variables
    messages = Messages([
        {"role": "system", "content": "Static message."}
    ])
    codeflash_output = messages.format(name="ShouldNotMatter"); formatted = codeflash_output # 9.53μs -> 9.39μs (1.57% faster)

def test_format_partial_kwargs():
    # Test with extra kwargs that are not in the template
    messages = Messages([
        {"role": "user", "content": "Hi ${name}."}
    ])
    codeflash_output = messages.format(name="Eve", unused="X"); formatted = codeflash_output # 15.0μs -> 14.5μs (3.83% faster)

def test_format_prompt_object():
    # Test with Prompt object as content
    messages = Messages([
        {"role": "user", "content": Prompt("Hi ${name}.")}
    ])
    codeflash_output = messages.format(name="Charlie"); formatted = codeflash_output # 16.0μs -> 15.9μs (0.634% faster)

# Edge Test Cases

def test_format_missing_variable():
    # Test with missing variable (should leave ${var} untouched)
    messages = Messages([
        {"role": "user", "content": "Hello ${name} and ${other}."}
    ])
    codeflash_output = messages.format(name="Dana"); formatted = codeflash_output # 16.7μs -> 16.0μs (4.40% faster)

def test_format_empty_string():
    # Test with empty message content
    messages = Messages([
        {"role": "user", "content": ""}
    ])
    codeflash_output = messages.format(anything="value"); formatted = codeflash_output # 9.48μs -> 9.16μs (3.46% faster)

def test_format_variable_with_special_characters():
    # Test variable values with special characters
    messages = Messages([
        {"role": "user", "content": "Special: ${sym}"}
    ])
    codeflash_output = messages.format(sym="!@#$%^&*()"); formatted = codeflash_output # 14.7μs -> 15.0μs (2.17% slower)

def test_format_variable_adjacent():
    # Test variables adjacent to each other
    messages = Messages([
        {"role": "user", "content": "${a}${b}${c}"}
    ])
    codeflash_output = messages.format(a="1", b="2", c="3"); formatted = codeflash_output # 16.0μs -> 15.6μs (2.20% faster)

def test_format_variable_with_number():
    # Test variable names with numbers
    messages = Messages([
        {"role": "user", "content": "ID: ${id1} and ${id2}"}
    ])
    codeflash_output = messages.format(id1="X", id2="Y"); formatted = codeflash_output # 16.0μs -> 15.2μs (5.51% faster)



def test_format_with_output_schema():
    # Test substitution of output_schema and xml_output_schema
    messages = Messages([
        {"role": "user", "content": "Schema: ${output_schema}, XML: ${xml_output_schema}"}
    ], output_schema="OUT", xml_output_schema="XML")
    codeflash_output = messages.format(); formatted = codeflash_output # 8.42μs -> 8.43μs (0.059% slower)

def test_format_with_escaped_dollar():
    # Test with escaped dollar signs
    messages = Messages([
        {"role": "user", "content": "Cost: $${price}"}
    ])
    codeflash_output = messages.format(price="100"); formatted = codeflash_output # 15.8μs -> 14.8μs (6.33% faster)

def test_format_with_unusual_variable_names():
    # Test with underscore and mixed-case variable names
    messages = Messages([
        {"role": "user", "content": "Hello ${First_Name}!"}
    ])
    codeflash_output = messages.format(First_Name="Alex"); formatted = codeflash_output # 14.3μs -> 14.0μs (1.83% faster)


def test_format_with_no_messages():
    # Test with empty messages list
    messages = Messages([])
    codeflash_output = messages.format(); formatted = codeflash_output # 1.32μs -> 1.82μs (27.6% slower)

# Large Scale Test Cases

def test_format_large_number_of_messages():
    # Test with 1000 messages, each with a variable
    n = 1000
    messages = Messages([
        {"role": "user", "content": f"Message {i}: ${'var'+str(i)}"} for i in range(n)
    ])
    # Prepare kwargs for all variables
    kwargs = {f"var{i}": f"value{i}" for i in range(n)}
    codeflash_output = messages.format(**kwargs); formatted = codeflash_output # 23.9ms -> 2.15ms (1014% faster)
    for i in range(n):
        expected = f"Message {i}: value{i}"

def test_format_large_message_content():
    # Test with a single message with a very large content and many variables
    n = 500
    content = " ".join([f"${{var{i}}}" for i in range(n)])
    messages = Messages([
        {"role": "user", "content": content}
    ])
    kwargs = {f"var{i}": str(i) for i in range(n)}
    codeflash_output = messages.format(**kwargs); formatted = codeflash_output # 1.09ms -> 369μs (197% faster)
    expected = " ".join([str(i) for i in range(n)])


def test_format_large_message_with_missing_vars():
    # Test with many variables, only some provided
    n = 200
    content = " ".join([f"${{var{i}}}" for i in range(n)])
    messages = Messages([
        {"role": "user", "content": content}
    ])
    # Only provide half the variables
    kwargs = {f"var{i}": f"VAL{i}" for i in range(n//2)}
    codeflash_output = messages.format(**kwargs); formatted = codeflash_output # 227μs -> 161μs (41.0% faster)
    expected = " ".join([f"VAL{i}" if i < n//2 else f"${{var{i}}}" for i in range(n)])

def test_format_performance_large_scale():
    # Test that formatting 1000 messages does not take excessive time
    import time
    n = 1000
    messages = Messages([
        {"role": "user", "content": f"Hi ${'name'+str(i)}"} for i in range(n)
    ])
    kwargs = {f"name{i}": f"User{i}" for i in range(n)}
    start = time.time()
    codeflash_output = messages.format(**kwargs); formatted = codeflash_output # 23.9ms -> 2.13ms (1021% faster)
    duration = time.time() - start
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import collections
# --- Messages class (from guardrails/prompt/messages.py) ---
import re
from string import Template
from typing import Dict, List, Optional, Union

# imports
import pytest
from guardrails.prompt.messages import Messages


# --- Mocks for dependencies ---
# Minimal mock for NamespaceTemplate
class NamespaceTemplate(Template):
    pass

# Minimal mock for constants
constants = {
    "FOO": "foo_value",
    "BAR": "bar_value",
    "LONG": "l" * 1000,
}

# Minimal mock for Prompt and Instructions
class Prompt:
    def __init__(self, source):
        self._source = source

class Instructions:
    def __init__(self, source):
        self._source = source
from guardrails.prompt.messages import Messages

# --- Unit Tests ---

# 1. Basic Test Cases

def test_basic_single_variable_substitution():
    """Test single variable substitution in one message."""
    msgs = Messages([{"role": "user", "content": "Hello, ${name}!"}])
    codeflash_output = msgs.format(name="Alice"); result = codeflash_output # 14.6μs -> 14.5μs (1.01% faster)

def test_basic_multiple_variable_substitution():
    """Test multiple variable substitution in one message."""
    msgs = Messages([{"role": "user", "content": "Hi ${first}, meet ${second}."}])
    codeflash_output = msgs.format(first="Bob", second="Carol"); result = codeflash_output # 16.2μs -> 15.4μs (4.97% faster)

def test_basic_multiple_messages():
    """Test formatting multiple messages in the list."""
    msgs = Messages([
        {"role": "user", "content": "Hi ${name}."},
        {"role": "assistant", "content": "Hello ${name}!"}
    ])
    codeflash_output = msgs.format(name="Dave"); result = codeflash_output # 18.5μs -> 17.4μs (6.50% faster)

def test_basic_no_variables():
    """Test message with no variables remains unchanged."""
    msgs = Messages([{"role": "user", "content": "Just text."}])
    codeflash_output = msgs.format(name="Ignored"); result = codeflash_output # 9.49μs -> 9.49μs (0.021% faster)

def test_basic_prompt_object():
    """Test substitution when content is a Prompt object."""
    msgs = Messages([{"role": "user", "content": Prompt("Hello ${name}!") }])
    codeflash_output = msgs.format(name="Eve"); result = codeflash_output # 16.1μs -> 15.9μs (1.24% faster)

# 2. Edge Test Cases

def test_edge_missing_variable():
    """Test missing variable is replaced with empty string."""
    msgs = Messages([{"role": "user", "content": "Hello, ${name}!"}])
    codeflash_output = msgs.format(); result = codeflash_output # 14.2μs -> 14.1μs (0.775% faster)

def test_edge_extra_kwargs():
    """Test extra kwargs are ignored if not present in template."""
    msgs = Messages([{"role": "user", "content": "Hi ${name}."}])
    codeflash_output = msgs.format(name="Frank", unused="should not appear"); result = codeflash_output # 14.7μs -> 14.4μs (2.36% faster)

def test_edge_empty_message_list():
    """Test formatting with empty message list returns empty Messages."""
    msgs = Messages([])
    codeflash_output = msgs.format(name="Ghost"); result = codeflash_output # 1.56μs -> 1.94μs (19.8% slower)

def test_edge_empty_string_content():
    """Test formatting with empty string content."""
    msgs = Messages([{"role": "user", "content": ""}])
    codeflash_output = msgs.format(any="thing"); result = codeflash_output # 9.72μs -> 9.52μs (2.07% faster)

def test_edge_nested_variable_names():
    """Test variable names that are substrings of other variable names."""
    msgs = Messages([{"role": "user", "content": "X: ${x}, XY: ${xy}"}])
    codeflash_output = msgs.format(x="1", xy="2"); result = codeflash_output # 16.1μs -> 15.4μs (4.49% faster)

def test_edge_variable_with_digits_and_underscore():
    """Test variable names with digits and underscores."""
    msgs = Messages([{"role": "user", "content": "ID: ${user_1_id}"}])
    codeflash_output = msgs.format(user_1_id="42"); result = codeflash_output # 14.1μs -> 13.9μs (0.918% faster)

def test_edge_content_is_instructions_object():
    """Test substitution when content is an Instructions object."""
    msgs = Messages([{"role": "user", "content": Instructions("Do ${action}.") }])
    codeflash_output = msgs.format(action="jump"); result = codeflash_output # 16.0μs -> 15.9μs (0.864% faster)



def test_edge_output_schema_substitution():
    """Test output_schema substitution on init."""
    msgs = Messages([{"role": "user", "content": "Schema: ${output_schema}."}], output_schema="schema_value")
    codeflash_output = msgs.format(); result = codeflash_output # 8.45μs -> 8.23μs (2.67% faster)

def test_edge_xml_output_schema_substitution():
    """Test xml_output_schema substitution on init."""
    msgs = Messages([{"role": "user", "content": "XML: ${xml_output_schema}."}], xml_output_schema="xml_value")
    codeflash_output = msgs.format(); result = codeflash_output # 7.99μs -> 7.98μs (0.125% faster)


def test_edge_variable_with_special_characters():
    """Test that variables with invalid names are not substituted."""
    msgs = Messages([{"role": "user", "content": "Special: ${na-me}"}])
    codeflash_output = msgs.format(**{"na-me": "fail"}); result = codeflash_output # 14.7μs -> 13.9μs (5.57% faster)

def test_edge_variable_adjacent_to_text():
    """Test variable directly adjacent to text."""
    msgs = Messages([{"role": "user", "content": "Hello${name}World"}])
    codeflash_output = msgs.format(name="Zed"); result = codeflash_output # 14.7μs -> 14.2μs (4.04% faster)

def test_edge_multiple_same_variable():
    """Test multiple occurrences of same variable in message."""
    msgs = Messages([{"role": "user", "content": "${x} + ${x} = ${y}"}])
    codeflash_output = msgs.format(x="2", y="4"); result = codeflash_output # 16.0μs -> 15.4μs (3.60% faster)

def test_edge_variable_not_in_kwargs():
    """Test variable present in template but not in kwargs."""
    msgs = Messages([{"role": "user", "content": "Missing: ${missing}"}])
    codeflash_output = msgs.format(); result = codeflash_output # 14.3μs -> 13.9μs (2.66% faster)

def test_edge_variable_with_none_value():
    """Test variable with None value is converted to empty string."""
    msgs = Messages([{"role": "user", "content": "None: ${val}"}])
    codeflash_output = msgs.format(val=None); result = codeflash_output # 14.5μs -> 13.9μs (4.18% faster)

# 3. Large Scale Test Cases

def test_large_many_messages():
    """Test formatting with a large number of messages (1000)."""
    msgs = Messages([
        {"role": "user", "content": "Hello ${name}!"}
        for _ in range(1000)
    ])
    codeflash_output = msgs.format(name="Big"); result = codeflash_output # 2.86ms -> 2.30ms (24.3% faster)
    for msg in result.source:
        pass

def test_large_many_variables():
    """Test formatting with a message containing many variables."""
    template = " ".join([f"${{var{i}}}" for i in range(1000)])
    msgs = Messages([{"role": "user", "content": template}])
    kwargs = {f"var{i}": str(i) for i in range(1000)}
    codeflash_output = msgs.format(**kwargs); result = codeflash_output # 3.57ms -> 704μs (406% faster)
    expected = " ".join(str(i) for i in range(1000))

def test_large_long_content():
    """Test formatting a message with very long content."""
    long_var = "x" * 1000
    msgs = Messages([{"role": "user", "content": f"Start {long_var} End"}])
    codeflash_output = msgs.format(); result = codeflash_output # 9.99μs -> 10.1μs (1.45% slower)

def test_large_long_variable_value():
    """Test substitution with very long variable value."""
    long_val = "y" * 1000
    msgs = Messages([{"role": "user", "content": "Long: ${longvar}"}])
    codeflash_output = msgs.format(longvar=long_val); result = codeflash_output # 15.4μs -> 14.7μs (4.36% faster)

def test_large_combined():
    """Test formatting with many messages and many variables."""
    template = " ".join([f"${{var{i}}}" for i in range(10)])
    msgs = Messages([
        {"role": "user", "content": template}
        for _ in range(100)
    ])
    kwargs = {f"var{i}": str(i) for i in range(10)}
    codeflash_output = msgs.format(**kwargs); result = codeflash_output # 1.11ms -> 792μs (39.8% faster)
    expected = " ".join(str(i) for i in range(10))
    for msg in result.source:
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

⏪ Replay Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_pytest_testsunit_teststest_guard_log_py_testsintegration_teststest_guard_py_testsunit_testsvalidator__replay_test_0.py::test_guardrails_prompt_messages_Messages_format	38.0μs	36.9μs	2.94%✅

To edit these changes git checkout codeflash/optimize-Messages.format-mh2jlxnj and push.