Skip to content

130 vectordb 유연성 개선 refactor #131

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 3 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
61 changes: 40 additions & 21 deletions .env.example
Original file line number Diff line number Diff line change
@@ -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=
Expand All @@ -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
Expand All @@ -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=
Expand All @@ -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
Expand Down Expand Up @@ -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
23 changes: 23 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ Lang2SQL은 자연어 쿼리를 최적화된 SQL 문으로 변환하는 오픈
- **🔍 스키마 인식**: DataHub 메타데이터를 활용한 정확한 컬럼 매핑
- **🛠️ 웹 인터페이스**: 대화형 Streamlit 앱을 통한 사용
- **📈 시각화**: 생성된 SQL 쿼리 결과를 다양한 차트와 그래프로 시각화하여 데이터 인사이트를 직관적으로 파악
- **🗄️ 유연한 VectorDB**: FAISS(로컬)와 pgvector(PostgreSQL) 중 선택 가능한 벡터 데이터베이스 지원

### 🤔 해결하는 문제

Expand Down Expand Up @@ -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 패키지 설치로 프로젝트 시작이 어려운 상황입니다.
Expand Down
65 changes: 62 additions & 3 deletions cli/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 명령 그룹입니다.
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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 쿼리로 변환하여 출력하는 명령어입니다.
Expand All @@ -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,
Expand Down
18 changes: 1 addition & 17 deletions llm_utils/retrieval.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand Down
7 changes: 7 additions & 0 deletions llm_utils/vectordb/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
"""
VectorDB 모듈 - FAISS와 pgvector를 지원하는 벡터 데이터베이스 추상화
"""

from .factory import get_vector_db

__all__ = ["get_vector_db"]
38 changes: 38 additions & 0 deletions llm_utils/vectordb/factory.py
Original file line number Diff line number Diff line change
@@ -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'를 사용하세요."
)
32 changes: 32 additions & 0 deletions llm_utils/vectordb/faiss_db.py
Original file line number Diff line number Diff line change
@@ -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
Loading