diff --git a/router/auth.py b/router/auth.py index 8c6cd04d..351477f4 100644 --- a/router/auth.py +++ b/router/auth.py @@ -370,7 +370,7 @@ async def revert_pay_for_request( async def adjust_payment_for_tokens( - key: ApiKey, response_data: dict, session: AsyncSession + key: ApiKey, response_data: dict, session: AsyncSession, deducted_max_cost: int ) -> dict: """ Adjusts the payment based on token usage in the response. @@ -378,20 +378,19 @@ async def adjust_payment_for_tokens( Returns cost data to be included in the response. """ model = response_data.get("model", "unknown") - max_cost = get_max_cost_for_model(model=model) logger.debug( "Starting payment adjustment for tokens", extra={ "key_hash": key.hashed_key[:8] + "...", "model": model, - "max_cost": max_cost, + "deducted_max_cost": deducted_max_cost, "current_balance": key.balance, "has_usage": "usage" in response_data, }, ) - match calculate_cost(response_data, max_cost): + match calculate_cost(response_data, deducted_max_cost): case MaxCostData() as cost: logger.debug( "Using max cost data (no token adjustment)", @@ -406,7 +405,7 @@ async def adjust_payment_for_tokens( case CostData() as cost: # If token-based pricing is enabled and base cost is 0, use token-based cost # Otherwise, token cost is additional to the base cost - cost_difference = cost.total_msats - max_cost + cost_difference = cost.total_msats - deducted_max_cost logger.info( "Calculated token-based cost", @@ -414,7 +413,7 @@ async def adjust_payment_for_tokens( "key_hash": key.hashed_key[:8] + "...", "model": model, "token_cost": cost.total_msats, - "max_cost": max_cost, + "deducted_max_cost": deducted_max_cost, "cost_difference": cost_difference, "input_msats": cost.input_msats, "output_msats": cost.output_msats, @@ -468,7 +467,7 @@ async def adjust_payment_for_tokens( await session.commit() if result.rowcount: - cost.total_msats = max_cost + cost_difference + cost.total_msats = deducted_max_cost + cost_difference await session.refresh(key) logger.info( @@ -513,7 +512,7 @@ async def adjust_payment_for_tokens( ) await session.exec(refund_stmt) # type: ignore[call-overload] await session.commit() - cost.total_msats = max_cost - refund + cost.total_msats = deducted_max_cost - refund await session.refresh(key) logger.info( diff --git a/router/payment/helpers.py b/router/payment/helpers.py index d7d0a5bb..6b58882b 100644 --- a/router/payment/helpers.py +++ b/router/payment/helpers.py @@ -47,7 +47,7 @@ def get_cost_per_request(model: str | None = None) -> int: return COST_PER_REQUEST -def check_token_balance(headers: dict, body: dict) -> None: +def check_token_balance(headers: dict, body: dict, max_cost_for_model: int) -> None: if x_cashu := headers.get("x-cashu", None): cashu_token = x_cashu logger.debug( @@ -90,20 +90,18 @@ def check_token_balance(headers: dict, body: dict) -> None: if cashu_token.startswith("sk-"): return - cost = get_cost_per_request(model=body.get("model", None)) - token_obj = deserialize_token_from_string(cashu_token) amount_msat = ( token_obj.amount if token_obj.unit == "msat" else token_obj.amount * 1000 ) - if cost > amount_msat: + if max_cost_for_model > amount_msat: raise HTTPException( status_code=413, detail={ "reason": "Insufficient balance", - "amount_required_msat": cost, + "amount_required_msat": max_cost_for_model, "model": body.get("model", "unknown"), "type": "minimum_balance_required", }, diff --git a/router/proxy.py b/router/proxy.py index e588c4f5..0e772b1b 100644 --- a/router/proxy.py +++ b/router/proxy.py @@ -19,6 +19,7 @@ from .payment.helpers import ( UPSTREAM_BASE_URL, check_token_balance, create_error_response, + get_cost_per_request, prepare_upstream_headers, ) from .payment.x_cashu import x_cashu_handler @@ -28,7 +29,7 @@ proxy_router = APIRouter() async def handle_streaming_chat_completion( - response: httpx.Response, key: ApiKey, session: AsyncSession + response: httpx.Response, key: ApiKey, max_cost_for_model: int ) -> StreamingResponse: """Handle streaming chat completion responses with token-based pricing.""" logger.info( @@ -40,7 +41,7 @@ async def handle_streaming_chat_completion( }, ) - async def stream_with_cost() -> AsyncGenerator[bytes, None]: + async def stream_with_cost(max_cost_for_model: int) -> AsyncGenerator[bytes, None]: # Store all chunks to analyze stored_chunks = [] @@ -103,7 +104,10 @@ async def handle_streaming_chat_completion( if fresh_key: try: cost_data = await adjust_payment_for_tokens( - fresh_key, data, new_session + fresh_key, + data, + new_session, + max_cost_for_model, ) logger.info( "Token adjustment completed for streaming", @@ -140,14 +144,17 @@ async def handle_streaming_chat_completion( ) return StreamingResponse( - stream_with_cost(), + stream_with_cost(max_cost_for_model), status_code=response.status_code, headers=dict(response.headers), ) async def handle_non_streaming_chat_completion( - response: httpx.Response, key: ApiKey, session: AsyncSession + response: httpx.Response, + key: ApiKey, + session: AsyncSession, + deducted_max_cost: int, ) -> Response: """Handle non-streaming chat completion responses with token-based pricing.""" logger.info( @@ -172,7 +179,9 @@ async def handle_non_streaming_chat_completion( }, ) - cost_data = await adjust_payment_for_tokens(key, response_json, session) + cost_data = await adjust_payment_for_tokens( + key, response_json, session, deducted_max_cost + ) response_json["cost"] = cost_data logger.info( @@ -239,6 +248,7 @@ async def forward_to_upstream( headers: dict, request_body: bytes | None, key: ApiKey, + max_cost_for_model: int, session: AsyncSession, ) -> Response | StreamingResponse: """Forward request to upstream and handle the response.""" @@ -338,7 +348,9 @@ async def forward_to_upstream( if is_streaming and response.status_code == 200: # Process streaming response and extract cost from the last chunk - result = await handle_streaming_chat_completion(response, key, session) + result = await handle_streaming_chat_completion( + response, key, max_cost_for_model + ) background_tasks = BackgroundTasks() background_tasks.add_task(response.aclose) background_tasks.add_task(client.aclose) @@ -349,7 +361,7 @@ async def forward_to_upstream( # Handle non-streaming response try: return await handle_non_streaming_chat_completion( - response, key, session + response, key, session, max_cost_for_model ) finally: await response.aclose() @@ -479,7 +491,10 @@ async def proxy( media_type="application/json", ) - check_token_balance(headers, request_body_dict) + max_cost_for_model = get_cost_per_request( + model=request_body_dict.get("model", None) + ) + check_token_balance(headers, request_body_dict, max_cost_for_model) # Handle authentication if x_cashu := headers.get("x-cashu", None): @@ -560,7 +575,7 @@ async def proxy( # Forward to upstream and handle response response = await forward_to_upstream( - request, path, headers, request_body, key, session + request, path, headers, request_body, key, max_cost_for_model, session ) if response.status_code != 200: