tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-18 04:28:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add generateContent cost tracking (#15014)

Browse Source
This commit is contained in:
Sameer Kankute
2025-10-01 09:03:45 -07:00
committed by GitHub
parent ab00ca2de9
commit 7ec7e5332c
5 changed files with 666 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
+29 -89
View File
@@ -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/<Model>/<endpoint>. Got: "
+ endpoint,
"error": "Model missing from endpoint. Expected format: /model/<Model>/<endpoint>. 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
litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py
+204
View File
@@ -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
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
+54 -141
View File
@@ -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)
litellm/proxy/pass_through_endpoints/success_handler.py
+92 -98
View File
@@ -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
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py
+287
View File
@@ -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"