Implement QueryRefinedAgainChain for enhanced query refinement process

Author: ehddnr301

llm_utils/chains.py

@@ -72,6 +72,61 @@ def create_query_refiner_chain(llm):
     return tool_choice_prompt | llm
 
 
+# QueryRefinedAgainChain
+def create_query_redefined_again_chain(llm):
+    query_redefined_again_prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                """
+                당신은 데이터 분석 전문가(데이터 분석가 페르소나)입니다.
+                사용자의 질문과 이미 구체화된 질문을 바탕으로, 실제 사용 가능한 테이블과 컬럼 정보를 검토하여
+                더욱 정교하게 질문을 재정의해 주세요.
+                
+                주의사항:
+                - 이전에 구체화된 질문을 기반으로 하되, 실제 DB 환경에서 사용 가능한 테이블/컬럼을 고려해 현실적인 분석 방향을 제시하세요.
+                - 불필요한 재질문 없이, 주어진 데이터로 최대한 분석 가능한 형태로 질문을 구체화하세요.
+                - 테이블 구조에 맞게 분석 질문을 조정하고, 필요한 가정을 추가하세요.
+                - 최종 출력 형식은 반드시 아래와 같아야 합니다.
+                
+                최종 형태 예시:
+                
+                <최종 구체화된 질문>
+                ```
+                최근 30일간 결제 금액이 10만원 이상인 사용자들의 서비스 이용 패턴과 이탈율을 분석하여, 
+                어떤 활동 패턴을 보이는 고액 결제자가 이탈하는지 파악
+                ```
+                
+                <분석 접근 방향>
+                ```
+                1. subscription_activities와 contract_activities 테이블을 조인하여 고액 결제자 식별
+                2. 해당 사용자들의 activity_type 분포 확인
+                3. 이탈 사용자(30일 이상 미접속)와 활성 사용자의 행동 패턴 비교 분석
+                4. 주요 이탈 지점 식별
+                ```
+                """,
+            ),
+            (
+                "system",
+                "아래는 사용자의 원래 질문 및 1차 구체화된 질문입니다:",
+            ),
+            MessagesPlaceholder(variable_name="user_input"),
+            MessagesPlaceholder(variable_name="refined_input"),
+            (
+                "system",
+                "다음은 사용자의 DB 환경정보와 실제 사용 가능한 테이블 및 컬럼 정보입니다:",
+            ),
+            MessagesPlaceholder(variable_name="user_database_env"),
+            MessagesPlaceholder(variable_name="searched_tables"),
+            (
+                "system",
+                "위 정보를 바탕으로 DB 구조에 맞게 더욱 구체화된 최종 질문과 분석 접근 방향을 최종 형태 예시와 같은 형식으로 작성해주세요.",
+            ),
+        ]
+    )
+    return query_redefined_again_prompt | llm
+
+
 # QueryMakerChain
 def create_query_maker_chain(llm):
     query_maker_prompt = ChatPromptTemplate.from_messages(
@@ -110,6 +165,7 @@ def create_query_maker_chain(llm):
             ),
             MessagesPlaceholder(variable_name="user_input"),
             MessagesPlaceholder(variable_name="refined_input"),
+            MessagesPlaceholder(variable_name="refined_input_again"),
             (
                 "system",
                 "다음은 사용자의 db 환경정보와 사용 가능한 테이블 및 컬럼 정보입니다:",
@@ -126,4 +182,5 @@ def create_query_maker_chain(llm):
 
 
 query_refiner_chain = create_query_refiner_chain(llm)
+query_redefined_again_chain = create_query_redefined_again_chain(llm)
 query_maker_chain = create_query_maker_chain(llm)

llm_utils/graph.py

@@ -10,13 +10,15 @@
 
 from llm_utils.chains import (
     query_refiner_chain,
+    query_redefined_again_chain,
     query_maker_chain,
 )
 
 from llm_utils.tools import get_info_from_db
 
 # 노드 식별자 정의
 QUERY_REFINER = "query_refiner"
+QUERY_REFINED_AGAIN = "query_redefined_again"
 GET_TABLE_INFO = "get_table_info"
 TOOL = "tool"
 TABLE_FILTER = "table_filter"
@@ -30,6 +32,7 @@ class QueryMakerState(TypedDict):
     searched_tables: dict[str, dict[str, str]]
     best_practice_query: str
     refined_input: str
+    refined_input_again: str
     generated_query: str
 
 
@@ -90,6 +93,20 @@ def get_table_info_node(state: QueryMakerState):
     return state
 
 
+def query_redefined_again_node(state: QueryMakerState):
+    res = query_redefined_again_chain.invoke(
+        input={
+            "user_input": [state["messages"][0].content],
+            "refined_input": [state["refined_input"]],
+            "user_database_env": [state["user_database_env"]],
+            "searched_tables": [json.dumps(state["searched_tables"])],
+        }
+    )
+    state["refined_input_again"] = res
+    print(state["refined_input_again"])
+    return state
+
+
 # 노드 함수: QUERY_MAKER 노드
 def query_maker_node(state: QueryMakerState):
     res = query_maker_chain.invoke(
@@ -121,7 +138,9 @@ def query_maker_node_with_db_guide(state: QueryMakerState):
     res = chain.invoke(
         input={
             "input": "\n\n---\n\n".join(
-                [state["messages"][0].content] + [state["refined_input"].content]
+                [state["messages"][0].content]
+                # + [state["refined_input"].content]
+                + [state["refined_input_again"].content]
             ),
             "table_info": [json.dumps(state["searched_tables"])],
             "top_k": 10,
@@ -143,10 +162,12 @@ def query_maker_node_with_db_guide(state: QueryMakerState):
 builder.add_node(
     QUERY_MAKER, query_maker_node_with_db_guide
 )  #  query_maker_node_with_db_guide
+builder.add_node(QUERY_REFINED_AGAIN, query_redefined_again_node)
 
 # 기본 엣지 설정
 builder.add_edge(QUERY_REFINER, GET_TABLE_INFO)
-builder.add_edge(GET_TABLE_INFO, QUERY_MAKER)
+builder.add_edge(GET_TABLE_INFO, QUERY_REFINED_AGAIN)
+builder.add_edge(QUERY_REFINED_AGAIN, QUERY_MAKER)
 
 # QUERY_MAKER 노드 후 종료
 builder.add_edge(QUERY_MAKER, END)