실습 1 — Retriever Node 분리 실습 따라 하기

# ==============================================================
# 24차시. PPT 24차시_RetrieverNode분리. 코드 폴더는 src/24
# PPT: 24차시 Retriever Node
# 역할: 24차시 Retriever와 25차시 Generator, fallback을 연결한 그래프입니다.
# 읽는 키: question, retrieval_ok, documents, answer, has_citation, log
# 쓰는 키: ask() 반환: answer, sources, cited, log, ok
# 그래프:
#   START --> retrieve
#     |-- ok ----> generate --> END
#     +-- empty -> fallback --> END
# --------------------------------------------------------------
# 파이썬만 아는 분을 위한 말
# State는 노드들이 함께 쓰는 딕셔너리입니다. 함수는 state를 읽고, 바꿀 키만 return 합니다.
# return에 적지 않은 키는 이전 값이 남습니다.
# log와 tried_queries는 리스트라서, 새로 return한 항목이 뒤에 이어 붙습니다.
# Retriever: 검색 담당. 질문과 비슷한 문서 조각만 고릅니다. 답 문장은 만들지 않습니다.
# Generator: 고른 문서를 근거로 답 문장을 만듭니다. 문서가 없으면 안내 문장만 돌려줍니다.
# Verifier: 답이 근거와 맞는지 검사하고 grade를 적습니다.
#   grade=pass     답을 확정하고 그래프를 끝냅니다.
#   grade=retry    같은 문서로 답 문장만 다시 씁니다. retries가 1 늘어납니다.
#   grade=research 문서가 부족합니다. 검색어를 바꾼 뒤 다시 찾습니다. rewrites가 늘어납니다.
#   grade=giveup   더 진행하지 않고 fallback으로 안내 문장만 돌려줍니다.
# fallback: 실패했을 때 답을 지어내지 않고 안내 문장을 반환하는 노드입니다.
# intent: 질문 종류. greeting 인사, calc 계산, scope 문서 밖, document 문서 검색.
# rewrite: 사용자가 친 question은 유지하고, 검색에 쓰는 query만 바꿉니다.
# safe_node: 노드를 try로 감쌉니다. 오류가 나도 그래프는 멈추지 않고 node_error에 적습니다.
# graph.py와 갈래는 같고, 답 생성 함수만 generator_node로 바뀌었습니다.
# ok 키는 검색 성공이 아니라 ask()가 항상 True로 넣는 실행 완료 표시입니다.
# cited는 Generator가 적은 has_citation 입니다. 인용 번호 [1]이 문서 개수 안에 있으면 참입니다.
# ==============================================================

# ============================================================
# Graph 24·25
# Retriever Node와 Generator Node를 LangGraph에 연결합니다.
# 검색 성공 여부에 따라 답변 생성 또는 안내 메시지로 이동합니다.
# ============================================================

import os
import sys
import warnings

# 현재 파일이 있는 폴더
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

# 필요한 차시 폴더를 Python 검색 경로에 추가합니다.
sys.path.insert(0, os.path.join(CURRENT_DIR, "..", "18"))
sys.path.insert(0, os.path.join(CURRENT_DIR, "..", "19"))
sys.path.insert(0, os.path.join(CURRENT_DIR, "..", "25"))

# 불필요한 경고 메시지를 숨깁니다.
warnings.filterwarnings("ignore")

# LangGraph의 그래프 구성 도구를 가져옵니다.
from langgraph.graph import StateGraph, START, END

# 그래프 State의 자료형을 가져옵니다.
from graph_state import RAGState

# 질문을 State 형태로 만들어 주는 함수를 가져옵니다.
from graph_state2 import make_initial_state

# 검색 노드를 가져옵니다.
from retriever import retriever_node

# 생성 노드를 가져옵니다.
from generator import generator_node

# 검색 실패 시 사용할 안내 노드를 가져옵니다.
from fallback import fallback_node

def route_after_retrieve(state: RAGState) -> str:
    """
    검색 결과에 따라 다음 노드를 결정합니다.

    검색 성공:
        generate 노드로 이동합니다.

    검색 실패:
        fallback 노드로 이동합니다.
    """

    # State에서 검색 성공 여부를 가져옵니다.
    retrieval_ok = state.get("retrieval_ok")

    # 검색에 성공했으면 "ok"를 반환합니다.
    if retrieval_ok:
        return "ok"

    # 검색에 실패했으면 "empty"를 반환합니다.
    return "empty"

def build_graph():
    """
    Retriever, Generator, Fallback 노드를 연결합니다.

    그래프의 흐름:
        START
          ↓
        retrieve
          ├─ 성공 → generate → END
          └─ 실패 → fallback → END
    """

    # RAGState를 사용하는 그래프를 만듭니다.
    graph = StateGraph(RAGState)

    # 검색 노드를 등록합니다.
    graph.add_node("retrieve", retriever_node)

    # 생성 노드를 등록합니다.
    graph.add_node("generate", generator_node)

    # 검색 실패 처리 노드를 등록합니다.
    graph.add_node("fallback", fallback_node)

    # 그래프 시작점에서 retrieve로 이동합니다.
    graph.add_edge(START, "retrieve")

    # retrieve 결과에 따라 다음 노드를 선택합니다.
    graph.add_conditional_edges(
        "retrieve",
        route_after_retrieve,
        {
            "ok": "generate",
            "empty": "fallback",
        },
    )

    # 답변 생성이 끝나면 그래프를 종료합니다.
    graph.add_edge("generate", END)

    # fallback 안내가 끝나면 그래프를 종료합니다.
    graph.add_edge("fallback", END)

    # 그래프를 실행할 수 있는 애플리케이션으로 컴파일합니다.
    application = graph.compile()

    return application

# 그래프를 한 번만 생성합니다.
app = build_graph()

def ask(question: str) -> dict:
    """
    질문 하나를 그래프에 전달하고 최종 결과를 반환합니다.

    입력:
        question: 사용자의 질문

    출력:
        answer: 최종 답변
        sources: 답변에 사용된 문서 출처
        cited: 인용 번호 존재 여부
        log: 전체 처리 기록
        ok: 실행 성공 여부
    """

    # 질문을 초기 State로 변환합니다.
    initial_state = make_initial_state(question)

    # 그래프를 실행합니다.
    final_state = app.invoke(
        initial_state,
        {
            "recursion_limit": 20
        },
    )

    # 최종 State에서 필요한 정보만 골라 반환합니다.
    return {
        "answer": final_state.get("answer", ""),
        "sources": get_sources(final_state),
        "cited": final_state.get("has_citation", False),
        "log": final_state.get("log", []),
        "ok": True,
    }

def get_sources(state: dict) -> list:
    """
    State의 documents에서 파일명과 페이지 정보를 추출합니다.
    """

    # 검색된 문서 목록을 가져옵니다.
    documents = state.get("documents", [])

    # 출처 목록을 저장할 변수입니다.
    sources = []

    # 문서를 하나씩 확인합니다.
    for document in documents:

        # 문서의 메타데이터를 가져옵니다.
        metadata = document.metadata

        # 파일명과 페이지 번호를 출처 목록에 추가합니다.
        sources.append(
            {
                "file": metadata.get("filename", "?"),
                "page": metadata.get("page_no", "?"),
            }
        )

    return sources

# 이 파일을 직접 실행했을 때만 테스트합니다.
if __name__ == "__main__":

    # 테스트할 질문 목록입니다.
    questions = [
        "환불은 며칠 이내인가요?",
        "반품하고 싶은데 언제까지?",
    ]

    # 질문을 하나씩 그래프에 전달합니다.
    for question in questions:

        # 전체 그래프를 실행합니다.
        result = ask(question)

        # 결과를 출력합니다.
        print("\n" + "=" * 55)
        print(f"질문: {question}")
        print(f"답변: {result['answer'][:110]}")
        print(f"출처: {result['sources']}")
        print(f"인용 정상: {result['cited']}")
        print(f"실행 성공: {result['ok']}")
        print(f"처리 기록: {result['log']}")

# ==============================================================
# 24차시. PPT 24차시_RetrieverNode분리. 코드 폴더는 src/24
# PPT: 24차시 Retriever Node
# 역할: 검색 다음 생성, 그 다음 검증. 검증이 실패하면 bump로 retries를 올리고 생성을 반복합니다.
# 읽는 키: retrieval_ok, grade, retries, documents, answer, reason
# 쓰는 키: retries, log, ask() 반환 answer, sources, grade, reason, retries, verified, log, ok
# 그래프:
#   START --> retrieve
#     |-- ok ----> generate --> verify --+-- done -----> END
#     |                                  |-- giveup ---> fallback --> END
#     |                                  +-- regenerate -> bump --> generate
#     +-- empty -> fallback --> END
# --------------------------------------------------------------
# 파이썬만 아는 분을 위한 말
# State는 노드들이 함께 쓰는 딕셔너리입니다. 함수는 state를 읽고, 바꿀 키만 return 합니다.
# return에 적지 않은 키는 이전 값이 남습니다.
# log와 tried_queries는 리스트라서, 새로 return한 항목이 뒤에 이어 붙습니다.
# Retriever: 검색 담당. 질문과 비슷한 문서 조각만 고릅니다. 답 문장은 만들지 않습니다.
# Generator: 고른 문서를 근거로 답 문장을 만듭니다. 문서가 없으면 안내 문장만 돌려줍니다.
# Verifier: 답이 근거와 맞는지 검사하고 grade를 적습니다.
#   grade=pass     답을 확정하고 그래프를 끝냅니다.
#   grade=retry    같은 문서로 답 문장만 다시 씁니다. retries가 1 늘어납니다.
#   grade=research 문서가 부족합니다. 검색어를 바꾼 뒤 다시 찾습니다. rewrites가 늘어납니다.
#   grade=giveup   더 진행하지 않고 fallback으로 안내 문장만 돌려줍니다.
# fallback: 실패했을 때 답을 지어내지 않고 안내 문장을 반환하는 노드입니다.
# intent: 질문 종류. greeting 인사, calc 계산, scope 문서 밖, document 문서 검색.
# rewrite: 사용자가 친 question은 유지하고, 검색에 쓰는 query만 바꿉니다.
# safe_node: 노드를 try로 감쌉니다. 오류가 나도 그래프는 멈추지 않고 node_error에 적습니다.
# bump_node는 검색을 다시 하지 않습니다. retries만 1 더하고 generate로 돌아갑니다.
# route_after_verify: pass면 done, giveup이거나 retries가 MAX_RETRY 이상이면 giveup, 아니면 regenerate.
# verified는 grade가 pass일 때만 True입니다.
# 이 파일은 24 폴더에 있지만, 재생성 루프를 미리 붙여 둔 그래프입니다. rewrite 분기는 아직 없습니다.
# ==============================================================

import os
import sys

# 18·19·25·26차시 폴더를 모듈 찾는 경로 맨 앞에 넣습니다.
# 그래야 아래 import 가 State, 각 노드 파일을 다른 차시 폴더에서 찾습니다.
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "18"))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "19"))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "25"))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "26"))

from langgraph.graph import StateGraph, START, END

# State 모양(18차시)과 첫 State를 만드는 함수입니다.
from graph_state import RAGState
from graph_state2 import make_initial_state

# 노드 네 개: 24차시 검색, 25차시 생성, 26차시 검증, 실패 안내.
from retriever import retriever_node
from generator import generator_node
from verifier import verifier_node
from fallback import fallback_node

import config

# config 에 MAX_RETRY 가 없으면 2를 씁니다. 재생성 루프의 상한입니다.
MAX_RETRY = getattr(config, "MAX_RETRY", 2)

# ===================================================
# 재시도 횟수 증가 노드
# ===================================================

def bump_node(state: RAGState) -> dict:
    # 검색은 다시 하지 않고 재시도 횟수만 1 올립니다.
    n = state.get("retries", 0) + 1

    # 바꾼 칸(retries)과 로그 한 줄만 돌려줍니다.
    return {
        "retries": n,
        "log": [f"재시도 {n}회차 진입"],
    }

# ===================================================
# 라우팅 함수
# ===================================================

def route_after_retrieve(state: RAGState) -> str:
    # retrieval_ok 가 참이면 ok, 아니면 empty. 문자열만 돌려주고 노드는 부르지 않습니다.
    return "ok" if state.get("retrieval_ok") else "empty"

def route_after_verify(state: RAGState) -> str:
    # State 에 grade 칸이 없으면 retry 로 봅니다.
    grade = state.get("grade", "retry")

    # 검증 통과
    if grade == "pass":
        return "done"

    # Verifier가 명시적으로 포기
    if grade == "giveup":
        return "giveup"

    # 최대 재시도 횟수 초과
    if state.get("retries", 0) >= MAX_RETRY:
        return "giveup"

    # 그 외에는 재생성
    return "regenerate"

# ===================================================
# 그래프 조립
# ===================================================

def build_graph():
    # RAGState 를 공유 노트로 쓰는 빈 그래프를 만듭니다.
    g = StateGraph(RAGState)

    # 노드 등록
    g.add_node("retrieve", retriever_node)
    g.add_node("generate", generator_node)
    g.add_node("verify", verifier_node)
    g.add_node("bump", bump_node)
    g.add_node("fallback", fallback_node)

    # 시작 → 검색
    g.add_edge(START, "retrieve")

    # 검색 결과에 따른 분기
    # 반환 문자열(ok/empty) → 다음 노드 이름을 경로맵으로 이어 줍니다.
    g.add_conditional_edges(
        "retrieve",
        route_after_retrieve,
        {
            "ok": "generate",
            "empty": "fallback",
        },
    )

    # 생성 → 검증
    g.add_edge("generate", "verify")

    # 검증 결과에 따른 분기
    # 경로맵 키는 route_after_verify 가 낼 수 있는 값 세 개와 똑같아야 합니다.
    g.add_conditional_edges(
        "verify",
        route_after_verify,
        {
            "done": END,
            "giveup": "fallback",
            "regenerate": "bump",
        },
    )

    # 재시도 횟수 증가 → 다시 생성
    g.add_edge("bump", "generate")

    # 검색 실패 또는 검증 포기
    g.add_edge("fallback", END)

    # 30차시에서 rewrite 및 재검색 루프 추가
    return g.compile()

# 파일을 불러올 때 그래프를 한 번만 compile 해 둡니다.
app = build_graph()

def ask(question: str, verbose: bool = False) -> dict:     
    # 루프가 있으므로 recursion_limit 25 로 비상 브레이크를 겁니다.
    final = app.invoke(make_initial_state(question),
                       {"recursion_limit": 25})
    
    # verbose 이면 노드들이 쌓은 log 를 한 줄씩 보여 줍니다.
    if verbose:
        print(f"\nQ: {question}")
        for line in final.get("log", []):
            print(f"   · {line}")
            
    # 15차시 ask() 키에 grade, reason, retries, verified 를 더해 돌려줍니다.
    return {
        "answer":   final["answer"],
        "sources":  [{"file": d.metadata.get("filename"),
                      "page": d.metadata.get("page_no")}
                     for d in final.get("documents", [])],
        "grade":    final.get("grade", ""),
        "reason":   final.get("reason", ""),
        "retries":  final.get("retries", 0),
        "verified": final.get("grade") == "pass",
        "log":      final.get("log", []),
        "ok":       True,
    }
    
if __name__ == "__main__":
    # 연결이 맞는지 문자 그림으로 먼저 확인합니다.
    print(app.get_graph().draw_ascii())
    for q in ["환불은 며칠 이내에 신청해야 하나요?",
              "대표이사가 누구인가요?"]:
        # 정상 질문 하나, 문서에 없는 질문 하나로 두 경로를 봅니다.
        r = ask(q, verbose=True)
        print(f"   A: {r['answer'][:70]}")
        print(f"   판정={r['grade']} 재시도={r['retries']}회 "
              f"검증통과={r['verified']}"
        )