main.py

import os
from contextlib import asynccontextmanager
from typing import Literal

from dotenv import load_dotenv
from fastapi import FastAPI
from pydantic import BaseModel, Field

from rag import RAGSystem

load_dotenv()

# RAG 인스턴스 생성
rag_system: RAGSystem = RAGSystem()


class AgentRequest(BaseModel):
    query: str = Field(..., description="The multiple-choice legal question")


class AgentResponse(BaseModel):
    answer: Literal["A", "B", "C", "D"] = Field(..., description="The correct option")


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    if os.path.exists("data/train.csv"):
        rag_system.ingest_data("data/train.csv")  # type: ignore[attr-defined]
    yield
    # Shutdown
    pass


app = FastAPI(lifespan=lifespan)


@app.post("/request", response_model=AgentResponse)
async def predict(request: AgentRequest) -> AgentResponse:
    try:
        raw_answer: str = rag_system.get_answer(request.query)  # type: ignore[attr-defined]

        if raw_answer not in ["A", "B", "C", "D"]:
            raise ValueError("Invalid answer")

        return AgentResponse(answer=raw_answer)  # type: ignore[arg-type]
    except Exception as e:
        print(f"Error: {e}")
        raise