Aether/_deprecated_py_src/services/billing/service.py

from __future__ import annotations

from datetime import datetime, timezone
from decimal import Decimal
from typing import Any

from sqlalchemy.orm import Session

from src.config.settings import config
from src.core.logger import logger
from src.services.billing.dimension_collector_service import DimensionCollectorService
from src.services.billing.formula_engine import BillingIncompleteError, FormulaEngine
from src.services.billing.precision import quantize_cost, to_decimal
from src.services.billing.rule_service import BillingRuleService

from .schema import BILLING_SNAPSHOT_SCHEMA_VERSION, BillingSnapshot, CostResult


class BillingService:
    """
    BillingService (pure-ish application helper for billing domain).

    Notes:
    - This service **does not** write Usage rows.
    - It may read billing rules & collectors from DB.
    """

    # FormulaEngine is stateless and safe to share within a process.
    _shared_formula_engine: FormulaEngine | None = None

    def __init__(self, db: Session):
        self.db = db
        self._formula_engine = self._get_formula_engine()
        # Lazy-init: most call sites already provide dimensions (hot path).
        self._dimension_collector: DimensionCollectorService | None = None

    @classmethod
    def _get_formula_engine(cls) -> FormulaEngine:
        if cls._shared_formula_engine is None:
            cls._shared_formula_engine = FormulaEngine()
        return cls._shared_formula_engine

    def _get_dimension_collector(self) -> DimensionCollectorService:
        if self._dimension_collector is None:
            self._dimension_collector = DimensionCollectorService(self.db)
        return self._dimension_collector

    def collect_dimensions(
        self,
        *,
        api_format: str | None,
        task_type: str | None,
        request: dict[str, Any] | None = None,
        response: dict[str, Any] | None = None,
        metadata: dict[str, Any] | None = None,
        base_dimensions: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        return self._get_dimension_collector().collect_dimensions(
            api_format=api_format,
            task_type=task_type,
            request=request,
            response=response,
            metadata=metadata,
            base_dimensions=base_dimensions,
        )

    def calculate(
        self,
        *,
        task_type: str,
        model: str,
        provider_id: str,
        dimensions: dict[str, Any],
        strict_mode: bool | None = None,
    ) -> CostResult:
        """
        Calculate cost for a task.

        Returns:
            CostResult (includes BillingSnapshot)

        Raises:
            BillingIncompleteError: when strict_mode=True and required dims missing.
        """
        strict = config.billing_strict_mode if strict_mode is None else bool(strict_mode)

        # Normalize & enrich dimensions (do not mutate caller dict)
        dims: dict[str, Any] = dict(dimensions or {})

        # Compatibility aliases (legacy fields in some call sites)
        if "cache_creation_tokens" not in dims and "cache_creation_input_tokens" in dims:
            dims["cache_creation_tokens"] = dims.get("cache_creation_input_tokens")
        if "cache_read_tokens" not in dims and "cache_read_input_tokens" in dims:
            dims["cache_read_tokens"] = dims.get("cache_read_input_tokens")

        # Default request_count=1 for per-request billing
        if "request_count" not in dims:
            dims["request_count"] = 1

        # total_input_context is the tier-key for legacy tiered pricing:
        # default: input_tokens + cache_creation_tokens + cache_read_tokens
        #
        # NOTE:
        # Some adapters (e.g. Claude) include cache_creation tokens in the tier context.
        # Making this the default avoids per-callsite inconsistency.
        if "total_input_context" not in dims:
            try:
                input_tokens_i = int(float(dims.get("input_tokens") or 0))
            except Exception:
                input_tokens_i = 0
            try:
                cache_creation_tokens_i = int(float(dims.get("cache_creation_tokens") or 0))
            except Exception:
                cache_creation_tokens_i = 0
            try:
                cache_read_tokens_i = int(float(dims.get("cache_read_tokens") or 0))
            except Exception:
                cache_read_tokens_i = 0
            dims["total_input_context"] = (
                input_tokens_i + cache_creation_tokens_i + cache_read_tokens_i
            )

        lookup = BillingRuleService.find_rule(
            self.db,
            provider_id=provider_id,
            model_name=model,
            task_type=task_type,
        )

        if lookup and lookup.rule and lookup.rule.expression:
            rule = lookup.rule
            result = self._formula_engine.evaluate(
                expression=rule.expression,
                variables=rule.variables or {},
                dimensions=dims,
                dimension_mappings=rule.dimension_mappings or {},
                strict_mode=strict,
            )
            # ------------------------------------------------------------
            # Quantize: component costs first, then total = sum(components)
            # ------------------------------------------------------------
            breakdown_dec: dict[str, Decimal] = {
                k: to_decimal(v) for k, v in (result.cost_breakdown or {}).items()
            }

            breakdown_quantized: dict[str, Decimal] = {
                k: quantize_cost(v) for k, v in breakdown_dec.items()
            }
            total_dec = (
                quantize_cost(sum(breakdown_quantized.values(), Decimal("0")))
                if breakdown_quantized
                else quantize_cost(to_decimal(result.cost))
            )

            cost_breakdown = {k: float(v) for k, v in breakdown_quantized.items()}
            total_cost = float(total_dec) if result.status == "complete" else 0.0

            # Filter resolved_variables for JSON safety + semantics clarity:
            # - remove dims (they live in resolved_dimensions)
            # - remove *_cost (they live in cost_breakdown)
            resolved_vars: dict[str, Any] = {}
            for k, v in (result.resolved_variables or {}).items():
                if k in (result.resolved_dimensions or {}):
                    continue
                if k.endswith("_cost"):
                    continue
                if isinstance(v, Decimal):
                    resolved_vars[k] = str(v)
                else:
                    resolved_vars[k] = v

            snapshot = BillingSnapshot(
                schema_version=BILLING_SNAPSHOT_SCHEMA_VERSION,
                rule_id=str(rule.id),
                rule_name=str(rule.name),
                scope=str(getattr(lookup, "scope", None) or ""),
                expression=str(rule.expression),
                resolved_dimensions=result.resolved_dimensions or dims,
                resolved_variables=resolved_vars,
                cost_breakdown=cost_breakdown,
                total_cost=total_cost,
                tier_index=result.tier_index,
                tier_info=result.tier_info,
                missing_required=result.missing_required,
                status=result.status,
                calculated_at=datetime.now(timezone.utc).isoformat(),
            )
            return CostResult(cost=total_cost, status=result.status, snapshot=snapshot)

        logger.warning(
            "No billing rule for task (task_type={}, model={}, provider_id={})",
            task_type,
            model,
            provider_id,
        )
        snapshot = BillingSnapshot(
            schema_version=BILLING_SNAPSHOT_SCHEMA_VERSION,
            rule_id=None,
            rule_name=None,
            scope=None,
            expression=None,
            resolved_dimensions=dims,
            resolved_variables={},
            cost_breakdown={},
            total_cost=0.0,
            missing_required=[],
            status="no_rule",
            calculated_at=datetime.now(timezone.utc).isoformat(),
        )
        return CostResult(cost=0.0, status="no_rule", snapshot=snapshot)

    def calculate_from_response(
        self,
        *,
        task_type: str,
        model: str,
        provider_id: str,
        api_format: str | None,
        request: dict[str, Any] | None = None,
        response: dict[str, Any] | None = None,
        metadata: dict[str, Any] | None = None,
        base_dimensions: dict[str, Any] | None = None,
        strict_mode: bool | None = None,
    ) -> CostResult:
        """
        Convenience wrapper:
        - collect dimensions from request/response/metadata
        - run billing calculation
        """
        dimensions = self.collect_dimensions(
            api_format=api_format,
            task_type=task_type,
            request=request,
            response=response,
            metadata=metadata,
            base_dimensions=base_dimensions,
        )
        return self.calculate(
            task_type=task_type,
            model=model,
            provider_id=provider_id,
            dimensions=dimensions,
            strict_mode=strict_mode,
        )