diff --git a/backend/app/services/cost_service.py b/backend/app/services/cost_service.py new file mode 100644 index 0000000..da811fe --- /dev/null +++ b/backend/app/services/cost_service.py @@ -0,0 +1,127 @@ +"""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 + ], + }