Initial release — OpenMontage: the first open-source agentic video production system
11 production pipelines, 47 tools, 124 agent skills. Supports cloud APIs (fal.ai, OpenAI, ElevenLabs, Suno, HeyGen, Runway) and free local providers (diffusers, Piper TTS, WAN 2.1, Hunyuan, CogVideo). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,212 @@
|
||||
"""Cost tracker core: estimate, reserve, reconcile, and persist to cost_log.json.
|
||||
|
||||
Implements the budget governance rules from the spec:
|
||||
- Every paid operation produces a preflight estimate
|
||||
- The orchestrator reserves estimated budget before execution
|
||||
- Budget overruns trigger pauses (in warn/cap mode)
|
||||
- Actual spend is reconciled when the tool finishes or fails
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from lib.config_model import BudgetMode
|
||||
|
||||
|
||||
class EntryStatus(str, Enum):
|
||||
ESTIMATED = "estimated"
|
||||
RESERVED = "reserved"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
REFUNDED = "refunded"
|
||||
|
||||
|
||||
class BudgetExceededError(Exception):
|
||||
"""Raised when an operation would exceed the budget in cap mode."""
|
||||
pass
|
||||
|
||||
|
||||
class ApprovalRequiredError(Exception):
|
||||
"""Raised when an operation needs user approval before proceeding."""
|
||||
pass
|
||||
|
||||
|
||||
class CostTracker:
|
||||
"""Tracks estimated, reserved, and actual costs for a pipeline project."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
budget_total_usd: float = 10.0,
|
||||
reserve_pct: float = 0.10,
|
||||
single_action_approval_usd: float = 0.50,
|
||||
require_approval_for_new_paid_tool: bool = True,
|
||||
mode: BudgetMode = BudgetMode.WARN,
|
||||
cost_log_path: Optional[Path] = None,
|
||||
) -> None:
|
||||
self.budget_total_usd = budget_total_usd
|
||||
self.reserve_pct = reserve_pct
|
||||
self.single_action_approval_usd = single_action_approval_usd
|
||||
self.require_approval_for_new_paid_tool = require_approval_for_new_paid_tool
|
||||
self.mode = mode
|
||||
self.cost_log_path = cost_log_path
|
||||
self.entries: list[dict[str, Any]] = []
|
||||
self._approved_tools: set[str] = set()
|
||||
|
||||
if cost_log_path and cost_log_path.exists():
|
||||
self._load()
|
||||
|
||||
# ---- Budget calculations ----
|
||||
|
||||
@property
|
||||
def budget_reserved_usd(self) -> float:
|
||||
return sum(
|
||||
e.get("reserved_usd", 0.0)
|
||||
for e in self.entries
|
||||
if e["status"] == EntryStatus.RESERVED.value
|
||||
)
|
||||
|
||||
@property
|
||||
def budget_spent_usd(self) -> float:
|
||||
return sum(
|
||||
e.get("actual_usd", 0.0)
|
||||
for e in self.entries
|
||||
if e["status"] in (EntryStatus.COMPLETED.value, EntryStatus.FAILED.value)
|
||||
)
|
||||
|
||||
@property
|
||||
def budget_remaining_usd(self) -> float:
|
||||
return self.budget_total_usd - self.budget_spent_usd - self.budget_reserved_usd
|
||||
|
||||
@property
|
||||
def usable_budget_usd(self) -> float:
|
||||
"""Budget minus the reserve holdback."""
|
||||
holdback = self.budget_total_usd * self.reserve_pct
|
||||
return max(0.0, self.budget_remaining_usd - holdback)
|
||||
|
||||
def cost_snapshot(self) -> dict[str, float]:
|
||||
return {
|
||||
"total_spent_usd": round(self.budget_spent_usd, 4),
|
||||
"total_reserved_usd": round(self.budget_reserved_usd, 4),
|
||||
"budget_remaining_usd": round(self.budget_remaining_usd, 4),
|
||||
}
|
||||
|
||||
# ---- Core operations ----
|
||||
|
||||
def estimate(self, tool: str, operation: str, estimated_usd: float) -> str:
|
||||
"""Record an estimate. Returns entry ID."""
|
||||
entry_id = self._new_id()
|
||||
self.entries.append({
|
||||
"id": entry_id,
|
||||
"tool": tool,
|
||||
"operation": operation,
|
||||
"status": EntryStatus.ESTIMATED.value,
|
||||
"estimated_usd": round(estimated_usd, 4),
|
||||
"reserved_usd": 0.0,
|
||||
"actual_usd": 0.0,
|
||||
"timestamp": self._now(),
|
||||
})
|
||||
self._save()
|
||||
return entry_id
|
||||
|
||||
def reserve(self, entry_id: str) -> None:
|
||||
"""Reserve budget for an estimated entry.
|
||||
|
||||
Raises BudgetExceededError in cap mode, or ApprovalRequiredError
|
||||
when the action exceeds the single-action approval threshold.
|
||||
"""
|
||||
entry = self._find(entry_id)
|
||||
estimated = entry["estimated_usd"]
|
||||
|
||||
# Check single-action approval threshold
|
||||
if estimated > self.single_action_approval_usd:
|
||||
if self.mode != BudgetMode.OBSERVE:
|
||||
raise ApprovalRequiredError(
|
||||
f"Action costs ${estimated:.2f}, exceeds "
|
||||
f"single-action threshold ${self.single_action_approval_usd:.2f}"
|
||||
)
|
||||
|
||||
# Check new paid tool approval
|
||||
if self.require_approval_for_new_paid_tool and estimated > 0:
|
||||
if entry["tool"] not in self._approved_tools:
|
||||
if self.mode != BudgetMode.OBSERVE:
|
||||
raise ApprovalRequiredError(
|
||||
f"First paid use of tool {entry['tool']!r} requires approval"
|
||||
)
|
||||
|
||||
# Check budget
|
||||
if estimated > self.usable_budget_usd:
|
||||
if self.mode == BudgetMode.CAP:
|
||||
raise BudgetExceededError(
|
||||
f"Reservation of ${estimated:.2f} exceeds usable budget "
|
||||
f"${self.usable_budget_usd:.2f}"
|
||||
)
|
||||
|
||||
entry["status"] = EntryStatus.RESERVED.value
|
||||
entry["reserved_usd"] = estimated
|
||||
entry["timestamp"] = self._now()
|
||||
self._save()
|
||||
|
||||
def approve_tool(self, tool: str) -> None:
|
||||
"""Mark a tool as approved for paid operations."""
|
||||
self._approved_tools.add(tool)
|
||||
|
||||
def reconcile(self, entry_id: str, actual_usd: float, success: bool = True) -> None:
|
||||
"""Reconcile actual spend after tool execution."""
|
||||
entry = self._find(entry_id)
|
||||
entry["status"] = EntryStatus.COMPLETED.value if success else EntryStatus.FAILED.value
|
||||
entry["actual_usd"] = round(actual_usd, 4)
|
||||
entry["reserved_usd"] = 0.0
|
||||
entry["timestamp"] = self._now()
|
||||
self._save()
|
||||
|
||||
def refund(self, entry_id: str) -> None:
|
||||
"""Cancel a reservation without executing."""
|
||||
entry = self._find(entry_id)
|
||||
entry["status"] = EntryStatus.REFUNDED.value
|
||||
entry["reserved_usd"] = 0.0
|
||||
entry["timestamp"] = self._now()
|
||||
self._save()
|
||||
|
||||
# ---- Persistence ----
|
||||
|
||||
def _save(self) -> None:
|
||||
if self.cost_log_path is None:
|
||||
return
|
||||
data = {
|
||||
"version": "1.0",
|
||||
"budget_total_usd": self.budget_total_usd,
|
||||
"budget_reserved_usd": round(self.budget_reserved_usd, 4),
|
||||
"budget_spent_usd": round(self.budget_spent_usd, 4),
|
||||
"entries": self.entries,
|
||||
}
|
||||
self.cost_log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(self.cost_log_path, "w") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
def _load(self) -> None:
|
||||
with open(self.cost_log_path) as f: # type: ignore[arg-type]
|
||||
data = json.load(f)
|
||||
self.entries = data.get("entries", [])
|
||||
self.budget_total_usd = data.get("budget_total_usd", self.budget_total_usd)
|
||||
|
||||
# ---- Helpers ----
|
||||
|
||||
def _find(self, entry_id: str) -> dict[str, Any]:
|
||||
for entry in self.entries:
|
||||
if entry["id"] == entry_id:
|
||||
return entry
|
||||
raise KeyError(f"Cost entry {entry_id!r} not found")
|
||||
|
||||
@staticmethod
|
||||
def _new_id() -> str:
|
||||
return uuid.uuid4().hex[:12]
|
||||
|
||||
@staticmethod
|
||||
def _now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
Reference in New Issue
Block a user