⚡️ Speed up method CohereRequestResponseResolver._resolve_chat_response by 18%

[codeflash-ai]

📄 18% (0.18x) speedup for CohereRequestResponseResolver._resolve_chat_response in wandb/integration/cohere/resolver.py

⏱️ Runtime : 85.5 microseconds → 72.8 microseconds (best of 91 runs)

📝 Explanation and details

The optimized code achieves a 17% speedup by eliminating function call overhead and reducing attribute lookups:

Key optimizations:

1. Inlined dictionary filtering: Instead of calling subset_dict(), the logic is embedded directly in _resolve_chat_response(), eliminating function call overhead and stack frame creation.

2. Cached attribute access: response.__dict__ is stored in a local variable response_dict and reused, avoiding repeated attribute lookups on the response object.

The line profiler shows the impact:

• Original: 564.7μs total time with 85.7% spent in the subset_dict function call
• Optimized: 494.3μs total time with the dictionary comprehension taking 85.5% of the reduced total

Performance characteristics by test case:

• Best gains (20-33%): Empty responses, missing keys, and objects with many irrelevant attributes benefit most from eliminating function overhead
• Consistent gains (11-24%): All typical use cases show meaningful improvement
• Scalable improvement: The 16.7% speedup on the 100-iteration test demonstrates the optimization scales well

This optimization is particularly effective for high-frequency logging scenarios where _resolve_chat_response is called repeatedly, as the eliminated function call overhead compounds across invocations.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Any, Dict, List, Sequence

# imports
import pytest
from wandb.integration.cohere.resolver import CohereRequestResponseResolver


class DummyResponse:
    """A dummy class to simulate the Response object with arbitrary attributes."""
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
from wandb.integration.cohere.resolver import CohereRequestResponseResolver

# --- Unit tests ---

@pytest.fixture
def resolver():
    """Fixture to provide a fresh instance of the resolver for each test."""
    return CohereRequestResponseResolver()

# 1. Basic Test Cases

def test_minimal_response(resolver):
    """Test with a minimal response containing only required keys."""
    resp = DummyResponse(
        response_id="resp1",
        generation_id="gen1",
        query="Hello?",
        text="Hi there!",
        conversation_id="conv1",
        prompt="Say hi",
        chatlog=[{"role": "user", "message": "Hello?"}],
        preamble="Start of conversation"
    )
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 2.31μs -> 1.90μs (21.5% faster)
    d = result[0]

def test_partial_response(resolver):
    """Test with a response missing some optional keys."""
    resp = DummyResponse(
        response_id="resp2",
        query="How are you?",
        text="I'm fine.",
        conversation_id="conv2"
    )
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 2.03μs -> 1.82μs (11.6% faster)
    d = result[0]

def test_extra_keys_ignored(resolver):
    """Test that extra keys in the response are ignored."""
    resp = DummyResponse(
        response_id="resp3",
        query="What's up?",
        text="All good.",
        conversation_id="conv3",
        extra_field1="should be ignored",
        extra_field2=123
    )
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 2.11μs -> 1.70μs (24.7% faster)
    d = result[0]

def test_empty_response(resolver):
    """Test with an empty response object."""
    resp = DummyResponse()
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 1.76μs -> 1.38μs (27.1% faster)
    d = result[0]

def test_non_string_values(resolver):
    """Test with non-string values for some fields."""
    resp = DummyResponse(
        response_id=12345,
        generation_id=None,
        query=["a", "b"],
        text={"content": "text"},
        conversation_id=False,
        prompt=0.0,
        chatlog=None,
        preamble=()
    )
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 2.32μs -> 1.92μs (20.5% faster)
    d = result[0]

# 2. Edge Test Cases

def test_all_keys_missing(resolver):
    """Test with a response object that has no relevant keys."""
    resp = DummyResponse(foo="bar", baz=123)
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 1.84μs -> 1.39μs (32.7% faster)
    d = result[0]

def test_some_keys_none(resolver):
    """Test with some keys explicitly set to None."""
    resp = DummyResponse(
        response_id=None,
        generation_id=None,
        query=None,
        text=None,
        conversation_id=None,
        prompt=None,
        chatlog=None,
        preamble=None
    )
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 2.26μs -> 1.87μs (20.6% faster)
    d = result[0]
    # All keys should be present with value None
    for key in [
        "response_id",
        "generation_id",
        "query",
        "text",
        "conversation_id",
        "prompt",
        "chatlog",
        "preamble",
    ]:
        pass

def test_unexpected_types_in_response(resolver):
    """Test with a response object that is not a DummyResponse but has __dict__."""
    class WeirdResponse:
        def __init__(self):
            self.response_id = [1, 2, 3]
            self.text = {"a": "b"}
    resp = WeirdResponse()
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 2.00μs -> 1.78μs (12.7% faster)
    d = result[0]

def test_response_with_nested_dicts(resolver):
    """Test with nested dicts and lists in values."""
    resp = DummyResponse(
        response_id="id",
        chatlog=[{"role": "user", "message": {"text": "Hi"}}],
        preamble={"instructions": ["do this", "do that"]}
    )
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 2.06μs -> 1.59μs (29.7% faster)
    d = result[0]

def test_response_with_unicode_and_special_characters(resolver):
    """Test with unicode, emoji, and special characters."""
    resp = DummyResponse(
        response_id="r\u2603",
        query="こんにちは",
        text="👋🌍",
        preamble="Résumé: naïve café"
    )
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 2.08μs -> 1.86μs (11.5% faster)
    d = result[0]

def test_response_object_with_slots(resolver):
    """Test with an object using __slots__, which does not have __dict__."""
    class SlotResponse:
        __slots__ = ("response_id", "text")
        def __init__(self):
            self.response_id = "slotid"
            self.text = "slottext"
    resp = SlotResponse()
    # Should raise AttributeError because __dict__ is missing
    with pytest.raises(AttributeError):
        resolver._resolve_chat_response(resp) # 1.29μs -> 1.24μs (3.86% faster)

# 3. Large Scale Test Cases

def test_large_response_object(resolver):
    """Test with a response object containing many irrelevant keys."""
    # 990 irrelevant keys, 8 relevant keys
    attrs = {f"irrelevant_{i}": i for i in range(990)}
    attrs.update({
        "response_id": "large1",
        "generation_id": "largegen",
        "query": "Big query?",
        "text": "Big answer.",
        "conversation_id": "bigconv",
        "prompt": "Big prompt",
        "chatlog": [{"role": "user", "message": "Big query?"}],
        "preamble": "Big preamble"
    })
    resp = DummyResponse(**attrs)
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 2.48μs -> 2.23μs (11.6% faster)
    d = result[0]

def test_large_chatlog(resolver):
    """Test with a large chatlog list."""
    large_chatlog = [{"role": "user", "message": f"msg {i}"} for i in range(1000)]
    resp = DummyResponse(
        response_id="r_largechat",
        chatlog=large_chatlog
    )
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 2.00μs -> 1.65μs (20.7% faster)
    d = result[0]

def test_many_response_objects(resolver):
    """Test calling the function many times in a row for scalability."""
    for i in range(100):
        resp = DummyResponse(
            response_id=f"id_{i}",
            text=f"text_{i}"
        )
        codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 56.7μs -> 48.6μs (16.7% faster)
        d = result[0]

def test_large_strings_in_fields(resolver):
    """Test with very large strings in each field."""
    bigstr = "x" * 10000
    resp = DummyResponse(
        response_id=bigstr,
        generation_id=bigstr,
        query=bigstr,
        text=bigstr,
        conversation_id=bigstr,
        prompt=bigstr,
        chatlog=[{"role": "user", "message": bigstr}],
        preamble=bigstr
    )
    codeflash_output = resolver._resolve_chat_response(resp); result = codeflash_output # 2.32μs -> 1.87μs (23.9% faster)
    d = result[0]
    for key in [
        "response_id", "generation_id", "query", "text",
        "conversation_id", "prompt", "preamble"
    ]:
        pass
# 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 wandb.integration.cohere.resolver import CohereRequestResponseResolver


def _resolve_chat_response(response: Any) -> List[Dict[str, Any]]:
    # The function expects an object with __dict__ containing relevant keys.
    return [
        subset_dict(
            response.__dict__,
            [
                "response_id",
                "generation_id",
                "query",
                "text",
                "conversation_id",
                "prompt",
                "chatlog",
                "preamble",
            ],
        )
    ]

# Helper class for simulating a Response object
class DummyResponse:
    def __init__(
        self,
        response_id=None,
        generation_id=None,
        query=None,
        text=None,
        conversation_id=None,
        prompt=None,
        chatlog=None,
        preamble=None,
        **extra
    ):
        self.response_id = response_id
        self.generation_id = generation_id
        self.query = query
        self.text = text
        self.conversation_id = conversation_id
        self.prompt = prompt
        self.chatlog = chatlog
        self.preamble = preamble
        # Allow for extra keys to simulate large/edge cases
        for k, v in extra.items():
            setattr(self, k, v)

# -------------------- UNIT TESTS --------------------

# 1. Basic Test Cases

To edit these changes git checkout codeflash/optimize-CohereRequestResponseResolver._resolve_chat_response-mhdg078w and push.