⚡️ Speed up function _make_span_response_attributes by 340%

[codeflash-ai]

📄 340% (3.40x) speedup for _make_span_response_attributes in src/xai_sdk/image.py

⏱️ Runtime : 8.43 milliseconds → 1.91 milliseconds (best of 190 runs)

📝 Explanation and details

The optimized code achieves a 340% speedup by eliminating redundant operations inside the loop and reducing the number of conditional branches executed.

Key optimizations:

1. Precomputed format string: The expensive image_pb2.ImageFormat.Name(request.format).removeprefix("IMG_FORMAT_").lower() operation is moved outside the loop and computed once, rather than being called repeatedly during dictionary initialization.

2. Hoisted conditional logic: Instead of checking request.format on every loop iteration (which was the biggest bottleneck at 48% + 24.3% = 72.3% of original runtime), the format check is moved outside the loop. This creates separate optimized loops for each format type.

3. Cached attribute access: request.format is stored in a local variable req_format to avoid repeated attribute lookups.

Performance impact by test case:

• Large-scale scenarios (1000+ responses) see the biggest gains: 493% and 257-489% speedups, as the loop optimizations compound significantly
• Multiple response cases show solid improvements: 18.6% and 16.5% faster
• Single response cases have modest gains: 1-4% improvements
• Empty response cases are slightly slower due to the additional upfront setup, but this is negligible in practice

The optimization is most effective for workloads with many responses, where eliminating the repeated conditional checks and format string computation provides substantial performance benefits.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from dataclasses import dataclass
from typing import Any, Literal, Sequence

# imports
import pytest
from xai_sdk.image import _make_span_response_attributes

# --- Mocks for image_pb2 and BaseImageResponse ---

class ImageFormatEnum:
    IMG_FORMAT_URL = 1
    IMG_FORMAT_BASE64 = 2

    _names = {
        1: "IMG_FORMAT_URL",
        2: "IMG_FORMAT_BASE64"
    }

    @classmethod
    def Name(cls, value):
        return cls._names.get(value, "UNKNOWN")

class image_pb2:
    ImageFormat = ImageFormatEnum

# Minimal BaseImageResponse for testing
@dataclass
class BaseImageResponse:
    prompt: str
    url: str = ""
    base64: str = ""
from xai_sdk.image import _make_span_response_attributes


# --- Helper for request mock ---
@dataclass
class MockRequest:
    model: str
    format: int

# --- Unit Tests ---

# 1. Basic Test Cases

def test_basic_url_format_single_response():
    # Test single response, URL format
    request = MockRequest(model="test-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    response = BaseImageResponse(prompt="a cat", url="http://img.com/cat.jpg")
    codeflash_output = _make_span_response_attributes(request, [response]); result = codeflash_output # 13.2μs -> 12.6μs (4.18% faster)

def test_basic_base64_format_single_response():
    # Test single response, BASE64 format
    request = MockRequest(model="test-model", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    response = BaseImageResponse(prompt="a dog", base64="b64string")
    codeflash_output = _make_span_response_attributes(request, [response]); result = codeflash_output # 7.53μs -> 7.69μs (2.09% slower)

def test_basic_multiple_responses_url_format():
    # Test multiple responses, URL format
    request = MockRequest(model="multi-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    responses = [
        BaseImageResponse(prompt="p1", url="u1"),
        BaseImageResponse(prompt="p2", url="u2"),
    ]
    codeflash_output = _make_span_response_attributes(request, responses); result = codeflash_output # 10.9μs -> 9.18μs (18.6% faster)

def test_basic_multiple_responses_base64_format():
    # Test multiple responses, BASE64 format
    request = MockRequest(model="multi-model", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    responses = [
        BaseImageResponse(prompt="p1", base64="b1"),
        BaseImageResponse(prompt="p2", base64="b2"),
    ]
    codeflash_output = _make_span_response_attributes(request, responses); result = codeflash_output # 8.25μs -> 7.56μs (9.18% faster)

# 2. Edge Test Cases

def test_empty_responses_list():
    # Test with no responses
    request = MockRequest(model="empty-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    responses = []
    codeflash_output = _make_span_response_attributes(request, responses); result = codeflash_output # 3.65μs -> 7.27μs (49.9% slower)

def test_empty_prompt_and_url_base64():
    # Test with empty prompt, url, base64
    request = MockRequest(model="edge-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    response = BaseImageResponse(prompt="", url="")
    codeflash_output = _make_span_response_attributes(request, [response]); result = codeflash_output # 8.02μs -> 7.88μs (1.76% faster)

    # Now test for BASE64 format
    request = MockRequest(model="edge-model", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    response = BaseImageResponse(prompt="", base64="")
    codeflash_output = _make_span_response_attributes(request, [response]); result = codeflash_output # 3.14μs -> 3.11μs (1.09% faster)


def test_special_characters_in_prompt_and_url():
    # Test with special characters
    request = MockRequest(model="specialchars", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    response = BaseImageResponse(prompt="!@#$%^&*()", url="http://test.com/!@#$%^&*().jpg")
    codeflash_output = _make_span_response_attributes(request, [response]); result = codeflash_output # 12.8μs -> 12.5μs (1.82% faster)

def test_non_ascii_characters_in_prompt_and_base64():
    # Test with non-ASCII characters
    request = MockRequest(model="unicode", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    response = BaseImageResponse(prompt="猫", base64="猫の画像")
    codeflash_output = _make_span_response_attributes(request, [response]); result = codeflash_output # 7.55μs -> 7.42μs (1.70% faster)

# 3. Large Scale Test Cases

def test_large_number_of_responses_url_format():
    # Test with a large number of responses (up to 1000)
    request = MockRequest(model="large-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    responses = [BaseImageResponse(prompt=f"prompt{i}", url=f"url{i}") for i in range(1000)]
    codeflash_output = _make_span_response_attributes(request, responses); result = codeflash_output # 1.74ms -> 292μs (493% faster)
    # Check a few random indices for correctness
    for i in [0, 499, 999]:
        pass

def test_large_number_of_responses_base64_format():
    # Test with a large number of responses (up to 1000)
    request = MockRequest(model="large-model", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    responses = [BaseImageResponse(prompt=f"prompt{i}", base64=f"b64_{i}") for i in range(1000)]
    codeflash_output = _make_span_response_attributes(request, responses); result = codeflash_output # 1.02ms -> 286μs (257% faster)
    # Check a few random indices for correctness
    for i in [0, 499, 999]:
        pass

def test_large_prompts_and_urls():
    # Test with very large strings in prompt and url/base64
    large_str = "x" * 1000
    request = MockRequest(model="bigstr", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    response = BaseImageResponse(prompt=large_str, url=large_str)
    codeflash_output = _make_span_response_attributes(request, [response]); result = codeflash_output # 9.41μs -> 10.3μs (8.32% slower)

    request = MockRequest(model="bigstr", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    response = BaseImageResponse(prompt=large_str, base64=large_str)
    codeflash_output = _make_span_response_attributes(request, [response]); result = codeflash_output # 3.23μs -> 3.21μs (0.622% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from dataclasses import dataclass
from typing import Any, Literal, Sequence

# imports
import pytest  # used for our unit tests
from xai_sdk.image import _make_span_response_attributes

# --- Mocks and Stubs for image_pb2 and BaseImageResponse ---

# Simulate image_pb2.ImageFormat Enum and Name function
class ImageFormatEnum:
    IMG_FORMAT_BASE64 = 1
    IMG_FORMAT_URL = 2

    _NAMES = {
        1: "IMG_FORMAT_BASE64",
        2: "IMG_FORMAT_URL"
    }

    @classmethod
    def Name(cls, value):
        return cls._NAMES.get(value, "UNKNOWN")

# Simulate image_pb2 module
class image_pb2:
    ImageFormat = ImageFormatEnum

# Simulate BaseImageResponse
@dataclass
class BaseImageResponse:
    prompt: str
    url: str = ""
    base64: str = ""
from xai_sdk.image import _make_span_response_attributes


# --- Helper class for request ---
@dataclass
class DummyRequest:
    model: str
    format: int

# --- Unit Tests ---

# 1. Basic Test Cases

def test_basic_url_format_single_response():
    # Test single response, format URL
    request = DummyRequest(model="test-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    responses = [BaseImageResponse(prompt="cat", url="http://img.com/cat.jpg")]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 6.16μs -> 6.38μs (3.46% slower)

def test_basic_base64_format_single_response():
    # Test single response, format BASE64
    request = DummyRequest(model="test-model", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    responses = [BaseImageResponse(prompt="dog", base64="b64dog")]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 7.51μs -> 7.63μs (1.68% slower)

def test_basic_multiple_responses_url():
    # Test multiple responses, format URL
    request = DummyRequest(model="multi-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    responses = [
        BaseImageResponse(prompt="cat", url="http://img.com/cat.jpg"),
        BaseImageResponse(prompt="dog", url="http://img.com/dog.jpg"),
    ]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 8.04μs -> 7.31μs (9.94% faster)

def test_basic_multiple_responses_base64():
    # Test multiple responses, format BASE64
    request = DummyRequest(model="multi-model", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    responses = [
        BaseImageResponse(prompt="cat", base64="b64cat"),
        BaseImageResponse(prompt="dog", base64="b64dog"),
    ]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 9.81μs -> 8.43μs (16.5% faster)

# 2. Edge Test Cases

def test_edge_empty_responses():
    # Test with empty responses list
    request = DummyRequest(model="empty-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    responses = []
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 3.36μs -> 5.58μs (39.9% slower)

def test_edge_empty_prompt_and_url_base64():
    # Test response with empty prompt and base64
    request = DummyRequest(model="edge-model", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    responses = [BaseImageResponse(prompt="", base64="")]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 7.87μs -> 7.71μs (2.18% faster)

def test_edge_empty_prompt_and_url_url():
    # Test response with empty prompt and url
    request = DummyRequest(model="edge-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    responses = [BaseImageResponse(prompt="", url="")]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 6.59μs -> 6.38μs (3.26% faster)


def test_edge_mixed_response_fields():
    # Test responses with both url and base64 present, but format is URL
    request = DummyRequest(model="mixed-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    responses = [BaseImageResponse(prompt="hybrid", url="http://img.com/hybrid.jpg", base64="b64hybrid")]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 11.6μs -> 11.0μs (6.14% faster)

def test_edge_mixed_response_fields_base64():
    # Test responses with both url and base64 present, but format is BASE64
    request = DummyRequest(model="mixed-model", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    responses = [BaseImageResponse(prompt="hybrid", url="http://img.com/hybrid.jpg", base64="b64hybrid")]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 9.19μs -> 9.04μs (1.60% faster)

def test_edge_non_ascii_prompt():
    # Test non-ASCII prompt
    request = DummyRequest(model="unicode-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    responses = [BaseImageResponse(prompt="猫", url="http://img.com/cat.jpg")]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 7.01μs -> 7.24μs (3.15% slower)

def test_edge_long_prompt():
    # Test very long prompt string
    long_prompt = "a" * 1000
    request = DummyRequest(model="long-model", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    responses = [BaseImageResponse(prompt=long_prompt, base64="longb64")]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 8.09μs -> 8.24μs (1.81% slower)

# 3. Large Scale Test Cases

def test_large_scale_many_responses_url():
    # Test with many responses (up to 1000), format URL
    request = DummyRequest(model="large-model", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    responses = [
        BaseImageResponse(prompt=f"prompt_{i}", url=f"http://img.com/img_{i}.jpg")
        for i in range(1000)
    ]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 1.03ms -> 289μs (256% faster)
    # Check a few random indices
    for i in [0, 499, 999]:
        pass

def test_large_scale_many_responses_base64():
    # Test with many responses (up to 1000), format BASE64
    request = DummyRequest(model="large-model", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    responses = [
        BaseImageResponse(prompt=f"prompt_{i}", base64=f"b64img_{i}")
        for i in range(1000)
    ]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 1.73ms -> 294μs (486% faster)
    for i in [0, 500, 999]:
        pass

def test_large_scale_empty_prompts():
    # Test with many responses, all empty prompts
    request = DummyRequest(model="large-empty", format=image_pb2.ImageFormat.IMG_FORMAT_URL)
    responses = [
        BaseImageResponse(prompt="", url=f"http://img.com/img_{i}.jpg")
        for i in range(1000)
    ]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 1.02ms -> 284μs (258% faster)
    for i in [0, 500, 999]:
        pass

def test_large_scale_mixed_formats():
    # Test with mixed formats, but function should only use request.format
    # Provide responses with both url and base64, but only base64 should be present
    request = DummyRequest(model="large-mixed", format=image_pb2.ImageFormat.IMG_FORMAT_BASE64)
    responses = [
        BaseImageResponse(prompt=f"prompt_{i}", url=f"http://img.com/img_{i}.jpg", base64=f"b64img_{i}")
        for i in range(1000)
    ]
    codeflash_output = _make_span_response_attributes(request, responses); attributes = codeflash_output # 1.73ms -> 293μs (489% faster)
    for i in [0, 500, 999]:
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from xai.api.v1.image_pb2 import GenerateImageRequest
from xai_sdk.image import _make_span_response_attributes

def test__make_span_response_attributes():
    _make_span_response_attributes(GenerateImageRequest(), ())

To edit these changes git checkout codeflash/optimize-_make_span_response_attributes-mgu0bgav and push.