Fix azure max retries error (BerriAI#8340)

* fix(azure.py): ensure max_retries=0 is respected Fixes BerriAI#6129 * fix(test_openai.py): add unit test to ensure openai sdk calls always respect max_retries = 0 * test(test_azure_openai.py): add unit testing for azure_text/ route * fix(azure.py): fix passing max retries on streaming * fix(azure.py): fix azure max retries on async completion + streaming * fix(completion/handler.py): fix azure text async completion + streaming * test(test_azure_openai.py): ensure azure openai max retries always respected * test(test_azure_o_series.py): add testing to ensure max retries always respected * Added gemini providers for 2.0-flash and 2.0-flash lite (BerriAI#8321) * Update model_prices_and_context_window.json added gemini providers for 2.0-flash and 2.0-flash light * Update model_prices_and_context_window.json fixed URL --------- Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com> * Convert tool use arguments to string before counting tokens (BerriAI#6989) In at least some cases the `messages["tool_calls"]["function"]["arguments"]` is a dict, not a string. In order to tokenize it properly it needs to be a string. In the case that it is already a string this is a noop, which is also fine. * build(model_prices_and_context_window.json): add gemini 2.0 flash lite pricing * build(model_prices_and_context_window.json): add gemini commercial rate limits * fix(utils.py): fix linting error * refactor(utils.py): refactor to maintain function size --------- Co-authored-by: Bardia Khosravi <bardiakhosravi95@gmail.com> Co-authored-by: Josh Morrow <josh@jcmorrow.com>

Author: 3 people

litellm/llms/azure/azure.py

@@ -9,6 +9,7 @@
 
 import litellm
 from litellm.caching.caching import DualCache
+from litellm.constants import DEFAULT_MAX_RETRIES
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.custom_httpx.http_handler import (
  AsyncHTTPHandler,
@@ -98,14 +99,6 @@ def map_openai_params_create_message_params(
 
 
 def select_azure_base_url_or_endpoint(azure_client_params: dict):
- # azure_client_params = {
- # "api_version": api_version,
- # "azure_endpoint": api_base,
- # "azure_deployment": model,
- # "http_client": litellm.client_session,
- # "max_retries": max_retries,
- # "timeout": timeout,
- # }
  azure_endpoint = azure_client_params.get("azure_endpoint", None)
  if azure_endpoint is not None:
  # see : https://github.com/openai/openai-python/blob/3d61ed42aba652b547029095a7eb269ad4e1e957/src/openai/lib/azure.py#L192
@@ -353,7 +346,9 @@ def completion( # noqa: PLR0915
  status_code=422, message="Missing model or messages"
  )
 
- max_retries = optional_params.pop("max_retries", 2)
+ max_retries = optional_params.pop("max_retries", None)
+ if max_retries is None:
+ max_retries = DEFAULT_MAX_RETRIES
  json_mode: Optional[bool] = optional_params.pop("json_mode", False)
 
  ### CHECK IF CLOUDFLARE AI GATEWAY ###
@@ -415,6 +410,7 @@ def completion( # noqa: PLR0915
  azure_ad_token_provider=azure_ad_token_provider,
  timeout=timeout,
  client=client,
+ max_retries=max_retries,
  )
  else:
  return self.acompletion(
@@ -430,6 +426,7 @@ def completion( # noqa: PLR0915
  timeout=timeout,
  client=client,
  logging_obj=logging_obj,
+ max_retries=max_retries,
  convert_tool_call_to_json_mode=json_mode,
  )
  elif "stream" in optional_params and optional_params["stream"] is True:
@@ -445,6 +442,7 @@ def completion( # noqa: PLR0915
  azure_ad_token_provider=azure_ad_token_provider,
  timeout=timeout,
  client=client,
+ max_retries=max_retries,
  )
  else:
  ## LOGGING
@@ -553,19 +551,14 @@ async def acompletion(
  dynamic_params: bool,
  model_response: ModelResponse,
  logging_obj: LiteLLMLoggingObj,
+ max_retries: int,
  azure_ad_token: Optional[str] = None,
  azure_ad_token_provider: Optional[Callable] = None,
  convert_tool_call_to_json_mode: Optional[bool] = None,
  client=None, # this is the AsyncAzureOpenAI
  ):
  response = None
  try:
- max_retries = data.pop("max_retries", 2)
- if not isinstance(max_retries, int):
- raise AzureOpenAIError(
- status_code=422, message="max retries must be an int"
- )
-
  # init AzureOpenAI Client
  azure_client_params = {
  "api_version": api_version,
@@ -671,15 +664,11 @@ def streaming(
  data: dict,
  model: str,
  timeout: Any,
+ max_retries: int,
  azure_ad_token: Optional[str] = None,
  azure_ad_token_provider: Optional[Callable] = None,
  client=None,
  ):
- max_retries = data.pop("max_retries", 2)
- if not isinstance(max_retries, int):
- raise AzureOpenAIError(
- status_code=422, message="max retries must be an int"
- )
  # init AzureOpenAI Client
  azure_client_params = {
  "api_version": api_version,
@@ -742,6 +731,7 @@ async def async_streaming(
  data: dict,
  model: str,
  timeout: Any,
+ max_retries: int,
  azure_ad_token: Optional[str] = None,
  azure_ad_token_provider: Optional[Callable] = None,
  client=None,
@@ -753,7 +743,7 @@ async def async_streaming(
  "azure_endpoint": api_base,
  "azure_deployment": model,
  "http_client": litellm.aclient_session,
- "max_retries": data.pop("max_retries", 2),
+ "max_retries": max_retries,
  "timeout": timeout,
  }
  azure_client_params = select_azure_base_url_or_endpoint(

litellm/llms/azure/completion/handler.py

@@ -131,6 +131,7 @@ def completion( # noqa: PLR0915
  timeout=timeout,
  client=client,
  logging_obj=logging_obj,
+ max_retries=max_retries,
  )
  elif "stream" in optional_params and optional_params["stream"] is True:
  return self.streaming(
@@ -236,17 +237,12 @@ async def acompletion(
  timeout: Any,
  model_response: ModelResponse,
  logging_obj: Any,
+ max_retries: int,
  azure_ad_token: Optional[str] = None,
  client=None, # this is the AsyncAzureOpenAI
  ):
  response = None
  try:
- max_retries = data.pop("max_retries", 2)
- if not isinstance(max_retries, int):
- raise AzureOpenAIError(
- status_code=422, message="max retries must be an int"
- )
-
  # init AzureOpenAI Client
  azure_client_params = {
  "api_version": api_version,

litellm/main.py

@@ -1222,6 +1222,8 @@ def completion( # type: ignore # noqa: PLR0915
 
  if extra_headers is not None:
  optional_params["extra_headers"] = extra_headers
+ if max_retries is not None:
+ optional_params["max_retries"] = max_retries
 
  if litellm.AzureOpenAIO1Config().is_o_series_model(model=model):
 

litellm/utils.py

@@ -1785,7 +1785,10 @@ def token_counter(
  for tool_call in message["tool_calls"]:
  if "function" in tool_call:
  function_arguments = tool_call["function"]["arguments"]
- text += function_arguments
+ text = (
+ text if isinstance(text, str) else "".join(text or [])
+ ) + (str(function_arguments) if function_arguments else "")
+
  else:
  raise ValueError("text and messages cannot both be None")
  elif isinstance(text, List):

tests/llm_translation/test_azure_o_series.py

@@ -152,3 +152,18 @@ def test_azure_o_series_routing():
  print(e)
  assert mock_create.call_count == 1
  assert "stream" not in mock_create.call_args.kwargs
+
+
+@patch("litellm.main.azure_o1_chat_completions._get_openai_client")
+def test_openai_o_series_max_retries_0(mock_get_openai_client):
+ import litellm
+
+ litellm.set_verbose = True
+ response = litellm.completion(
+ model="azure/o1-preview",
+ messages=[{"role": "user", "content": "hi"}],
+ max_retries=0,
+ )
+
+ mock_get_openai_client.assert_called_once()
+ assert mock_get_openai_client.call_args.kwargs["max_retries"] == 0

tests/llm_translation/test_azure_openai.py

@@ -436,3 +436,130 @@ def test_map_openai_params():
  optional_params = azure_openai_config.map_openai_params(**received_args)
  assert "tools" in optional_params
  assert len(optional_params["tools"]) > 1
+
+
+@pytest.mark.parametrize("max_retries", [0, 4])
+@pytest.mark.parametrize("stream", [True, False])
+@patch(
+ "litellm.main.azure_chat_completions.make_sync_azure_openai_chat_completion_request"
+)
+def test_azure_max_retries_0(
+ mock_make_sync_azure_openai_chat_completion_request, max_retries, stream
+):
+ from litellm import completion
+
+ try:
+ completion(
+ model="azure/gpt-4o",
+ messages=[{"role": "user", "content": "Hello world"}],
+ max_retries=max_retries,
+ stream=stream,
+ )
+ except Exception as e:
+ print(e)
+
+ mock_make_sync_azure_openai_chat_completion_request.assert_called_once()
+ assert (
+ mock_make_sync_azure_openai_chat_completion_request.call_args.kwargs[
+ "azure_client"
+ ].max_retries
+ == max_retries
+ )
+
+
+@pytest.mark.parametrize("max_retries", [0, 4])
+@pytest.mark.parametrize("stream", [True, False])
+@patch("litellm.main.azure_chat_completions.make_azure_openai_chat_completion_request")
+@pytest.mark.asyncio
+async def test_async_azure_max_retries_0(
+ make_azure_openai_chat_completion_request, max_retries, stream
+):
+ from litellm import acompletion
+
+ try:
+ await acompletion(
+ model="azure/gpt-4o",
+ messages=[{"role": "user", "content": "Hello world"}],
+ max_retries=max_retries,
+ stream=stream,
+ )
+ except Exception as e:
+ print(e)
+
+ make_azure_openai_chat_completion_request.assert_called_once()
+ assert (
+ make_azure_openai_chat_completion_request.call_args.kwargs[
+ "azure_client"
+ ].max_retries
+ == max_retries
+ )
+
+
+@pytest.mark.parametrize("max_retries", [0, 4])
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.parametrize("sync_mode", [True, False])
+@patch("litellm.llms.azure.completion.handler.select_azure_base_url_or_endpoint")
+@pytest.mark.asyncio
+async def test_azure_instruct(
+ mock_select_azure_base_url_or_endpoint, max_retries, stream, sync_mode
+):
+ from litellm import completion, acompletion
+
+ args = {
+ "model": "azure_text/instruct-model",
+ "messages": [
+ {"role": "user", "content": "What is the weather like in Boston?"}
+ ],
+ "max_tokens": 10,
+ "max_retries": max_retries,
+ }
+
+ try:
+ if sync_mode:
+ completion(**args)
+ else:
+ await acompletion(**args)
+ except Exception:
+ pass
+
+ mock_select_azure_base_url_or_endpoint.assert_called_once()
+ assert (
+ mock_select_azure_base_url_or_endpoint.call_args.kwargs["azure_client_params"][
+ "max_retries"
+ ]
+ == max_retries
+ )
+
+
+@pytest.mark.parametrize("max_retries", [0, 4])
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.parametrize("sync_mode", [True, False])
+@patch("litellm.llms.azure.azure.select_azure_base_url_or_endpoint")
+@pytest.mark.asyncio
+async def test_azure_embedding_max_retries_0(
+ mock_select_azure_base_url_or_endpoint, max_retries, stream, sync_mode
+):
+ from litellm import aembedding, embedding
+
+ args = {
+ "model": "azure/azure-embedding-model",
+ "input": "Hello world",
+ "max_retries": max_retries,
+ "stream": stream,
+ }
+
+ try:
+ if sync_mode:
+ embedding(**args)
+ else:
+ await aembedding(**args)
+ except Exception as e:
+ print(e)
+
+ mock_select_azure_base_url_or_endpoint.assert_called_once()
+ assert (
+ mock_select_azure_base_url_or_endpoint.call_args.kwargs["azure_client_params"][
+ "max_retries"
+ ]
+ == max_retries
+ )

tests/llm_translation/test_openai.py

@@ -314,3 +314,18 @@ def test_completion_bad_org():
  os.environ["OPENAI_ORGANIZATION"] = _old_org
  else:
  del os.environ["OPENAI_ORGANIZATION"]
+
+
+@patch("litellm.main.openai_chat_completions._get_openai_client")
+def test_openai_max_retries_0(mock_get_openai_client):
+ import litellm
+
+ 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()
+ assert mock_get_openai_client.call_args.kwargs["max_retries"] == 0