src/utils.py

"""
Utility functions for the Gemini AI Agent.
"""

import re
import logging
from typing import Optional


logger = logging.getLogger(__name__)


def sanitize_input(text: str) -> str:
    """
    Sanitize user input by removing potentially harmful content.

    Args:
        text: Raw input text

    Returns:
        str: Sanitized text
    """
    if not isinstance(text, str):
        return str(text)

    # Remove excessive whitespace
    text = re.sub(r'\s+', ' ', text.strip())

    # Remove potential injection attempts (basic protection)
    text = text.replace('\\n', '\n').replace('\\t', '\t')

    return text


def format_response(response: str) -> str:
    """
    Format the AI response for better readability.

    Args:
        response: Raw response from the AI model

    Returns:
        str: Formatted response
    """
    if not response:
        return "No response generated."

    # Clean up the response
    response = response.strip()

    # Ensure proper spacing after periods
    response = re.sub(r'\.([A-Z])', r'. \1', response)

    # Fix common formatting issues
    response = re.sub(r'\n\s*\n\s*\n', '\n\n', response)  # Multiple newlines
    response = re.sub(r'([.!?])\s*([A-Z])', r'\1 \2', response)  # Spacing after punctuation

    return response


def truncate_text(text: str, max_length: int = 1000) -> str:
    """
    Truncate text to a maximum length while preserving word boundaries.

    Args:
        text: Text to truncate
        max_length: Maximum length allowed

    Returns:
        str: Truncated text
    """
    if len(text) <= max_length:
        return text

    # Find the last space before the max_length
    truncated = text[:max_length]
    last_space = truncated.rfind(' ')

    if last_space > max_length * 0.8:  # If space is reasonably close to end
        return truncated[:last_space] + "..."
    else:
        return truncated + "..."


def extract_code_blocks(text: str) -> list:
    """
    Extract code blocks from markdown-formatted text.

    Args:
        text: Text containing potential code blocks

    Returns:
        list: List of code blocks found
    """
    code_pattern = r'```(\w+)?\n(.*?)\n```'
    matches = re.findall(code_pattern, text, re.DOTALL)

    return [{'language': match[0] or 'text', 'code': match[1]} for match in matches]


def validate_question(question: str) -> tuple[bool, Optional[str]]:
    """
    Validate if a question is appropriate and well-formed.

    Args:
        question: The question to validate

    Returns:
        tuple: (is_valid, error_message)
    """
    if not question or not question.strip():
        return False, "Question cannot be empty."

    if len(question.strip()) < 3:
        return False, "Question is too short. Please provide more detail."

    if len(question) > 5000:
        return False, "Question is too long. Please keep it under 5000 characters."

    return True, None


def format_error_message(error: Exception) -> str:
    """
    Format error messages for user display.

    Args:
        error: The exception that occurred

    Returns:
        str: User-friendly error message
    """
    error_type = type(error).__name__

    # Map common errors to user-friendly messages
    error_messages = {
        'ConnectionError': 'Unable to connect to the AI service. Please check your internet connection.',
        'TimeoutError': 'The request timed out. Please try again with a shorter question.',
        'ValueError': 'Invalid input provided. Please check your question format.',
        'KeyError': 'Configuration error. Please check your API settings.',
        'PermissionError': 'Access denied. Please check your API key permissions.'
    }

    user_message = error_messages.get(error_type, f"An unexpected error occurred: {str(error)}")

    logger.error(f"Error formatted for user: {error_type} - {str(error)}")
    return user_message