app.py

import gradio as gr
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import os
import warnings
from dotenv import load_dotenv

# --- 1. 환경 변수 및 모듈 로드 ---
# 이 부분은 기존 코드와 동일합니다.
load_dotenv()
warnings.filterwarnings('ignore')

from prediction import single_prediction
from chatbot import process_chatbot_query_with_llm
from ui import create_ui


# --- 2. FastAPI 애플리케이션 생성 및 설정 ---

# FastAPI 앱을 생성합니다. 이제부터 이 앱이 전체 서버를 관리합니다.
app = FastAPI()

# CORS (Cross-Origin Resource Sharing) 미들웨어를 추가합니다.
# 이 부분이 WebSocket 403 Forbidden 오류를 해결하는 핵심입니다.
# 외부(로컬 PC, 다른 웹사이트 등)에서의 모든 종류의 접속을 허용합니다.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],          # 모든 출처에서의 접속을 허용합니다.
    allow_credentials=True,       # 인증 정보 (쿠키 등)를 허용합니다.
    allow_methods=["*"],          # 모든 HTTP 메소드 (GET, POST 등)를 허용합니다.
    allow_headers=["*"],          # 모든 HTTP 헤더를 허용합니다.
)


# --- 3. Gradio UI 생성 및 FastAPI에 통합 ---

# 기존과 동일하게 UI 생성 함수를 호출하여 Gradio 데모 객체를 만듭니다.
gradio_demo = create_ui(
    prediction_handler=single_prediction,
    chatbot_handler=process_chatbot_query_with_llm
)

# FastAPI 앱에 Gradio 데모를 탑재(mount)합니다.
# 이 함수는 WebSocket 경로를 포함한 모든 것을 올바르게 처리하는 표준 방식입니다.
app = gr.mount_gradio_app(
    app,                # 마운트할 FastAPI 앱
    gradio_demo,        # 마운트할 Gradio 데모
    path="/gradio_api"  # UI와 API가 서비스될 기본 경로
)


# --- 4. (선택사항) 서버 상태 확인용 루트 경로 ---
# 스페이스 주소로 바로 접속했을 때 서버가 살아있는지 확인할 수 있습니다.
@app.get("/")
def read_root():
    return {
        "status": "ok",
        "message": "FastAPI server is running.",
        "ui_path": "/gradio_api"
    }


# --- 5. (참고) 로컬 개발 환경에서 실행하기 위한 코드 ---
# Hugging Face Spaces에서는 이 부분이 직접 사용되지 않을 수 있지만,
# 로컬에서 'python app.py'로 테스트할 때 유용합니다.
if __name__ == "__main__":
    print("로컬 개발 서버를 시작합니다. http://127.0.0.1:7860/gradio_api 에서 확인하세요.")
    uvicorn.run(app, host="127.0.0.1", port=7860)