scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

Add passthrough cost tracking for veo

Browse source
This commit is contained in:
Sameer Kankute 2025-12-01 14:33:03 +05:30
parent b949ec90db
commit 7dac498efb
7 changed files with 191 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/_types.py
View file

@ -47,6 +47,7 @@ from litellm.types.utils import (
StandardPassThroughResponseObject,
TextCompletionResponse,
)
from litellm.types.videos.main import VideoObject
from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type
@ -3275,6 +3276,7 @@ PassThroughEndpointLoggingResultValues = Union[
TextCompletionResponse,
ImageResponse,
EmbeddingResponse,
VideoObject,
StandardPassThroughResponseObject,
]

9
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -187,13 +187,10 @@ 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"
)
# Get LiteLLM API key from Authorization header for authentication
api_key_to_use = get_litellm_virtual_key(request=request)
user_api_key_dict = await user_api_key_auth(
request=request, api_key=f"Bearer {google_ai_studio_api_key}"
request=request, api_key=api_key_to_use
)
base_target_url = (

38
litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py
View file

@ -7,6 +7,7 @@ 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.gemini.videos.transformation import GeminiVideoConfig
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator as GeminiModelResponseIterator,
)
@ -39,6 +40,43 @@ class GeminiPassthroughLoggingHandler:
request_body: dict,
**kwargs,
) -> PassThroughEndpointLoggingTypedDict:
if "predictLongRunning" in url_route:
model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route)
gemini_video_config = GeminiVideoConfig()
litellm_video_response = gemini_video_config.transform_video_create_response(
model=model,
raw_response=httpx_response,
logging_obj=logging_obj,
custom_llm_provider="gemini",
request_data=request_body,
)
logging_obj.model = model
logging_obj.model_call_details["model"] = model
logging_obj.model_call_details["custom_llm_provider"] = "gemini"
logging_obj.custom_llm_provider = "gemini"
response_cost = litellm.completion_cost(
completion_response=litellm_video_response,
model=model,
custom_llm_provider="gemini",
call_type="create_video",
)
# Set response_cost in _hidden_params to prevent recalculation
if not hasattr(litellm_video_response, "_hidden_params"):
litellm_video_response._hidden_params = {}
litellm_video_response._hidden_params["response_cost"] = response_cost
kwargs["response_cost"] = response_cost
kwargs["model"] = model
kwargs["custom_llm_provider"] = "gemini"
logging_obj.model_call_details["response_cost"] = response_cost
return {
"result": litellm_video_response,
"kwargs": kwargs,
}
if "generateContent" in url_route:
model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route)

43
litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py
View file

@ -14,6 +14,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
from litellm.llms.vertex_ai.vector_stores.search_api.transformation import (
VertexSearchAPIVectorStoreConfig,
)
from litellm.llms.vertex_ai.videos.transformation import VertexAIVideoConfig
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.types.utils import (
Choices,
@ -49,9 +50,49 @@ class VertexPassthroughLoggingHandler:
start_time: datetime,
end_time: datetime,
cache_hit: bool,
request_body: Optional[dict] = None,
**kwargs,
) -> PassThroughEndpointLoggingTypedDict:
if "generateContent" in url_route:
if "predictLongRunning" in url_route:
model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route)
vertex_video_config = VertexAIVideoConfig()
litellm_video_response = vertex_video_config.transform_video_create_response(
model=model,
raw_response=httpx_response,
logging_obj=logging_obj,
custom_llm_provider="vertex_ai",
request_data=request_body,
)
logging_obj.model = model
logging_obj.model_call_details["model"] = model
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai"
logging_obj.custom_llm_provider = "vertex_ai"
response_cost = litellm.completion_cost(
completion_response=litellm_video_response,
model=model,
custom_llm_provider="vertex_ai",
call_type="create_video",
)
# Set response_cost in _hidden_params to prevent recalculation
if not hasattr(litellm_video_response, "_hidden_params"):
litellm_video_response._hidden_params = {}
litellm_video_response._hidden_params["response_cost"] = response_cost
kwargs["response_cost"] = response_cost
kwargs["model"] = model
kwargs["custom_llm_provider"] = "vertex_ai"
logging_obj.model_call_details["response_cost"] = response_cost
return {
"result": litellm_video_response,
"kwargs": kwargs,
}
elif "generateContent" in url_route:
model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route)
instance_of_vertex_llm = litellm.VertexGeminiConfig()

26
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -412,6 +412,31 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
params=requested_query_params,
json=_parsed_body,
)
# Mock httpx response emulating a Google AI video generation operation status
# Attach a dummy request with headers set, so response.request.headers is always present
dummy_request = httpx.Request(
method=request.method,
url=str(url),
headers=headers or {},
params=requested_query_params,
json=_parsed_body,
)
# Ensure the .headers attribute exists and is a dict (httpx will normalize it)
mock_headers = httpx.Headers({"content-type": "application/json"})
response = httpx.Response(
status_code=200,
headers=mock_headers,
json={
"name": "operations/1234567890123456789",
"metadata": {
"@type": "type.googleapis.com/google.ai.generativelanguage.v1beta.GenerateVideoMetadata",
"state": "RUNNING",
"createTime": "2025-01-01T12:00:00Z"
},
"done": False
},
request=dummy_request
)
return response
@staticmethod
@ -737,7 +762,6 @@ async def pass_through_request( # noqa: PLR0915
# Store custom_llm_provider in kwargs and logging object if provided
if custom_llm_provider:
kwargs["custom_llm_provider"] = custom_llm_provider
logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
logging_obj.model_call_details["litellm_params"] = kwargs.get("litellm_params", {})

4
litellm/proxy/pass_through_endpoints/success_handler.py
View file

@ -42,6 +42,7 @@ class PassThroughEndpointLogging:
"streamRawPredict",
"search",
"batchPredictionJobs",
"predictLongRunning",
]
# Anthropic
@ -57,7 +58,7 @@ class PassThroughEndpointLogging:
self.TRACKED_LANGFUSE_ROUTES = ["/langfuse/"]
# Gemini
self.TRACKED_GEMINI_ROUTES = ["generateContent", "streamGenerateContent"]
self.TRACKED_GEMINI_ROUTES = ["generateContent", "streamGenerateContent", "predictLongRunning"]
# Vertex AI Live API WebSocket
self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"]
@ -149,6 +150,7 @@ class PassThroughEndpointLogging:
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
request_body=request_body,
**kwargs,
)
)

79
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py
View file

@ -75,7 +75,9 @@ class TestGeminiPassthroughLoggingHandler:
def test_is_gemini_route(self):
"""Test that Gemini routes are correctly identified"""
from litellm.proxy.pass_through_endpoints.success_handler import PassThroughEndpointLogging
from litellm.proxy.pass_through_endpoints.success_handler import (
PassThroughEndpointLogging,
)
handler = PassThroughEndpointLogging()
@ -285,3 +287,78 @@ class TestGeminiPassthroughLoggingHandler:
assert call_kwargs["response_cost"] is not None
assert call_kwargs["model"] == "gemini-1.5-flash"
assert call_kwargs["custom_llm_provider"] == "gemini"
@patch("litellm.completion_cost")
def test_veo3_passthrough_cost_tracking(self, mock_completion_cost):
"""Test Veo3 video generation cost tracking for passthrough requests"""
# Mock the completion_cost to return the expected video generation cost
# For veo-2.0-generate-001 with 8 seconds: 0.35 * 8 = 2.8
expected_cost = 0.35 * 8.0 # $2.80
mock_completion_cost.return_value = expected_cost
# Mock Veo3 predictLongRunning response
mock_veo_response = {
"name": "operations/1234567890123456789"
}
mock_httpx_response = MagicMock(spec=httpx.Response)
mock_httpx_response.status_code = 200
mock_httpx_response.json.return_value = mock_veo_response
mock_httpx_response.headers = {"content-type": "application/json"}
mock_logging_obj = self._create_mock_logging_obj()
# Request body with durationSeconds
request_body = {
"instances": [{"prompt": "A close up of two people staring at a cryptic drawing on a wall,"}],
"parameters": {"durationSeconds": 8}
}
kwargs = {
"passthrough_logging_payload": PassthroughStandardLoggingPayload(
url="https://generativelanguage.googleapis.com/v1beta/models/veo-2.0-generate-001:predictLongRunning",
request_body=request_body,
request_method="POST",
),
}
# Act
result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler(
httpx_response=mock_httpx_response,
response_body=mock_veo_response,
logging_obj=mock_logging_obj,
url_route="https://generativelanguage.googleapis.com/v1beta/models/veo-2.0-generate-001:predictLongRunning",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body=request_body,
**kwargs,
)
# Assert
assert result is not None
assert "result" in result
assert "kwargs" in result
# Verify the cost is calculated correctly
assert result["kwargs"]["response_cost"] == expected_cost
assert result["kwargs"]["model"] == "veo-2.0-generate-001"
assert result["kwargs"]["custom_llm_provider"] == "gemini"
# Verify completion_cost was called with create_video call_type
mock_completion_cost.assert_called_once()
call_args = mock_completion_cost.call_args
assert call_args.kwargs.get("call_type") == "create_video"
assert call_args.kwargs.get("custom_llm_provider") == "gemini"
assert call_args.kwargs.get("model") == "veo-2.0-generate-001"
# Verify the response object has _hidden_params with response_cost
video_response = result["result"]
assert hasattr(video_response, "_hidden_params")
assert video_response._hidden_params.get("response_cost") == expected_cost
# Verify logging object was updated
assert mock_logging_obj.model_call_details["response_cost"] == expected_cost
assert mock_logging_obj.model_call_details["model"] == "veo-2.0-generate-001"
assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini"