[Feat] Passthrough - Add support for setting custom cost per pass through request (#11870)

* use cost_per_request

* fix cost_per_request

* fixes cost_per_request

* fixes for cost per request for pass through

* ui fix param name

* fixes for _set_cost_per_request

* test cost per request pass through endpoints

* Update tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* tests pass through endpoints

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Ishaan Jaff
2025-06-18 16:24:38 -07:00
committed by GitHub
parent e1fbdde289
commit dfdbfdd71c
7 changed files with 316 additions and 10 deletions
+1 -1
View File
@@ -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.",
)
@@ -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)
@@ -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
@@ -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.
"""
@@ -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"])
@@ -55,7 +55,7 @@ const AddPassThroughEndpoint: React.FC<AddFallbacksProps> = ({
"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<AddFallbacksProps> = ({
</Tooltip>
</span>
}
name="input_cost_per_request"
name="cost_per_request"
>
<InputNumber
min={0}
@@ -79,7 +79,7 @@ export interface passThroughItem {
target: string
headers: object
include_subpath?: boolean
input_cost_per_request?: number
cost_per_request?: number
}
// Password field component for headers