app/routers/upload.py

from fastapi import APIRouter, File, UploadFile, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, RedirectResponse
from fastapi.templating import Jinja2Templates
from app.modules.embedding import process_and_store_file
import shutil
import os
import uuid
import asyncio

router = APIRouter()

# 업로드된 파일을 저장할 디렉토리 설정
UPLOAD_DIRECTORY = "./uploaded_files"
if not os.path.exists(UPLOAD_DIRECTORY):
    os.makedirs(UPLOAD_DIRECTORY, exist_ok=True)
    os.chmod(UPLOAD_DIRECTORY, 0o777)

# 템플릿 디렉토리 설정
templates = Jinja2Templates(directory="app/templates")

# WebSocket 연결 관리
connections = {}

@router.get("/", response_class=HTMLResponse)
async def main(request: Request):
    return templates.TemplateResponse("upload.html", {"request": request})

@router.post("/uploadfile/{user_id}/")
async def upload_file(user_id: str, file: UploadFile = File(...)):
    # 고유 파일 이름 생성
    unique_filename = f"{uuid.uuid4()}_{file.filename}"
    file_location = os.path.join(UPLOAD_DIRECTORY, unique_filename)

    # 파일 저장
    try:
        with open(file_location, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        print(f"File saved at {file_location}")
    except Exception as e:
        return {"error": f"파일 저장 오류: {e}"}

    # 비동기 작업 실행
    asyncio.create_task(process_and_store_file(file_location, user_id, connections.get(user_id), UPLOAD_DIRECTORY))

    return {"info": f"File '{file.filename}' uploaded successfully."}

@router.websocket("/ws/{user_id}")
async def websocket_endpoint(websocket: WebSocket, user_id: str):
    await websocket.accept()
    connections[user_id] = websocket  # WebSocket 연결 저장
    try:
        while True:
            data = await websocket.receive_text()
            print(f"Received from {user_id}: {data}")
    except WebSocketDisconnect:
        print(f"User {user_id} disconnected")
        del connections[user_id]  # 연결 제거
    except Exception as e:
        print(f"Error with user {user_id}: {e}")
        del connections[user_id]  # 연결 제거