텍스트 PDF로 RAG 동작 확인

RAG.py

import os
from dotenv import load_dotenv
from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

load_dotenv()

def create_rag_chain(pdf_path: str):
    # 1. PDF 읽기
    loader = PyPDFLoader(pdf_path)
    documents = loader.load()

    # 2. 텍스트 청크로 쪼개기
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
    chunks = splitter.split_documents(documents)

    # 3. 임베딩 → 벡터DB 저장
    embeddings = GoogleGenerativeAIEmbeddings(
        model="gemini-embedding-001",
        google_api_key=os.getenv("GOOGLE_API_KEY")
    )
    vectorstore = FAISS.from_documents(chunks, embeddings)

    # 4. LLM 설정
    llm = ChatGoogleGenerativeAI(
        model="gemini-2.5-flash",
        google_api_key=os.getenv("GOOGLE_API_KEY")
    )

    # 5. 프롬프트 설정 - LLM에게 어떻게 답변할지 지시하는 것
    prompt = ChatPromptTemplate.from_template("""
아래 문서 내용을 바탕으로 질문에 답변해주세요.
문서에 없는 내용은 "해당 내용은 문서에서 찾을 수 없습니다"라고 답변하세요.

문서 내용:
{context}

질문: {question}
""")

    # 6. RAG 체인 완성
    retriever = vectorstore.as_retriever()
    chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )
    return chain

Main.py

from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import shutil
import os
from app.rag import create_rag_chain

app = FastAPI()

# CORS 설정 - 다른 주소에서 오는 요청을 별도의 인증 없이 허용하는 설정
# Spring의 @CrossOrigin 이랑 똑같은 역할
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

rag_chain = None

class QuestionRequest(BaseModel):
    question: str

@app.post("/upload-pdf")
async def upload_pdf(file: UploadFile = File(...)):
    # PDF 파일을 data 폴더에 저장
    global rag_chain
    file_path = f"data/{file.filename}"
    with open(file_path, "wb") as f:
        shutil.copyfileobj(file.file, f)
    
    # RAG 체인 생성
    rag_chain = create_rag_chain(file_path)
    return {"message": f"{file.filename} 업로드 완료. 질문할 수 있어요!"}

@app.post("/ask")
async def ask_question(request: QuestionRequest):
    if rag_chain is None:
        return {"error": "등록된 PDF가 없습니다."}
    
    result = rag_chain.invoke(request.question)
    return {"answer": result}

현장 관리 데이터 PDF 만들기 create_sample_pdf.py

pip install reportlab

from reportlab.lib.pagesizes import A4
from reportlab.pdfgen import canvas
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
import os

def create_sample_pdf():
    # 윈도우 기본 한글 폰트 등록
    font_path = "C:/Windows/Fonts/malgun.ttf"
    pdfmetrics.registerFont(TTFont("Korean", font_path))

    file_path = "data/sample_safety.pdf"
    pdf = canvas.Canvas(file_path, pagesize=A4)
    width, height = A4
    pdf.setFont("Korean", 12)

    lines = [
        "현장 안전수칙 가이드",
        "",
        "1. 화재 대피 요령",
        "- 화재 발생 시 즉시 비상벨을 누르세요.",
        "- 엘리베이터 사용을 금지하고 비상계단을 이용하세요.",
        "- 연기가 많을 경우 낮은 자세로 이동하세요.",
        "",
        "2. 전기 안전수칙",
        "- 젖은 손으로 전기 기기를 만지지 마세요.",
        "- 전선이 끊어진 경우 즉시 관리자에게 보고하세요.",
        "",
        "3. 추락 방지 수칙",
        "- 2미터 이상 고소 작업 시 안전벨트를 착용하세요.",
        "- 안전모는 현장 진입 시 반드시 착용하세요.",
        "",
        "4. 비상 연락처",
        "- 현장 관리자: 010-1234-5678",
        "- 소방서: 119",
        "- 응급실: 112",
    ]

    y = height - 50
    for line in lines:
        pdf.drawString(50, y, line)
        y -= 25

    pdf.save()
    print(f"PDF 생성 완료: {file_path}")

create_sample_pdf()

python create_sample_pdf.py

서버실행해서 확인하기

uvicorn app.main:app --reload

http://127.0.0.1:8000/docs

질문에 대한 대답