서브그래프와 레거시 개조 — 운영 중인 그래프를 안 깨고 확장하기

def create_export_graph(
    ignore_new_line_in_text=True,
    show_image_in_markdown=False,
    verbose=True,
    visualize=False,
):
    """문서 내보내기 그래프를 생성한다.

    Returns:
        CompiledGraph: 컴파일된 내보내기 그래프
    """
    export_workflow = StateGraph(ParseState)

    export_image = export.ExportImage(verbose=verbose)
    export_html = export.ExportHTML(
        ignore_new_line_in_text=ignore_new_line_in_text, verbose=verbose
    )
    export_markdown = export.ExportMarkdown(
        ignore_new_line_in_text=ignore_new_line_in_text,
        show_image=show_image_in_markdown,
        verbose=verbose,
    )
    export_table_csv = export.ExportTableCSV(verbose=verbose)

    export_workflow.add_node("export_image", export_image)
    export_workflow.add_node("export_html", export_html)
    export_workflow.add_node("export_markdown", export_markdown)
    export_workflow.add_node("export_table_to_csv", export_table_csv)

    export_workflow.add_edge("export_image", "export_html")
    export_workflow.add_edge("export_image", "export_markdown")
    export_workflow.add_edge("export_image", "export_table_to_csv")

    export_workflow.add_edge("export_html", END)
    export_workflow.add_edge("export_markdown", END)
    export_workflow.add_edge("export_table_to_csv", END)

    export_workflow.set_entry_point("export_image")

    export_graph = export_workflow.compile()
    if visualize:
        visualize_graph(export_graph)
    return export_graph

from langgraph.checkpoint.memory import MemorySaver

parent_workflow = StateGraph(ParseState)

# 서브그래프를 "노드"로 등록한다
parent_workflow.add_node("upstage_parser", upstage_parser_graph)
parent_workflow.add_node("export_output", export_graph)

parent_workflow.add_edge("upstage_parser", "export_output")
parent_workflow.set_entry_point("upstage_parser")

parent_graph = parent_workflow.compile(checkpointer=MemorySaver())

# xray 옵션으로 서브그래프 내부까지 펼쳐 본다
visualize_graph(parent_graph, xray=True)

# 자식 그래프는 자기만의 스키마를 갖는다
class TranslatorState(TypedDict):
    source_text: str
    target_lang: str
    translated_text: str


translator_graph = build_translator_graph()


# 부모(ParseState) <-> 자식(TranslatorState) 사이를 잇는 래퍼 노드
def translate_bridge(state: ParseState):
    results = []
    for element in state["elements_from_parser"]:
        # 1) 부모 State -> 자식 입력으로 변환 (입력 매핑)
        child_input = {
            "source_text": element["content"]["markdown"],
            "target_lang": "ko",
        }
        # 2) 자식 그래프 실행
        child_output = translator_graph.invoke(child_input)
        # 3) 자식 출력 -> 부모 State로 변환 (출력 매핑)
        element["content"]["markdown"] = child_output["translated_text"]
        results.append(element)

    return {"elements_from_parser": results}


parent_workflow.add_node("translate", translate_bridge)

class TranslatedText(BaseModel):
    translated_text: str = Field(description="The translated text of the given text")


prompt = PromptTemplate.from_template(
    """You are a translation expert. Translate the <given_text> into Korean.
[IMPORTANT] Keep the <given_text>'s markdown format.

###

<given_text>
{text}
</given_text>"""
)

llm = ChatOpenAI(model=MODEL_NAME, temperature=0).with_structured_output(TranslatedText)
chain = prompt | llm


def add_translation_module(state: ParseState):
    """상태의 텍스트 요소들을 한국어로 번역한다."""
    translated_elements = []
    # 번역이 필요한 카테고리만 선택한다
    for element in state["elements_from_parser"]:
        if element["category"] in [
            "paragraph", "index", "heading1", "header",
            "footer", "caption", "list", "footnote",
        ]:
            translated_elements.append(element)

    BATCH_SIZE = 50
    all_translated_results = []

    for i in range(0, len(translated_elements), BATCH_SIZE):
        batch = translated_elements[i : i + BATCH_SIZE]
        batch_data = [{"text": text["content"]["markdown"]} for text in batch]

        trial = 3
        while trial > 0:
            try:
                batch_results = chain.batch(batch_data)
                break
            except Exception as e:
                print(e)
                trial -= 1
                continue

        all_translated_results.extend(batch_results)

    # 번역 결과를 원본 요소에 반영한다
    for i, result in enumerate(all_translated_results):
        translated_elements[i]["content"]["markdown"] = result.translated_text

    return {"elements_from_parser": state["elements_from_parser"]}

def contextualize_text(state: ParseState):
    """페이지별로 요소를 그룹화하고 배치 처리해 문맥을 추가한다."""
    elements_by_page = {}
    for element in state["elements_from_parser"]:
        if element["category"] in [
            "paragraph", "index", "heading1", "header",
            "footer", "caption", "list", "footnote",
        ]:
            page = element["page"]
            elements_by_page.setdefault(page, []).append(element)

    BATCH_SIZE = 10
    for page, elements in elements_by_page.items():
        for i in range(0, len(elements), BATCH_SIZE):
            batch = elements[i : i + BATCH_SIZE]
            # 배치의 모든 텍스트를 배경 정보로 결합한다
            background = " ".join([elem["content"]["text"] for elem in batch])
            batch_data = [
                {
                    "text": elem["content"]["markdown"],
                    "background_information": background,
                }
                for elem in batch
            ]

            contextualized_results = chain.batch(batch_data)
            for elem, result in zip(batch, contextualized_results):
                elem["content"]["markdown"] = result.contextualized_text

    return {"elements_from_parser": state["elements_from_parser"]}

# 앞단만 실행해 State 확보
parser_graph.invoke(inputs, config=config)
previous_state = parser_graph.get_state(config).values

# 신규 모듈을 함수로 직접 호출
translated_state = add_translation_module(previous_state)
translated_state["elements_from_parser"][-30:-25]

# 대안 모듈도 같은 방식으로 검증
contextualized_state = contextualize_text(previous_state)
contextualized_state["elements_from_parser"][-30:-25]

parent_workflow = StateGraph(ParseState)

parent_workflow.add_node("upstage_parser", upstage_parser_graph)   # PDF 파싱
parent_workflow.add_node("add_translation", add_translation_module)  # 번역
# parent_workflow.add_node("contextualize_text", contextualize_text)  # 문맥화 (비활성)
parent_workflow.add_node("export_output", export_graph)            # 내보내기

parent_workflow.add_edge("upstage_parser", "add_translation")
# parent_workflow.add_edge("upstage_parser", "contextualize_text")  # (비활성)
parent_workflow.add_edge("add_translation", "export_output")

parent_workflow.set_entry_point("upstage_parser")
parent_graph = parent_workflow.compile(checkpointer=MemorySaver())

visualize_graph(parent_graph, xray=True)