5. 실습 ② — Rewrite Node 구현

import config
from prompts import REWRITE_PROMPT
from nodes.generator import _llm
from langchain_core.output_parsers import StrOutputParser

_rewrite = REWRITE_PROMPT | _llm | StrOutputParser()

MAX_REWRITE = getattr(config, "MAX_REWRITE", 2)

def rewrite_node(state) -> dict:
    question = state.get("question", "")
    tried    = state.get("tried_queries", [])
    reason   = state.get("reason") or state.get("fail_reason", "")
    n        = state.get("rewrites", 0)

    # 상한 도달 시 재작성하지 않는다
    if n >= MAX_REWRITE:
        return {"rewrites": n,
                "log": [f"재작성 상한({MAX_REWRITE}) 도달"]}

    try:
        new_query = _rewrite.invoke({
            "question": question,
            "tried":    ", ".join(tried) if tried else "(없음)",
            "reason":   reason or "관련 자료를 찾지 못함",
        }).strip().strip('"').strip("'")
    except Exception as e:
        return {"rewrites": n + 1,
                "log": [f"재작성 오류({type(e).__name__})"]}

    # 빈 결과나 중복이면 재작성 실패로 처리
    if not new_query or new_query in tried:
        return {"rewrites": n + 1,
                "log": [f"재작성 실패(중복/빈값): {new_query}"]}

    return {
        "query":         new_query,
        "rewrites":      n + 1,
        "tried_queries": [new_query],      # Reducer로 누적
        "retrieval_ok":  False,            # 재검색 전 초기화
        "log": [f"재작성 {n+1}회: '{question[:20]}...' "
                f"-> '{new_query}'"],
    }

# ==============================================================
# 30차시. PPT 30차시_질의재작성재검색. 코드 폴더는 src/30
# PPT: 30차시 검색어 재작성
# 역할: Rewrite Node. 검색에 실패한 질문을 검색용 문장으로 바꿉니다. 검색 자체는 하지 않습니다.
# 읽는 키: question, tried_queries, rewrites, reason, fail_reason
# 쓰는 키: 성공 시 query, rewrites, tried_queries, retrieval_ok=False, rewrite_ok, log
# 그래프:
#   rewrites >= 상한 --> rewrite_ok=False
#   LLM 오류·빈 문장·중복 문장 --> rewrite_ok=False, rewrites+1
#   새 문장 성공 --> query=새 검색어, rewrites+1, tried_queries에 추가
# --------------------------------------------------------------
# 파이썬만 아는 분을 위한 말
# 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에 적습니다.
# question은 그대로 두고 query만 바꿉니다. 다음 Retriever가 query를 우선 사용합니다.
# retrieval_ok를 False로 되돌려, 이어지는 검색이 새 query를 쓰게 합니다.
# MAX_REWRITE 기본값은 config에 없으면 2입니다.
# reason이 없으면 fail_reason을 쓰고, 둘 다 없으면 '관련 자료를 찾지 못함'을 넣습니다.
# ==============================================================

# ============================================================
# 30차시 - Rewrite Node
# 검색에 실패한 질문을 검색하기 좋은 표현으로 다시 작성합니다.
# 원본 question은 유지하고, 검색용 query만 변경합니다.
# ============================================================

import os
import sys
import warnings

# ------------------------------------------------------------
# 필요한 모듈의 경로를 설정합니다.
# ------------------------------------------------------------

CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
SRC_DIR = os.path.abspath(os.path.join(CURRENT_DIR, ".."))

warnings.filterwarnings("ignore")

# 공통 모듈이 있는 src 폴더
sys.path.insert(0, SRC_DIR)

# 이전 차시에서 만든 모듈
sys.path.append(os.path.join(SRC_DIR, "13"))
sys.path.append(os.path.join(SRC_DIR, "18"))
sys.path.append(os.path.join(SRC_DIR, "24"))
sys.path.append(os.path.join(SRC_DIR, "25"))

# 현재 30차시 모듈을 가장 먼저 찾도록 합니다.
sys.path.insert(0, CURRENT_DIR)

# ------------------------------------------------------------
# 필요한 모듈을 가져옵니다.
# ------------------------------------------------------------

import config

from prompts import REWRITE_PROMPT
from generator import _llm
from langchain_core.output_parsers import StrOutputParser

# ------------------------------------------------------------
# Rewrite Chain을 만듭니다.
#
# Prompt → LLM → 문자열
# ------------------------------------------------------------

rewrite_chain = REWRITE_PROMPT | _llm | StrOutputParser()

# 최대 재작성 횟수입니다.
MAX_REWRITE = getattr(config, "MAX_REWRITE", 2)

# ============================================================
# Rewrite Node
# 검색에 실패한 질문을 새로운 검색어로 바꿉니다.
# 직접 검색하지 않고 query만 새로 만듭니다.
# ============================================================

def rewrite_node(state) -> dict:

    # 원본 사용자 질문을 가져옵니다.
    question = state.get("question", "")

    # 지금까지 사용한 검색어를 가져옵니다.
    tried = state.get("tried_queries", [])

    # 현재까지 재작성한 횟수를 가져옵니다.
    rewrites = state.get("rewrites", 0)

    # 검색 또는 검증 단계의 실패 이유를 가져옵니다.
    reason = state.get("reason")

    if not reason:
        reason = state.get("fail_reason")

    if not reason:
        reason = "관련 자료를 찾지 못함"

    # --------------------------------------------------------
    # 최대 재작성 횟수에 도달하면 중단합니다.
    # --------------------------------------------------------

    if rewrites >= MAX_REWRITE:

        return {
            "rewrite_ok": False,
            "log": [f"재작성 상한({MAX_REWRITE}) 도달"],
        }

    # --------------------------------------------------------
    # 이미 사용한 검색어 목록을 문자열로 만듭니다.
    # --------------------------------------------------------

    if tried:
        tried_text = ", ".join(tried)
    else:
        tried_text = "(없음)"

    # --------------------------------------------------------
    # LLM에게 새로운 검색어를 요청합니다.
    # --------------------------------------------------------

    try:

        new_query = rewrite_chain.invoke({
            "question": question,
            "tried": tried_text,
            "reason": reason,
        })

        # 앞뒤 공백과 따옴표를 제거합니다.
        new_query = new_query.strip()
        new_query = new_query.strip('"')
        new_query = new_query.strip("'")

    # LLM 호출 중 오류가 발생한 경우입니다.
    except Exception as e:

        return {
            "rewrite_ok": False,
            "rewrites": rewrites + 1,
            "log": [f"재작성 오류({type(e).__name__})"],
        }

    # --------------------------------------------------------
    # 새로운 검색어가 비어 있으면 실패입니다.
    # --------------------------------------------------------

    if not new_query:

        return {
            "rewrite_ok": False,
            "rewrites": rewrites + 1,
            "log": ["재작성 실패(빈값)"],
        }

    # --------------------------------------------------------
    # 이미 사용했던 검색어라면 다시 검색하지 않습니다.
    # --------------------------------------------------------

    if new_query in tried:

        return {
            "rewrite_ok": False,
            "rewrites": rewrites + 1,
            "log": [f"재작성 실패(중복): {new_query}"],
        }

    # --------------------------------------------------------
    # 새로운 검색어 생성에 성공했습니다.
    # --------------------------------------------------------

    return {
        # Retriever가 다음 검색에서 사용할 검색어입니다.
        "query": new_query,

        # 재작성 횟수를 1 증가시킵니다.
        "rewrites": rewrites + 1,

        # 새로운 검색어를 검색 이력에 추가합니다.
        "tried_queries": [new_query],

        # 새로운 query는 아직 검색하지 않은 상태입니다.
        "retrieval_ok": False,

        # Rewrite 성공 여부입니다.
        "rewrite_ok": True,

        # 실행 기록을 남깁니다.
        "log": [
            f"재작성 {rewrites + 1}회: "
            f"'{question[:20]}...' -> '{new_query}'"
        ],
    }