import asyncio # used to run async functions
import json

Patch the actual TextCompletionResponse and ModelResponse for the test

import sys
import types
from typing import Optional

import litellm
import pytest # used for our unit tests
from litellm.llms.base import BaseLLM
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.openai.completion.handler import OpenAITextCompletion
from litellm.llms.openai.completion.transformation import
OpenAITextCompletionConfig
from litellm.types.utils import ModelResponse, TextCompletionResponse
from openai import AsyncOpenAI

Mocks and helpers for testing

class DummyAsyncOpenAIClient:
"""A dummy AsyncOpenAI client that simulates the OpenAI API."""
class completions:
class with_raw_response:
@staticmethod
async def create(**data):
# Simulate a successful response with a parse() method
class DummyRawResponse:
def parse(self):
class DummyParsedResponse:
def model_dump(self):
# Return a dict with keys expected by TextCompletionResponse
return {
"id": "dummy_id",
"choices": [{"text": "Hello, world!", "index": 0}],
"model": data.get("model", "test-model"),
"object": "text_completion",
"created": 1234567890,
}
return DummyParsedResponse()
return DummyRawResponse()

class DummyLoggingObj:
"""A dummy logging object that simulates the Logging class."""
def post_call(self, api_key, original_response, additional_args):
# Simulate logging without side effects
self.logged = True

class DummyModelResponse:
"""A dummy ModelResponse for passing to the function."""
pass

class DummyTextCompletionResponse:
"""A dummy TextCompletionResponse for type checking."""
def init(self, **kwargs):
self.id = kwargs.get("id")
self.choices = kwargs.get("choices")
self.model = kwargs.get("model")
self.object = kwargs.get("object")
self.created = kwargs.get("created")
self._hidden_params = type("HiddenParams", (), {})()
self._hidden_params.original_response = kwargs.get("original_response", "")
from litellm.llms.openai.completion.handler import OpenAITextCompletion

-------------- UNIT TESTS --------------

1. Basic Test Cases

@pytest.mark.asyncio

async def test_acompletion_basic_async_await():
"""Test that acompletion is truly async and can be awaited."""
completion = OpenAITextCompletion()
logging_obj = DummyLoggingObj()
api_base = "https://dummy.openai.api"
data = {"model": "test-model", "prompt": "Say hello"}
headers = {"Authorization": "Bearer dummy"}
model_response = DummyModelResponse()
api_key = "dummy-key"
model = "test-model"
timeout = 10.0
max_retries = 3
client = DummyAsyncOpenAIClient()

# Await the coroutine and check result type
result = await completion.acompletion(
    logging_obj, api_base, data, headers, model_response, api_key, model, timeout, max_retries, client=client
)

2. Edge Test Cases

@pytest.mark.asyncio

async def test_acompletion_edge_missing_fields():
"""Test that acompletion handles missing fields gracefully."""
completion = OpenAITextCompletion()
logging_obj = DummyLoggingObj()
api_base = "https://dummy.openai.api"
# Data missing 'model' field
data = {"prompt": "Say hello"}
headers = {"Authorization": "Bearer dummy"}
model_response = DummyModelResponse()
api_key = "dummy-key"
model = "test-model"
timeout = 10.0
max_retries = 3
client = DummyAsyncOpenAIClient()

result = await completion.acompletion(
    logging_obj, api_base, data, headers, model_response, api_key, model, timeout, max_retries, client=client
)

3. Large Scale Test Cases

@pytest.mark.asyncio
async def test_acompletion_large_scale_concurrent():
"""Test acompletion with 50 concurrent executions."""
completion = OpenAITextCompletion()
logging_obj = DummyLoggingObj()
api_base = "https://dummy.openai.api"
headers = {"Authorization": "Bearer dummy"}
model_response = DummyModelResponse()
api_key = "dummy-key"
model = "test-model"
timeout = 10.0
max_retries = 3
client = DummyAsyncOpenAIClient()

async def call_acompletion(prompt):
    data = {"model": model, "prompt": prompt}
    return await completion.acompletion(
        logging_obj, api_base, data, headers, model_response, api_key, model, timeout, max_retries, client=client
    )

prompts = [f"Prompt {i}" for i in range(50)]
results = await asyncio.gather(*(call_acompletion(p) for p in prompts))
for res in results:
    pass

4. Throughput Test Cases

@pytest.mark.asyncio
async def test_acompletion_throughput_small_load():
"""Test throughput with a small load of 10 concurrent calls."""
completion = OpenAITextCompletion()
logging_obj = DummyLoggingObj()
api_base = "https://dummy.openai.api"
headers = {"Authorization": "Bearer dummy"}
model_response = DummyModelResponse()
api_key = "dummy-key"
model = "test-model"
timeout = 10.0
max_retries = 3
client = DummyAsyncOpenAIClient()

async def call_acompletion(prompt):
    data = {"model": model, "prompt": prompt}
    return await completion.acompletion(
        logging_obj, api_base, data, headers, model_response, api_key, model, timeout, max_retries, client=client
    )

prompts = [f"Prompt {i}" for i in range(10)]
results = await asyncio.gather(*(call_acompletion(p) for p in prompts))
for res in results:
    pass

@pytest.mark.asyncio
async def test_acompletion_throughput_medium_load():
"""Test throughput with a medium load of 100 concurrent calls."""
completion = OpenAITextCompletion()
logging_obj = DummyLoggingObj()
api_base = "https://dummy.openai.api"
headers = {"Authorization": "Bearer dummy"}
model_response = DummyModelResponse()
api_key = "dummy-key"
model = "test-model"
timeout = 10.0
max_retries = 3
client = DummyAsyncOpenAIClient()

async def call_acompletion(prompt):
    data = {"model": model, "prompt": prompt}
    return await completion.acompletion(
        logging_obj, api_base, data, headers, model_response, api_key, model, timeout, max_retries, client=client
    )

prompts = [f"Prompt {i}" for i in range(100)]
results = await asyncio.gather(*(call_acompletion(p) for p in prompts))
for res in results:
    pass

@pytest.mark.asyncio
async def test_acompletion_throughput_high_volume():
"""Test throughput with a high volume of 250 concurrent calls."""
completion = OpenAITextCompletion()
logging_obj = DummyLoggingObj()
api_base = "https://dummy.openai.api"
headers = {"Authorization": "Bearer dummy"}
model_response = DummyModelResponse()
api_key = "dummy-key"
model = "test-model"
timeout = 10.0
max_retries = 3
client = DummyAsyncOpenAIClient()

async def call_acompletion(prompt):
    data = {"model": model, "prompt": prompt}
    return await completion.acompletion(
        logging_obj, api_base, data, headers, model_response, api_key, model, timeout, max_retries, client=client
    )

prompts = [f"Prompt {i}" for i in range(250)]
results = await asyncio.gather(*(call_acompletion(p) for p in prompts))
for res in results:
    pass

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

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

import pytest # used for our unit tests
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.openai.completion.handler import OpenAITextCompletion
from litellm.types.utils import ModelResponse, TextCompletionResponse

Helper function to construct a minimal valid logging object

class DummyLoggingObj:
def post_call(self, **kwargs):
self.called = True
self.kwargs = kwargs

@pytest.mark.asyncio

async def test_acompletion_edge_invalid_client_raises():
"""
Edge case: If client raises an exception, ensure OpenAIError is raised and contains details.
"""
logging_obj = DummyLoggingObj()
dummy_client = MagicMock()
dummy_client.completions.with_raw_response.create = AsyncMock(side_effect=Exception("API failure"))
api_base = "https://api.openai.com/v1"
data = {"model": "gpt-3.5-turbo", "prompt": "Fail", "max_tokens": 5}
headers = {"Authorization": "Bearer testkey"}
model_response = MagicMock()
api_key = "testkey"
model = "gpt-3.5-turbo"
timeout = 10.0
max_retries = 1

with pytest.raises(OpenAIError) as excinfo:
    await OpenAITextCompletion().acompletion(
        logging_obj=logging_obj,
        api_base=api_base,
        data=data,
        headers=headers,
        model_response=model_response,
        api_key=api_key,
        model=model,
        timeout=timeout,
        max_retries=max_retries,
        client=dummy_client,
    )

@pytest.mark.asyncio

async def test_acompletion_edge_exception_with_status_code_and_headers():
"""
Edge case: Exception with status_code and headers should propagate correctly.
"""
class DummyException(Exception):
def init(self):
self.status_code = 403
self.headers = {"X-Error": "Forbidden"}
self.text = "Forbidden"
self.response = None

logging_obj = DummyLoggingObj()
dummy_client = MagicMock()
dummy_client.completions.with_raw_response.create = AsyncMock(side_effect=DummyException())
api_base = "https://api.openai.com/v1"
data = {"model": "gpt-3.5-turbo", "prompt": "Forbidden", "max_tokens": 5}
headers = {"Authorization": "Bearer testkey"}
model_response = MagicMock()
api_key = "testkey"
model = "gpt-3.5-turbo"
timeout = 10.0
max_retries = 1

with pytest.raises(OpenAIError) as excinfo:
    await OpenAITextCompletion().acompletion(
        logging_obj=logging_obj,
        api_base=api_base,
        data=data,
        headers=headers,
        model_response=model_response,
        api_key=api_key,
        model=model,
        timeout=timeout,
        max_retries=max_retries,
        client=dummy_client,
    )

@pytest.mark.asyncio