Merge pull request #131 from #130

130 vectordb 유연성 개선 refactor

Author: ehddnr301

.env.example

@@ -1,20 +1,29 @@
 ###############################################
 ############## LLM API SELECTION ##############
 ###############################################
+LLM_PROVIDER=openai
+
+OPENAI_API_KEY=sk-proj-
+LANGCHAIN_TRACING_V2=true
+LANGCHAIN_PROJECT=langgraph_tutorial
+LANGCHAIN_ENDPOINT=https://api.smith.langchain.com
+LANGCHAIN_API_KEY=lsv2_
+
+
 
 # LLM_PROVIDER=openai
 # OPEN_AI_KEY=sk-proj-----
-# OPEN_AI_LLM_MODEL=gpt-4.1
+OPEN_AI_LLM_MODEL=gpt-4.1
 
 # LLM_PROVIDER=gemini
 # GEMINI_API_KEY=
 # GEMINI_LLM_MODEL=gemini-2.0-flash-lite
 
-LLM_PROVIDER=azure
-AZURE_OPENAI_LLM_ENDPOINT=https://-------.openai.azure.com/
-AZURE_OPENAI_LLM_KEY=-
-AZURE_OPENAI_LLM_MODEL=gpt4o
-AZURE_OPENAI_LLM_API_VERSION=2024-07-01-preview
+# LLM_PROVIDER=azure
+# AZURE_OPENAI_LLM_ENDPOINT=https://-------.openai.azure.com/
+# AZURE_OPENAI_LLM_KEY=-
+# AZURE_OPENAI_LLM_MODEL=gpt4o
+# AZURE_OPENAI_LLM_API_VERSION=2024-07-01-preview
 
 # LLM_PROVIDER=ollama
 # OLLAMA_LLM_BASE_URL=
@@ -36,8 +45,8 @@ AZURE_OPENAI_LLM_API_VERSION=2024-07-01-preview
 ########### Embedding API SElECTION ###########
 ###############################################
 # Only used if you are using an LLM that does not natively support embedding (openai or Azure)
-# EMBEDDING_PROVIDER='openai'
-# OPEN_AI_EMBEDDING_MODEL='text-embedding-ada-002'
+EMBEDDING_PROVIDER='openai'
+OPEN_AI_EMBEDDING_MODEL='text-embedding-ada-002'
 
 # EMBEDDING_PROVIDER=azure
 # AZURE_OPENAI_EMBEDDING_ENDPOINT=https://-------.openai.azure.com/openai/deployments
@@ -50,11 +59,11 @@ AZURE_OPENAI_LLM_API_VERSION=2024-07-01-preview
 # EMBEDDING_MODEL='nomic-embed-text:latest'
 # EMBEDDING_MODEL_MAX_CHUNK_LENGTH=8192
 
-EMBEDDING_PROVIDER='bedrock'
-AWS_BEDROCK_EMBEDDING_ACCESS_KEY_ID=--
-AWS_BEDROCK_EMBEDDING_SECRET_ACCESS_KEY=-/-+-+-
-AWS_BEDROCK_EMBEDDING_REGION=us-west-2
-AWS_BEDROCK_EMBEDDING_MODEL=amazon.titan-embed-text-v2:0
+# EMBEDDING_PROVIDER='bedrock'
+# AWS_BEDROCK_EMBEDDING_ACCESS_KEY_ID=--
+# AWS_BEDROCK_EMBEDDING_SECRET_ACCESS_KEY=-/-+-+-
+# AWS_BEDROCK_EMBEDDING_REGION=us-west-2
+# AWS_BEDROCK_EMBEDDING_MODEL=amazon.titan-embed-text-v2:0
 
 # EMBEDDING_PROVIDER='gemini'
 # GEMINI_EMBEDDING_API_KEY=
@@ -63,23 +72,22 @@ AWS_BEDROCK_EMBEDDING_MODEL=amazon.titan-embed-text-v2:0
 # EMBEDDING_PROVIDER='huggingface'
 # HUGGING_FACE_EMBEDDING_REPO_ID=
 # HUGGING_FACE_EMBEDDING_MODEL=
-
 # HUGGING_FACE_EMBEDDING_API_TOKEN=
 
-DATAHUB_SERVER = 'http://-.-.-.-:-'
+DATAHUB_SERVER = 'http://localhost:8080'
 
 
 ###############################################
 ######## Database Connector SELECTION #########
 ###############################################
 
 # clickhouse
-# DB_TYPE=clickhouse
-# CLICKHOUSE_HOST=_._._._
-# CLICKHOUSE_PORT=9000
-# CLICKHOUSE_USER=_
-# CLICKHOUSE_PASSWORD=_
-# CLICKHOUSE_DATABASE=_
+DB_TYPE=clickhouse
+CLICKHOUSE_HOST=localhost
+CLICKHOUSE_PORT=9001
+CLICKHOUSE_USER=clickhouse
+CLICKHOUSE_PASSWORD=clickhouse
+CLICKHOUSE_DATABASE=default
 
 # databricks
 # DB_TYPE=databricks
@@ -134,3 +142,14 @@ DATAHUB_SERVER = 'http://-.-.-.-:-'
 # DB_TYPE=sqlite
 # SQLITE_PATH=./data/sqlite.db
 
+
+# pgvector 설정 (VECTORDB_TYPE=pgvector일 때 사용)
+PGVECTOR_HOST=localhost
+PGVECTOR_PORT=5432
+PGVECTOR_USER=postgres
+PGVECTOR_PASSWORD=postgres
+PGVECTOR_DATABASE=postgres
+PGVECTOR_COLLECTION=table_info_db
+
+# VectorDB 설정
+VECTORDB_TYPE=faiss  # faiss 또는 pgvector

README.md

@@ -34,6 +34,7 @@ Lang2SQL은 자연어 쿼리를 최적화된 SQL 문으로 변환하는 오픈
 - **🔍 스키마 인식**: DataHub 메타데이터를 활용한 정확한 컬럼 매핑
 - **🛠️ 웹 인터페이스**: 대화형 Streamlit 앱을 통한 사용
 - **📈 시각화**: 생성된 SQL 쿼리 결과를 다양한 차트와 그래프로 시각화하여 데이터 인사이트를 직관적으로 파악
+- **🗄️ 유연한 VectorDB**: FAISS(로컬)와 pgvector(PostgreSQL) 중 선택 가능한 벡터 데이터베이스 지원
 
 ### 🤔 해결하는 문제
 
@@ -85,6 +86,28 @@ lang2sql run-streamlit
 lang2sql --datahub_server http://your-datahub-server:8080 run-streamlit -p 8888
 ```
 
+### VectorDB 선택
+
+FAISS(로컬) 또는 pgvector(PostgreSQL) 중 선택:
+
+```bash
+# FAISS 사용 (기본값)
+lang2sql --vectordb-type faiss run-streamlit
+
+# pgvector 사용
+lang2sql --vectordb-type pgvector run-streamlit
+```
+
+### 자연어 쿼리 실행
+
+```bash
+# 기본 FAISS 사용
+lang2sql query "고객 데이터를 기반으로 유니크한 유저 수를 카운트하는 쿼리"
+
+# pgvector 사용
+lang2sql query "고객 데이터를 기반으로 유니크한 유저 수를 카운트하는 쿼리" --vectordb-type pgvector --vectordb-location "postgresql://postgres:postgres@localhost:5432/postgres"
+```
+
 ### 환경 설정
 
 - 현재는 pip 패키지 설치로 프로젝트 시작이 어려운 상황입니다.

cli/__init__.py

@@ -64,14 +64,31 @@
     type=click.Path(exists=True, file_okay=False, dir_okay=True, readable=True),
     help="프롬프트 템플릿(.md 파일)이 저장된 디렉토리 경로를 지정합니다. 지정하지 않으면 기본 경로를 사용합니다.",
 )
+@click.option(
+    "--vectordb-type",
+    type=click.Choice(["faiss", "pgvector"]),
+    default="faiss",
+    help="사용할 벡터 데이터베이스 타입 (기본값: faiss)",
+)
+@click.option(
+    "--vectordb-location",
+    help=(
+        "VectorDB 위치 설정\n"
+        "- FAISS: 디렉토리 경로 (예: ./my_vectordb)\n"
+        "- pgvector: 연결 문자열 (예: postgresql://user:pass@host:port/db)\n"
+        "기본값: FAISS는 './table_info_db', pgvector는 환경변수 사용"
+    ),
+)
 # 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,
+    env_file_path: str | None = None,
+    prompt_dir_path: str | None = None,
+    vectordb_type: str = "faiss",
+    vectordb_location: str = None,
 ) -> None:
     """
     Datahub GMS 서버 URL을 설정하고, Streamlit 애플리케이션을 실행할 수 있는 CLI 명령 그룹입니다.
@@ -117,6 +134,23 @@ def cli(
             click.secho(f"프롬프트 디렉토리 환경변수 설정 실패: {str(e)}", fg="red")
             ctx.exit(1)
 
+    # VectorDB 타입을 환경 변수로 설정
+    try:
+        os.environ["VECTORDB_TYPE"] = vectordb_type
+        click.secho(f"VectorDB 타입 설정됨: {vectordb_type}", fg="green")
+    except Exception as e:
+        click.secho(f"VectorDB 타입 설정 실패: {str(e)}", fg="red")
+        ctx.exit(1)
+
+    # VectorDB 경로를 환경 변수로 설정
+    if vectordb_location:
+        try:
+            os.environ["VECTORDB_LOCATION"] = vectordb_location
+            click.secho(f"VectorDB 경로 설정됨: {vectordb_location}", fg="green")
+        except Exception as e:
+            click.secho(f"VectorDB 경로 설정 실패: {str(e)}", fg="red")
+            ctx.exit(1)
+
     logger.info(
         "Initialization started: GMS server = %s, run_streamlit = %s, port = %d",
         datahub_server,
@@ -129,7 +163,7 @@ def cli(
         logger.info("GMS server URL successfully set: %s", datahub_server)
     else:
         logger.error("GMS server health check failed. URL: %s", datahub_server)
-        ctx.exit(1)
+        # ctx.exit(1)
 
     if run_streamlit:
         run_streamlit_command(port)
@@ -234,6 +268,21 @@ def run_streamlit_cli_command(port: int) -> None:
     is_flag=True,
     help="단순화된 그래프(QUERY_REFINER 제거) 사용 여부",
 )
+@click.option(
+    "--vectordb-type",
+    type=click.Choice(["faiss", "pgvector"]),
+    default="faiss",
+    help="사용할 벡터 데이터베이스 타입 (기본값: faiss)",
+)
+@click.option(
+    "--vectordb-location",
+    help=(
+        "VectorDB 위치 설정\n"
+        "- FAISS: 디렉토리 경로 (예: ./my_vectordb)\n"
+        "- pgvector: 연결 문자열 (예: postgresql://user:pass@host:port/db)\n"
+        "기본값: FAISS는 './table_info_db', pgvector는 환경변수 사용"
+    ),
+)
 def query_command(
     question: str,
     database_env: str,
@@ -242,6 +291,8 @@ def query_command(
     device: str,
     use_enriched_graph: bool,
     use_simplified_graph: bool,
+    vectordb_type: str = "faiss",
+    vectordb_location: str = None,
 ) -> None:
     """
     자연어 질문을 SQL 쿼리로 변환하여 출력하는 명령어입니다.
@@ -260,11 +311,19 @@ def query_command(
     예시:
         lang2sql query "고객 데이터를 기반으로 유니크한 유저 수를 카운트하는 쿼리"
         lang2sql query "고객 데이터를 기반으로 유니크한 유저 수를 카운트하는 쿼리" --use-enriched-graph
+        lang2sql query "고객 데이터를 기반으로 유니크한 유저 수를 카운트하는 쿼리" --vectordb-type pgvector
     """
 
     try:
         from llm_utils.query_executor import execute_query, extract_sql_from_result
 
+        # VectorDB 타입을 환경 변수로 설정
+        os.environ["VECTORDB_TYPE"] = vectordb_type
+
+        # VectorDB 위치를 환경 변수로 설정
+        if vectordb_location:
+            os.environ["VECTORDB_LOCATION"] = vectordb_location
+
         # 공용 함수를 사용하여 쿼리 실행
         res = execute_query(
             query=question,

llm_utils/retrieval.py

@@ -8,23 +8,7 @@
 
 from .tools import get_info_from_db
 from .llm_factory import get_embeddings
-
-
-def get_vector_db():
-    """벡터 데이터베이스를 로드하거나 생성합니다."""
-    embeddings = get_embeddings()
-    try:
-        db = FAISS.load_local(
-            os.getcwd() + "/table_info_db",
-            embeddings,
-            allow_dangerous_deserialization=True,
-        )
-    except:
-        documents = get_info_from_db()
-        db = FAISS.from_documents(documents, embeddings)
-        db.save_local(os.getcwd() + "/table_info_db")
-        print("table_info_db not found")
-    return db
+from .vectordb import get_vector_db
 
 
 def load_reranker_model(device: str = "cpu"):

llm_utils/vectordb/__init__.py

@@ -0,0 +1,7 @@
+"""
+VectorDB 모듈 - FAISS와 pgvector를 지원하는 벡터 데이터베이스 추상화
+"""
+
+from .factory import get_vector_db
+
+__all__ = ["get_vector_db"]

llm_utils/vectordb/factory.py

@@ -0,0 +1,38 @@
+"""
+VectorDB 팩토리 모듈 - 환경 변수에 따라 적절한 VectorDB 인스턴스를 생성
+"""
+
+import os
+from typing import Optional
+
+from llm_utils.vectordb.faiss_db import get_faiss_vector_db
+from llm_utils.vectordb.pgvector_db import get_pgvector_db
+
+
+def get_vector_db(
+    vectordb_type: Optional[str] = None, vectordb_location: Optional[str] = None
+):
+    """
+    VectorDB 타입과 위치에 따라 적절한 VectorDB 인스턴스를 반환합니다.
+
+    Args:
+        vectordb_type: VectorDB 타입 ("faiss" 또는 "pgvector"). None인 경우 환경 변수에서 읽음.
+        vectordb_location: VectorDB 위치 (FAISS: 디렉토리 경로, pgvector: 연결 문자열). None인 경우 환경 변수에서 읽음.
+
+    Returns:
+        VectorDB 인스턴스 (FAISS 또는 PGVector)
+    """
+    if vectordb_type is None:
+        vectordb_type = os.getenv("VECTORDB_TYPE", "faiss").lower()
+
+    if vectordb_location is None:
+        vectordb_location = os.getenv("VECTORDB_LOCATION")
+
+    if vectordb_type == "faiss":
+        return get_faiss_vector_db(vectordb_location)
+    elif vectordb_type == "pgvector":
+        return get_pgvector_db(vectordb_location)
+    else:
+        raise ValueError(
+            f"지원하지 않는 VectorDB 타입: {vectordb_type}. 'faiss' 또는 'pgvector'를 사용하세요."
+        )

llm_utils/vectordb/faiss_db.py

@@ -0,0 +1,32 @@
+"""
+FAISS VectorDB 구현
+"""
+
+import os
+from langchain_community.vectorstores import FAISS
+from typing import Optional
+
+from llm_utils.tools import get_info_from_db
+from llm_utils.llm_factory import get_embeddings
+
+
+def get_faiss_vector_db(vectordb_path: Optional[str] = None):
+    """FAISS 벡터 데이터베이스를 로드하거나 생성합니다."""
+    embeddings = get_embeddings()
+
+    # 기본 경로 설정
+    if vectordb_path is None:
+        vectordb_path = os.path.join(os.getcwd(), "table_info_db")
+
+    try:
+        db = FAISS.load_local(
+            vectordb_path,
+            embeddings,
+            allow_dangerous_deserialization=True,
+        )
+    except:
+        documents = get_info_from_db()
+        db = FAISS.from_documents(documents, embeddings)
+        db.save_local(vectordb_path)
+        print(f"VectorDB를 새로 생성했습니다: {vectordb_path}")
+    return db