PMFlow/backend/app/services/cost_optimization_service.py

"""成本优化服务（P1）"""

from typing import Any, Dict, Optional

from sqlalchemy import and_, func, select
from sqlalchemy.ext.asyncio import AsyncSession

from ..models import AICallLog, CostTracking, ModelRequirementBinding, ModelVersion, TruthNode


class CostOptimizationService:
    """面向宪章维度的成本优化建议。"""

    def __init__(self, db: AsyncSession):
        self.db = db

    async def analyze_cost_optimization(
        self,
        truth_node_id: int,
        tenant_id: int,
        period: str = "weekly",
        project_id: Optional[int] = None,
    ) -> Dict[str, Any]:
        if project_id is not None:
            await self.validate_truth_node_in_project(
                truth_node_id=truth_node_id,
                tenant_id=tenant_id,
                project_id=project_id,
            )
        tracking = await self._get_tracking(truth_node_id, tenant_id, period)
        if not tracking:
            return {
                "truth_node_id": truth_node_id,
                "period": period,
                "monthly_cost": 0.0,
                "recommendations": [],
                "total_potential_savings_percent": 0,
            }

        monthly_cost = float(tracking.actual_cost_usd or 0.0) * self._period_multiplier(period)
        recs = []
        recs.extend(await self._recommend_model_switch(truth_node_id, tenant_id))
        recs.extend(await self._recommend_cache_and_batching(truth_node_id, tenant_id))

        total_pct = sum(int(item.get("estimated_savings_percent", 0)) for item in recs)
        total_pct = min(total_pct, 60)

        result = {
            "truth_node_id": truth_node_id,
            "period": period,
            "monthly_cost": round(monthly_cost, 2),
            "recommendations": recs,
            "total_potential_savings_percent": total_pct,
            "estimated_monthly_savings_usd": round(monthly_cost * total_pct / 100, 2),
        }
        tracking.optimization_suggestions = result
        await self.db.commit()
        return result

    async def validate_truth_node_in_project(
        self, truth_node_id: int, tenant_id: int, project_id: int
    ) -> TruthNode:
        result = await self.db.execute(
            select(TruthNode).where(
                and_(
                    TruthNode.id == truth_node_id,
                    TruthNode.tenant_id == tenant_id,
                )
            )
        )
        truth_node = result.scalar_one_or_none()
        if not truth_node:
            raise ValueError(f"宪章 {truth_node_id} 不存在")
        if truth_node.project_id != project_id:
            raise ValueError(
                f"宪章 {truth_node_id} 不属于项目 {project_id}，拒绝跨项目分析"
            )
        return truth_node

    async def _get_tracking(
        self, truth_node_id: int, tenant_id: int, period: str
    ) -> Optional[CostTracking]:
        result = await self.db.execute(
            select(CostTracking).where(
                and_(
                    CostTracking.truth_node_id == truth_node_id,
                    CostTracking.tenant_id == tenant_id,
                    CostTracking.period == period,
                )
            )
        )
        return result.scalar_one_or_none()

    async def _recommend_model_switch(
        self, truth_node_id: int, tenant_id: int
    ) -> list[Dict[str, Any]]:
        result = await self.db.execute(
            select(ModelRequirementBinding).where(
                and_(
                    ModelRequirementBinding.truth_node_id == truth_node_id,
                    ModelRequirementBinding.tenant_id == tenant_id,
                )
            )
        )
        bindings = result.scalars().all()
        if not bindings:
            return []

        current = bindings[0]
        model = await self._get_model(current.model_version_id)
        if not model or not model.cost_per_1k_tokens:
            return []

        cheaper = await self._find_cheaper_compatible_model(float(model.cost_per_1k_tokens))
        if not cheaper:
            return []

        return [
            {
                "type": "model_switch",
                "priority": "high",
                "description": f"可从 {model.name} 切换到 {cheaper.name}（更低 token 成本）",
                "estimated_savings_percent": 25,
            }
        ]

    async def _recommend_cache_and_batching(
        self, truth_node_id: int, tenant_id: int
    ) -> list[Dict[str, Any]]:
        # 当前日志没有强关联 truth_node 字段，按租户统计调用频次近似分析。
        call_count_res = await self.db.execute(
            select(func.count(AICallLog.id)).where(AICallLog.tenant_id == str(tenant_id))
        )
        call_count = int(call_count_res.scalar() or 0)

        recs: list[Dict[str, Any]] = [
            {
                "type": "cache",
                "priority": "high",
                "description": "启用/提升相似请求缓存命中率（去重 prompt + TTL 分层）",
                "estimated_savings_percent": 15,
            }
        ]
        if call_count > 200:
            recs.append(
                {
                    "type": "batching",
                    "priority": "medium",
                    "description": "批量聚合低优先级推理请求，减少并发和重复上下文开销",
                    "estimated_savings_percent": 10,
                }
            )
        return recs

    async def _get_model(self, model_version_id: int) -> Optional[ModelVersion]:
        result = await self.db.execute(
            select(ModelVersion).where(ModelVersion.id == model_version_id)
        )
        return result.scalar_one_or_none()

    async def _find_cheaper_compatible_model(
        self, current_cost_per_1k_tokens: float
    ) -> Optional[ModelVersion]:
        result = await self.db.execute(
            select(ModelVersion)
            .where(ModelVersion.is_active == True)  # noqa: E712
            .order_by(ModelVersion.cost_per_1k_tokens.asc())
        )
        candidates = result.scalars().all()
        for m in candidates:
            cost = float(m.cost_per_1k_tokens or 0.0)
            if cost > 0 and cost < current_cost_per_1k_tokens:
                return m
        return None

    def _period_multiplier(self, period: str) -> int:
        if period == "daily":
            return 30
        if period == "weekly":
            return 4
        if period == "monthly":
            return 1
        return 1