""" Usage Metering Service """ from datetime import datetime, date from typing import Optional, Dict, Any from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy import select, func import logging from app.models.billing import UsageLog, UsageType, UsageSummary from app.models.client import Client from app.models.subscription import Subscription logger = logging.getLogger(__name__) class UsageService: """Track and meter usage for billing""" # Cost per unit (in INR) COST_PER_TOKEN = 0.0001 # ₹0.0001 per token COST_PER_MESSAGE = 0.05 # ₹0.05 per message COST_PER_VISION_REQUEST = 0.50 # ₹0.50 per vision request COST_PER_WHATSAPP_MESSAGE = 0.10 # ₹0.10 per WhatsApp message async def log_chat_usage( self, db: AsyncSession, client_id: int, conversation_id: int, input_tokens: int, output_tokens: int, model: str, api_key_id: Optional[int] = None ) -> UsageLog: """ Log chat message usage Args: db: Database session client_id: Client ID conversation_id: Conversation ID input_tokens: Input tokens used output_tokens: Output tokens used model: AI model used api_key_id: API key ID used Returns: Created usage log """ total_tokens = input_tokens + output_tokens total_cost = total_tokens * self.COST_PER_TOKEN usage_log = UsageLog( client_id=client_id, usage_type=UsageType.CHAT_TOKEN, quantity=total_tokens, conversation_id=conversation_id, api_key_id=api_key_id, model=model, input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=total_tokens, unit_cost=self.COST_PER_TOKEN, total_cost=total_cost, billing_period=datetime.now().strftime("%Y-%m") ) db.add(usage_log) await db.commit() # Update daily summary await self._update_daily_summary(db, client_id, UsageType.CHAT_TOKEN, total_tokens, total_cost) return usage_log async def log_message_usage( self, db: AsyncSession, client_id: int, conversation_id: int, channel: str = "web" ) -> UsageLog: """ Log a single message (for message-based billing) Args: db: Database session client_id: Client ID conversation_id: Conversation ID channel: Channel (web, whatsapp, etc.) Returns: Created usage log """ if channel == "whatsapp": usage_type = UsageType.WHATSAPP_MESSAGE unit_cost = self.COST_PER_WHATSAPP_MESSAGE else: usage_type = UsageType.CHAT_MESSAGE unit_cost = self.COST_PER_MESSAGE usage_log = UsageLog( client_id=client_id, usage_type=usage_type, quantity=1, conversation_id=conversation_id, unit_cost=unit_cost, total_cost=unit_cost, billing_period=datetime.now().strftime("%Y-%m") ) db.add(usage_log) await db.commit() return usage_log async def log_vision_usage( self, db: AsyncSession, client_id: int, images_processed: int, job_type: str ) -> UsageLog: """ Log computer vision usage Args: db: Database session client_id: Client ID images_processed: Number of images processed job_type: Type of vision job Returns: Created usage log """ total_cost = images_processed * self.COST_PER_VISION_REQUEST usage_log = UsageLog( client_id=client_id, usage_type=UsageType.VISION_REQUEST, quantity=images_processed, unit_cost=self.COST_PER_VISION_REQUEST, total_cost=total_cost, billing_period=datetime.now().strftime("%Y-%m"), metadata={"job_type": job_type} ) db.add(usage_log) await db.commit() return usage_log async def get_usage_summary( self, db: AsyncSession, client_id: int, period: Optional[str] = None ) -> Dict[str, Any]: """ Get usage summary for a client Args: db: Database session client_id: Client ID period: Billing period (YYYY-MM format), defaults to current month Returns: Usage summary dict """ if not period: period = datetime.now().strftime("%Y-%m") # Get totals by usage type result = await db.execute( select( UsageLog.usage_type, func.sum(UsageLog.quantity).label("total_quantity"), func.sum(UsageLog.total_cost).label("total_cost") ) .where( UsageLog.client_id == client_id, UsageLog.billing_period == period ) .group_by(UsageLog.usage_type) ) usage_by_type = {} total_cost = 0 for row in result: usage_type = row.usage_type.value if row.usage_type else "unknown" usage_by_type[usage_type] = { "quantity": int(row.total_quantity or 0), "cost": float(row.total_cost or 0) } total_cost += float(row.total_cost or 0) # Get subscription limits sub_result = await db.execute( select(Subscription) .where(Subscription.client_id == client_id) .order_by(Subscription.created_at.desc()) .limit(1) ) subscription = sub_result.scalar_one_or_none() return { "period": period, "usage_by_type": usage_by_type, "total_cost": total_cost, "messages_used": subscription.messages_used if subscription else 0, "messages_limit": 5000, # Get from plan "is_over_limit": False } async def check_usage_limits( self, db: AsyncSession, client_id: int ) -> Dict[str, Any]: """ Check if client is within usage limits Args: db: Database session client_id: Client ID Returns: Limit check result """ # Get client result = await db.execute(select(Client).where(Client.id == client_id)) client = result.scalar_one_or_none() if not client: return {"allowed": False, "reason": "Client not found"} # Get current month usage period = datetime.now().strftime("%Y-%m") usage_result = await db.execute( select(func.count(UsageLog.id)) .where( UsageLog.client_id == client_id, UsageLog.billing_period == period, UsageLog.usage_type == UsageType.CHAT_MESSAGE ) ) messages_used = usage_result.scalar() or 0 # Check against limit if messages_used >= client.monthly_message_limit: return { "allowed": False, "reason": "Monthly message limit reached", "messages_used": messages_used, "messages_limit": client.monthly_message_limit } return { "allowed": True, "messages_used": messages_used, "messages_limit": client.monthly_message_limit, "messages_remaining": client.monthly_message_limit - messages_used } async def _update_daily_summary( self, db: AsyncSession, client_id: int, usage_type: UsageType, quantity: int, cost: float ): """Update daily usage summary""" today = date.today().isoformat() result = await db.execute( select(UsageSummary) .where( UsageSummary.client_id == client_id, UsageSummary.period_type == "daily", UsageSummary.period_date == today ) ) summary = result.scalar_one_or_none() if summary: if usage_type == UsageType.CHAT_TOKEN: summary.total_tokens += quantity elif usage_type == UsageType.CHAT_MESSAGE: summary.total_messages += 1 summary.total_cost += cost else: summary = UsageSummary( client_id=client_id, period_type="daily", period_date=today, total_messages=1 if usage_type == UsageType.CHAT_MESSAGE else 0, total_tokens=quantity if usage_type == UsageType.CHAT_TOKEN else 0, total_cost=cost ) db.add(summary) await db.commit() # Singleton instance usage_service = UsageService()