9. 오늘의 완성 코드

# ============================================================
# 13차시 | PPT: 13차시_프롬프트최적화.pptx
# 역할: V3 prompt로 답을 만들고, 인용 검사 결과를 딕셔너리로 돌려줌
# 입력: 질문 3개 (환불 기한, 대표이사, 환불 방법과 수수료)
# 출력: 검색에 성공하면 Q만 출력. 0건일 때만 A와 경고를 출력
# 흐름:
#   search -> 없으면 안내 dict
#     -> 있으면 V3 체인 -> check_citation -> dict 반환 (답은 print하지 않음)
# ============================================================
# prompt: RAG_PROMPT_V3. 있는 내용만 답하고 문장 끝에 citation [1]을 요구합니다.
# citation: [숫자] 근거 번호. check_citation이 문서 개수 안인지 확인합니다.
# chunk: search가 찾은 문서 조각. sources에는 그 metadata 딕셔너리를 담습니다.

import os
import re
import sys

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser

# 11차시 폴더를 모듈 검색 경로에 넣어 search와 build_context를 불러옵니다.
CURRENT_DIR = os.path.dirname(__file__)
RETRIEVER_DIR = os.path.join(CURRENT_DIR, "..", "11")

sys.path.append(RETRIEVER_DIR)

from retriever import search, build_context
from prompts import RAG_PROMPT_V3
# 인용 검사는 validators.py의 함수를 그대로 다시 씁니다.
from validators import check_citation

# .env의 API 키를 읽습니다.
load_dotenv()

# 답이 흔들리지 않게 temperature=0으로 둡니다.
llm = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0
)

# V3 틀, 모델, 문자열 변환기를 미리 이어 둡니다.
_chain = RAG_PROMPT_V3 | llm | StrOutputParser()

# V3가 모를 때 쓰는 문장입니다. 답에 이 문장이 있으면 근거가 부족했다고 봅니다.
NO_INFO = "자료에서 확인할 수 없습니다"

# 질문 하나를 받아 답, 출처, 인용 여부, 근거 부족 여부를 딕셔너리로 돌려줍니다.
def ask(question, k=3, verbose=True):
    # search는 MIN_SCORE를 넘은 조각만 돌려줍니다. 하나도 없으면 빈 리스트입니다.
    documents = search(question, k=k)

    # verbose=True면 질문을 화면에 찍습니다.
    if verbose:
        print("Q:", question)

    # 통과한 조각이 없으면 모델을 호출하지 않고 바로 돌려줍니다.
    if not documents:
        result = {
            "answer": "관련 자료를 찾지 못했습니다.",
            "sources": [],
            "cited": False,
            "insufficient": True
        }

        if verbose:
            print("A:", result["answer"])
            print("⚠️ 검색 결과 없음")
            print("-" * 55)

        return result

    # 조각에 번호와 출처를 붙인 자료 문자열을 만듭니다.
    context = build_context(documents)

    # 미리 이어 둔 체인에 자료와 질문을 넣어 답 문자열을 받습니다.
    answer = _chain.invoke({
        "context": context,
        "question": question
    })

    # check_citation은 (통과 여부, 설명) 두 값을 돌려줍니다. 여기서는 통과 여부만 씁니다.
    cited, citation_message = check_citation(
        answer,
        len(documents)
    )

    # 출처로 각 조각의 metadata(파일 이름, 쪽 번호 등)를 그대로 담습니다.
    sources = []

    for document in documents:
        sources.append(document.metadata)

    # 답, 출처, 인용 성공 여부, '자료에서 확인할 수 없습니다' 포함 여부.
    # 이 딕셔너리는 반환만 하고, 성공 분기에서는 print하지 않습니다.
    return {
        "answer": answer,
        "sources": sources,
        "cited": cited,
        "insufficient": NO_INFO in answer
    }

# 직접 실행할 때만 세 질문을 돌립니다. 성공한 답은 출력하지 않으므로, 보고 싶으면 반환값을 print 해 보세요.
if __name__ == "__main__":
    ask("환불은 며칠 이내에 신청해야 하나요?")
    ask("대표이사가 누구인가요?")
    ask("환불 방법과 수수료를 알려주세요")