from pydantic import BaseModel
from fastapi import FastAPI, UploadFile, File
from app.pdf_handler import process_pdf
from app.rag_pipeline import generate_answer
from app.test import query_test


app = FastAPI()


class ChatRequest(BaseModel):
    query: str
    top_k: int = 3


@app.post("/chat/")
async def chat(request: ChatRequest):
    return generate_answer(request.query, request.top_k)


@app.post("/upload/")
async def upload_pdf(file: UploadFile = File(...)):
    return process_pdf(file)


@app.post("/test/")
async def chat(query: str, top_k: int = 3):
    return query_test(query)