feat(enriched_graph) : 프로파일 추출, 컨텍스트 보강 노드를 추가한 새로운 그래프 파일을 추가

Author: seyoung4503

llm_utils/chains.py

@@ -142,6 +142,7 @@ def create_profile_extraction_chain(llm):
 
 query_refiner_chain = create_query_refiner_chain(llm)
 query_maker_chain = create_query_maker_chain(llm)
+profile_extraction_chain = create_profile_extraction_chain(llm)
 query_refiner_with_profile_chain = create_query_refiner_with_profile_chain(llm)
 
 if __name__ == "__main__":

llm_utils/enriched_graph.py

@@ -0,0 +1,41 @@
+import json
+
+from langgraph.graph import StateGraph, END
+from llm_utils.graph import (
+    QueryMakerState,
+    GET_TABLE_INFO,
+    PROFILE_EXTRACTION,
+    QUERY_REFINER,
+    CONTEXT_ENRICHMENT,
+    QUERY_MAKER,
+    get_table_info_node,
+    profile_extraction_node,
+    query_refiner_with_profile_node,
+    context_enrichment_node,
+    query_maker_node,
+)
+
+"""
+기본 워크플로우에 '프로파일 추출(PROFILE_EXTRACTION)'과 '컨텍스트 보강(CONTEXT_ENRICHMENT)'를 
+추가한 확장된 그래프입니다.   
+"""
+
+# StateGraph 생성 및 구성
+builder = StateGraph(QueryMakerState)
+builder.set_entry_point(GET_TABLE_INFO)
+
+# 노드 추가
+builder.add_node(GET_TABLE_INFO, get_table_info_node)
+builder.add_node(QUERY_REFINER, query_refiner_with_profile_node)
+builder.add_node(PROFILE_EXTRACTION, profile_extraction_node)
+builder.add_node(CONTEXT_ENRICHMENT, context_enrichment_node)
+builder.add_node(QUERY_MAKER, query_maker_node)
+
+# 기본 엣지 설정
+builder.add_edge(GET_TABLE_INFO, PROFILE_EXTRACTION)
+builder.add_edge(PROFILE_EXTRACTION, QUERY_REFINER)
+builder.add_edge(QUERY_REFINER, CONTEXT_ENRICHMENT)
+builder.add_edge(CONTEXT_ENRICHMENT, QUERY_MAKER)
+
+# QUERY_MAKER 노드 후 종료
+builder.add_edge(QUERY_MAKER, END)

llm_utils/graph.py

@@ -12,6 +12,7 @@
     query_refiner_chain,
     query_maker_chain,
     query_refiner_with_profile_chain,
+    profile_extraction_chain,
 )
 
 from llm_utils.tools import get_info_from_db
@@ -24,6 +25,8 @@
 TOOL = "tool"
 TABLE_FILTER = "table_filter"
 QUERY_MAKER = "query_maker"
+PROFILE_EXTRACTION = "profile_extraction"
+CONTEXT_ENRICHMENT = "context_enrichment"
 
 
 # 상태 타입 정의 (추가 상태 정보와 메시지들을 포함)
@@ -43,11 +46,10 @@ class QueryMakerState(TypedDict):
 # 노드 함수: PROFILE_EXTRACTION 노드
 def profile_extraction_node(state: QueryMakerState):
 
-    result = query_refiner_with_profile_chain.invoke(
-        {"question": state["messages"][0].content}
-    )
+    result = profile_extraction_chain.invoke({"question": state["messages"][0].content})
 
     state["question_profile"] = result
+    print("profile_extraction_node : ", result)
     return state
 
 
@@ -132,6 +134,7 @@ def context_enrichment_node(state: QueryMakerState):
     from langchain_core.messages import HumanMessage
 
     state["messages"].append(HumanMessage(content=enriched_text.content))
+    print("After context enrichment : ", enriched_text.content)
 
     return state