diff --git a/litellm/__init__.py b/litellm/__init__.py index acae4fc70e6..98c36dc582a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1614,6 +1614,7 @@ from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_k ### PASSTHROUGH ### from .passthrough import allm_passthrough_route, llm_passthrough_route from .google_genai import agenerate_content +from .google_genai import generate_content as generate_content # pyright: ignore[reportUnknownVariableType] # native request types require the optional google-genai SDK ### GLOBAL CONFIG ### global_bitbucket_config: Optional[Mapping[str, object]] = None diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index db090c9b007..94d6c0ce344 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -8,12 +8,13 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.google_genai.adapters import GenerateContentCompletionKwargs from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import ModelResponse +from litellm.types.utils import CustomPricingLiteLLMParams, ModelResponse from .transformation import GoogleGenAIAdapter # Initialize adapter GOOGLE_GENAI_ADAPTER: Final = GoogleGenAIAdapter() +CUSTOM_PRICING_PARAMS: Final = frozenset(CustomPricingLiteLLMParams.model_fields) class _TranslatedCompletionKwargs(TypedDict, total=False, extra_items=object): @@ -48,6 +49,7 @@ class GenerateContentToCompletionHandler: stream: bool = False, litellm_params: GenericLiteLLMParams | None = None, extra_kwargs: Mapping[str, object] | None = None, + custom_logging_obj: LiteLLMLoggingObj | None = None, ) -> GenerateContentCompletionKwargs: """Prepare kwargs for litellm.completion/acompletion""" @@ -75,6 +77,13 @@ class GenerateContentToCompletionHandler: if "litellm_metadata" in extra_kwargs: completion_kwargs["litellm_metadata"] = extra_kwargs["litellm_metadata"] + if isinstance(custom_logging_obj, LiteLLMLoggingObj): + completion_kwargs["litellm_logging_obj"] = custom_logging_obj + completion_kwargs["litellm_call_id"] = custom_logging_obj.litellm_call_id + completion_kwargs.update( + {key: value for key, value in (extra_kwargs or {}).items() if key in CUSTOM_PRICING_PARAMS} + ) + if stream: completion_kwargs["stream"] = stream stream_options: Final = _adapter_stream_options((extra_kwargs or {}).get("stream_options")) @@ -90,6 +99,7 @@ class GenerateContentToCompletionHandler: litellm_params: GenericLiteLLMParams, config: dict[str, object] | None = None, stream: bool = False, + custom_logging_obj: LiteLLMLoggingObj | None = None, **kwargs: object, ) -> dict[str, object] | AsyncIterator[bytes]: """Handle generate_content call asynchronously using completion adapter""" @@ -101,6 +111,7 @@ class GenerateContentToCompletionHandler: stream=stream, litellm_params=litellm_params, extra_kwargs=kwargs, + custom_logging_obj=custom_logging_obj, ) try: @@ -141,6 +152,7 @@ class GenerateContentToCompletionHandler: config: dict[str, object] | None = None, stream: bool = False, _is_async: bool = False, + custom_logging_obj: LiteLLMLoggingObj | None = None, **kwargs: object, ) -> dict[str, object] | AsyncIterator[bytes] | Coroutine[None, None, dict[str, object] | AsyncIterator[bytes]]: """Handle generate_content call using completion adapter""" @@ -152,6 +164,7 @@ class GenerateContentToCompletionHandler: config=config, stream=stream, litellm_params=litellm_params, + custom_logging_obj=custom_logging_obj, **kwargs, ) @@ -162,6 +175,7 @@ class GenerateContentToCompletionHandler: stream=stream, litellm_params=litellm_params, extra_kwargs=kwargs, + custom_logging_obj=custom_logging_obj, ) try: diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index c1895376675..2ea2a849b70 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -115,7 +115,13 @@ _RESPONSE_MIME_TYPE_KEYS: Final = ("responseMimeType", "response_mime_type") _RESPONSE_SCHEMA_KEYS: Final = ("responseJsonSchema", "response_json_schema", "responseSchema", "response_schema") _TOOL_PARAMETERS_KEYS: Final = ("parametersJsonSchema", "parameters") _JSON_MIME_TYPE: Final = "application/json" +_TEXT_MIME_TYPE: Final = "text/plain" +_ADAPTER_SCALAR_CONFIG_KEYS: Final = frozenset({"temperature", "maxOutputTokens", "topP", "stopSequences"}) +_RESPONSE_FORMAT_CONFIG_KEYS: Final = frozenset({*_RESPONSE_MIME_TYPE_KEYS, *_RESPONSE_SCHEMA_KEYS}) _GEMINI_ONLY_SCHEMA_KEYS: Final = frozenset({"propertyOrdering", "property_ordering"}) +_ADAPTER_DEFAULT_GENERATION_CONFIG: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"responseModalities": ["TEXT"], "candidateCount": 1} +) _CONFIG_FIELDS: Final = TypeAdapter(Mapping[str, object]) _JSON_OBJECT_SCHEMA: Final = TypeAdapter(dict[str, JsonValue]) _Validated: Final = TypeVar("_Validated") @@ -127,10 +133,29 @@ _GENAI_FUNCTION_CALL: Final = TypeAdapter(_GenAIRequestFunctionCall) _GenAIMessagePart: TypeAlias = ChatCompletionTextObject | ChatCompletionImageObject | ChatCompletionFileObject +def adapter_untranslated_generation_config(config: Mapping[str, object]) -> tuple[str, ...]: + keeps_response_format: Final = _asks_for_plain_text(config) or _translate_response_format(config) is not None + return tuple(sorted(key for key, value in config.items() if not _adapter_honors(key, value, keeps_response_format))) + + +def _adapter_honors(key: str, value: object, keeps_response_format: bool) -> bool: + if key in _RESPONSE_FORMAT_CONFIG_KEYS: + return keeps_response_format + return key in _ADAPTER_SCALAR_CONFIG_KEYS or ( + key in _ADAPTER_DEFAULT_GENERATION_CONFIG and _ADAPTER_DEFAULT_GENERATION_CONFIG[key] == value + ) + + def _first_present(config: Mapping[str, object], keys: Sequence[str]) -> object | None: return next((config[key] for key in keys if config.get(key) is not None), None) +def _asks_for_plain_text(config: Mapping[str, object]) -> bool: + return _first_present(config, _RESPONSE_SCHEMA_KEYS) is None and _first_present( + config, _RESPONSE_MIME_TYPE_KEYS + ) in (None, _TEXT_MIME_TYPE) + + def _validated(adapter: TypeAdapter[_Validated], value: object) -> _Validated | None: try: return adapter.validate_python(value) diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index 71000fec81a..128a7b263e3 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -1,25 +1,33 @@ import asyncio +import base64 import contextvars +from collections.abc import Awaitable, Mapping, Sequence from functools import partial from typing import TYPE_CHECKING, Any, ClassVar, Final import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, JsonValue, TypeAdapter +from pydantic_core import to_jsonable_python import litellm -from litellm.constants import request_timeout +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH, request_timeout # Import the adapter for fallback to completion format -from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler +from litellm.google_genai.adapters.handler import CUSTOM_PRICING_PARAMS, GenerateContentToCompletionHandler +from litellm.google_genai.adapters.transformation import adapter_untranslated_generation_config from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.per_call_pricing import per_call_model_info from litellm.llms.base_llm.google_genai.transformation import ( + GENERATE_CONTENT_REQUEST_TOP_LEVEL_FIELDS, BaseGoogleGenAIGenerateContentConfig, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.custom_llm import CustomLLM +from litellm.llms.vertex_ai.gemini.transformation import snake_to_camel from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import CallTypes +from litellm.types.utils import CallTypes, LlmProviders from litellm.utils import ProviderConfigManager, client if TYPE_CHECKING: @@ -40,10 +48,28 @@ else: # Initialize any necessary instances or variables here base_llm_http_handler = BaseLLMHTTPHandler() _GENERATE_CONTENT_RESULT: Final = TypeAdapter(object) +_CUSTOM_GENERATE_CONTENT_BODY: Final = TypeAdapter(dict[str, object]) +_CUSTOM_NATIVE_REQUEST_FIELDS: Final = ( + *GENERATE_CONTENT_REQUEST_TOP_LEVEL_FIELDS, + "systemInstruction", + "system_instruction", +) +_BUILT_IN_PROVIDERS: Final = frozenset(provider.value for provider in LlmProviders) +_CALLER_DATA_FIELDS: Final = frozenset({"args", "response"}) +_SDK_ONLY_RESPONSE_FIELDS: Final = frozenset({"sdkHttpResponse", "automaticFunctionCallingHistory", "parsed"}) +_STR_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECT_SEQUENCE: Final = TypeAdapter(Sequence[object]) +_JSON_VALUE: Final = TypeAdapter[JsonValue](JsonValue) _GENERATE_CONTENT_PAYLOAD: Final = TypeAdapter(object) ################################################# +def _custom_generate_content_handler(provider: str) -> CustomLLM | None: + if provider in _BUILT_IN_PROVIDERS: + return None + return ProviderConfigManager.get_custom_llm_handler(provider) + + def _mark_async_entrypoint(logging_obj: LiteLLMLoggingObj | None, marker: str, is_async: bool) -> None: if logging_obj is not None: logging_obj.model_call_details.setdefault("litellm_params", {})[marker] = is_async @@ -142,6 +168,7 @@ class GenerateContentHelper: if litellm_params.custom_llm_provider is None: litellm_params.custom_llm_provider = custom_llm_provider + custom_handler: Final = _custom_generate_content_handler(custom_llm_provider) # get provider config generate_content_provider_config: Final[BaseGoogleGenAIGenerateContentConfig | None] = ( ProviderConfigManager.get_provider_google_genai_generate_content_config( @@ -149,6 +176,8 @@ class GenerateContentHelper: provider=litellm.LlmProviders(custom_llm_provider), litellm_params=litellm_params.model_dump(include={"base_model", "api_base"}), ) + if custom_handler is None + else None ) if generate_content_provider_config is None: @@ -156,6 +185,21 @@ class GenerateContentHelper: # Signal that we should use the adapter by returning special result if litellm_logging_obj is None: raise ValueError("litellm_logging_obj is required, but got None") + if custom_handler is not None: + request_kwargs: Final = _STR_KEYED_MAPPING.validate_python(kwargs) + litellm_logging_obj.update_from_kwargs( # pyright: ignore[reportUnknownMemberType] # legacy logging interface is untyped + kwargs=kwargs, + model=model, + optional_params={"generationConfig": _CUSTOM_GENERATE_CONTENT_BODY.validate_python(config or {})}, + litellm_params={ + "litellm_call_id": litellm_call_id, + **{key: value for key, value in request_kwargs.items() if key in CUSTOM_PRICING_PARAMS}, + }, + custom_llm_provider=custom_llm_provider, + ) + litellm_logging_obj.set_per_call_pricing( # pyright: ignore[reportUnknownMemberType] # ModelInfo has untyped fields + per_call_model_info(model, custom_llm_provider, request_kwargs, "chat") + ) return GenerateContentSetupResult( model=model, custom_llm_provider=custom_llm_provider, @@ -237,6 +281,107 @@ def _merge_native_request_fields( return {**native_request_fields, **(extra_body or {})} +def _custom_adapter_logging_obj(setup_result: GenerateContentSetupResult) -> LiteLLMLoggingObj | None: + if _custom_generate_content_handler(setup_result.custom_llm_provider) is None: + return None + return setup_result.litellm_logging_obj + + +def _require_custom_stream_hook(setup_result: GenerateContentSetupResult, is_async: bool) -> None: + custom_handler: Final = _custom_generate_content_handler(setup_result.custom_llm_provider) + if custom_handler is None: + return + hook_name: Final = "astreaming" if is_async else "streaming" + if getattr(type(custom_handler), hook_name) is getattr(CustomLLM, hook_name): + raise litellm.APIConnectionError( + message=f"{setup_result.custom_llm_provider} does not implement {hook_name}, so it cannot stream generateContent", + llm_provider=setup_result.custom_llm_provider, + model=setup_result.model, + ) + + +def _require_adapter_generation_config(setup_result: GenerateContentSetupResult) -> None: + if _custom_generate_content_handler(setup_result.custom_llm_provider) is None: + return + if setup_result.litellm_params.drop_params or litellm.drop_params: + return + untranslated: Final = adapter_untranslated_generation_config(setup_result.generate_content_config_dict) + if untranslated: + raise litellm.UnsupportedParamsError( + message=( + f"{setup_result.custom_llm_provider} has no generateContent hook, and the completion adapter cannot " + f"send generationConfig {list(untranslated)}. Set drop_params=True to drop them." + ), + llm_provider=setup_result.custom_llm_provider, + model=setup_result.model, + ) + + +def _plain_json_value(value: object) -> JsonValue: + return _JSON_VALUE.validate_python(to_jsonable_python(value, fallback=str)) + + +def _gemini_wire_value(value: object, camel_case_keys: bool, depth: int = 0) -> JsonValue: + if depth > DEFAULT_MAX_RECURSE_DEPTH: + return _plain_json_value(value) + match value: + case bytes(): + return base64.b64encode(value).decode() + case str(): + return value + case Mapping(): + return { + (snake_to_camel(key) if camel_case_keys else key): _gemini_wire_value( + item, camel_case_keys and key not in _CALLER_DATA_FIELDS, depth + 1 + ) + for key, item in _STR_KEYED_MAPPING.validate_python(value).items() + if item is not None or not camel_case_keys + } + case Sequence(): + return [ + _gemini_wire_value(item, camel_case_keys, depth + 1) for item in _OBJECT_SEQUENCE.validate_python(value) + ] + case _: + return _plain_json_value(value) + + +def _custom_generate_content_wire_body(response: object) -> Mapping[str, JsonValue]: + wire_body: Final = _gemini_wire_value( + _GENERATE_CONTENT_RESULT.validate_python( + _GENERATE_CONTENT_RESULT.dump_python(response, by_alias=True, exclude_none=True) + ), + camel_case_keys=True, + ) + if not isinstance(wire_body, dict): + return {} + return {key: value for key, value in wire_body.items() if key not in _SDK_ONLY_RESPONSE_FIELDS} + + +def _log_custom_generate_content_response( + response: object, logging_obj: LiteLLMLoggingObj, contents: object, api_key: str | None +) -> Mapping[str, JsonValue]: + wire_body: Final = _custom_generate_content_wire_body(response) + raw_response: Final = httpx.Response(status_code=200, json={"usageMetadata": {}, **wire_body}) + logging_obj.model_call_details["httpx_response"] = raw_response # rebind-ok: success logging needs raw response + logging_obj.post_call(input=contents, api_key=api_key, original_response=raw_response.text) # pyright: ignore[reportUnknownMemberType] # legacy logging interface is untyped + return wire_body + + +async def _await_custom_generate_content_response( + response: Awaitable[object], logging_obj: LiteLLMLoggingObj, contents: object, api_key: str | None +) -> Mapping[str, JsonValue]: + result: Final = await response + return _log_custom_generate_content_response(result, logging_obj, contents, api_key) + + +def _finish_custom_generate_content_response( + response: object, logging_obj: LiteLLMLoggingObj, contents: object, api_key: str | None +) -> object: + if asyncio.iscoroutine(response): + return _await_custom_generate_content_response(response, logging_obj, contents, api_key) + return _log_custom_generate_content_response(response, logging_obj, contents, api_key) + + @client async def agenerate_content( model: str, @@ -348,13 +493,56 @@ def generate_content( **kwargs, ) + custom_handler: Final = _custom_generate_content_handler(setup_result.custom_llm_provider) + native_hook_overridden: Final = custom_handler is not None and ( + type(custom_handler).agenerate_content is not CustomLLM.agenerate_content + if _is_async + else type(custom_handler).generate_content is not CustomLLM.generate_content + ) + if custom_handler is not None and native_hook_overridden: + logging_obj: Final = setup_result.litellm_logging_obj + native_contents: Final = _GENERATE_CONTENT_RESULT.validate_python(contents) + native_tools: Final = _GENERATE_CONTENT_RESULT.validate_python(tools) + native_optional_params: Final[Mapping[str, object]] = _CUSTOM_GENERATE_CONTENT_BODY.validate_python( + { + **{field: kwargs[field] for field in _CUSTOM_NATIVE_REQUEST_FIELDS if field in kwargs}, + "extra_headers": extra_headers, + "extra_query": extra_query, + "extra_body": extra_body, + } + ) + logging_obj.pre_call( # pyright: ignore[reportUnknownMemberType] # legacy logging interface is untyped + input=native_contents, + api_key=setup_result.litellm_params.api_key, + additional_args={ + "complete_input_dict": {"contents": contents, "generationConfig": config, "tools": tools} + }, + ) + hook: Final = custom_handler.agenerate_content if _is_async else custom_handler.generate_content + native_response: Final = hook( + model=setup_result.model, + contents=native_contents, + config=setup_result.generate_content_config_dict if config is not None else None, + tools=native_tools, + litellm_params=setup_result.litellm_params, + logging_obj=logging_obj, + optional_params=native_optional_params, + api_key=setup_result.litellm_params.api_key, + api_base=setup_result.litellm_params.api_base, + timeout=timeout or request_timeout, + ) + return _finish_custom_generate_content_response( + native_response, logging_obj, native_contents, setup_result.litellm_params.api_key + ) + # Extract systemInstruction from kwargs to pass to handler system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction") # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: + _require_adapter_generation_config(setup_result) # Use the adapter to convert to completion format - return GenerateContentToCompletionHandler.generate_content_handler( + adapter_response: Final = GenerateContentToCompletionHandler.generate_content_handler( model=model, contents=contents, config=setup_result.generate_content_config_dict, @@ -362,8 +550,17 @@ def generate_content( _is_async=_is_async, litellm_params=setup_result.litellm_params, extra_headers=extra_headers, + custom_logging_obj=setup_result.litellm_logging_obj if custom_handler is not None else None, **kwargs, ) + if custom_handler is not None: + return _finish_custom_generate_content_response( + adapter_response, + setup_result.litellm_logging_obj, + _GENERATE_CONTENT_RESULT.validate_python(contents), + setup_result.litellm_params.api_key, + ) + return adapter_response # Call the standard handler response: Final = _GENERATE_CONTENT_RESULT.validate_python( @@ -448,6 +645,8 @@ async def agenerate_content_stream( if setup_result.generate_content_provider_config is None: if "stream" in kwargs: kwargs.pop("stream", None) + _require_custom_stream_hook(setup_result, is_async=True) + _require_adapter_generation_config(setup_result) # Use the adapter to convert to completion format return await GenerateContentToCompletionHandler.async_generate_content_handler( @@ -458,6 +657,7 @@ async def agenerate_content_stream( tools=tools, stream=True, extra_headers=extra_headers, + custom_logging_obj=_custom_adapter_logging_obj(setup_result), **kwargs, ) @@ -540,6 +740,8 @@ def generate_content_stream( if setup_result.generate_content_provider_config is None: if "stream" in kwargs: kwargs.pop("stream", None) + _require_custom_stream_hook(setup_result, is_async=_is_async is True) + _require_adapter_generation_config(setup_result) # Use the adapter to convert to completion format return GenerateContentToCompletionHandler.generate_content_handler( @@ -550,6 +752,7 @@ def generate_content_stream( litellm_params=setup_result.litellm_params, stream=True, extra_headers=extra_headers, + custom_logging_obj=_custom_adapter_logging_obj(setup_result), **kwargs, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index bd2fb262810..1c87a48f1a0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3320,6 +3320,8 @@ class Logging(LiteLLMLoggingBaseClass): if call_type == CallTypes.anthropic_messages.value: return self._anthropic_messages_logged_response(result=result) if call_type in (CallTypes.generate_content.value, CallTypes.agenerate_content.value): + if isinstance(result, ModelResponse): + return result return self._handle_non_streaming_google_genai_generate_content_response_logging(result=result) if call_type in (CallTypes.asend_message.value, CallTypes.send_message.value): return self._handle_a2a_response_logging(result=result) diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index 03e5b8f6769..1e8df2a2849 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -22,6 +22,14 @@ else: from litellm.types.router import GenericLiteLLMParams +GENERATE_CONTENT_REQUEST_TOP_LEVEL_FIELDS: Final = ( + "safetySettings", + "toolConfig", + "cachedContent", + "labels", + "continuationToken", +) + class BaseGoogleGenAIGenerateContentConfig(ABC): """Base configuration class for Google GenAI generate_content functionality""" @@ -73,7 +81,7 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): ``systemInstruction`` (dedicated extraction) and ``generationConfig`` (mapped to ``config``). """ - return ("safetySettings", "toolConfig", "cachedContent", "labels", "continuationToken") + return GENERATE_CONTENT_REQUEST_TOP_LEVEL_FIELDS def transform_top_level_request_fields( self, diff --git a/litellm/llms/custom_llm.py b/litellm/llms/custom_llm.py index fd09d300ae2..6e47f90da5b 100644 --- a/litellm/llms/custom_llm.py +++ b/litellm/llms/custom_llm.py @@ -8,7 +8,7 @@ - async_streaming """ -from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Sequence +from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping, Sequence from typing import ( TYPE_CHECKING, Union, @@ -26,6 +26,8 @@ from .base import BaseLLM if TYPE_CHECKING: from litellm import CustomStreamWrapper from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.google_genai.main import GenerateContentResponse + from litellm.types.router import GenericLiteLLMParams class CustomLLMError(Exception): # use this for all your exceptions @@ -43,6 +45,36 @@ class CustomLLM(BaseLLM): def __init__(self) -> None: super().__init__() + def generate_content( + self, + model: str, + contents: object, + config: Mapping[str, object] | None, + tools: object | None, + litellm_params: "GenericLiteLLMParams", + logging_obj: "LiteLLMLoggingObj", + optional_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> "GenerateContentResponse | Mapping[str, object]": + raise CustomLLMError(status_code=500, message="Not implemented yet!") + + async def agenerate_content( + self, + model: str, + contents: object, + config: Mapping[str, object] | None, + tools: object | None, + litellm_params: "GenericLiteLLMParams", + logging_obj: "LiteLLMLoggingObj", + optional_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> "GenerateContentResponse | Mapping[str, object]": + raise CustomLLMError(status_code=500, message="Not implemented yet!") + def completion( self, model: str, diff --git a/litellm/main.py b/litellm/main.py index 2cd9f12c649..595612cb508 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8430,7 +8430,7 @@ async def aadapter_generate_content( coro: Final = cast( Coroutine[object, object, dict[str, object] | AsyncIterator[bytes]], - GenerateContentToCompletionHandler.generate_content_handler(**kwargs, _is_async=True), + GenerateContentToCompletionHandler.generate_content_handler(**kwargs, _is_async=True, custom_logging_obj=None), ) return await coro diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index ab39c2ad774..fc153ee9447 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -100,6 +100,7 @@ IGNORE_FUNCTIONS = [ "to_wire", # router_rust model info: bounded by the nesting depth of a stored deployment or cost-map entry (finite JSON trees: pydantic refuses a cycle when it dumps a deployment). "from_wire", # router_rust model info: walks only what router-core merged from to_wire's output, so it never goes deeper than to_wire did. "_valid_frontmatter_node", # max depth 32, manifest size capped at 1 MiB, and YAML aliases rejected before traversal. + "_gemini_wire_value", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); returns the subtree as plain JSON at the cap. ] diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index ae350c98075..5eba0ae1770 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -3,17 +3,29 @@ Test to verify the Google GenAI generate_content adapter functionality """ +import asyncio import json -from collections.abc import Iterator, Mapping +from collections.abc import AsyncIterator, Iterator, Mapping +from datetime import datetime +from threading import Event from types import MappingProxyType from typing import Final, Literal +import httpx import pytest -from pydantic import TypeAdapter +import respx +from pydantic import BaseModel, ConfigDict, TypeAdapter import litellm +from litellm.google_genai.main import agenerate_content_stream, generate_content_stream +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.custom_llm import CustomLLM, CustomLLMError from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig from litellm.llms.vertex_ai.google_genai.transformation import VertexAIGoogleGenAIConfig +from litellm.types.google_genai.main import GenerateContentResponse from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES from litellm.types.router import GenericLiteLLMParams @@ -387,6 +399,827 @@ async def test_vertex_generate_content_spend_follows_the_deployment_location( assert recorded == pytest.approx(expected), (location, stream) +class NativeContentLLM(CustomLLM): + def __init__(self) -> None: + super().__init__() + self.requests: tuple[dict[str, object], ...] = () + self.response: dict[str, object] = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"inlineData": {"mimeType": "image/png", "data": "aW1hZ2U="}}], + }, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 5, "totalTokenCount": 8}, + } + + def generate_content(self, **kwargs: object) -> dict[str, object]: + self.requests = (*self.requests, kwargs) + return self.response + + async def agenerate_content(self, **kwargs: object) -> dict[str, object]: + self.requests = (*self.requests, kwargs) + return self.response + + +class CompletionContentLLM(CustomLLM): + def completion(self, **kwargs: object) -> litellm.ModelResponse: + return litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "adapter result"}}]) + + async def acompletion(self, **kwargs: object) -> litellm.ModelResponse: + return self.completion(**kwargs) + + +class StreamingContentLLM(NativeContentLLM): + def streaming(self, **kwargs: object) -> Iterator[litellm.GenericStreamingChunk]: + yield litellm.GenericStreamingChunk( + text="stream adapter", + tool_use=None, + is_finished=True, + finish_reason="stop", + index=0, + usage={"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + ) + + async def astreaming(self, **kwargs: object) -> AsyncIterator[litellm.GenericStreamingChunk]: + for chunk in self.streaming(**kwargs): + yield chunk + + +@pytest.fixture +def native_content_handler(monkeypatch: pytest.MonkeyPatch) -> NativeContentLLM: + handler: Final = NativeContentLLM() + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": handler}]) + monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers]) + monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list]) + litellm.utils.custom_llm_setup() + return handler + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_native_content_preserves_images_and_config( + native_content_handler: NativeContentLLM, is_async: bool +) -> None: + contents: Final = [{"role": "user", "parts": [{"inlineData": {"mimeType": "image/png", "data": "aW5wdXQ="}}]}] + config: Final = {"imageConfig": {"aspectRatio": "1:1"}, "responseModalities": ["IMAGE"], "seed": 17} + tools: Final = [{"googleSearch": {}}] + request: Final = dict( + model="native_image_test/image-model", + contents=contents, + config=config, + tools=tools, + safetySettings=[], + systemInstruction={"parts": [{"text": "draw"}]}, + secret_fields={"raw_headers": {"authorization": "Bearer sk-caller-key"}}, + metadata={"user_api_key_alias": "caller"}, + ) + result: Final = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + assert result == native_content_handler.response + received: Final = native_content_handler.requests[0] + assert received["model"] == "image-model" + assert received["contents"] is contents + assert received["config"] == config + assert received["tools"] is tools + optional_params: Final = received["optional_params"] + assert isinstance(optional_params, dict) + assert optional_params["safetySettings"] == [] + assert optional_params["systemInstruction"] == request["systemInstruction"] + assert "sk-caller-key" not in str(optional_params) + assert "metadata" not in optional_params + params: Final = received["litellm_params"] + assert isinstance(params, litellm.types.router.GenericLiteLLMParams) + assert params.custom_llm_provider == "native_image_test" + assert received["logging_obj"] is not None + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("config", [None, {}]) +@pytest.mark.asyncio +async def test_custom_native_content_preserves_empty_request( + native_content_handler: NativeContentLLM, is_async: bool, config: object +) -> None: + request: Final = dict(model="native_image_test/image-model", contents=[], config=config) + result: Final = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + assert result == native_content_handler.response + received: Final = native_content_handler.requests[0] + assert received["contents"] == [] + assert received["config"] == config + assert received["tools"] is None + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_without_native_hook_uses_completion( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, is_async: bool +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": CompletionContentLLM()}] + ) + request: Final = dict(model="native_image_test/chat-model", contents=[{"parts": [{"text": "hello"}]}]) + result: Final = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + assert result["candidates"][0]["content"]["parts"] == [{"text": "adapter result"}] + + +class FailingCompletionContentLLM(CustomLLM): + def completion(self, **kwargs: object) -> litellm.ModelResponse: + raise CustomLLMError(status_code=429, message="adapter quota exhausted") + + async def acompletion(self, **kwargs: object) -> litellm.ModelResponse: + return self.completion(**kwargs) + + +class ContentFailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.calls: tuple[str, ...] = () + self.asynchronous_calls: tuple[str, ...] = () + self.logged: Final = Event() + + def log_failure_event(self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object) -> None: + self.calls = (*self.calls, str(kwargs["call_type"])) + self.logged.set() + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.asynchronous_calls = (*self.asynchronous_calls, str(kwargs["call_type"])) + self.log_failure_event(kwargs, response_obj, start_time, end_time) + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("use_logging_context", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_adapter_delivers_per_request_failure_callbacks_once( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, + is_async: bool, use_logging_context: bool, +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": FailingCompletionContentLLM()}] + ) + global_recorder: Final = ContentFailureRecorder() + request_recorder: Final = ContentFailureRecorder() + monkeypatch.setattr(litellm, "callbacks", [global_recorder]) + monkeypatch.setattr(litellm, "failure_callback", [global_recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [global_recorder]) + callback_args: Final = ( + {"litellm_logging_obj": Logging( + model="native_image_test/chat-model", messages=[], stream=False, + call_type="agenerate_content" if is_async else "generate_content", + start_time=datetime(2026, 1, 1), litellm_call_id="failure-test", function_id="failure-test", + dynamic_failure_callbacks=[request_recorder], dynamic_async_failure_callbacks=[request_recorder], + )} + if use_logging_context else {"failure_callback": [request_recorder.log_failure_event]} + ) + request: Final = dict( + model="native_image_test/chat-model", contents=[{"parts": [{"text": "hello"}]}], + **callback_args, + ) + if is_async: + with pytest.raises(litellm.RateLimitError, match="adapter quota exhausted"): + await litellm.agenerate_content(**request) + else: + with pytest.raises(litellm.RateLimitError, match="adapter quota exhausted"): + litellm.generate_content(**request) + assert await asyncio.to_thread(request_recorder.logged.wait, 5) + assert await asyncio.to_thread(global_recorder.logged.wait, 5) + expected: Final = ("agenerate_content" if is_async else "generate_content",) + assert request_recorder.calls == expected + assert global_recorder.calls == expected + assert request_recorder.asynchronous_calls == (expected if is_async and use_logging_context else ()) + assert global_recorder.asynchronous_calls == (expected if is_async else ()) + + +@pytest.mark.asyncio +async def test_router_custom_native_content_leg(native_content_handler: NativeContentLLM) -> None: + router: Final = litellm.Router( + model_list=[ + {"model_name": "image-group", "litellm_params": {"model": "native_image_test/image-model", "order": 0}}, + { + "model_name": "image-group", + "litellm_params": {"model": "gemini/test-model", "api_key": "test-key", "order": 1}, + }, + ] + ) + result: Final = await router.agenerate_content(model="image-group", contents=[], config={}) + assert isinstance(result, dict) + assert {key: value for key, value in result.items() if key != "_hidden_params"} == native_content_handler.response + assert "_hidden_params" not in native_content_handler.response + assert native_content_handler.requests[0]["contents"] == [] + assert native_content_handler.requests[0]["config"] == {} + + +class FailingContentLLM(CustomLLM): + def generate_content(self, **kwargs: object) -> dict[str, object]: + raise CustomLLMError(status_code=429, message="native quota exhausted") + + async def agenerate_content(self, **kwargs: object) -> dict[str, object]: + raise CustomLLMError(status_code=429, message="native quota exhausted") + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_native_content_maps_errors( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, is_async: bool +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": FailingContentLLM()}] + ) + if is_async: + with pytest.raises(litellm.RateLimitError, match="native quota exhausted"): + await litellm.agenerate_content(model="native_image_test/image-model", contents=[]) + else: + with pytest.raises(litellm.RateLimitError, match="native quota exhausted"): + litellm.generate_content(model="native_image_test/image-model", contents=[]) + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_stream_uses_completion_adapter( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, is_async: bool +) -> None: + handler: Final = StreamingContentLLM() + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": handler}]) + stream: Final = ( + await agenerate_content_stream(model="native_image_test/chat-model", contents=[]) + if is_async + else generate_content_stream(model="native_image_test/chat-model", contents=[]) + ) + chunks: Final = [chunk async for chunk in stream] + assert any(b"stream adapter" in chunk for chunk in chunks) + assert handler.requests == () + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_native_config_cannot_overwrite_accounting_context( + native_content_handler: NativeContentLLM, is_async: bool +) -> None: + config: Final = { + "litellm_params": {"metadata": {}}, + "model": "spoofed", + "custom_llm_provider": "spoofed", + "call_type": "spoofed", + "seed": 17, + } + metadata: Final = {"user_api_key": "trusted-key-id", "user_api_key_team_id": "trusted-team"} + request: Final = dict(model="native_image_test/image-model", contents=[], config=config, metadata=metadata) + result: Final = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + received: Final = native_content_handler.requests[0] + logged: Final = received["logging_obj"].model_call_details + assert result == native_content_handler.response + assert received["config"] == config + assert logged["litellm_params"]["metadata"]["user_api_key"] == "trusted-key-id" + assert logged["litellm_params"]["metadata"]["user_api_key_team_id"] == "trusted-team" + assert logged["model"] == "image-model" + assert logged["custom_llm_provider"] == "native_image_test" + assert logged["call_type"] == ("agenerate_content" if is_async else "generate_content") + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_builtin_native_content_keeps_priority_over_colliding_custom_handler( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, is_async: bool +) -> None: + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "gemini", "custom_handler": native_content_handler}]) + litellm.utils.custom_llm_setup() + request: Final = dict( + model="gemini/test-model", + contents=[{"parts": [{"text": "draw"}]}], + config={"responseModalities": ["IMAGE"]}, + api_key="test-key", + ) + with respx.mock: + route: Final = respx.post( + "https://generativelanguage.googleapis.com/v1alpha/models/test-model:generateContent" + ).respond(200, json=native_content_handler.response) + if is_async: + client: Final = AsyncHTTPHandler(transport=httpx.AsyncHTTPTransport()) + try: + result = await litellm.agenerate_content(**request, client=client) + finally: + await client.close() + else: + result = litellm.generate_content(**request) + assert result.candidates == native_content_handler.response["candidates"] + assert result.usageMetadata == native_content_handler.response["usageMetadata"] + assert route.call_count == 1 + assert json.loads(route.calls[0].request.content)["generationConfig"] == request["config"] + assert native_content_handler.requests == () + + +class ContentSuccessRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.sync_calls: tuple[str, ...] = () + self.async_calls: tuple[str, ...] = () + self.call_ids: tuple[str, ...] = () + self.sync_logged = Event() + self.async_logged = asyncio.Event() + + def log_success_event(self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object) -> None: + self.sync_calls = (*self.sync_calls, str(kwargs["call_type"])) + self.call_ids = (*self.call_ids, str(kwargs["litellm_call_id"])) + self.sync_logged.set() + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.async_calls = (*self.async_calls, str(kwargs["call_type"])) + self.call_ids = (*self.call_ids, str(kwargs["litellm_call_id"])) + self.async_logged.set() + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_adapter_accounts_for_one_success( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, is_async: bool +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": CompletionContentLLM()}] + ) + recorder: Final = ContentSuccessRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + outer_logging: Final = Logging( + model="native_image_test/chat-model", messages=[], stream=False, + call_type="agenerate_content" if is_async else "generate_content", + start_time=datetime(2026, 1, 1), litellm_call_id="adapter-accounting", function_id="adapter-accounting", + ) + request: Final = dict( + model="native_image_test/chat-model", contents=[{"parts": [{"text": "hello"}]}], + litellm_logging_obj=outer_logging, + ) + result: Final = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + if is_async: + await asyncio.wait_for(recorder.async_logged.wait(), timeout=5) + await GLOBAL_LOGGING_WORKER.flush() + await outer_logging.async_success_handler(result=result) + assert recorder.async_calls == ("agenerate_content",) + else: + assert await asyncio.to_thread(recorder.sync_logged.wait, 5) + outer_logging.success_handler(result=result) + assert recorder.sync_calls == ("generate_content",) + assert result["candidates"][0]["content"]["parts"] == [{"text": "adapter result"}] + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_adapter_delivers_per_request_success_callbacks_once( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, is_async: bool +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": CompletionContentLLM()}] + ) + global_recorder: Final = ContentSuccessRecorder() + request_recorder: Final = ContentSuccessRecorder() + monkeypatch.setattr(litellm, "callbacks", [global_recorder]) + monkeypatch.setattr(litellm, "success_callback", [global_recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [global_recorder]) + request: Final = dict( + model="native_image_test/chat-model", contents=[{"parts": [{"text": "hello"}]}], + success_callback=[request_recorder.async_log_success_event if is_async else request_recorder.log_success_event], + ) + result: Final = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + if is_async: + await asyncio.wait_for( + asyncio.gather(request_recorder.async_logged.wait(), global_recorder.async_logged.wait()), timeout=5 + ) + await GLOBAL_LOGGING_WORKER.flush() + assert request_recorder.async_calls == ("agenerate_content",) + assert global_recorder.async_calls == ("agenerate_content",) + else: + assert await asyncio.to_thread(request_recorder.sync_logged.wait, 5) + assert await asyncio.to_thread(global_recorder.sync_logged.wait, 5) + assert request_recorder.sync_calls == ("generate_content",) + assert global_recorder.sync_calls == ("generate_content",) + assert result["candidates"][0]["content"]["parts"] == [{"text": "adapter result"}] + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_stream_delivers_per_request_success_callbacks_once( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, is_async: bool +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": StreamingContentLLM()}] + ) + request_recorder: Final = ContentSuccessRecorder() + callback: Final = request_recorder.async_log_success_event if is_async else request_recorder.log_success_event + request: Final = dict(model="native_image_test/chat-model", contents=[], success_callback=[callback]) + stream: Final = await agenerate_content_stream(**request) if is_async else generate_content_stream(**request) + chunks: Final = [chunk async for chunk in stream] + assert any(b"stream adapter" in chunk for chunk in chunks) + if is_async: + await asyncio.wait_for(request_recorder.async_logged.wait(), timeout=5) + await GLOBAL_LOGGING_WORKER.flush() + assert request_recorder.async_calls == ("agenerate_content_stream",) + else: + assert await asyncio.to_thread(request_recorder.sync_logged.wait, 5) + assert request_recorder.sync_calls == ("generate_content_stream",) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_adapter_logs_under_the_callers_call_id( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, is_async: bool, stream: bool +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": StreamingContentLLM()}] + ) + recorder: Final = ContentSuccessRecorder() + callback: Final = recorder.async_log_success_event if is_async else recorder.log_success_event + request: Final = dict( + model="native_image_test/chat-model", contents=[], success_callback=[callback], litellm_call_id="client-call" + ) + if stream: + stream_response: Final = ( + await agenerate_content_stream(**request) if is_async else generate_content_stream(**request) + ) + _ = [chunk async for chunk in stream_response] + else: + _ = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + if is_async: + await asyncio.wait_for(recorder.async_logged.wait(), timeout=5) + await GLOBAL_LOGGING_WORKER.flush() + else: + assert await asyncio.to_thread(recorder.sync_logged.wait, 5) + assert recorder.call_ids == ("client-call",) + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_stream_without_stream_hook_fails_before_streaming( + native_content_handler: NativeContentLLM, is_async: bool +) -> None: + request: Final = dict(model="native_image_test/image-model", contents=[]) + with pytest.raises(litellm.APIConnectionError, match="astreaming" if is_async else "streaming") as error: + _ = await agenerate_content_stream(**request) if is_async else generate_content_stream(**request) + assert error.value.status_code == 500 + assert native_content_handler.requests == () + + +@pytest.mark.asyncio +async def test_custom_content_stream_without_stream_hook_keeps_router_legs_serving_chat( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": CompletionContentLLM()}] + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "chat-group", + "litellm_params": {"model": f"native_image_test/{leg}", "order": order}, + "model_info": {"id": leg}, + } + for order, leg in enumerate(("leg-a", "leg-b"), start=1) + ], + num_retries=0, + allowed_fails=0, + ) + for _ in range(2): + with pytest.raises(litellm.APIConnectionError, match="astreaming"): + _ = await router.agenerate_content_stream(model="chat-group", contents=[{"parts": [{"text": "hi"}]}]) + response: Final = await router.acompletion(model="chat-group", messages=[{"role": "user", "content": "hi"}]) + assert response.choices[0].message.content == "adapter result" + + +class CamelCaseContentResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + candidates: tuple[dict[str, object], ...] + usageMetadata: dict[str, int] + promptFeedback: dict[str, object] | None = None + + +_IMAGE_CANDIDATE: Final = { + "content": {"role": "model", "parts": [{"inlineData": {"mimeType": "image/png", "data": "aW1hZ2U="}}]}, + "finishReason": "STOP", + "index": 0, +} +_IMAGE_USAGE: Final = {"promptTokenCount": 3, "candidatesTokenCount": 5, "totalTokenCount": 8} +_IMAGE_WIRE_BODY: Final = {"candidates": [_IMAGE_CANDIDATE], "usageMetadata": _IMAGE_USAGE} +_DICT_WITHOUT_USAGE: Final = {"candidates": [_IMAGE_CANDIDATE]} +_NULL_TOOL_ARGS: Final = {"end_date": None, "trip": {"return_date": None, "seats": [None, 2]}} +_HOOK_RESPONSE_SHAPES: Final = ( + pytest.param( + GenerateContentResponse( + candidates=[ + { + "content": {"role": "model", "parts": [{"inline_data": {"mime_type": "image/png", "data": b"image"}}]}, + "finish_reason": "STOP", + "index": 0, + } + ], + usage_metadata={"prompt_token_count": 3, "candidates_token_count": 5, "total_token_count": 8}, + ), + (3, 5), + _IMAGE_WIRE_BODY, + id="snake-case-sdk-object-with-bytes", + ), + pytest.param( + CamelCaseContentResponse(candidates=(_IMAGE_CANDIDATE,), usageMetadata=_IMAGE_USAGE), + (3, 5), + _IMAGE_WIRE_BODY, + id="camel-case-model-with-unset-fields", + ), + pytest.param( + GenerateContentResponse( + candidates=[ + { + "content": {"role": "model", "parts": [{"function_call": {"name": "book", "args": _NULL_TOOL_ARGS}}]}, + "index": 0, + } + ], + usage_metadata={"prompt_token_count": 3, "candidates_token_count": 5, "total_token_count": 8}, + ), + (3, 5), + { + "candidates": [ + { + "content": {"role": "model", "parts": [{"functionCall": {"name": "book", "args": _NULL_TOOL_ARGS}}]}, + "index": 0, + } + ], + "usageMetadata": _IMAGE_USAGE, + }, + id="sdk-object-with-null-tool-args", + ), + pytest.param( + GenerateContentResponse( + candidates=[_IMAGE_CANDIDATE], + usage_metadata={"prompt_token_count": 3, "candidates_token_count": 5, "total_token_count": 8}, + sdk_http_response={"headers": {"x-upstream-trace": "abc"}}, + automatic_function_calling_history=[{"role": "user", "parts": [{"text": "draw"}]}], + parsed={"answer": "draw"}, + ), + (3, 5), + _IMAGE_WIRE_BODY, + id="sdk-object-with-client-side-fields", + ), + pytest.param( + GenerateContentResponse( + candidates=[_IMAGE_CANDIDATE], + usage_metadata={"prompt_token_count": 3, "candidates_token_count": 5, "total_token_count": 8}, + sdk_http_response={"headers": {"x-upstream-trace": "abc"}}, + automatic_function_calling_history=[{"role": "user", "parts": [{"text": "draw"}]}], + parsed={"answer": "draw"}, + ).model_dump(exclude_none=True), + (3, 5), + _IMAGE_WIRE_BODY, + id="dumped-sdk-dict-with-client-side-fields", + ), + pytest.param(_DICT_WITHOUT_USAGE, (0, 0), _DICT_WITHOUT_USAGE, id="dict-without-usage"), +) + + +class ContentUsageRecorder: + def __init__(self) -> None: + self.usage: tuple[tuple[object, object], ...] = () + self.costs: tuple[object, ...] = () + self.logged: Final = Event() + + def log(self, kwargs: dict[str, object], *_: object) -> None: + standard_logging_object: Final = kwargs["standard_logging_object"] + assert isinstance(standard_logging_object, dict) + self.usage = ( + *self.usage, + (standard_logging_object["prompt_tokens"], standard_logging_object["completion_tokens"]), + ) + self.costs = (*self.costs, standard_logging_object["response_cost"]) + self.logged.set() + + async def alog(self, kwargs: dict[str, object], *rest: object) -> None: + self.log(kwargs, *rest) + + +@pytest.mark.parametrize("hook_response, expected_usage, expected_result", _HOOK_RESPONSE_SHAPES) +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_native_content_logs_each_hook_response_shape_once( + native_content_handler: NativeContentLLM, + is_async: bool, + hook_response: object, + expected_usage: tuple[int, int], + expected_result: dict[str, object], +) -> None: + native_content_handler.response = hook_response + recorder: Final = ContentUsageRecorder() + request: Final = dict( + model="native_image_test/image-model", contents=[{"parts": [{"text": "draw"}]}], + success_callback=[recorder.alog if is_async else recorder.log], + ) + result: Final = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + assert await asyncio.to_thread(recorder.logged.wait, 5) + if is_async: + await GLOBAL_LOGGING_WORKER.flush() + assert result == expected_result + assert recorder.usage == (expected_usage,) + + +@pytest.mark.asyncio +async def test_router_tags_custom_native_content_objects_with_their_deployment( + native_content_handler: NativeContentLLM, +) -> None: + native_content_handler.response = CamelCaseContentResponse(candidates=(_IMAGE_CANDIDATE,), usageMetadata=_IMAGE_USAGE) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "image-group", + "litellm_params": {"model": "native_image_test/image-model"}, + "model_info": {"id": "native-leg"}, + } + ] + ) + result: Final = await router.agenerate_content(model="image-group", contents=[], config={}) + assert isinstance(result, dict) + hidden_params: Final = result.pop("_hidden_params") + assert result == _IMAGE_WIRE_BODY + assert hidden_params["model_id"] == "native-leg" + assert hidden_params["additional_headers"]["x-litellm-model-group"] == "image-group" + + +class AdapterOnlyContentLLM(CompletionContentLLM): + streaming = StreamingContentLLM.streaming + astreaming = StreamingContentLLM.astreaming + + +@pytest.mark.parametrize("drop_params", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_adapter_rejects_generation_config_it_cannot_send( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, + is_async: bool, stream: bool, drop_params: bool, +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": AdapterOnlyContentLLM()}] + ) + request: Final = dict( + model="native_image_test/chat-model", contents=[{"parts": [{"text": "draw"}]}], + config={"temperature": 0.2, "imageConfig": {"aspectRatio": "16:9"}}, drop_params=drop_params, + ) + + async def call() -> object: + if stream: + response: Final = ( + await agenerate_content_stream(**request) if is_async else generate_content_stream(**request) + ) + return [chunk async for chunk in response] + return await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + + if not drop_params: + with pytest.raises(litellm.BadRequestError, match="imageConfig"): + _ = await call() + return + result: Final = await call() + if stream: + assert b"stream adapter" in b"".join(result) + else: + assert result["candidates"][0]["content"]["parts"] == [{"text": "adapter result"}] + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_adapter_accepts_generation_config_it_already_satisfies( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, is_async: bool +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": AdapterOnlyContentLLM()}] + ) + request: Final = dict( + model="native_image_test/chat-model", contents=[{"parts": [{"text": "hi"}]}], + config={"responseModalities": ["TEXT"], "candidateCount": 1}, + ) + result: Final = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + assert result["candidates"][0]["content"]["parts"] == [{"text": "adapter result"}] + + +@pytest.mark.parametrize( + "config", + [ + {"responseModalities": ["TEXT", "IMAGE"]}, + {"candidateCount": 2}, + {"responseMimeType": "application/json"}, + {"responseMimeType": "application/json", "responseSchema": {"type": "ARRAY", "items": {"type": "STRING"}}}, + {"responseMimeType": "text/x.enum", "responseSchema": {"type": "STRING", "enum": ["a", "b"]}}, + ], +) +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_adapter_rejects_generation_config_values_it_cannot_satisfy( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, + is_async: bool, config: dict[str, object], +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": AdapterOnlyContentLLM()}] + ) + request: Final = dict(model="native_image_test/chat-model", contents=[{"parts": [{"text": "hi"}]}], config=config) + with pytest.raises(litellm.BadRequestError, match=next(iter(config))): + _ = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + + +class ResponseFormatRecorderLLM(CompletionContentLLM): + def __init__(self) -> None: + super().__init__() + self.response_formats: tuple[object, ...] = () + + def completion(self, **kwargs: object) -> litellm.ModelResponse: + optional_params: Final = kwargs["optional_params"] + assert isinstance(optional_params, dict) + self.response_formats = (*self.response_formats, optional_params.get("response_format")) + return super().completion(**kwargs) + + +@pytest.mark.parametrize( + ("config", "response_format"), + [ + ({"responseMimeType": "text/plain"}, None), + ( + { + "responseMimeType": "application/json", + "responseSchema": {"type": "OBJECT", "properties": {"color": {"type": "STRING"}}}, + }, + { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": {"type": "object", "properties": {"color": {"type": "string"}}}, + }, + }, + ), + ], +) +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_content_adapter_sends_the_response_format_it_accepts( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch, + is_async: bool, config: dict[str, object], response_format: object, +) -> None: + handler: Final = ResponseFormatRecorderLLM() + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": handler}]) + request: Final = dict(model="native_image_test/chat-model", contents=[{"parts": [{"text": "hi"}]}], config=config) + _ = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + assert handler.response_formats == (response_format,) + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_custom_native_content_bills_per_request_prices( + native_content_handler: NativeContentLLM, is_async: bool +) -> None: + recorder: Final = ContentUsageRecorder() + request: Final = dict( + model="native_image_test/image-model", contents=[{"parts": [{"text": "draw"}]}], + input_cost_per_token=0.5, output_cost_per_token=2.0, + success_callback=[recorder.alog if is_async else recorder.log], + ) + _ = await litellm.agenerate_content(**request) if is_async else litellm.generate_content(**request) + assert await asyncio.to_thread(recorder.logged.wait, 5) + if is_async: + await GLOBAL_LOGGING_WORKER.flush() + assert recorder.usage == ((3, 5),) + assert recorder.costs == (pytest.approx(3 * 0.5 + 5 * 2.0),) + + +@pytest.mark.asyncio +async def test_custom_content_adapter_holds_success_logs_while_the_caller_defers_them( + native_content_handler: NativeContentLLM, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "native_image_test", "custom_handler": CompletionContentLLM()}] + ) + recorder: Final = ContentSuccessRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + outer_logging: Final = Logging( + model="native_image_test/chat-model", messages=[], stream=False, call_type="agenerate_content", + start_time=datetime(2026, 1, 1), litellm_call_id="deferred-adapter", function_id="deferred-adapter", + ) + outer_logging.defer_async_logging = True + _ = await litellm.agenerate_content( + model="native_image_test/chat-model", contents=[{"parts": [{"text": "hello"}]}], + litellm_logging_obj=outer_logging, + ) + await GLOBAL_LOGGING_WORKER.flush() + assert (recorder.sync_calls, recorder.async_calls) == ((), ()) + enqueue_deferred_logging: Final = outer_logging.enqueue_deferred_logging + assert enqueue_deferred_logging is not None + enqueue_deferred_logging() + await asyncio.wait_for(recorder.async_logged.wait(), timeout=5) + await GLOBAL_LOGGING_WORKER.flush() + assert recorder.async_calls == ("agenerate_content",) + assert set(recorder.call_ids) == {"deferred-adapter"} + + _LABEL_FIELDS: Final = ("user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata") _IDENTITY: Final = MappingProxyType({"user_api_key_alias": "chargeback-key", "user_api_key_team_alias": "ai-platform"}) _PROXY_METADATA: Final = MappingProxyType(