llm/utils.py

import os
from datetime import datetime, timedelta
from threading import Lock
from typing import Any, Dict, Optional

import boto3
from botocore.exceptions import ClientError

from llm.connection_manager import get_checkpointer

# ------------------------- Agent's tool related utils -------------------------
# User-specific storage for tool results (thread-safe)
_user_tool_results: Dict[str, Dict] = {}
_storage_lock = Lock()
# Track when results were stored for cleanup
_result_timestamps: Dict[str, datetime] = {}


def store_tool_result(user_id: str, tool_name: str, result: Dict[str, Any]) -> None:
    """
    Store tool result for a specific user.

    Args:
        user_id: Unique identifier for the user
        tool_name: Name of the tool that produced the result
        result: The result data to store
    """
    with _storage_lock:
        if user_id not in _user_tool_results:
            _user_tool_results[user_id] = {}
        _user_tool_results[user_id][tool_name] = result
        _result_timestamps[f"{user_id}:{tool_name}"] = datetime.now()
        print(f"[STORAGE] Stored {tool_name} result for user {user_id}: {result}")


def get_tool_result(user_id: str, tool_name: str) -> Optional[Dict[str, Any]]:
    """
    Get tool result for a specific user and clear it.

    Args:
        user_id: Unique identifier for the user
        tool_name: Name of the tool to get result for

    Returns:
        The tool result if found, None otherwise
    """
    with _storage_lock:
        if user_id in _user_tool_results and tool_name in _user_tool_results[user_id]:
            result = _user_tool_results[user_id].pop(tool_name)
            timestamp_key = f"{user_id}:{tool_name}"
            if timestamp_key in _result_timestamps:
                del _result_timestamps[timestamp_key]
            print(f"[STORAGE] Retrieved {tool_name} result for user {user_id}: {result}")
            return result
        return None


def clear_user_tool_results(user_id: str) -> None:
    """
    Clear all tool results for a specific user.

    Args:
        user_id: Unique identifier for the user
    """
    with _storage_lock:
        if user_id in _user_tool_results:
            # Remove all timestamps for this user
            keys_to_remove = [k for k in _result_timestamps.keys() if k.startswith(f"{user_id}:")]
            for key in keys_to_remove:
                del _result_timestamps[key]

            del _user_tool_results[user_id]
            print(f"[STORAGE] Cleared all tool results for user {user_id}")


def cleanup_old_tool_results(max_age_hours: int = 24) -> None:
    """
    Clean up tool results older than the specified age.

    Args:
        max_age_hours: Maximum age in hours before cleanup (default: 24 hours)
    """
    cutoff_time = datetime.now() - timedelta(hours=max_age_hours)

    with _storage_lock:
        keys_to_remove = []
        for timestamp_key, timestamp in _result_timestamps.items():
            if timestamp < cutoff_time:
                keys_to_remove.append(timestamp_key)

        for timestamp_key in keys_to_remove:
            user_id, tool_name = timestamp_key.split(":", 1)
            if user_id in _user_tool_results and tool_name in _user_tool_results[user_id]:
                del _user_tool_results[user_id][tool_name]
                # Clean up empty user entries
                if not _user_tool_results[user_id]:
                    del _user_tool_results[user_id]
            del _result_timestamps[timestamp_key]

        if keys_to_remove:
            print(f"[STORAGE] Cleaned up {len(keys_to_remove)} old tool results")


# ------------------------- S3 Upload of images -------------------------
def upload_generated_image_to_s3(image_data: bytes, image_id: str, user_id: str, prompt: str, title: str = "Generated Image") -> Dict[str, Any]:
    """
    Upload a generated image to S3.

    Args:
        image_data: The image data as bytes
        image_id: Unique identifier for the image
        user_id: User identifier
        prompt: The prompt used to generate the image
        title: Custom title for the image

    Returns:
        Dict with success status, URL, and metadata or error message
    """
    try:
        # Initialize S3 client
        s3_client = boto3.client(
            "s3",
            region_name=os.environ.get("AWS_REGION", "us-east-1"),
            aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
            aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
        )

        # Generate S3 key with userId and imageId for organization
        key = f"users/{user_id}/images/{image_id}"
        bucket_name = os.environ.get("AWS_S3_BUCKET_NAME")

        if not bucket_name:
            return {"success": False, "error": "AWS_S3_BUCKET_NAME environment variable is not set"}

        # Upload to S3
        s3_client.put_object(
            Bucket=bucket_name,
            Key=key,
            Body=image_data,
            ContentType="image/png",
            Metadata={
                "title": title,
                "imageId": image_id,
                "userId": user_id,
                "uploadedAt": datetime.now().isoformat(),
                "type": "generated",
                "generationPrompt": prompt,
            },
        )

        # Generate presigned URL for reading the uploaded file (valid for 2 hours)
        presigned_url = s3_client.generate_presigned_url(
            "get_object",
            Params={"Bucket": bucket_name, "Key": key},
            ExpiresIn=7200,  # 2 hours
        )

        return {"success": True, "url": presigned_url, "image_id": image_id}

    except ClientError as e:
        return {"success": False, "error": str(e)}
    except Exception as e:
        return {"success": False, "error": str(e)}


# ------------------------- IP Generation Count and Guardrails -------------------------


def get_ip_generation_count(ip_address: str) -> int:
    """
    Query the database for IP address generation count for the current week.

    Args:
        ip_address: The IP address to query

    Returns:
        generation_count
        If no data found, returns 0
    """
    try:
        checkpointer = get_checkpointer()

        # Get the start of the current week (Monday)
        now = datetime.now()
        start_of_week = now - timedelta(days=now.weekday())
        start_of_week = start_of_week.replace(hour=0, minute=0, second=0, microsecond=0)

        # Query the rate_limits table for this IP in current week
        # Using a simple SQL query to get the data
        with checkpointer.conn.cursor() as cursor:
            cursor.execute(
                """
                SELECT generation_count
                FROM rate_limits
                WHERE ip_address = %s AND week_start = %s
            """,
                (ip_address, start_of_week.date()),
            )

            row = cursor.fetchone()

        if row:
            count = row.get("generation_count")
            print(f"[UTILS] IP {ip_address}: {count} generations already made this week")
            return int(count)

        print(f"[UTILS] IP {ip_address}: No data found")
        return 0

    except Exception as e:
        print(f"[UTILS] Error querying IP generation data: {e}")
        return 0


def create_rate_limits_table():
    """
    Create the rate_limits table if it doesn't exist.
    """
    try:
        from llm.connection_manager import get_checkpointer

        checkpointer = get_checkpointer()

        with checkpointer.conn.cursor() as cursor:
            cursor.execute(
                """
                CREATE TABLE IF NOT EXISTS rate_limits (
                    id SERIAL PRIMARY KEY,
                    ip_address VARCHAR(45) NOT NULL,
                    week_start DATE NOT NULL,
                    generation_count INTEGER DEFAULT 0,
                    last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(ip_address, week_start)
                )
            """
            )
            checkpointer.conn.commit()
            print("[UTILS] Rate limits table created/verified successfully")

    except Exception as e:
        print(f"[UTILS] Error creating rate limits table: {e}")


def create_or_update_ip_generation_count(ip_address: str) -> bool:
    """
    Update the generation count for an IP address, or create a new one if it doesn't exist.

    Args:
        ip_address: The IP address to update

    Returns:
        True if successful, False otherwise
    """
    try:
        from llm.connection_manager import get_checkpointer

        checkpointer = get_checkpointer()

        # Get the start of the current week (Monday)
        now = datetime.now()
        start_of_week = now - timedelta(days=now.weekday())
        start_of_week = start_of_week.replace(hour=0, minute=0, second=0, microsecond=0)

        with checkpointer.conn.cursor() as cursor:
            # Use UPSERT to either insert new record or update existing one
            cursor.execute(
                """
                INSERT INTO rate_limits (ip_address, week_start, generation_count, last_updated)
                VALUES (%s, %s, 1, %s)
                ON CONFLICT (ip_address, week_start)
                DO UPDATE SET
                    generation_count = rate_limits.generation_count + 1,
                    last_updated = EXCLUDED.last_updated
            """,
                (ip_address, start_of_week.date(), now.isoformat()),
            )

            checkpointer.conn.commit()
            print(f"[UTILS] Created or Updated generation count for IP {ip_address}")
            return True

    except Exception as e:
        print(f"[UTILS] Error updating IP generation count: {e}")
        return False