Agentic-OS/backend/app/services/balance.py

"""Poll OpenRouter and DeepSeek for credits/balance and store snapshots."""
from __future__ import annotations

import asyncio
import logging

import httpx
from sqlalchemy import select

from app.config import settings
from app.database import SessionLocal
from app.models.balance import BalanceSnapshot

logger = logging.getLogger(__name__)


async def fetch_openrouter() -> dict:
    """Return {total_credits, total_usage, remaining} for OpenRouter.

    Prefers the provisioning key against /credits (account totals); falls back to
    /key (per-key limit/usage) when only the regular API key is available.
    """
    key = settings.openrouter_provisioning_key or settings.openrouter_api_key
    if not key:
        return {"error": "no OpenRouter key configured"}

    async with httpx.AsyncClient(timeout=20) as client:
        # Try account-wide credits first
        if settings.openrouter_provisioning_key:
            r = await client.get(
                "https://openrouter.ai/api/v1/credits",
                headers={"Authorization": f"Bearer {settings.openrouter_provisioning_key}"},
            )
            if r.status_code == 200:
                data = r.json().get("data", {})
                total = data.get("total_credits")
                usage = data.get("total_usage")
                remaining = (total - usage) if total is not None and usage is not None else None
                return {
                    "total_credits": total,
                    "total_usage": usage,
                    "remaining": remaining,
                    "raw": data,
                }

        # Fall back to per-key info
        r = await client.get(
            "https://openrouter.ai/api/v1/key",
            headers={"Authorization": f"Bearer {settings.openrouter_api_key}"},
        )
        r.raise_for_status()
        data = r.json().get("data", {})
        limit = data.get("limit")
        usage = data.get("usage")
        remaining = data.get("limit_remaining")
        if remaining is None and limit is not None and usage is not None:
            remaining = limit - usage
        return {
            "total_credits": limit,
            "total_usage": usage,
            "remaining": remaining,
            "raw": data,
        }


async def fetch_deepseek() -> dict:
    if not settings.deepseek_api_key:
        return {"error": "no DeepSeek key configured"}
    async with httpx.AsyncClient(timeout=20) as client:
        r = await client.get(
            "https://api.deepseek.com/user/balance",
            headers={"Authorization": f"Bearer {settings.deepseek_api_key}"},
        )
        r.raise_for_status()
        data = r.json()
        infos = data.get("balance_infos", [])
        info = infos[0] if infos else {}
        total = float(info.get("total_balance")) if info.get("total_balance") else None
        return {
            "total_credits": None,
            "total_usage": None,
            "remaining": total,
            "currency": info.get("currency", "USD"),
            "raw": data,
        }


async def _store(provider: str, data: dict) -> None:
    async with SessionLocal() as db:
        db.add(
            BalanceSnapshot(
                provider=provider,
                total_credits=data.get("total_credits"),
                total_usage=data.get("total_usage"),
                remaining=data.get("remaining"),
                currency=data.get("currency", "USD"),
                raw=data.get("raw"),
                error=data.get("error"),
            )
        )
        await db.commit()


async def poll_once() -> None:
    for provider, fetcher in (("openrouter", fetch_openrouter), ("deepseek", fetch_deepseek)):
        try:
            data = await fetcher()
        except Exception as exc:  # noqa: BLE001
            data = {"error": str(exc)}
            logger.warning("Balance poll failed for %s: %s", provider, exc)
        if data.get("error"):
            # Keep last good snapshot visible when DNS/network blips.
            async with SessionLocal() as db:
                res = await db.execute(
                    select(BalanceSnapshot)
                    .where(BalanceSnapshot.provider == provider)
                    .where(BalanceSnapshot.error.is_(None))
                    .where(BalanceSnapshot.remaining.is_not(None))
                    .order_by(BalanceSnapshot.created_at.desc())
                    .limit(1)
                )
                if res.scalar_one_or_none():
                    logger.info("Skipping error balance snapshot for %s: %s", provider, data["error"])
                    continue
        await _store(provider, data)


async def latest_remaining_usd() -> float | None:
    """Best estimate of remaining USD across configured providers (max)."""
    async with SessionLocal() as db:
        remainings: list[float] = []
        for provider in ("openrouter", "deepseek"):
            res = await db.execute(
                select(BalanceSnapshot)
                .where(BalanceSnapshot.provider == provider)
                .order_by(BalanceSnapshot.created_at.desc())
                .limit(1)
            )
            snap = res.scalar_one_or_none()
            if snap and snap.remaining is not None:
                remainings.append(snap.remaining)
        return max(remainings) if remainings else None


async def balance_poller_loop() -> None:
    while True:
        await poll_once()
        await asyncio.sleep(settings.balance_poll_interval_seconds)