Agent/examples/tencent_realtime_control/today_spend_query.py

"""Read-only Tencent daily spend summaries for Feishu commands."""

from __future__ import annotations

import logging
from datetime import datetime
from typing import Any

from operator_commands import (
    SCOPE_ACCOUNTS,
    SCOPE_ALL,
    SCOPE_AUTOMATION,
    ParsedCommand,
)
from storage import load_all_spend_accounts, load_automation_spend_accounts
from tencent_client import TencentClient


logger = logging.getLogger("tencent_realtime_control.today_spend")


def _select_accounts(parsed: ParsedCommand) -> tuple[str, list[int]]:
    all_accounts = {
        int(row["account_id"]) for row in load_all_spend_accounts()
    }
    if parsed.scope_type == SCOPE_ALL:
        return "全部账户", sorted(all_accounts)
    if parsed.scope_type == SCOPE_AUTOMATION:
        automation_ids = {
            int(row["account_id"]) for row in load_automation_spend_accounts()
        }
        return "自动化账户", sorted(automation_ids & set(all_accounts))
    if parsed.scope_type == SCOPE_ACCOUNTS:
        missing = [value for value in parsed.account_ids if value not in all_accounts]
        if missing:
            raise ValueError(f"账户不在启用白名单: {missing}")
        return "指定账户", list(parsed.account_ids)
    raise ValueError(f"不支持的今日消耗查询范围: {parsed.scope_type}")


def query_today_spend(
    parsed: ParsedCommand,
    *,
    now: datetime,
    tencent: TencentClient | None = None,
) -> dict[str, Any]:
    scope_label, account_ids = _select_accounts(parsed)
    if not account_ids:
        raise ValueError(f"{scope_label}范围内没有可查询账户")

    results: list[dict[str, int]] = []
    failed_account_ids: list[int] = []
    client = tencent or TencentClient()
    for account_id in account_ids:
        try:
            results.append(
                client.get_today_account_metrics(account_id, now.date())
            )
        except Exception:
            failed_account_ids.append(account_id)
            logger.exception("Today spend query failed account=%s", account_id)

    return {
        "scope_label": scope_label,
        "data_date": now.date(),
        "queried_at": now,
        "account_count": len(account_ids),
        "successful_account_count": len(results),
        "failed_account_count": len(failed_account_ids),
        "spending_account_count": sum(
            1 for row in results if int(row.get("cost_fen") or 0) > 0
        ),
        "cost_fen": sum(int(row.get("cost_fen") or 0) for row in results),
        "impressions": sum(int(row.get("impressions") or 0) for row in results),
        "clicks": sum(int(row.get("clicks") or 0) for row in results),
        "conversions": sum(int(row.get("conversions") or 0) for row in results),
        "complete": not failed_account_ids,
        "failed_account_ids": failed_account_ids,
    }