Update app.py

app.py

@@ -1,6 +1,8 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
+from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks, Query
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse
 import whisper
 import torch
 import tempfile
@@ -10,49 +12,84 @@ import logging
 import hashlib
 import json
 import sqlite3
-from datetime import datetime
+from datetime import datetime, timedelta
 import threading
 import time
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from typing import Optional, Dict, Any
+from pathlib import Path
+import aiofiles
+from contextlib import asynccontextmanager
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-app = FastAPI()
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("app.log"),
+        logging.StreamHandler()
+    ]
 )
+logger = logging.getLogger(__name__)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-logger.info(f"Loading Whisper model on {device}")
-whisper_model = whisper.load_model("large-v3", device=device)
-logger.info("Whisper model loaded successfully")
-
-logger.info("Loading translation model...")
+# Global variables
+whisper_model = None
 translation_tokenizer = None
 translation_model = None
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Lifespan manager for startup and shutdown events"""
+    # Startup
+    global whisper_model, translation_tokenizer, translation_model
+    
+    try:
+        logger.info(f"Loading Whisper model on {device}")
+        whisper_model = whisper.load_model("large-v3", device=device)
+        logger.info("Whisper model loaded successfully")
+    except Exception as e:
+        logger.error(f"Failed to load Whisper model: {e}")
+        raise RuntimeError(f"Whisper model loading failed: {e}")
 
-def load_translation_model():
-    global translation_tokenizer, translation_model
     try:
+        logger.info("Loading translation model...")
         model_name = "facebook/nllb-200-distilled-600M"
         translation_tokenizer = AutoTokenizer.from_pretrained(model_name)
         translation_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-        if device == "cuda" and torch.cuda.is_available():
+        if device == "cuda":
             translation_model = translation_model.to(device)
         logger.info(f"Translation model loaded on {device}")
-        return True
     except Exception as e:
-        logger.error(f"Failed to load translation model: {e}")
-        return False
+        logger.warning(f"Translation model not available: {e}")
+        translation_model = None
 
-translation_loaded = load_translation_model()
+    # Initialize database
+    init_cache_db()
+    
+    yield
+    
+    # Shutdown
+    logger.info("Shutting down application...")
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+app = FastAPI(
+    title="Whisper Transcription API",
+    description="API for audio transcription and translation using OpenAI Whisper",
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+# CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
+# Language codes mapping
 LANGUAGE_CODES = {
     'persian': 'pes_Arab',
     'farsi': 'pes_Arab', 
@@ -76,160 +113,197 @@ LANGUAGE_CODES = {
     'finnish': 'fin_Latn'
 }
 
-def translate_text(text, target_language):
-    if not translation_loaded or not translation_model or not translation_tokenizer:
-        return None
+class DatabaseManager:
+    """Database management class with connection pooling"""
     
-    try:
-        target_code = LANGUAGE_CODES.get(target_language.lower())
-        if not target_code:
-            return None
+    def __init__(self, db_path: str = 'transcription_cache.db'):
+        self.db_path = db_path
+        self._init_db()
+    
+    def _init_db(self):
+        """Initialize database tables"""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
             
-        inputs = translation_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
-        if device == "cuda" and torch.cuda.is_available():
-            inputs = {k: v.to(device) for k, v in inputs.items()}
-        
-        translated_tokens = translation_model.generate(
-            **inputs, 
-            forced_bos_token_id=translation_tokenizer.lang_code_to_id[target_code],
-            max_length=512,
-            num_beams=5,
-            early_stopping=True
-        )
-        
-        translated_text = translation_tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
-        return translated_text.strip()
-        
-    except Exception as e:
-        logger.error(f"Translation error: {e}")
-        return None
+            # Cache table
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS cache (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    file_hash TEXT UNIQUE,
+                    filename TEXT,
+                    file_size INTEGER,
+                    transcription TEXT,
+                    language TEXT,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    last_accessed TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            ''')
+            
+            # Processing status table
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS processing_status (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    file_hash TEXT UNIQUE,
+                    filename TEXT,
+                    file_size INTEGER,
+                    status TEXT DEFAULT 'processing',
+                    progress INTEGER DEFAULT 0,
+                    estimated_time INTEGER DEFAULT 0,
+                    started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            ''')
+            
+            # Translation cache table
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS translation_cache (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    text_hash TEXT,
+                    target_language TEXT,
+                    translated_text TEXT,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    last_accessed TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    UNIQUE(text_hash, target_language)
+                )
+            ''')
+            
+            # Create indexes for better performance
+            cursor.execute('CREATE INDEX IF NOT EXISTS idx_cache_hash ON cache(file_hash)')
+            cursor.execute('CREATE INDEX IF NOT EXISTS idx_cache_created ON cache(created_at)')
+            cursor.execute('CREATE INDEX IF NOT EXISTS idx_status_hash ON processing_status(file_hash)')
+            cursor.execute('CREATE INDEX IF NOT EXISTS idx_translation_hash ON translation_cache(text_hash, target_language)')
+            
+            conn.commit()
+    
+    def get_connection(self):
+        """Get database connection"""
+        return sqlite3.connect(self.db_path)
+
+# Global database manager
+db_manager = DatabaseManager()
 
 def init_cache_db():
-    conn = sqlite3.connect('transcription_cache.db')
-    cursor = conn.cursor()
-    cursor.execute('''
-        CREATE TABLE IF NOT EXISTS cache (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            file_hash TEXT UNIQUE,
-            filename TEXT,
-            file_size INTEGER,
-            transcription TEXT,
-            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-        )
-    ''')
-    
-    cursor.execute('''
-        CREATE TABLE IF NOT EXISTS processing_status (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            file_hash TEXT UNIQUE,
-            filename TEXT,
-            file_size INTEGER,
-            status TEXT DEFAULT 'processing',
-            progress INTEGER DEFAULT 0,
-            estimated_time INTEGER DEFAULT 0,
-            started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-        )
-    ''')
-    
-    cursor.execute('''
-        CREATE TABLE IF NOT EXISTS translation_cache (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            text_hash TEXT,
-            target_language TEXT,
-            translated_text TEXT,
-            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-            UNIQUE(text_hash, target_language)
-        )
-    ''')
-    
-    conn.commit()
-    conn.close()
+    """Initialize cache database"""
+    # Already handled by DatabaseManager
+    pass
 
-def calculate_file_hash(content, filename, file_size):
+def calculate_file_hash(content: bytes, filename: str, file_size: int) -> str:
+    """Calculate hash for file identification"""
     hash_input = f"{filename}_{file_size}_{len(content)}"
-    file_hash = hashlib.md5(content[:1024] + content[-1024:] + hash_input.encode()).hexdigest()
-    return file_hash
+    return hashlib.md5(content[:1024] + content[-1024:] + hash_input.encode()).hexdigest()
 
-def calculate_text_hash(text):
+def calculate_text_hash(text: str) -> str:
+    """Calculate hash for text"""
     return hashlib.md5(text.encode('utf-8')).hexdigest()
 
-def get_from_cache(file_hash):
+async def get_from_cache(file_hash: str) -> Optional[str]:
+    """Get transcription from cache"""
     try:
-        conn = sqlite3.connect('transcription_cache.db')
-        cursor = conn.cursor()
-        cursor.execute('SELECT transcription FROM cache WHERE file_hash = ?', (file_hash,))
-        result = cursor.fetchone()
-        conn.close()
-        return result[0] if result else None
-    except:
+        with db_manager.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                'SELECT transcription FROM cache WHERE file_hash = ?',
+                (file_hash,)
+            )
+            result = cursor.fetchone()
+            
+            if result:
+                # Update last accessed time
+                cursor.execute(
+                    'UPDATE cache SET last_accessed = CURRENT_TIMESTAMP WHERE file_hash = ?',
+                    (file_hash,)
+                )
+                conn.commit()
+                return result[0]
+        return None
+    except Exception as e:
+        logger.error(f"Error getting from cache: {e}")
         return None
 
-def get_translation_from_cache(text_hash, target_language):
+async def get_translation_from_cache(text_hash: str, target_language: str) -> Optional[str]:
+    """Get translation from cache"""
     try:
-        conn = sqlite3.connect('transcription_cache.db')
-        cursor = conn.cursor()
-        cursor.execute('SELECT translated_text FROM translation_cache WHERE text_hash = ? AND target_language = ?', 
-                      (text_hash, target_language))
-        result = cursor.fetchone()
-        conn.close()
-        return result[0] if result else None
-    except:
+        with db_manager.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                '''SELECT translated_text FROM translation_cache 
+                   WHERE text_hash = ? AND target_language = ?''',
+                (text_hash, target_language)
+            )
+            result = cursor.fetchone()
+            
+            if result:
+                # Update last accessed time
+                cursor.execute(
+                    '''UPDATE translation_cache SET last_accessed = CURRENT_TIMESTAMP 
+                       WHERE text_hash = ? AND target_language = ?''',
+                    (text_hash, target_language)
+                )
+                conn.commit()
+                return result[0]
+        return None
+    except Exception as e:
+        logger.error(f"Error getting translation from cache: {e}")
         return None
 
-def save_to_cache(file_hash, filename, file_size, transcription):
+async def save_to_cache(file_hash: str, filename: str, file_size: int, transcription: str, language: str = None):
+    """Save transcription to cache"""
     try:
-        conn = sqlite3.connect('transcription_cache.db')
-        cursor = conn.cursor()
-        cursor.execute('''
-            INSERT OR REPLACE INTO cache (file_hash, filename, file_size, transcription)
-            VALUES (?, ?, ?, ?)
-        ''', (file_hash, filename, file_size, transcription))
-        conn.commit()
-        conn.close()
+        with db_manager.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                '''INSERT OR REPLACE INTO cache 
+                   (file_hash, filename, file_size, transcription, language) 
+                   VALUES (?, ?, ?, ?, ?)''',
+                (file_hash, filename, file_size, transcription, language)
+            )
+            conn.commit()
     except Exception as e:
         logger.error(f"Error saving to cache: {e}")
 
-def save_translation_to_cache(text_hash, target_language, translated_text):
+async def save_translation_to_cache(text_hash: str, target_language: str, translated_text: str):
+    """Save translation to cache"""
     try:
-        conn = sqlite3.connect('transcription_cache.db')
-        cursor = conn.cursor()
-        cursor.execute('''
-            INSERT OR REPLACE INTO translation_cache (text_hash, target_language, translated_text)
-            VALUES (?, ?, ?)
-        ''', (text_hash, target_language, translated_text))
-        conn.commit()
-        conn.close()
+        with db_manager.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                '''INSERT OR REPLACE INTO translation_cache 
+                   (text_hash, target_language, translated_text) 
+                   VALUES (?, ?, ?)''',
+                (text_hash, target_language, translated_text)
+            )
+            conn.commit()
     except Exception as e:
         logger.error(f"Error saving translation to cache: {e}")
 
-def get_processing_status(file_hash):
+async def get_processing_status(file_hash: str) -> Optional[Dict[str, Any]]:
+    """Get processing status for a file"""
     try:
-        conn = sqlite3.connect('transcription_cache.db')
-        cursor = conn.cursor()
-        cursor.execute('''
-            SELECT status, progress, estimated_time, 
+        with db_manager.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                '''SELECT status, progress, estimated_time, 
                    (julianday('now') - julianday(started_at)) * 24 * 60 as elapsed_minutes
-            FROM processing_status WHERE file_hash = ?
-        ''', (file_hash,))
-        result = cursor.fetchone()
-        conn.close()
-        if result:
-            return {
-                'status': result[0],
-                'progress': result[1], 
-                'estimated_time': result[2],
-                'elapsed_minutes': int(result[3])
-            }
+                   FROM processing_status WHERE file_hash = ?''',
+                (file_hash,)
+            )
+            result = cursor.fetchone()
+            
+            if result:
+                return {
+                    'status': result[0],
+                    'progress': result[1], 
+                    'estimated_time': result[2],
+                    'elapsed_minutes': int(result[3] or 0)
+                }
         return None
-    except:
+    except Exception as e:
+        logger.error(f"Error getting processing status: {e}")
         return None
 
-def update_processing_status(file_hash, status=None, progress=None, estimated_time=None):
+async def update_processing_status(file_hash: str, status: str = None, progress: int = None, estimated_time: int = None):
+    """Update processing status"""
     try:
-        conn = sqlite3.connect('transcription_cache.db')
-        cursor = conn.cursor()
-        
         updates = []
         params = []
         
@@ -247,70 +321,82 @@ def update_processing_status(file_hash, status=None, progress=None, estimated_ti
         params.append(file_hash)
         
         query = f"UPDATE processing_status SET {', '.join(updates)} WHERE file_hash = ?"
-        cursor.execute(query, params)
-        conn.commit()
-        conn.close()
+        
+        with db_manager.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(query, params)
+            conn.commit()
     except Exception as e:
         logger.error(f"Error updating status: {e}")
 
-def add_processing_status(file_hash, filename, file_size, estimated_time):
+async def add_processing_status(file_hash: str, filename: str, file_size: int, estimated_time: int):
+    """Add new processing status entry"""
     try:
-        conn = sqlite3.connect('transcription_cache.db')
-        cursor = conn.cursor()
-        cursor.execute('''
-            INSERT OR REPLACE INTO processing_status 
-            (file_hash, filename, file_size, status, progress, estimated_time)
-            VALUES (?, ?, ?, 'processing', 0, ?)
-        ''', (file_hash, filename, file_size, estimated_time))
-        conn.commit()
-        conn.close()
+        with db_manager.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                '''INSERT OR REPLACE INTO processing_status 
+                   (file_hash, filename, file_size, status, progress, estimated_time)
+                   VALUES (?, ?, ?, 'processing', 0, ?)''',
+                (file_hash, filename, file_size, estimated_time)
+            )
+            conn.commit()
     except Exception as e:
         logger.error(f"Error adding processing status: {e}")
 
-def remove_processing_status(file_hash):
+async def remove_processing_status(file_hash: str):
+    """Remove processing status entry"""
     try:
-        conn = sqlite3.connect('transcription_cache.db')
-        cursor = conn.cursor()
-        cursor.execute('DELETE FROM processing_status WHERE file_hash = ?', (file_hash,))
-        conn.commit()
-        conn.close()
+        with db_manager.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                'DELETE FROM processing_status WHERE file_hash = ?',
+                (file_hash,)
+            )
+            conn.commit()
     except Exception as e:
         logger.error(f"Error removing processing status: {e}")
 
-def estimate_processing_time(file_size_mb):
-    estimated_seconds = file_size_mb * 30
-    return int(estimated_seconds / 60) + 1
+def estimate_processing_time(file_size_mb: float) -> int:
+    """Estimate processing time in minutes"""
+    estimated_seconds = file_size_mb * 30  # 30 seconds per MB
+    return max(1, int(estimated_seconds / 60))
 
-def background_transcription(file_path, file_hash, filename, file_size, translate_to_english=False):
+async def background_transcription(file_path: str, file_hash: str, filename: str, file_size: int, translate_to_english: bool = False):
+    """Background task for transcription"""
     try:
         logger.info(f"Starting background transcription for {filename}")
         
-        update_processing_status(file_hash, status='processing', progress=10)
+        await update_processing_status(file_hash, status='processing', progress=10)
         
+        # Transcribe audio
         result = whisper_model.transcribe(
             file_path, 
-            fp16=False if device == "cpu" else True,
+            fp16=(device != "cpu"),
             language=None,
             task="transcribe",
             verbose=False,
             word_timestamps=False
         )
         
-        update_processing_status(file_hash, progress=60)
+        await update_processing_status(file_hash, progress=60)
         
-        text = result["text"].strip()
-        if not text:
-            text = "No text detected"
-            
-        response_data = {"text": text, "from_cache": False}
+        text = result["text"].strip() or "No text detected"
+        detected_language = result.get("language", "unknown")
         
-        if translate_to_english and result.get("language") != "en":
-            update_processing_status(file_hash, progress=80)
-            logger.info("Background: Translating to English...")
+        response_data = {
+            "text": text,
+            "language": detected_language,
+            "from_cache": False
+        }
+        
+        # Translate if requested and needed
+        if translate_to_english and detected_language != "en":
+            await update_processing_status(file_hash, progress=80)
             
             english_result = whisper_model.transcribe(
                 file_path, 
-                fp16=False if device == "cpu" else True,
+                fp16=(device != "cpu"),
                 language=None,
                 task="translate",
                 verbose=False,
@@ -320,102 +406,132 @@ def background_transcription(file_path, file_hash, filename, file_size, translat
             english_text = english_result["text"].strip()
             if english_text:
                 response_data["english_text"] = english_text
-                logger.info("Background: English translation completed")
         
-        save_to_cache(file_hash, filename, file_size, json.dumps(response_data))
-        
-        update_processing_status(file_hash, status='completed', progress=100)
+        # Save to cache
+        await save_to_cache(
+            file_hash, filename, file_size, 
+            json.dumps(response_data), detected_language
+        )
         
+        await update_processing_status(file_hash, status='completed', progress=100)
         logger.info(f"Background transcription completed for {filename}")
         
     except Exception as e:
         logger.error(f"Error in background transcription: {e}")
-        update_processing_status(file_hash, status='error', progress=0)
+        await update_processing_status(file_hash, status='error', progress=0)
     
     finally:
-        if os.path.exists(file_path):
-            try:
+        # Clean up temporary file
+        try:
+            if os.path.exists(file_path):
                 os.unlink(file_path)
-            except:
-                pass
+        except Exception as e:
+            logger.error(f"Error deleting temp file: {e}")
 
-def cleanup_old_cache():
+async def cleanup_old_cache(days: int = 30):
+    """Clean up old cache entries"""
     try:
-        conn = sqlite3.connect('transcription_cache.db')
-        cursor = conn.cursor()
-        cursor.execute("DELETE FROM cache WHERE created_at < datetime('now', '-30 days')")
-        cursor.execute("DELETE FROM processing_status WHERE started_at < datetime('now', '-1 days')")
-        cursor.execute("DELETE FROM translation_cache WHERE created_at < datetime('now', '-7 days')")
-        conn.commit()
-        conn.close()
+        cutoff_date = (datetime.now() - timedelta(days=days)).strftime('%Y-%m-%d %H:%M:%S')
+        
+        with db_manager.get_connection() as conn:
+            cursor = conn.cursor()
+            
+            # Clean cache
+            cursor.execute(
+                "DELETE FROM cache WHERE last_accessed < ?",
+                (cutoff_date,)
+            )
+            
+            # Clean processing status (keep only 1 day)
+            cursor.execute(
+                "DELETE FROM processing_status WHERE started_at < datetime('now', '-1 day')"
+            )
+            
+            # Clean translation cache
+            cursor.execute(
+                "DELETE FROM translation_cache WHERE last_accessed < ?",
+                (cutoff_date,)
+            )
+            
+            conn.commit()
+            logger.info(f"Cache cleanup completed. Removed entries older than {days} days.")
+            
     except Exception as e:
         logger.error(f"Error cleaning cache: {e}")
 
-init_cache_db()
-
 @app.get("/")
 async def root():
-    conn = sqlite3.connect('transcription_cache.db')
-    cursor = conn.cursor()
-    cursor.execute('SELECT COUNT(*) FROM cache')
-    cache_count = cursor.fetchone()[0]
-    cursor.execute('SELECT COUNT(*) FROM processing_status WHERE status = "processing"')
-    processing_count = cursor.fetchone()[0]
-    cursor.execute('SELECT COUNT(*) FROM translation_cache')
-    translation_cache_count = cursor.fetchone()[0]
-    conn.close()
-    
-    return {
-        "message": "Whisper API with Translation is running", 
-        "device": device,
-        "cached_files": cache_count,
-        "translation_cache": translation_cache_count,
-        "currently_processing": processing_count,
-        "translation_available": translation_loaded,
-        "supported_languages": list(LANGUAGE_CODES.keys())
-    }
+    """Root endpoint with system information"""
+    try:
+        with db_manager.get_connection() as conn:
+            cursor = conn.cursor()
+            
+            cursor.execute('SELECT COUNT(*) FROM cache')
+            cache_count = cursor.fetchone()[0] or 0
+            
+            cursor.execute('SELECT COUNT(*) FROM processing_status WHERE status = "processing"')
+            processing_count = cursor.fetchone()[0] or 0
+            
+            cursor.execute('SELECT COUNT(*) FROM translation_cache')
+            translation_cache_count = cursor.fetchone()[0] or 0
+        
+        return {
+            "message": "Whisper API with Translation is running", 
+            "device": device,
+            "cuda_available": torch.cuda.is_available(),
+            "cached_files": cache_count,
+            "translation_cache": translation_cache_count,
+            "currently_processing": processing_count,
+            "translation_available": translation_model is not None,
+            "supported_languages": list(LANGUAGE_CODES.keys())
+        }
+    except Exception as e:
+        logger.error(f"Error in root endpoint: {e}")
+        return {"error": "Unable to retrieve system information"}
 
 @app.post("/transcribe")
-async def transcribe_audio(background_tasks: BackgroundTasks, file: UploadFile = File(...), translate_to_english: bool = False):
+async def transcribe_audio(
+    background_tasks: BackgroundTasks,
+    file: UploadFile = File(...),
+    translate_to_english: bool = False,
+    language: Optional[str] = Query(None, description="Specify language code for better accuracy")
+):
+    """Transcribe audio file to text"""
     tmp_file_path = None
     
     try:
-        logger.info(f"Received file: {file.filename}, size: {file.size}, translate_to_english: {translate_to_english}")
+        logger.info(f"Received file: {file.filename}, size: {file.size}")
         
         if not file or not file.filename:
             raise HTTPException(status_code=400, detail="No valid file provided")
         
+        # Read file content
         contents = await file.read()
         file_size = len(contents)
         file_size_mb = file_size / (1024 * 1024)
         
-        logger.info(f"File read successfully, size: {file_size} bytes ({file_size_mb:.1f} MB)")
+        logger.info(f"File size: {file_size} bytes ({file_size_mb:.1f} MB)")
         
-        if file_size > 50 * 1024 * 1024:
-            raise HTTPException(status_code=413, detail="File too large")
+        # Validate file size
+        if file_size > 100 * 1024 * 1024:  # 100MB limit
+            raise HTTPException(status_code=413, detail="File too large (max 100MB)")
         
         if file_size == 0:
             raise HTTPException(status_code=400, detail="Empty file")
         
+        # Calculate file hash for caching
         file_hash = calculate_file_hash(contents, file.filename, file_size)
         logger.info(f"File hash: {file_hash}")
         
-        cached_result = get_from_cache(file_hash)
+        # Check cache
+        cached_result = await get_from_cache(file_hash)
         if cached_result:
-            logger.info("Found in cache, returning cached result")
-            remove_processing_status(file_hash)
-            try:
-                cached_data = json.loads(cached_result)
-                cached_data["from_cache"] = True
-                return JSONResponse(cached_data)
-            except:
-                return JSONResponse({
-                    "text": cached_result,
-                    "from_cache": True,
-                    "message": "Result returned from cache"
-                })
-        
-        processing_status = get_processing_status(file_hash)
+            logger.info("Cache hit - returning cached result")
+            await remove_processing_status(file_hash)
+            return JSONResponse(json.loads(cached_result))
+        
+        # Check if already processing
+        processing_status = await get_processing_status(file_hash)
         if processing_status:
             logger.info("File is currently being processed")
             return JSONResponse({
@@ -423,104 +539,119 @@ async def transcribe_audio(background_tasks: BackgroundTasks, file: UploadFile =
                 "progress": processing_status['progress'],
                 "estimated_time": processing_status['estimated_time'],
                 "elapsed_minutes": processing_status['elapsed_minutes'],
-                "message": f"File is being processed. Please wait {processing_status['estimated_time'] - processing_status['elapsed_minutes']} minutes"
+                "message": f"File is being processed. Estimated time remaining: {processing_status['estimated_time'] - processing_status['elapsed_minutes']} minutes"
             })
         
         logger.info("Starting new processing...")
         
-        file_ext = os.path.splitext(file.filename)[1].lower()
-        if not file_ext:
-            file_ext = ".wav"
-        
+        # Create temporary file
+        file_ext = os.path.splitext(file.filename)[1].lower() or ".wav"
         with tempfile.NamedTemporaryFile(delete=False, suffix=file_ext) as tmp_file:
             tmp_file.write(contents)
             tmp_file_path = tmp_file.name
         
-        logger.info(f"Temp file created: {tmp_file_path}")
+        logger.info(f"Created temp file: {tmp_file_path}")
         
+        # Estimate processing time
         estimated_time = estimate_processing_time(file_size_mb)
         
-        if file_size_mb < 5:
-            result = whisper_model.transcribe(
-                tmp_file_path, 
-                fp16=False if device == "cpu" else True,
-                language=None,
-                task="transcribe",
-                verbose=False,
-                word_timestamps=False
-            )
-            
-            text = result["text"].strip()
-            if not text:
-                text = "No text detected"
+        # Process small files immediately
+        if file_size_mb < 10:  # Increased threshold to 10MB
+            try:
+                # Transcribe with optional language hint
+                transcribe_args = {
+                    'fp16': (device != "cpu"),
+                    'language': language,
+                    'task': "transcribe",
+                    'verbose': False,
+                    'word_timestamps': False
+                }
                 
-            response_data = {"text": text, "from_cache": False}
-            
-            if translate_to_english and result.get("language") != "en":
-                english_result = whisper_model.transcribe(
-                    tmp_file_path, 
-                    fp16=False if device == "cpu" else True,
-                    language=None,
-                    task="translate",
-                    verbose=False,
-                    word_timestamps=False
+                result = whisper_model.transcribe(tmp_file_path, **transcribe_args)
+                
+                text = result["text"].strip() or "No text detected"
+                detected_language = result.get("language", "unknown")
+                
+                response_data = {
+                    "text": text,
+                    "language": detected_language,
+                    "from_cache": False
+                }
+                
+                # Translate if requested
+                if translate_to_english and detected_language != "en":
+                    transcribe_args['task'] = "translate"
+                    english_result = whisper_model.transcribe(tmp_file_path, **transcribe_args)
+                    english_text = english_result["text"].strip()
+                    if english_text:
+                        response_data["english_text"] = english_text
+                
+                # Save to cache
+                await save_to_cache(
+                    file_hash, file.filename, file_size, 
+                    json.dumps(response_data), detected_language
                 )
                 
-                english_text = english_result["text"].strip()
-                if english_text:
-                    response_data["english_text"] = english_text
-            
-            save_to_cache(file_hash, file.filename, file_size, json.dumps(response_data))
-            
-            return JSONResponse(response_data)
+                return JSONResponse(response_data)
+                
+            except Exception as e:
+                logger.error(f"Error in immediate transcription: {e}")
+                raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")
         
         else:
-            add_processing_status(file_hash, file.filename, file_size, estimated_time)
+            # Large file - process in background
+            await add_processing_status(file_hash, file.filename, file_size, estimated_time)
             
-            background_tasks.add_task(background_transcription, tmp_file_path, file_hash, file.filename, file_size, translate_to_english)
+            background_tasks.add_task(
+                background_transcription, 
+                tmp_file_path, file_hash, file.filename, file_size, translate_to_english
+            )
             
             return JSONResponse({
                 "status": "processing_started",
                 "estimated_time": estimated_time,
                 "file_hash": file_hash,
-                "message": f"Processing started. It will take about {estimated_time} minutes. You can check the result later"
+                "message": f"Processing started. Estimated time: {estimated_time} minutes."
             })
     
+    except HTTPException:
+        raise
     except Exception as e:
-        logger.error(f"Error in transcription: {str(e)}")
-        if "No module named" in str(e):
-            raise HTTPException(status_code=500, detail="Missing required modules")
-        elif "CUDA" in str(e):
-            raise HTTPException(status_code=500, detail="GPU error")
-        elif "FFmpeg" in str(e):
-            raise HTTPException(status_code=500, detail="Audio processing error")
-        else:
-            raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")
+        logger.error(f"Error in transcription endpoint: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")
     
     finally:
-        if tmp_file_path and os.path.exists(tmp_file_path) and file_size < 5 * 1024 * 1024:
+        # Clean up temporary file for small immediate processing
+        if tmp_file_path and os.path.exists(tmp_file_path) and file_size_mb < 10:
             try:
                 os.unlink(tmp_file_path)
-                logger.info(f"Temp file deleted: {tmp_file_path}")
-            except:
-                pass
+            except Exception as e:
+                logger.error(f"Error deleting temp file: {e}")
 
 @app.post("/translate")
-async def translate_endpoint(text: str, target_language: str):
+async def translate_endpoint(
+    text: str = Query(..., min_length=1, description="Text to translate"),
+    target_language: str = Query(..., description="Target language")
+):
     """Translate text to target language"""
     
-    if not translation_loaded:
+    if not translation_model:
         raise HTTPException(status_code=503, detail="Translation service not available")
     
-    if not text or not text.strip():
+    text = text.strip()
+    if not text:
         raise HTTPException(status_code=400, detail="Text is required")
     
-    if target_language.lower() not in LANGUAGE_CODES:
-        raise HTTPException(status_code=400, detail=f"Unsupported language. Supported: {list(LANGUAGE_CODES.keys())}")
+    target_language_lower = target_language.lower()
+    if target_language_lower not in LANGUAGE_CODES:
+        raise HTTPException(
+            status_code=400, 
+            detail=f"Unsupported language. Supported: {list(LANGUAGE_CODES.keys())}"
+        )
     
+    # Check cache
     text_hash = calculate_text_hash(text)
-    
-    cached_translation = get_translation_from_cache(text_hash, target_language.lower())
+    cached_translation = await get_translation_from_cache(text_hash, target_language_lower)
     if cached_translation:
         return JSONResponse({
             "text": text,
@@ -529,48 +660,80 @@ async def translate_endpoint(text: str, target_language: str):
             "from_cache": True
         })
     
-    translated_text = translate_text(text, target_language)
-    
-    if not translated_text:
-        raise HTTPException(status_code=500, detail="Translation failed")
-    
-    save_translation_to_cache(text_hash, target_language.lower(), translated_text)
-    
-    return JSONResponse({
-        "text": text,
-        "translated_text": translated_text,
-        "target_language": target_language,
-        "from_cache": False
-    })
+    # Perform translation
+    try:
+        target_code = LANGUAGE_CODES[target_language_lower]
+        
+        # Tokenize input
+        inputs = translation_tokenizer(
+            text, 
+            return_tensors="pt", 
+            padding=True, 
+            truncation=True, 
+            max_length=512
+        )
+        
+        if device == "cuda":
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+        
+        # Generate translation
+        translated_tokens = translation_model.generate(
+            **inputs, 
+            forced_bos_token_id=translation_tokenizer.lang_code_to_id[target_code],
+            max_length=512,
+            num_beams=5,
+            early_stopping=True
+        )
+        
+        # Decode output
+        translated_text = translation_tokenizer.batch_decode(
+            translated_tokens, 
+            skip_special_tokens=True
+        )[0].strip()
+        
+        if not translated_text:
+            raise HTTPException(status_code=500, detail="Translation returned empty result")
+        
+        # Save to cache
+        await save_translation_to_cache(text_hash, target_language_lower, translated_text)
+        
+        return JSONResponse({
+            "text": text,
+            "translated_text": translated_text,
+            "target_language": target_language,
+            "from_cache": False
+        })
+        
+    except Exception as e:
+        logger.error(f"Translation error: {e}")
+        raise HTTPException(status_code=500, detail=f"Translation failed: {str(e)}")
 
 @app.get("/languages")
 async def get_supported_languages():
     """Get list of supported languages for translation"""
     return JSONResponse({
-        "supported_languages": LANGUAGE_CODES,
-        "translation_available": translation_loaded
+        "supported_languages": list(LANGUAGE_CODES.keys()),
+        "language_codes": LANGUAGE_CODES,
+        "translation_available": translation_model is not None
     })
 
 @app.get("/status/{file_hash}")
 async def check_status(file_hash: str):
-    cached_result = get_from_cache(file_hash)
+    """Check processing status for a file"""
+    # Check cache first
+    cached_result = await get_from_cache(file_hash)
     if cached_result:
-        remove_processing_status(file_hash)
-        try:
-            cached_data = json.loads(cached_result)
-            cached_data["status"] = "completed"
-            cached_data["from_cache"] = True
-            cached_data["message"] = "Processing completed and result is ready"
-            return JSONResponse(cached_data)
-        except:
-            return JSONResponse({
-                "status": "completed",
-                "text": cached_result,
-                "from_cache": True,
-                "message": "Processing completed and result is ready"
-            })
+        await remove_processing_status(file_hash)
+        cached_data = json.loads(cached_result)
+        cached_data.update({
+            "status": "completed",
+            "from_cache": True,
+            "message": "Processing completed and result is ready"
+        })
+        return JSONResponse(cached_data)
     
-    processing_status = get_processing_status(file_hash)
+    # Check processing status
+    processing_status = await get_processing_status(file_hash)
     if processing_status:
         remaining_time = max(0, processing_status['estimated_time'] - processing_status['elapsed_minutes'])
         return JSONResponse({
@@ -584,42 +747,70 @@ async def check_status(file_hash: str):
     
     return JSONResponse({
         "status": "not_found",
-        "message": "File not found"
-    })
+        "message": "File not found in cache or processing queue"
+    }, status_code=404)
 
 @app.get("/cache/stats")
 async def cache_stats():
+    """Get cache statistics"""
     try:
-        conn = sqlite3.connect('transcription_cache.db')
-        cursor = conn.cursor()
-        
-        cursor.execute('SELECT COUNT(*) FROM cache')
-        total_count = cursor.fetchone()[0]
-        
-        cursor.execute('SELECT COUNT(*) FROM cache WHERE created_at >= datetime("now", "-1 day")')
-        today_count = cursor.fetchone()[0]
-        
-        cursor.execute('SELECT AVG(LENGTH(transcription)) FROM cache')
-        avg_text_length = cursor.fetchone()[0] or 0
-        
-        cursor.execute('SELECT COUNT(*) FROM processing_status WHERE status = "processing"')
-        processing_count = cursor.fetchone()[0]
-        
-        cursor.execute('SELECT COUNT(*) FROM translation_cache')
-        translation_count = cursor.fetchone()[0]
-        
-        conn.close()
-        
+        with db_manager.get_connection() as conn:
+            cursor = conn.cursor()
+            
+            cursor.execute('SELECT COUNT(*) FROM cache')
+            total_count = cursor.fetchone()[0] or 0
+            
+            cursor.execute('SELECT COUNT(*) FROM cache WHERE created_at >= datetime("now", "-1 day")')
+            today_count = cursor.fetchone()[0] or 0
+            
+            cursor.execute('SELECT AVG(LENGTH(transcription)) FROM cache')
+            avg_text_length = cursor.fetchone()[0] or 0
+            
+            cursor.execute('SELECT COUNT(*) FROM processing_status WHERE status = "processing"')
+            processing_count = cursor.fetchone()[0] or 0
+            
+            cursor.execute('SELECT COUNT(*) FROM translation_cache')
+            translation_count = cursor.fetchone()[0] or 0
+            
+            cursor.execute('SELECT SUM(file_size) FROM cache')
+            total_cache_size = cursor.fetchone()[0] or 0
+            
         return {
             "total_cached_files": total_count,
             "cached_today": today_count,
+            "total_cache_size_mb": round(total_cache_size / (1024 * 1024), 2),
             "average_text_length": int(avg_text_length),
             "currently_processing": processing_count,
             "translation_cache_count": translation_count,
-            "translation_available": translation_loaded
+            "translation_available": translation_model is not None
         }
     except Exception as e:
-        return {"error": str(e)}
+        logger.error(f"Error getting cache stats: {e}")
+        return {"error": "Unable to retrieve cache statistics"}
+
+@app.post("/cache/cleanup")
+async def manual_cache_cleanup(days: int = 30):
+    """Manually trigger cache cleanup"""
+    await cleanup_old_cache(days)
+    return {"message": f"Cache cleanup initiated for entries older than {days} days"}
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "timestamp": datetime.now().isoformat(),
+        "device": device,
+        "cuda_available": torch.cuda.is_available(),
+        "whisper_loaded": whisper_model is not None,
+        "translation_loaded": translation_model is not None
+    }
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860, timeout_keep_alive=900)
+    uvicorn.run(
+        app, 
+        host="0.0.0.0", 
+        port=7860, 
+        timeout_keep_alive=900,
+        log_config=None
+    )