Skip to content
Back to top

# Source code for mindroot.coreplugins.usage.handlers

    from typing import Optional
    from .models import UsageEvent
    from .storage import UsageStorage
    from loguru import logger
    import traceback

    class UsageTracker:
        def __init__(self, storage: UsageStorage):
            self.storage = storage

        async def get_cost(self, plugin_id: str, cost_type_id: str, model_id: Optional[str] = None) -> float:
            """Get the cost for a specific usage type and optional model"""
            costs = await self.storage.load_costs()
            try:

                plugin_costs = costs[plugin_id]
                type_costs = plugin_costs[cost_type_id]

                # Try model-specific cost first
                if model_id and 'model_specific' in type_costs:
                    if model_id in type_costs['model_specific']:
                        return type_costs['model_specific'][model_id]

                # Fall back to default cost
                return type_costs.get('default', 0.0)
            except Exception as e:
                trace = traceback.format_exc()
                logger.error(f"Error getting cost: {e}\n\n{trace}")
                return 0.0

        async def track_usage(self, event: UsageEvent):
            """Track a usage event and store it with calculated cost"""
            unit_cost = await self.get_cost(event.plugin_id, event.cost_type_id, event.model_id)
            total_cost = unit_cost * event.quantity

            await self.storage.store_event(event, total_cost)

        async def get_usage(self, username: str, start_date=None, end_date=None):
            """Get usage records for a user"""
            return await self.storage.get_usage(username, start_date, end_date)

        async def get_total_cost(self, username: str, start_date=None, end_date=None):
            """Get total cost for a user"""
            return await self.storage.get_total_cost(username, start_date, end_date)