Merge remote-tracking branch 'upstream/main' into kwsantiago/62-comprehensive-tests

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(

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",
             },

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: