diff --git a/.env.example b/.env.example index 58a43a4..02e0b30 100644 --- a/.env.example +++ b/.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,10 +72,9 @@ 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' ############################################### @@ -74,12 +82,12 @@ DATAHUB_SERVER = 'http://-.-.-.-:-' ############################################### # 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 diff --git a/README.md b/README.md index 14f7ca7..d6aed28 100644 --- a/README.md +++ b/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 패키지 설치로 프로젝트 시작이 어려운 상황입니다. diff --git a/cli/__init__.py b/cli/__init__.py index 13347f3..23fd59e 100644 --- a/cli/__init__.py +++ b/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, diff --git a/llm_utils/retrieval.py b/llm_utils/retrieval.py index 0fa78a8..620b6a1 100644 --- a/llm_utils/retrieval.py +++ b/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"): diff --git a/llm_utils/vectordb/__init__.py b/llm_utils/vectordb/__init__.py new file mode 100644 index 0000000..31674ad --- /dev/null +++ b/llm_utils/vectordb/__init__.py @@ -0,0 +1,7 @@ +""" +VectorDB 모듈 - FAISS와 pgvector를 지원하는 벡터 데이터베이스 추상화 +""" + +from .factory import get_vector_db + +__all__ = ["get_vector_db"] diff --git a/llm_utils/vectordb/factory.py b/llm_utils/vectordb/factory.py new file mode 100644 index 0000000..b09ae1d --- /dev/null +++ b/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'를 사용하세요." + ) diff --git a/llm_utils/vectordb/faiss_db.py b/llm_utils/vectordb/faiss_db.py new file mode 100644 index 0000000..f79ed6a --- /dev/null +++ b/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 diff --git a/llm_utils/vectordb/pgvector_db.py b/llm_utils/vectordb/pgvector_db.py new file mode 100644 index 0000000..ade634a --- /dev/null +++ b/llm_utils/vectordb/pgvector_db.py @@ -0,0 +1,81 @@ +""" +pgvector VectorDB 구현 +""" + +import os +from typing import Optional +import psycopg2 +from sqlalchemy.orm import Session +from langchain_postgres.vectorstores import PGVector + +from llm_utils.tools import get_info_from_db +from llm_utils.llm_factory import get_embeddings + + +def _check_collection_exists(connection_string: str, collection_name: str) -> bool: + """PostgreSQL에서 collection이 존재하는지 확인합니다.""" + try: + # 연결 문자열에서 연결 정보 추출 + conn = psycopg2.connect(connection_string) + cursor = conn.cursor() + + # langchain_pg_embedding 테이블에서 collection_name이 존재하는지 확인 + cursor.execute( + "SELECT COUNT(*) FROM langchain_pg_embedding WHERE collection_name = %s", + (collection_name,), + ) + result = cursor.fetchone() + count = result[0] if result else 0 + + cursor.close() + conn.close() + + return count > 0 + except Exception as e: + print(f"Collection 존재 여부 확인 중 오류: {e}") + return False + + +def get_pgvector_db( + connection_string: Optional[str] = None, collection_name: Optional[str] = None +): + """pgvector 벡터 데이터베이스를 로드하거나 생성합니다.""" + embeddings = get_embeddings() + + if connection_string is None: + # 환경 변수에서 연결 정보 읽기 (기존 방식) + host = os.getenv("PGVECTOR_HOST", "localhost") + port = os.getenv("PGVECTOR_PORT", "5432") + user = os.getenv("PGVECTOR_USER", "postgres") + password = os.getenv("PGVECTOR_PASSWORD", "postgres") + database = os.getenv("PGVECTOR_DATABASE", "postgres") + connection_string = f"postgresql://{user}:{password}@{host}:{port}/{database}" + + if collection_name is None: + collection_name = os.getenv("PGVECTOR_COLLECTION", "lang2sql_table_info_db") + try: + vector_store = PGVector( + embeddings=embeddings, + collection_name=collection_name, + connection=connection_string, + ) + + results = vector_store.similarity_search("test", k=1) + if not results: + raise RuntimeError(f"Collection '{collection_name}' is empty") + + # 컬렉션이 존재하면 실제 검색도 진행해 볼 수 있습니다. + vector_store.similarity_search("test", k=1) + return vector_store + + except Exception as e: + print(f"exception: {e}") + # 컬렉션이 없거나 불러오기에 실패한 경우, 문서를 다시 인덱싱 + documents = get_info_from_db() + vector_store = PGVector.from_documents( + documents=documents, + embedding=embeddings, + connection=connection_string, + collection_name=collection_name, + ) + return vector_store diff --git a/pgvector.sh b/pgvector.sh new file mode 100644 index 0000000..6d122f4 --- /dev/null +++ b/pgvector.sh @@ -0,0 +1,5 @@ +docker run -d \ + --name pgvector \ + -e POSTGRES_PASSWORD=postgres \ + -p 5431:5432 \ + pgvector/pgvector:pg17 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index aa4c963..614613c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -33,4 +33,6 @@ psycopg2>=2.9.10,<3.0.0 pyodbc>=5.1.0,<6.0.0 crate>=0.29.0,<1.0.0 pyhive>=0.6.6,<1.0.0 -google-cloud-bigquery>=3.20.1,<4.0.0 \ No newline at end of file +google-cloud-bigquery>=3.20.1,<4.0.0 +pgvector==0.3.6 +langchain-postgres==0.0.15 \ No newline at end of file