From 7ec7e5332c648fc7871e1084a2c7bb74b60a0204 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 1 Oct 2025 21:33:45 +0530 Subject: [PATCH] Add generateContent cost tracking (#15014) --- .../llm_passthrough_endpoints.py | 118 ++----- .../gemini_passthrough_logging_handler.py | 204 +++++++++++++ .../pass_through_endpoints.py | 195 ++++-------- .../pass_through_endpoints/success_handler.py | 190 ++++++------ ...test_gemini_passthrough_logging_handler.py | 287 ++++++++++++++++++ 5 files changed, 666 insertions(+), 328 deletions(-) create mode 100644 litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index f25ed8a7bf..8aa3b90d95 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -57,9 +57,7 @@ def create_request_copy(request: Request): } -def is_passthrough_request_using_router_model( - request_body: dict, llm_router: Optional[litellm.Router] -) -> bool: +def is_passthrough_request_using_router_model(request_body: dict, llm_router: Optional[litellm.Router]) -> bool: """ Returns True if the model is in the llm_router model names """ @@ -95,16 +93,12 @@ async def llm_passthrough_factory_proxy_route( model=None, ) if provider_config is None: - raise HTTPException( - status_code=404, detail=f"Provider {custom_llm_provider} not found" - ) + raise HTTPException(status_code=404, detail=f"Provider {custom_llm_provider} not found") base_target_url = provider_config.get_api_base() if base_target_url is None: - raise HTTPException( - status_code=404, detail=f"Provider {custom_llm_provider} api base not found" - ) + raise HTTPException(status_code=404, detail=f"Provider {custom_llm_provider} api base not found") encoded_endpoint = httpx.URL(endpoint).path @@ -183,17 +177,11 @@ async def gemini_proxy_route( [Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio) """ ## CHECK FOR LITELLM API KEY IN THE QUERY PARAMS - ?..key=LITELLM_API_KEY - google_ai_studio_api_key = request.query_params.get("key") or request.headers.get( - "x-goog-api-key" - ) + google_ai_studio_api_key = request.query_params.get("key") or request.headers.get("x-goog-api-key") - user_api_key_dict = await user_api_key_auth( - request=request, api_key=f"Bearer {google_ai_studio_api_key}" - ) + user_api_key_dict = await user_api_key_auth(request=request, api_key=f"Bearer {google_ai_studio_api_key}") - base_target_url = ( - os.getenv("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com" - ) + base_target_url = os.getenv("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com" encoded_endpoint = httpx.URL(endpoint).path # Ensure endpoint starts with '/' for proper URL construction @@ -226,6 +214,7 @@ async def gemini_proxy_route( endpoint_func = create_pass_through_route( endpoint=endpoint, target=str(updated_url), + custom_llm_provider="gemini", ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, @@ -310,9 +299,7 @@ async def vllm_proxy_route( from litellm.proxy.proxy_server import llm_router request_body = await get_request_body(request) - is_router_model = is_passthrough_request_using_router_model( - request_body, llm_router - ) + is_router_model = is_passthrough_request_using_router_model(request_body, llm_router) is_streaming_request = is_passthrough_request_streaming(request_body) if is_router_model and llm_router: result = cast( @@ -327,11 +314,7 @@ async def vllm_proxy_route( content=None, data=None, files=None, - json=( - request_body - if request.headers.get("content-type") == "application/json" - else None - ), + json=(request_body if request.headers.get("content-type") == "application/json" else None), params=None, headers=None, cookies=None, @@ -509,9 +492,7 @@ async def handle_bedrock_count_tokens( # Extract model from request body model = request_body.get("model") if not model: - raise HTTPException( - status_code=400, detail={"error": "Model is required in request body"} - ) + raise HTTPException(status_code=400, detail={"error": "Model is required in request body"}) # Get model parameters from router litellm_params = {"user_api_key_dict": user_api_key_dict} @@ -550,9 +531,7 @@ async def handle_bedrock_count_tokens( raise except Exception as e: verbose_proxy_logger.error(f"Error in handle_bedrock_count_tokens: {str(e)}") - raise HTTPException( - status_code=500, detail={"error": f"CountTokens processing error: {str(e)}"} - ) + raise HTTPException(status_code=500, detail={"error": f"CountTokens processing error: {str(e)}"}) async def bedrock_llm_proxy_route( @@ -604,8 +583,7 @@ async def bedrock_llm_proxy_route( raise HTTPException( status_code=400, detail={ - "error": "Model missing from endpoint. Expected format: /model//. Got: " - + endpoint, + "error": "Model missing from endpoint. Expected format: /model//. Got: " + endpoint, }, ) @@ -669,9 +647,7 @@ async def bedrock_proxy_route( aws_region_name = litellm.utils.get_secret(secret_name="AWS_REGION_NAME") if _is_bedrock_agent_runtime_route(endpoint=endpoint): # handle bedrock agents - base_target_url = ( - f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" - ) + base_target_url = f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" else: return await bedrock_llm_proxy_route( endpoint=endpoint, @@ -701,9 +677,7 @@ async def bedrock_proxy_route( data = await request.json() except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) - _request = AWSRequest( - method="POST", url=str(updated_url), data=json.dumps(data), headers=headers - ) + _request = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers) sigv4.add_auth(_request) prepped = _request.prepare() @@ -764,14 +738,8 @@ async def assemblyai_proxy_route( [Docs](https://api.assemblyai.com) """ # Set base URL based on the route - assembly_region = AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url( - url=str(request.url) - ) - base_target_url = ( - AssemblyAIPassthroughLoggingHandler._get_assembly_base_url_from_region( - region=assembly_region - ) - ) + assembly_region = AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url(url=str(request.url)) + base_target_url = AssemblyAIPassthroughLoggingHandler._get_assembly_base_url_from_region(region=assembly_region) encoded_endpoint = httpx.URL(endpoint).path # Ensure endpoint starts with '/' for proper URL construction if not encoded_endpoint.startswith("/"): @@ -829,18 +797,14 @@ async def azure_proxy_route( """ base_target_url = get_secret_str(secret_name="AZURE_API_BASE") if base_target_url is None: - raise Exception( - "Required 'AZURE_API_BASE' in environment to make pass-through calls to Azure." - ) + raise Exception("Required 'AZURE_API_BASE' in environment to make pass-through calls to Azure.") # Add or update query parameters azure_api_key = passthrough_endpoint_router.get_credentials( custom_llm_provider=litellm.LlmProviders.AZURE.value, region_name=None, ) if azure_api_key is None: - raise Exception( - "Required 'AZURE_API_KEY' in environment to make pass-through calls to Azure." - ) + raise Exception("Required 'AZURE_API_KEY' in environment to make pass-through calls to Azure.") return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler( endpoint=endpoint, @@ -864,9 +828,7 @@ class BaseVertexAIPassThroughHandler(ABC): @staticmethod @abstractmethod - def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] - ) -> str: + def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str: pass @@ -876,9 +838,7 @@ class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler): return "https://discoveryengine.googleapis.com/" @staticmethod - def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] - ) -> str: + def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str: return base_target_url @@ -888,9 +848,7 @@ class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler): return get_vertex_base_url(vertex_location) @staticmethod - def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] - ) -> str: + def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str: return get_vertex_base_url(vertex_location) @@ -956,18 +914,14 @@ async def _base_vertex_proxy_route( location=vertex_location, ) - base_target_url = get_vertex_pass_through_handler.get_default_base_target_url( - vertex_location - ) + base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location) headers_passed_through = False # Use headers from the incoming request if no vertex credentials are found if vertex_credentials is None or vertex_credentials.vertex_project is None: headers = dict(request.headers) or {} headers_passed_through = True - verbose_proxy_logger.debug( - "default_vertex_config not set, incoming request headers %s", headers - ) + verbose_proxy_logger.debug("default_vertex_config not set, incoming request headers %s", headers) headers.pop("content-length", None) headers.pop("host", None) else: @@ -1133,9 +1087,7 @@ async def openai_proxy_route( region_name=None, ) if openai_api_key is None: - raise Exception( - "Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI." - ) + raise Exception("Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI.") return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler( endpoint=endpoint, @@ -1181,9 +1133,7 @@ class BaseOpenAIPassThroughHandler: endpoint_func = create_pass_through_route( endpoint=endpoint, target=str(updated_url), - custom_headers=BaseOpenAIPassThroughHandler._assemble_headers( - api_key=api_key, request=request - ), + custom_headers=BaseOpenAIPassThroughHandler._assemble_headers(api_key=api_key, request=request), ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, @@ -1200,10 +1150,7 @@ class BaseOpenAIPassThroughHandler: """ Appends the OpenAI-Beta header to the headers if the request is an OpenAI Assistants API request """ - if ( - RouteChecks._is_assistants_api_request(request) is True - and "OpenAI-Beta" not in headers - ): + if RouteChecks._is_assistants_api_request(request) is True and "OpenAI-Beta" not in headers: headers["OpenAI-Beta"] = "assistants=v2" return headers @@ -1219,9 +1166,7 @@ class BaseOpenAIPassThroughHandler: ) @staticmethod - def _join_url_paths( - base_url: httpx.URL, path: str, custom_llm_provider: litellm.LlmProviders - ) -> str: + def _join_url_paths(base_url: httpx.URL, path: str, custom_llm_provider: litellm.LlmProviders) -> str: """ Properly joins a base URL with a path, preserving any existing path in the base URL. """ @@ -1237,14 +1182,9 @@ class BaseOpenAIPassThroughHandler: joined_path_str = str(base_url.copy_with(path=full_path)) # Apply OpenAI-specific path handling for both branches - if ( - custom_llm_provider == litellm.LlmProviders.OPENAI - and "/v1/" not in joined_path_str - ): + if custom_llm_provider == litellm.LlmProviders.OPENAI and "/v1/" not in joined_path_str: # Insert v1 after api.openai.com for OpenAI requests - joined_path_str = joined_path_str.replace( - "api.openai.com/", "api.openai.com/v1/" - ) + joined_path_str = joined_path_str.replace("api.openai.com/", "api.openai.com/v1/") return joined_path_str diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py new file mode 100644 index 0000000000..8c96c2ab96 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -0,0 +1,204 @@ +import json +import re +from datetime import datetime +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +import httpx + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator as GeminiModelResponseIterator, +) +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.types.utils import ( + ModelResponse, + TextCompletionResponse, +) + +if TYPE_CHECKING: + from ..success_handler import PassThroughEndpointLogging + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType +else: + PassThroughEndpointLogging = Any + EndpointType = Any + + +class GeminiPassthroughLoggingHandler: + @staticmethod + def gemini_passthrough_handler( + httpx_response: httpx.Response, + response_body: dict, + logging_obj: LiteLLMLoggingObj, + url_route: str, + result: str, + start_time: datetime, + end_time: datetime, + cache_hit: bool, + request_body: dict, + **kwargs, + ) -> PassThroughEndpointLoggingTypedDict: + if "generateContent" in url_route: + model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) + + # Use Gemini config for transformation + instance_of_gemini_llm = litellm.GoogleAIStudioGeminiConfig() + litellm_model_response: ModelResponse = instance_of_gemini_llm.transform_response( + model=model, + messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}], + raw_response=httpx_response, + model_response=litellm.ModelResponse(), + logging_obj=logging_obj, + optional_params={}, + litellm_params={}, + api_key="", + request_data={}, + encoding=litellm.encoding, + ) + kwargs = GeminiPassthroughLoggingHandler._create_gemini_response_logging_payload_for_generate_content( + litellm_model_response=litellm_model_response, + model=model, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + logging_obj=logging_obj, + custom_llm_provider="gemini", + ) + + return { + "result": litellm_model_response, + "kwargs": kwargs, + } + else: + return { + "result": None, + "kwargs": kwargs, + } + + @staticmethod + def _handle_logging_gemini_collected_chunks( + litellm_logging_obj: LiteLLMLoggingObj, + passthrough_success_handler_obj: PassThroughEndpointLogging, + url_route: str, + request_body: dict, + endpoint_type: EndpointType, + start_time: datetime, + all_chunks: List[str], + model: Optional[str], + end_time: datetime, + ) -> PassThroughEndpointLoggingTypedDict: + """ + Takes raw chunks from Gemini passthrough endpoint and logs them in litellm callbacks + + - Builds complete response from chunks + - Creates standard logging object + - Logs in litellm callbacks + """ + kwargs: Dict[str, Any] = {} + model = model or GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) + complete_streaming_response = GeminiPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + url_route=url_route, + ) + + if complete_streaming_response is None: + verbose_proxy_logger.error( + "Unable to build complete streaming response for Gemini passthrough endpoint, not logging..." + ) + return { + "result": None, + "kwargs": kwargs, + } + + kwargs = GeminiPassthroughLoggingHandler._create_gemini_response_logging_payload_for_generate_content( + litellm_model_response=complete_streaming_response, + model=model, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + logging_obj=litellm_logging_obj, + custom_llm_provider="gemini", + ) + + return { + "result": complete_streaming_response, + "kwargs": kwargs, + } + + @staticmethod + def _build_complete_streaming_response( + all_chunks: List[str], + litellm_logging_obj: LiteLLMLoggingObj, + model: str, + url_route: str, + ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: + parsed_chunks = [] + if "generateContent" in url_route or "streamGenerateContent" in url_route: + gemini_iterator: Any = GeminiModelResponseIterator( + streaming_response=None, + sync_stream=False, + logging_obj=litellm_logging_obj, + ) + chunk_parsing_logic: Any = gemini_iterator._common_chunk_parsing_logic + parsed_chunks = [chunk_parsing_logic(chunk) for chunk in all_chunks] + else: + return None + + if len(parsed_chunks) == 0: + return None + + all_openai_chunks = [] + for parsed_chunk in parsed_chunks: + if parsed_chunk is None: + continue + all_openai_chunks.append(parsed_chunk) + + complete_streaming_response = litellm.stream_chunk_builder(chunks=all_openai_chunks) + + return complete_streaming_response + + @staticmethod + def extract_model_from_url(url: str) -> str: + pattern = r"/models/([^:]+)" + match = re.search(pattern, url) + if match: + return match.group(1) + return "unknown" + + @staticmethod + def _create_gemini_response_logging_payload_for_generate_content( + litellm_model_response: Union[ModelResponse, TextCompletionResponse], + model: str, + kwargs: dict, + start_time: datetime, + end_time: datetime, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str, + ): + """ + Create the standard logging object for Gemini passthrough generateContent (streaming and non-streaming) + """ + + response_cost = litellm.completion_cost( + completion_response=litellm_model_response, + model=model, + custom_llm_provider="gemini", + ) + + kwargs["response_cost"] = response_cost + kwargs["model"] = model + kwargs["custom_llm_provider"] = custom_llm_provider + + # pretty print standard logging object + verbose_proxy_logger.debug("kwargs= %s", json.dumps(kwargs, indent=4)) + + # set litellm_call_id to logging response object + litellm_model_response.id = logging_obj.litellm_call_id + logging_obj.model = litellm_model_response.model or model + logging_obj.model_call_details["model"] = logging_obj.model + logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + logging_obj.model_call_details["response_cost"] = response_cost + return kwargs diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3f28ba92d3..e001b27ad3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -96,13 +96,9 @@ async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optiona # langfuse requires b64 encoded headers - we construct that here _langfuse_public_key = custom_headers["LANGFUSE_PUBLIC_KEY"] _langfuse_secret_key = custom_headers["LANGFUSE_SECRET_KEY"] - if isinstance( - _langfuse_public_key, str - ) and _langfuse_public_key.startswith("os.environ/"): + if isinstance(_langfuse_public_key, str) and _langfuse_public_key.startswith("os.environ/"): _langfuse_public_key = get_secret_str(_langfuse_public_key) - if isinstance( - _langfuse_secret_key, str - ) and _langfuse_secret_key.startswith("os.environ/"): + if isinstance(_langfuse_secret_key, str) and _langfuse_secret_key.startswith("os.environ/"): _langfuse_secret_key = get_secret_str(_langfuse_secret_key) headers["Authorization"] = "Basic " + b64encode( f"{_langfuse_public_key}:{_langfuse_secret_key}".encode("utf-8") @@ -111,9 +107,7 @@ async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optiona # for all other headers headers[key] = value if isinstance(value, str) and "os.environ/" in value: - verbose_proxy_logger.debug( - "pass through endpoint - looking up 'os.environ/' variable" - ) + verbose_proxy_logger.debug("pass through endpoint - looking up 'os.environ/' variable") # get string section that is os.environ/ start_index = value.find("os.environ/") _variable_name = value[start_index:] @@ -206,9 +200,7 @@ async def chat_completion_pass_through_endpoint( # noqa: PLR0915 # skip router if user passed their key if "api_key" in data: llm_response = asyncio.create_task(litellm.aadapter_completion(**data)) - elif ( - llm_router is not None and data["model"] in router_model_names - ): # model in router model list + elif llm_router is not None and data["model"] in router_model_names: # model in router model list llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) elif ( llm_router is not None @@ -237,10 +229,7 @@ async def chat_completion_pass_through_endpoint( # noqa: PLR0915 else: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "completion: Invalid model name passed in model=" - + data.get("model", "") - }, + detail={"error": "completion: Invalid model name passed in model=" + data.get("model", "")}, ) # Await the llm_response task @@ -254,9 +243,7 @@ async def chat_completion_pass_through_endpoint( # noqa: PLR0915 ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) verbose_proxy_logger.debug("final response: %s", response) @@ -278,11 +265,7 @@ async def chat_completion_pass_through_endpoint( # noqa: PLR0915 await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.completion(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - {}".format(str(e))) error_msg = f"{str(e)}" raise ProxyException( message=getattr(e, "message", error_msg), @@ -301,11 +284,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): ) -> dict: excluded_headers = {"transfer-encoding", "content-encoding"} - return_headers = { - key: value - for key, value in headers.items() - if key.lower() not in excluded_headers - } + return_headers = {key: value for key, value in headers.items() if key.lower() not in excluded_headers} if litellm_call_id: return_headers["x-litellm-call-id"] = litellm_call_id if custom_headers: @@ -432,10 +411,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): for field_name, field_value in form_data.items(): if isinstance(field_value, (StarletteUploadFile, UploadFile)): - files[field_name] = ( - await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( - upload_file=field_value - ) + files[field_name] = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( + upload_file=field_value ) else: form_data_dict[field_name] = field_value @@ -485,9 +462,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): user_api_key_spend=user_api_key_dict.spend, user_api_key_max_budget=user_api_key_dict.max_budget, user_api_key_budget_reset_at=( - user_api_key_dict.budget_reset_at.isoformat() - if user_api_key_dict.budget_reset_at - else None + user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None ), ) ) @@ -521,16 +496,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): "passthrough_logging_payload": passthrough_logging_payload, } - logging_obj.model_call_details["passthrough_logging_payload"] = ( - passthrough_logging_payload - ) + logging_obj.model_call_details["passthrough_logging_payload"] = passthrough_logging_payload return kwargs @staticmethod - def construct_target_url_with_subpath( - base_target: str, subpath: str, include_subpath: Optional[bool] - ) -> str: + def construct_target_url_with_subpath(base_target: str, subpath: str, include_subpath: Optional[bool]) -> str: """ Helper function to construct the full target URL with subpath handling. @@ -581,6 +552,7 @@ async def pass_through_request( # noqa: PLR0915 query_params: Optional[dict] = None, stream: Optional[bool] = None, cost_per_request: Optional[float] = None, + custom_llm_provider: Optional[str] = None, ): """ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called @@ -632,9 +604,7 @@ async def pass_through_request( # noqa: PLR0915 ).encode("ascii") ) - endpoint_type: EndpointType = HttpPassThroughEndpointHelpers.get_endpoint_type( - str(url) - ) + endpoint_type: EndpointType = HttpPassThroughEndpointHelpers.get_endpoint_type(str(url)) if custom_body: _parsed_body = custom_body @@ -701,9 +671,7 @@ async def pass_through_request( # noqa: PLR0915 requested_query_params_str = None if requested_query_params: - requested_query_params_str = "&".join( - f"{k}={v}" for k, v in requested_query_params.items() - ) + requested_query_params_str = "&".join(f"{k}={v}" for k, v in requested_query_params.items()) logging_url = str(url) if requested_query_params_str: @@ -721,11 +689,9 @@ async def pass_through_request( # noqa: PLR0915 "headers": headers, }, ) - stream = ( - HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( - parsed_body=_parsed_body, - stream=stream, - ) + stream = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( + parsed_body=_parsed_body, + stream=stream, ) if stream: @@ -742,9 +708,7 @@ async def pass_through_request( # noqa: PLR0915 try: response.raise_for_status() except httpx.HTTPStatusError as e: - raise HTTPException( - status_code=e.response.status_code, detail=await e.response.aread() - ) + raise HTTPException(status_code=e.response.status_code, detail=await e.response.aread()) return StreamingResponse( PassThroughStreamingHandler.chunk_processor( @@ -766,20 +730,16 @@ async def pass_through_request( # noqa: PLR0915 verbose_proxy_logger.debug("request method: {}".format(request.method)) verbose_proxy_logger.debug("request url: {}".format(url)) verbose_proxy_logger.debug("request headers: {}".format(headers)) - verbose_proxy_logger.debug( - "requested_query_params={}".format(requested_query_params) - ) + verbose_proxy_logger.debug("requested_query_params={}".format(requested_query_params)) verbose_proxy_logger.debug("request body: {}".format(_parsed_body)) - response = ( - await HttpPassThroughEndpointHelpers.non_streaming_http_request_handler( - request=request, - async_client=async_client, - url=url, - headers=headers, - requested_query_params=requested_query_params, - _parsed_body=_parsed_body, - ) + response = await HttpPassThroughEndpointHelpers.non_streaming_http_request_handler( + request=request, + async_client=async_client, + url=url, + headers=headers, + requested_query_params=requested_query_params, + _parsed_body=_parsed_body, ) verbose_proxy_logger.debug("response.headers= %s", response.headers) @@ -787,9 +747,7 @@ async def pass_through_request( # noqa: PLR0915 try: response.raise_for_status() except httpx.HTTPStatusError as e: - raise HTTPException( - status_code=e.response.status_code, detail=await e.response.aread() - ) + raise HTTPException(status_code=e.response.status_code, detail=await e.response.aread()) return StreamingResponse( PassThroughStreamingHandler.chunk_processor( @@ -811,9 +769,7 @@ async def pass_through_request( # noqa: PLR0915 try: response.raise_for_status() except httpx.HTTPStatusError as e: - raise HTTPException( - status_code=e.response.status_code, detail=e.response.text - ) + raise HTTPException(status_code=e.response.status_code, detail=e.response.text) if response.status_code >= 300: raise HTTPException(status_code=response.status_code, detail=response.text) @@ -835,6 +791,7 @@ async def pass_through_request( # noqa: PLR0915 logging_obj=logging_obj, cache_hit=False, request_body=_parsed_body, + custom_llm_provider=custom_llm_provider, **kwargs, ) ) @@ -865,9 +822,7 @@ async def pass_through_request( # noqa: PLR0915 api_base=str(url._uri_reference) if url else None, ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.pass_through_endpoint(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.pass_through_endpoint(): Exception occured - {}".format(str(e)) ) ######################################################### @@ -930,6 +885,7 @@ def create_pass_through_route( dependencies: Optional[List] = None, include_subpath: Optional[bool] = False, cost_per_request: Optional[float] = None, + custom_llm_provider: Optional[str] = None, ): # check if target is an adapter.py or a url from litellm._uuid import uuid @@ -965,16 +921,12 @@ def create_pass_through_route( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), query_params: Optional[dict] = None, custom_body: Optional[dict] = None, - stream: Optional[ - bool - ] = None, # if pass-through endpoint is a streaming request + stream: Optional[bool] = None, # if pass-through endpoint is a streaming request subpath: str = "", # captures sub-paths when include_subpath=True ): # Construct the full target URL with subpath if needed - full_target = ( - HttpPassThroughEndpointHelpers.construct_target_url_with_subpath( - base_target=target, subpath=subpath, include_subpath=include_subpath - ) + full_target = HttpPassThroughEndpointHelpers.construct_target_url_with_subpath( + base_target=target, subpath=subpath, include_subpath=include_subpath ) return await pass_through_request( # type: ignore @@ -988,6 +940,7 @@ def create_pass_through_route( stream=stream, custom_body=custom_body, cost_per_request=cost_per_request, + custom_llm_provider=custom_llm_provider, ) return endpoint_func @@ -1644,15 +1597,11 @@ class InitPassThroughEndpointHelpers: def remove_endpoint_routes(endpoint_id: str): """Remove all routes for a specific endpoint ID from the registry""" keys_to_remove = [ - key - for key, value in _registered_pass_through_routes.items() - if value["endpoint_id"] == endpoint_id + key for key, value in _registered_pass_through_routes.items() if value["endpoint_id"] == endpoint_id ] for key in keys_to_remove: del _registered_pass_through_routes[key] - verbose_proxy_logger.debug( - "Removed pass-through route from registry: %s", key - ) + verbose_proxy_logger.debug("Removed pass-through route from registry: %s", key) async def initialize_pass_through_endpoints( @@ -1689,9 +1638,7 @@ async def initialize_pass_through_endpoints( if _path is None: raise ValueError("Path is required for pass-through endpoint") _custom_headers = endpoint.get("headers", None) - _custom_headers = await set_env_variables_in_header( - custom_headers=_custom_headers - ) + _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) _forward_headers = endpoint.get("forward_headers", None) _merge_query_params = endpoint.get("merge_query_params", None) _auth = endpoint.get("auth", None) @@ -1710,9 +1657,7 @@ async def initialize_pass_through_endpoints( continue # Add exact path route - verbose_proxy_logger.debug( - "Initializing pass through endpoint: %s (ID: %s)", _path, endpoint_id - ) + verbose_proxy_logger.debug("Initializing pass through endpoint: %s (ID: %s)", _path, endpoint_id) InitPassThroughEndpointHelpers.add_exact_path_route( app=app, path=_path, @@ -1739,9 +1684,7 @@ async def initialize_pass_through_endpoints( endpoint_id=endpoint_id, ) - verbose_proxy_logger.debug( - "Added new pass through endpoint: %s (ID: %s)", _path, endpoint_id - ) + verbose_proxy_logger.debug("Added new pass through endpoint: %s (ID: %s)", _path, endpoint_id) async def _get_pass_through_endpoints_from_db( @@ -1845,11 +1788,7 @@ async def update_pass_through_endpoints( # Find the index for updating the list endpoint_index = None for idx, endpoint in enumerate(pass_through_endpoint_data): - _endpoint = ( - PassThroughGenericEndpoint(**endpoint) - if isinstance(endpoint, dict) - else endpoint - ) + _endpoint = PassThroughGenericEndpoint(**endpoint) if isinstance(endpoint, dict) else endpoint if _endpoint.id == endpoint_id: endpoint_index = idx break @@ -1857,9 +1796,7 @@ async def update_pass_through_endpoints( if endpoint_index is None: raise HTTPException( status_code=404, - detail={ - "error": f"Could not find index for endpoint with ID '{endpoint_id}'" - }, + detail={"error": f"Could not find index for endpoint with ID '{endpoint_id}'"}, ) # Get the update data as dict, excluding None values for partial updates @@ -1890,13 +1827,9 @@ async def update_pass_through_endpoints( field_value=pass_through_endpoint_data, config_type="general_settings", ) - await update_config_general_settings( - data=updated_data, user_api_key_dict=user_api_key_dict - ) + await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict) - return PassThroughEndpointResponse( - endpoints=[updated_endpoint] if updated_endpoint else [] - ) + return PassThroughEndpointResponse(endpoints=[updated_endpoint] if updated_endpoint else []) @router.post( @@ -1923,9 +1856,7 @@ async def create_pass_through_endpoints( field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict ) except Exception: - response = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=None - ) + response = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None) ## Auto-generate ID if not provided data_dict = data.model_dump() @@ -1943,9 +1874,7 @@ async def create_pass_through_endpoints( field_value=response.field_value, config_type="general_settings", ) - await update_config_general_settings( - data=updated_data, user_api_key_dict=user_api_key_dict - ) + await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict) # Return the created endpoint with the generated ID created_endpoint = PassThroughGenericEndpoint(**data_dict) @@ -1978,9 +1907,7 @@ async def delete_pass_through_endpoints( field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict ) except Exception: - response = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=None - ) + response = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None) ## Update field by removing endpoint pass_through_endpoint_data: Optional[List] = response.field_value @@ -1996,21 +1923,13 @@ async def delete_pass_through_endpoints( if found_endpoint is None: raise HTTPException( status_code=400, - detail={ - "error": "Endpoint with ID '{}' was not found in pass-through endpoint list.".format( - endpoint_id - ) - }, + detail={"error": "Endpoint with ID '{}' was not found in pass-through endpoint list.".format(endpoint_id)}, ) # Find the index for deleting from the list endpoint_index = None for idx, endpoint in enumerate(pass_through_endpoint_data): - _endpoint = ( - PassThroughGenericEndpoint(**endpoint) - if isinstance(endpoint, dict) - else endpoint - ) + _endpoint = PassThroughGenericEndpoint(**endpoint) if isinstance(endpoint, dict) else endpoint if _endpoint.id == endpoint_id: endpoint_index = idx break @@ -2018,9 +1937,7 @@ async def delete_pass_through_endpoints( if endpoint_index is None: raise HTTPException( status_code=400, - detail={ - "error": f"Could not find index for endpoint with ID '{endpoint_id}'" - }, + detail={"error": f"Could not find index for endpoint with ID '{endpoint_id}'"}, ) # Remove the endpoint @@ -2036,9 +1953,7 @@ async def delete_pass_through_endpoints( field_value=pass_through_endpoint_data, config_type="general_settings", ) - await update_config_general_settings( - data=updated_data, user_api_key_dict=user_api_key_dict - ) + await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict) return PassThroughEndpointResponse(endpoints=[response_obj]) @@ -2076,6 +1991,4 @@ async def initialize_pass_through_endpoints_in_db(): Gets all pass-through endpoints from db and initializes them in the proxy server. """ pass_through_endpoints = await _get_pass_through_endpoints_from_db() - await initialize_pass_through_endpoints( - pass_through_endpoints=pass_through_endpoints - ) + await initialize_pass_through_endpoints(pass_through_endpoints=pass_through_endpoints) diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 94517235a0..a819c429f1 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -25,6 +25,9 @@ from .llm_provider_handlers.cohere_passthrough_logging_handler import ( from .llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) +from .llm_provider_handlers.gemini_passthrough_logging_handler import ( + GeminiPassthroughLoggingHandler, +) cohere_passthrough_logging_handler = CoherePassthroughLoggingHandler() @@ -44,13 +47,14 @@ class PassThroughEndpointLogging: # Cohere self.TRACKED_COHERE_ROUTES = ["/v2/chat"] - self.assemblyai_passthrough_logging_handler = ( - AssemblyAIPassthroughLoggingHandler() - ) + self.assemblyai_passthrough_logging_handler = AssemblyAIPassthroughLoggingHandler() # Langfuse self.TRACKED_LANGFUSE_ROUTES = ["/langfuse/"] + # Gemini + self.TRACKED_GEMINI_ROUTES = ["generateContent", "streamGenerateContent"] + # Vertex AI Live API WebSocket self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"] @@ -81,11 +85,7 @@ class PassThroughEndpointLogging: # Handle async logging await logging_obj.async_success_handler( - result=( - json.dumps(result) - if isinstance(result, dict) - else standard_logging_response_object - ), + result=(json.dumps(result) if isinstance(result, dict) else standard_logging_response_object), start_time=start_time, end_time=end_time, cache_hit=False, @@ -103,6 +103,7 @@ class PassThroughEndpointLogging: start_time: datetime, end_time: datetime, cache_hit: bool, + custom_llm_provider: Optional[str] = None, **kwargs, ): return_dict = { @@ -110,22 +111,34 @@ class PassThroughEndpointLogging: "kwargs": kwargs, } standard_logging_response_object: Optional[Any] = None - if self.is_vertex_route(url_route): - vertex_passthrough_logging_handler_result = ( - VertexPassthroughLoggingHandler.vertex_passthrough_handler( - httpx_response=httpx_response, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - **kwargs, - ) + + if self.is_gemini_route(url_route, custom_llm_provider): + gemini_passthrough_logging_handler_result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler( + httpx_response=httpx_response, + response_body=response_body or {}, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, ) - standard_logging_response_object = ( - vertex_passthrough_logging_handler_result["result"] + standard_logging_response_object = gemini_passthrough_logging_handler_result["result"] + kwargs = gemini_passthrough_logging_handler_result["kwargs"] + elif self.is_vertex_route(url_route): + vertex_passthrough_logging_handler_result = VertexPassthroughLoggingHandler.vertex_passthrough_handler( + httpx_response=httpx_response, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + **kwargs, ) + standard_logging_response_object = vertex_passthrough_logging_handler_result["result"] kwargs = vertex_passthrough_logging_handler_result["kwargs"] elif self.is_anthropic_route(url_route): anthropic_passthrough_logging_handler_result = ( @@ -142,28 +155,22 @@ class PassThroughEndpointLogging: ) ) - standard_logging_response_object = ( - anthropic_passthrough_logging_handler_result["result"] - ) + standard_logging_response_object = anthropic_passthrough_logging_handler_result["result"] kwargs = anthropic_passthrough_logging_handler_result["kwargs"] elif self.is_cohere_route(url_route): - cohere_passthrough_logging_handler_result = ( - cohere_passthrough_logging_handler.passthrough_chat_handler( - httpx_response=httpx_response, - response_body=response_body or {}, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - request_body=request_body, - **kwargs, - ) - ) - standard_logging_response_object = ( - cohere_passthrough_logging_handler_result["result"] + cohere_passthrough_logging_handler_result = cohere_passthrough_logging_handler.passthrough_chat_handler( + httpx_response=httpx_response, + response_body=response_body or {}, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, ) + standard_logging_response_object = cohere_passthrough_logging_handler_result["result"] kwargs = cohere_passthrough_logging_handler_result["kwargs"] elif self.is_openai_route(url_route) and self._is_supported_openai_endpoint( url_route @@ -172,24 +179,21 @@ class PassThroughEndpointLogging: OpenAIPassthroughLoggingHandler, ) - openai_passthrough_logging_handler_result = ( - OpenAIPassthroughLoggingHandler.openai_passthrough_handler( - httpx_response=httpx_response, - response_body=response_body or {}, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - request_body=request_body, - **kwargs, - ) - ) - standard_logging_response_object = ( - openai_passthrough_logging_handler_result["result"] + openai_passthrough_logging_handler_result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=httpx_response, + response_body=response_body or {}, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, ) + standard_logging_response_object = openai_passthrough_logging_handler_result["result"] kwargs = openai_passthrough_logging_handler_result["kwargs"] + elif self.is_vertex_ai_live_route(url_route): from .llm_provider_handlers.vertex_ai_live_passthrough_logging_handler import ( VertexAILivePassthroughLoggingHandler, @@ -216,6 +220,7 @@ class PassThroughEndpointLogging: return_dict[ "standard_logging_response_object" ] = standard_logging_response_object + return_dict["kwargs"] = kwargs return return_dict @@ -231,21 +236,13 @@ class PassThroughEndpointLogging: cache_hit: bool, request_body: dict, passthrough_logging_payload: PassthroughStandardLoggingPayload, + custom_llm_provider: Optional[str] = None, **kwargs, ): - standard_logging_response_object: Optional[ - PassThroughEndpointLoggingResultValues - ] = None - logging_obj.model_call_details[ - "passthrough_logging_payload" - ] = passthrough_logging_payload + standard_logging_response_object: Optional[PassThroughEndpointLoggingResultValues] = None + logging_obj.model_call_details["passthrough_logging_payload"] = passthrough_logging_payload if self.is_assemblyai_route(url_route): - if ( - AssemblyAIPassthroughLoggingHandler._should_log_request( - httpx_response.request.method - ) - is not True - ): + if AssemblyAIPassthroughLoggingHandler._should_log_request(httpx_response.request.method) is not True: return self.assemblyai_passthrough_logging_handler.assemblyai_passthrough_logging_handler( httpx_response=httpx_response, @@ -263,30 +260,25 @@ class PassThroughEndpointLogging: # Don't log langfuse pass-through requests return else: - normalized_llm_passthrough_logging_payload = ( - self.normalize_llm_passthrough_logging_payload( - httpx_response=httpx_response, - response_body=response_body, - request_body=request_body, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - **kwargs, - ) - ) - standard_logging_response_object = ( - normalized_llm_passthrough_logging_payload[ - "standard_logging_response_object" - ] + normalized_llm_passthrough_logging_payload = self.normalize_llm_passthrough_logging_payload( + httpx_response=httpx_response, + response_body=response_body, + request_body=request_body, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + custom_llm_provider=custom_llm_provider, + **kwargs, ) + standard_logging_response_object = normalized_llm_passthrough_logging_payload[ + "standard_logging_response_object" + ] kwargs = normalized_llm_passthrough_logging_payload["kwargs"] if standard_logging_response_object is None: - standard_logging_response_object = StandardPassThroughResponseObject( - response=httpx_response.text - ) + standard_logging_response_object = StandardPassThroughResponseObject(response=httpx_response.text) kwargs = self._set_cost_per_request( logging_obj=logging_obj, @@ -352,10 +344,16 @@ class PassThroughEndpointLogging: return False parsed_url = urlparse(url_route) return parsed_url.hostname and ( - "api.openai.com" in parsed_url.hostname - or "openai.azure.com" in parsed_url.hostname + "api.openai.com" in parsed_url.hostname or "openai.azure.com" in parsed_url.hostname ) + def is_gemini_route(self, url_route: str, custom_llm_provider: Optional[str] = None): + """Check if the URL route is a Gemini API route.""" + for route in self.TRACKED_GEMINI_ROUTES: + if route in url_route and custom_llm_provider == "gemini": + return True + return False + def _is_supported_openai_endpoint(self, url_route: str) -> bool: """Check if the OpenAI endpoint is supported by the passthrough logging handler.""" from .llm_provider_handlers.openai_passthrough_logging_handler import ( @@ -386,11 +384,7 @@ class PassThroughEndpointLogging: # Check if cost per request is set ######################################################### if passthrough_logging_payload.get("cost_per_request") is not None: - kwargs["response_cost"] = passthrough_logging_payload.get( - "cost_per_request" - ) - logging_obj.model_call_details[ - "response_cost" - ] = passthrough_logging_payload.get("cost_per_request") + kwargs["response_cost"] = passthrough_logging_payload.get("cost_per_request") + logging_obj.model_call_details["response_cost"] = passthrough_logging_payload.get("cost_per_request") return kwargs diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py new file mode 100644 index 0000000000..6f87d8f6ab --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py @@ -0,0 +1,287 @@ +import json +import os +import sys +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.gemini_passthrough_logging_handler import ( + GeminiPassthroughLoggingHandler, +) +from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, +) +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + PassthroughStandardLoggingPayload, +) + + +class TestGeminiPassthroughLoggingHandler: + """Test the Gemini passthrough logging handler for cost tracking.""" + + def setup_method(self): + """Set up test fixtures""" + self.start_time = datetime.now() + self.end_time = datetime.now() + self.handler = GeminiPassthroughLoggingHandler() + + # Mock Gemini generateContent response + self.mock_gemini_response = { + "candidates": [ + { + "content": {"parts": [{"text": "Hello! How can I help you today?"}], "role": "model"}, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + {"category": "HARM_CATEGORY_HARASSMENT", "probability": "NEGLIGIBLE"}, + {"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}, + {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "probability": "NEGLIGIBLE"}, + {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "probability": "NEGLIGIBLE"}, + ], + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 8, "totalTokenCount": 18}, + } + + def _create_mock_httpx_response(self) -> httpx.Response: + """Create a mock httpx.Response for testing""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.text = json.dumps(self.mock_gemini_response) + mock_response.json.return_value = self.mock_gemini_response + mock_response.headers = {"content-type": "application/json"} + return mock_response + + def _create_mock_logging_obj(self) -> LiteLLMLoggingObj: + """Create a mock logging object for testing""" + mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {} + mock_logging_obj.optional_params = {} + mock_logging_obj.litellm_call_id = "test-call-id-123" + return mock_logging_obj + + def _create_passthrough_logging_payload(self) -> PassthroughStandardLoggingPayload: + """Create a mock passthrough logging payload for testing""" + return PassthroughStandardLoggingPayload( + url="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", + request_body={"contents": [{"parts": [{"text": "Hello"}]}]}, + request_method="POST", + ) + + def test_is_gemini_route(self): + """Test that Gemini routes are correctly identified""" + from litellm.proxy.pass_through_endpoints.success_handler import PassThroughEndpointLogging + + handler = PassThroughEndpointLogging() + + # Test generateContent endpoint + assert ( + handler.is_gemini_route( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", + custom_llm_provider="gemini", + ) + is True + ) + + # Test streamGenerateContent endpoint + assert ( + handler.is_gemini_route( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:streamGenerateContent", + custom_llm_provider="gemini", + ) + is True + ) + + # Test non-Gemini endpoint + assert ( + handler.is_gemini_route("https://api.openai.com/v1/chat/completions", custom_llm_provider="openai") is False + ) + + def test_extract_model_from_url(self): + """Test that model is correctly extracted from Gemini URLs""" + # Test generateContent endpoint + model = GeminiPassthroughLoggingHandler.extract_model_from_url( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent" + ) + assert model == "gemini-1.5-flash" + + # Test streamGenerateContent endpoint + model = GeminiPassthroughLoggingHandler.extract_model_from_url( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro:streamGenerateContent" + ) + assert model == "gemini-1.5-pro" + + @patch("litellm.completion_cost") + @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload") + def test_gemini_passthrough_handler_success(self, mock_get_standard_logging, mock_completion_cost): + """Test successful cost tracking for Gemini generateContent endpoint""" + # Arrange + mock_completion_cost.return_value = 0.000045 + mock_get_standard_logging.return_value = {"test": "logging_payload"} + + mock_httpx_response = self._create_mock_httpx_response() + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gemini-1.5-flash", + } + + # Act + result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=self.mock_gemini_response, + logging_obj=mock_logging_obj, + url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"contents": [{"parts": [{"text": "Hello"}]}]}, + **kwargs, + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == 0.000045 + assert result["kwargs"]["model"] == "gemini-1.5-flash" + assert result["kwargs"]["custom_llm_provider"] == "gemini" + + # Verify cost calculation was called + mock_completion_cost.assert_called_once() + + # Verify logging object was updated + assert mock_logging_obj.model_call_details["response_cost"] == 0.000045 + assert mock_logging_obj.model_call_details["model"] == "gemini-1.5-flash" + assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini" + + @patch("litellm.completion_cost") + def test_gemini_passthrough_handler_streaming(self, mock_completion_cost): + """Test cost tracking for Gemini streaming endpoint""" + # Arrange + mock_completion_cost.return_value = 0.000030 + + # Mock streaming response chunks + mock_chunks = [ + {"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]}, + {"candidates": [{"content": {"parts": [{"text": " there!"}]}}]}, + ] + + mock_httpx_response = self._create_mock_httpx_response() + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gemini-1.5-flash", + } + + # Act - Use generateContent URL since that's what the handler processes + result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=mock_chunks, + logging_obj=mock_logging_obj, + url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"contents": [{"parts": [{"text": "Hello"}]}]}, + **kwargs, + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == 0.000030 + assert result["kwargs"]["model"] == "gemini-1.5-flash" + assert result["kwargs"]["custom_llm_provider"] == "gemini" + + # Verify cost calculation was called + mock_completion_cost.assert_called_once() + + def test_gemini_passthrough_handler_non_gemini_route(self): + """Test that non-Gemini routes return None""" + mock_httpx_response = self._create_mock_httpx_response() + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gpt-4o", + } + + # Act + result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=self.mock_gemini_response, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/chat/completions", # Non-Gemini route (no generateContent) + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}, + **kwargs, + ) + + # Assert - the handler should return a dict with None result for non-Gemini routes + assert result is not None + assert result["result"] is None + assert "kwargs" in result + + @pytest.mark.asyncio + async def test_pass_through_success_handler_gemini_routing(self): + """Test that the success handler correctly routes Gemini requests to the Gemini handler""" + handler = PassThroughEndpointLogging() + + # Mock the logging object + mock_logging_obj = self._create_mock_logging_obj() + + # Mock the _handle_logging method to capture the call + handler._handle_logging = AsyncMock() + + # Mock httpx response + mock_response = self._create_mock_httpx_response() + + # Create passthrough logging payload + passthrough_logging_payload = self._create_passthrough_logging_payload() + + # Call the success handler with Gemini route and provider + result = await handler.pass_through_async_success_handler( + httpx_response=mock_response, + response_body=self.mock_gemini_response, + logging_obj=mock_logging_obj, + url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"contents": [{"parts": [{"text": "Hello"}]}]}, + passthrough_logging_payload=passthrough_logging_payload, + custom_llm_provider="gemini", + ) + + # Assert - The success handler returns None on success (following the pattern from other tests) + assert result is None + + # Verify that the logging object has the cost set (from Gemini handler) + assert mock_logging_obj.model_call_details["response_cost"] is not None + assert mock_logging_obj.model_call_details["model"] == "gemini-1.5-flash" + assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini" + + # Verify that _handle_logging was called with the correct kwargs + handler._handle_logging.assert_called_once() + call_kwargs = handler._handle_logging.call_args[1] + assert call_kwargs["response_cost"] is not None + assert call_kwargs["model"] == "gemini-1.5-flash" + assert call_kwargs["custom_llm_provider"] == "gemini"