128 lines
4.5 KiB
Python
128 lines
4.5 KiB
Python
"""Cost service: budget tracking and enforcement."""
|
|
from sqlalchemy.orm import Session
|
|
from sqlalchemy import func
|
|
from typing import Optional, List
|
|
from app.models.cost import CostLog
|
|
from app.models.task import Task
|
|
from app.models.worker import AIWorker
|
|
from app.models.project import Project
|
|
from app.core.exceptions import BudgetExceededError
|
|
|
|
|
|
class CostService:
|
|
@staticmethod
|
|
def get_task_cost(db: Session, tenant_id: int, task_id: int) -> int:
|
|
"""Get total cost (cents) for a task."""
|
|
result = (
|
|
db.query(func.sum(CostLog.cost_cents))
|
|
.filter(CostLog.tenant_id == tenant_id, CostLog.task_id == task_id)
|
|
.scalar()
|
|
)
|
|
return result or 0
|
|
|
|
@staticmethod
|
|
def get_project_cost(db: Session, tenant_id: int, project_id: int) -> int:
|
|
"""Get total cost (cents) for a project."""
|
|
result = (
|
|
db.query(func.sum(CostLog.cost_cents))
|
|
.filter(CostLog.tenant_id == tenant_id, CostLog.project_id == project_id)
|
|
.scalar()
|
|
)
|
|
return result or 0
|
|
|
|
@staticmethod
|
|
def get_worker_monthly_cost(
|
|
db: Session, tenant_id: int, worker_id: int
|
|
) -> int:
|
|
"""Get current month's cost for a worker."""
|
|
from datetime import datetime
|
|
month_start = datetime.utcnow().replace(day=1, hour=0, minute=0, second=0)
|
|
result = (
|
|
db.query(func.sum(CostLog.cost_cents))
|
|
.filter(
|
|
CostLog.tenant_id == tenant_id,
|
|
CostLog.worker_id == worker_id,
|
|
CostLog.created_at >= month_start,
|
|
)
|
|
.scalar()
|
|
)
|
|
return result or 0
|
|
|
|
@staticmethod
|
|
def check_budget(
|
|
db: Session, tenant_id: int, task: Task, worker: AIWorker
|
|
) -> None:
|
|
"""Check if executing this task would exceed budget limits."""
|
|
# Worker per-task limit
|
|
if worker.max_cost_per_task_cents:
|
|
task_cost = CostService.get_task_cost(db, tenant_id, task.id)
|
|
if task_cost >= worker.max_cost_per_task_cents:
|
|
raise BudgetExceededError(
|
|
f"Worker '{worker.name}' per-task budget exceeded: "
|
|
f"${task_cost/100:.2f} >= ${worker.max_cost_per_task_cents/100:.2f}"
|
|
)
|
|
|
|
# Worker monthly limit
|
|
if worker.max_cost_per_month_cents:
|
|
monthly_cost = CostService.get_worker_monthly_cost(db, tenant_id, worker.id)
|
|
if monthly_cost >= worker.max_cost_per_month_cents:
|
|
raise BudgetExceededError(
|
|
f"Worker '{worker.name}' monthly budget exceeded: "
|
|
f"${monthly_cost/100:.2f} >= ${worker.max_cost_per_month_cents/100:.2f}"
|
|
)
|
|
|
|
# Project budget limit
|
|
project = (
|
|
db.query(Project)
|
|
.filter(Project.id == task.project_id, Project.tenant_id == tenant_id)
|
|
.first()
|
|
)
|
|
if project and project.budget_limit_cents:
|
|
project_cost = CostService.get_project_cost(db, tenant_id, project.id)
|
|
if project_cost >= project.budget_limit_cents:
|
|
raise BudgetExceededError(
|
|
f"Project '{project.name}' budget exceeded: "
|
|
f"${project_cost/100:.2f} >= ${project.budget_limit_cents/100:.2f}"
|
|
)
|
|
|
|
@staticmethod
|
|
def get_tenant_cost_summary(db: Session, tenant_id: int) -> dict:
|
|
"""Get cost summary for tenant dashboard."""
|
|
total_cost = (
|
|
db.query(func.sum(CostLog.cost_cents))
|
|
.filter(CostLog.tenant_id == tenant_id)
|
|
.scalar()
|
|
) or 0
|
|
|
|
total_calls = (
|
|
db.query(CostLog)
|
|
.filter(CostLog.tenant_id == tenant_id)
|
|
.count()
|
|
)
|
|
|
|
by_model = (
|
|
db.query(
|
|
CostLog.model_name,
|
|
func.sum(CostLog.cost_cents).label("cost"),
|
|
func.sum(CostLog.total_tokens).label("tokens"),
|
|
func.count().label("calls"),
|
|
)
|
|
.filter(CostLog.tenant_id == tenant_id)
|
|
.group_by(CostLog.model_name)
|
|
.all()
|
|
)
|
|
|
|
return {
|
|
"total_cost_cents": total_cost,
|
|
"total_calls": total_calls,
|
|
"by_model": [
|
|
{
|
|
"model": r.model_name,
|
|
"cost_cents": r.cost or 0,
|
|
"tokens": r.tokens or 0,
|
|
"calls": r.calls,
|
|
}
|
|
for r in by_model
|
|
],
|
|
}
|