기존 get_info_from 함수 수정 (get_table_schema)

Author: ehddnr301

utils/llm/README.md

@@ -146,7 +146,7 @@ utils/llm/
 **목적**: DataHub 메타데이터 수집 및 LangGraph ChatBot용 Tool 함수 제공
 
 **주요 기능:**
-- `get_info_from_db()`: DataHub에서 테이블 메타데이터를 LangChain Document로 수집
+- `get_table_schema()`: DataHub에서 테이블 메타데이터를 dictionary 형태로 반환
 - `get_metadata_from_db()`: 전체 메타데이터 딕셔너리 반환
 - `search_database_tables()`: 벡터 검색 기반 테이블 정보 검색 Tool
 - `get_glossary_terms()`: 용어집 정보 조회 Tool
@@ -301,7 +301,7 @@ engine/query_executor.py
 │       └── utils/llm/retrieval.py
 │           └── utils/llm/vectordb/get_vector_db()
 │               ├── utils/llm/core/get_embeddings()
-│               └── utils/llm/tools/get_info_from_db()
+│               └── utils/llm/tools/get_table_schema()
 └── utils/llm/llm_response_parser.py
 ```
 
@@ -319,8 +319,8 @@ engine/query_executor.py
 - `retrieval.py` → `vectordb/get_vector_db()` 사용
 
 **vectordb 모듈:**
-- `vectordb/faiss_db.py` → `core/get_embeddings()`, `tools/get_info_from_db()` 사용
-- `vectordb/pgvector_db.py` → `core/get_embeddings()`, `tools/get_info_from_db()` 사용
+- `vectordb/faiss_db.py` → `core/get_embeddings()`, `tools/get_table_schema()` 사용
+- `vectordb/pgvector_db.py` → `core/get_embeddings()`, `tools/get_table_schema()` 사용
 
 **tools 모듈:**
 - `tools/datahub.py` → DataHub 메타데이터 수집

utils/llm/tools/README.md

@@ -21,7 +21,7 @@ utils/llm/tools/
 
 **datahub 모듈에서**:
 - `set_gms_server`: GMS 서버 설정
-- `get_info_from_db`: LangChain Document 리스트로 테이블/컬럼 정보 반환
+- `get_table_schema`: LangChain Document 리스트로 테이블/컬럼 정보 반환
 - `get_metadata_from_db`: 전체 메타데이터 딕셔너리 리스트 반환
 
 **chatbot_tool 모듈에서**:
@@ -39,7 +39,7 @@ utils/llm/tools/
    - 환경변수 `DATAHUB_SERVER`를 설정하고 DatahubMetadataFetcher 초기화
    - 유효하지 않은 서버 URL 시 ValueError 발생
 
-2. **`get_info_from_db(max_workers: int = 8) -> List[Document]`**
+2. **`get_table_schema(max_workers: int = 8) -> List[Document]`**
    - DataHub에서 모든 테이블 메타데이터를 수집하여 LangChain Document 리스트 반환
    - 각 Document에는 테이블명, 설명, 컬럼 정보가 포함
    - 형식: `"{테이블명}: {설명}\nColumns:\n {컬럼명}: {컬럼설명}"`
@@ -157,10 +157,10 @@ utils/llm/tools/
 #### 1. DataHub 메타데이터 수집 (vectorDB 초기화)
 
 ```python
-from utils.llm.tools import get_info_from_db
+from utils.llm.tools import get_table_schema
 
 # 모든 테이블 메타데이터를 LangChain Document로 수집
-documents = get_info_from_db(max_workers=8)
+documents = get_table_schema(max_workers=8)
 
 # 각 document는 다음과 같은 형식:
 # "테이블명: 설명\nColumns:\n 컬럼1: 설명1\n 컬럼2: 설명2"
@@ -224,8 +224,8 @@ queries = get_query_examples(
 
 **import하는 파일**:
 - `utils/llm/chatbot.py`: `from utils.llm.tools import search_database_tables, get_glossary_terms, get_query_examples`
-- `utils/llm/vectordb/faiss_db.py`: `from utils.llm.tools import get_info_from_db`
-- `utils/llm/vectordb/pgvector_db.py`: `from utils.llm.tools import get_info_from_db`
+- `utils/llm/vectordb/faiss_db.py`: `from utils.llm.tools import get_table_schema`
+- `utils/llm/vectordb/pgvector_db.py`: `from utils.llm.tools import get_table_schema`
 - `interface/core/config/settings.py`: `from utils.llm.tools import set_gms_server`
 
 **내부 의존성**:
@@ -258,7 +258,7 @@ queries = get_query_examples(
 
 #### 메타데이터 수집 흐름 (벡터DB 초기화 시)
 
-1. `get_info_from_db()` 호출
+1. `get_table_schema()` 호출
 2. `_get_fetcher()`로 DatahubMetadataFetcher 인스턴스 생성
 3. `parallel_process()`로 병렬 테이블 정보 수집
 4. 각 테이블별로 컬럼 정보 추가 수집

utils/llm/tools/__init__.py

@@ -1,5 +1,5 @@
 from utils.llm.tools.datahub import (
-    get_info_from_db,
+    get_table_schema,
     get_metadata_from_db,
     set_gms_server,
 )
@@ -12,7 +12,7 @@
 
 __all__ = [
     "set_gms_server",
-    "get_info_from_db",
+    "get_table_schema",
     "get_metadata_from_db",
     "search_database_tables",
     "get_glossary_terms",

utils/llm/tools/datahub.py

@@ -76,7 +76,7 @@ def _get_table_info(max_workers: int = 8) -> Dict[str, str]:
 
 
 def _get_column_info(
-    table_name: str, urn_table_mapping: Dict[str, str], max_workers: int = 8
+    table_name: str, urn_table_mapping: Dict[str, str]
 ) -> List[Dict[str, str]]:
     target_urn = urn_table_mapping.get(table_name)
     if not target_urn:
@@ -103,7 +103,21 @@ def _extract_dataset_name_from_urn(urn: str) -> Optional[str]:
     return None
 
 
-def get_info_from_db(max_workers: int = 8) -> List[Document]:
+def get_metadata_from_db() -> List[Dict]:
+    fetcher = _get_fetcher()
+    urns = list(fetcher.get_urns())
+
+    metadata = []
+    total = len(urns)
+    for idx, urn in enumerate(urns, 1):
+        print(f"[{idx}/{total}] Processing URN: {urn}")
+        table_metadata = fetcher.build_table_metadata(urn)
+        metadata.append(table_metadata)
+
+    return metadata
+
+
+def _prepare_datahub_metadata_mappings(max_workers: int = 8):
     table_info = _get_table_info(max_workers=max_workers)
 
     fetcher = _get_fetcher()
@@ -118,20 +132,31 @@ def get_info_from_db(max_workers: int = 8) -> List[Document]:
             if parsed_name:
                 display_name_by_table[original_name] = parsed_name
 
-    def process_table_info(item: tuple[str, str, str]) -> str:
-        original_table_name, table_description, display_table_name = item
-        # 컬럼 조회는 기존 테이블 이름으로 수행 (urn_table_mapping과 일치)
-        column_info = _get_column_info(
-            original_table_name, urn_table_mapping, max_workers=max_workers
-        )
-        column_info_str = "\n".join(
-            [
-                f"{col['column_name']}: {col['column_description']}"
-                for col in column_info
-            ]
-        )
-        used_name = display_table_name or original_table_name
-        return f"{used_name}: {table_description}\nColumns:\n {column_info_str}"
+    return table_info, urn_table_mapping, display_name_by_table
+
+
+def _format_datahub_table_info(
+    item: tuple[str, str, str], urn_table_mapping: Dict[str, str]
+) -> Dict:
+    original_table_name, table_description, display_table_name = item
+    # 컬럼 조회는 기존 테이블 이름으로 수행 (urn_table_mapping과 일치)
+    column_info = _get_column_info(original_table_name, urn_table_mapping)
+
+    columns = {col["column_name"]: col["column_description"] for col in column_info}
+
+    used_name = display_table_name or original_table_name
+    return {
+        used_name: {
+            "table_description": table_description,
+            "columns": columns,
+        }
+    }
+
+
+def get_table_schema(max_workers: int = 8) -> List[Dict]:
+    table_info, urn_table_mapping, display_name_by_table = (
+        _prepare_datahub_metadata_mappings(max_workers)
+    )
 
     # 표시용 이름을 세 번째 파라미터로 함께 전달
     items_with_display = [
@@ -143,25 +168,15 @@ def process_table_info(item: tuple[str, str, str]) -> str:
         for name, desc in table_info.items()
     ]
 
-    table_info_str_list = parallel_process(
+    # parallel_process에 전달할 함수 래핑
+    def process_fn(item):
+        return _format_datahub_table_info(item, urn_table_mapping)
+
+    table_info_list = parallel_process(
         items_with_display,
-        process_table_info,
+        process_fn,
         max_workers=max_workers,
         desc="컬럼 정보 수집 중",
     )
 
-    return [Document(page_content=info) for info in table_info_str_list]
-
-
-def get_metadata_from_db() -> List[Dict]:
-    fetcher = _get_fetcher()
-    urns = list(fetcher.get_urns())
-
-    metadata = []
-    total = len(urns)
-    for idx, urn in enumerate(urns, 1):
-        print(f"[{idx}/{total}] Processing URN: {urn}")
-        table_metadata = fetcher.build_table_metadata(urn)
-        metadata.append(table_metadata)
-
-    return metadata
+    return table_info_list

utils/llm/vectordb/README.md

@@ -52,13 +52,13 @@ utils/llm/vectordb/
    - `vectordb_path`: 저장 경로 (기본: `dev/table_info_db`)
    - 동작 방식:
      - 기존 DB가 있으면 `FAISS.load_local()`로 로드
-     - 없으면 `get_info_from_db()`로 문서 수집 후 `FAISS.from_documents()` 생성 및 저장
+     - 없으면 `get_table_schema()`로 문서 수집 후 `FAISS.from_documents()` 생성 및 저장
    - 반환: FAISS 벡터스토어 인스턴스
 
 **의존성**:
 - `langchain_community.vectorstores.FAISS`: LangChain FAISS 래퍼
 - `utils.llm.core.get_embeddings`: 임베딩 모델 로드
-- `utils.llm.tools.get_info_from_db`: DataHub에서 테이블 메타데이터 수집
+- `utils.llm.tools.get_table_schema`: DataHub에서 테이블 메타데이터 수집
 
 **특징**:
 - 로컬 디스크에 저장되어 네트워크 연결 불필요
@@ -84,7 +84,7 @@ utils/llm/vectordb/
      - `PGVECTOR_COLLECTION`: "lang2sql_table_info_db"
    - 동작 방식:
      - 기존 컬렉션이 있고 비어있지 않으면 로드
-     - 없거나 비어있으면 `get_info_from_db()`로 문서 수집 후 `PGVector.from_documents()` 생성
+     - 없거나 비어있으면 `get_table_schema()`로 문서 수집 후 `PGVector.from_documents()` 생성
    - 반환: PGVector 벡터스토어 인스턴스
 
 2. **`_check_collection_exists(connection_string, collection_name)`**
@@ -96,7 +96,7 @@ utils/llm/vectordb/
 - `langchain_postgres.vectorstores.PGVector`: LangChain pgvector 래퍼
 - `psycopg2`: PostgreSQL 연결
 - `utils.llm.core.get_embeddings`: 임베딩 모델 로드
-- `utils.llm.tools.get_info_from_db`: DataHub에서 테이블 메타데이터 수집
+- `utils.llm.tools.get_table_schema`: DataHub에서 테이블 메타데이터 수집
 
 **특징**:
 - PostgreSQL 데이터베이스에 저장되어 다중 서버 환경에 적합
@@ -181,7 +181,7 @@ export PGVECTOR_COLLECTION=lang2sql_table_info_db
 
 **내부 의존성**:
 - `utils/llm/core/factory.py`: `get_embeddings()` - 임베딩 모델 로드
-- `utils/llm/tools/datahub.py`: `get_info_from_db()` - DataHub 메타데이터 수집
+- `utils/llm/tools/datahub.py`: `get_table_schema()` - DataHub 메타데이터 수집
 
 **외부 의존성**:
 - `langchain_community.vectorstores.FAISS`: FAISS 벡터스토어

utils/llm/vectordb/faiss_db.py

@@ -6,9 +6,10 @@
 from typing import Optional
 
 from langchain_community.vectorstores import FAISS
+from langchain.schema import Document
 
 from utils.llm.core import get_embeddings
-from utils.llm.tools import get_info_from_db
+from utils.llm.tools import get_table_schema
 
 
 def get_faiss_vector_db(vectordb_path: Optional[str] = None):
@@ -26,7 +27,15 @@ def get_faiss_vector_db(vectordb_path: Optional[str] = None):
             allow_dangerous_deserialization=True,
         )
     except:
-        documents = get_info_from_db()
+        raw_data = get_table_schema()
+        documents = []
+        for item in raw_data:
+            for table_name, table_info in item.items():
+                column_info_str = "\n".join(
+                    [f"{k}: {v}" for k, v in table_info["columns"].items()]
+                )
+                page_content = f"{table_name}: {table_info['table_description']}\nColumns:\n {column_info_str}"
+                documents.append(Document(page_content=page_content))
         db = FAISS.from_documents(documents, embeddings)
         db.save_local(vectordb_path)
         print(f"VectorDB를 새로 생성했습니다: {vectordb_path}")

utils/llm/vectordb/pgvector_db.py

@@ -7,9 +7,10 @@
 
 import psycopg2
 from langchain_postgres.vectorstores import PGVector
+from langchain.schema import Document
 
 from utils.llm.core import get_embeddings
-from utils.llm.tools import get_info_from_db
+from utils.llm.tools import get_table_schema
 
 
 def _check_collection_exists(connection_string: str, collection_name: str) -> bool:
@@ -71,7 +72,15 @@ def get_pgvector_db(
     except Exception as e:
         print(f"exception: {e}")
         # 컬렉션이 없거나 불러오기에 실패한 경우, 문서를 다시 인덱싱
-        documents = get_info_from_db()
+        raw_data = get_table_schema()
+        documents = []
+        for item in raw_data:
+            for table_name, table_info in item.items():
+                column_info_str = "\n".join(
+                    [f"{k}: {v}" for k, v in table_info["columns"].items()]
+                )
+                page_content = f"{table_name}: {table_info['table_description']}\nColumns:\n {column_info_str}"
+                documents.append(Document(page_content=page_content))
         vector_store = PGVector.from_documents(
             documents=documents,
             embedding=embeddings,