Đến bài này, bạn đã đi qua gần hết các khái niệm về cost management trong module này: unit economics, cost attribution, model routing, caching... Bài hands-on này sẽ ghép tất cả lại thành một hệ thống chạy được thật — một token budget framework (framework quản lý ngân sách token) hoàn chỉnh mà một team engineering có thể đưa vào production ngay tuần sau.

Bối cảnh: bạn là kỹ sư platform tại Acme Corp, một công ty phần mềm ở quy mô trung bình. Tháng trước, hóa đơn LLM API của công ty tăng gấp 2.3 lần so với dự kiến, và không ai biết chính xác workflow nào gây ra chuyện đó cho tới khi finance gửi email hỏi. Ban lãnh đạo engineering giao cho team bạn một yêu cầu rất cụ thể: "Tôi không cần các bạn cắt giảm chi phí AI ngay lập tức, tôi cần các bạn làm cho chi phí đó dự đoán được." Đó chính xác là bài toán mà một budget framework giải quyết — không phải để cấm dùng AI, mà để mọi workflow agentic tự biết giới hạn của mình và tự điều chỉnh trước khi vượt ngưỡng.

Chúng ta sẽ xây dựng một package Python tên acme_ai_budget, gồm 5 module (config.py, budget.py, router.py, client.py, report.py) cùng một bộ integration test. Yêu cầu môi trường: Python 3.11+ (dùng cú pháp dict[str, float], tuple[str, ModelTier]), pytest để chạy test. Toàn bộ code dưới đây là code mẫu thực tế, bạn có thể copy nguyên và chạy được — chỉ cần thay phần gọi LLM provider thật ở client.py.

Hãy bắt đầu từ bước quan trọng nhất mà rất nhiều team hay bỏ qua: đo trước khi xây.

Phase 1: Rà Soát Workflow Và Ước Tính Chi Phí Baseline

Sai lầm phổ biến nhất khi làm budget framework là viết code enforcement (thực thi giới hạn) trước, rồi mới đi tìm số liệu để điền vào cap. Làm ngược lại thứ tự này gần như luôn dẫn đến cap sai — quá chặt thì chặn nhầm workflow quan trọng, quá lỏng thì framework tồn tại chỉ để làm cảnh. Phase 1 không viết một dòng code enforcement nào cả. Việc duy nhất cần làm là: liệt kê toàn bộ workflow agentic đang chạy trong tổ chức, và tính ra baseline cost estimate (ước tính chi phí nền) dựa trên dữ liệu thật.

Bước 1: Kiểm Kê (Inventory) Các Workflow

Ngồi lại với từng team lead sở hữu một agentic workflow, hoặc tốt hơn, kéo số liệu trực tiếp từ LLM gateway/API log nếu công ty đã có logging tập trung. Với mỗi workflow, bạn cần thu thập tối thiểu 6 thông tin: tên workflow, team sở hữu, số lần chạy trung bình mỗi ngày, số token input trung bình, số token output trung bình, model đang dùng, và mức độ ưu tiên (priority) — tức là nếu phải giảm chất lượng hoặc chặn tạm, workflow này chấp nhận được đến đâu.

Ở Acme Corp, sau hai ngày phỏng vấn các team, bảng kiểm kê trông như sau:

Vài điểm cần chú ý khi đọc bảng này. code-review-agent chạy 400 lần/ngày vì nó review từng commit trong PR, không chỉ một lần lúc mở PR — đây là chi tiết rất dễ bị bỏ sót nếu chỉ hỏi "workflow này chạy bao nhiêu lần" mà không hỏi kỹ "chạy theo trigger nào". support-ticket-agent được đánh priority "critical" vì nó đối mặt trực tiếp với khách hàng: nếu framework sau này chặn nhầm workflow này, khách hàng sẽ không nhận được phản hồi — hậu quả nghiêm trọng hơn nhiều so với việc release-notes-agent bị chặn một hôm.

Một chi tiết dễ gây sai số: support-ticket-agent chạy cả 7 ngày/tuần (khách hàng không nghỉ cuối tuần), còn các workflow còn lại chỉ chạy vào ngày làm việc (22 ngày/tháng theo lịch trung bình). Ghi rõ giả định này ngay trong bảng kiểm kê, vì nó ảnh hưởng trực tiếp đến bước tính baseline tiếp theo.

Bước 2: Tính Toán Baseline Cost Estimate

Có bảng kiểm kê rồi, viết một script một lần chạy (one-off script) để chuyển số liệu thô thành ước tính chi phí hàng tháng. Đừng viết baseline estimate bằng Excel nếu tránh được — viết bằng chính hàm tính cost mà framework sẽ dùng sau này, để đảm bảo con số baseline và con số enforcement sau này nhất quán với nhau.

"""
baseline_estimate.py — one-off script to produce an initial monthly cost
projection BEFORE any enforcement code exists. Run this first, walk the
numbers past team leads, then use the result to set monthly_cap_usd values
in acme_ai_budget/config.py.

Run: python baseline_estimate.py
"""

from acme_ai_budget.budget import calculate_cost_usd

WORKFLOW_INVENTORY = [
    ("code-review-agent",    400, 25_000, 1_500, "gpt-4o",           22),
    ("pr-summary-agent",     400,  4_000,   500, "gpt-4o-mini",      22),
    ("support-ticket-agent", 2500, 3_000,   400, "claude-haiku-4-5-20251001", 30),
    ("doc-qa-agent",         1200, 20_000,  600, "gpt-4o-mini",      22),
    ("test-gen-agent",       200, 30_000, 4_000, "gpt-4o",           22),
    ("release-notes-agent",   20, 10_000, 1_500, "gpt-3.5-turbo",    22),
]


def project_monthly_cost():
    print(f"{'Workflow':<24}{'Runs/mo':>10}{'Est. cost/mo':>16}")
    print("-" * 52)
    grand_total = 0.0
    for name, runs_per_day, avg_in, avg_out, model, days_per_month in WORKFLOW_INVENTORY:
        runs_per_month = runs_per_day * days_per_month
        cost_per_call = calculate_cost_usd(model, avg_in, avg_out)
        monthly_cost = cost_per_call * runs_per_month
        grand_total += monthly_cost
        print(f"{name:<24}{runs_per_month:>10,}{monthly_cost:>16,.2f}")
    print("-" * 52)
    print(f"{'TOTAL':<24}{'':>10}{grand_total:>16,.2f}")
    return grand_total


if __name__ == "__main__":
    project_monthly_cost()

Chạy script này (giả sử calculate_cost_usd đã được viết ở Phase 2, bạn có thể tạm inline bảng giá để chạy trước), Acme Corp nhận được kết quả tổng baseline khoảng $1,588/tháng. Con số này thấp hơn nhiều so với dự đoán ban đầu của mọi người trong phòng — và đó chính là giá trị của việc đo đạc thật: hóa đơn tăng vọt tháng trước không phải vì baseline quá cao, mà vì có một workflow tăng volume đột biến (một chiến dịch marketing khiến support-ticket-agent chạy gấp 4 lần bình thường trong 3 ngày) mà không ai theo dõi theo thời gian thực. Baseline cho bạn biết "bình thường trông như thế nào" — đó là tiền đề để phát hiện bất thường.

Đừng chốt baseline này làm cap ngay. Baseline là con số trung bình của hiện tại, không phải giới hạn cho tương lai. Bước tiếp theo trong Phase 2 sẽ cộng thêm buffer (khoảng đệm) hợp lý cho từng workflow.

Mẹo

• Lấy số liệu thật từ log của LLM gateway hoặc APM (Application Performance Monitoring) nếu có, đừng chỉ dựa vào ước lượng miệng của team lead — con số thực tế và con số "cảm giác" có thể lệch nhau vài lần.
• Ghi lại rõ giả định về đơn vị thời gian (ngày làm việc vs ngày lịch) ngay trong bảng kiểm kê — đây là nguồn sai số phổ biến nhất khi nhân ra chi phí tháng.
• Gán priority ngay từ Phase 1, đừng để tới lúc viết enforcement mới tranh luận workflow nào "quan trọng hơn" — lúc đó thường đã muộn và mang tính cảm tính.
• Lưu lại ngày thu thập dữ liệu. Baseline sẽ lệch dần theo thời gian khi team thêm workflow mới hoặc đổi model, nên bạn cần biết baseline này "hết hạn" từ khi nào để làm lại.

Phase 2: Triển Khai Budget Framework

Có baseline rồi, giờ là lúc viết code enforcement thật. Nguyên tắc thiết kế quan trọng nhất ở phase này: tách biệt hoàn toàn cấu hình (số liệu, ngưỡng, cap) khỏi logic (cách tính, cách quyết định). Toàn bộ số liệu sẽ nằm trong config.py; toàn bộ logic tính toán và enforcement nằm trong budget.py. Nhờ vậy, khi finance yêu cầu tăng cap cho một workflow, bạn chỉ sửa một số trong config.py và gửi pull request — không cần đụng vào logic đã test kỹ.

Cấu Trúc Project

acme_ai_budget/
├── __init__.py
├── config.py            # Budget caps + model routing config (Phase 2)
├── budget.py             # Core cost calculation + enforcement (Phase 2)
├── router.py             # Model routing decision logic (Phase 3)
├── client.py             # Main entry point wiring everything (Phase 4)
├── report.py             # Monthly report generator (Phase 5)
└── data/
    └── usage_log.jsonl   # Append-only ledger, one line per LLM call

baseline_estimate.py       # One-off script from Phase 1 (not part of package)

tests/
└── test_budget_framework.py   # Integration tests (Phase 6)

Lưu ý usage_log.jsonl là file JSONL (mỗi dòng một JSON object) — chọn định dạng append-only vì nó ghi rất nhanh, không cần lock file phức tạp, và dễ replay lại lịch sử chi tiêu nếu cần audit. Với quy mô của Acme Corp (vài chục nghìn call/tháng), một file JSONL đơn giản là đủ; nếu tổ chức bạn có hàng triệu call/ngày, hãy thay bằng một time-series database hoặc bảng trong data warehouse — kiến trúc logic phía trên (BudgetTracker) không cần đổi, chỉ đổi lớp lưu trữ.

config.py — Cấu Hình Budget

config.py là nguồn sự thật duy nhất (single source of truth) cho hai thứ: cap ngân sách của từng workflow, và cấu hình routing model theo tier (bậc model, ví dụ premium/standard/economy). File này thuần khai báo, không chứa logic if/else phức tạp.

"""
config.py — Master budget configuration for acme_ai_budget.

This module is the single source of truth for:
- Monthly budget caps (org-wide and per-workflow)
- Model routing tiers and thresholds
- Alerting thresholds

Everything else in the framework (budget.py, router.py, client.py, report.py)
reads from this file. Keep it declarative — no business logic here.
"""

from dataclasses import dataclass
from enum import Enum


class ModelTier(str, Enum):
    PREMIUM = "premium"
    STANDARD = "standard"
    ECONOMY = "economy"


@dataclass(frozen=True)
class WorkflowBudget:
    name: str
    team: str
    monthly_cap_usd: float
    priority: str  # "critical" | "high" | "normal" | "low"
    default_tier: ModelTier
    allow_downgrade: bool = True
    hard_block_at_exceeded: bool = True


ORG_MONTHLY_BUDGET_USD = 1_800.00

WORKFLOW_BUDGETS = {
    "code-review-agent": WorkflowBudget(
        name="code-review-agent",
        team="platform-eng",
        monthly_cap_usd=850.00,       # baseline $682 + ~25% buffer
        priority="high",
        default_tier=ModelTier.PREMIUM,
    ),
    "pr-summary-agent": WorkflowBudget(
        name="pr-summary-agent",
        team="platform-eng",
        monthly_cap_usd=25.00,        # baseline $7.92, generous margin — cheap workflow
        priority="low",
        default_tier=ModelTier.STANDARD,
    ),
    "support-ticket-agent": WorkflowBudget(
        name="support-ticket-agent",
        team="customer-support",
        monthly_cap_usd=420.00,       # baseline $300 + 40% buffer — customer-facing
        priority="critical",
        default_tier=ModelTier.STANDARD,
        hard_block_at_exceeded=False,  # never fully block customer support
    ),
    "doc-qa-agent": WorkflowBudget(
        name="doc-qa-agent",
        team="platform-eng",
        monthly_cap_usd=120.00,       # baseline $88.70 + ~35% buffer
        priority="normal",
        default_tier=ModelTier.STANDARD,
    ),
    "test-gen-agent": WorkflowBudget(
        name="test-gen-agent",
        team="qa",
        monthly_cap_usd=650.00,       # baseline $506 + ~28% buffer
        priority="normal",
        default_tier=ModelTier.PREMIUM,
    ),
    "release-notes-agent": WorkflowBudget(
        name="release-notes-agent",
        team="platform-eng",
        monthly_cap_usd=15.00,        # baseline $3.19, wide margin — cheap, low-priority
        priority="low",
        default_tier=ModelTier.ECONOMY,
    ),
}

MODEL_BY_TIER = {
    ModelTier.PREMIUM: "gpt-4o",
    ModelTier.STANDARD: "gpt-4o-mini",
    ModelTier.ECONOMY: "gpt-3.5-turbo",
}

TIER_DOWNGRADE_PATH = {
    ModelTier.PREMIUM: ModelTier.STANDARD,
    ModelTier.STANDARD: ModelTier.ECONOMY,
    ModelTier.ECONOMY: None,  # no cheaper tier left
}

WARNING_THRESHOLD = 0.80    # send alert, keep current tier
DOWNGRADE_THRESHOLD = 0.90  # force downgrade one tier
CRITICAL_THRESHOLD = 0.97   # force downgrade to economy regardless of default tier
EXCEEDED_THRESHOLD = 1.00   # block, unless hard_block_at_exceeded=False

Chú ý cách đặt ORG_MONTHLY_BUDGET_USD = 1_800.00 thấp hơn tổng các cap con ($850+$25+$420+$120+$650+$15 = $2,080). Đây không phải lỗi, mà là chủ đích: giả định rằng không phải workflow nào cũng chạm cap tối đa cùng một lúc trong cùng một tháng. Nếu bạn cộng thẳng tổng các cap con làm org cap, bạn đang giả định trường hợp xấu nhất tuyệt đối xảy ra đồng thời ở mọi workflow — thường không thực tế, và sẽ khiến org cap quá lỏng để cảnh báo sớm.

budget.py — Core Logic Enforcement Budget

budget.py chịu trách nhiệm cho 4 việc: tính chi phí USD từ số token, ghi lại usage vào log, tính trạng thái ngân sách hiện tại (OK/warning/downgrade/critical/exceeded), và raise exception khi cần chặn cứng (hard block).

"""
budget.py — Core budget enforcement for acme_ai_budget.

Responsibilities:
1. Track USD cost of every LLM call, per workflow.
2. Persist usage so cost survives process restarts (append-only JSONL log).
3. Answer "what status is this workflow in right now?" for router.py.
4. Raise BudgetExceededError when a hard cap is hit.
"""

import json
import datetime as dt
from pathlib import Path
from dataclasses import dataclass

from acme_ai_budget.config import (
    WORKFLOW_BUDGETS,
    WARNING_THRESHOLD,
    DOWNGRADE_THRESHOLD,
    CRITICAL_THRESHOLD,
    EXCEEDED_THRESHOLD,
)

COST_PER_MILLION_TOKENS = {
    "gpt-4o":              {"input": 2.50, "output": 10.00},
    "gpt-4o-mini":         {"input": 0.15, "output": 0.60},
    "gpt-3.5-turbo":       {"input": 0.50, "output": 1.50},
    "claude-sonnet-5":   {"input": 3.00, "output": 15.00},
    "claude-haiku-4-5-20251001":    {"input": 0.80, "output": 4.00},
    "llama-3.1-8b-hosted": {"input": 0.05, "output": 0.05},
}

USAGE_LOG_PATH = Path("acme_ai_budget/data/usage_log.jsonl")


class BudgetExceededError(RuntimeError):
    """Raised when a workflow's hard cap is exceeded and hard-block is on."""


class BudgetStatus:
    OK = "ok"
    WARNING = "warning"
    DOWNGRADE = "downgrade"
    CRITICAL = "critical"
    EXCEEDED = "exceeded"


def calculate_cost_usd(model: str, input_tokens: int, output_tokens: int) -> float:
    """Compute the USD cost of one call, given raw token counts."""
    if model not in COST_PER_MILLION_TOKENS:
        raise ValueError(f"Unknown model '{model}' — add it to COST_PER_MILLION_TOKENS first.")
    price = COST_PER_MILLION_TOKENS[model]
    cost = (input_tokens / 1_000_000) * price["input"]
    cost += (output_tokens / 1_000_000) * price["output"]
    return round(cost, 6)


@dataclass
class UsageRecord:
    timestamp: str
    workflow: str
    model: str
    input_tokens: int
    output_tokens: int
    cost_usd: float


class BudgetTracker:
    """
    In-memory + on-disk ledger of spend per workflow for the current
    calendar month. One instance is meant to be shared (or re-hydrated)
    across the whole process — client.py holds a single BudgetTracker.
    """

    def __init__(self, log_path: Path = USAGE_LOG_PATH):
        self.log_path = log_path
        self.log_path.parent.mkdir(parents=True, exist_ok=True)
        self._spend_by_workflow: dict[str, float] = {}
        self._load_current_month()

    def _load_current_month(self) -> None:
        if not self.log_path.exists():
            return
        month_prefix = dt.date.today().strftime("%Y-%m")
        with self.log_path.open("r", encoding="utf-8") as f:
            for line in f:
                record = json.loads(line)
                if record["timestamp"].startswith(month_prefix):
                    self._spend_by_workflow[record["workflow"]] = (
                        self._spend_by_workflow.get(record["workflow"], 0.0)
                        + record["cost_usd"]
                    )

    def record_usage(self, workflow: str, model: str, input_tokens: int, output_tokens: int) -> UsageRecord:
        cost = calculate_cost_usd(model, input_tokens, output_tokens)
        record = UsageRecord(
            timestamp=dt.datetime.utcnow().isoformat(),
            workflow=workflow,
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost_usd=cost,
        )
        self._spend_by_workflow[workflow] = self._spend_by_workflow.get(workflow, 0.0) + cost
        with self.log_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(record.__dict__) + "\n")
        return record

    def get_spend(self, workflow: str) -> float:
        return round(self._spend_by_workflow.get(workflow, 0.0), 4)

    def get_status(self, workflow: str) -> str:
        budget = WORKFLOW_BUDGETS.get(workflow)
        if budget is None:
            raise ValueError(f"Workflow '{workflow}' is not registered in config.WORKFLOW_BUDGETS")

        spent_ratio = self.get_spend(workflow) / budget.monthly_cap_usd

        if spent_ratio >= EXCEEDED_THRESHOLD:
            return BudgetStatus.EXCEEDED
        if spent_ratio >= CRITICAL_THRESHOLD:
            return BudgetStatus.CRITICAL
        if spent_ratio >= DOWNGRADE_THRESHOLD:
            return BudgetStatus.DOWNGRADE
        if spent_ratio >= WARNING_THRESHOLD:
            return BudgetStatus.WARNING
        return BudgetStatus.OK

    def enforce(self, workflow: str) -> None:
        """Call this before making an LLM request. Raises if hard-blocked."""
        budget = WORKFLOW_BUDGETS[workflow]
        status = self.get_status(workflow)
        if status == BudgetStatus.EXCEEDED and budget.hard_block_at_exceeded:
            raise BudgetExceededError(
                f"Workflow '{workflow}' has spent ${self.get_spend(workflow):.2f} "
                f"of its ${budget.monthly_cap_usd:.2f} monthly cap. Hard-blocked."
            )

Bảng cost-per-token ở trên (COST_PER_MILLION_TOKENS) là phần dễ bị "chết" (outdated) nhanh nhất trong toàn bộ framework — nhà cung cấp model đổi giá thường xuyên hơn bạn nghĩ. Coi bảng này như dữ liệu cần review hàng tháng, không phải constant viết một lần rồi quên.

Mẹo

• Đừng để logic tính toán rải rác nhiều nơi — mọi ngưỡng, mọi cap chỉ nằm trong config.py. Khi cần đổi số, PR review sẽ chỉ đụng một file.
• Luôn set buffer 20-40% trên baseline khi đặt cap. Cap = baseline gần như chắc chắn sẽ bị vượt trong tháng đầu vì biến động tự nhiên (traffic, seasonality).
• Version-control mọi thay đổi cap trong config.py, và yêu cầu ghi rõ lý do trong PR description (ví dụ "tăng cap doc-qa-agent vì thêm 2 sản phẩm mới vào RAG index") — đây chính là audit trail cho finance sau này.
• hard_block_at_exceeded=False không có nghĩa là "bỏ qua budget" — nó có nghĩa là "đừng chặn hoàn toàn, nhưng router.py vẫn phải downgrade xuống model rẻ nhất". Đọc kỹ Phase 3 để hiểu cách hai cờ này phối hợp.

Phase 3: Triển Khai Model Router

Có budget.py biết "workflow này đang ở trạng thái nào", giờ cần một thành phần quyết định "vậy thì gọi model nào". Đó là việc của router.py — một model router (bộ định tuyến model) độc lập, không tự gọi LLM, không tự ghi log; nó chỉ nhận input là tên workflow và trả về quyết định routing.

Tách router thành module riêng (thay vì nhúng logic if/else vào client.py) mang lại lợi ích rất thực tế: bạn có thể unit test toàn bộ ma trận quyết định routing (OK → giữ tier mặc định, warning → giữ nguyên nhưng cảnh báo, downgrade → hạ 1 bậc, critical → hạ thẳng xuống economy) mà không cần gọi LLM thật hay tốn một đồng nào.

"""
router.py — Model router for acme_ai_budget.

Decides which concrete model a workflow should call *right now*, based on:
- The workflow's default tier (from config.py)
- How much of its monthly budget is already spent (via BudgetTracker)
- Whether the workflow allows downgrading at all (some critical workflows
  should never be silently downgraded to a weaker model)
"""

from acme_ai_budget.config import (
    WORKFLOW_BUDGETS,
    MODEL_BY_TIER,
    TIER_DOWNGRADE_PATH,
    ModelTier,
)
from acme_ai_budget.budget import BudgetTracker, BudgetStatus


class ModelRouter:
    def __init__(self, tracker: BudgetTracker):
        self.tracker = tracker

    def choose_model(self, workflow: str) -> tuple[str, ModelTier]:
        budget = WORKFLOW_BUDGETS[workflow]
        status = self.tracker.get_status(workflow)
        tier = budget.default_tier

        if not budget.allow_downgrade:
            return MODEL_BY_TIER[tier], tier

        if status in (BudgetStatus.CRITICAL, BudgetStatus.EXCEEDED):
            tier = ModelTier.ECONOMY
        elif status == BudgetStatus.DOWNGRADE:
            downgraded = TIER_DOWNGRADE_PATH.get(budget.default_tier)
            tier = downgraded or budget.default_tier
        # status == WARNING or OK -> keep default_tier unchanged

        model = MODEL_BY_TIER[tier]
        return model, tier

    def explain(self, workflow: str) -> str:
        """Human-readable reason for the current routing decision — useful
        for logging and for the monthly report."""
        model, tier = self.choose_model(workflow)
        status = self.tracker.get_status(workflow)
        spend = self.tracker.get_spend(workflow)
        cap = WORKFLOW_BUDGETS[workflow].monthly_cap_usd
        return (
            f"[{workflow}] status={status} spend=${spend:.2f}/${cap:.2f} "
            f"-> routed to tier={tier.value} model={model}"
        )

Đây là lúc bạn thấy giá trị của cờ hard_block_at_exceeded=False trên support-ticket-agent từ Phase 2. Khi workflow này chạm status EXCEEDED, budget.enforce() sẽ không raise BudgetExceededError (vì hard block đã tắt), nhưng router.choose_model() vẫn hạ nó xuống ModelTier.ECONOMY. Kết quả: khách hàng vẫn nhận được phản hồi, nhưng bằng model rẻ hơn nhiều — một dạng graceful degradation (suy giảm có kiểm soát) hợp lý hơn nhiều so với việc chặn cứng một kênh hỗ trợ khách hàng.

Mẹo

• Luôn log lại router.explain() mỗi khi có quyết định routing — khi có ai hỏi "tại sao hôm qua code review dùng gpt-4o-mini thay vì gpt-4o", bạn cần trả lời bằng log, không phải bằng suy đoán.
• Với workflow priority "critical", hãy cân nhắc kỹ trước khi cho allow_downgrade=True — với một số workflow, hạ tier âm thầm còn tệ hơn là dừng lại và báo lỗi rõ ràng cho người vận hành.
• Test ModelRouter hoàn toàn tách biệt với LLM call thật (dùng BudgetTracker giả lập số liệu) — Phase 6 sẽ làm chính xác điều này.
• Cẩn thận hiện tượng "flap" (dao động qua lại liên tục quanh một ngưỡng) quanh ngưỡng: nếu spend dao động lên xuống ngay tại 90%, workflow có thể nhảy qua nhảy lại giữa 2 tier liên tục trong ngày. Nếu gặp trường hợp này, thêm hysteresis (ví dụ hạ tier ở 90% nhưng chỉ nâng lại ở 85%) là hướng xử lý hợp lý cho bản v2.

Phase 4: Kết Nối Main Client

client.py là điểm vào duy nhất mà code ứng dụng còn lại nên import từ acme_ai_budget. Nó ghép 3 thành phần đã viết ở trên theo đúng thứ tự bắt buộc: enforce budget trước khi gọi LLM, chọn model qua router, gọi LLM thật, rồi ghi lại usage thật (không phải usage ước tính) sau khi có response.

"""
client.py — Main entry point that application code calls.

Wires together: ModelRouter (which model to use) + a real LLM API call +
BudgetTracker (record actual usage after the call). This is the only
module the rest of the codebase should import from acme_ai_budget.
"""

import logging

from acme_ai_budget.budget import BudgetTracker, BudgetExceededError
from acme_ai_budget.router import ModelRouter

logger = logging.getLogger("acme_ai_budget.client")


class AcmeAgentClient:
    def __init__(self):
        self.tracker = BudgetTracker()
        self.router = ModelRouter(self.tracker)

    def run(self, workflow: str, prompt: str, **llm_kwargs) -> str:
        """
        Run one agentic call under budget enforcement.

        Raises BudgetExceededError if the workflow is hard-blocked.
        For workflows configured with hard_block_at_exceeded=False (e.g.
        support-ticket-agent), enforce() never raises — the router already
        downgraded the model tier instead.
        """
        self.tracker.enforce(workflow)  # raises BudgetExceededError if hard-blocked

        model, tier = self.router.choose_model(workflow)
        logger.info(self.router.explain(workflow))

        response_text, usage = self._call_llm(model=model, prompt=prompt, **llm_kwargs)

        record = self.tracker.record_usage(
            workflow=workflow,
            model=model,
            input_tokens=usage["input_tokens"],
            output_tokens=usage["output_tokens"],
        )
        logger.info(
            "%s call cost $%.4f (model=%s, in=%d, out=%d)",
            workflow, record.cost_usd, model, usage["input_tokens"], usage["output_tokens"],
        )
        return response_text

    def _call_llm(self, model: str, prompt: str, **kwargs) -> tuple[str, dict]:
        """
        Thin wrapper around your actual LLM SDK call.
        Swap this body for the real provider client (OpenAI, Anthropic, ...).
        Kept as a separate method so it's trivial to monkeypatch in tests.
        """
        # Example with a generic provider SDK:
        # response = provider_client.chat.completions.create(
        #     model=model, messages=[{"role": "user", "content": prompt}], **kwargs
        # )
        # return response.choices[0].message.content, {
        #     "input_tokens": response.usage.prompt_tokens,
        #     "output_tokens": response.usage.completion_tokens,
        # }
        raise NotImplementedError("Wire this up to your actual LLM provider SDK.")

Ở tầng gọi AcmeAgentClient.run(), code ứng dụng cần bọc lời gọi trong try/except BudgetExceededError để quyết định UX phù hợp khi bị chặn — ví dụ báo cho user "hệ thống đang bận, thử lại sau", đưa request vào hàng đợi cho tháng sau, hoặc gửi cảnh báo cho on-call engineer. Đừng để BudgetExceededError rơi tự do lên tận UI — nó là một exception có kế hoạch xử lý, không phải bug.

Một điểm rất quan trọng: record_usage() luôn dùng usage["input_tokens"] và usage["output_tokens"] lấy từ response thật của provider, không phải số ước tính trước khi gọi. Nhiều đội mắc lỗi ước tính token bằng cách đếm ký tự / 4 rồi ghi số đó vào ledger — sai số này tích lũy theo tháng có thể lệch 10-20% so với hóa đơn thật.

Mẹo

• enforce() luôn phải chạy trước lời gọi LLM, không phải sau — nếu enforce sau, bạn đã tốn tiền cho call vượt cap rồi mới biết.
• Tách _call_llm() thành method riêng chỉ để dễ mock — đây là kỹ thuật giúp Phase 6 viết test không cần gọi API thật, không tốn tiền, không cần network.
• Luôn dùng usage token thật từ response provider để ghi log, không dùng số ước tính trước khi gọi.
• Nếu ứng dụng của bạn có nhiều luồng/xử lý song song, cân nhắc thêm lock hoặc dùng queue khi ghi vào usage_log.jsonl để tránh interleaving ghi file — ở quy mô Acme Corp hiện tại điều này chưa cần thiết, nhưng đáng ghi chú lại cho tương lai.

Phase 5: Xây Dựng Monthly Budget Report

Enforcement chạy đúng chỉ là một nửa câu chuyện. Nửa còn lại là visibility (khả năng quan sát) — leadership, finance, và các team owner cần nhìn thấy tình hình chi tiêu mà không phải tự query file JSONL. report.py đọc lại usage_log.jsonl, tổng hợp theo workflow và theo model, rồi tính luôn một con số quan trọng: projected end-of-month spend (dự báo chi tiêu cuối tháng) — ngoại suy tuyến tính đơn giản từ tốc độ chi tiêu hiện tại.

"""
report.py — Monthly budget report generator.

Reads the usage log for the current month and produces:
- A per-workflow spend vs. cap table
- A per-model breakdown
- A projected end-of-month total (linear extrapolation)
- A list of workflows at WARNING/CRITICAL/EXCEEDED status
"""

import json
import datetime as dt
from pathlib import Path
from collections import defaultdict

from acme_ai_budget.budget import USAGE_LOG_PATH, BudgetTracker
from acme_ai_budget.config import WORKFLOW_BUDGETS, ORG_MONTHLY_BUDGET_USD


def _days_elapsed_in_month() -> int:
    return dt.date.today().day


def _days_in_month() -> int:
    today = dt.date.today()
    if today.month == 12:
        next_month = today.replace(year=today.year + 1, month=1, day=1)
    else:
        next_month = today.replace(month=today.month + 1, day=1)
    return (next_month - today.replace(day=1)).days


def build_monthly_report(log_path: Path = USAGE_LOG_PATH) -> dict:
    month_prefix = dt.date.today().strftime("%Y-%m")
    spend_by_workflow = defaultdict(float)
    spend_by_model = defaultdict(float)
    calls_by_workflow = defaultdict(int)

    if log_path.exists():
        with log_path.open("r", encoding="utf-8") as f:
            for line in f:
                rec = json.loads(line)
                if not rec["timestamp"].startswith(month_prefix):
                    continue
                spend_by_workflow[rec["workflow"]] += rec["cost_usd"]
                spend_by_model[rec["model"]] += rec["cost_usd"]
                calls_by_workflow[rec["workflow"]] += 1

    elapsed = _days_elapsed_in_month()
    total_days = _days_in_month()
    total_spend = sum(spend_by_workflow.values())
    projected_total = (total_spend / elapsed) * total_days if elapsed else 0.0

    tracker = BudgetTracker(log_path)
    rows = []
    for workflow, budget in WORKFLOW_BUDGETS.items():
        spent = spend_by_workflow.get(workflow, 0.0)
        rows.append({
            "workflow": workflow,
            "team": budget.team,
            "spent_usd": round(spent, 2),
            "cap_usd": budget.monthly_cap_usd,
            "pct_used": round(spent / budget.monthly_cap_usd * 100, 1) if budget.monthly_cap_usd else 0,
            "status": tracker.get_status(workflow),
            "calls": calls_by_workflow.get(workflow, 0),
        })

    return {
        "month": month_prefix,
        "rows": rows,
        "spend_by_model": dict(spend_by_model),
        "total_spend_usd": round(total_spend, 2),
        "org_cap_usd": ORG_MONTHLY_BUDGET_USD,
        "projected_total_usd": round(projected_total, 2),
        "at_risk": [r["workflow"] for r in rows if r["status"] in ("critical", "exceeded")],
    }


def render_report_markdown(report: dict) -> str:
    lines = [f"# Budget Report — {report['month']}", ""]
    lines.append(f"**Total spend so far:** ${report['total_spend_usd']:.2f} / ${report['org_cap_usd']:.2f} org cap")
    lines.append(f"**Projected end-of-month:** ${report['projected_total_usd']:.2f}")
    lines.append("")
    lines.append("| Workflow | Team | Spent | Cap | % Used | Status | Calls |")
    lines.append("|---|---|---|---|---|---|---|")
    for r in report["rows"]:
        lines.append(
            f"| {r['workflow']} | {r['team']} | ${r['spent_usd']:.2f} | "
            f"${r['cap_usd']:.2f} | {r['pct_used']}% | {r['status']} | {r['calls']} |"
        )
    if report["at_risk"]:
        lines.append("")
        lines.append(f"**At risk this month:** {', '.join(report['at_risk'])}")
    return "\n".join(lines)


if __name__ == "__main__":
    report = build_monthly_report()
    print(render_report_markdown(report))

projected_total_usd là con số quan trọng nhất trong toàn bộ report, quan trọng hơn cả total_spend_usd. Lý do: nếu hôm nay mới ngày 10 trong tháng mà support-ticket-agent đã chi $180/$420 cap (43%), con số tuyệt đối trông vẫn ổn — nhưng nếu ngoại suy tuyến tính, cuối tháng sẽ là khoảng $540, vượt cap gần 30%. Chỉ nhìn số tuyệt đối, bạn phát hiện vấn đề vào ngày 25; nhìn số projected, bạn phát hiện ngay từ ngày 10 — còn đủ thời gian để điều chỉnh cap hoặc điều tra nguyên nhân tăng đột biến trước khi nó thành hóa đơn đã chốt.

Ở Acme Corp, report.py được lên cron job chạy mỗi 8 giờ sáng, đẩy bản Markdown này vào channel Slack #ai-budget-daily. Toàn bộ team lead chỉ cần nhìn channel đó mỗi sáng, không cần hỏi ai.

Mẹo

• Luôn tính và hiển thị projected_total_usd, không chỉ tổng đã chi — đây là chỉ số cảnh báo sớm thực sự hữu dụng, số đã chi chỉ là lịch sử.
• Chạy report.py qua cron hàng ngày và đẩy digest vào Slack/email — đừng bắt stakeholder phải nhớ tự chạy script.
• Lưu report dạng file (Markdown hoặc CSV) theo ngày để có lịch sử tra cứu, phục vụ đối chiếu với hóa đơn thật cuối tháng — sai lệch giữa usage_log.jsonl và invoice thật là dấu hiệu cần điều tra ngay (có thể do call lỗi không log được, hoặc giá model đã đổi mà COST_PER_MILLION_TOKENS chưa cập nhật).
• Khi số lượng workflow vượt quá vài chục, chuyển từ bảng Markdown sang dashboard trực quan (Grafana, Metabase, hoặc nội bộ) — bảng text sẽ khó đọc ở quy mô lớn.

Phase 6: Viết Integration Test Cho Framework

Một budget framework mà không có test tự động là một framework bạn không thể tự tin sửa. Vì client.py phụ thuộc vào một LLM provider thật (không nên gọi trong test), toàn bộ test ở đây tập trung vào budget.py và router.py — hai module chứa logic quan trọng nhất và hoàn toàn deterministic (có thể kiểm chứng được với input cố định).

"""
tests/test_budget_framework.py — integration tests for acme_ai_budget.

Run with: pytest tests/test_budget_framework.py -v
"""

import json
import pytest

from acme_ai_budget import config
from acme_ai_budget.budget import (
    BudgetTracker,
    BudgetExceededError,
    calculate_cost_usd,
)
from acme_ai_budget.router import ModelRouter


@pytest.fixture
def usage_log(tmp_path):
    return tmp_path / "usage_log.jsonl"


@pytest.fixture
def tracker(usage_log):
    return BudgetTracker(log_path=usage_log)


def test_calculate_cost_usd_basic():
    cost = calculate_cost_usd("gpt-4o-mini", input_tokens=1_000_000, output_tokens=500_000)
    assert cost == pytest.approx(0.15 + 0.30, rel=1e-3)


def test_calculate_cost_usd_unknown_model_raises():
    with pytest.raises(ValueError):
        calculate_cost_usd("not-a-real-model", 100, 100)


def test_record_usage_persists_to_log(tracker, usage_log):
    tracker.record_usage("pr-summary-agent", "gpt-4o-mini", 10_000, 2_000)
    assert usage_log.exists()
    lines = usage_log.read_text().strip().splitlines()
    assert len(lines) == 1
    record = json.loads(lines[0])
    assert record["workflow"] == "pr-summary-agent"


def test_status_ok_below_warning_threshold(tracker):
    tracker.record_usage("release-notes-agent", "gpt-3.5-turbo", 50_000, 10_000)
    assert tracker.get_status("release-notes-agent") == "ok"


def test_status_downgrades_at_90_percent(tracker):
    cap = config.WORKFLOW_BUDGETS["release-notes-agent"].monthly_cap_usd
    # Force spend to 92% of cap using a cheap, output-heavy synthetic call.
    tokens_needed = int((cap * 0.92) / (1.50 / 1_000_000))
    tracker.record_usage("release-notes-agent", "gpt-3.5-turbo", 0, tokens_needed)
    assert tracker.get_status("release-notes-agent") == "downgrade"


def test_enforce_raises_when_hard_capped(tracker):
    cap = config.WORKFLOW_BUDGETS["test-gen-agent"].monthly_cap_usd
    tokens_needed = int((cap * 1.10) / (10.00 / 1_000_000))
    tracker.record_usage("test-gen-agent", "gpt-4o", 0, tokens_needed)
    with pytest.raises(BudgetExceededError):
        tracker.enforce("test-gen-agent")


def test_soft_workflow_never_hard_blocks(tracker):
    """support-ticket-agent has hard_block_at_exceeded=False — must not raise."""
    cap = config.WORKFLOW_BUDGETS["support-ticket-agent"].monthly_cap_usd
    tokens_needed = int((cap * 1.20) / (4.00 / 1_000_000))
    tracker.record_usage("support-ticket-agent", "claude-haiku-4-5-20251001", 0, tokens_needed)
    tracker.enforce("support-ticket-agent")  # must not raise


def test_router_downgrades_model_when_over_threshold(tracker):
    router = ModelRouter(tracker)
    cap = config.WORKFLOW_BUDGETS["code-review-agent"].monthly_cap_usd
    tokens_needed = int((cap * 0.95) / (10.00 / 1_000_000))
    tracker.record_usage("code-review-agent", "gpt-4o", 0, tokens_needed)
    model, tier = router.choose_model("code-review-agent")
    assert tier.value in ("standard", "economy")

Có hai kỹ thuật đáng chú ý trong bộ test này. Thứ nhất, fixture usage_log dùng tmp_path (fixture built-in của pytest) để mỗi test chạy trên một file log tạm, riêng biệt, tự động dọn dẹp sau khi test xong — không bao giờ đụng vào usage_log.jsonl thật của production. Thứ hai, các test tính tokens_needed bằng công thức ngược từ cap thực tế trong config.WORKFLOW_BUDGETS, thay vì hard-code một số token cụ thể — nhờ vậy nếu sau này bạn đổi cap trong config.py, test vẫn đúng logic mà không cần sửa số trong test.

Chạy pytest tests/test_budget_framework.py -v, cả 7 test này chạy xong trong chưa tới một giây — không cần network, không tốn một token thật nào, và bạn có thể chạy nó trong CI trên mọi pull request đụng vào acme_ai_budget.

Mẹo

• Luôn dùng tmp_path (hoặc tương đương) để cách ly file log trong test — không bao giờ để test tự động đụng vào file log production.
• Test cả hai nhánh hard-block: workflow có hard_block_at_exceeded=True phải raise, workflow có hard_block_at_exceeded=False (như support-ticket-agent) không bao giờ được raise.
• Viết thêm test ở vùng biên ngưỡng (ví dụ 79.9% vs 80.1%) để chắc chắn logic so sánh >= trong get_status() hoạt động đúng hướng — off-by-one ở biên ngưỡng là lỗi rất dễ lọt qua review bằng mắt.
• Đừng bỏ qua test cho trường hợp lỗi (calculate_cost_usd với model không tồn tại) — đây thường là lỗi thật gặp trong production khi có người thêm workflow mới dùng model chưa được khai báo trong bảng giá.

Phase 7: Deploy Và Validate

Framework chạy đúng trên máy bạn và pass hết test không có nghĩa là an toàn để bật enforcement cứng cho toàn bộ 6 workflow ngay lập tức. budget.enforce() có khả năng chặn hẳn một workflow đang chạy production — nếu cap bị đặt sai (quá thấp so với baseline thật, hoặc công thức tính sai ở một góc nào đó bạn chưa test tới), bạn có thể vô tình làm dừng code-review-agent giữa lúc cả team đang review PR. Rollout (triển khai dần) cho một hệ thống có khả năng chặn traffic luôn cần một checklist cẩn thận, không "deploy một lần rồi xong".

Checklist rollout cho Acme Corp, chia làm 3 giai đoạn kéo dài khoảng 4-5 tuần:

Giai đoạn 1 — Shadow mode (chạy nền, chỉ quan sát và ghi log, không chặn gì cả; khoảng 2 tuần). Deploy toàn bộ acme_ai_budget, nhưng gọi tracker.record_usage() và tracker.get_status() song song với luồng cũ mà không gọi enforce() thật (comment tạm dòng raise, hoặc thêm flag ENFORCEMENT_ENABLED = False trong config.py để tắt toàn cục). Mục tiêu duy nhất của giai đoạn này: xác nhận usage_log.jsonl khớp với hóa đơn thật từ provider, sai lệch nên dưới 2-3%. Nếu sai lệch lớn hơn, đừng qua giai đoạn 2 — quay lại kiểm tra COST_PER_MILLION_TOKENS hoặc cách bạn đọc usage field từ response.

Giai đoạn 2 — Soft enforcement / cảnh báo (khoảng 1-2 tuần). Bật enforce() thật nhưng tạm đặt hard_block_at_exceeded=False cho tất cả workflow (không chỉ support-ticket-agent), để enforcement chạy đúng luồng, ghi log đầy đủ trạng thái, gửi cảnh báo Slack khi có workflow chạm WARNING/CRITICAL, nhưng chưa raise BudgetExceededError cho ai cả. Đây là lúc đội ngũ quan sát xem cap đặt ở Phase 2 có thực tế không, và điều chỉnh nếu cần — hầu như luôn có ít nhất một cap cần chỉnh trong giai đoạn này.

Giai đoạn 3 — Hard enforcement (rollout dần theo workflow). Chuyển từng workflow một sang cấu hình cap thật trong config.py (theo đúng priority: bắt đầu từ priority "low" như release-notes-agent, pr-summary-agent, rồi mới tới "normal", cuối cùng mới tới "high"/"critical"). Không bật hard block cho tất cả 6 workflow trong cùng một ngày. Giữa mỗi lần chuyển, quan sát ít nhất 2-3 ngày trước khi chuyển workflow tiếp theo.

Về monitoring (giám sát) khi đã live: theo dõi tối thiểu 3 chỉ số — projected_total_usd từ report.py mỗi ngày, số lần BudgetExceededError được raise (nếu con số này tăng bất thường, có thể cap đặt sai hoặc có bug ở nơi khác gây gọi lặp), và độ lệch giữa usage_log.jsonl với invoice thật cuối tháng.

Về rollback plan (kế hoạch khôi phục): luôn giữ một cách tắt enforcement ngay lập tức mà không cần deploy lại code — cách đơn giản nhất là thêm biến ENFORCEMENT_ENABLED (đọc từ environment variable hoặc feature flag service) kiểm tra ngay đầu BudgetTracker.enforce(); nếu False, hàm return ngay không raise gì. Khi có sự cố (ví dụ enforcement chặn nhầm một workflow quan trọng giữa giờ làm), on-call chỉ cần lật flag này, không cần chờ pipeline CI/CD chạy lại.

import os

ENFORCEMENT_ENABLED = os.getenv("ACME_BUDGET_ENFORCEMENT_ENABLED", "true").lower() == "true"

def enforce(self, workflow: str) -> None:
    if not ENFORCEMENT_ENABLED:
        return  # global kill switch — see Phase 7 rollback plan
    budget = WORKFLOW_BUDGETS[workflow]
    status = self.get_status(workflow)
    if status == BudgetStatus.EXCEEDED and budget.hard_block_at_exceeded:
        raise BudgetExceededError(
            f"Workflow '{workflow}' has spent ${self.get_spend(workflow):.2f} "
            f"of its ${budget.monthly_cap_usd:.2f} monthly cap. Hard-blocked."
        )

Mẹo

• Không bao giờ bật hard enforcement cho tất cả workflow cùng lúc, dù test pass 100% — rollout theo từng workflow, ưu tiên bật cho priority thấp trước.
• Luôn có kill switch (cơ chế tắt khẩn cấp, ví dụ ENFORCEMENT_ENABLED hoặc tương đương) triển khai từ trước khi cần dùng, không phải viết vội lúc đang có sự cố.
• Thông báo trước cho các team bị ảnh hưởng (đặc biệt team customer-support) trước khi chuyển bất kỳ workflow nào từ soft sang hard enforcement — đừng để họ phát hiện qua việc agent bỗng "ngừng chạy".
• Đặt lịch review cap định kỳ (gợi ý: đầu mỗi quý), dựa trên số liệu thật từ report.py chứ không dựa vào baseline Phase 1 mãi mãi — workflow mới ra đời, workflow cũ tăng traffic, giá model thay đổi, cap phải sống theo những thay đổi đó.

Tổng Kết

Bạn vừa xây xong một token budget framework hoàn chỉnh, từ số 0 tới có thể deploy: kiểm kê workflow thật và tính baseline (Phase 1), cấu hình cap tách biệt khỏi logic enforcement (Phase 2), model router tự động hạ tier khi cần (Phase 3), client wiring gắn enforcement đúng vị trí trong luồng gọi LLM (Phase 4), report hàng ngày với chỉ số dự báo cuối tháng (Phase 5), bộ test tự động không tốn tiền và không cần network (Phase 6), và một quy trình rollout 3 giai đoạn có kill switch rõ ràng (Phase 7).

Điểm cốt lõi cần nhớ, xuyên suốt cả 7 phase: một budget framework tốt không phải là một framework chặn được nhiều nhất, mà là một framework mà mọi cap, mọi ngưỡng, mọi quyết định đều giải thích được bằng số liệu thật — từ bảng kiểm kê Phase 1 cho tới log routing ở Phase 3. Khi finance hoặc leadership hỏi "tại sao cap của workflow này là $850 mà không phải $500 hay $1,000", bạn luôn có câu trả lời bằng số, không phải bằng cảm giác.

acme_ai_budget như đã viết ở đây đủ cho một team engineering vài chục người với vài workflow agentic. Khi tổ chức lớn hơn, những hướng mở rộng tự nhiên tiếp theo là: budget theo từng user/team con thay vì chỉ theo workflow, phát hiện bất thường (anomaly detection) tự động thay vì chỉ dựa vào ngưỡng % cố định, và tích hợp thẳng report.py vào hệ thống kế toán nội bộ để finance không cần đọc Slack digest thủ công nữa.

Mẹo

• Đừng coi framework này là "xong rồi để đó". Cap, threshold, và cả bảng giá model trong budget.py đều là dữ liệu sống, cần review định kỳ.
• Lưu lại toàn bộ quyết định thay đổi cap/threshold trong lịch sử git của config.py — đây chính là cuốn "nhật ký" giải trình chi phí AI của tổ chức bạn, hữu ích hơn bất kỳ báo cáo viết tay nào.
• Khi mở rộng framework, luôn giữ nguyên nguyên tắc tách config khỏi logic đã áp dụng từ Phase 2 — đây là lý do framework này dễ mở rộng mà không cần viết lại từ đầu.