microdao-daarion/services/agents-service/quotas.py

"""
Quotas & Rate Limits — Обмеження використання агентів
"""

import time
from typing import Dict, Optional
from datetime import datetime, timedelta
from collections import defaultdict

# ============================================================================
# Quota Configuration
# ============================================================================

class QuotaConfig:
    """Конфігурація квот"""
    
    def __init__(
        self,
        tokens_per_minute: int = 1000,
        runs_per_day: int = 100,
        users_per_day: int = 50,
        max_concurrent_runs: int = 5
    ):
        self.tokens_per_minute = tokens_per_minute
        self.runs_per_day = runs_per_day
        self.users_per_day = users_per_day
        self.max_concurrent_runs = max_concurrent_runs


# Default quota configurations
DEFAULT_QUOTAS = {
    "free": QuotaConfig(
        tokens_per_minute=500,
        runs_per_day=50,
        users_per_day=20,
        max_concurrent_runs=2
    ),
    "pro": QuotaConfig(
        tokens_per_minute=2000,
        runs_per_day=500,
        users_per_day=200,
        max_concurrent_runs=10
    ),
    "enterprise": QuotaConfig(
        tokens_per_minute=10000,
        runs_per_day=5000,
        users_per_day=1000,
        max_concurrent_runs=50
    )
}


# ============================================================================
# Quota Tracker
# ============================================================================

class QuotaTracker:
    """
    Трекер використання ресурсів агентами
    
    Тримає в пам'яті (in-memory) лічильники для:
    - Токенів за хвилину
    - Запусків за день
    - Унікальних користувачів за день
    - Паралельних запусків
    """
    
    def __init__(self):
        # Agent ID → tokens used in current minute
        self._tokens_minute: Dict[str, list[tuple[float, int]]] = defaultdict(list)
        
        # Agent ID → runs count today
        self._runs_today: Dict[str, int] = defaultdict(int)
        self._runs_today_date: Optional[str] = None
        
        # Agent ID → unique users today
        self._users_today: Dict[str, set[str]] = defaultdict(set)
        
        # Agent ID → concurrent runs count
        self._concurrent_runs: Dict[str, int] = defaultdict(int)
    
    def check_tokens_quota(self, agent_id: str, tokens: int, quota: QuotaConfig) -> bool:
        """
        Перевірити, чи не перевищено квоту токенів за хвилину
        
        Args:
            agent_id: ID агента
            tokens: Кількість токенів для використання
            quota: Конфігурація квот
        
        Returns:
            True, якщо квота дозволяє використання
        """
        now = time.time()
        one_minute_ago = now - 60
        
        # Видалити старі записи (старше 1 хвилини)
        self._tokens_minute[agent_id] = [
            (ts, count) for ts, count in self._tokens_minute[agent_id]
            if ts > one_minute_ago
        ]
        
        # Підрахувати використані токени за останню хвилину
        used_tokens = sum(count for _, count in self._tokens_minute[agent_id])
        
        # Перевірити квоту
        if used_tokens + tokens > quota.tokens_per_minute:
            return False
        
        return True
    
    def record_tokens(self, agent_id: str, tokens: int) -> None:
        """
        Записати використані токени
        
        Args:
            agent_id: ID агента
            tokens: Кількість використаних токенів
        """
        now = time.time()
        self._tokens_minute[agent_id].append((now, tokens))
    
    def check_runs_quota(self, agent_id: str, quota: QuotaConfig) -> bool:
        """
        Перевірити, чи не перевищено квоту запусків за день
        
        Args:
            agent_id: ID агента
            quota: Конфігурація квот
        
        Returns:
            True, якщо квота дозволяє запуск
        """
        today = datetime.utcnow().strftime("%Y-%m-%d")
        
        # Скинути лічильники, якщо новий день
        if self._runs_today_date != today:
            self._runs_today.clear()
            self._users_today.clear()
            self._runs_today_date = today
        
        # Перевірити квоту
        if self._runs_today[agent_id] >= quota.runs_per_day:
            return False
        
        return True
    
    def record_run(self, agent_id: str, user_id: Optional[str] = None) -> None:
        """
        Записати запуск
        
        Args:
            agent_id: ID агента
            user_id: ID користувача (опційно)
        """
        today = datetime.utcnow().strftime("%Y-%m-%d")
        
        # Скинути лічильники, якщо новий день
        if self._runs_today_date != today:
            self._runs_today.clear()
            self._users_today.clear()
            self._runs_today_date = today
        
        self._runs_today[agent_id] += 1
        
        if user_id:
            self._users_today[agent_id].add(user_id)
    
    def check_users_quota(self, agent_id: str, user_id: str, quota: QuotaConfig) -> bool:
        """
        Перевірити, чи не перевищено квоту унікальних користувачів за день
        
        Args:
            agent_id: ID агента
            user_id: ID користувача
            quota: Конфігурація квот
        
        Returns:
            True, якщо квота дозволяє користувача
        """
        today = datetime.utcnow().strftime("%Y-%m-%d")
        
        # Скинути лічильники, якщо новий день
        if self._runs_today_date != today:
            self._runs_today.clear()
            self._users_today.clear()
            self._runs_today_date = today
        
        # Якщо користувач вже був — завжди дозволяємо
        if user_id in self._users_today[agent_id]:
            return True
        
        # Перевірити квоту нових користувачів
        if len(self._users_today[agent_id]) >= quota.users_per_day:
            return False
        
        return True
    
    def check_concurrent_runs(self, agent_id: str, quota: QuotaConfig) -> bool:
        """
        Перевірити, чи не перевищено квоту паралельних запусків
        
        Args:
            agent_id: ID агента
            quota: Конфігурація квот
        
        Returns:
            True, якщо квота дозволяє запуск
        """
        if self._concurrent_runs[agent_id] >= quota.max_concurrent_runs:
            return False
        
        return True
    
    def start_run(self, agent_id: str) -> None:
        """
        Почати запуск (збільшити лічильник паралельних запусків)
        
        Args:
            agent_id: ID агента
        """
        self._concurrent_runs[agent_id] += 1
    
    def finish_run(self, agent_id: str) -> None:
        """
        Завершити запуск (зменшити лічильник паралельних запусків)
        
        Args:
            agent_id: ID агента
        """
        if self._concurrent_runs[agent_id] > 0:
            self._concurrent_runs[agent_id] -= 1
    
    def get_usage_stats(self, agent_id: str) -> Dict:
        """
        Отримати статистику використання агента
        
        Args:
            agent_id: ID агента
        
        Returns:
            Dict зі статистикою
        """
        now = time.time()
        one_minute_ago = now - 60
        
        # Токени за останню хвилину
        tokens_minute = sum(
            count for ts, count in self._tokens_minute[agent_id]
            if ts > one_minute_ago
        )
        
        today = datetime.utcnow().strftime("%Y-%m-%d")
        if self._runs_today_date != today:
            runs_today = 0
            users_today = 0
        else:
            runs_today = self._runs_today[agent_id]
            users_today = len(self._users_today[agent_id])
        
        return {
            "tokens_minute": tokens_minute,
            "runs_today": runs_today,
            "users_today": users_today,
            "concurrent_runs": self._concurrent_runs[agent_id]
        }


# ============================================================================
# Global Quota Tracker Instance
# ============================================================================

_global_tracker = QuotaTracker()

def get_quota_tracker() -> QuotaTracker:
    """Отримати глобальний екземпляр QuotaTracker"""
    return _global_tracker