refactor: 쿼리 실행을 위한 공용 모듈 추가 및 기존 코드 리팩토링

ehddnr301 · ehddnr301 · commit d7e60e602408 · 2025-06-29T02:43:28.000Z
- llm_utils/query_executor.py 파일 추가: CLI와 Streamlit에서 공통으로 사용할 수 있는 쿼리 실행 함수 구현
- query_command 함수에서 그래프 실행 로직을 execute_query 함수로 변경
- lang2sql.py에서 그래프 실행 로직을 execute_query 함수로 통합하여 코드 간소화
diff --git a/cli/__init__.py b/cli/__init__.py
@@ -217,45 +217,32 @@ def query_command(
     """
 
     try:
-        if use_enriched_graph:
-            from llm_utils.graph_utils.enriched_graph import builder
-        else:
-            from llm_utils.graph_utils.basic_graph import builder
-        from llm_utils.llm_response_parser import LLMResponseParser
-        from langchain_core.messages import HumanMessage
-
-        logger.info("Processing query: %s", question)
-        logger.info("Using %s graph", "enriched" if use_enriched_graph else "basic")
-
-        # 그래프 컴파일 및 실행
-        graph = builder.compile()
-        res = graph.invoke(
-            input={
-                "messages": [HumanMessage(content=question)],
-                "user_database_env": database_env,
-                "best_practice_query": "",
-                "retriever_name": retriever_name,
-                "top_n": top_n,
-                "device": device,
-            }
+        from llm_utils.query_executor import execute_query, extract_sql_from_result
+
+        # 공용 함수를 사용하여 쿼리 실행
+        res = execute_query(
+            query=question,
+            database_env=database_env,
+            retriever_name=retriever_name,
+            top_n=top_n,
+            device=device,
+            use_enriched_graph=use_enriched_graph,
         )
 
         # SQL 추출 및 출력
-        generated_query = res.get("generated_query")
-        if generated_query:
-            query_text = (
-                generated_query.content
-                if hasattr(generated_query, "content")
-                else str(generated_query)
-            )
-            try:
-                sql = LLMResponseParser.extract_sql(query_text)
-                print(sql)
-            except ValueError:
-                logger.error("SQL을 추출할 수 없습니다.")
-                print(query_text)
+        sql = extract_sql_from_result(res)
+        if sql:
+            print(sql)
         else:
-            logger.error("생성된 쿼리가 없습니다.")
+            # SQL 추출 실패 시 원본 쿼리 텍스트 출력
+            generated_query = res.get("generated_query")
+            if generated_query:
+                query_text = (
+                    generated_query.content
+                    if hasattr(generated_query, "content")
+                    else str(generated_query)
+                )
+                print(query_text)
 
     except Exception as e:
         logger.error("쿼리 처리 중 오류 발생: %s", e)
diff --git a/interface/lang2sql.py b/interface/lang2sql.py
@@ -7,14 +7,15 @@
 
 import streamlit as st
 from langchain.chains.sql_database.prompt import SQL_PROMPTS
-from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.messages import AIMessage
 
 from llm_utils.connect_db import ConnectDB
 from llm_utils.display_chart import DisplayChart
-from llm_utils.graph_utils.enriched_graph import builder as enriched_builder
-from llm_utils.graph_utils.basic_graph import builder
+from llm_utils.query_executor import execute_query as execute_query_common
 from llm_utils.llm_response_parser import LLMResponseParser
 from llm_utils.token_utils import TokenUtils
+from llm_utils.graph_utils.enriched_graph import builder as enriched_builder
+from llm_utils.graph_utils.basic_graph import builder
 
 TITLE = "Lang2SQL"
 DEFAULT_QUERY = "고객 데이터를 기반으로 유니크한 유저 수를 카운트하는 쿼리"
@@ -40,9 +41,8 @@ def execute_query(
     """
     자연어 쿼리를 SQL로 변환하고 실행 결과를 반환하는 Lang2SQL 그래프 인터페이스 함수입니다.
 
-    이 함수는 Lang2SQL 파이프라인(graph)을 세션 상태에서 가져오거나 새로 컴파일한 뒤,
-    사용자의 자연어 질문을 SQL 쿼리로 변환하고 관련 메타데이터와 함께 결과를 반환합니다.
-    내부적으로 LangChain의 `graph.invoke` 메서드를 호출합니다.
+    이 함수는 공용 execute_query 함수를 호출하여 Lang2SQL 파이프라인을 실행합니다.
+    Streamlit 세션 상태를 활용하여 그래프를 재사용합니다.
 
     Args:
         query (str): 사용자가 입력한 자연어 기반 질문.
@@ -59,27 +59,16 @@ def execute_query(
             - "searched_tables": 참조된 테이블 목록 등 추가 정보
     """
 
-    graph = st.session_state.get("graph")
-    if graph is None:
-        graph_builder = (
-            enriched_builder if st.session_state.get("use_enriched") else builder
-        )
-        graph = graph_builder.compile()
-        st.session_state["graph"] = graph
-
-    res = graph.invoke(
-        input={
-            "messages": [HumanMessage(content=query)],
-            "user_database_env": database_env,
-            "best_practice_query": "",
-            "retriever_name": retriever_name,
-            "top_n": top_n,
-            "device": device,
-        }
+    return execute_query_common(
+        query=query,
+        database_env=database_env,
+        retriever_name=retriever_name,
+        top_n=top_n,
+        device=device,
+        use_enriched_graph=st.session_state.get("use_enriched", False),
+        session_state=st.session_state,
     )
 
-    return res
-
 
 def display_result(
     *,
@@ -120,40 +109,50 @@ def should_show(_key: str) -> bool:
     if should_show("show_sql"):
         st.markdown("---")
         generated_query = res.get("generated_query")
-        query_text = (
-            generated_query.content
-            if isinstance(generated_query, AIMessage)
-            else str(generated_query)
-        )
+        if generated_query:
+            query_text = (
+                generated_query.content
+                if isinstance(generated_query, AIMessage)
+                else str(generated_query)
+            )
 
-        try:
-            sql = LLMResponseParser.extract_sql(query_text)
-            st.markdown("**생성된 SQL 쿼리:**")
-            st.code(sql, language="sql")
-        except ValueError:
-            st.warning("SQL 블록을 추출할 수 없습니다.")
-            st.text(query_text)
-
-        interpretation = LLMResponseParser.extract_interpretation(query_text)
-        if interpretation:
-            st.markdown("**결과 해석:**")
-            st.code(interpretation)
+            # query_text가 문자열인지 확인
+            if isinstance(query_text, str):
+                try:
+                    sql = LLMResponseParser.extract_sql(query_text)
+                    st.markdown("**생성된 SQL 쿼리:**")
+                    st.code(sql, language="sql")
+                except ValueError:
+                    st.warning("SQL 블록을 추출할 수 없습니다.")
+                    st.text(query_text)
+
+                interpretation = LLMResponseParser.extract_interpretation(query_text)
+                if interpretation:
+                    st.markdown("**결과 해석:**")
+                    st.code(interpretation)
+            else:
+                st.warning("쿼리 텍스트가 문자열이 아닙니다.")
+                st.text(str(query_text))
 
     if should_show("show_result_description"):
         st.markdown("---")
         st.markdown("**결과 설명:**")
         result_message = res["messages"][-1].content
 
-        try:
-            sql = LLMResponseParser.extract_sql(result_message)
-            st.code(sql, language="sql")
-        except ValueError:
-            st.warning("SQL 블록을 추출할 수 없습니다.")
-            st.text(result_message)
-
-        interpretation = LLMResponseParser.extract_interpretation(result_message)
-        if interpretation:
-            st.code(interpretation, language="plaintext")
+        if isinstance(result_message, str):
+            try:
+                sql = LLMResponseParser.extract_sql(result_message)
+                st.code(sql, language="sql")
+            except ValueError:
+                st.warning("SQL 블록을 추출할 수 없습니다.")
+                st.text(result_message)
+
+            interpretation = LLMResponseParser.extract_interpretation(result_message)
+            if interpretation:
+                st.code(interpretation, language="plaintext")
+        else:
+            st.warning("결과 메시지가 문자열이 아닙니다.")
+            st.text(str(result_message))
 
     if should_show("show_question_reinterpreted_by_ai"):
         st.markdown("---")
@@ -173,26 +172,41 @@ def should_show(_key: str) -> bool:
                 if isinstance(res["generated_query"], AIMessage)
                 else str(res["generated_query"])
             )
-            sql = LLMResponseParser.extract_sql(sql_raw)
-            df = database.run_sql(sql)
-            st.dataframe(df.head(10) if len(df) > 10 else df)
+            if isinstance(sql_raw, str):
+                sql = LLMResponseParser.extract_sql(sql_raw)
+                df = database.run_sql(sql)
+                st.dataframe(df.head(10) if len(df) > 10 else df)
+            else:
+                st.error("SQL 원본이 문자열이 아닙니다.")
         except Exception as e:
             st.error(f"쿼리 실행 중 오류 발생: {e}")
 
     if should_show("show_chart"):
         st.markdown("---")
-        df = database.run_sql(sql)
-        st.markdown("**쿼리 결과 시각화:**")
-        display_code = DisplayChart(
-            question=res["refined_input"].content,
-            sql=sql,
-            df_metadata=f"Running df.dtypes gives:\n{df.dtypes}",
-        )
-        # plotly_code 변수도 따로 보관할 필요 없이 바로 그려도 됩니다
-        fig = display_code.get_plotly_figure(
-            plotly_code=display_code.generate_plotly_code(), df=df
-        )
-        st.plotly_chart(fig)
+        try:
+            sql_raw = (
+                res["generated_query"].content
+                if isinstance(res["generated_query"], AIMessage)
+                else str(res["generated_query"])
+            )
+            if isinstance(sql_raw, str):
+                sql = LLMResponseParser.extract_sql(sql_raw)
+                df = database.run_sql(sql)
+                st.markdown("**쿼리 결과 시각화:**")
+                display_code = DisplayChart(
+                    question=res["refined_input"].content,
+                    sql=sql,
+                    df_metadata=f"Running df.dtypes gives:\n{df.dtypes}",
+                )
+                # plotly_code 변수도 따로 보관할 필요 없이 바로 그려도 됩니다
+                fig = display_code.get_plotly_figure(
+                    plotly_code=display_code.generate_plotly_code(), df=df
+                )
+                st.plotly_chart(fig)
+            else:
+                st.error("SQL 원본이 문자열이 아닙니다.")
+        except Exception as e:
+            st.error(f"차트 생성 중 오류 발생: {e}")
 
 
 db = ConnectDB()
diff --git a/llm_utils/query_executor.py b/llm_utils/query_executor.py
@@ -0,0 +1,111 @@
+"""
+Lang2SQL 쿼리 실행을 위한 공용 모듈입니다.
+
+이 모듈은 CLI와 Streamlit 인터페이스에서 공통으로 사용할 수 있는
+쿼리 실행 함수를 제공합니다.
+"""
+
+import logging
+from typing import Dict, Any, Optional, Union
+
+from langchain_core.messages import HumanMessage
+
+from llm_utils.graph_utils.enriched_graph import builder as enriched_builder
+from llm_utils.graph_utils.basic_graph import builder as basic_builder
+from llm_utils.llm_response_parser import LLMResponseParser
+
+logger = logging.getLogger(__name__)
+
+
+def execute_query(
+    *,
+    query: str,
+    database_env: str,
+    retriever_name: str = "기본",
+    top_n: int = 5,
+    device: str = "cpu",
+    use_enriched_graph: bool = False,
+    session_state: Optional[Union[Dict[str, Any], Any]] = None,
+) -> Dict[str, Any]:
+    """
+    자연어 쿼리를 SQL로 변환하고 실행 결과를 반환하는 공용 함수입니다.
+
+    이 함수는 Lang2SQL 파이프라인(graph)을 사용하여 사용자의 자연어 질문을
+    SQL 쿼리로 변환하고 관련 메타데이터와 함께 결과를 반환합니다.
+    CLI와 Streamlit 인터페이스에서 공통으로 사용할 수 있습니다.
+
+    Args:
+        query (str): 사용자가 입력한 자연어 기반 질문.
+        database_env (str): 사용할 데이터베이스 환경 이름 또는 키 (예: "dev", "prod").
+        retriever_name (str, optional): 테이블 검색기 이름. 기본값은 "기본".
+        top_n (int, optional): 검색된 상위 테이블 수 제한. 기본값은 5.
+        device (str, optional): LLM 실행에 사용할 디바이스 ("cpu" 또는 "cuda"). 기본값은 "cpu".
+        use_enriched_graph (bool, optional): 확장된 그래프 사용 여부. 기본값은 False.
+        session_state (Optional[Union[Dict[str, Any], Any]], optional): Streamlit 세션 상태 (Streamlit에서만 사용).
+
+    Returns:
+        Dict[str, Any]: 다음 정보를 포함한 Lang2SQL 실행 결과 딕셔너리:
+            - "generated_query": 생성된 SQL 쿼리 (`AIMessage`)
+            - "messages": 전체 LLM 응답 메시지 목록
+            - "refined_input": AI가 재구성한 입력 질문
+            - "searched_tables": 참조된 테이블 목록 등 추가 정보
+    """
+
+    logger.info("Processing query: %s", query)
+    logger.info("Using %s graph", "enriched" if use_enriched_graph else "basic")
+
+    # 그래프 선택 및 컴파일
+    if session_state is not None:
+        # Streamlit 환경: 세션 상태에서 그래프 재사용
+        graph = session_state.get("graph")
+        if graph is None:
+            graph_builder = enriched_builder if use_enriched_graph else basic_builder
+            graph = graph_builder.compile()
+            session_state["graph"] = graph
+    else:
+        # CLI 환경: 매번 새로운 그래프 컴파일
+        graph_builder = enriched_builder if use_enriched_graph else basic_builder
+        graph = graph_builder.compile()
+
+    # 그래프 실행
+    res = graph.invoke(
+        input={
+            "messages": [HumanMessage(content=query)],
+            "user_database_env": database_env,
+            "best_practice_query": "",
+            "retriever_name": retriever_name,
+            "top_n": top_n,
+            "device": device,
+        }
+    )
+
+    return res
+
+
+def extract_sql_from_result(res: Dict[str, Any]) -> Optional[str]:
+    """
+    Lang2SQL 실행 결과에서 SQL 쿼리를 추출합니다.
+
+    Args:
+        res (Dict[str, Any]): execute_query 함수의 반환 결과
+
+    Returns:
+        Optional[str]: 추출된 SQL 쿼리 문자열. 추출 실패 시 None
+    """
+    generated_query = res.get("generated_query")
+    if not generated_query:
+        logger.error("생성된 쿼리가 없습니다.")
+        return None
+
+    query_text = (
+        generated_query.content
+        if hasattr(generated_query, "content")
+        else str(generated_query)
+    )
+
+    try:
+        sql = LLMResponseParser.extract_sql(query_text)
+        return sql
+    except ValueError:
+        logger.error("SQL을 추출할 수 없습니다.")
+        return None
