"""
OpenAI Image Variations Handler
"""
from typing import Callable, Optional
import httpx
from openai import AsyncOpenAI, OpenAI
import litellm
from litellm.types.utils import FileTypes, ImageResponse, LlmProviders
from litellm.utils import ProviderConfigManager
from ...base_llm.image_variations.transformation import BaseImageVariationConfig
from ...custom_httpx.llm_http_handler import LiteLLMLoggingObj
from ..common_utils import OpenAIError
class OpenAIImageVariationsHandler:
def get_sync_client(
self,
client: Optional[OpenAI],
init_client_params: dict,
):
if client is None:
openai_client = OpenAI(
**init_client_params,
)
else:
openai_client = client
return openai_client
def get_async_client(
self, client: Optional[AsyncOpenAI], init_client_params: dict
) -> AsyncOpenAI:
if client is None:
openai_client = AsyncOpenAI(
**init_client_params,
)
else:
openai_client = client
return openai_client
async def async_image_variations(
self,
api_key: str,
api_base: str,
organization: Optional[str],
client: Optional[AsyncOpenAI],
data: dict,
headers: dict,
model: Optional[str],
timeout: float,
max_retries: int,
logging_obj: LiteLLMLoggingObj,
model_response: ImageResponse,
optional_params: dict,
litellm_params: dict,
image: FileTypes,
provider_config: BaseImageVariationConfig,
) -> ImageResponse:
try:
init_client_params = {
"api_key": api_key,
"base_url": api_base,
"http_client": litellm.client_session,
"timeout": timeout,
"max_retries": max_retries, # type: ignore
"organization": organization,
}
client = self.get_async_client(
client=client, init_client_params=init_client_params
)
raw_response = await client.images.with_raw_response.create_variation(**data) # type: ignore
response = raw_response.parse()
response_json = response.model_dump()
## LOGGING
logging_obj.post_call(
api_key=api_key,
original_response=response_json,
additional_args={
"headers": headers,
"api_base": api_base,
},
)
## RESPONSE OBJECT
return provider_config.transform_response_image_variation(
model=model,
model_response=ImageResponse(**response_json),
raw_response=httpx.Response(
status_code=200,
request=httpx.Request(
method="GET", url="https://litellm.ai"
), # mock request object
),
logging_obj=logging_obj,
request_data=data,
image=image,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=None,
api_key=api_key,
)
except Exception as e:
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)
raise OpenAIError(
status_code=status_code, message=error_text, headers=error_headers
)
def image_variations(
self,
model_response: ImageResponse,
api_key: str,
api_base: 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,
logger_fn=None,
client=None,
organization: Optional[str] = None,
headers: Optional[dict] = None,
) -> ImageResponse:
try:
provider_config = ProviderConfigManager.get_provider_image_variation_config(
model=model or "", # openai defaults to dall-e-2
provider=LlmProviders.OPENAI,
)
if provider_config is None:
raise ValueError(
f"image variation provider not found: {custom_llm_provider}."
)
max_retries = optional_params.pop("max_retries", 2)
data = provider_config.transform_request_image_variation(
model=model,
image=image,
optional_params=optional_params,
headers=headers or {},
)
json_data = data.get("data")
if not json_data:
raise ValueError(
f"data field is required, for openai image variations. Got={data}"
)
## LOGGING
logging_obj.pre_call(
input="",
api_key=api_key,
additional_args={
"headers": headers,
"api_base": api_base,
"complete_input_dict": data,
},
)
if litellm_params.get("async_call", False):
return self.async_image_variations(
api_base=api_base,
data=json_data,
headers=headers or {},
model_response=model_response,
api_key=api_key,
logging_obj=logging_obj,
model=model,
timeout=timeout,
max_retries=max_retries,
organization=organization,
client=client,
provider_config=provider_config,
image=image,
optional_params=optional_params,
litellm_params=litellm_params,
) # type: ignore
init_client_params = {
"api_key": api_key,
"base_url": api_base,
"http_client": litellm.client_session,
"timeout": timeout,
"max_retries": max_retries, # type: ignore
"organization": organization,
}
client = self.get_sync_client(
client=client, init_client_params=init_client_params
)
raw_response = client.images.with_raw_response.create_variation(**json_data) # type: ignore
response = raw_response.parse()
response_json = response.model_dump()
## LOGGING
logging_obj.post_call(
api_key=api_key,
original_response=response_json,
additional_args={
"headers": headers,
"api_base": api_base,
},
)
## RESPONSE OBJECT
return provider_config.transform_response_image_variation(
model=model,
model_response=ImageResponse(**response_json),
raw_response=httpx.Response(
status_code=200,
request=httpx.Request(
method="GET", url="https://litellm.ai"
), # mock request object
),
logging_obj=logging_obj,
request_data=json_data,
image=image,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=None,
api_key=api_key,
)
except Exception as e:
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)
raise OpenAIError(
status_code=status_code, message=error_text, headers=error_headers
)
