diff --git a/api/api_key/schemas.py b/api/api_key/schemas.py index 3e646941..bab38bda 100644 --- a/api/api_key/schemas.py +++ b/api/api_key/schemas.py @@ -61,8 +61,10 @@ def validate_object_type(_, __, type_): """ Limit which types of objects we can manipulate with API keys. """ - if type_ not in ("images", "chutes", "invocations"): - raise ValueError("Invalid object_type, must be one of images, chutes, invocations") + if type_ not in ("images", "chutes", "invocations", "billing"): + raise ValueError( + "Invalid object_type, must be one of images, chutes, invocations, billing" + ) return type_ diff --git a/api/invocation/router.py b/api/invocation/router.py index e78b2a63..7efdb135 100644 --- a/api/invocation/router.py +++ b/api/invocation/router.py @@ -766,6 +766,7 @@ async def _stream_with_first_chunk(): request, { "X-Chutes-InvocationID": parent_invocation_id, + "Inference-Id": parent_invocation_id, "Cache-Control": "no-cache, no-transform", "X-Accel-Buffering": "no", }, @@ -815,6 +816,7 @@ async def _streamfile(): request, { "X-Chutes-InvocationID": parent_invocation_id, + "Inference-Id": parent_invocation_id, "Cache-Control": "no-cache, no-transform", "X-Accel-Buffering": "no", }, @@ -825,7 +827,11 @@ async def _streamfile(): content=result["text"], media_type=result["content_type"], headers=build_response_headers( - request, {"X-Chutes-InvocationID": parent_invocation_id} + request, + { + "X-Chutes-InvocationID": parent_invocation_id, + "Inference-Id": parent_invocation_id, + }, ), ) else: @@ -837,6 +843,7 @@ async def _streamfile(): { "Content-type": "application/json", "X-Chutes-InvocationID": parent_invocation_id, + "Inference-Id": parent_invocation_id, }, ), ) diff --git a/api/main.py b/api/main.py index a50a4e51..357f21e3 100644 --- a/api/main.py +++ b/api/main.py @@ -365,7 +365,11 @@ async def host_router_middleware(request: Request, call_next): if request.state.auth_method != "invoke": # Handle /users/me/* paths specially for OAuth scope checking - if request.url.path.startswith("/users/me"): + if request.url.path == "/partners/huggingface/billing": + request.state.auth_method = "read" + request.state.auth_object_type = "billing" + request.state.auth_object_id = "__self__" + elif request.url.path.startswith("/users/me"): if "/balance" in request.url.path: request.state.auth_object_type = "billing" elif "/quota" in request.url.path: diff --git a/api/payment/router.py b/api/payment/router.py index f899ea9b..570eae78 100644 --- a/api/payment/router.py +++ b/api/payment/router.py @@ -3,14 +3,15 @@ """ import orjson as json +from decimal import Decimal, InvalidOperation, ROUND_HALF_UP from typing import Optional -from pydantic import BaseModel +from pydantic import BaseModel, Field from fastapi import APIRouter, status, HTTPException, Depends, Request from sqlalchemy.ext.asyncio import AsyncSession from api.gpu import SUPPORTED_GPUS, COMPUTE_UNIT_PRICE_BASIS, COMPUTE_MIN from api.config import settings from api.fmv.fetcher import get_fetcher -from api.database import get_db_session +from api.database import get_db_session, get_inv_session from api.user.util import refund_deposit from api.user.schemas import User from api.user.service import get_current_user @@ -24,6 +25,34 @@ class ReturnDepositArgs(BaseModel): address: str +class HuggingFaceBillingArgs(BaseModel): + requestIds: list[str] = Field(default_factory=list, max_length=10000) + + +class HuggingFaceBillingItem(BaseModel): + requestId: str + costNanoUsd: int + + +class HuggingFaceBillingResponse(BaseModel): + requests: list[HuggingFaceBillingItem] + + +NANO_USD = Decimal("1000000000") + + +def usd_to_nano_usd(value) -> int: + """ + Convert a USD amount to HF's integer nano-USD billing unit. + """ + try: + usd_amount = Decimal(str(value if value is not None else 0)) + except (InvalidOperation, ValueError): + return 0 + nano_usd = (usd_amount * NANO_USD).quantize(Decimal("1"), rounding=ROUND_HALF_UP) + return max(int(nano_usd), 0) + + @router.get("/daily_revenue_summary") async def get_daily_revenue_summary( days: Optional[int] = 90, db: AsyncSession = Depends(get_db_session) @@ -145,6 +174,65 @@ async def get_pricing(): } +@router.post( + "/partners/huggingface/billing", + response_model=HuggingFaceBillingResponse, +) +async def huggingface_billing( + args: HuggingFaceBillingArgs, + current_user: User = Depends(get_current_user()), +): + """ + Hugging Face billing callback. + + HF sends request IDs copied from the Inference-Id response header. Those IDs + are Chutes parent_invocation_id values, and the persisted metrics["b"] value + is the actual USD amount charged to the authenticated Chutes account. + Missing/unknown request IDs intentionally return 0. + """ + request_ids = args.requestIds or [] + unique_ids = list(dict.fromkeys(request_ids)) + if not unique_ids: + return HuggingFaceBillingResponse(requests=[]) + + async with get_inv_session() as session: + result = await session.execute( + text( + """ + SELECT + parent_invocation_id, + SUM( + CASE + WHEN jsonb_typeof(metrics->'b') = 'number' + THEN GREATEST((metrics->>'b')::numeric, 0) + ELSE 0 + END + ) AS usd_amount + FROM invocations + WHERE user_id = :user_id + AND parent_invocation_id = ANY(:request_ids) + GROUP BY parent_invocation_id + """ + ), + { + "user_id": current_user.user_id, + "request_ids": unique_ids, + }, + ) + rows = result.fetchall() + + costs_by_id = {row.parent_invocation_id: usd_to_nano_usd(row.usd_amount) for row in rows} + return HuggingFaceBillingResponse( + requests=[ + HuggingFaceBillingItem( + requestId=request_id, + costNanoUsd=costs_by_id.get(request_id, 0), + ) + for request_id in request_ids + ] + ) + + @router.post("/return_developer_deposit") async def return_developer_deposit( args: ReturnDepositArgs, diff --git a/tests/unit/test_huggingface_billing.py b/tests/unit/test_huggingface_billing.py new file mode 100644 index 00000000..487ca753 --- /dev/null +++ b/tests/unit/test_huggingface_billing.py @@ -0,0 +1,25 @@ +from api.api_key.schemas import APIKeyScope, Action +from api.payment.router import usd_to_nano_usd + + +def test_usd_to_nano_usd_is_deterministic(): + assert usd_to_nano_usd("0") == 0 + assert usd_to_nano_usd("0.000218875") == 218875 + assert usd_to_nano_usd("1.234567891") == 1234567891 + assert usd_to_nano_usd("0.0000000005") == 1 + assert usd_to_nano_usd("-1") == 0 + assert usd_to_nano_usd("not-a-number") == 0 + + +def test_api_keys_can_be_scoped_to_billing_read(): + scope = APIKeyScope(object_type="billing", action=Action.READ) + assert scope.object_type == "billing" + + +def test_unknown_api_key_scope_is_still_rejected(): + try: + APIKeyScope(object_type="huggingface", action=Action.READ) + except ValueError as exc: + assert "billing" in str(exc) + else: + raise AssertionError("unknown API key scope should be rejected")