Merge pull request #109 from #108

feat: enhance CLI with options for environment file and prompt directory

Author: ehddnr301

cli/__init__.py

@@ -5,10 +5,12 @@
 명령어 예시: lang2sql --datahub_server http://localhost:8080 --run-streamlit
 """
 
+import os
 import logging
 import subprocess
 
 import click
+import dotenv
 
 from llm_utils.check_server import CheckServer
 from llm_utils.tools import set_gms_server
@@ -52,12 +54,24 @@
         "기본 포트는 8501이며, 포트 충돌을 피하거나 여러 인스턴스를 실행할 때 변경할 수 있습니다."
     ),
 )
+@click.option(
+    "--env-file-path",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True),
+    help="환경 변수를 로드할 .env 파일의 경로를 지정합니다. 지정하지 않으면 기본 경로를 사용합니다.",
+)
+@click.option(
+    "--prompt-dir-path",
+    type=click.Path(exists=True, file_okay=False, dir_okay=True, readable=True),
+    help="프롬프트 템플릿(.md 파일)이 저장된 디렉토리 경로를 지정합니다. 지정하지 않으면 기본 경로를 사용합니다.",
+)
 # pylint: disable=redefined-outer-name
 def cli(
     ctx: click.Context,
     datahub_server: str,
     run_streamlit: bool,
     port: int,
+    env_file_path: str = None,
+    prompt_dir_path: str = None,
 ) -> None:
     """
     Datahub GMS 서버 URL을 설정하고, Streamlit 애플리케이션을 실행할 수 있는 CLI 명령 그룹입니다.
@@ -66,17 +80,43 @@ def cli(
     - 전달받은 'datahub_server' URL을 바탕으로 GMS 서버 연결을 설정합니다.
     - 설정 과정 중 오류가 발생하면 오류 메시지를 출력하고 프로그램을 종료합니다.
     - '--run-streamlit' 옵션이 활성화된 경우, 지정된 포트에서 Streamlit 웹 앱을 즉시 실행합니다.
+    - '--env-file-path' 옵션이 지정된 경우, 해당 .env 파일에서 환경 변수를 로드합니다.
+    - '--prompt-dir-path' 옵션이 지정된 경우, 해당 디렉토리에서 프롬프트 템플릿을 로드합니다.
 
     매개변수:
         ctx (click.Context): 명령어 실행 컨텍스트 객체입니다.
         datahub_server (str): 설정할 Datahub GMS 서버의 URL입니다.
         run_streamlit (bool): Streamlit 앱을 실행할지 여부를 나타내는 플래그입니다.
         port (int): Streamlit 서버가 바인딩될 포트 번호입니다.
+        env_file_path (str, optional): 환경 변수를 로드할 .env 파일 경로입니다.
+        prompt_dir_path (str, optional): 프롬프트 템플릿을 로드할 디렉토리 경로입니다.
 
     주의:
         'set_gms_server' 함수에서 ValueError가 발생할 경우, 프로그램은 비정상 종료(exit code 1)합니다.
     """
 
+    # 환경 변수 파일 로드
+    if env_file_path:
+        try:
+            if not dotenv.load_dotenv(env_file_path, override=True):
+                click.secho(f"환경 변수 파일 로드 실패: {env_file_path}", fg="yellow")
+            else:
+                click.secho(f"환경 변수 파일 로드 성공: {env_file_path}", fg="green")
+        except Exception as e:
+            click.secho(f"환경 변수 로드 중 오류 발생: {str(e)}", fg="red")
+            ctx.exit(1)
+
+    # 프롬프트 디렉토리를 환경 변수로 설정
+    if prompt_dir_path:
+        try:
+            os.environ["PROMPT_TEMPLATES_DIR"] = prompt_dir_path
+            click.secho(
+                f"프롬프트 디렉토리 환경변수 설정됨: {prompt_dir_path}", fg="green"
+            )
+        except Exception as e:
+            click.secho(f"프롬프트 디렉토리 환경변수 설정 실패: {str(e)}", fg="red")
+            ctx.exit(1)
+
     logger.info(
         "Initialization started: GMS server = %s, run_streamlit = %s, port = %d",
         datahub_server,

db_utils/__init__.py

@@ -3,8 +3,6 @@
 from .config import DBConfig
 from .logger import logger
 
-from dotenv import load_dotenv
-
 from .base_connector import BaseConnector
 
 from .clickhouse_connector import ClickHouseConnector
@@ -18,12 +16,6 @@
 
 env_path = os.path.join(os.getcwd(), ".env")
 
-if os.path.exists(env_path):
-    load_dotenv(env_path, override=True)
-    print(f"✅ 환경변수 파일(.env)이 {os.getcwd()}에 로드되었습니다!")
-else:
-    print(f"⚠️  환경변수 파일(.env)이 {os.getcwd()}에 없습니다!")
-
 
 def get_db_connector(db_type: Optional[str] = None, config: Optional[DBConfig] = None):
     """

llm_utils/chains.py

@@ -8,17 +8,9 @@
 
 from .llm_factory import get_llm
 
-from dotenv import load_dotenv
 from prompt.template_loader import get_prompt_template
 
 
-env_path = os.path.join(os.getcwd(), ".env")
-
-if os.path.exists(env_path):
-    load_dotenv(env_path)
-else:
-    print(f"⚠️  환경변수 파일(.env)이 {os.getcwd()}에 없습니다!")
-
 llm = get_llm()
 
 

llm_utils/connect_db.py

@@ -14,9 +14,6 @@
 
 import pandas as pd
 from clickhouse_driver import Client
-from dotenv import load_dotenv
-
-load_dotenv()
 
 logging.basicConfig(
     level=logging.INFO,

llm_utils/display_chart.py

@@ -1,9 +1,5 @@
 import re
-from llm_utils import llm_factory
-from dotenv import load_dotenv
-from langchain.chains.llm import LLMChain
 from langchain_openai import ChatOpenAI
-from langchain_core.prompts import PromptTemplate
 from langchain_core.messages import HumanMessage, SystemMessage
 import pandas as pd
 import os
@@ -13,10 +9,6 @@
 import plotly.graph_objects as go
 
 
-# .env 파일 로딩
-load_dotenv()
-
-
 class DisplayChart:
     """
     SQL쿼리가 실행된 결과를 그래프로 시각화하는 Class입니다.

llm_utils/llm_factory.py

@@ -2,7 +2,6 @@
 import os
 from typing import Optional
 
-from dotenv import load_dotenv
 from langchain.llms.base import BaseLanguageModel
 from langchain_aws import ChatBedrockConverse, BedrockEmbeddings
 from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
@@ -19,14 +18,6 @@
     OpenAIEmbeddings,
 )
 
-env_path = os.path.join(os.getcwd(), ".env")
-
-if os.path.exists(env_path):
-    load_dotenv(env_path, override=True)
-    print(f"✅ 환경변수 파일(.env)이 {os.getcwd()}에 로드되었습니다!")
-else:
-    print(f"⚠️  환경변수 파일(.env)이 {os.getcwd()}에 없습니다!")
-
 
 def get_llm(**kwargs) -> BaseLanguageModel:
     """

prompt/template_loader.py

@@ -1,14 +1,20 @@
+"""
+이 모듈은 프롬프트 템플릿을 로드하는 기능을 제공합니다.
+- 프롬프트 템플릿은 마크다운 파일로 관리되고 있으며, 환경변수에서 템플릿 디렉토리를 가져오거나, 없으면 현재 파일 위치 기준으로 설정합니다.
+"""
+
 import os
 
 
 def get_prompt_template(prompt_name: str) -> str:
+    # 환경변수에서 템플릿 디렉토리를 가져오거나, 없으면 현재 파일 위치 기준으로 설정
+    templates_dir = os.environ.get("PROMPT_TEMPLATES_DIR", os.path.dirname(__file__))
+
     try:
-        with open(
-            os.path.join(os.path.dirname(__file__), f"{prompt_name}.md"),
-            "r",
-            encoding="utf-8",
-        ) as f:
+        template_path = os.path.join(templates_dir, f"{prompt_name}.md")
+        with open(template_path, "r", encoding="utf-8") as f:
             template = f.read()
     except FileNotFoundError:
         raise FileNotFoundError(f"경고: '{prompt_name}.md' 파일을 찾을 수 없습니다.")
+
     return template