Initial commit

opc/layer6_observability/cost_tracker.py

@@ -0,0 +1,107 @@
+"""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, ""