⚡️ Speed up method ChunkProcessor._get_chunk_id by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for ChunkProcessor._get_chunk_id in litellm/litellm_core_utils/streaming_chunk_builder_utils.py

⏱️ Runtime : 286 microseconds → 260 microseconds (best of 68 runs)

📝 Explanation and details

The optimized code achieves a 9% speedup through two key optimizations:

1. Method Lookup Caching in _get_chunk_id:

• The original code calls chunk.get("id") in every loop iteration, which performs attribute lookup on each dictionary object
• The optimized version caches dict.get as a local variable chunk_get and calls chunk_get(chunk, "id"), avoiding repeated attribute lookups
• This optimization is most effective for large datasets - the test results show significant improvements (11-16% faster) when processing 1000+ chunks, while smaller datasets show slight overhead due to the setup cost

2. Added _sort_chunks Implementation:

• Provides efficient chunk sorting based on created_at timestamps when available
• Uses optimized attribute access patterns with local variables to minimize lookups in the sorting key function
• Falls back to returning unsorted chunks when no timestamp data is present

Performance Characteristics:

• Small datasets (1-10 chunks): Slight overhead (1-19% slower) due to local variable setup cost
• Large datasets (500-1000+ chunks): Significant gains (11-16% faster) where the method lookup optimization pays off
• Best case scenarios: When the target chunk with a valid ID is located later in the list, allowing the optimization to compound over many iterations

The optimization follows the classic Python performance pattern of moving attribute lookups outside tight loops, which becomes increasingly beneficial as the iteration count grows.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Any, Dict, List, Optional

# imports
import pytest  # used for our unit tests
from litellm.litellm_core_utils.streaming_chunk_builder_utils import \
    ChunkProcessor

# unit tests

# --- Basic Test Cases ---

def test_single_chunk_with_id():
    # Single chunk with a valid id
    chunks = [{"id": "abc123"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 753ns -> 822ns (8.39% slower)

def test_single_chunk_with_empty_id():
    # Single chunk with empty id
    chunks = [{"id": ""}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 655ns -> 739ns (11.4% slower)

def test_multiple_chunks_first_has_id():
    # Multiple chunks, first has id
    chunks = [{"id": "first"}, {"id": "second"}, {"id": "third"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 747ns -> 807ns (7.43% slower)

def test_multiple_chunks_first_empty_second_has_id():
    # Multiple chunks, first empty id, second has id
    chunks = [{"id": ""}, {"id": "second"}, {"id": "third"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 859ns -> 963ns (10.8% slower)

def test_multiple_chunks_all_empty_ids():
    # Multiple chunks, all ids empty
    chunks = [{"id": ""}, {"id": ""}, {"id": ""}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 657ns -> 796ns (17.5% slower)

def test_multiple_chunks_some_missing_id_key():
    # Some chunks missing 'id' key
    chunks = [{"id": ""}, {}, {"id": "present"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 947ns -> 1.01μs (6.52% slower)

def test_multiple_chunks_all_missing_id_key():
    # All chunks missing 'id' key
    chunks = [{}, {}, {}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 794ns -> 807ns (1.61% slower)

def test_empty_chunks_list():
    # Empty list of chunks
    chunks = []
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 482ns -> 560ns (13.9% slower)

# --- Edge Test Cases ---

def test_chunk_id_is_none():
    # id is None (should be treated as falsy)
    chunks = [{"id": None}, {"id": "not_none"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 787ns -> 889ns (11.5% slower)

def test_chunk_id_is_zero():
    # id is 0 (should be treated as falsy)
    chunks = [{"id": 0}, {"id": "nonzero"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 900ns -> 894ns (0.671% faster)

def test_chunk_id_is_false():
    # id is False (should be treated as falsy)
    chunks = [{"id": False}, {"id": "truthy"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 873ns -> 880ns (0.795% slower)

def test_chunk_id_is_non_string_type():
    # id is int, float, tuple, etc.
    chunks = [{"id": 123}, {"id": ""}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 736ns -> 824ns (10.7% slower)

    chunks = [{"id": (1, 2)}, {"id": "tuple"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 448ns -> 473ns (5.29% slower)

def test_chunk_id_is_list():
    # id is a non-empty list (truthy)
    chunks = [{"id": [1,2,3]}, {"id": "list"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 753ns -> 834ns (9.71% slower)

def test_chunk_id_is_dict():
    # id is a non-empty dict (truthy)
    chunks = [{"id": {"key": "value"}}, {"id": "dict"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 794ns -> 871ns (8.84% slower)

def test_chunk_id_is_empty_list():
    # id is an empty list (falsy)
    chunks = [{"id": []}, {"id": "nonempty"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 951ns -> 968ns (1.76% slower)

def test_chunk_id_is_empty_dict():
    # id is an empty dict (falsy)
    chunks = [{"id": {}}, {"id": "nonempty"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 970ns -> 998ns (2.81% slower)

def test_chunk_id_is_whitespace_string():
    # id is a whitespace string (truthy)
    chunks = [{"id": "   "}, {"id": "notspace"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 741ns -> 768ns (3.52% slower)

def test_chunk_id_is_falsey_but_not_empty():
    # id is "", None, 0, False, [], {}, but only non-empty is returned
    chunks = [{"id": ""}, {"id": None}, {"id": 0}, {"id": False}, {"id": []}, {"id": {}}, {"id": "valid"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 1.27μs -> 1.33μs (4.15% slower)

def test_chunk_id_is_boolean_true():
    # id is True (truthy)
    chunks = [{"id": True}, {"id": "nottrue"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 770ns -> 836ns (7.89% slower)

def test_chunk_id_is_boolean_false():
    # id is False (falsy)
    chunks = [{"id": False}, {"id": "notfalse"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 875ns -> 1.04μs (16.1% slower)

def test_chunk_id_is_zero_float():
    # id is 0.0 (falsy)
    chunks = [{"id": 0.0}, {"id": "notzero"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 929ns -> 973ns (4.52% slower)

def test_chunk_id_is_negative_number():
    # id is negative number (truthy)
    chunks = [{"id": -1}, {"id": "notneg"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 767ns -> 830ns (7.59% slower)

def test_chunk_id_is_empty_string_and_none():
    # id is "" and None, should skip both
    chunks = [{"id": ""}, {"id": None}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 780ns -> 816ns (4.41% slower)

# --- Large Scale Test Cases ---

def test_large_number_of_chunks_first_id_at_start():
    # Large list, id at start
    chunks = [{"id": "start"}] + [{"id": ""} for _ in range(999)]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 763ns -> 847ns (9.92% slower)

def test_large_number_of_chunks_id_at_end():
    # Large list, id at end
    chunks = [{"id": ""} for _ in range(999)] + [{"id": "end"}]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 30.3μs -> 26.9μs (12.7% faster)

def test_large_number_of_chunks_id_in_middle():
    # Large list, id in the middle
    chunks = [{"id": ""} for _ in range(500)] + [{"id": "middle"}] + [{"id": ""} for _ in range(499)]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 15.6μs -> 14.0μs (11.4% faster)

def test_large_number_of_chunks_all_empty():
    # Large list, all ids empty
    chunks = [{"id": ""} for _ in range(1000)]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 30.0μs -> 26.8μs (11.8% faster)

def test_large_number_of_chunks_some_missing_id_key():
    # Large list, some chunks missing 'id'
    chunks = [{"id": ""} for _ in range(500)] + [{} for _ in range(500)]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 29.2μs -> 25.2μs (15.6% faster)

def test_large_number_of_chunks_various_types():
    # Large list, various types for id
    chunks = [{"id": ""} for _ in range(250)] + [{"id": None} for _ in range(250)] + [{"id": 0} for _ in range(250)] + [{"id": "found"}] + [{"id": ""} for _ in range(249)]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 22.9μs -> 20.3μs (13.1% faster)

def test_large_number_of_chunks_first_truthy_non_string():
    # Large list, first truthy id is a non-string type
    chunks = [{"id": ""} for _ in range(500)] + [{"id": [1,2,3]}] + [{"id": ""} for _ in range(499)]
    codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 15.7μs -> 14.1μs (11.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 typing import Any, Dict, List, Optional

# imports
import pytest  # used for our unit tests
from litellm.litellm_core_utils.streaming_chunk_builder_utils import \
    ChunkProcessor

# unit tests

class TestGetChunkId:
    # 1. Basic Test Cases

    def test_single_chunk_with_id(self):
        # Single chunk with non-empty id
        chunks = [{"id": "abc123"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 764ns -> 829ns (7.84% slower)

    def test_single_chunk_with_empty_id(self):
        # Single chunk with empty id
        chunks = [{"id": ""}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 683ns -> 772ns (11.5% slower)

    def test_multiple_chunks_first_has_id(self):
        # First chunk has non-empty id
        chunks = [{"id": "xyz"}, {"id": "123"}, {"id": "456"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 813ns -> 852ns (4.58% slower)

    def test_multiple_chunks_first_empty_second_has_id(self):
        # First chunk has empty id, second has non-empty id
        chunks = [{"id": ""}, {"id": "second"}, {"id": "third"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 863ns -> 876ns (1.48% slower)

    def test_multiple_chunks_all_empty_id(self):
        # All chunks have empty id
        chunks = [{"id": ""}, {"id": ""}, {"id": ""}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 777ns -> 810ns (4.07% slower)

    def test_multiple_chunks_mixed_ids(self):
        # Mixed empty and non-empty ids
        chunks = [{"id": ""}, {"id": None}, {"id": "found"}, {"id": ""}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 953ns -> 984ns (3.15% slower)

    # 2. Edge Test Cases

    def test_empty_chunks_list(self):
        # Empty list of chunks
        chunks = []
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 480ns -> 598ns (19.7% slower)

    def test_chunk_missing_id_key(self):
        # Chunk missing 'id' key entirely
        chunks = [{"not_id": "value"}, {"id": "real_id"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 899ns -> 949ns (5.27% slower)

    def test_chunk_with_id_none(self):
        # Chunk with id as None
        chunks = [{"id": None}, {"id": "not_none"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 832ns -> 927ns (10.2% slower)

    def test_chunk_with_id_zero(self):
        # Chunk with id as integer zero (should be falsy, so skip)
        chunks = [{"id": 0}, {"id": "nonzero"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 911ns -> 961ns (5.20% slower)

    def test_chunk_with_id_false(self):
        # Chunk with id as boolean False (should be falsy, so skip)
        chunks = [{"id": False}, {"id": "truthy"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 871ns -> 900ns (3.22% slower)

    def test_chunk_with_id_true(self):
        # Chunk with id as boolean True (should be truthy, so return True)
        chunks = [{"id": True}, {"id": "other"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 722ns -> 882ns (18.1% slower)

    def test_chunk_with_id_empty_list(self):
        # Chunk with id as empty list (should be falsy, so skip)
        chunks = [{"id": []}, {"id": "list_id"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 908ns -> 974ns (6.78% slower)

    def test_chunk_with_id_nonempty_list(self):
        # Chunk with id as non-empty list (should be truthy, so return list)
        chunks = [{"id": [1, 2, 3]}, {"id": "other"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 758ns -> 842ns (9.98% slower)

    def test_chunk_with_id_empty_dict(self):
        # Chunk with id as empty dict (should be falsy, so skip)
        chunks = [{"id": {}}, {"id": "dict_id"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 905ns -> 969ns (6.60% slower)

    def test_chunk_with_id_nonempty_dict(self):
        # Chunk with id as non-empty dict (should be truthy, so return dict)
        chunks = [{"id": {"key": "value"}}, {"id": "other"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 748ns -> 837ns (10.6% slower)

    def test_chunk_with_id_zero_string(self):
        # Chunk with id as string "0" (should be truthy, so return "0")
        chunks = [{"id": "0"}, {"id": "other"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 746ns -> 819ns (8.91% slower)

    def test_chunk_with_id_whitespace_string(self):
        # Chunk with id as whitespace string (should be truthy, so return whitespace)
        chunks = [{"id": "   "}, {"id": "other"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 760ns -> 827ns (8.10% slower)

    def test_chunk_with_id_special_characters(self):
        # Chunk with id as special characters
        chunks = [{"id": "@#$_"}, {"id": "other"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 757ns -> 848ns (10.7% slower)

    # 3. Large Scale Test Cases

    def test_large_number_of_chunks_first_id(self):
        # Large list, first chunk has id
        chunks = [{"id": "first"}] + [{"id": ""} for _ in range(999)]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 790ns -> 828ns (4.59% slower)

    def test_large_number_of_chunks_last_id(self):
        # Large list, last chunk has id
        chunks = [{"id": ""} for _ in range(999)] + [{"id": "last"}]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 30.4μs -> 26.9μs (12.9% faster)

    def test_large_number_of_chunks_middle_id(self):
        # Large list, middle chunk has id
        chunks = [{"id": ""} for _ in range(499)] + [{"id": "middle"}] + [{"id": ""} for _ in range(500)]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 15.7μs -> 13.7μs (14.5% faster)

    def test_large_number_of_chunks_all_empty(self):
        # Large list, all chunks have empty id
        chunks = [{"id": ""} for _ in range(1000)]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 30.6μs -> 26.8μs (14.3% faster)

    def test_large_number_of_chunks_varied_types(self):
        # Large list, varied id types, only one is truthy and unique
        chunks = [{"id": ""}, {"id": None}, {"id": 0}, {"id": False}] * 250
        chunks.insert(500, {"id": "truthy_id"})
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 15.3μs -> 13.5μs (12.8% faster)

    def test_large_number_of_chunks_multiple_truthy(self):
        # Large list, multiple truthy ids, should return the first
        chunks = [{"id": ""} for _ in range(400)]
        chunks += [{"id": "first_truthy"}]
        chunks += [{"id": "second_truthy"} for _ in range(599)]
        codeflash_output = ChunkProcessor._get_chunk_id(chunks) # 12.8μs -> 11.3μs (13.6% faster)
# 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.streaming_chunk_builder_utils import ChunkProcessor

def test_ChunkProcessor__get_chunk_id():
    ChunkProcessor._get_chunk_id([{}, {'id': '\x00'}, {}, {}])

def test_ChunkProcessor__get_chunk_id_2():
    ChunkProcessor._get_chunk_id([])

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_ualnrfea/tmp7n2punt_/test_concolic_coverage.py::test_ChunkProcessor__get_chunk_id	1.06μs	1.14μs	-6.76%⚠️
codeflash_concolic_ualnrfea/tmp7n2punt_/test_concolic_coverage.py::test_ChunkProcessor__get_chunk_id_2	493ns	615ns	-19.8%⚠️

To edit these changes git checkout codeflash/optimize-ChunkProcessor._get_chunk_id-mh2p0w8d and push.