diff --git a/litellm/images/main.py b/litellm/images/main.py index dd84de83e2e..50c05bcc778 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -569,12 +569,8 @@ async def aimage_variation( response_format: Literal["url", "b64_json"] = "url", size: str | None = None, user: str | None = None, - # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. - extra_headers: dict[str, object] | None = None, - extra_query: dict[str, object] | None = None, - extra_body: dict[str, object] | None = None, - timeout: float | httpx.Timeout | None = None, - # LiteLLM specific params, + # LiteLLM specific params; extra_headers/extra_query/extra_body/timeout + # forward to image_variation through **kwargs like every other kwarg. custom_llm_provider: str | None = None, **kwargs, ) -> ImageResponse: @@ -593,6 +589,11 @@ async def aimage_variation( loop: Final = asyncio.get_event_loop() ### PASS ARGS TO Image Generation ### kwargs["async_call"] = True + resolved_provider: str = ( + custom_llm_provider + if custom_llm_provider is not None + else get_llm_provider(model=model, api_base=kwargs.get("api_base", None))[1] + ) try: # Use a partial function to pass your keyword arguments func: Final = partial( @@ -603,11 +604,7 @@ async def aimage_variation( response_format=response_format, size=size, user=user, - extra_headers=extra_headers, - extra_query=extra_query, - extra_body=extra_body, - timeout=timeout, - custom_llm_provider=custom_llm_provider, + custom_llm_provider=resolved_provider, **kwargs, ) @@ -615,9 +612,6 @@ async def aimage_variation( ctx: Final = contextvars.copy_context() func_with_context: Final = partial(ctx.run, func) - if custom_llm_provider is None: - _, 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, ImageResponse): ## CACHING SCENARIO @@ -631,10 +625,9 @@ async def aimage_variation( 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, + custom_llm_provider=resolved_provider or "openai", original_exception=e, completion_kwargs=local_vars, extra_kwargs=kwargs, diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 48649190ffa..aa6c9ff2fd9 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -1,13 +1,9 @@ import asyncio import io -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Sequence from itertools import chain from types import MappingProxyType -from typing import ( - Final, - cast, # noqa: TID251 # cast narrows the untyped proxy_server module globals to base_process_llm_request's declared parameter types - get_type_hints, -) +from typing import Final, get_type_hints import orjson from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, UploadFile, status @@ -443,7 +439,8 @@ async def image_variation_api( # noqa: ANN201 # response type mirrors image_ed ) form_fields: Final = coerce_numeric_form_fields( - parsed_body=cast("Mapping[str, object]", await read_request_body(request=request)), + # read_request_body's return is partially typed; the form coercion only reads mapping keys + parsed_body=await read_request_body(request=request), # pyright: ignore[reportUnknownArgumentType] # proxy_server helper's return is partially typed; the coercion only reads mapping keys numeric_fields=IMAGE_VARIATION_NUMERIC_FORM_FIELDS, ) data: Final = { @@ -469,15 +466,15 @@ async def image_variation_api( # noqa: ANN201 # response type mirrors image_ed route_type="aimage_variation", proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, - general_settings=cast("dict[str, object]", general_settings), + general_settings=general_settings, # pyright: ignore[reportUnknownArgumentType] # proxy_server global; base_process_llm_request declares dict[str, object] proxy_config=proxy_config, - select_data_generator=cast("Callable[..., object] | None", select_data_generator), + select_data_generator=select_data_generator, # pyright: ignore[reportUnknownArgumentType] # proxy_server global; the parameter is Callable[..., object] | None model=model, - user_model=cast("str | None", user_model), - user_temperature=cast("float | None", user_temperature), + user_model=user_model, # pyright: ignore[reportUnknownArgumentType] # proxy_server global; the parameter is str | None + user_temperature=user_temperature, # pyright: ignore[reportUnknownArgumentType] # proxy_server global; the parameter is float | None user_request_timeout=user_request_timeout, user_max_tokens=user_max_tokens, - user_api_base=cast("str | None", user_api_base), + user_api_base=user_api_base, # pyright: ignore[reportUnknownArgumentType] # proxy_server global; the parameter is str | None version=version, ) except Exception as e: # noqa: BLE001 # provider-agnostic error surfacing mirrors image_edit_api