Create app.py

app.py

@@ -0,0 +1,372 @@
+"""
+This module provides a server that's compatible with OpenAI's TTS API format.
+Main entry point for the Flask application.
+"""
+
+import os
+import json
+import logging
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, Any, Tuple, Optional
+from flask import Flask, request, jsonify, send_file, Response, send_from_directory
+from flask_cors import CORS
+from dotenv import load_dotenv
+from celery.result import AsyncResult
+from celery_worker import celery, process_tts_request
+from werkzeug.exceptions import HTTPException
+from werkzeug.utils import secure_filename
+import os.path
+import re
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+load_dotenv()
+
+# App configuration from environment variables
+HOST = os.getenv("HOST", "localhost")
+PORT = int(os.getenv("PORT", "7000"))
+VERIFY_SSL = os.getenv("VERIFY_SSL", "true").lower() != "false"
+MAX_QUEUE_SIZE = int(os.getenv("MAX_QUEUE_SIZE", "100"))
+
+# Security configuration
+ALLOWED_ORIGINS = os.getenv("ALLOWED_ORIGINS", "*").split(",")
+MAX_CONTENT_LENGTH = 16 * 1024 * 1024  # 16MB max request size
+
+# Create Flask app
+app = Flask(__name__, 
+            static_folder='static',
+            template_folder='templates')
+
+# Configure CORS with specific routes and security
+CORS(app, resources={
+    r"/v1/*": {
+        "origins": ALLOWED_ORIGINS,
+        "methods": ["POST", "OPTIONS"],
+        "allow_headers": ["Content-Type", "Authorization"],
+        "max_age": 3600
+    },
+    r"/api/*": {
+        "origins": ALLOWED_ORIGINS,
+        "methods": ["GET", "OPTIONS"],
+        "allow_headers": ["Content-Type"],
+        "max_age": 3600
+    }
+})
+
+# Set maximum content length
+app.config['MAX_CONTENT_LENGTH'] = MAX_CONTENT_LENGTH
+
+# Voice samples directory
+VOICE_SAMPLES_DIR = Path('voices')
+
+def _sanitize_input(text: str) -> str:
+    """Sanitize user input to prevent XSS and other attacks"""
+    # Remove any HTML tags
+    text = re.sub(r'<[^>]+>', '', text)
+    # Remove any script tags
+    text = re.sub(r'<script.*?</script>', '', text, flags=re.DOTALL)
+    # Remove any potentially dangerous characters
+    text = re.sub(r'[<>"\']', '', text)
+    return text.strip()
+
+# Error handlers
+@app.errorhandler(HTTPException)
+def handle_http_error(error: HTTPException) -> Tuple[Dict[str, str], int]:
+    """Handle HTTP errors"""
+    logger.warning(f"HTTP error: {error.code} - {error.description}")
+    return jsonify({"error": error.description}), error.code
+
+@app.errorhandler(Exception)
+def handle_generic_error(error: Exception) -> Tuple[Dict[str, str], int]:
+    """Handle unexpected errors"""
+    logger.error(f"Unexpected error: {str(error)}", exc_info=True)
+    return jsonify({"error": "Internal Server Error"}), 500
+
+@app.after_request
+def add_security_headers(response: Response) -> Response:
+    """Add security headers to all responses"""
+    response.headers['X-Content-Type-Options'] = 'nosniff'
+    response.headers['X-Frame-Options'] = 'DENY'
+    response.headers['X-XSS-Protection'] = '1; mode=block'
+    response.headers['Strict-Transport-Security'] = 'max-age=31536000; includeSubDomains'
+    response.headers['Content-Security-Policy'] = (
+        "default-src 'self'; "
+        "style-src 'self' https://cdnjs.cloudflare.com https://fonts.googleapis.com; "
+        "script-src 'self' https://cdnjs.cloudflare.com; "
+        "font-src 'self' https://cdnjs.cloudflare.com https://fonts.gstatic.com; "
+        "img-src 'self' data:; "
+        "media-src 'self' blob:; "
+        "connect-src 'self'"
+    )
+    return response
+
+@app.route('/')
+def index() -> Response:
+    """Render the main index page"""
+    return send_from_directory('static', 'index.html')
+
+@app.route('/static/<path:filename>')
+def serve_static(filename: str) -> Response:
+    """Serve static files with correct MIME types"""
+    if filename.endswith('.css'):
+        return send_from_directory('static', filename, mimetype='text/css')
+    elif filename.endswith('.js'):
+        return send_from_directory('static', filename, mimetype='application/javascript')
+    return send_from_directory('static', filename)
+
+def validate_tts_request(body: Dict[str, Any]) -> Tuple[Optional[Dict[str, Any]], Optional[str], Optional[int]]:
+    """Validate TTS request parameters"""
+    try:
+        # Validate required parameters
+        if 'input' not in body or 'voice' not in body:
+            return None, "Missing required parameters: input and voice", 400
+        
+        # Sanitize input
+        sanitized_input = _sanitize_input(body['input'])
+        if not sanitized_input:
+            return None, "Input text cannot be empty", 400
+        
+        # Validate voice parameter
+        if not isinstance(body['voice'], str) or not body['voice']:
+            return None, "Invalid voice parameter", 400
+        
+        openai_fm_data = {
+            'input': sanitized_input,
+            'voice': body['voice']
+        }
+        
+        # Validate and sanitize instructions if provided
+        if 'instructions' in body:
+            if not isinstance(body['instructions'], str):
+                return None, "Instructions must be a string", 400
+            openai_fm_data['prompt'] = _sanitize_input(body['instructions'])
+        
+        # Validate response format
+        content_type = "audio/mpeg"
+        if 'response_format' in body:
+            format_mapping = {
+                'mp3': 'audio/mpeg',
+                'opus': 'audio/opus',
+                'aac': 'audio/aac',
+                'flac': 'audio/flac',
+                'wav': 'audio/wav',
+                'pcm': 'audio/pcm'
+            }
+            requested_format = body['response_format'].lower()
+            if requested_format not in format_mapping:
+                return None, f"Unsupported response format: {requested_format}. Supported formats are: {', '.join(format_mapping.keys())}", 400
+            content_type = format_mapping[requested_format]
+            openai_fm_data['format'] = requested_format
+        
+        return openai_fm_data, None, None
+    except Exception as e:
+        logger.error(f"Error validating request: {str(e)}")
+        return None, "Invalid request format", 400
+
+def get_queue_details() -> Dict[str, Any]:
+    """Get detailed queue counts from Celery workers."""
+    details = {
+        'active': 0,
+        'reserved': 0,
+        'scheduled': 0,
+        'total_reported_by_workers': 0,
+        'error': None
+    }
+    try:
+        i = celery.control.inspect(timeout=1.0) # Add timeout
+        if not i:
+             details['error'] = "Could not connect to Celery workers for inspection."
+             return details
+             
+        active_tasks = i.active()
+        reserved_tasks = i.reserved()
+        scheduled_tasks = i.scheduled()
+        
+        if active_tasks:
+            details['active'] = sum(len(tasks) for tasks in active_tasks.values())
+        if reserved_tasks:
+            details['reserved'] = sum(len(tasks) for tasks in reserved_tasks.values())
+        if scheduled_tasks:
+            details['scheduled'] = sum(len(tasks) for tasks in scheduled_tasks.values())
+            
+        details['total_reported_by_workers'] = details['active'] + details['reserved'] + details['scheduled']
+            
+    except Exception as e:
+        logger.error(f"Error calculating queue details: {str(e)}")
+        details['error'] = f"Failed to inspect Celery workers: {str(e)}"
+        # Reset counts to 0 on error to avoid misleading data
+        details['active'] = 0
+        details['reserved'] = 0
+        details['scheduled'] = 0
+        details['total_reported_by_workers'] = 0
+        
+    return details
+
+@app.route('/v1/audio/speech', methods=['POST'])
+def openai_speech() -> Response:
+    """Handle POST requests to /v1/audio/speech (OpenAI compatible API)"""
+    try:
+        # Check queue size from Celery worker reports
+        queue_details = get_queue_details()
+        current_total = queue_details['total_reported_by_workers']
+        
+        # Check for inspection errors
+        if queue_details['error']:
+             logger.warning(f"Could not determine queue size due to inspection error: {queue_details['error']}. Allowing request.")
+             # Optionally, you could reject here, but allowing might be safer if inspection is flaky
+        
+        elif current_total >= MAX_QUEUE_SIZE:
+            logger.warning(f"Queue is full based on worker reports. Current total: {current_total}, Max size: {MAX_QUEUE_SIZE}")
+            return jsonify({
+                "error": "Queue is full. Please try again later.",
+                "queue_details": queue_details, # Provide detailed counts
+                "max_queue_size_limit": MAX_QUEUE_SIZE
+            }), 429
+
+        # Read and validate JSON data
+        body = request.get_json()
+        openai_fm_data, error, status_code = validate_tts_request(body)
+        if error:
+            return jsonify({"error": error}), status_code
+        
+        # Determine content type from validation or default
+        validated_content_type = "audio/mpeg" # Default
+        if 'response_format' in body:
+            format_mapping = {
+                'mp3': 'audio/mpeg',
+                'opus': 'audio/opus',
+                'aac': 'audio/aac',
+                'flac': 'audio/flac',
+                'wav': 'audio/wav',
+                'pcm': 'audio/pcm'
+            }
+            requested_format = body['response_format'].lower()
+            if requested_format in format_mapping:
+                validated_content_type = format_mapping[requested_format]
+
+        # Create task data
+        task_data = {
+            'data': openai_fm_data,
+            'timestamp': datetime.now().isoformat()
+        }
+        
+        # Submit task to Celery
+        task = process_tts_request.delay(task_data)
+        
+        # Wait for result with timeout
+        try:
+            audio_data, error, status_code = task.get(timeout=30)
+            if error:
+                logger.error(f"Task error: {error}")
+                return jsonify({"error": error}), status_code
+            
+            return Response(
+                audio_data,
+                mimetype=validated_content_type # Use the correctly determined content type
+            )
+        except TimeoutError:
+            logger.error(f"Task timeout: {task.id}")
+            return jsonify({
+                "error": "Request timed out. Please try again.",
+                "task_id": task.id
+            }), 408
+                
+    except json.JSONDecodeError:
+        logger.error("Invalid JSON in request body")
+        return jsonify({"error": "Invalid JSON in request body"}), 400
+    except Exception as e:
+        logger.error(f"Unexpected error in speech endpoint: {str(e)}")
+        return jsonify({"error": "Internal Server Error"}), 500
+
+@app.route('/api/queue-size', methods=['GET'])
+def queue_size() -> Response:
+    """Handle GET requests to /api/queue-size with detailed counts"""
+    try:
+        queue_details = get_queue_details()
+        
+        response_data = {
+            "active_tasks": queue_details['active'],
+            "reserved_tasks": queue_details['reserved'],
+            "scheduled_tasks": queue_details['scheduled'],
+            "total_reported_by_workers": queue_details['total_reported_by_workers'],
+            "max_queue_size_limit": MAX_QUEUE_SIZE,
+            "error": queue_details['error']
+        }
+        
+        # Determine status code based on whether there was an inspection error
+        status_code = 500 if queue_details['error'] else 200
+        
+        return jsonify(response_data), status_code
+        
+    except Exception as e:
+        # This handles errors in the route handler itself, not inspection errors
+        logger.error(f"Error in /api/queue-size endpoint: {str(e)}")
+        return jsonify({
+            "active_tasks": 0,
+            "reserved_tasks": 0,
+            "scheduled_tasks": 0,
+            "total_reported_by_workers": 0,
+            "max_queue_size_limit": MAX_QUEUE_SIZE,
+            "error": "Failed to process queue status request"
+        }), 500
+
+@app.route('/api/voice-sample/<voice>', methods=['GET'])
+def voice_sample(voice: str) -> Response:
+    """Handle GET requests for voice samples"""
+    try:
+        if not voice:
+            return jsonify({
+                "error": "Voice parameter is required"
+            }), 400
+            
+        # Secure the voice parameter and prevent path traversal
+        secure_voice = secure_filename(voice)
+        if not secure_voice or secure_voice != voice:
+            logger.warning(f"Invalid voice parameter: {voice}")
+            return jsonify({
+                "error": "Invalid voice parameter"
+            }), 400
+            
+        # Normalize and validate the path
+        base_path = os.path.abspath(VOICE_SAMPLES_DIR)
+        sample_path = os.path.normpath(os.path.join(base_path, f"{secure_voice}_sample.mp3"))
+        
+        # Ensure the path is within the voice samples directory
+        if not sample_path.startswith(base_path):
+            logger.warning(f"Path traversal attempt: {sample_path}")
+            return jsonify({
+                "error": "Invalid path"
+            }), 400
+            
+        if not os.path.exists(sample_path):
+            logger.warning(f"Sample not found for voice: {voice}")
+            return jsonify({
+                "error": f"Sample not found for voice: {voice}"
+            }), 404
+            
+        return send_file(
+            sample_path,
+            mimetype="audio/mpeg",
+            as_attachment=False,
+            download_name=f"{secure_voice}_sample.mp3"
+        )
+        
+    except Exception as e:
+        logger.error(f"Error serving voice sample: {str(e)}")
+        return jsonify({
+            "error": "Internal Server Error"
+        }), 500
+
+@app.route('/api/version', methods=['GET'])
+def get_version() -> Response:
+    """Handle GET requests for API version"""
+    return jsonify({
+        "version": "v2.0.0-alpha_x"
+    })