test_openai_o1.py

importjson
importos
importsys
fromdatetimeimportdatetime
fromunittest.mockimportAsyncMock, patch, MagicMock

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, BaseOSeriesModelsTest


@pytest.mark.parametrize("model", ["o1-preview", "o1-mini", "o1"])
@pytest.mark.asyncio
asyncdeftest_o1_handle_system_role(model):
"""
 Tests that:
 - max_tokens is translated to 'max_completion_tokens'
 - role 'system' is translated to 'user'
 """
fromopenaiimportAsyncOpenAI
fromlitellm.utilsimportsupports_system_messages

os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] ="True"
litellm.model_cost=litellm.get_model_cost_map(url="")

litellm.set_verbose=True

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

withpatch.object(
client.chat.completions.with_raw_response, "create"
 ) asmock_client:
try:
awaitlitellm.acompletion(
model=model,
max_tokens=10,
messages=[{"role": "system", "content": "Be a good bot!"}],
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["model"] ==model
assertrequest_body["max_completion_tokens"] ==10
ifsupports_system_messages(model, "openai"):
assertrequest_body["messages"] == [
 {"role": "system", "content": "Be a good bot!"}
 ]
else:
assertrequest_body["messages"] == [
 {"role": "user", "content": "Be a good bot!"}
 ]


@pytest.mark.parametrize(
"model, expected_tool_calling_support",
 [("o1-preview", False), ("o1-mini", False), ("o1", True)],
)
@pytest.mark.asyncio
asyncdeftest_o1_handle_tool_calling_optional_params(
model, expected_tool_calling_support
):
"""
 Tests that:
 - max_tokens is translated to 'max_completion_tokens'
 - role 'system' is translated to 'user'
 """
fromopenaiimportAsyncOpenAI
fromlitellm.utilsimportProviderConfigManager
fromlitellm.types.utilsimportLlmProviders

os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] ="True"
litellm.model_cost=litellm.get_model_cost_map(url="")

config=ProviderConfigManager.get_provider_chat_config(
model=model, provider=LlmProviders.OPENAI
 )

supported_params=config.get_supported_openai_params(model=model)

assertexpected_tool_calling_support== ("tools"insupported_params)


@pytest.mark.asyncio
@pytest.mark.parametrize("model", ["gpt-4", "gpt-4-0314", "gpt-4-32k", "o1-preview"])
asyncdeftest_o1_max_completion_tokens(model: str):
"""
 Tests that:
 - max_completion_tokens is passed directly to OpenAI chat completion models
 """
fromopenaiimportAsyncOpenAI

litellm.set_verbose=True

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

withpatch.object(
client.chat.completions.with_raw_response, "create"
 ) asmock_client:
try:
awaitlitellm.acompletion(
model=model,
max_completion_tokens=10,
messages=[{"role": "user", "content": "Hello!"}],
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["model"] ==model
assertrequest_body["max_completion_tokens"] ==10
assertrequest_body["messages"] == [{"role": "user", "content": "Hello!"}]


deftest_litellm_responses():
"""
 ensures that type of completion_tokens_details is correctly handled / returned
 """
fromlitellmimportModelResponse
fromlitellm.types.utilsimportCompletionTokensDetails

response=ModelResponse(
usage={
"completion_tokens": 436,
"prompt_tokens": 14,
"total_tokens": 450,
"completion_tokens_details": {"reasoning_tokens": 0},
 }
 )

print("response: ", response)

assertisinstance(response.usage.completion_tokens_details, CompletionTokensDetails)


classTestOpenAIO1(BaseOSeriesModelsTest, BaseLLMChatTest):
defget_base_completion_call_args(self):
return {
"model": "o1",
 }

defget_client(self):
fromopenaiimportOpenAI

returnOpenAI(api_key="fake-api-key")

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):
"""Temporary override. o1 prompt caching is not working."""
pass


classTestOpenAIO3(BaseOSeriesModelsTest, BaseLLMChatTest):
defget_base_completion_call_args(self):
return {
"model": "o3-mini",
 }

defget_client(self):
fromopenaiimportOpenAI

returnOpenAI(api_key="fake-api-key")

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_o1_supports_vision():
"""Test that o1 supports vision"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] ="True"
litellm.model_cost=litellm.get_model_cost_map(url="")
fork, vinlitellm.model_cost.items():
ifk.startswith("o1") andv.get("litellm_provider") =="openai":
assertv.get("supports_vision") isTrue, f"{k} does not support vision"


deftest_o3_reasoning_effort():
resp=litellm.completion(
model="o3-mini",
messages=[{"role": "user", "content": "Hello!"}],
reasoning_effort="high",
 )
assertresp.choices[0].message.contentisnotNone


@pytest.mark.parametrize("model", ["o1-preview", "o1-mini", "o1", "o3-mini"])
deftest_streaming_response(model):
"""Test that streaming response is returned correctly"""
fromlitellmimportcompletion

response=completion(
model=model,
messages=[
 {"role": "system", "content": "Be a good bot!"},
 {"role": "user", "content": "Hello!"},
 ],
stream=True,
 )

assertresponseisnotNone

chunks= []
forchunkinresponse:
chunks.append(chunk)

resp=litellm.stream_chunk_builder(chunks=chunks)
print(resp)