From d81c55f4cc86d0303dd89ebb7dba0a745ce425b4 Mon Sep 17 00:00:00 2001 From: mika-lt <349045450+mika-lt@users.noreply.github.com> Date: Sat, 14 Nov 2026 02:57:58 -0800 Subject: [PATCH] feat(custom_llm): let custom providers serve native generateContent (#51083) * feat: support native generateContent in CustomLLM Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: preserve custom generateContent accounting and built-in priority Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: forward custom generateContent success callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: preserve custom generateContent failure callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: retain async callbacks across generateContent fallback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(google_genai): scope custom generateContent hooks to native fields and stream only with a stream hook * fix(google_genai): log custom adapter calls on the caller's logging object and reject untranslatable config Custom adapter calls now share the caller's logging object, so the proxy's deferred logging covers them and one request logs one row under the caller's call id and call type. Native hook responses log from any pydantic, SDK, or dict shape, per-request prices apply to custom handlers, and a missing stream hook or a generationConfig key the adapter cannot send answers 400 * fix(google_genai): return custom hook objects as their camelCase dict so the router tags the deployment * fix(google_genai): keep tool args, drop SDK-only fields, and let routers retry custom legs that cannot stream Tool call args and function responses keep the caller's keys and null values. SDK objects no longer send sdkHttpResponse, automaticFunctionCallingHistory or parsed to the client or the logs. A custom leg without a stream hook answers 500 like CustomLLM's other unimplemented hooks, so a mixed router group retries the request on another leg. The wire-value walk is depth capped for the recursive function check * fix(google_genai): keep router legs serving chat after a custom stream without a hook, and accept config the adapter already meets * fix(google_genai): reject JSON mode a custom leg's completion adapter cannot send The custom-leg generationConfig check now accepts responseMimeType and response schema keys only when the adapter sends them as response_format, or when they ask for plain text. JSON mode without an object schema answers 400 unless drop_params is set * fix(google_genai): send a custom hook's dict back without SDK-only fields A hook that returns a dumped google-genai response as a dict sent its snake_case keys, parsed, sdk_http_response and automatic_function_calling_history to the client, past the guardrail that rewrites only candidate text. Dicts now go back as the same filtered camelCase body as every other hook return, and the router tags that copy instead of writing into the hook's own dict --------- Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 1 + litellm/google_genai/adapters/handler.py | 16 +- .../google_genai/adapters/transformation.py | 25 + litellm/google_genai/main.py | 213 ++++- litellm/litellm_core_utils/litellm_logging.py | 2 + .../base_llm/google_genai/transformation.py | 10 +- litellm/llms/custom_llm.py | 34 +- litellm/main.py | 2 +- .../code_coverage_tests/recursive_detector.py | 1 + .../google_genai/test_google_genai_main.py | 837 +++++++++++++++++- 10 files changed, 1130 insertions(+), 11 deletions(-) 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(