from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Union, cast
import aiohttp
import httpx # type: ignore
from aiohttp import ClientSession, FormData
import litellm
import litellm.litellm_core_utils
import litellm.types
import litellm.types.utils
from litellm.llms.base_llm.chat.transformation import BaseConfig
from litellm.llms.base_llm.image_variations.transformation import (
BaseImageVariationConfig,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
_get_httpx_client,
)
from litellm.types.llms.openai import FileTypes
from litellm.types.utils import HttpHandlerRequestFields, ImageResponse, LlmProviders
from litellm.utils import CustomStreamWrapper, ModelResponse, ProviderConfigManager
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
DEFAULT_TIMEOUT = 600
class BaseLLMAIOHTTPHandler:
def __init__(self):
self.client_session: Optional[aiohttp.ClientSession] = None
def _get_async_client_session(
self, dynamic_client_session: Optional[ClientSession] = None
) -> ClientSession:
if dynamic_client_session:
return dynamic_client_session
elif self.client_session:
return self.client_session
else:
# init client session, and then return new session
self.client_session = aiohttp.ClientSession()
return self.client_session
async def _make_common_async_call(
self,
async_client_session: Optional[ClientSession],
provider_config: BaseConfig,
api_base: str,
headers: dict,
data: Optional[dict],
timeout: Union[float, httpx.Timeout],
litellm_params: dict,
form_data: Optional[FormData] = None,
stream: bool = False,
) -> aiohttp.ClientResponse:
"""Common implementation across stream + non-stream calls. Meant to ensure consistent error-handling."""
max_retry_on_unprocessable_entity_error = (
provider_config.max_retry_on_unprocessable_entity_error
)
response: Optional[aiohttp.ClientResponse] = None
async_client_session = self._get_async_client_session(
dynamic_client_session=async_client_session
)
for i in range(max(max_retry_on_unprocessable_entity_error, 1)):
try:
response = await async_client_session.post(
url=api_base,
headers=headers,
json=data,
data=form_data,
)
if not response.ok:
response.raise_for_status()
except aiohttp.ClientResponseError as e:
setattr(e, "text", e.message)
raise self._handle_error(e=e, provider_config=provider_config)
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
break
if response is None:
raise provider_config.get_error_class(
error_message="No response from the API",
status_code=422,
headers={},
)
return response
def _make_common_sync_call(
self,
sync_httpx_client: HTTPHandler,
provider_config: BaseConfig,
api_base: str,
headers: dict,
data: dict,
timeout: Union[float, httpx.Timeout],
litellm_params: dict,
stream: bool = False,
files: Optional[dict] = None,
content: Any = None,
params: Optional[dict] = None,
) -> httpx.Response:
max_retry_on_unprocessable_entity_error = (
provider_config.max_retry_on_unprocessable_entity_error
)
response: Optional[httpx.Response] = None
for i in range(max(max_retry_on_unprocessable_entity_error, 1)):
try:
response = sync_httpx_client.post(
url=api_base,
headers=headers,
data=data, # do not json dump the data here. let the individual endpoint handle this.
timeout=timeout,
stream=stream,
files=files,
content=content,
params=params,
)
except httpx.HTTPStatusError as e:
hit_max_retry = i + 1 == max_retry_on_unprocessable_entity_error
should_retry = provider_config.should_retry_llm_api_inside_llm_translation_on_http_error(
e=e, litellm_params=litellm_params
)
if should_retry and not hit_max_retry:
data = (
provider_config.transform_request_on_unprocessable_entity_error(
e=e, request_data=data
)
)
continue
else:
raise self._handle_error(e=e, provider_config=provider_config)
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
break
if response is None:
raise provider_config.get_error_class(
error_message="No response from the API",
status_code=422, # don't retry on this error
headers={},
)
return response
async def async_completion(
self,
custom_llm_provider: str,
provider_config: BaseConfig,
api_base: str,
headers: dict,
data: dict,
timeout: Union[float, httpx.Timeout],
model: str,
model_response: ModelResponse,
logging_obj: LiteLLMLoggingObj,
messages: list,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
client: Optional[ClientSession] = None,
):
_response = await self._make_common_async_call(
async_client_session=client,
provider_config=provider_config,
api_base=api_base,
headers=headers,
data=data,
timeout=timeout,
litellm_params=litellm_params,
stream=False,
)
_transformed_response = await provider_config.transform_response( # type: ignore
model=model,
raw_response=_response, # type: ignore
model_response=model_response,
logging_obj=logging_obj,
api_key=api_key,
request_data=data,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
)
return _transformed_response
def completion(
self,
model: str,
messages: list,
api_base: str,
custom_llm_provider: str,
model_response: ModelResponse,
encoding,
logging_obj: LiteLLMLoggingObj,
optional_params: dict,
timeout: Union[float, httpx.Timeout],
litellm_params: dict,
acompletion: bool,
stream: Optional[bool] = False,
fake_stream: bool = False,
api_key: Optional[str] = None,
headers: Optional[dict] = {},
client: Optional[Union[HTTPHandler, AsyncHTTPHandler, ClientSession]] = None,
):
provider_config = ProviderConfigManager.get_provider_chat_config(
model=model, provider=litellm.LlmProviders(custom_llm_provider)
)
# get config from model, custom llm provider
headers = provider_config.validate_environment(
api_key=api_key,
headers=headers or {},
model=model,
messages=messages,
optional_params=optional_params,
api_base=api_base,
)
api_base = provider_config.get_complete_url(
api_base=api_base,
model=model,
optional_params=optional_params,
litellm_params=litellm_params,
stream=stream,
)
data = provider_config.transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
## LOGGING
logging_obj.pre_call(
input=messages,
api_key=api_key,
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": headers,
},
)
if acompletion is True:
return self.async_completion(
custom_llm_provider=custom_llm_provider,
provider_config=provider_config,
api_base=api_base,
headers=headers,
data=data,
timeout=timeout,
model=model,
model_response=model_response,
logging_obj=logging_obj,
api_key=api_key,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
client=(
client
if client is not None and isinstance(client, ClientSession)
else None
),
)
if stream is True:
if fake_stream is not True:
data["stream"] = stream
completion_stream, headers = self.make_sync_call(
provider_config=provider_config,
api_base=api_base,
headers=headers, # type: ignore
data=data,
model=model,
messages=messages,
logging_obj=logging_obj,
timeout=timeout,
fake_stream=fake_stream,
client=(
client
if client is not None and isinstance(client, HTTPHandler)
else None
),
litellm_params=litellm_params,
)
return CustomStreamWrapper(
completion_stream=completion_stream,
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client()
else:
sync_httpx_client = client
response = self._make_common_sync_call(
sync_httpx_client=sync_httpx_client,
provider_config=provider_config,
api_base=api_base,
headers=headers,
timeout=timeout,
litellm_params=litellm_params,
data=data,
)
return provider_config.transform_response(
model=model,
raw_response=response,
model_response=model_response,
logging_obj=logging_obj,
api_key=api_key,
request_data=data,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
)
def make_sync_call(
self,
provider_config: BaseConfig,
api_base: str,
headers: dict,
data: dict,
model: str,
messages: list,
logging_obj,
litellm_params: dict,
timeout: Union[float, httpx.Timeout],
fake_stream: bool = False,
client: Optional[HTTPHandler] = None,
) -> Tuple[Any, dict]:
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client()
else:
sync_httpx_client = client
stream = True
if fake_stream is True:
stream = False
response = self._make_common_sync_call(
sync_httpx_client=sync_httpx_client,
provider_config=provider_config,
api_base=api_base,
headers=headers,
data=data,
timeout=timeout,
litellm_params=litellm_params,
stream=stream,
)
if fake_stream is True:
completion_stream = provider_config.get_model_response_iterator(
streaming_response=response.json(), sync_stream=True
)
else:
completion_stream = provider_config.get_model_response_iterator(
streaming_response=response.iter_lines(), sync_stream=True
)
# LOGGING
logging_obj.post_call(
input=messages,
api_key="",
original_response="first stream response received",
additional_args={"complete_input_dict": data},
)
return completion_stream, dict(response.headers)
async def async_image_variations(
self,
client: Optional[ClientSession],
provider_config: BaseImageVariationConfig,
api_base: str,
headers: dict,
data: HttpHandlerRequestFields,
timeout: float,
litellm_params: dict,
model_response: ImageResponse,
logging_obj: LiteLLMLoggingObj,
api_key: str,
model: Optional[str],
image: FileTypes,
optional_params: dict,
) -> ImageResponse:
# create aiohttp form data if files in data
form_data: Optional[FormData] = None
if "files" in data and "data" in data:
form_data = FormData()
for k, v in data["files"].items():
form_data.add_field(k, v[1], filename=v[0], content_type=v[2])
for key, value in data["data"].items():
form_data.add_field(key, value)
_response = await self._make_common_async_call(
async_client_session=client,
provider_config=provider_config,
api_base=api_base,
headers=headers,
data=None if form_data is not None else cast(dict, data),
form_data=form_data,
timeout=timeout,
litellm_params=litellm_params,
stream=False,
)
## LOGGING
logging_obj.post_call(
api_key=api_key,
original_response=_response.text,
additional_args={
"headers": headers,
"api_base": api_base,
},
)
## RESPONSE OBJECT
return await provider_config.async_transform_response_image_variation(
model=model,
model_response=model_response,
raw_response=_response,
logging_obj=logging_obj,
request_data=cast(dict, data),
image=image,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=None,
api_key=api_key,
)
def image_variations(
self,
model_response: ImageResponse,
api_key: str,
model: Optional[str],
image: FileTypes,
timeout: float,
custom_llm_provider: str,
logging_obj: LiteLLMLoggingObj,
optional_params: dict,
litellm_params: dict,
print_verbose: Optional[Callable] = None,
api_base: Optional[str] = None,
aimage_variation: bool = False,
logger_fn=None,
client=None,
organization: Optional[str] = None,
headers: Optional[dict] = None,
) -> ImageResponse:
if model is None:
raise ValueError("model is required for non-openai image variations")
provider_config = ProviderConfigManager.get_provider_image_variation_config(
model=model, # openai defaults to dall-e-2
provider=LlmProviders(custom_llm_provider),
)
if provider_config is None:
raise ValueError(
f"image variation provider not found: {custom_llm_provider}."
)
api_base = provider_config.get_complete_url(
api_base=api_base,
model=model,
optional_params=optional_params,
litellm_params=litellm_params,
stream=False,
)
headers = provider_config.validate_environment(
api_key=api_key,
headers=headers or {},
model=model,
messages=[{"role": "user", "content": "test"}],
optional_params=optional_params,
api_base=api_base,
)
data = provider_config.transform_request_image_variation(
model=model,
image=image,
optional_params=optional_params,
headers=headers,
)
## LOGGING
logging_obj.pre_call(
input="",
api_key=api_key,
additional_args={
"headers": headers,
"api_base": api_base,
"complete_input_dict": data.copy(),
},
)
if litellm_params.get("async_call", False):
return self.async_image_variations(
api_base=api_base,
data=data,
headers=headers,
model_response=model_response,
api_key=api_key,
logging_obj=logging_obj,
model=model,
timeout=timeout,
client=client,
optional_params=optional_params,
litellm_params=litellm_params,
image=image,
provider_config=provider_config,
) # type: ignore
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client()
else:
sync_httpx_client = client
response = self._make_common_sync_call(
sync_httpx_client=sync_httpx_client,
provider_config=provider_config,
api_base=api_base,
headers=headers,
timeout=timeout,
litellm_params=litellm_params,
stream=False,
data=data.get("data") or {},
files=data.get("files"),
content=data.get("content"),
params=data.get("params"),
)
## LOGGING
logging_obj.post_call(
api_key=api_key,
original_response=response.text,
additional_args={
"headers": headers,
"api_base": api_base,
},
)
## RESPONSE OBJECT
return provider_config.transform_response_image_variation(
model=model,
model_response=model_response,
raw_response=response,
logging_obj=logging_obj,
request_data=cast(dict, data),
image=image,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=None,
api_key=api_key,
)
def _handle_error(self, e: Exception, provider_config: BaseConfig):
status_code = getattr(e, "status_code", 500)
error_headers = getattr(e, "headers", None)
error_text = getattr(e, "text", str(e))
error_response = getattr(e, "response", None)
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
if error_response and hasattr(error_response, "text"):
error_text = getattr(error_response, "text", error_text)
if error_headers:
error_headers = dict(error_headers)
else:
error_headers = {}
raise provider_config.get_error_class(
error_message=error_text,
status_code=status_code,
headers=error_headers,
)
