⚡️ Speed up method AsyncRawModelsClient.get by 21%

[codeflash-ai]

📄 21% (0.21x) speedup for AsyncRawModelsClient.get in src/cohere/models/raw_client.py

⏱️ Runtime : 10.9 seconds → 9.02 seconds (best of 16 runs)

📝 Explanation and details

The optimized code achieves a 21% runtime speedup and 14.3% throughput improvement through strategic early short-circuiting optimizations in two critical data processing functions:

Key Optimizations:

1. jsonable_encoder Fast Path: Added an immediate isinstance(obj, (str, int, float, type(None))) check at the beginning to bypass all other processing for the most common primitive types. This eliminates expensive checks for Pydantic models, dataclasses, and other complex types when dealing with simple data.

2. construct_type Primitive Fast Track: Introduced early handling of basic types (str, int, float, bool, type(None)) right after the None check, avoiding costly generic type introspection via get_origin() and get_args() for simple cases.

Why This Works:

• Hot Path Optimization: The line profiler shows that jsonable_encoder had significant time spent on isinstance checks for complex types when most objects are likely simple primitives (strings, numbers, None).
• Reduced Type Introspection Overhead: construct_type was spending considerable time in get_origin() and get_args() calls even for primitive types that don't need complex type analysis.
• Early Returns: Both optimizations follow the principle of handling the most common cases first and immediately returning, avoiding the execution of subsequent expensive code paths.

Test Case Performance:
The optimizations excel particularly well in high-throughput scenarios (concurrent requests, medium/high load tests) where these functions are called repeatedly during JSON processing and type construction. The concurrent test cases benefit most because they amplify the per-operation savings across many simultaneous requests.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import asyncio

import pytest
from cohere.models.raw_client import AsyncRawModelsClient


# Mocks and helpers for testing
class MockResponse:
    def __init__(self, status_code, json_data=None, headers=None, text=None):
        self.status_code = status_code
        self._json_data = json_data or {}
        self.headers = headers or {}
        self.text = text or ""

    def json(self):
        if self._json_data == "RAISE_JSON":
            raise Exception("JSONDecodeError")
        return self._json_data

class MockAsyncClient:
    def __init__(self, responses):
        self._responses = responses
        self._call_count = 0

    async def request(self, *args, **kwargs):
        # Return next response in sequence
        if self._call_count < len(self._responses):
            resp = self._responses[self._call_count]
            self._call_count += 1
            return resp
        # If out of responses, return last one
        return self._responses[-1]

# Minimal stubs for required types
class GetModelResponse:
    def __init__(self, name=None, version=None, details=None):
        self.name = name
        self.version = version
        self.details = details

class AsyncHttpResponse:
    def __init__(self, response, data):
        self.response = response
        self.data = data

class ApiError(Exception):
    def __init__(self, status_code, headers, body):
        self.status_code = status_code
        self.headers = headers
        self.body = body

class BadRequestError(ApiError): pass
class UnauthorizedError(ApiError): pass
class ForbiddenError(ApiError): pass
class NotFoundError(ApiError): pass
class UnprocessableEntityError(ApiError): pass
class TooManyRequestsError(ApiError): pass
class InvalidTokenError(ApiError): pass
class ClientClosedRequestError(ApiError): pass
class InternalServerError(ApiError): pass
class NotImplementedError(ApiError): pass
class ServiceUnavailableError(ApiError): pass
class GatewayTimeoutError(ApiError): pass
from cohere.models.raw_client import AsyncRawModelsClient


class AsyncClientWrapper:
    def __init__(self, httpx_client):
        self.httpx_client = httpx_client

# ---- UNIT TESTS ----

# Basic Test Cases

@pytest.mark.asyncio
async def test_get_success_basic():
    """Test basic successful response (200 OK)."""
    mock_resp = MockResponse(200, json_data={"name": "test-model", "version": "1.0", "details": "details"})
    client = AsyncRawModelsClient(client_wrapper=AsyncClientWrapper(MockAsyncClient([mock_resp])))
    result = await client.get("test-model")

@pytest.mark.asyncio
async def test_get_success_with_different_model_name():
    """Test successful response with a different model name."""
    mock_resp = MockResponse(200, json_data={"name": "other-model", "version": "2.1", "details": "other details"})
    client = AsyncRawModelsClient(client_wrapper=AsyncClientWrapper(MockAsyncClient([mock_resp])))
    result = await client.get("other-model")

@pytest.mark.asyncio
async def test_get_success_empty_details():
    """Test successful response with empty details."""
    mock_resp = MockResponse(200, json_data={"name": "empty-model", "version": "0.0", "details": ""})
    client = AsyncRawModelsClient(client_wrapper=AsyncClientWrapper(MockAsyncClient([mock_resp])))
    result = await client.get("empty-model")

# Edge Test Cases

@pytest.mark.asyncio










async def test_get_concurrent_success():
    """Test multiple concurrent successful calls."""
    responses = [
        MockResponse(200, json_data={"name": f"model{i}", "version": f"{i}.0", "details": f"details{i}"})
        for i in range(10)
    ]
    client = AsyncRawModelsClient(client_wrapper=AsyncClientWrapper(MockAsyncClient(responses)))
    # Run 10 concurrent gets
    tasks = [client.get(f"model{i}") for i in range(10)]
    results = await asyncio.gather(*tasks)
    for i, result in enumerate(results):
        pass

@pytest.mark.asyncio


async def test_get_throughput_small_load():
    """Test throughput under small load."""
    responses = [
        MockResponse(200, json_data={"name": f"model{i}", "version": f"{i}.0", "details": f"details{i}"})
        for i in range(5)
    ]
    client = AsyncRawModelsClient(client_wrapper=AsyncClientWrapper(MockAsyncClient(responses)))
    tasks = [client.get(f"model{i}") for i in range(5)]
    results = await asyncio.gather(*tasks)
    for i, result in enumerate(results):
        pass

@pytest.mark.asyncio
async def test_get_throughput_medium_load():
    """Test throughput under medium load."""
    responses = [
        MockResponse(200, json_data={"name": f"model{i}", "version": f"{i}.0", "details": f"details{i}"})
        for i in range(50)
    ]
    client = AsyncRawModelsClient(client_wrapper=AsyncClientWrapper(MockAsyncClient(responses)))
    tasks = [client.get(f"model{i}") for i in range(50)]
    results = await asyncio.gather(*tasks)
    for i, result in enumerate(results):
        pass

@pytest.mark.asyncio
async def test_get_throughput_high_load():
    """Test throughput under high load (max 100)."""
    responses = [
        MockResponse(200, json_data={"name": f"model{i}", "version": f"{i}.0", "details": f"details{i}"})
        for i in range(100)
    ]
    client = AsyncRawModelsClient(client_wrapper=AsyncClientWrapper(MockAsyncClient(responses)))
    tasks = [client.get(f"model{i}") for i in range(100)]
    results = await asyncio.gather(*tasks)
    for i, result in enumerate(results):
        pass

@pytest.mark.asyncio

#------------------------------------------------
import asyncio  # used to run async functions
from unittest.mock import AsyncMock, MagicMock, patch

import pytest  # used for our unit tests
from cohere.models.raw_client import AsyncRawModelsClient

# --- Minimal stubs for dependent classes and errors (to allow the test suite to run) ---

class DummyResponse:
    def __init__(self, status_code=200, json_data=None, text="", headers=None):
        self.status_code = status_code
        self._json_data = json_data or {}
        self.text = text
        self.headers = headers or {}

    def json(self):
        if self._json_data is not None:
            return self._json_data
        raise ValueError("No JSON data")

class DummyAsyncHttpResponse:
    def __init__(self, response, data):
        self.response = response
        self.data = data

class DummyGetModelResponse:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

# Error classes (minimal stubs)
class ApiError(Exception):
    def __init__(self, status_code, headers, body):
        super().__init__(f"ApiError: {status_code}")
        self.status_code = status_code
        self.headers = headers
        self.body = body

class BadRequestError(ApiError): pass
class UnauthorizedError(ApiError): pass
class ForbiddenError(ApiError): pass
class NotFoundError(ApiError): pass
class UnprocessableEntityError(ApiError): pass
class TooManyRequestsError(ApiError): pass
class InvalidTokenError(ApiError): pass
class ClientClosedRequestError(ApiError): pass
class InternalServerError(ApiError): pass
class NotImplementedError(ApiError): pass
class ServiceUnavailableError(ApiError): pass
class GatewayTimeoutError(ApiError): pass
from cohere.models.raw_client import AsyncRawModelsClient

# --- Helper for constructing a mock client_wrapper ---

class DummyHttpxClient:
    def __init__(self, response):
        self._response = response
        self.request_called = False
        self.request_args = None

    async def request(self, *args, **kwargs):
        self.request_called = True
        self.request_args = (args, kwargs)
        # Simulate async delay
        await asyncio.sleep(0)
        # If self._response is a callable, call it to allow dynamic responses
        if callable(self._response):
            return self._response(*args, **kwargs)
        return self._response

class DummyClientWrapper:
    def __init__(self, httpx_client):
        self.httpx_client = httpx_client

# --- TESTS START HERE ---

# 1. BASIC TEST CASES

@pytest.mark.asyncio
async def test_get_success_basic():
    """Test get returns expected AsyncHttpResponse on 200 OK."""
    response_json = {"id": "abc", "name": "test-model"}
    dummy_response = DummyResponse(status_code=200, json_data=response_json)
    client = AsyncRawModelsClient(
        client_wrapper=DummyClientWrapper(DummyHttpxClient(dummy_response))
    )
    result = await client.get("test-model")

@pytest.mark.asyncio
async def test_get_success_with_request_options():
    """Test get returns expected AsyncHttpResponse with request_options."""
    response_json = {"id": "xyz", "name": "option-model"}
    dummy_response = DummyResponse(status_code=200, json_data=response_json)
    client = AsyncRawModelsClient(
        client_wrapper=DummyClientWrapper(DummyHttpxClient(dummy_response))
    )
    request_options = {"timeout_in_seconds": 10}
    result = await client.get("option-model", request_options=request_options)

@pytest.mark.asyncio
async def test_get_success_different_model_names():
    """Test get returns correct data for different model names."""
    for model_name in ["foo", "bar", "baz"]:
        response_json = {"id": model_name, "name": f"model-{model_name}"}
        dummy_response = DummyResponse(status_code=200, json_data=response_json)
        client = AsyncRawModelsClient(
            client_wrapper=DummyClientWrapper(DummyHttpxClient(dummy_response))
        )
        result = await client.get(model_name)

# 2. EDGE TEST CASES

@pytest.mark.asyncio






async def test_get_concurrent_requests():
    """Test multiple concurrent get requests return correct results."""
    def response_factory(*args, **kwargs):
        # args[0] is the path, which contains the model name
        model_name = args[0].split("/")[-1]
        return DummyResponse(status_code=200, json_data={"id": model_name, "name": f"model-{model_name}"})

    client = AsyncRawModelsClient(
        client_wrapper=DummyClientWrapper(DummyHttpxClient(response_factory))
    )
    model_names = [f"model{i}" for i in range(10)]
    results = await asyncio.gather(*(client.get(name) for name in model_names))
    for i, result in enumerate(results):
        pass

@pytest.mark.asyncio
async def test_get_handles_empty_json_response():
    """Test get handles empty JSON response on 200."""
    dummy_response = DummyResponse(status_code=200, json_data={})
    client = AsyncRawModelsClient(
        client_wrapper=DummyClientWrapper(DummyHttpxClient(dummy_response))
    )
    result = await client.get("empty-model")

# 3. LARGE SCALE TEST CASES

@pytest.mark.asyncio
async def test_get_large_scale_concurrent():
    """Test get handles 100 concurrent requests correctly."""
    def response_factory(*args, **kwargs):
        model_name = args[0].split("/")[-1]
        return DummyResponse(status_code=200, json_data={"id": model_name, "name": f"model-{model_name}"})

    client = AsyncRawModelsClient(
        client_wrapper=DummyClientWrapper(DummyHttpxClient(response_factory))
    )
    model_names = [f"large_model_{i}" for i in range(100)]
    results = await asyncio.gather(*(client.get(name) for name in model_names))
    for i, result in enumerate(results):
        pass

@pytest.mark.asyncio

async def test_get_throughput_small_load():
    """Throughput test: small load (10 concurrent requests)."""
    def response_factory(*args, **kwargs):
        model_name = args[0].split("/")[-1]
        return DummyResponse(status_code=200, json_data={"id": model_name, "name": f"model-{model_name}"})

    client = AsyncRawModelsClient(
        client_wrapper=DummyClientWrapper(DummyHttpxClient(response_factory))
    )
    model_names = [f"throughput_small_{i}" for i in range(10)]
    results = await asyncio.gather(*(client.get(name) for name in model_names))
    for i, result in enumerate(results):
        pass

@pytest.mark.asyncio
async def test_get_throughput_medium_load():
    """Throughput test: medium load (100 concurrent requests)."""
    def response_factory(*args, **kwargs):
        model_name = args[0].split("/")[-1]
        return DummyResponse(status_code=200, json_data={"id": model_name, "name": f"model-{model_name}"})

    client = AsyncRawModelsClient(
        client_wrapper=DummyClientWrapper(DummyHttpxClient(response_factory))
    )
    model_names = [f"throughput_medium_{i}" for i in range(100)]
    results = await asyncio.gather(*(client.get(name) for name in model_names))
    for i, result in enumerate(results):
        pass

@pytest.mark.asyncio

async def test_get_throughput_high_volume():
    """Throughput test: high volume (500 concurrent requests)."""
    def response_factory(*args, **kwargs):
        model_name = args[0].split("/")[-1]
        return DummyResponse(status_code=200, json_data={"id": model_name, "name": f"model-{model_name}"})

    client = AsyncRawModelsClient(
        client_wrapper=DummyClientWrapper(DummyHttpxClient(response_factory))
    )
    model_names = [f"throughput_high_{i}" for i in range(500)]
    results = await asyncio.gather(*(client.get(name) for name in model_names))
    # Spot check a few
    for idx in [0, 99, 199, 299, 399, 499]:
        pass
# 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-AsyncRawModelsClient.get-mh11mex7 and push.