Merge branch 'master' into feature/77-add-isort-flake8

Author: nonegom

cli/__init__.py

@@ -1,3 +1,6 @@
+"""
+Datahub GMS 서버 URL을 설정하고, 필요 시 Streamlit 인터페이스를 실행하는 CLI 프로그램입니다.
+"""
 import subprocess
 
 import click
@@ -9,29 +12,123 @@
 @click.version_option(version="0.1.4")
 @click.pass_context
 @click.option(
-    "--datahub_server", default="http://localhost:8080", help="Datahub GMS 서버 URL"
+    "--datahub_server",
+    default="http://localhost:8080",
+    help=(
+        "Datahub GMS 서버의 URL을 설정합니다. "
+        "기본값은 'http://localhost:8080'이며, "
+        "운영 환경 또는 테스트 환경에 맞게 변경할 수 있습니다."
+    ),
 )
-@click.option("--run-streamlit", is_flag=True, help="Run the Streamlit app.")
-@click.option("-p", "--port", type=int, default=8501, help="Streamlit port")
-def cli(ctx, datahub_server, run_streamlit, port):
+@click.option(
+    "--run-streamlit",
+    is_flag=True,
+    help=(
+        "이 옵션을 지정하면 CLI 실행 시 Streamlit 애플리케이션을 바로 실행합니다. "
+        "별도의 명령어 입력 없이 웹 인터페이스를 띄우고 싶을 때 사용합니다."
+    ),
+)
+@click.option(
+    "-p",
+    "--port",
+    type=int,
+    default=8501,
+    help=(
+        "Streamlit 서버가 바인딩될 포트 번호를 지정합니다. "
+        "기본 포트는 8501이며, 포트 충돌을 피하거나 여러 인스턴스를 실행할 때 변경할 수 있습니다."
+    ),
+)
+# pylint: disable=redefined-outer-name
+def cli(
+    ctx: click.Context,
+    datahub_server: str,
+    run_streamlit: bool,
+    port: int,
+) -> None:
+    """
+    Datahub GMS 서버 URL을 설정하고, Streamlit 애플리케이션을 실행할 수 있는 CLI 명령 그룹입니다.
+
+    이 함수는 다음 역할을 수행합니다:
+    - 전달받은 'datahub_server' URL을 바탕으로 GMS 서버 연결을 설정합니다.
+    - 설정 과정 중 오류가 발생하면 오류 메시지를 출력하고 프로그램을 종료합니다.
+    - '--run-streamlit' 옵션이 활성화된 경우, 지정된 포트에서 Streamlit 웹 앱을 즉시 실행합니다.
+
+    매개변수:
+        ctx (click.Context): 명령어 실행 컨텍스트 객체입니다.
+        datahub_server (str): 설정할 Datahub GMS 서버의 URL입니다.
+        run_streamlit (bool): Streamlit 앱을 실행할지 여부를 나타내는 플래그입니다.
+        port (int): Streamlit 서버가 바인딩될 포트 번호입니다.
+
+    주의:
+        'set_gms_server' 함수에서 ValueError가 발생할 경우, 프로그램은 비정상 종료(exit code 1)합니다.
+    """
+
     try:
         set_gms_server(datahub_server)
     except ValueError as e:
-        click.echo(str(e))
+        click.secho(f"GMS 서버 URL 설정 실패: {str(e)}", fg="red")
         ctx.exit(1)
     if run_streamlit:
         run_streamlit_command(port)
 
 
-def run_streamlit_command(port):
-    """Run the Streamlit app."""
-    subprocess.run(
-        ["streamlit", "run", "interface/streamlit_app.py", "--server.port", str(port)]
-    )
+def run_streamlit_command(port: int) -> None:
+    """
+    지정된 포트에서 Streamlit 애플리케이션을 실행하는 함수입니다.
+
+    이 함수는 subprocess를 통해 'streamlit run' 명령어를 실행하여
+    'interface/streamlit_app.py' 파일을 웹 서버 형태로 구동합니다.
+    사용자가 지정한 포트 번호를 Streamlit 서버의 포트로 설정합니다.
+
+    매개변수:
+        port (int): Streamlit 서버가 바인딩될 포트 번호입니다.
+
+    주의:
+        - Streamlit이 시스템에 설치되어 있어야 정상 동작합니다.
+        - subprocess 호출 실패 시 예외가 발생할 수 있습니다.
+    """
+
+    try:
+        subprocess.run(
+            [
+                "streamlit",
+                "run",
+                "interface/streamlit_app.py",
+                "--server.port",
+                str(port),
+            ],
+            check=True,
+        )
+    except subprocess.CalledProcessError as e:
+        click.echo(f"Streamlit 실행 실패: {e}")
+        raise
+
+
+@cli.command(name="run-streamlit")
+@click.option(
+    "-p",
+    "--port",
+    type=int,
+    default=8501,
+    help=(
+        "Streamlit 애플리케이션이 바인딩될 포트 번호를 지정합니다. "
+        "기본 포트는 8501이며, 필요 시 포트 충돌을 피하거나 "
+        "여러 인스턴스를 동시에 실행할 때 다른 포트 번호를 설정할 수 있습니다."
+    ),
+)
+def run_streamlit_cli_command(port: int) -> None:
+    """
+    CLI 명령어를 통해 Streamlit 애플리케이션을 실행하는 함수입니다.
+
+    이 명령은 'interface/streamlit_app.py' 파일을 Streamlit 서버로 구동하며,
+    사용자가 지정한 포트 번호를 바인딩하여 웹 인터페이스를 제공합니다.
+
+    매개변수:
+        port (int): Streamlit 서버가 사용할 포트 번호입니다. 기본값은 8501입니다.
 
+    주의:
+        - Streamlit이 시스템에 설치되어 있어야 정상적으로 실행됩니다.
+        - Streamlit 실행에 실패할 경우 subprocess 호출에서 예외가 발생할 수 있습니다.
+    """
 
-@cli.command()
-@click.option("-p", "--port", type=int, default=8501, help="Streamlit port")
-def run_streamlit(port):
-    """Run the Streamlit app."""
     run_streamlit_command(port)

interface/lang2sql.py

@@ -1,80 +1,103 @@
-import os
-from typing import Union
+"""
+Lang2SQL Streamlit 애플리케이션.
+
+자연어로 입력된 질문을 SQL 쿼리로 변환하고,
+ClickHouse 데이터베이스에 실행한 결과를 출력합니다.
+"""
 
-import pandas as pd
 import streamlit as st
-from clickhouse_driver import Client
-from dotenv import load_dotenv
 from langchain.chains.sql_database.prompt import SQL_PROMPTS
 from langchain_core.messages import HumanMessage
 
 from llm_utils.connect_db import ConnectDB
 from llm_utils.graph import builder
 
-# Clickhouse 연결
-db = ConnectDB()
-db.connect_to_clickhouse()
+DEFAULT_QUERY = "고객 데이터를 기반으로 유니크한 유저 수를 카운트하는 쿼리"
+SIDEBAR_OPTIONS = {
+    "show_total_token_usage": "Show Total Token Usage",
+    "show_result_description": "Show Result Description",
+    "show_sql": "Show SQL",
+    "show_question_reinterpreted_by_ai": "Show User Question Reinterpreted by AI",
+    "show_referenced_tables": "Show List of Referenced Tables",
+    "show_table": "Show Table",
+    "show_chart": "Show Chart",
+}
+def summarize_total_tokens(data: list) -> int:
+    """
+    메시지 데이터에서 총 토큰 사용량을 집계합니다.
 
-# Streamlit 앱 제목
-st.title("Lang2SQL")
+    Args:
+        data (list): usage_metadata를 포함하는 객체들의 리스트.
 
-# 사용자 입력 받기
-user_query = st.text_area(
-    "쿼리를 입력하세요:",
-    value="고객 데이터를 기반으로 유니크한 유저 수를 카운트하는 쿼리",
-)
+    Returns:
+        int: 총 토큰 사용량 합계.
+    """
 
-user_database_env = st.selectbox(
-    "db 환경정보를 입력하세요:",
-    options=SQL_PROMPTS.keys(),
-    index=0,
-)
-st.sidebar.title("Output Settings")
-st.sidebar.checkbox("Show Total Token Usage", value=True, key="show_total_token_usage")
-st.sidebar.checkbox(
-    "Show Result Description", value=True, key="show_result_description"
-)
-st.sidebar.checkbox("Show SQL", value=True, key="show_sql")
-st.sidebar.checkbox(
-    "Show User Question Reinterpreted by AI",
-    value=True,
-    key="show_question_reinterpreted_by_ai",
-)
-st.sidebar.checkbox(
-    "Show List of Referenced Tables", value=True, key="show_referenced_tables"
-)
-st.sidebar.checkbox("Show Table", value=True, key="show_table")
-st.sidebar.checkbox("Show Chart", value=True, key="show_chart")
-
-
-# Token usage 집계 함수 정의
-def summarize_total_tokens(data):
     total_tokens = 0
     for item in data:
         token_usage = getattr(item, "usage_metadata", {})
         total_tokens += token_usage.get("total_tokens", 0)
     return total_tokens
 
 
-# 버튼 클릭 시 실행
-if st.button("쿼리 실행"):
-    # 그래프 컴파일 및 쿼리 실행
-    graph = builder.compile()
+def execute_query(
+    *,
+    query: str,
+    database_env: str,
+) -> dict:
+    """
+    Lang2SQL 그래프를 실행하여 자연어 쿼리를 SQL 쿼리로 변환하고 결과를 반환합니다.
+
+    Args:
+        query (str): 자연어로 작성된 사용자 쿼리.
+        database_env (str): 사용할 데이터베이스 환경 설정 이름.
+
+    Returns:
+        dict: 변환된 SQL 쿼리 및 관련 메타데이터를 포함하는 결과 딕셔너리.
+    """
+    # 세션 상태에서 그래프 가져오기
+    graph = st.session_state.get("graph")
+    if graph is None:
+        graph = builder.compile()
+        st.session_state["graph"] = graph
 
     res = graph.invoke(
         input={
-            "messages": [HumanMessage(content=user_query)],
-            "user_database_env": user_database_env,
+            "messages": [HumanMessage(content=query)],
+            "user_database_env": database_env,
             "best_practice_query": "",
         }
     )
+
+    return res
+
+
+def display_result(
+    *,
+    res: dict,
+    database: ConnectDB,
+) -> None:
+    """
+    Lang2SQL 실행 결과를 Streamlit 화면에 출력합니다.
+
+    Args:
+        res (dict): Lang2SQL 실행 결과 딕셔너리.
+        database (ConnectDB): SQL 쿼리 실행을 위한 데이터베이스 연결 객체.
+
+    출력 항목:
+        - 총 토큰 사용량
+        - 생성된 SQL 쿼리
+        - 결과 설명
+        - AI가 재해석한 사용자 질문
+        - 참조된 테이블 목록
+        - 쿼리 실행 결과 테이블
+    """
     total_tokens = summarize_total_tokens(res["messages"])
 
-    # 결과 출력
     if st.session_state.get("show_total_token_usage", True):
         st.write("총 토큰 사용량:", total_tokens)
     if st.session_state.get("show_sql", True):
-        st.write("결과:", "\n\n```sql\n" + res["generated_query"] + "\n```")
+        st.write("결과:", "\n\n```sql\n" + res["generated_query"].content + "\n```")
     if st.session_state.get("show_result_description", True):
         st.write("결과 설명:\n\n", res["messages"][-1].content)
     if st.session_state.get("show_question_reinterpreted_by_ai", True):
@@ -83,8 +106,42 @@ def summarize_total_tokens(data):
         st.write("참고한 테이블 목록:", res["searched_tables"])
     if st.session_state.get("show_table", True):
         sql = res["generated_query"]
-        df = db.run_sql(sql)
-        if len(df) > 10:
-            st.dataframe(df.head(10))
-        else:
-            st.dataframe(df)
+        df = database.run_sql(sql)
+        st.dataframe(df.head(10) if len(df) > 10 else df)
+
+
+db = ConnectDB()
+db.connect_to_clickhouse()
+
+st.title("Lang2SQL")
+
+# 세션 상태 초기화
+if "graph" not in st.session_state:
+    st.session_state["graph"] = builder.compile()
+    st.info("Lang2SQL이 성공적으로 시작되었습니다.")
+
+# 새로고침 버튼 추가
+if st.sidebar.button("Lang2SQL 새로고침"):
+    st.session_state["graph"] = builder.compile()
+    st.sidebar.success("Lang2SQL이 성공적으로 새로고침되었습니다.")
+
+user_query = st.text_area(
+    "쿼리를 입력하세요:",
+    value=DEFAULT_QUERY,
+)
+user_database_env = st.selectbox(
+    "DB 환경정보를 입력하세요:",
+    options=SQL_PROMPTS.keys(),
+    index=0,
+)
+
+st.sidebar.title("Output Settings")
+for key, label in SIDEBAR_OPTIONS.items():
+    st.sidebar.checkbox(label, value=True, key=key)
+
+if st.button("쿼리 실행"):
+    result = execute_query(
+        query=user_query,
+        database_env=user_database_env,
+    )
+    display_result(res=result, database=db)