⚡️ Speed up function extract_between_tags by 8%

[codeflash-ai]

📄 8% (0.08x) speedup for extract_between_tags in litellm/litellm_core_utils/prompt_templates/factory.py

⏱️ Runtime : 1.68 milliseconds → 1.56 milliseconds (best of 175 runs)

📝 Explanation and details

The optimization introduces regex pattern caching using a module-level dictionary _tag_re_cache to store pre-compiled regex patterns for each unique tag.

Key changes:

• Pattern compilation caching: Instead of calling re.findall() which compiles the regex pattern on every function call, the optimized version compiles the pattern once using re.compile() and stores it in _tag_re_cache
• Cache lookup: Subsequent calls with the same tag reuse the cached compiled pattern via pattern.findall(string)

Why this leads to speedup:
Regex compilation is computationally expensive. The original code's re.findall(f"<{tag}>(.+?)</{tag}>", string, re.DOTALL) recompiles the same pattern every time, even for repeated tag names. By caching compiled patterns, we eliminate redundant compilation overhead.

Performance characteristics from test results:

• Small-scale operations: 50-200% speedup on basic test cases with simple tags and content
• Repeated tag usage: Maximum benefit when the same tag is used multiple times (common in real applications)
• Large-scale operations: Modest 1-8% improvements on large datasets, as the regex matching itself (not compilation) dominates runtime for very large strings
• Edge cases: Consistent improvements across all edge cases including empty strings, special characters, and unicode tags

The caching strategy is most effective for workloads with repeated tag extractions, which appears to be the common usage pattern based on the test suite.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import re
from typing import List

# imports
import pytest  # used for our unit tests
from litellm.litellm_core_utils.prompt_templates.factory import \
    extract_between_tags

# unit tests

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

def test_single_tag_basic():
    # Basic extraction of a single tag
    s = "<foo>Hello World</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 4.13μs -> 2.20μs (87.5% faster)

def test_multiple_tags_basic():
    # Multiple tags of same type
    s = "<foo>First</foo> and <foo>Second</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 3.80μs -> 2.16μs (75.9% faster)

def test_strip_option_basic():
    # Extraction with extra whitespace and strip=True
    s = "<foo>   padded   </foo><foo>\nsecond\n</foo>"
    codeflash_output = extract_between_tags("foo", s, strip=True); result = codeflash_output # 5.61μs -> 3.64μs (54.1% faster)

def test_different_tag_basic():
    # Extraction of a different tag
    s = "<bar>Bar1</bar><foo>Foo1</foo><bar>Bar2</bar>"
    codeflash_output = extract_between_tags("bar", s); result = codeflash_output # 3.97μs -> 2.23μs (77.5% faster)

def test_no_tags_basic():
    # String with no matching tags
    s = "No tags here"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 2.79μs -> 1.06μs (162% faster)

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

def test_empty_string_edge():
    # Empty string input
    s = ""
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 2.72μs -> 907ns (200% faster)

def test_empty_tag_content_edge():
    # Tag with empty content
    s = "<foo></foo><foo> </foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 3.67μs -> 1.97μs (86.7% faster)

def test_nested_tags_edge():
    # Nested tags (should not match across tags)
    s = "<foo>outer <foo>inner</foo> outer2</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 3.62μs -> 2.07μs (74.9% faster)

def test_overlapping_tags_edge():
    # Overlapping tags (should not match)
    s = "<foo>first<foo>second</foo>third</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 3.54μs -> 1.93μs (83.9% faster)

def test_tags_with_newlines_edge():
    # Content with newlines
    s = "<foo>\nline1\nline2\n</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 3.47μs -> 1.73μs (101% faster)

def test_tags_with_special_characters_edge():
    # Content with special characters
    s = "<foo>!@#$%^&*()</foo><foo>中文</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 4.83μs -> 3.19μs (51.3% faster)

def test_tags_with_numbers_edge():
    # Tag name is numeric
    s = "<123>abc</123><123>def</123>"
    codeflash_output = extract_between_tags("123", s); result = codeflash_output # 3.85μs -> 2.20μs (75.2% faster)

def test_strip_false_with_whitespace_edge():
    # strip=False should leave whitespace
    s = "<foo>   padded   </foo>"
    codeflash_output = extract_between_tags("foo", s, strip=False); result = codeflash_output # 4.05μs -> 2.19μs (85.3% faster)

def test_strip_true_with_only_whitespace_edge():
    # strip=True with only whitespace content
    s = "<foo>   </foo>"
    codeflash_output = extract_between_tags("foo", s, strip=True); result = codeflash_output # 4.80μs -> 2.81μs (70.6% faster)

def test_tag_case_sensitivity_edge():
    # Tag name is case sensitive
    s = "<Foo>upper</Foo><foo>lower</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 3.53μs -> 1.88μs (87.5% faster)

def test_tag_with_attributes_edge():
    # Tag with attributes should not match (since regex is strict)
    s = "<foo bar='baz'>should not match</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 2.73μs -> 1.03μs (165% faster)

def test_tag_with_different_closing_edge():
    # Tag with mismatched closing tag
    s = "<foo>hello</bar>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 3.54μs -> 1.93μs (83.2% faster)

def test_tag_at_string_edges_edge():
    # Tag at very start and end of string
    s = "<foo>start</foo>middle<foo>end</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 3.78μs -> 2.06μs (83.6% faster)

def test_tag_with_empty_tag_name_edge():
    # Empty tag name
    s = "<>no_tag</>"
    codeflash_output = extract_between_tags("", s); result = codeflash_output # 3.57μs -> 1.71μs (109% faster)

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

def test_large_number_of_tags_large_scale():
    # Many tags in a large string
    s = "".join([f"<foo>{i}</foo>" for i in range(1000)])
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 106μs -> 105μs (1.41% faster)

def test_large_content_in_tag_large_scale():
    # Very large content inside a tag
    large_content = "A" * 10000
    s = f"<foo>{large_content}</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 84.8μs -> 83.9μs (1.03% faster)

def test_large_mixed_tags_large_scale():
    # Large string with mixed tags, only some match
    s = "".join([f"<foo>{i}</foo><bar>{i}</bar>" for i in range(500)])
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 62.6μs -> 59.4μs (5.40% faster)
    codeflash_output = extract_between_tags("bar", s); result_bar = codeflash_output # 59.6μs -> 58.3μs (2.17% faster)

def test_large_strip_large_scale():
    # Large number of tags with whitespace, strip=True
    s = "".join([f"<foo>   {i}   </foo>" for i in range(1000)])
    codeflash_output = extract_between_tags("foo", s, strip=True); result = codeflash_output # 237μs -> 222μs (7.00% faster)

def test_large_tags_with_newlines_large_scale():
    # Large number of tags with newlines
    s = "".join([f"<foo>\n{i}\n</foo>" for i in range(1000)])
    codeflash_output = extract_between_tags("foo", s, strip=True); result = codeflash_output # 169μs -> 162μs (4.25% faster)

# ===========================
# ADDITIONAL EDGE CASES
# ===========================

def test_tag_with_regex_special_characters_edge():
    # Tag name contains regex special characters
    s = "<f.o+o>abc</f.o+o>"
    codeflash_output = extract_between_tags("f.o+o", s); result = codeflash_output # 3.71μs -> 1.72μs (116% faster)

def test_tag_with_unicode_edge():
    # Tag name is unicode
    s = "<фу>привет</фу>"
    codeflash_output = extract_between_tags("фу", s); result = codeflash_output # 5.09μs -> 2.65μs (91.8% faster)

def test_tag_with_partial_match_edge():
    # Partial tag name match should not match
    s = "<foobar>abc</foobar>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 2.92μs -> 1.08μs (172% faster)

def test_tag_with_multiple_lines_edge():
    # Tag content spans multiple lines
    s = "<foo>line1\nline2\nline3</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 3.65μs -> 1.98μs (84.4% faster)

def test_tag_with_interleaved_tags_edge():
    # Interleaved tags of different types
    s = "<foo>foo1</foo><bar>bar1</bar><foo>foo2</foo><bar>bar2</bar>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 3.90μs -> 2.27μs (71.9% faster)
    codeflash_output = extract_between_tags("bar", s); result_bar = codeflash_output # 1.95μs -> 1.19μs (63.4% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import re
from typing import List

# imports
import pytest  # used for our unit tests
from litellm.litellm_core_utils.prompt_templates.factory import \
    extract_between_tags

# unit tests

# -------------------------------
# Basic Test Cases
# -------------------------------

def test_single_occurrence_basic():
    # Test extracting a single tag occurrence
    s = "<foo>Hello World</foo>"
    codeflash_output = extract_between_tags("foo", s) # 3.60μs -> 1.90μs (89.9% faster)

def test_multiple_occurrences_basic():
    # Test extracting multiple tag occurrences
    s = "<foo>One</foo> and <foo>Two</foo>"
    codeflash_output = extract_between_tags("foo", s) # 3.73μs -> 1.94μs (92.1% faster)

def test_strip_basic():
    # Test with strip=True to remove leading/trailing whitespace
    s = "<foo>  One  </foo> <foo>\nTwo\t</foo>"
    codeflash_output = extract_between_tags("foo", s, strip=True) # 5.30μs -> 3.54μs (49.5% faster)

def test_no_occurrence_basic():
    # Test when tag does not exist
    s = "<bar>Test</bar>"
    codeflash_output = extract_between_tags("foo", s) # 2.71μs -> 1.04μs (160% faster)

def test_different_tags_basic():
    # Test extracting a different tag
    s = "<foo>Test1</foo><bar>Test2</bar>"
    codeflash_output = extract_between_tags("bar", s) # 3.36μs -> 1.94μs (73.1% faster)

# -------------------------------
# Edge Test Cases
# -------------------------------

def test_empty_string_edge():
    # Test with empty input string
    codeflash_output = extract_between_tags("foo", "") # 2.49μs -> 915ns (172% faster)

def test_empty_tag_edge():
    # Test with empty tag name
    s = "<>Empty tag</>"
    codeflash_output = extract_between_tags("", s) # 3.70μs -> 1.85μs (100% faster)

def test_nested_tags_edge():
    # Test nested tags (should not extract inner tags)
    s = "<foo>Outer <foo>Inner</foo> OuterEnd</foo>"
    # Only matches the first closing tag for each opening tag
    codeflash_output = extract_between_tags("foo", s) # 3.72μs -> 1.87μs (98.8% faster)

def test_overlapping_tags_edge():
    # Test overlapping tags (not valid XML, but regex still matches)
    s = "<foo>First</foo><foo>Second</foo><foo>Third</foo>"
    codeflash_output = extract_between_tags("foo", s) # 4.24μs -> 2.33μs (81.9% faster)

def test_tag_with_special_characters_edge():
    # Test tag with special regex characters
    s = "<f.o+o>Special</f.o+o>"
    codeflash_output = extract_between_tags("f.o+o", s) # 3.50μs -> 1.76μs (98.8% faster)

def test_tag_with_newlines_edge():
    # Test content with newlines
    s = "<foo>\nHello\nWorld\n</foo>"
    codeflash_output = extract_between_tags("foo", s) # 3.59μs -> 1.88μs (90.8% faster)

def test_tag_with_empty_content_edge():
    # Test tag with empty content
    s = "<foo></foo>"
    codeflash_output = extract_between_tags("foo", s) # 2.47μs -> 919ns (169% faster)

def test_tag_with_whitespace_content_edge():
    # Test tag with only whitespace content
    s = "<foo>   </foo>"
    codeflash_output = extract_between_tags("foo", s) # 3.40μs -> 1.71μs (99.6% faster)
    codeflash_output = extract_between_tags("foo", s, strip=True) # 2.96μs -> 2.10μs (40.8% faster)

def test_tag_with_attributes_edge():
    # Test tags with attributes (should not match, since regex expects exact tag)
    s = '<foo id="1">Test</foo>'
    codeflash_output = extract_between_tags("foo", s) # 2.65μs -> 933ns (183% faster)

def test_partial_tag_edge():
    # Test incomplete tag (should not match)
    s = "<foo>Test"
    codeflash_output = extract_between_tags("foo", s) # 2.42μs -> 891ns (171% faster)

def test_multiple_different_tags_edge():
    # Test multiple different tags in sequence
    s = "<foo>One</foo><bar>Two</bar><foo>Three</foo>"
    codeflash_output = extract_between_tags("foo", s) # 3.77μs -> 2.10μs (79.5% faster)
    codeflash_output = extract_between_tags("bar", s) # 1.69μs -> 982ns (71.8% faster)

def test_tag_case_sensitivity_edge():
    # Test case sensitivity (should be case-sensitive)
    s = "<Foo>Uppercase</Foo><foo>Lowercase</foo>"
    codeflash_output = extract_between_tags("foo", s) # 3.41μs -> 1.84μs (85.6% faster)
    codeflash_output = extract_between_tags("Foo", s) # 1.96μs -> 1.03μs (90.3% faster)

def test_tag_with_numbers_edge():
    # Test tag with numbers
    s = "<foo1>Test1</foo1><foo2>Test2</foo2>"
    codeflash_output = extract_between_tags("foo1", s) # 3.58μs -> 1.75μs (104% faster)
    codeflash_output = extract_between_tags("foo2", s) # 1.76μs -> 957ns (84.2% faster)

def test_strip_with_newlines_and_tabs_edge():
    # Test strip with newlines and tabs
    s = "<foo>\n\tTest\t\n</foo>"
    codeflash_output = extract_between_tags("foo", s, strip=True) # 4.74μs -> 2.95μs (60.6% faster)

def test_strip_false_with_whitespace_edge():
    # Test strip=False with whitespace
    s = "<foo>   spaced   </foo>"
    codeflash_output = extract_between_tags("foo", s, strip=False) # 3.63μs -> 1.96μs (85.2% faster)

def test_strip_true_with_whitespace_edge():
    # Test strip=True with whitespace
    s = "<foo>   spaced   </foo>"
    codeflash_output = extract_between_tags("foo", s, strip=True) # 4.71μs -> 2.90μs (62.4% faster)

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

def test_large_number_of_tags_large_scale():
    # Test with a large number of tags (1000)
    s = "".join(f"<foo>{i}</foo>" for i in range(1000))
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 109μs -> 106μs (3.28% faster)

def test_large_content_large_scale():
    # Test with a single tag containing a large amount of content
    large_content = "A" * 10000
    s = f"<foo>{large_content}</foo>"
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 83.7μs -> 81.9μs (2.29% faster)

def test_large_mixed_tags_large_scale():
    # Test with 500 <foo> and 500 <bar> tags mixed
    s = "".join(f"<foo>{i}</foo><bar>{i}</bar>" for i in range(500))
    codeflash_output = extract_between_tags("foo", s); foo_result = codeflash_output # 62.4μs -> 57.4μs (8.74% faster)
    codeflash_output = extract_between_tags("bar", s); bar_result = codeflash_output # 56.3μs -> 56.7μs (0.865% slower)

def test_large_strip_large_scale():
    # Test strip=True with large number of tags and whitespace
    s = "".join(f"<foo>   {i}   </foo>" for i in range(1000))
    codeflash_output = extract_between_tags("foo", s, strip=True); result = codeflash_output # 237μs -> 235μs (0.647% faster)

def test_large_empty_tags_large_scale():
    # Test with a large number of empty tags
    s = "".join("<foo></foo>" for _ in range(1000))
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 100μs -> 99.1μs (1.23% faster)

def test_large_tags_with_newlines_large_scale():
    # Test with large number of tags containing newlines
    s = "".join(f"<foo>\n{i}\n</foo>" for i in range(1000))
    codeflash_output = extract_between_tags("foo", s); result = codeflash_output # 126μs -> 131μs (3.51% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from litellm.litellm_core_utils.prompt_templates.factory import extract_between_tags

def test_extract_between_tags():
    extract_between_tags('', '', strip=True)

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic__vli01p5/tmpjc3b71zz/test_concolic_coverage.py::test_extract_between_tags	3.88μs	1.80μs	116%✅

To edit these changes git checkout codeflash/optimize-extract_between_tags-mhbopf4k and push.