LZH-YS1998
108 lines
3.4 KiB
Python
108 lines
3.4 KiB
Python
"""Cost tracking for LLM API usage."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from loguru import logger
|
|
|
|
from opc.database.store import OPCStore
|
|
from opc.core.events import EventBus
|
|
from opc.core.models import OPCEvent, CostEvent
|
|
|
|
|
|
@dataclass
|
|
class CostEntry:
|
|
task_id: str | None = None
|
|
agent_id: str | None = None
|
|
org_id: str | None = None
|
|
model: str = ""
|
|
tokens_in: int = 0
|
|
tokens_out: int = 0
|
|
cost: float = 0.0
|
|
timestamp: datetime = field(default_factory=datetime.now)
|
|
|
|
|
|
class CostTracker:
|
|
"""Tracks LLM API costs per task and agent."""
|
|
|
|
def __init__(self, store: OPCStore, event_bus: EventBus | None = None) -> None:
|
|
self.store = store
|
|
self.event_bus = event_bus
|
|
self._session_total = 0.0
|
|
|
|
async def record(self, entry: CostEntry) -> None:
|
|
await self.store.record_cost(
|
|
task_id=entry.task_id,
|
|
agent_id=entry.agent_id,
|
|
model=entry.model,
|
|
tokens_in=entry.tokens_in,
|
|
tokens_out=entry.tokens_out,
|
|
cost=entry.cost,
|
|
)
|
|
# Also record CostEvent for cost_events table (org-scoped tracking)
|
|
event = CostEvent(
|
|
org_id=entry.org_id,
|
|
agent_id=entry.agent_id,
|
|
task_id=entry.task_id,
|
|
model=entry.model,
|
|
tokens_in=entry.tokens_in,
|
|
tokens_out=entry.tokens_out,
|
|
cost_usd=entry.cost,
|
|
timestamp=entry.timestamp,
|
|
)
|
|
await self.store.record_cost_event(event)
|
|
self._session_total += entry.cost
|
|
|
|
if self.event_bus:
|
|
await self.event_bus.publish(OPCEvent(
|
|
event_type="cost_update",
|
|
payload={
|
|
"task_id": entry.task_id,
|
|
"cost": entry.cost,
|
|
"session_total": self._session_total,
|
|
},
|
|
))
|
|
|
|
async def check_budget(
|
|
self,
|
|
agent_id: str | None = None,
|
|
org_id: str | None = None,
|
|
) -> tuple[bool, str]:
|
|
"""Check if agent/org is within budget. Returns (allowed, reason)."""
|
|
return await check_budget(self.store, agent_id=agent_id, org_id=org_id)
|
|
|
|
async def get_summary(self, project_id: str | None = None) -> dict[str, Any]:
|
|
db_totals = await self.store.get_total_cost(project_id)
|
|
return {
|
|
**db_totals,
|
|
"session_cost": self._session_total,
|
|
}
|
|
|
|
@property
|
|
def session_total(self) -> float:
|
|
return self._session_total
|
|
|
|
|
|
async def check_budget(
|
|
store: OPCStore,
|
|
agent_id: str | None = None,
|
|
org_id: str | None = None,
|
|
) -> tuple[bool, str]:
|
|
"""Check if agent/org is within budget. Returns (allowed, reason)."""
|
|
if org_id:
|
|
org = await store.get_organization(org_id)
|
|
if org and org.budget_monthly_cents > 0:
|
|
if org.spent_monthly_cents >= org.budget_monthly_cents:
|
|
return False, f"Organization '{org.name}' has exceeded its monthly budget"
|
|
if agent_id and org_id:
|
|
agents = await store.list_org_agents(org_id)
|
|
for agent in agents:
|
|
if agent.agent_id == agent_id or agent.role_id == agent_id:
|
|
if agent.budget_monthly_cents > 0 and agent.spent_monthly_cents >= agent.budget_monthly_cents:
|
|
return False, f"Agent '{agent.name}' has exceeded its monthly budget"
|
|
break
|
|
return True, ""
|