diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index 026f06fb86..208c026785 100644 --- a/litellm/llms/azure.py +++ b/litellm/llms/azure.py @@ -464,6 +464,42 @@ class AzureChatCompletion(BaseLLM): import traceback raise AzureOpenAIError(status_code=500, message=traceback.format_exc()) + async def aimage_generation( + self, + data: dict, + model_response: ModelResponse, + azure_client_params: dict, + api_key: str, + input: list, + client=None, + logging_obj=None + ): + response = None + try: + if client is None: + openai_aclient = AsyncAzureOpenAI(**azure_client_params) + else: + openai_aclient = client + response = await openai_aclient.images.generate(**data) + stringified_response = response.model_dump_json() + ## LOGGING + logging_obj.post_call( + input=input, + api_key=api_key, + additional_args={"complete_input_dict": data}, + original_response=stringified_response, + ) + return convert_to_model_response_object(response_object=json.loads(stringified_response), model_response_object=model_response, response_type="embedding") + except Exception as e: + ## LOGGING + logging_obj.post_call( + input=input, + api_key=api_key, + additional_args={"complete_input_dict": data}, + original_response=str(e), + ) + raise e + def image_generation(self, prompt: str, timeout: float, @@ -492,9 +528,9 @@ class AzureChatCompletion(BaseLLM): if not isinstance(max_retries, int): raise AzureOpenAIError(status_code=422, message="max retries must be an int") - # if aembedding == True: - # response = self.aembedding(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries) # type: ignore - # return response + if aimg_generation == True: + response = self.aimage_generation(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries) # type: ignore + return response if client is None: client_session = litellm.client_session or httpx.Client(transport=CustomHTTPTransport(),) diff --git a/litellm/llms/openai.py b/litellm/llms/openai.py index 0731bd5096..e6b5352953 100644 --- a/litellm/llms/openai.py +++ b/litellm/llms/openai.py @@ -445,6 +445,43 @@ class OpenAIChatCompletion(BaseLLM): import traceback raise OpenAIError(status_code=500, message=traceback.format_exc()) + async def aimage_generation( + self, + prompt: str, + data: dict, + model_response: ModelResponse, + timeout: float, + api_key: Optional[str]=None, + api_base: Optional[str]=None, + client=None, + max_retries=None, + logging_obj=None + ): + response = None + try: + if client is None: + openai_aclient = AsyncOpenAI(api_key=api_key, base_url=api_base, http_client=litellm.aclient_session, timeout=timeout, max_retries=max_retries) + else: + openai_aclient = client + response = await openai_aclient.images.generate(**data) # type: ignore + stringified_response = response.model_dump_json() + ## LOGGING + logging_obj.post_call( + input=prompt, + api_key=api_key, + additional_args={"complete_input_dict": data}, + original_response=stringified_response, + ) + return convert_to_model_response_object(response_object=json.loads(stringified_response), model_response_object=model_response, response_type="embedding") # type: ignore + except Exception as e: + ## LOGGING + logging_obj.post_call( + input=input, + api_key=api_key, + original_response=str(e), + ) + raise e + def image_generation(self, model: Optional[str], prompt: str, diff --git a/litellm/main.py b/litellm/main.py index 318ba8ffbe..b2ed72f7fd 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2230,6 +2230,49 @@ def moderation(input: str, api_key: Optional[str]=None): return response ##### Image Generation ####################### +@client +async def aimage_generation(*args, **kwargs): + """ + Asynchronously calls the `image_generation` function with the given arguments and keyword arguments. + + Parameters: + - `args` (tuple): Positional arguments to be passed to the `embedding` function. + - `kwargs` (dict): Keyword arguments to be passed to the `embedding` function. + + Returns: + - `response` (Any): The response returned by the `embedding` function. + """ + loop = asyncio.get_event_loop() + model = args[0] if len(args) > 0 else kwargs["model"] + ### PASS ARGS TO Image Generation ### + kwargs["aimg_generation"] = True + custom_llm_provider = None + try: + # Use a partial function to pass your keyword arguments + func = partial(image_generation, *args, **kwargs) + + # Add the context to the function + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + + _, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None)) + + # Await normally + init_response = await loop.run_in_executor(None, func_with_context) + if isinstance(init_response, dict) or isinstance(init_response, ModelResponse): ## CACHING SCENARIO + response = init_response + elif asyncio.iscoroutine(init_response): + response = await init_response + else: + # Call the synchronous function using run_in_executor + response = await loop.run_in_executor(None, func_with_context) + return response + except Exception as e: + custom_llm_provider = custom_llm_provider or "openai" + raise exception_type( + model=model, custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=args, + ) + @client def image_generation(prompt: str, model: Optional[str]=None, @@ -2251,6 +2294,7 @@ def image_generation(prompt: str, Currently supports just Azure + OpenAI. """ + aimg_generation = kwargs.get("aimg_generation", False) litellm_call_id = kwargs.get("litellm_call_id", None) logger_fn = kwargs.get("logger_fn", None) proxy_server_request = kwargs.get('proxy_server_request', None) @@ -2264,7 +2308,7 @@ def image_generation(prompt: str, model = "dall-e-2" custom_llm_provider = "openai" # default to dall-e-2 on openai openai_params = ["user", "request_timeout", "api_base", "api_version", "api_key", "deployment_id", "organization", "base_url", "default_headers", "timeout", "max_retries", "n", "quality", "size", "style"] - litellm_params = ["metadata", "aembedding", "caching", "mock_response", "api_key", "api_version", "api_base", "force_timeout", "logger_fn", "verbose", "custom_llm_provider", "litellm_logging_obj", "litellm_call_id", "use_client", "id", "fallbacks", "azure", "headers", "model_list", "num_retries", "context_window_fallback_dict", "roles", "final_prompt_value", "bos_token", "eos_token", "request_timeout", "complete_response", "self", "client", "rpm", "tpm", "input_cost_per_token", "output_cost_per_token", "hf_model_name", "proxy_server_request", "model_info", "preset_cache_key", "caching_groups"] + litellm_params = ["metadata", "aimg_generation", "caching", "mock_response", "api_key", "api_version", "api_base", "force_timeout", "logger_fn", "verbose", "custom_llm_provider", "litellm_logging_obj", "litellm_call_id", "use_client", "id", "fallbacks", "azure", "headers", "model_list", "num_retries", "context_window_fallback_dict", "roles", "final_prompt_value", "bos_token", "eos_token", "request_timeout", "complete_response", "self", "client", "rpm", "tpm", "input_cost_per_token", "output_cost_per_token", "hf_model_name", "proxy_server_request", "model_info", "preset_cache_key", "caching_groups"] default_params = openai_params + litellm_params non_default_params = {k: v for k,v in kwargs.items() if k not in default_params} # model-specific params - pass them straight to the model/provider optional_params = get_optional_params_image_gen(n=n, @@ -2307,9 +2351,9 @@ def image_generation(prompt: str, get_secret("AZURE_AD_TOKEN") ) - model_response = azure_chat_completions.image_generation(model=model, prompt=prompt, timeout=timeout, api_key=api_key, api_base=api_base, logging_obj=litellm_logging_obj, optional_params=optional_params, model_response = model_response, api_version = api_version) + model_response = azure_chat_completions.image_generation(model=model, prompt=prompt, timeout=timeout, api_key=api_key, api_base=api_base, logging_obj=litellm_logging_obj, optional_params=optional_params, model_response = model_response, api_version = api_version, aimg_generation=aimage_generation) elif custom_llm_provider == "openai": - model_response = openai_chat_completions.image_generation(model=model, prompt=prompt, timeout=timeout, api_key=api_key, api_base=api_base, logging_obj=litellm_logging_obj, optional_params=optional_params, model_response = model_response) + model_response = openai_chat_completions.image_generation(model=model, prompt=prompt, timeout=timeout, api_key=api_key, api_base=api_base, logging_obj=litellm_logging_obj, optional_params=optional_params, model_response = model_response, aimg_generation=aimage_generation) return model_response