"""Usage domain — business logic: Redis counters, snapshot queries, AI usage ingestion.""" from __future__ import annotations from datetime import date, datetime, timezone from decimal import Decimal from typing import Optional from uuid import UUID import structlog from src.domain.usage.models import WorkspaceUsageSnapshot from src.domain.usage.repository import ( ProjectHostedAppUsageSnapshotRepository, WorkspaceUsageSnapshotRepository, ) from src.domain.usage.schemas import ( AIUsageAckResponse, AIUsageReportRequest, CurrentUsageResponse, ProjectUsageBreakdownResponse, ProjectHostedAppUsageHistoryResponse, ProjectHostedAppUsageSnapshotResponse, UsageMetrics, UsageSnapshotResponse, WorkspaceUsageHistoryResponse, ) logger = structlog.get_logger() # Redis key pattern for real-time counters _REDIS_KEY = "usage:current:{workspace_id}:{period}:{metric}" # Credit cost per unit (configurable — stub values, will move to DB config) _CREDIT_COST: dict[str, float] = { "embed_tokens": 0.00001, "completion_input_tokens": 0.00003, "completion_output_tokens": 0.00006, "search_queries": 0.001, "job_executions": 0.005, "ext_queries": 0.001, } class UsageService: def __init__( self, *, snapshot_repo: WorkspaceUsageSnapshotRepository, hosted_app_snapshot_repo: Optional[ProjectHostedAppUsageSnapshotRepository] = None, ) -> None: self._snapshots = snapshot_repo self._hosted_app_snapshots = hosted_app_snapshot_repo # ------------------------------------------------------------------ # Current-period usage (Redis) # ------------------------------------------------------------------ async def get_current_usage( self, workspace_id: UUID, billing_period_start: date, billing_period_end: date, plan_credits: Optional[float] = None, ) -> CurrentUsageResponse: metrics = await _read_redis_counters(workspace_id, billing_period_start) credits_used_pct: Optional[float] = None if plan_credits and plan_credits > 0: credits_used_pct = round(metrics.credits_consumed / plan_credits * 100, 2) return CurrentUsageResponse( workspace_id=workspace_id, billing_period_start=billing_period_start, billing_period_end=billing_period_end, metrics=metrics, plan_credits=plan_credits, credits_used_pct=credits_used_pct, ) # ------------------------------------------------------------------ # AI usage ingestion (called by POST /v1/cloud/usage/ai) # ------------------------------------------------------------------ async def record_ai_usage( self, body: AIUsageReportRequest, ) -> AIUsageAckResponse: workspace_id = UUID(body.workspace_id) period = _current_period_key() metric_map = { "embed_tokens": "ai_embed_tokens", "completion_input_tokens": "ai_completion_tokens", "completion_output_tokens": "ai_completion_tokens", "search_queries": "ai_searches", } redis_metric = metric_map.get(body.usage_type) if redis_metric: await _incr_redis_counter(workspace_id, period, redis_metric, body.count) if not body.is_byok: cost = _CREDIT_COST.get(body.usage_type, 0) * body.count if cost > 0: await _incr_redis_counter_float(workspace_id, period, "credits_consumed", cost) logger.info( "ai_usage_recorded", workspace_id=str(workspace_id), usage_type=body.usage_type, count=body.count, is_byok=body.is_byok, ) return AIUsageAckResponse(ok=True) # ------------------------------------------------------------------ # Snapshot queries (historical) # ------------------------------------------------------------------ async def get_workspace_usage_history( self, workspace_id: UUID, from_date: Optional[date] = None, to_date: Optional[date] = None, ) -> WorkspaceUsageHistoryResponse: rows = await self._snapshots.list_for_workspace( workspace_id, from_date=from_date, to_date=to_date ) return WorkspaceUsageHistoryResponse( workspace_id=workspace_id, snapshots=[_snapshot_response(r) for r in rows], ) async def get_latest_snapshot( self, workspace_id: UUID, period_start: date ) -> Optional[WorkspaceUsageSnapshot]: return await self._snapshots.get_latest_for_period(workspace_id, period_start) async def get_project_hosted_app_usage_history( self, project_id: UUID, *, from_date: Optional[date] = None, to_date: Optional[date] = None, ) -> ProjectHostedAppUsageHistoryResponse: if self._hosted_app_snapshots is None: return ProjectHostedAppUsageHistoryResponse(project_id=project_id, snapshots=[]) rows = await self._hosted_app_snapshots.list_for_project( project_id, from_date=from_date, to_date=to_date ) return ProjectHostedAppUsageHistoryResponse( project_id=project_id, snapshots=[ ProjectHostedAppUsageSnapshotResponse( project_id=row.project_id, snapshot_date=row.snapshot_date, requests=row.requests, egress_bytes=row.egress_bytes, created_at=row.created_at, ) for row in rows ], ) # --------------------------------------------------------------------------- # Redis helpers # --------------------------------------------------------------------------- def _current_period_key() -> str: today = date.today() return f"{today.year}-{today.month:02d}" async def _incr_redis_counter( workspace_id: UUID, period: str, metric: str, amount: int ) -> None: from src.core.redis import get_redis key = _REDIS_KEY.format( workspace_id=str(workspace_id), period=period, metric=metric ) async for redis in get_redis(): await redis.incrbyfloat(key, amount) break async def _incr_redis_counter_float( workspace_id: UUID, period: str, metric: str, amount: float ) -> None: from src.core.redis import get_redis key = _REDIS_KEY.format( workspace_id=str(workspace_id), period=period, metric=metric ) async for redis in get_redis(): await redis.incrbyfloat(key, amount) break async def _read_redis_counters(workspace_id: UUID, period_start: date) -> UsageMetrics: from src.core.redis import get_redis period = f"{period_start.year}-{period_start.month:02d}" metrics_keys = [ "cpu_minutes", "db_reads", "db_writes", "storage_gb", "cdn_gb", "hosted_app_requests", "hosted_app_egress_bytes", "auth_users_active", "ai_embed_tokens", "ai_completion_tokens", "ai_searches", "job_executions", "realtime_minutes", "ext_queries", "credits_consumed", ] values: dict[str, float] = {} async for redis in get_redis(): pipe = redis.pipeline() for metric in metrics_keys: key = _REDIS_KEY.format( workspace_id=str(workspace_id), period=period, metric=metric ) pipe.get(key) results = await pipe.execute() for metric, val in zip(metrics_keys, results): values[metric] = float(val or 0) break return UsageMetrics( cpu_minutes=values.get("cpu_minutes", 0), db_reads=int(values.get("db_reads", 0)), db_writes=int(values.get("db_writes", 0)), storage_gb=values.get("storage_gb", 0), cdn_gb=values.get("cdn_gb", 0), hosted_app_requests=int(values.get("hosted_app_requests", 0)), hosted_app_egress_bytes=int(values.get("hosted_app_egress_bytes", 0)), auth_users_active=int(values.get("auth_users_active", 0)), ai_embed_tokens=int(values.get("ai_embed_tokens", 0)), ai_completion_tokens=int(values.get("ai_completion_tokens", 0)), ai_searches=int(values.get("ai_searches", 0)), job_executions=int(values.get("job_executions", 0)), realtime_minutes=values.get("realtime_minutes", 0), ext_queries=int(values.get("ext_queries", 0)), credits_consumed=values.get("credits_consumed", 0), ) # --------------------------------------------------------------------------- # Response builders # --------------------------------------------------------------------------- def _snapshot_response(row: WorkspaceUsageSnapshot) -> UsageSnapshotResponse: return UsageSnapshotResponse( workspace_id=row.workspace_id, snapshot_date=row.snapshot_date, metrics=UsageMetrics( cpu_minutes=float(row.cpu_minutes), db_reads=row.db_reads, db_writes=row.db_writes, storage_gb=float(row.storage_gb), cdn_gb=float(row.cdn_gb), hosted_app_requests=row.hosted_app_requests, hosted_app_egress_bytes=row.hosted_app_egress_bytes, auth_users_active=row.auth_users_active, ai_embed_tokens=row.ai_embed_tokens, ai_completion_tokens=row.ai_completion_tokens, ai_searches=row.ai_searches, job_executions=row.job_executions, realtime_minutes=float(row.realtime_minutes), ext_queries=row.ext_queries, credits_consumed=float(row.credits_consumed), ), created_at=row.created_at, )