test_azure_ai.py

# What is this?
## Unit tests for Azure AI integration

importasyncio
importos
importsys
importtraceback

fromdotenvimportload_dotenv

importlitellm.types
importlitellm.types.utils
fromlitellm.llms.anthropic.chatimportModelResponseIterator
importhttpx
importjson
fromlitellm.llms.custom_httpx.http_handlerimportHTTPHandler
# from base_rerank_unit_tests import BaseLLMRerankTest

load_dotenv()
importio
importos

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

importpytest

importlitellm
fromlitellmimportcompletion


@pytest.mark.parametrize(
"model_group_header, expected_model",
 [
 ("offer-cohere-embed-multili-paygo", "Cohere-embed-v3-multilingual"),
 ("offer-cohere-embed-english-paygo", "Cohere-embed-v3-english"),
 ],
)
deftest_map_azure_model_group(model_group_header, expected_model):
fromlitellm.llms.azure_ai.embed.cohere_transformationimportAzureAICohereConfig

config=AzureAICohereConfig()
assertconfig._map_azure_model_group(model_group_header) ==expected_model


@pytest.mark.asyncio
asyncdeftest_azure_ai_with_image_url():
"""
 Important test:

 Test that Azure AI studio can handle image_url passed when content is a list containing both text and image_url
 """
fromlitellm.llms.custom_httpx.http_handlerimportAsyncHTTPHandler

litellm.set_verbose=True

client=AsyncHTTPHandler()

withpatch.object(client, "post") asmock_client:
try:
awaitlitellm.acompletion(
model="azure_ai/Phi-3-5-vision-instruct-dcvov",
api_base="https://Phi-3-5-vision-instruct-dcvov.eastus2.models.ai.azure.com",
messages=[
 {
"role": "user",
"content": [
 {
"type": "text",
"text": "What is in this image?",
 },
 {
"type": "image_url",
"image_url": {
"url": "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
 },
 },
 ],
 },
 ],
api_key="fake-api-key",
client=client,
 )
exceptExceptionase:
traceback.print_exc()
print(f"Error: {e}")

# Verify the request was made
mock_client.assert_called_once()

print(f"mock_client.call_args.kwargs: {mock_client.call_args.kwargs}")
# Check the request body
request_body=json.loads(mock_client.call_args.kwargs["data"])
assertrequest_body["model"] =="Phi-3-5-vision-instruct-dcvov"
assertrequest_body["messages"] == [
 {
"role": "user",
"content": [
 {"type": "text", "text": "What is in this image?"},
 {
"type": "image_url",
"image_url": {
"url": "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
 },
 },
 ],
 }
 ]


@pytest.mark.parametrize(
"api_base, expected_url",
 [
 (
"https://litellm8397336933.services.ai.azure.com/models/chat/completions?api-version=2024-05-01-preview",
"https://litellm8397336933.services.ai.azure.com/models/chat/completions?api-version=2024-05-01-preview",
 ),
 (
"https://litellm8397336933.services.ai.azure.com/models/chat/completions",
"https://litellm8397336933.services.ai.azure.com/models/chat/completions",
 ),
 (
"https://litellm8397336933.services.ai.azure.com/models",
"https://litellm8397336933.services.ai.azure.com/models/chat/completions",
 ),
 (
"https://litellm8397336933.services.ai.azure.com",
"https://litellm8397336933.services.ai.azure.com/models/chat/completions",
 ),
 ],
)
deftest_azure_ai_services_handler(api_base, expected_url):
fromlitellm.llms.custom_httpx.http_handlerimportHTTPHandler

litellm.set_verbose=True

client=HTTPHandler()

withpatch.object(client, "post") asmock_client:
try:
response=litellm.completion(
model="azure_ai/Meta-Llama-3.1-70B-Instruct",
messages=[{"role": "user", "content": "Hello, how are you?"}],
api_key="my-fake-api-key",
api_base=api_base,
client=client,
 )

print(response)

exceptExceptionase:
print(f"Error: {e}")

mock_client.assert_called_once()
assertmock_client.call_args.kwargs["headers"]["api-key"] =="my-fake-api-key"
assertmock_client.call_args.kwargs["url"] ==expected_url


deftest_azure_ai_services_with_api_version():
fromlitellm.llms.custom_httpx.http_handlerimportHTTPHandler, AsyncHTTPHandler

client=HTTPHandler()

withpatch.object(client, "post") asmock_client:
try:
response=litellm.completion(
model="azure_ai/Meta-Llama-3.1-70B-Instruct",
messages=[{"role": "user", "content": "Hello, how are you?"}],
api_key="my-fake-api-key",
api_version="2024-05-01-preview",
api_base="https://litellm8397336933.services.ai.azure.com/models",
client=client,
 )
exceptExceptionase:
print(f"Error: {e}")

mock_client.assert_called_once()
assertmock_client.call_args.kwargs["headers"]["api-key"] =="my-fake-api-key"
assert (
mock_client.call_args.kwargs["url"]
=="https://litellm8397336933.services.ai.azure.com/models/chat/completions?api-version=2024-05-01-preview"
 )


deftest_completion_azure_ai_command_r():
try:
importos

litellm.set_verbose=True

os.environ["AZURE_AI_API_BASE"] =os.getenv("AZURE_COHERE_API_BASE", "")
os.environ["AZURE_AI_API_KEY"] =os.getenv("AZURE_COHERE_API_KEY", "")

response=completion(
model="azure_ai/command-r-plus",
messages=[
 {
"role": "user",
"content": [
 {"type": "text", "text": "What is the meaning of life?"}
 ],
 }
 ],
 ) # type: ignore

assert"azure_ai"inresponse.model
exceptlitellm.Timeoutase:
pass
exceptExceptionase:
pytest.fail(f"Error occurred: {e}")


deftest_azure_deepseek_reasoning_content():
importjson

client=HTTPHandler()

withpatch.object(client, "post") asmock_post:
mock_response=MagicMock()

mock_response.text=json.dumps(
 {
"choices": [
 {
"finish_reason": "stop",
"index": 0,
"message": {
"content": "<think>I am thinking here</think>\n\nThe sky is a canvas of blue",
"role": "assistant",
 },
 }
 ],
 }
 )

mock_response.status_code=200
# Add required response attributes
mock_response.headers= {"Content-Type": "application/json"}
mock_response.json=lambda: json.loads(mock_response.text)
mock_post.return_value=mock_response

response=litellm.completion(
model="azure_ai/deepseek-r1",
messages=[{"role": "user", "content": "Hello, world!"}],
api_base="https://litellm8397336933.services.ai.azure.com/models/chat/completions",
api_key="my-fake-api-key",
client=client,
 )

print(response)
assertresponse.choices[0].message.reasoning_content=="I am thinking here"
assertresponse.choices[0].message.content=="\n\nThe sky is a canvas of blue"


# skipping due to cohere rbac issues
# class TestAzureAIRerank(BaseLLMRerankTest):
# def get_custom_llm_provider(self) -> litellm.LlmProviders:
# return litellm.LlmProviders.AZURE_AI

# def get_base_rerank_call_args(self) -> dict:
# return {
# "model": "azure_ai/cohere-rerank-v3-english",
# "api_base": os.getenv("AZURE_AI_COHERE_API_BASE"),
# "api_key": os.getenv("AZURE_AI_COHERE_API_KEY"),
# }


@pytest.mark.asyncio
asyncdeftest_azure_ai_request_format():
"""
 Test that Azure AI requests are formatted correctly with the proper endpoint and parameters
 for both synchronous and asynchronous calls
 """
fromopenaiimportAsyncAzureOpenAI, AzureOpenAI

litellm._turn_on_debug()

# Set up the test parameters
api_key=os.getenv("AZURE_API_KEY")
api_base=f"{os.getenv('AZURE_API_BASE')}/openai/deployments/gpt-4o-new-test/chat/completions?api-version=2024-08-01-preview"
model="azure_ai/gpt-4o"
messages= [
 {"role": "user", "content": "hi"},
 {"role": "assistant", "content": "Hello! How can I assist you today?"},
 {"role": "user", "content": "hi"},
 ]

awaitlitellm.acompletion(
custom_llm_provider="azure_ai",
api_key=api_key,
api_base=api_base,
model=model,
messages=messages,
 )