Refactor query handling in llm_utils

- Add Reranker
- Added new dependencies: langchain-huggingface==0.1.2 and transformers==4.51.2 to requirements.txt.
- Removed the QueryRefinedAgainChain and its associated logic from chains.py and graph.py to streamline the query refinement process.

Author: ehddnr301

llm_utils/chains.py

@@ -40,23 +40,28 @@ def create_query_refiner_chain(llm):
                 예시:
                 사용자가 "유저 이탈 원인이 궁금해요"라고 했다면,
                 재질문 형식이 아니라
-                "최근 1개월 간의 접속·결제 로그를 기준으로,
+                "접속·결제 로그를 기준으로,
                 주로 어떤 사용자가 어떤 과정을 거쳐 이탈하는지를 분석해야 한다"처럼
                 분석 방향이 명확해진 질문 한 문장(또는 한 문단)으로 정리해 주세요.
 
                 최종 출력 형식 예시:
                 ------------------------------
                 구체화된 질문:
-                "최근 1개월 동안 고액 결제 경험이 있는 유저가 
+                "고액 결제 경험이 있는 유저가 
                 행동 로그에서 이탈 전 어떤 패턴을 보였는지 분석"
 
                 가정한 조건:
-                - 최근 1개월치 행동 로그와 결제 로그 중심
+                - 행동 로그와 결제 로그 중심
                 - 고액 결제자(월 결제액 10만 원 이상) 그룹 대상으로 한정
                 ------------------------------
                 """,
             ),
             MessagesPlaceholder(variable_name="user_input"),
+            (
+                "system",
+                "다음은 사용자의 실제 사용 가능한 테이블 및 컬럼 정보입니다:",
+            ),
+            MessagesPlaceholder(variable_name="searched_tables"),
             (
                 "system",
                 """
@@ -72,61 +77,6 @@ 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(
@@ -165,7 +115,6 @@ def create_query_maker_chain(llm):
             ),
             MessagesPlaceholder(variable_name="user_input"),
             MessagesPlaceholder(variable_name="refined_input"),
-            MessagesPlaceholder(variable_name="refined_input_again"),
             (
                 "system",
                 "다음은 사용자의 db 환경정보와 사용 가능한 테이블 및 컬럼 정보입니다:",
@@ -182,5 +131,4 @@ 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,15 +10,13 @@
 
 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"
@@ -32,7 +30,6 @@ class QueryMakerState(TypedDict):
     searched_tables: dict[str, dict[str, str]]
     best_practice_query: str
     refined_input: str
-    refined_input_again: str
     generated_query: str
 
 
@@ -43,6 +40,7 @@ def query_refiner_node(state: QueryMakerState):
             "user_input": [state["messages"][0].content],
             "user_database_env": [state["user_database_env"]],
             "best_practice_query": [state["best_practice_query"]],
+            "searched_tables": [json.dumps(state["searched_tables"])],
         }
     )
     state["messages"].append(res)
@@ -66,9 +64,42 @@ def get_table_info_node(state: QueryMakerState):
         db = FAISS.from_documents(documents, embeddings)
         db.save_local(os.getcwd() + "/table_info_db")
         print("table_info_db not found")
-    doc_res = db.similarity_search(state["messages"][-1].content)
-    documents_dict = {}
 
+    retriever = db.as_retriever(search_kwargs={"k": 10})
+
+    from langchain.retrievers import ContextualCompressionRetriever
+    from langchain.retrievers.document_compressors import CrossEncoderReranker
+    from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+    from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+    # Reranking 적용 여부 설정
+    use_rerank = True  # 필요에 따라 True 또는 False로 설정
+
+    if use_rerank:
+        local_model_path = os.path.join(os.getcwd(), "ko_reranker_local")
+
+        # 로컬에 저장된 모델이 있으면 불러오고, 없으면 다운로드 후 저장
+        if os.path.exists(local_model_path) and os.path.isdir(local_model_path):
+            print("🔄 ko-reranker 모델 로컬에서 로드 중...")
+        else:
+            print("⬇️ ko-reranker 모델 다운로드 및 저장 중...")
+            model = AutoModelForSequenceClassification.from_pretrained(
+                "Dongjin-kr/ko-reranker"
+            )
+            tokenizer = AutoTokenizer.from_pretrained("Dongjin-kr/ko-reranker")
+            model.save_pretrained(local_model_path)
+            tokenizer.save_pretrained(local_model_path)
+        model = HuggingFaceCrossEncoder(model_name=local_model_path)
+        compressor = CrossEncoderReranker(model=model, top_n=3)
+        retriever = db.as_retriever(search_kwargs={"k": 10})
+        compression_retriever = ContextualCompressionRetriever(
+            base_compressor=compressor, base_retriever=retriever
+        )
+
+        doc_res = compression_retriever.invoke(state["messages"][0].content)
+    else:  # Reranking 미적용
+        doc_res = db.similarity_search(state["messages"][0].content, k=10)
+    documents_dict = {}
     for doc in doc_res:
         lines = doc.page_content.split("\n")
 
@@ -93,19 +124,6 @@ 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
-    return state
-
-
 # 노드 함수: QUERY_MAKER 노드
 def query_maker_node(state: QueryMakerState):
     res = query_maker_chain.invoke(
@@ -137,9 +155,7 @@ 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["refined_input_again"].content]
+                [state["messages"][0].content] + [state["refined_input"].content]
             ),
             "table_info": [json.dumps(state["searched_tables"])],
             "top_k": 10,
@@ -152,21 +168,16 @@ def query_maker_node_with_db_guide(state: QueryMakerState):
 
 # StateGraph 생성 및 구성
 builder = StateGraph(QueryMakerState)
-builder.set_entry_point(QUERY_REFINER)
+builder.set_entry_point(GET_TABLE_INFO)
 
 # 노드 추가
-builder.add_node(QUERY_REFINER, query_refiner_node)
 builder.add_node(GET_TABLE_INFO, get_table_info_node)
-# builder.add_node(QUERY_MAKER, query_maker_node)  #  query_maker_node_with_db_guide
-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_node(QUERY_REFINER, query_refiner_node)
+builder.add_node(QUERY_MAKER, query_maker_node_with_db_guide)
 
 # 기본 엣지 설정
-builder.add_edge(QUERY_REFINER, GET_TABLE_INFO)
-builder.add_edge(GET_TABLE_INFO, QUERY_REFINED_AGAIN)
-builder.add_edge(QUERY_REFINED_AGAIN, QUERY_MAKER)
+builder.add_edge(GET_TABLE_INFO, QUERY_REFINER)
+builder.add_edge(QUERY_REFINER, QUERY_MAKER)
 
 # QUERY_MAKER 노드 후 종료
 builder.add_edge(QUERY_MAKER, END)

requirements.txt

@@ -11,3 +11,5 @@ pre_commit==4.1.0
 setuptools
 wheel
 twine
+langchain-huggingface==0.1.2
+transformers==4.51.2