excel_translator/app.py

import os
import shutil
import logging
import asyncio
import tempfile
from fastapi import FastAPI, UploadFile, File, BackgroundTasks, HTTPException, Form
from fastapi.responses import FileResponse
from pydantic import BaseModel
from typing import Optional, List, Dict
from enum import Enum
import uuid
from datetime import datetime

# Import translation functionality from main.py
from main import translate_excel

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

app = FastAPI(
    title="Excel Translation API",
    description="API for translating Excel files with Google Translate or LLM",
    version="1.0.0"
)

# Directory to store uploaded and translated files
UPLOAD_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "uploads")
RESULTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results")

# Create directories if they don't exist
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(RESULTS_DIR, exist_ok=True)

# Dictionary to track translation jobs
translation_jobs = {}

class TranslationMethod(str, Enum):
    google = "google"
    llm = "llm"

class TranslationRequest(BaseModel):
    file_id: str
    target_language: str
    translation_method: TranslationMethod = TranslationMethod.google
    llm_model: Optional[str] = "llama3.1:8b"

class TranslationStatus(BaseModel):
    job_id: str
    file_name: str
    status: str
    target_language: str
    translation_method: str
    created_at: str
    completed_at: Optional[str] = None
    result_file: Optional[str] = None
    error: Optional[str] = None

@app.get("/")
def read_root():
    """Root endpoint with API information"""
    return {
        "message": "Excel Translation API",
        "docs": "/docs",
        "redoc": "/redoc"
    }

@app.post("/upload/", response_model=dict)
async def upload_file(file: UploadFile = File(...)):
    """Upload an Excel file for translation"""
    if not file.filename.endswith(('.xlsx', '.xls')):
        raise HTTPException(status_code=400, detail="Only Excel files (.xlsx, .xls) are supported")
    
    try:
        file_id = str(uuid.uuid4())
        file_path = os.path.join(UPLOAD_DIR, f"{file_id}_{file.filename}")
        
        # Save uploaded file
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        return {
            "file_id": file_id,
            "file_name": file.filename,
            "message": "File uploaded successfully"
        }
    except Exception as e:
        logging.error(f"Error uploading file: {e}")
        raise HTTPException(status_code=500, detail=f"Error uploading file: {str(e)}")

@app.post("/translate/", response_model=TranslationStatus)
async def translate_file(request: TranslationRequest, background_tasks: BackgroundTasks):
    """Start a translation job for the uploaded Excel file"""
    # Find the uploaded file
    uploaded_files = [f for f in os.listdir(UPLOAD_DIR) if f.startswith(f"{request.file_id}_")]
    if not uploaded_files:
        raise HTTPException(status_code=404, detail=f"File with ID {request.file_id} not found")
    
    file_path = os.path.join(UPLOAD_DIR, uploaded_files[0])
    file_name = uploaded_files[0][len(request.file_id) + 1:]  # Remove the UUID prefix
    
    # Create job ID
    job_id = str(uuid.uuid4())
    
    # Create initial job status
    job_status = TranslationStatus(
        job_id=job_id,
        file_name=file_name,
        status="pending",
        target_language=request.target_language,
        translation_method=request.translation_method,
        created_at=datetime.now().isoformat()
    )
    
    # Store job status
    translation_jobs[job_id] = job_status
    
    # Start background translation task
    background_tasks.add_task(
        process_translation,
        job_id=job_id,
        file_path=file_path,
        target_language=request.target_language,
        translation_method=request.translation_method,
        llm_model=request.llm_model
    )
    
    return job_status

async def process_translation(job_id: str, file_path: str, target_language: str, 
                             translation_method: str, llm_model: str):
    """Background task to process the translation"""
    try:
        # Update job status
        translation_jobs[job_id].status = "processing"
        
        # Perform translation
        result_file_path = await translate_excel(
            file_path=file_path,
            target_language=target_language,
            translation_method=translation_method,
            llm_model=llm_model
        )
        
        # Move the result file to results directory
        file_name = os.path.basename(result_file_path)
        new_result_path = os.path.join(RESULTS_DIR, f"{job_id}_{file_name}")
        shutil.copy2(result_file_path, new_result_path)
        
        # Update job status
        translation_jobs[job_id].status = "completed"
        translation_jobs[job_id].completed_at = datetime.now().isoformat()
        translation_jobs[job_id].result_file = f"{job_id}_{file_name}"
        
        logging.info(f"Translation job {job_id} completed successfully")
    except Exception as e:
        logging.error(f"Error in translation job {job_id}: {e}")
        translation_jobs[job_id].status = "failed"
        translation_jobs[job_id].error = str(e)

@app.get("/jobs/{job_id}", response_model=TranslationStatus)
async def get_job_status(job_id: str):
    """Get the status of a translation job"""
    if job_id not in translation_jobs:
        raise HTTPException(status_code=404, detail=f"Job with ID {job_id} not found")
    
    return translation_jobs[job_id]

@app.get("/jobs/", response_model=List[TranslationStatus])
async def list_jobs():
    """List all translation jobs"""
    return list(translation_jobs.values())

@app.get("/download/{job_id}")
async def download_result(job_id: str):
    """Download the translated Excel file"""
    if job_id not in translation_jobs:
        raise HTTPException(status_code=404, detail=f"Job with ID {job_id} not found")
    
    job = translation_jobs[job_id]
    
    if job.status != "completed":
        raise HTTPException(status_code=400, detail=f"Job {job_id} is not completed yet")
    
    if not job.result_file:
        raise HTTPException(status_code=400, detail=f"No result file available for job {job_id}")
    
    result_path = os.path.join(RESULTS_DIR, job.result_file)
    if not os.path.exists(result_path):
        raise HTTPException(status_code=404, detail=f"Result file not found")
    
    return FileResponse(
        path=result_path,
        filename=job.result_file[len(job_id) + 1:],  # Remove the UUID prefix
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    )

@app.post("/translate-form/")
async def translate_file_form(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    target_language: str = Form(...),
    translation_method: TranslationMethod = Form(TranslationMethod.google),
    llm_model: str = Form("llama3.1:8b")
):
    """Form-based endpoint for file upload and translation in one step"""
    if not file.filename.endswith(('.xlsx', '.xls')):
        raise HTTPException(status_code=400, detail="Only Excel files (.xlsx, .xls) are supported")
    
    try:
        # Generate IDs
        file_id = str(uuid.uuid4())
        job_id = str(uuid.uuid4())
        
        # Save uploaded file
        file_path = os.path.join(UPLOAD_DIR, f"{file_id}_{file.filename}")
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        # Create job status
        job_status = TranslationStatus(
            job_id=job_id,
            file_name=file.filename,
            status="pending",
            target_language=target_language,
            translation_method=translation_method,
            created_at=datetime.now().isoformat()
        )
        
        # Store job status
        translation_jobs[job_id] = job_status
        
        # Start translation in background
        background_tasks.add_task(
            process_translation,
            job_id=job_id,
            file_path=file_path,
            target_language=target_language,
            translation_method=translation_method,
            llm_model=llm_model
        )
        
        return {
            "job_id": job_id,
            "message": "Translation job started",
            "status_endpoint": f"/jobs/{job_id}"
        }
    except Exception as e:
        logging.error(f"Error processing translation request: {e}")
        raise HTTPException(status_code=500, detail=f"Error processing translation request: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)