update: 테이블 정보를 가져오는 방식 수정

Author: DShomin

evaluation/gen_persona.py

@@ -1,39 +1,23 @@
-from dotenv import load_dotenv
 import os
 
-from datahub_cls.metadata_fetcher import (
-    DatahubMetadataFetcher,
-    get_all_tables_info,
-)
-
 from utils import save_persona_json, pretty_print_persona
 from persona_class import PersonaList
 
-
+from llm_utils.tools import _get_table_info
 from langchain_openai.chat_models import ChatOpenAI
 from langchain_core.prompts import ChatPromptTemplate
 from argparse import ArgumentParser
 
-load_dotenv()
-
-
-def drop_empty_tables(tables_df):
-    drop_empty_tables = tables_df[
-        tables_df["table_description"].apply(lambda x: x != "")
-    ]
-    return drop_empty_tables[["table_name", "table_description"]]
-
 
-def get_table_des_string(tables_df):
+def get_table_des_string(tables_desc):
     return_string = "table name : table description\n---\n"
-    for _, row in tables_df.iterrows():
-        return_string += f"{row['table_name']} : {row['table_description']}\n---\n"
+    for table_name, table_desc in tables_desc.items():
+        return_string += f"{table_name} : {table_desc}\n---\n"
     return return_string
 
 
-def generate_persona(tables_df):
-    tables_df = drop_empty_tables(tables_df)
-    description_string = get_table_des_string(tables_df)
+def generate_persona(tables_desc):
+    description_string = get_table_des_string(tables_desc)
 
     llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
     system_prompt = """주어진 Tabel description들을 참고하여 Text2SQL 서비스로 질문을 할만한 패르소나를 생성하세요"""
@@ -50,9 +34,8 @@ def generate_persona(tables_df):
 
 def main(output_path):
     # 데이터허브 서버 연결
-    fetcher = DatahubMetadataFetcher(gms_server=os.getenv("DATAHUB_SERVER"))
-    tables_df = get_all_tables_info(fetcher)
-    personas = generate_persona(tables_df)
+    tables_desc = _get_table_info()
+    personas = generate_persona(tables_desc)
 
     for persona in personas.personas:
         print(pretty_print_persona(persona))