src/usage.py

"""Token usage extraction helpers for litellm responses."""


def _extract_usage(response) -> dict:
    """Extract token usage from a litellm response object."""
    if hasattr(response, "usage") and response.usage:
        return {
            "prompt_tokens": response.usage.prompt_tokens or 0,
            "completion_tokens": response.usage.completion_tokens or 0,
            "total_tokens": response.usage.total_tokens or 0,
        }
    return _empty_usage()


def _empty_usage() -> dict:
    """Return a zeroed usage dict."""
    return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}


def _sum_usage(*usages: dict) -> dict:
    """Sum multiple usage dicts."""
    return {
        "prompt_tokens": sum(u.get("prompt_tokens", 0) for u in usages),
        "completion_tokens": sum(u.get("completion_tokens", 0) for u in usages),
        "total_tokens": sum(u.get("total_tokens", 0) for u in usages),
    }