diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9ed2323a7b..0516a0f0c8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1456,7 +1456,7 @@ class PassThroughGenericEndpoint(LiteLLMPydanticObjectBase): default=False, description="If True, requests to subpaths of the path will be forwarded to the target endpoint. For example, if the path is /bria and include_subpath is True, requests to /bria/v1/text-to-image/base/2.3 will be forwarded to the target endpoint.", ) - input_cost_per_request: float = Field( + cost_per_request: float = Field( default=0.0, description="The USD cost per request to the target endpoint. This is used to calculate the cost of the request to the target endpoint.", ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 53c9c49a47..9ab7c2c7e6 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -12,6 +12,7 @@ import httpx from fastapi import ( APIRouter, Depends, + FastAPI, HTTPException, Request, Response, @@ -519,9 +520,22 @@ async def pass_through_request( # noqa: PLR0915 merge_query_params: Optional[bool] = False, query_params: Optional[dict] = None, stream: Optional[bool] = None, + cost_per_request: Optional[float] = None, ): """ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called + + Args: + request: The incoming request + target: The target URL + custom_headers: The custom headers + user_api_key_dict: The user API key dictionary + custom_body: The custom body + forward_headers: Whether to forward headers + merge_query_params: Whether to merge query params + query_params: The query params + stream: Whether to stream the response + cost_per_request: Optional field - cost per request to the target endpoint """ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_logging_obj @@ -599,6 +613,7 @@ async def pass_through_request( # noqa: PLR0915 url=str(url), request_body=_parsed_body, request_method=getattr(request, "method", None), + cost_per_request=cost_per_request, ) kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( user_api_key_dict=user_api_key_dict, @@ -848,6 +863,7 @@ def create_pass_through_route( _merge_query_params: Optional[bool] = False, dependencies: Optional[List] = None, include_subpath: Optional[bool] = False, + cost_per_request: Optional[float] = None, ): # check if target is an adapter.py or a url import uuid @@ -906,6 +922,7 @@ def create_pass_through_route( query_params=query_params, stream=stream, custom_body=custom_body, + cost_per_request=cost_per_request, ) return endpoint_func @@ -921,13 +938,14 @@ def _is_streaming_response(response: httpx.Response) -> bool: class InitPassThroughEndpointHelpers: @staticmethod def add_exact_path_route( - app, + app: FastAPI, path: str, target: str, custom_headers: Optional[dict], forward_headers: Optional[bool], merge_query_params: Optional[bool], dependencies: Optional[List], + cost_per_request: Optional[float], ): """Add exact path route for pass-through endpoint""" verbose_proxy_logger.debug( @@ -936,15 +954,16 @@ class InitPassThroughEndpointHelpers: dependencies, ) - app.add_api_route( # type: ignore + app.add_api_route( path=path, - endpoint=create_pass_through_route( # type: ignore + endpoint=create_pass_through_route( path, target, custom_headers, forward_headers, merge_query_params, dependencies, + cost_per_request=cost_per_request, ), methods=["GET", "POST", "PUT", "DELETE", "PATCH"], dependencies=dependencies, @@ -952,13 +971,14 @@ class InitPassThroughEndpointHelpers: @staticmethod def add_subpath_route( - app, + app: FastAPI, path: str, target: str, custom_headers: Optional[dict], forward_headers: Optional[bool], merge_query_params: Optional[bool], dependencies: Optional[List], + cost_per_request: Optional[float], ): """Add wildcard route for sub-paths""" wildcard_path = f"{path}/{{subpath:path}}" @@ -968,9 +988,9 @@ class InitPassThroughEndpointHelpers: dependencies, ) - app.add_api_route( # type: ignore + app.add_api_route( path=wildcard_path, - endpoint=create_pass_through_route( # type: ignore + endpoint=create_pass_through_route( path, target, custom_headers, @@ -978,6 +998,7 @@ class InitPassThroughEndpointHelpers: merge_query_params, dependencies, include_subpath=True, + cost_per_request=cost_per_request, ), methods=["GET", "POST", "PUT", "DELETE", "PATCH"], dependencies=dependencies, @@ -1038,6 +1059,7 @@ async def initialize_pass_through_endpoints( forward_headers=_forward_headers, merge_query_params=_merge_query_params, dependencies=_dependencies, + cost_per_request=endpoint.get("cost_per_request", None), ) # Add wildcard route for sub-paths @@ -1050,6 +1072,7 @@ async def initialize_pass_through_endpoints( forward_headers=_forward_headers, merge_query_params=_merge_query_params, dependencies=_dependencies, + cost_per_request=endpoint.get("cost_per_request", None), ) verbose_proxy_logger.debug("Added new pass through endpoint: %s", _path) diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index b1908acdf8..4075f4097d 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -88,6 +88,8 @@ class PassThroughEndpointLogging: cache_hit=False, **kwargs, ) + + async def pass_through_async_success_handler( self, @@ -192,6 +194,13 @@ class PassThroughEndpointLogging: standard_logging_response_object = StandardPassThroughResponseObject( response=httpx_response.text ) + + kwargs = self._set_cost_per_request( + logging_obj=logging_obj, + passthrough_logging_payload=passthrough_logging_payload, + kwargs=kwargs, + ) + await self._handle_logging( logging_obj=logging_obj, @@ -200,6 +209,7 @@ class PassThroughEndpointLogging: start_time=start_time, end_time=end_time, cache_hit=cache_hit, + standard_pass_through_logging_payload=passthrough_logging_payload, **kwargs, ) @@ -234,3 +244,29 @@ class PassThroughEndpointLogging: if route in parsed_url.path: return True return False + + + def _set_cost_per_request( + self, + logging_obj: LiteLLMLoggingObj, + passthrough_logging_payload: PassthroughStandardLoggingPayload, + kwargs: dict, + ): + """ + Helper function to set the cost per request in the logging object + + Only set the cost per request if it's set in the passthrough logging payload. + If it's not set, don't set it in the logging object. + """ + ######################################################### + # 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" + ) + + return kwargs diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index ef4e2a0a25..8a91dd14eb 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -32,3 +32,10 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False): """ The body of the response """ + + cost_per_request: Optional[float] + """ + The cost per request to the target endpoint + + Optional field, we use this for cost tracking only if it's set. + """ diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d83ae59ac5..0565c8e5df 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -343,6 +343,7 @@ def test_add_exact_path_route(): forward_headers=forward_headers, merge_query_params=merge_query_params, dependencies=dependencies, + cost_per_request=None, ) # Verify add_api_route was called with correct parameters @@ -383,6 +384,7 @@ def test_add_subpath_route(): forward_headers=forward_headers, merge_query_params=merge_query_params, dependencies=dependencies, + cost_per_request=None, ) # Verify add_api_route was called with correct parameters @@ -492,3 +494,241 @@ async def test_initialize_pass_through_endpoints_without_include_subpath(): exact_call_args = mock_add_exact_route.call_args[1] assert exact_call_args["path"] == "/test/endpoint" assert exact_call_args["target"] == "http://example.com" + + +def test_set_cost_per_request(): + """ + Test that _set_cost_per_request correctly sets the cost in logging object and kwargs + """ + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + PassthroughStandardLoggingPayload, + ) + + handler = PassThroughEndpointLogging() + + # Mock the logging object + mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {} + + # Test with cost_per_request set + passthrough_logging_payload = PassthroughStandardLoggingPayload( + url="http://example.com/api", + request_body={"test": "data"}, + request_method="POST", + cost_per_request=0.50, + ) + + kwargs = {"some_existing_key": "value"} + + # Call the method + result_kwargs = handler._set_cost_per_request( + logging_obj=mock_logging_obj, + passthrough_logging_payload=passthrough_logging_payload, + kwargs=kwargs, + ) + + # Verify that response_cost is set in kwargs and logging object + assert result_kwargs["response_cost"] == 0.50 + assert mock_logging_obj.model_call_details["response_cost"] == 0.50 + assert result_kwargs["some_existing_key"] == "value" # Existing kwargs preserved + + +def test_set_cost_per_request_none(): + """ + Test that _set_cost_per_request does nothing when cost_per_request is None + """ + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + PassthroughStandardLoggingPayload, + ) + + handler = PassThroughEndpointLogging() + + # Mock the logging object + mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {} + + # Test with cost_per_request not set (None) + passthrough_logging_payload = PassthroughStandardLoggingPayload( + url="http://example.com/api", + request_body={"test": "data"}, + request_method="POST", + cost_per_request=None, + ) + + kwargs = {"some_existing_key": "value"} + + # Call the method + result_kwargs = handler._set_cost_per_request( + logging_obj=mock_logging_obj, + passthrough_logging_payload=passthrough_logging_payload, + kwargs=kwargs, + ) + + # Verify that response_cost is not set + assert "response_cost" not in result_kwargs + assert "response_cost" not in mock_logging_obj.model_call_details + assert result_kwargs["some_existing_key"] == "value" # Existing kwargs preserved + + +@pytest.mark.asyncio +async def test_pass_through_success_handler_with_cost_per_request(): + """ + Test that the success handler correctly processes cost_per_request + """ + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + PassthroughStandardLoggingPayload, + ) + + handler = PassThroughEndpointLogging() + + # Mock the logging object + mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {} + + # Mock the _handle_logging method to capture the call + handler._handle_logging = AsyncMock() + + # Mock httpx response + mock_response = MagicMock(spec=httpx.Response) + mock_response.text = '{"status": "success", "data": "test"}' + + # Create passthrough logging payload with cost_per_request + passthrough_logging_payload = PassthroughStandardLoggingPayload( + url="http://example.com/api", + request_body={"test": "data"}, + request_method="POST", + cost_per_request=1.25, + ) + + start_time = datetime.now() + end_time = datetime.now() + + # Call the success handler + result = await handler.pass_through_async_success_handler( + httpx_response=mock_response, + response_body={"status": "success", "data": "test"}, + logging_obj=mock_logging_obj, + url_route="http://example.com/api", + result="", + start_time=start_time, + end_time=end_time, + cache_hit=False, + request_body={"test": "data"}, + passthrough_logging_payload=passthrough_logging_payload, + ) + + # Verify that the logging object has the cost set + assert mock_logging_obj.model_call_details["response_cost"] == 1.25 + + # 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"] == 1.25 + + +@pytest.mark.asyncio +async def test_create_pass_through_route_with_cost_per_request(): + """ + Test that create_pass_through_route correctly passes cost_per_request to the endpoint function + """ + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + create_pass_through_route, + ) + + # Create the endpoint function with cost_per_request + endpoint_func = create_pass_through_route( + endpoint="/test/path", + target="http://example.com", + custom_headers={}, + _forward_headers=True, + _merge_query_params=False, + dependencies=[], + cost_per_request=3.75, + ) + + # Mock the pass_through_request function to capture its call + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request" + ) as mock_pass_through: + mock_pass_through.return_value = MagicMock() + + # Create mock request + mock_request = MagicMock(spec=Request) + mock_request.path_params = {} + mock_request.query_params = QueryParams({}) + + # Call the endpoint function + await endpoint_func( + request=mock_request, + user_api_key_dict={}, + fastapi_response=MagicMock(), + ) + + # Verify that pass_through_request was called with cost_per_request + mock_pass_through.assert_called_once() + call_kwargs = mock_pass_through.call_args[1] + assert call_kwargs["cost_per_request"] == 3.75 + + +def test_initialize_pass_through_endpoints_with_cost_per_request(): + """ + Test that initialize_pass_through_endpoints correctly passes cost_per_request to route creation + """ + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + + # Mock FastAPI app + mock_app = MagicMock() + + # Test exact path route with cost_per_request + InitPassThroughEndpointHelpers.add_exact_path_route( + app=mock_app, + path="/test/path", + target="http://example.com", + custom_headers={}, + forward_headers=True, + merge_query_params=False, + dependencies=[], + cost_per_request=5.00, + ) + + # Verify add_api_route was called + mock_app.add_api_route.assert_called_once() + call_args = mock_app.add_api_route.call_args[1] + + # Verify the endpoint function was created with cost_per_request + # We can't directly test the internal cost_per_request value, but we can verify + # that the endpoint function was created properly + assert call_args["path"] == "/test/path" + assert callable(call_args["endpoint"]) + + # Reset mock for subpath test + mock_app.reset_mock() + + # Test subpath route with cost_per_request + InitPassThroughEndpointHelpers.add_subpath_route( + app=mock_app, + path="/test/path", + target="http://example.com", + custom_headers={}, + forward_headers=True, + merge_query_params=False, + dependencies=[], + cost_per_request=7.50, + ) + + # Verify add_api_route was called for subpath + mock_app.add_api_route.assert_called_once() + call_args = mock_app.add_api_route.call_args[1] + + # Verify the wildcard path and endpoint function + assert call_args["path"] == "/test/path/{subpath:path}" + assert callable(call_args["endpoint"]) diff --git a/ui/litellm-dashboard/src/components/add_pass_through.tsx b/ui/litellm-dashboard/src/components/add_pass_through.tsx index c5060261f5..efea49e86b 100644 --- a/ui/litellm-dashboard/src/components/add_pass_through.tsx +++ b/ui/litellm-dashboard/src/components/add_pass_through.tsx @@ -55,7 +55,7 @@ const AddPassThroughEndpoint: React.FC = ({ "path": formValues["path"], "target": formValues["target"], "include_subpath": formValues["include_subpath"] || false, - "input_cost_per_request": formValues["input_cost_per_request"] || 0 + "cost_per_request": formValues["cost_per_request"] || 0 } await createPassThroughEndpoint(accessToken, formValues); @@ -185,7 +185,7 @@ const AddPassThroughEndpoint: React.FC = ({ } - name="input_cost_per_request" + name="cost_per_request" >