Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 40 additions & 0 deletions cli/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,11 @@
Datahub GMS 서버 URL을 설정하고, 필요 시 Streamlit 인터페이스를 실행하는 CLI 프로그램입니다.
"""

import os
import subprocess

import click
import dotenv

from llm_utils.tools import set_gms_server

Expand Down Expand Up @@ -39,12 +41,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 명령 그룹입니다.
Expand All @@ -53,17 +67,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)

try:
set_gms_server(datahub_server)
except ValueError as e:
Expand Down
7 changes: 0 additions & 7 deletions llm_utils/chains.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,6 @@
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()


Expand Down
3 changes: 0 additions & 3 deletions llm_utils/connect_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,6 @@

import pandas as pd
from clickhouse_driver import Client
from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(
level=logging.INFO,
Expand Down
4 changes: 0 additions & 4 deletions llm_utils/llm_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -20,9 +19,6 @@
)
from langchain_community.llms.bedrock import Bedrock

# .env 파일 로딩
load_dotenv()


def get_llm() -> BaseLanguageModel:
"""
Expand Down
16 changes: 11 additions & 5 deletions prompt/template_loader.py
Original file line number Diff line number Diff line change
@@ -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