main.py

from fastapi import FastAPI, HTTPException
import logging
from models import ReviewRequest
from redis_client import r
from multiprocessing import Process, Queue
from contextlib import asynccontextmanager
import worker

MAX_QUEUE_SIZE = 50

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
)

logger = logging.getLogger("tabelog_app")

result_queue = Queue()

@asynccontextmanager
async def lifespan(app: FastAPI):
    p = Process(target=worker.run, args=(result_queue,))
    p.start()

    yield

    p.terminate()
    p.join()


app = FastAPI(lifespan=lifespan)


    
@app.post("/enqueue")
def enqueue_url(request: ReviewRequest):
    if r.llen("task_queue") >= MAX_QUEUE_SIZE:
        raise HTTPException(status_code=429, detail="대기열이 가득 찼습니다.")
    r.rpush("task_queue", request.model_dump_json())
    return {"message": "작업이 대기열에 추가되었습니다.", "queue_size": r.llen("task_queue")}

@app.get("/result")
def get_result():
    if result_queue.empty():
        return None
    
    return result_queue.get()

@app.get("/queue/status")
def get_queue_status():
    return {
        "queue_size": r.llen("task_queue"),
        "max_queue_size": MAX_QUEUE_SIZE
    }

@app.get("/result_size")
def get_queue_status():
    return {
        "queue_size": result_queue.qsize()
    }


@app.get("/")
def root():
    return {"message": "KcELECTRA NER API is running!"}