adaptive-pricing/projects/coulomb-pricing/observatory/tuning.py

from __future__ import annotations

from decimal import Decimal
from typing import Any

from ._repo_root import ensure_repo_root_on_syspath
from .boundary import build_boundary_policy
from .ltv import _configuration, _ltv_policy, _profile, _usage_unit_cost
from .models import EconomicsSnapshot, PricingModel

ensure_repo_root_on_syspath()

from adaptive_pricing_core.customer_tuning import (  # noqa: E402
    CustomerTuningRequest,
    UsagePriceSearchPolicy,
    solve_customer_tuning,
)


def _serialize(value: Any) -> Any:
    if isinstance(value, Decimal):
        return str(value)
    if hasattr(value, "__dataclass_fields__"):
        return {key: _serialize(getattr(value, key)) for key in value.__dataclass_fields__}
    if isinstance(value, tuple):
        return [_serialize(item) for item in value]
    if isinstance(value, list):
        return [_serialize(item) for item in value]
    if isinstance(value, dict):
        return {key: _serialize(item) for key, item in value.items()}
    return value


def _decimal(value: Decimal | str | int | float | None) -> Decimal | None:
    if value in (None, ""):
        return None
    return Decimal(str(value))


def _customer_tunable_keys(model: PricingModel) -> set[str]:
    return {
        parameter.key
        for parameter in model.tunable_parameters
        if parameter.parameter_class == "customer_tunable"
    }


def _validate_request_surface(model: PricingModel, selected_tunables: dict[str, Any]) -> tuple[str, ...]:
    tunable_keys = _customer_tunable_keys(model)
    issues: list[str] = []
    for key in selected_tunables:
        if key not in tunable_keys:
            issues.append(f"{key} is not customer-tunable on {model.id}")
    return tuple(issues)


def _request(raw: dict[str, Any]) -> CustomerTuningRequest:
    selected = raw.get("selected_tunables", {})
    return CustomerTuningRequest(
        included_units=_decimal(selected.get("included_tokens")),
        contract_duration_months=(
            int(selected["contract_duration_months"])
            if selected.get("contract_duration_months") is not None
            else None
        ),
        minimum_monthly_turnover=_decimal(selected.get("minimum_monthly_turnover")) or Decimal("0"),
        prepaid_amount=_decimal(selected.get("prepaid_amount")) or Decimal("0"),
        guaranteed_platform_fee=_decimal(selected.get("guaranteed_platform_fee")) or Decimal("0"),
        customer_funded_onboarding=_decimal(selected.get("customer_funded_onboarding")) or Decimal("0"),
        reduced_cancellation_flexibility=raw.get("reduced_cancellation_flexibility"),
        preference=raw.get("preference", "lower_usage_price"),
        approval_mode=raw.get("approval_mode", "self_serve_only"),
    )


def _search_policy(raw: dict[str, Any]) -> UsagePriceSearchPolicy | None:
    if not raw:
        return None
    return UsagePriceSearchPolicy(
        min_usage_unit_price=_decimal(raw.get("min_usage_unit_price")),
        max_usage_unit_price=_decimal(raw.get("max_usage_unit_price")),
        usage_unit_price_step=_decimal(raw.get("usage_unit_price_step")) or Decimal("0.0001"),
        max_usage_price_multiplier=_decimal(raw.get("max_usage_price_multiplier")) or Decimal("4"),
    )


def build_customer_tuning_pilot(
    snapshot: EconomicsSnapshot,
    models: list[PricingModel],
    usage_records: list[dict[str, Any]],
    scenario_catalog: dict[str, Any],
    request_catalog: dict[str, Any] | None = None,
) -> dict[str, Any]:
    request_catalog = request_catalog or {}
    if not request_catalog.get("requests"):
        return {
            "period": snapshot.period,
            "currency": snapshot.currency,
            "requests": [],
            "notes": [
                "No customer-tuning pilot requests are configured for this observatory deployment.",
            ],
        }

    profile_index = {item["id"]: _profile(item) for item in scenario_catalog.get("profiles", [])}
    model_index = {model.id: model for model in models if model.status in ("active", "candidate")}
    policy = _ltv_policy(scenario_catalog)
    boundary_policy = build_boundary_policy(snapshot)
    observed_usage_unit_cost = _usage_unit_cost(usage_records, snapshot.period)

    results = []
    for raw_request in request_catalog.get("requests", []):
        model = model_index[raw_request["model_id"]]
        profile = profile_index[raw_request["profile_id"]]
        selected_tunables = raw_request.get("selected_tunables", {})
        issues = _validate_request_surface(model, selected_tunables)

        if issues:
            results.append(
                {
                    "id": raw_request["id"],
                    "name": raw_request["name"],
                    "decision": "rejected",
                    "issues": list(issues),
                    "profile_id": profile.id,
                    "model_id": model.id,
                }
            )
            continue

        base_configuration = _configuration(model, profile, snapshot, observed_usage_unit_cost)
        reference_configurations = [
            _configuration(candidate, profile, snapshot, observed_usage_unit_cost)
            for candidate in models
            if candidate.status in ("active", "candidate")
        ]
        outcome = solve_customer_tuning(
            base_configuration,
            reference_configurations,
            profile,
            boundary_policy,
            policy,
            _request(raw_request),
            search_policy=_search_policy(raw_request.get("search_policy", {})),
        )
        results.append(
            {
                "id": raw_request["id"],
                "name": raw_request["name"],
                "profile_id": profile.id,
                "profile_name": profile.name,
                "model_id": model.id,
                "model_name": model.name,
                "selected_tunables": selected_tunables,
                "result": outcome,
            }
        )

    accepted = [
        item["id"]
        for item in results
        if item.get("result") is not None and getattr(item["result"], "decision", None) == "accepted"
    ]

    return _serialize(
        {
            "period": snapshot.period,
            "currency": snapshot.currency,
            "request_count": len(results),
            "accepted_request_ids": accepted,
            "requests": results,
            "notes": [
                request_catalog.get("notes", ""),
                "Pilot requests map product-level tunables into canonical pricing configuration fields before running the generic solver.",
                "For Coulomb's hybrid prototype, selected included token values are treated as total package allowances rather than per-seat multipliers.",
            ],
        }
    )