api for streamlit app
This commit is contained in:
parent
dafd3fba38
commit
3592737fbd
253
app.py
Normal file
253
app.py
Normal file
@ -0,0 +1,253 @@
|
||||
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)
|
||||
Loading…
x
Reference in New Issue
Block a user