phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(handler.py): same as last commit

Browse source
This commit is contained in:
Krrish Dholakia 2025-03-11 18:48:12 -07:00
parent 23bf7b5700
commit 145cd483b9
1 changed files with 16 additions and 65 deletions
Download patch file Download diff file Expand all files Collapse all files

81
litellm/llms/azure/completion/handler.py
View file

@ -8,7 +8,7 @@ from litellm.utils import CustomStreamWrapper, ModelResponse, TextCompletionResp
from ...base import BaseLLM from ...base import BaseLLM
from ...openai.completion.transformation import OpenAITextCompletionConfig from ...openai.completion.transformation import OpenAITextCompletionConfig
from ..common_utils import AzureOpenAIError from ..common_utils import AzureOpenAIError, BaseAzureLLM
openai_text_completion_config = OpenAITextCompletionConfig() openai_text_completion_config = OpenAITextCompletionConfig()
@ -25,7 +25,7 @@ def select_azure_base_url_or_endpoint(azure_client_params: dict):
return azure_client_params return azure_client_params
class AzureTextCompletion(BaseLLM): class AzureTextCompletion(BaseAzureLLM):
def __init__(self) -> None: def __init__(self) -> None:
super().__init__() super().__init__()
@ -60,7 +60,6 @@ class AzureTextCompletion(BaseLLM):
headers: Optional[dict] = None, headers: Optional[dict] = None,
client=None, client=None,
): ):
super().completion()
try: try:
if model is None or messages is None: if model is None or messages is None:
raise AzureOpenAIError( raise AzureOpenAIError(
@ -72,6 +71,14 @@ class AzureTextCompletion(BaseLLM):
messages=messages, model=model, custom_llm_provider="azure_text" messages=messages, model=model, custom_llm_provider="azure_text"
) )
azure_client_params = self.initialize_azure_sdk_client(
litellm_params=litellm_params or {},
api_key=api_key,
model_name=model,
api_version=api_version,
api_base=api_base,
)
### CHECK IF CLOUDFLARE AI GATEWAY ### ### CHECK IF CLOUDFLARE AI GATEWAY ###
### if so - set the model as part of the base url ### if so - set the model as part of the base url
if "gateway.ai.cloudflare.com" in api_base: if "gateway.ai.cloudflare.com" in api_base:
@ -118,6 +125,7 @@ class AzureTextCompletion(BaseLLM):
azure_ad_token=azure_ad_token, azure_ad_token=azure_ad_token,
timeout=timeout, timeout=timeout,
client=client, client=client,
azure_client_params=azure_client_params,
) )
else: else:
return self.acompletion( return self.acompletion(
@ -132,6 +140,7 @@ class AzureTextCompletion(BaseLLM):
client=client, client=client,
logging_obj=logging_obj, logging_obj=logging_obj,
max_retries=max_retries, max_retries=max_retries,
azure_client_params=azure_client_params,
) )
elif "stream" in optional_params and optional_params["stream"] is True: elif "stream" in optional_params and optional_params["stream"] is True:
return self.streaming( return self.streaming(
@ -144,6 +153,7 @@ class AzureTextCompletion(BaseLLM):
azure_ad_token=azure_ad_token, azure_ad_token=azure_ad_token,
timeout=timeout, timeout=timeout,
client=client, client=client,
azure_client_params=azure_client_params,
) )
else: else:
## LOGGING ## LOGGING
@ -165,22 +175,6 @@ class AzureTextCompletion(BaseLLM):
status_code=422, message="max retries must be an int" status_code=422, message="max retries must be an int"
) )
# init AzureOpenAI Client # init AzureOpenAI Client
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_ad_token_provider": azure_ad_token_provider,
}
azure_client_params = select_azure_base_url_or_endpoint(
azure_client_params=azure_client_params
)
if api_key is not None:
azure_client_params["api_key"] = api_key
elif azure_ad_token is not None:
azure_client_params["azure_ad_token"] = azure_ad_token
if client is None: if client is None:
azure_client = AzureOpenAI(**azure_client_params) azure_client = AzureOpenAI(**azure_client_params)
else: else:
@ -240,26 +234,11 @@ class AzureTextCompletion(BaseLLM):
max_retries: int, max_retries: int,
azure_ad_token: Optional[str] = None, azure_ad_token: Optional[str] = None,
client=None, # this is the AsyncAzureOpenAI client=None, # this is the AsyncAzureOpenAI
azure_client_params: dict = {},
): ):
response = None response = None
try: try:
# init AzureOpenAI Client # init AzureOpenAI Client
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_client_params = select_azure_base_url_or_endpoint(
azure_client_params=azure_client_params
)
if api_key is not None:
azure_client_params["api_key"] = api_key
elif azure_ad_token is not None:
azure_client_params["azure_ad_token"] = azure_ad_token
# setting Azure client # setting Azure client
if client is None: if client is None:
azure_client = AsyncAzureOpenAI(**azure_client_params) azure_client = AsyncAzureOpenAI(**azure_client_params)
@ -312,6 +291,7 @@ class AzureTextCompletion(BaseLLM):
timeout: Any, timeout: Any,
azure_ad_token: Optional[str] = None, azure_ad_token: Optional[str] = None,
client=None, client=None,
azure_client_params: dict = {},
): ):
max_retries = data.pop("max_retries", 2) max_retries = data.pop("max_retries", 2)
if not isinstance(max_retries, int): if not isinstance(max_retries, int):
@ -319,21 +299,6 @@ class AzureTextCompletion(BaseLLM):
status_code=422, message="max retries must be an int" status_code=422, message="max retries must be an int"
) )
# init AzureOpenAI Client # init AzureOpenAI Client
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_client_params = select_azure_base_url_or_endpoint(
azure_client_params=azure_client_params
)
if api_key is not None:
azure_client_params["api_key"] = api_key
elif azure_ad_token is not None:
azure_client_params["azure_ad_token"] = azure_ad_token
if client is None: if client is None:
azure_client = AzureOpenAI(**azure_client_params) azure_client = AzureOpenAI(**azure_client_params)
else: else:
@ -375,24 +340,10 @@ class AzureTextCompletion(BaseLLM):
timeout: Any, timeout: Any,
azure_ad_token: Optional[str] = None, azure_ad_token: Optional[str] = None,
client=None, client=None,
azure_client_params: dict = {},
): ):
try: try:
# init AzureOpenAI Client # init AzureOpenAI Client
azure_client_params = {
"api_version": api_version,
"azure_endpoint": api_base,
"azure_deployment": model,
"http_client": litellm.client_session,
"max_retries": data.pop("max_retries", 2),
"timeout": timeout,
}
azure_client_params = select_azure_base_url_or_endpoint(
azure_client_params=azure_client_params
)
if api_key is not None:
azure_client_params["api_key"] = api_key
elif azure_ad_token is not None:
azure_client_params["azure_ad_token"] = azure_ad_token
if client is None: if client is None:
azure_client = AsyncAzureOpenAI(**azure_client_params) azure_client = AsyncAzureOpenAI(**azure_client_params)
else: else: