Calculate Streaming Token Usage Costs

Task Description

Language model APIs bill users based on token consumption, tracking prompt (input) and completion (output) tokens separately. In production, usage data arrives via two ingestion paths:

1. Non-streaming requests: Emit a single terminal payload containing aggregate token counts.
2. Streaming requests: Emit multiple incremental chunk deltas over time as text is generated.

Implement the TokenBillingManager class to ingest usage events, enforce validation rules, and calculate exact costs without floating-point inaccuracies.

Unit of Currency

To avoid floating-point billing errors, all monetary amounts must be stored and computed in nano-USD (1 USD = 1,000,000,000 nano-USD). For example, a rate of $3.00 per 1,000,000 tokens corresponds to 3,000 nano-USD per token.

Class Interface

class TokenBillingManager:
    def __init__(self, rates: dict[str, dict[str, int]]):
        """
        rates: Mapping from model_name to a dict containing:
               - 'input_nano_per_token': int (cost in nano-USD per prompt token)
               - 'output_nano_per_token': int (cost in nano-USD per completion token)
        """
        pass

    def record_event(self, event: dict) -> None:
        """
        Ingests a usage event. Raises ValueError if the event is malformed or invalid.
        """
        pass

    def get_request_cost(self, request_id: str) -> int:
        """
        Returns the accumulated cost for request_id in nano-USD.
        Raises KeyError if request_id does not exist.
        """
        pass

    def get_request_tokens(self, request_id: str) -> tuple[int, int]:
        """
        Returns (prompt_tokens, completion_tokens) for request_id.
        Raises KeyError if request_id does not exist.
        """
        pass

    def get_total_cost(self) -> int:
        """
        Returns the total accumulated cost across all recorded requests in nano-USD.
        """
        pass

Event Schema

Each event is a dict with the following fields:

• "request_id": non-empty str
• "model": str (must exist in rates)
• "is_stream": bool
• "prompt_tokens": int (non-negative)
• "completion_tokens": int (non-negative)
• "is_final": bool (optional, default False; if True, marks a streaming request as complete)

Validation Rules

Raise a ValueError if:

1. Any required field is missing, has an incorrect type (e.g., bool passed for an integer count, or float tokens), or string fields are empty.
2. Either token count is strictly negative (< 0).
3. The model is not registered in rates.
4. A non-streaming event (is_stream = False) is recorded for a request_id that already exists, or any new event arrives for a finalized stream.
5. An event arrives with is_stream = False on a request previously registered as streaming, or vice versa.
6. An event attempts to change the model associated with an active request_id.

Example

rates = {
    "claude-sonnet": {
        "input_nano_per_token": 3000,    # $3.00 / 1M tokens
        "output_nano_per_token": 15000   # $15.00 / 1M tokens
    }
}

manager = TokenBillingManager(rates)

# Non-streaming request
manager.record_event({
    "request_id": "req-1",
    "model": "claude-sonnet",
    "is_stream": False,
    "prompt_tokens": 100,
    "completion_tokens": 50
})
# Cost: 100 * 3000 + 50 * 15000 = 300_000 + 750_000 = 1_050_000 nano-USD ($0.00105)
assert manager.get_request_cost("req-1") == 1_050_000

# Streaming request in two deltas
manager.record_event({
    "request_id": "req-2",
    "model": "claude-sonnet",
    "is_stream": True,
    "prompt_tokens": 200,
    "completion_tokens": 10,
    "is_final": False
})
manager.record_event({
    "request_id": "req-2",
    "model": "claude-sonnet",
    "is_stream": True,
    "prompt_tokens": 0,
    "completion_tokens": 40,
    "is_final": True
})
# Tokens: 200 prompt, 50 completion
# Cost: 200 * 3000 + 50 * 15000 = 600_000 + 750_000 = 1_350_000 nano-USD
assert manager.get_request_tokens("req-2") == (200, 50)
assert manager.get_request_cost("req-2") == 1_350_000

# Total cost: 1_050_000 + 1_350_000 = 2_400_000 nano-USD
assert manager.get_total_cost() == 2_400_000

Constraints

• Number of events: up to 100,000.
• Token counts per event: 0 <= tokens <= 1,000,000,000.
• request_id: length between 1 and 128 characters.
• record_event, get_request_cost, and get_request_tokens should run in O(1) average time.