⚡️ Speed up function _get_mode_intro_message by 32%

[codeflash-ai]

📄 32% (0.32x) speedup for _get_mode_intro_message in marimo/_server/ai/prompts.py

⏱️ Runtime : 735 microseconds → 558 microseconds (best of 126 runs)

📝 Explanation and details

The optimization achieves a 31% speedup by eliminating redundant string construction and concatenation operations on each function call.

Key Changes:

• Pre-computed constants: Moved the complete intro messages to module-level constants _MANUAL_MODE_INTRO and _ASK_MODE_INTRO instead of constructing them dynamically
• Eliminated string operations: Removed the base_intro variable creation, f-string formatting (f"{base_intro}"), and multi-line string concatenation that occurred on every function call

Why This Speeds Up Performance:
The original code performed expensive string operations every time the function was called:

1. Created a base_intro string (806ns + 719ns per call from profiler)
2. Used f-string formatting to concatenate base_intro with additional content (480ns + 493ns per call)
3. Built multi-line strings through concatenation

The optimized version simply returns pre-built string constants, requiring only a conditional check and direct return - no string construction overhead.

Performance Impact:

• Single calls: 58-75% faster per individual call (534ns→311ns for manual mode)
• Batch operations: ~30-35% improvement for repeated calls (81.7μs→62.7μs for 500 calls)
• Memory efficiency: Reduces temporary string object creation and garbage collection pressure

Test Case Performance:
The optimization is particularly effective for:

• Repeated function calls (batch tests show consistent 30%+ improvements)
• Applications that frequently switch between modes
• Any scenario where this function is called in loops or hot paths

This optimization trades a small amount of module initialization time and memory for significantly faster runtime performance, making it ideal for functions called frequently during application execution.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Any, Literal

imports

import pytest
from marimo._server.ai.prompts import _get_mode_intro_message

Simulate CopilotMode typing for tests

CopilotMode = Literal["manual", "ask"]
from marimo._server.ai.prompts import _get_mode_intro_message

unit tests

------------------------

Basic Test Cases

------------------------

def test_manual_mode_message_basic():
"""Test that 'manual' mode returns the expected message."""
codeflash_output = _get_mode_intro_message("manual"); result = codeflash_output # 534ns -> 311ns (71.7% faster)

def test_ask_mode_message_basic():
"""Test that 'ask' mode returns the expected message."""
codeflash_output = _get_mode_intro_message("ask"); result = codeflash_output # 559ns -> 342ns (63.5% faster)

def test_manual_and_ask_modes_are_distinct():
"""Test that manual and ask modes produce different outputs."""
codeflash_output = _get_mode_intro_message("manual"); manual_msg = codeflash_output # 550ns -> 333ns (65.2% faster)
codeflash_output = _get_mode_intro_message("ask"); ask_msg = codeflash_output # 301ns -> 190ns (58.4% faster)

------------------------

Edge Test Cases

------------------------

def test_large_batch_of_manual_modes():
"""Test that multiple calls with 'manual' mode are consistent and efficient."""
codeflash_output = _get_mode_intro_message("manual"); expected = codeflash_output # 775ns -> 442ns (75.3% faster)
for _ in range(500): # reasonable batch size
codeflash_output = _get_mode_intro_message("manual") # 81.7μs -> 62.7μs (30.2% faster)

def test_large_batch_of_ask_modes():
"""Test that multiple calls with 'ask' mode are consistent and efficient."""
codeflash_output = _get_mode_intro_message("ask"); expected = codeflash_output # 629ns -> 383ns (64.2% faster)
for _ in range(500):
codeflash_output = _get_mode_intro_message("ask") # 88.2μs -> 67.7μs (30.3% faster)

def test_stress_with_alternating_modes():
"""Test alternating valid modes for consistency."""
codeflash_output = _get_mode_intro_message("manual"); manual_msg = codeflash_output # 704ns -> 439ns (60.4% faster)
codeflash_output = _get_mode_intro_message("ask"); ask_msg = codeflash_output # 364ns -> 228ns (59.6% faster)
for i in range(1000):
if i % 2 == 0:
codeflash_output = _get_mode_intro_message("manual")
else:
codeflash_output = _get_mode_intro_message("ask")

------------------------

Additional Robustness Tests

------------------------

def test_message_formatting_is_consistent():
"""Test that the message always starts with the base intro."""
for mode in ["manual", "ask"]:
codeflash_output = _get_mode_intro_message(mode); msg = codeflash_output # 1.01μs -> 526ns (92.0% faster)

def test_message_contains_sections():
"""Test that both modes contain 'Capabilities' and 'Limitations' sections."""
for mode in ["manual", "ask"]:
codeflash_output = _get_mode_intro_message(mode); msg = codeflash_output # 816ns -> 502ns (62.5% faster)

#------------------------------------------------
from typing import Literal

imports

import pytest # used for our unit tests
from marimo._server.ai.prompts import _get_mode_intro_message

Simulate the CopilotMode type used in the function

CopilotMode = Literal["manual", "ask"]
from marimo._server.ai.prompts import _get_mode_intro_message

unit tests

---- Basic Test Cases ----

def test_manual_mode_returns_expected_message():
# Test that "manual" mode returns the correct message
expected = (
"You are Marimo Copilot, an AI assistant integrated into the marimo notebook code editor.\n"
"Your primary function is to help users create, analyze, and improve data science notebooks using marimo's reactive programming model.\n"
"## Capabilities\n"
"- Answer questions and provide guidance using only your internal knowledge and the notebook context provided by the user.\n"
"\n"
"## Limitations\n"
"- You do NOT have access to any external tools, plugins, or APIs.\n"
"- You may not perform any actions beyond generating text and code suggestions.\n"
)
codeflash_output = _get_mode_intro_message("manual"); result = codeflash_output # 712ns -> 406ns (75.4% faster)

def test_ask_mode_returns_expected_message():
# Test that "ask" mode returns the correct message
expected = (
"You are Marimo Copilot, an AI assistant integrated into the marimo notebook code editor.\n"
"Your primary function is to help users create, analyze, and improve data science notebooks using marimo's reactive programming model.\n"
"## Capabilities\n"
"- You can use a set of read-only tools to gather additional context from the notebook or environment (e.g., searching code, summarizing data, or reading documentation).\n"
"- You may use these tools ONLY to gather information, not to modify code or state.\n"
"\n"
"## Limitations\n"
"- All tool use is strictly read-only. You may not perform write, edit, or execution actions.\n"
"- You must always explain to the user why you are using a tool before invoking it.\n"
)
codeflash_output = _get_mode_intro_message("ask"); result = codeflash_output # 652ns -> 375ns (73.9% faster)

def test_manual_and_ask_modes_are_distinct():
# Ensure the outputs for "manual" and "ask" are different
codeflash_output = _get_mode_intro_message("manual"); manual_msg = codeflash_output # 591ns -> 372ns (58.9% faster)
codeflash_output = _get_mode_intro_message("ask"); ask_msg = codeflash_output # 290ns -> 204ns (42.2% faster)

---- Edge Test Cases ----

def test_message_starts_with_base_intro():
# Ensure both modes start with the base_intro string
base_intro = (
"You are Marimo Copilot, an AI assistant integrated into the marimo notebook code editor.\n"
"Your primary function is to help users create, analyze, and improve data science notebooks using marimo's reactive programming model.\n"
)

def test_manual_mode_contains_expected_keywords():
# Check that specific keywords are present in manual mode
codeflash_output = _get_mode_intro_message("manual"); msg = codeflash_output # 697ns -> 421ns (65.6% faster)

def test_ask_mode_contains_expected_keywords():
# Check that specific keywords are present in ask mode
codeflash_output = _get_mode_intro_message("ask"); msg = codeflash_output # 635ns -> 361ns (75.9% faster)

---- Large Scale Test Cases ----

def test_large_batch_of_manual_modes():
# Test calling the function with "manual" mode 1000 times
codeflash_output = _get_mode_intro_message("manual"); expected = codeflash_output # 615ns -> 308ns (99.7% faster)
for _ in range(1000):
codeflash_output = _get_mode_intro_message("manual") # 163μs -> 123μs (32.2% faster)

def test_large_batch_of_ask_modes():
# Test calling the function with "ask" mode 1000 times
codeflash_output = _get_mode_intro_message("ask"); expected = codeflash_output # 632ns -> 383ns (65.0% faster)
for _ in range(1000):
codeflash_output = _get_mode_intro_message("ask") # 177μs -> 134μs (31.7% faster)

def test_all_modes_in_list():
# Test the function over a mixed list of valid and invalid modes
modes = ["manual", "ask", "Manual", "ASK", "", None, 123, "invalid"]
for mode in modes:
if mode in ("manual", "ask"):
# Should not raise
codeflash_output = _get_mode_intro_message(mode); msg = codeflash_output
else:
with pytest.raises(Exception):
_get_mode_intro_message(mode) # type: ignore

def test_message_length_consistency():
# Ensure that the message length for manual and ask is consistent across calls
manual_len = len(_get_mode_intro_message("manual")) # 714ns -> 407ns (75.4% faster)
ask_len = len(_get_mode_intro_message("ask")) # 366ns -> 222ns (64.9% faster)
for _ in range(100):
codeflash_output = len(_get_mode_intro_message("manual")) # 17.5μs -> 13.0μs (35.1% faster)
codeflash_output = len(_get_mode_intro_message("ask"))

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-_get_mode_intro_message-mhvjmtan and push.