Cost Governance and Budget Allocation Across Agent Types: Token Spend Is Infrastructure Spend

# Wrong way: single-layer per-request limit# No per-agent-type daily tracking. No per-pipeline cap.# Catches individual runaway calls. Misses routing bugs, loop bugs,# and compounding costs across a multi-agent pipeline.MAX_TOKENS_PER_REQUEST = 2000def call_llm_naive(prompt: str, model: str) -> str:    response = openai_client.chat.completions.create(        model=model,        messages=[{"role": "user", "content": prompt}],        max_tokens=MAX_TOKENS_PER_REQUEST,  # only protection    )    return response.choices[0].message.content# What this misses:# 1. Enrichment agent misconfigured to run on 100% of requests instead of 20%:#    each individual call is under limit, daily spend is 5x budget, nobody notices# 2. Pipeline loop bug: agent calls itself 50 times, each within the per-request limit,#    total pipeline cost 50x expected, no alert fires# 3. Cost attribution: no idea which team, product, or pipeline is driving spend# 4. No forecast: spend only visible on the monthly invoice

# budget_enforcer.py# Three-layer budget enforcement with real-time tracking.# Integrates with the OTel fleet telemetry from Part 1.from __future__ import annotationsimport timefrom dataclasses import dataclass, fieldfrom enum import Enumfrom threading import Lockfrom typing import Optionalfrom opentelemetry import metricsclass BudgetAction(Enum):    ALLOW = "allow"    WARN = "warn"               # usage above warning threshold, proceed    HALT = "halt"               # hard limit hit, stop execution    ESCALATE_TO_HUMAN = "escalate"  # soft limit hit, require approval@dataclassclass BudgetResult:    action: BudgetAction    reason: str    current_spend_usd: float    limit_usd: float    utilization_pct: float@dataclassclass AgentBudgetConfig:    """Budget configuration per agent type. Set by platform team, not product teams."""    agent_type: str    per_request_limit_usd: float       # hard limit per single LLM call    daily_limit_usd: float             # hard limit per agent type per day    warn_threshold_pct: float = 0.80   # warn at 80% of any limit@dataclassclass PipelineBudgetConfig:    """Budget configuration per pipeline execution."""    pipeline_type: str    per_execution_limit_usd: float     # max cost for one end-to-end pipeline run    warn_threshold_pct: float = 0.80class BudgetEnforcer:    """    Enforces three-layer budgets: per-request, per-agent-type, per-pipeline.    Thread-safe. Emits OTel metrics on every decision for fleet dashboard.    """    def __init__(        self,        agent_configs: dict[str, AgentBudgetConfig],        pipeline_configs: dict[str, PipelineBudgetConfig],        meter: metrics.Meter,    ) -> None:        self._agent_configs = agent_configs        self._pipeline_configs = pipeline_configs        self._lock = Lock()        # Daily spend tracking per agent type - reset at midnight UTC        self._agent_daily_spend: dict[str, float] = {}        self._daily_reset_timestamp: float = self._today_midnight_utc()        # Per-pipeline-execution spend tracking        # session_id -> cumulative spend        self._pipeline_spend: dict[str, float] = {}        # OTel metric instruments        self._budget_decision_counter = meter.create_counter(            "gen_ai.fleet.budget_decisions",            description="Budget enforcement decisions by agent, action, and level",        )        self._spend_histogram = meter.create_histogram(            "gen_ai.fleet.spend_usd",            description="Token spend in USD per agent invocation",            unit="USD",        )    def check_and_record(        self,        agent_type: str,        session_id: str,        pipeline_type: str,        estimated_cost_usd: float,    ) -> BudgetResult:        """        Check all three budget layers before allowing an agent to execute.        Call BEFORE the LLM call, not after.        """        with self._lock:            self._maybe_reset_daily(agent_type)            config = self._agent_configs.get(agent_type)            pipeline_config = self._pipeline_configs.get(pipeline_type)            # Layer 1: Per-request limit            if config and estimated_cost_usd > config.per_request_limit_usd:                return self._decision(                    BudgetAction.HALT,                    f"Request cost ${estimated_cost_usd:.4f} exceeds per-request limit ${config.per_request_limit_usd:.4f}",                    estimated_cost_usd,                    config.per_request_limit_usd,                    agent_type, session_id, "per_request",                )            # Layer 2: Per-agent-type daily limit            if config:                daily_spend = self._agent_daily_spend.get(agent_type, 0)                projected_daily = daily_spend + estimated_cost_usd                utilization = projected_daily / config.daily_limit_usd                if projected_daily > config.daily_limit_usd:                    return self._decision(                        BudgetAction.HALT,                        f"Daily limit for {agent_type} would be exceeded: ${projected_daily:.2f} > ${config.daily_limit_usd:.2f}",                        projected_daily,                        config.daily_limit_usd,                        agent_type, session_id, "daily",                    )                if utilization > config.warn_threshold_pct:                    # Record the cost and warn - don't halt                    self._agent_daily_spend[agent_type] = projected_daily                    self._spend_histogram.record(estimated_cost_usd, {"gen_ai.agent.name": agent_type})                    return self._decision(                        BudgetAction.WARN,                        f"Daily budget {utilization*100:.0f}% utilized for {agent_type}",                        projected_daily,                        config.daily_limit_usd,                        agent_type, session_id, "daily",                    )            # Layer 3: Per-pipeline execution limit            if pipeline_config:                pipeline_spend = self._pipeline_spend.get(session_id, 0)                projected_pipeline = pipeline_spend + estimated_cost_usd                if projected_pipeline > pipeline_config.per_execution_limit_usd:                    return self._decision(                        BudgetAction.HALT,                        f"Pipeline execution budget would be exceeded: ${projected_pipeline:.4f} > ${pipeline_config.per_execution_limit_usd:.4f}",                        projected_pipeline,                        pipeline_config.per_execution_limit_usd,                        agent_type, session_id, "pipeline",                    )            # All layers cleared - record spend and allow            self._agent_daily_spend[agent_type] = self._agent_daily_spend.get(agent_type, 0) + estimated_cost_usd            self._pipeline_spend[session_id] = self._pipeline_spend.get(session_id, 0) + estimated_cost_usd            self._spend_histogram.record(estimated_cost_usd, {"gen_ai.agent.name": agent_type})            return BudgetResult(                action=BudgetAction.ALLOW,                reason="Within all budget limits",                current_spend_usd=estimated_cost_usd,                limit_usd=config.per_request_limit_usd if config else float("inf"),                utilization_pct=0.0,            )    def release_session(self, session_id: str) -> float:        """Call when a pipeline session completes. Returns total session spend."""        with self._lock:            return self._pipeline_spend.pop(session_id, 0.0)    def _decision(        self,        action: BudgetAction,        reason: str,        current: float,        limit: float,        agent_type: str,        session_id: str,        level: str,    ) -> BudgetResult:        utilization = current / limit if limit > 0 else 1.0        self._budget_decision_counter.add(1, {            "gen_ai.agent.name": agent_type,            "budget.action": action.value,            "budget.level": level,            "session.id": session_id,        })        return BudgetResult(action, reason, current, limit, utilization)    def _maybe_reset_daily(self, agent_type: str) -> None:        now_midnight = self._today_midnight_utc()        if now_midnight > self._daily_reset_timestamp:            self._agent_daily_spend = {}            self._daily_reset_timestamp = now_midnight    @staticmethod    def _today_midnight_utc() -> float:        import datetime        today = datetime.datetime.now(datetime.timezone.utc).date()        midnight = datetime.datetime.combine(today, datetime.time(), tzinfo=datetime.timezone.utc)        return midnight.timestamp()

# budget_configs.py# Fleet-wide budget configuration. Owned by platform team.# Enforced by BudgetEnforcer at runtime.AGENT_BUDGET_CONFIGS = {    # Frontier model agents: high value, high cost - strict limits    "invoice_extractor": AgentBudgetConfig(        agent_type="invoice_extractor",        per_request_limit_usd=0.05,    # ~3,300 tokens at GPT-4o pricing        daily_limit_usd=50.00,        warn_threshold_pct=0.80,    ),    # Mid-tier agents: standard tasks, cheaper models    "vendor_enricher": AgentBudgetConfig(        agent_type="vendor_enricher",        per_request_limit_usd=0.01,    # uses a smaller model        daily_limit_usd=20.00,        warn_threshold_pct=0.80,    ),    # Classification agents: should use smallest viable model    "invoice_classifier": AgentBudgetConfig(        agent_type="invoice_classifier",        per_request_limit_usd=0.002,   # classification: gpt-4o-mini or equivalent        daily_limit_usd=5.00,        warn_threshold_pct=0.90,       # tighter warning for cheap tasks    ),}PIPELINE_BUDGET_CONFIGS = {    "invoice_processing": PipelineBudgetConfig(        pipeline_type="invoice_processing",        per_execution_limit_usd=0.10,  # full pipeline should cost <10 cents        warn_threshold_pct=0.75,    ),}

# cost_attribution.py# Business unit tagging for cost attribution.# Tags flow through OTel spans to the metrics backend.# Grafana can then answer: "cost by team, by product, by pipeline type."from dataclasses import dataclass@dataclassclass CostAttributionContext:    """Carries business unit metadata into every OTel span."""    team: str           # e.g. "accounts-payable"    product: str        # e.g. "invoice-automation"    cost_center: str    # e.g. "CC-4402"    environment: str    # "production" | "staging" | "development"def attribution_attributes(ctx: CostAttributionContext) -> dict[str, str]:    """Returns OTel span attributes for cost attribution."""    return {        "business.team": ctx.team,        "business.product": ctx.product,        "business.cost_center": ctx.cost_center,        "fleet.environment": ctx.environment,    }# Usage: merge into every span's attribute set alongside gen_ai.agent.* attributes.# In Grafana: sum(gen_ai.fleet.spend_usd) by business.team# Result: monthly spend broken down by team, billable to each cost center.

# model_escalation.py# Tiered model selection: attempt cheap model first, escalate on failure.# Reduces frontier model usage to cases where cheaper models genuinely fail.from langchain_openai import ChatOpenAIfrom typing import Optionalclass TieredModelSelector:    """    Attempts task with a cheaper model first.    Escalates to a frontier model only if:    - Cheap model returns low-confidence output    - Cheap model fails schema validation    - Budget enforcer allows the escalation cost    """    def __init__(        self,        cheap_model: str = "gpt-4o-mini",        frontier_model: str = "gpt-4o-2024-11-20",        confidence_threshold: float = 0.85,    ) -> None:        self.cheap = ChatOpenAI(model=cheap_model, temperature=0)        self.frontier = ChatOpenAI(model=frontier_model, temperature=0)        self.confidence_threshold = confidence_threshold    def invoke(self, prompt: str, budget_enforcer: BudgetEnforcer,               agent_type: str, session_id: str, pipeline_type: str) -> tuple[str, str]:        """        Returns (response_content, model_used).        Falls back to frontier model only if cheap model output is below threshold.        """        # Try cheap model first        cheap_cost_estimate = 0.001  # estimate before call        cheap_check = budget_enforcer.check_and_record(            agent_type, session_id, pipeline_type, cheap_cost_estimate        )        if cheap_check.action == BudgetAction.HALT:            raise RuntimeError(f"Budget halted cheap model call: {cheap_check.reason}")        cheap_response = self.cheap.invoke(prompt)        confidence = self._estimate_confidence(cheap_response.content)        if confidence >= self.confidence_threshold:            return cheap_response.content, "cheap"        # Cheap model confidence too low - escalate        frontier_cost_estimate = 0.01        frontier_check = budget_enforcer.check_and_record(            agent_type, session_id, pipeline_type, frontier_cost_estimate        )        if frontier_check.action == BudgetAction.HALT:            # Can't escalate - return cheap model output with low-confidence flag            return cheap_response.content + "\n[LOW_CONFIDENCE]", "cheap_unescalated"        frontier_response = self.frontier.invoke(prompt)        return frontier_response.content, "frontier"    def _estimate_confidence(self, content: str) -> float:        """        Heuristic confidence estimation.        In production: use logprobs if available, or a dedicated confidence model.        """        if "[UNCERTAIN]" in content or "I'm not sure" in content:            return 0.5        if len(content.strip()) < 10:            return 0.4        return 0.9