orchestrata_ultimate.providers.cost_tracker

Reference › orchestrata_ultimate › providers › cost_tracker

cost_tracker prices token usage and enforces budgets (specification 011, requirements 010 and 045). CostTracker is used by provider_manager. calculate_cost turns a token usage and a model's pricing into a cost record. set_redis attaches a Redis client and loads the Lua scripts that make budget checks atomic; set_budget sets a limit in dollars for a scope key such as a volk or project; check_budget reports whether an estimated cost fits, and get_remaining_budget returns the remaining amount. record_cost stores a cost with its volk, session, project, provider and model so get_usage can aggregate spend by those scopes. budget_key produces the canonical Redis key for a scope's budget counter.

At a glance

Class diagram

Dependency diagram

API

Classes

CostTracker

Decorators: traces(spec='011', story='US-002', reqs=['REQ-011-010', 'REQ-011-045']) · Line 71

Redis-backed cost tracking with atomic budget enforcement.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/providers/cost_tracker.py

# mypy: ignore-errors
"""Cost tracking and budget enforcement.

SPEC-011 FR-011/FR-012: Redis-backed cost tracking with atomic
budget enforcement via Lua script. Aggregates by volk, session,
project, and portfolio.
"""

import logging
from collections.abc import Awaitable
from datetime import UTC, datetime
from typing import cast

import redis.asyncio as redis

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.protocols import (
    BudgetStatus,
    CostRecord,
    ModelPricing,
    TokenUsage,
    UsageReport,
)

logger = logging.getLogger(__name__)


def _redis_result[T](value: Awaitable[T] | T) -> Awaitable[T]:
    """Narrow a redis-py command result back to its true async type.

    redis-py types command methods (``evalsha``, ``hgetall``, ``hget``, ...)
    as ``Awaitable[T] | T`` because the same command mixin backs both the
    sync and async clients (see ``redis.typing.ResponseT``). ``CostTracker``
    only ever holds an async client, which always returns an awaitable at
    runtime -- this documents that fact for pyright instead of suppressing
    the check.
    """
    return cast(Awaitable[T], value)


_COST_PREFIX = "orchestrata:provider:cost"

# Lua script for atomic budget check (read-only). Enforcement and recording
# share a single canonical budget key (``{_COST_PREFIX}:{scope_key}``) whose
# ``total`` field holds *actual* spend (written by ``record_cost``). The check
# only verifies that adding the estimate stays within the limit; it does NOT
# mutate the key, so the budget always reflects real recorded spend rather than
# accumulated estimates.
_BUDGET_CHECK_SCRIPT = """
local current = tonumber(redis.call('HGET', KEYS[1], 'total') or '0')
local estimate = tonumber(ARGV[1])
local limit = tonumber(ARGV[2])
if current + estimate > limit then
    return 0
end
return 1
"""


@traces(spec="011", story="US-002", reqs=["REQ-011-045"])
def budget_key(scope_key: str) -> str:
    """Canonical Redis key for a scope's budget enforcement counter.

    The same key is read by ``check_budget``/``get_remaining_budget`` and
    incremented with *actual* cost by ``record_cost``.
    """
    return f"{_COST_PREFIX}:{scope_key}"


@traces(spec="011", story="US-002", reqs=["REQ-011-010", "REQ-011-045"])
class CostTracker:
    """Redis-backed cost tracking with atomic budget enforcement.

    Tracks costs per request, per volk, per session, and per project.
    Budget enforcement uses a Lua script for TOCTOU-safe atomicity.
    """

    def __init__(self):
        self._redis: redis.Redis | None = None
        self._budget_limits: dict[str, float] = {}
        self._budget_script_sha: str | None = None

    @traces(spec="011", story="US-002", reqs=["REQ-011-045"])
    async def set_redis(self, client: redis.Redis) -> None:
        """Set Redis client and load Lua scripts."""
        self._redis = client
        self._budget_script_sha = await client.script_load(_BUDGET_CHECK_SCRIPT)

    @traces(spec="011", story="US-006", reqs=["REQ-011-045"])
    def set_budget(self, scope_key: str, limit_usd: float) -> None:
        """Set a budget limit for a scope (volk_id, project_id, etc.)."""
        self._budget_limits[scope_key] = limit_usd

    @traces(spec="011", story="US-002", reqs=["REQ-011-019"])
    def calculate_cost(self, usage: TokenUsage, pricing: ModelPricing) -> CostRecord:
        """Calculate cost from token usage and model pricing."""
        input_cost = (usage.input_tokens / 1_000_000) * pricing.input_per_1m
        output_cost = (usage.output_tokens / 1_000_000) * pricing.output_per_1m
        return CostRecord(
            input_cost=input_cost,
            output_cost=output_cost,
            total_cost=input_cost + output_cost,
            currency="USD",
        )

    async def check_budget(self, scope_key: str, estimated_cost: float) -> bool:
        """Check if estimated cost fits within budget. Returns True if allowed."""
        limit = self._budget_limits.get(scope_key)
        if limit is None:
            return True  # No budget limit set
        if not self._redis or not self._budget_script_sha:
            return True  # No Redis = no enforcement
        cost_key = budget_key(scope_key)
        result = await _redis_result(
            self._redis.evalsha(
                self._budget_script_sha,
                1,
                cost_key,
                str(estimated_cost),
                str(limit),
            )
        )
        return bool(int(result))

    @traces(
        spec="011", story="US-002", reqs=["REQ-011-019", "REQ-011-020", "REQ-011-021"]
    )
    async def record_cost(
        self,
        cost: CostRecord,
        *,
        volk_id: str | None = None,
        session_id: str | None = None,
        project_id: str | None = None,
        provider: str = "",
        model: str = "",
    ) -> None:
        """Record a cost against all applicable scopes."""
        if not self._redis:
            return
        pipe = self._redis.pipeline()
        today = datetime.now(UTC).strftime("%Y-%m-%d")
        # Per-volk daily cost
        if volk_id:
            key = f"{_COST_PREFIX}:{volk_id}:{today}"
            pipe.hincrbyfloat(key, "total", cost.total_cost)
            pipe.hincrbyfloat(key, "input", cost.input_cost)
            pipe.hincrbyfloat(key, "output", cost.output_cost)
            pipe.hincrby(key, "requests", 1)
            # Canonical budget key — actual spend that check_budget enforces
            # against and get_remaining_budget reports. Must agree with the key
            # read in check_budget()/get_remaining_budget().
            budget = budget_key(volk_id)
            pipe.hincrbyfloat(budget, "total", cost.total_cost)
        # Per-project canonical budget key (project_id can also be budgeted)
        if project_id:
            budget = budget_key(project_id)
            pipe.hincrbyfloat(budget, "total", cost.total_cost)
        # Per-session cost
        if session_id:
            key = f"{_COST_PREFIX}:session:{session_id}"
            pipe.hincrbyfloat(key, "total", cost.total_cost)
            pipe.hincrbyfloat(key, "input", cost.input_cost)
            pipe.hincrbyfloat(key, "output", cost.output_cost)
            pipe.hincrby(key, "requests", 1)
        # Per-project cost
        if project_id:
            key = f"{_COST_PREFIX}:project:{project_id}"
            pipe.hincrbyfloat(key, "total", cost.total_cost)
            pipe.hincrbyfloat(key, "input", cost.input_cost)
            pipe.hincrbyfloat(key, "output", cost.output_cost)
            pipe.hincrby(key, "requests", 1)
        # Per-provider cost
        if provider:
            key = f"{_COST_PREFIX}:provider:{provider}:{today}"
            pipe.hincrbyfloat(key, "total", cost.total_cost)
            pipe.hincrby(key, "requests", 1)
        # Per-model cost
        if model:
            key = f"{_COST_PREFIX}:model:{model}:{today}"
            pipe.hincrbyfloat(key, "total", cost.total_cost)
            pipe.hincrby(key, "requests", 1)
        await pipe.execute()

    @traces(spec="011", story="US-002", reqs=["REQ-011-020", "REQ-011-021"])
    async def get_usage(
        self,
        *,
        volk_id: str | None = None,
        session_id: str | None = None,
        project_id: str | None = None,
    ) -> UsageReport:
        """Query aggregated usage for a scope."""
        report = UsageReport()
        if not self._redis:
            return report
        # Find the most specific scope
        if volk_id:
            pattern = f"{_COST_PREFIX}:{volk_id}:*"
        elif session_id:
            pattern = f"{_COST_PREFIX}:session:{session_id}"
        elif project_id:
            pattern = f"{_COST_PREFIX}:project:{project_id}"
        else:
            return report
        async for key in self._redis.scan_iter(match=pattern, count=100):
            data = await _redis_result(self._redis.hgetall(key))
            if not data:
                continue
            total = float(data.get("total", 0))
            requests = int(data.get("requests", 0))
            report.total_cost += total
            report.total_requests += requests
        return report

    @traces(spec="011", story="US-006", reqs=["REQ-011-045"])
    async def get_remaining_budget(
        self,
        scope_key: str,
    ) -> BudgetStatus:
        """Query remaining budget for a scope."""
        limit = self._budget_limits.get(scope_key, 0.0)
        spent = 0.0
        if self._redis:
            cost_key = budget_key(scope_key)
            raw = await _redis_result(self._redis.hget(cost_key, "total"))
            if raw:
                spent = float(raw)
        remaining = max(0.0, limit - spent)
        usage_pct = (spent / limit * 100.0) if limit > 0 else 0.0
        return BudgetStatus(
            scope="custom",
            scope_id=scope_key,
            budget_limit_usd=limit,
            spent_usd=spent,
            remaining_usd=remaining,
            usage_percent=usage_pct,
            warning_threshold=0.8,
            is_warning=usage_pct >= 80.0,
            is_exceeded=spent >= limit and limit > 0,
        )