feat: AI Worker 平台 MVP v1.0.0 - 多租户/项目管理/AI Worker/任务编排/HITL审核/成本治理

2026-08-12 13:31:00 +08:00
parent 44748b11ca
commit aacb9528f4
1 changed files with 127 additions and 0 deletions
@@ -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
+            ],
+        }