diff --git a/agents/text/modules/nodes.py b/agents/text/modules/nodes.py index b793de5..d3271ee 100644 --- a/agents/text/modules/nodes.py +++ b/agents/text/modules/nodes.py @@ -99,3 +99,26 @@ def execute(self, state: TextState) -> dict: result = self.chain.invoke(input_data) state.update(result) return result + + +class TextRegenerateRouterNode(BaseNode): + """ + 텍스트 체크 결과를 기반으로 다음 노드를 결정하는 라우터 노드 + """ + + def execute(self, state: TextState) -> dict: + """ + 텍스트 체크 결과를 평가하여 다음 실행할 노드를 결정합니다. + + Args: + state (TextState): 현재 워크플로우 상태 + + Returns: + dict: 다음 노드 정보를 포함한 상태 업데이트 + """ + check_result = state.get("text_content_checker_result", {}) + next_node = ( + "__end__" if check_result.get("success", False) else "text_generation" + ) + + return {"next": next_node} diff --git a/agents/text/workflow.py b/agents/text/workflow.py index e8e9163..fea72bf 100644 --- a/agents/text/workflow.py +++ b/agents/text/workflow.py @@ -5,6 +5,8 @@ GenTextNode, PersonaExtractionNode, TextContentCheckNode, + TextRegenerateRouterNode +, ) from agents.text.modules.state import TextState @@ -42,22 +44,27 @@ def build(self): # 텍스트 컨텐츠 체커 노드 추가 builder.add_node("text_content_check", TextContentCheckNode()) + # 텍스트 재생성 노드 추가 + builder.add_node("text_regenerate_router", TextRegenerateRouterNode()) + # 시작 노드에서 페르소나 추출 노드로 연결 builder.add_edge("__start__", "persona_extraction") # 페르소나 추출 노드에서 텍스트 생성 노드로 연결 builder.add_edge("persona_extraction", "text_generation") # 텍스트 생성 노드에서 텍스트 컨텐츠 체커 노드로 연결 builder.add_edge("text_generation", "text_content_check") - # 텍스트 컨텐츠 체커 노드에서 종료 노드로 연결 - builder.add_edge("text_content_check", "__end__") - - # 조건부 에지 추가 예시 - # builder.add_conditional_edges( - # "call_model", - # # call_model 실행이 완료된 후, 다음 노드(들)는 - # # router의 출력을 기반으로 예약됩니다 - # router, - # ) + # 텍스트 컨텐츠 체커 노드에서 텍스트 재생성 노드로 연결 + builder.add_edge("text_content_check", "text_regenerate_router") + + # 조건부 에지 추가 + builder.add_conditional_edges( + "text_regenerate_router", + lambda state: state["next"], # state에서 next 키의 값을 사용 + { + "text_generation": "text_content_check", # 실패시 text_generation으로 돌아감 + "__end__": "__end__", # 성공시 종료 + }, + ) workflow = builder.compile() # 그래프 컴파일 workflow.name = self.name # Workflow 이름 설정