mbart-translation/main.py

from fastapi import FastAPI, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
import logging

from config import config
from translator import translator
from models import (
    TranslationRequest,
    TranslationResponse,
    HealthResponse,
    LanguagesResponse,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """애플리케이션 시작 시 모델을 로드합니다."""
    logger.info("Starting up: Loading mBART model...")
    try:
        translator.load_model()
        logger.info("Model loaded successfully")
    except Exception as e:
        logger.error(f"Failed to load model: {str(e)}")
        raise
    yield
    logger.info("Shutting down...")


app = FastAPI(
    title="mBART Translation API",
    description="mBART 모델을 사용한 다국어 번역 API 서비스",
    version="1.0.0",
    lifespan=lifespan,
)

# CORS 설정
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/", tags=["Root"])
async def root():
    """API 루트 엔드포인트"""
    return {
        "message": "mBART Translation API",
        "docs": "/docs",
        "health": "/health",
    }


@app.get("/health", response_model=HealthResponse, tags=["Health"])
async def health_check():
    """서비스 헬스 체크"""
    return HealthResponse(
        status="healthy" if translator.is_ready() else "not ready",
        model_loaded=translator.is_ready(),
        device=translator.device,
    )


@app.get("/languages", response_model=LanguagesResponse, tags=["Languages"])
async def get_supported_languages():
    """지원하는 언어 목록 조회"""
    return LanguagesResponse(supported_languages=config.SUPPORTED_LANGUAGES)


@app.post(
    "/translate",
    response_model=TranslationResponse,
    tags=["Translation"],
    status_code=status.HTTP_200_OK,
)
async def translate_text(request: TranslationRequest):
    """
    텍스트 번역 API

    - **text**: 번역할 텍스트
    - **source_lang**: 소스 언어 코드 (예: ko, en, ja)
    - **target_lang**: 타겟 언어 코드 (예: en, ko, ja)
    """
    if not translator.is_ready():
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Translation model is not ready",
        )

    # 언어 코드 검증
    if request.source_lang not in config.SUPPORTED_LANGUAGES:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Unsupported source language: {request.source_lang}",
        )

    if request.target_lang not in config.SUPPORTED_LANGUAGES:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Unsupported target language: {request.target_lang}",
        )

    try:
        translated_text = translator.translate(
            text=request.text,
            source_lang=request.source_lang,
            target_lang=request.target_lang,
        )

        return TranslationResponse(
            translated_text=translated_text,
            source_lang=request.source_lang,
            target_lang=request.target_lang,
            original_text=request.text,
        )

    except ValueError as e:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)
        )
    except Exception as e:
        logger.error(f"Translation failed: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Translation failed",
        )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "main:app",
        host=config.HOST,
        port=config.PORT,
        reload=True,
    )