test_openai.py

importjson
importos
importsys
fromdatetimeimportdatetime
fromunittest.mockimportAsyncMock, patch
fromtypingimportOptional

sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path


importhttpx
importpytest
fromrespximportMockRouter

importlitellm
fromlitellmimportChoices, Message, ModelResponse
frombase_llm_unit_testsimportBaseLLMChatTest
importasyncio
fromlitellm.types.llms.openaiimport (
ChatCompletionAnnotation,
ChatCompletionAnnotationURLCitation,
)
frombase_audio_transcription_unit_testsimportBaseLLMAudioTranscriptionTest


deftest_openai_prediction_param():
litellm.set_verbose=True
code="""
 /// <summary>
 /// Represents a user with a first name, last name, and username.
 /// </summary>
 public class User
 {
 /// <summary>
 /// Gets or sets the user's first name.
 /// </summary>
 public string FirstName { get; set; }

 /// <summary>
 /// Gets or sets the user's last name.
 /// </summary>
 public string LastName { get; set; }

 /// <summary>
 /// Gets or sets the user's username.
 /// </summary>
 public string Username { get; set; }
 }
 """

completion=litellm.completion(
model="gpt-4o-mini",
messages=[
 {
"role": "user",
"content": "Replace the Username property with an Email property. Respond only with code, and with no markdown formatting.",
 },
 {"role": "user", "content": code},
 ],
prediction={"type": "content", "content": code},
 )

print(completion)

assert (
completion.usage.completion_tokens_details.accepted_prediction_tokens>0
orcompletion.usage.completion_tokens_details.rejected_prediction_tokens>0
 )


@pytest.mark.asyncio
asyncdeftest_openai_prediction_param_mock():
"""
 Tests that prediction parameter is correctly passed to the API
 """
litellm.set_verbose=True

code="""
 /// <summary>
 /// Represents a user with a first name, last name, and username.
 /// </summary>
 public class User
 {
 /// <summary>
 /// Gets or sets the user's first name.
 /// </summary>
 public string FirstName { get; set; }

 /// <summary>
 /// Gets or sets the user's last name.
 /// </summary>
 public string LastName { get; set; }

 /// <summary>
 /// Gets or sets the user's username.
 /// </summary>
 public string Username { get; set; }
 }
 """
fromopenaiimportAsyncOpenAI

client=AsyncOpenAI(api_key="fake-api-key")

withpatch.object(
client.chat.completions.with_raw_response, "create"
 ) asmock_client:
try:
awaitlitellm.acompletion(
model="gpt-4o-mini",
messages=[
 {
"role": "user",
"content": "Replace the Username property with an Email property. Respond only with code, and with no markdown formatting.",
 },
 {"role": "user", "content": code},
 ],
prediction={"type": "content", "content": code},
client=client,
 )
exceptExceptionase:
print(f"Error: {e}")

mock_client.assert_called_once()
request_body=mock_client.call_args.kwargs

# Verify the request contains the prediction parameter
assert"prediction"inrequest_body
# verify prediction is correctly sent to the API
assertrequest_body["prediction"] == {"type": "content", "content": code}


@pytest.mark.asyncio
asyncdeftest_openai_prediction_param_with_caching():
"""
 Tests using `prediction` parameter with caching
 """
fromlitellm.caching.cachingimportLiteLLMCacheType
importlogging
fromlitellm._loggingimportverbose_logger

verbose_logger.setLevel(logging.DEBUG)
importtime

litellm.set_verbose=True
litellm.cache=litellm.Cache(type=LiteLLMCacheType.LOCAL)
code="""
 /// <summary>
 /// Represents a user with a first name, last name, and username.
 /// </summary>
 public class User
 {
 /// <summary>
 /// Gets or sets the user's first name.
 /// </summary>
 public string FirstName { get; set; }

 /// <summary>
 /// Gets or sets the user's last name.
 /// </summary>
 public string LastName { get; set; }

 /// <summary>
 /// Gets or sets the user's username.
 /// </summary>
 public string Username { get; set; }
 }
 """

completion_response_1=litellm.completion(
model="gpt-4o-mini",
messages=[
 {
"role": "user",
"content": "Replace the Username property with an Email property. Respond only with code, and with no markdown formatting.",
 },
 {"role": "user", "content": code},
 ],
prediction={"type": "content", "content": code},
 )

time.sleep(0.5)

# cache hit
completion_response_2=litellm.completion(
model="gpt-4o-mini",
messages=[
 {
"role": "user",
"content": "Replace the Username property with an Email property. Respond only with code, and with no markdown formatting.",
 },
 {"role": "user", "content": code},
 ],
prediction={"type": "content", "content": code},
 )

assertcompletion_response_1.id==completion_response_2.id

completion_response_3=litellm.completion(
model="gpt-4o-mini",
messages=[
 {"role": "user", "content": "What is the first name of the user?"},
 ],
prediction={"type": "content", "content": code+"FirstName"},
 )

assertcompletion_response_3.id!=completion_response_1.id


@pytest.mark.asyncio()
asyncdeftest_vision_with_custom_model():
"""
 Tests that an OpenAI compatible endpoint when sent an image will receive the image in the request

 """
importbase64
importrequests
fromopenaiimportAsyncOpenAI

client=AsyncOpenAI(api_key="fake-api-key")

litellm.set_verbose=True
api_base="https://my-custom.api.openai.com"

# Fetch and encode a test image
url="https://dummyimage.com/100/100/fff&text=Test+image"
response=requests.get(url)
file_data=response.content
encoded_file=base64.b64encode(file_data).decode("utf-8")
base64_image=f"data:image/png;base64,{encoded_file}"

withpatch.object(
client.chat.completions.with_raw_response, "create"
 ) asmock_client:
try:
response=awaitlitellm.acompletion(
model="openai/my-custom-model",
max_tokens=10,
api_base=api_base, # use the mock api
messages=[
 {
"role": "user",
"content": [
 {"type": "text", "text": "What's in this image?"},
 {
"type": "image_url",
"image_url": {"url": base64_image},
 },
 ],
 }
 ],
client=client,
 )
exceptExceptionase:
print(f"Error: {e}")

mock_client.assert_called_once()
request_body=mock_client.call_args.kwargs

print("request_body: ", request_body)

assertrequest_body["messages"] == [
 {
"role": "user",
"content": [
 {"type": "text", "text": "What's in this image?"},
 {
"type": "image_url",
"image_url": {
"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAGQAAABkBAMAAACCzIhnAAAAG1BMVEURAAD///+ln5/h39/Dv79qX18uHx+If39MPz9oMSdmAAAACXBIWXMAAA7EAAAOxAGVKw4bAAABDElEQVRYhe2SzWqEMBRGPyQTfQxJsc5jBKGzFmlslyFIZxsCQ7sUaWd87EanpdpIrbtC71mE/NyTm9wEIAiCIAiC+N/otQBxU2Sf/aeh4enqptHXri+/yxIq63jlKCw6cXssnr3ObdzdGYFYCJ2IzHKXLygHXCB98Gm4DE+ZZemu5EisQSyZTmyg+AuzQbkezCuIy7EI0k9Ig3FtruwydY+qniqtV5yQyo8qpUIl2fc90KVzJWohWf2qu75vlw52rdfjVDHg8vLWwixW7PChqLkSyUadwfSS0uQZhEvRuIkS53uJvrK8cGWYaPwpGt8efvw+vlo8TPMzcmP8w7lrNypc1RsNgiAIgiD+Iu/RyDYhCaWrgQAAAABJRU5ErkJggg=="
 },
 },
 ],
 },
 ]
assertrequest_body["model"] =="my-custom-model"
assertrequest_body["max_tokens"] ==10


classTestOpenAIChatCompletion(BaseLLMChatTest):
defget_base_completion_call_args(self) ->dict:
return {"model": "gpt-4o-mini"}

deftest_tool_call_no_arguments(self, tool_call_no_arguments):
"""Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833"""
pass

deftest_prompt_caching(self):
"""
 Test that prompt caching works correctly.
 Skip for now, as it's working locally but not in CI
 """
pass

deftest_multilingual_requests(self):
"""
 Tests that the provider can handle multilingual requests and invalid utf-8 sequences

 Context: https://github.com/openai/openai-python/issues/1921
 """
base_completion_call_args=self.get_base_completion_call_args()
try:
response=self.completion_function(
**base_completion_call_args,
messages=[{"role": "user", "content": "你好世界！\ud83e, ö"}],
 )
assertresponseisnotNone
exceptlitellm.InternalServerError:
pytest.skip("Skipping test due to InternalServerError")

deftest_prompt_caching(self):
"""
 Works locally but CI/CD is failing this test. Temporary skip to push out a new release.
 """
pass


deftest_completion_bad_org():
importlitellm

litellm.set_verbose=True
_old_org=os.environ.get("OPENAI_ORGANIZATION", None)
os.environ["OPENAI_ORGANIZATION"] ="bad-org"
messages= [{"role": "user", "content": "hi"}]

withpytest.raises(Exception) asexc_info:
comp=litellm.completion(
model="gpt-4o-mini", messages=messages, organization="bad-org"
 )

print(exc_info.value)
assert"header should match organization for API key"instr(exc_info.value)

if_old_orgisnotNone:
os.environ["OPENAI_ORGANIZATION"] =_old_org
else:
delos.environ["OPENAI_ORGANIZATION"]


@patch("litellm.main.openai_chat_completions._get_openai_client")
deftest_openai_max_retries_0(mock_get_openai_client):
importlitellm

litellm.set_verbose=True
response=litellm.completion(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
max_retries=0,
 )

mock_get_openai_client.assert_called_once()
assertmock_get_openai_client.call_args.kwargs["max_retries"] ==0


@pytest.mark.parametrize("model", ["o1", "o1-preview", "o1-mini", "o3-mini"])
deftest_o1_parallel_tool_calls(model):
litellm.completion(
model=model,
messages=[
 {
"role": "user",
"content": "foo",
 }
 ],
parallel_tool_calls=True,
drop_params=True,
 )


deftest_openai_chat_completion_streaming_handler_reasoning_content():
fromlitellm.llms.openai.chat.gpt_transformationimport (
OpenAIChatCompletionStreamingHandler,
 )
fromunittest.mockimportMagicMock

streaming_handler=OpenAIChatCompletionStreamingHandler(
streaming_response=MagicMock(),
sync_stream=True,
 )
response=streaming_handler.chunk_parser(
chunk={
"id": "e89b6501-8ac2-464c-9550-7cd3daf94350",
"object": "chat.completion.chunk",
"created": 1741037890,
"model": "deepseek-reasoner",
"system_fingerprint": "fp_5417b77867_prod0225",
"choices": [
 {
"index": 0,
"delta": {"content": None, "reasoning_content": "."},
"logprobs": None,
"finish_reason": None,
 }
 ],
 }
 )

assertresponse.choices[0].delta.reasoning_content=="."


defvalidate_response_url_citation(url_citation: ChatCompletionAnnotationURLCitation):
assert"end_index"inurl_citation
assert"start_index"inurl_citation
assert"url"inurl_citation


defvalidate_web_search_annotations(annotations: ChatCompletionAnnotation):
"""validates litellm response contains web search annotations"""
print("annotations: ", annotations)
assertannotationsisnotNone
assertisinstance(annotations, list)
forannotationinannotations:
assertannotation["type"] =="url_citation"
url_citation: ChatCompletionAnnotationURLCitation=annotation["url_citation"]
validate_response_url_citation(url_citation)


@pytest.mark.flaky(reruns=3)
deftest_openai_web_search():
"""Makes a simple web search request and validates the response contains web search annotations and all expected fields are present"""
litellm._turn_on_debug()
response=litellm.completion(
model="openai/gpt-4o-search-preview",
messages=[
 {
"role": "user",
"content": "What was a positive news story from today?",
 }
 ],
 )
print("litellm response: ", response.model_dump_json(indent=4))
message=response.choices[0].message
annotations: ChatCompletionAnnotation=message.annotations
validate_web_search_annotations(annotations)


deftest_openai_web_search_streaming():
"""Makes a simple web search request and validates the response contains web search annotations and all expected fields are present"""
# litellm._turn_on_debug()
test_openai_web_search: Optional[ChatCompletionAnnotation] =None
response=litellm.completion(
model="openai/gpt-4o-search-preview",
messages=[
 {
"role": "user",
"content": "What was a positive news story from today?",
 }
 ],
stream=True,
 )
forchunkinresponse:
print("litellm response chunk: ", chunk)
if (
hasattr(chunk.choices[0].delta, "annotations")
andchunk.choices[0].delta.annotationsisnotNone
 ):
test_openai_web_search=chunk.choices[0].delta.annotations

# Assert this request has at-least one web search annotation
asserttest_openai_web_searchisnotNone
validate_web_search_annotations(test_openai_web_search)


classTestOpenAIGPT4OAudioTranscription(BaseLLMAudioTranscriptionTest):
defget_base_audio_transcription_call_args(self) ->dict:
return {
"model": "openai/gpt-4o-transcribe",
 }

defget_custom_llm_provider(self) ->litellm.LlmProviders:
returnlitellm.LlmProviders.OPENAI