UsageKiosk/app/providers/devin.py

import math
import time
from datetime import datetime
from typing import Any
from urllib.parse import quote
from zoneinfo import ZoneInfo

import httpx
from pydantic import BaseModel, ConfigDict, Field, StrictBool, StrictInt, StrictStr, ValidationError

from app.config import Settings
from app.models import ProviderUsage, UsageLimit
from app.providers.base import GENERIC_ERROR, UsageProvider, utc_now


class Cycle(BaseModel):
    model_config = ConfigDict(extra='ignore')
    after: StrictInt
    before: StrictInt


class CyclesResponse(BaseModel):
    model_config = ConfigDict(extra='ignore')
    items: list[Cycle]
    has_next_page: StrictBool = False
    end_cursor: StrictStr | None = None


class DailyResponse(BaseModel):
    model_config = ConfigDict(extra='ignore')
    total_acus: Any
    consumption_by_date: list[Any] = Field(default_factory=list)


class OrgResponse(BaseModel):
    model_config = ConfigDict(extra='ignore')
    max_cycle_acu_limit: Any = None


def _acu(value: Any) -> float:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError
    number = float(value)
    if not math.isfinite(number) or number < 0:
        raise ValueError
    return number


class DevinProvider(UsageProvider):
    def __init__(self, settings: Settings, client: httpx.AsyncClient | None = None):
        super().__init__('devin', settings.provider_cache_ttl, 'Devin API v3')
        self.settings = settings
        self._client = client or httpx.AsyncClient(base_url=settings.devin_api_url.rstrip('/'), timeout=settings.devin_timeout,
                                                   headers={'Authorization': f'Bearer {settings.devin_api_key.get_secret_value()}'})
        self._owns_client = client is None

    async def _fetch_usage(self) -> ProviderUsage:
        now = int(time.time())
        params = {'first': 100}
        seen_cursors: set[str] = set()
        cycle = None
        for page in range(100):
            cycles_response = await self._client.get('/v3/enterprise/consumption/cycles', params=params)
            cycles_response.raise_for_status()
            try:
                cycles_page = CyclesResponse.model_validate(cycles_response.json())
            except (ValidationError, ValueError, TypeError):
                raise ValueError
            cycle = next((item for item in cycles_page.items if item.after <= now < item.before), None)
            if cycle is not None:
                break
            if not cycles_page.has_next_page:
                break
            cursor = cycles_page.end_cursor
            if not cursor or cursor in seen_cursors:
                raise ValueError
            seen_cursors.add(cursor)
            params = {'first': 100, 'after': cursor}
        else:
            raise ValueError
        if cycle is None:
            raise ValueError
        user_id = self.settings.devin_user_id.strip()
        org_id = self.settings.devin_org_id.strip()
        if user_id:
            path = f'/v3/enterprise/consumption/daily/users/{quote(user_id, safe="")}'
            org_scope = False
        elif org_id:
            path = f'/v3/enterprise/consumption/daily/organizations/{quote(org_id, safe="")}'
            org_scope = True
        else:
            path = '/v3/enterprise/consumption/daily'
            org_scope = False
        daily_response = await self._client.get(path, params={'time_after': cycle.after, 'time_before': min(now, cycle.before)})
        daily_response.raise_for_status()
        try:
            daily = DailyResponse.model_validate(daily_response.json())
            used = _acu(daily.total_acus)
        except (ValidationError, ValueError, TypeError):
            raise ValueError
        limit = remaining = used_percent = remaining_percent = None
        if org_scope:
            org_path = f'/v3/enterprise/organizations/{quote(org_id, safe="")}'
            try:
                org_response = await self._client.get(org_path)
                if org_response.is_success:
                    org_data = OrgResponse.model_validate(org_response.json())
                    if org_data.max_cycle_acu_limit is not None:
                        limit = _acu(org_data.max_cycle_acu_limit)
                        if limit > 0:
                            remaining = max(limit - used, 0)
                            used_percent = min(used / limit * 100, 100)
                            remaining_percent = max(100 - used_percent, 0)
                        else:
                            remaining = 0
                            used_percent = 100 if used else 0
                            remaining_percent = 0 if used else 100
            except (httpx.HTTPError, ValidationError, ValueError, TypeError):
                pass
        updated = utc_now()
        return ProviderUsage(provider='devin', status='ok', limits=[UsageLimit(
            id='cycle-acu-usage', name='Cycle ACU Usage', used=used, remaining=remaining, limit=limit,
            used_percent=used_percent, remaining_percent=remaining_percent, unit='ACUs',
            reset_at=datetime.fromtimestamp(cycle.before, tz=ZoneInfo(self.settings.timezone)),
            window='billing cycle', updated_at=updated)], last_successful_update=updated, source=self.source)

    async def aclose(self) -> None:
        if self._owns_client:
            await self._client.aclose()