lightweight_embeddings/analytics.py

import logging
import asyncio
import redis.asyncio as redis
import redis.exceptions
from datetime import datetime
from collections import defaultdict
from typing import Dict

logger = logging.getLogger(__name__)

class Analytics:
    def __init__(self, redis_url: str, sync_interval: int = 60, max_retries: int = 5):
        """
        Initializes the Analytics class with an async Redis connection and sync interval.

        Parameters:
        - redis_url: Redis connection URL (e.g., 'redis://localhost:6379/0')
        - sync_interval: Interval in seconds for syncing with Redis.
        - max_retries: Maximum number of retries for reconnecting to Redis.
        """
        self.redis_url = redis_url
        self.sync_interval = sync_interval
        self.max_retries = max_retries
        self.redis_client = self._create_redis_client()
        self.local_buffer = {
            "access": defaultdict(
                lambda: defaultdict(int)
            ),  # {period: {model_id: access_count}}
            "tokens": defaultdict(
                lambda: defaultdict(int)
            ),  # {period: {model_id: tokens_count}}
        }
        self.lock = asyncio.Lock()  # Async lock for thread-safe updates
        asyncio.create_task(self._start_sync_task())

        logger.info("Initialized Analytics with Redis connection: %s", redis_url)

    def _create_redis_client(self) -> redis.Redis:
        """
        Creates and returns a new Redis client.
        """
        return redis.from_url(
            self.redis_url,
            decode_responses=True,
            health_check_interval=10,
            socket_connect_timeout=5,
            retry_on_timeout=True,
            socket_keepalive=True,
        )

    def _get_period_keys(self) -> tuple:
        """
        Returns keys for day, week, month, and year based on the current date.
        """
        now = datetime.utcnow()
        day_key = now.strftime("%Y-%m-%d")
        week_key = f"{now.year}-W{now.strftime('%U')}"
        month_key = now.strftime("%Y-%m")
        year_key = now.strftime("%Y")
        return day_key, week_key, month_key, year_key

    async def access(self, model_id: str, tokens: int):
        """
        Records an access and token usage for a specific model_id.

        Parameters:
        - model_id: The ID of the model being accessed.
        - tokens: Number of tokens used in this access.
        """
        day_key, week_key, month_key, year_key = self._get_period_keys()

        async with self.lock:
            # Increment access count
            self.local_buffer["access"][day_key][model_id] += 1
            self.local_buffer["access"][week_key][model_id] += 1
            self.local_buffer["access"][month_key][model_id] += 1
            self.local_buffer["access"][year_key][model_id] += 1
            self.local_buffer["access"]["total"][model_id] += 1

            # Increment token count
            self.local_buffer["tokens"][day_key][model_id] += tokens
            self.local_buffer["tokens"][week_key][model_id] += tokens
            self.local_buffer["tokens"][month_key][model_id] += tokens
            self.local_buffer["tokens"][year_key][model_id] += tokens
            self.local_buffer["tokens"]["total"][model_id] += tokens

    async def stats(self) -> Dict[str, Dict[str, Dict[str, int]]]:
        """
        Returns statistics for all models from the local buffer.

        Returns:
        - A dictionary with access counts and token usage for each period.
        """
        async with self.lock:
            return {
                "access": {
                    period: dict(models)
                    for period, models in self.local_buffer["access"].items()
                },
                "tokens": {
                    period: dict(models)
                    for period, models in self.local_buffer["tokens"].items()
                },
            }

    async def _sync_to_redis(self):
        """
        Synchronizes local buffer data with Redis.
        """
        async with self.lock:
            try:
                pipeline = self.redis_client.pipeline()

                # Sync access counts
                for period, models in self.local_buffer["access"].items():
                    for model_id, count in models.items():
                        redis_key = f"analytics:access:{period}"
                        pipeline.hincrby(redis_key, model_id, count)

                # Sync token counts
                for period, models in self.local_buffer["tokens"].items():
                    for model_id, count in models.items():
                        redis_key = f"analytics:tokens:{period}"
                        pipeline.hincrby(redis_key, model_id, count)

                pipeline.execute()
                self.local_buffer["access"].clear()  # Clear access buffer after sync
                self.local_buffer["tokens"].clear()  # Clear tokens buffer after sync
                logger.info("Synced analytics data to Redis.")

            except redis.exceptions.ConnectionError as e:
                logger.error("Redis connection error during sync: %s", e)
                raise e
            except Exception as e:
                logger.error("Unexpected error during Redis sync: %s", e)
                raise e

    async def _start_sync_task(self):
        """
        Starts a background task that periodically syncs data to Redis.
        Implements retry logic with exponential backoff on connection failures.
        """
        retry_delay = 1  # Initial retry delay in seconds

        while True:
            await asyncio.sleep(self.sync_interval)
            try:
                await self._sync_to_redis()
                retry_delay = 1  # Reset retry delay after successful sync
            except redis.exceptions.ConnectionError as e:
                logger.error("Redis connection error: %s", e)
                await self._handle_redis_reconnection()
            except Exception as e:
                logger.error("Error during sync: %s", e)
                # Depending on the error, you might want to handle differently

    async def _handle_redis_reconnection(self):
        """
        Handles Redis reconnection with exponential backoff.
        """
        retry_count = 0
        delay = 1  # Start with 1 second delay

        while retry_count < self.max_retries:
            try:
                logger.info("Attempting to reconnect to Redis (Attempt %d)...", retry_count + 1)
                self.redis_client.close()
                self.redis_client = self._create_redis_client()
                # Optionally, perform a simple command to check connection
                self.redis_client.ping()
                logger.info("Successfully reconnected to Redis.")
                return
            except redis.exceptions.ConnectionError as e:
                logger.error("Reconnection attempt %d failed: %s", retry_count + 1, e)
                retry_count += 1
                await asyncio.sleep(delay)
                delay *= 2  # Exponential backoff

        logger.critical("Max reconnection attempts reached. Unable to reconnect to Redis.")
        # Depending on your application's requirements, you might choose to exit or keep retrying indefinitely
        # For example, to keep retrying:
        while True:
            try:
                logger.info("Retrying to reconnect to Redis...")
                self.redis_client.close()
                self.redis_client = self._create_redis_client()
                self.redis_client.ping()
                logger.info("Successfully reconnected to Redis.")
                break
            except redis.exceptions.ConnectionError as e:
                logger.error("Reconnection attempt failed: %s", e)
                await asyncio.sleep(delay)
                delay = min(delay * 2, 60)  # Cap the delay to 60 seconds

    async def close(self):
        """
        Closes the Redis connection gracefully.
        """
        self.redis_client.close()
        logger.info("Closed Redis connection.")