mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-12 23:05:52 +00:00
[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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user