from sqlalchemy import create_engine
from llama_index.core import SQLDatabase
from llama_index.core.llms import ChatMessage
from llama_index.llms.ollama import Ollama
import pandas as pd

db_user = "shenzhen_ai_for_vibemate_eson"
db_password = "dBsnc7OrM0MVi0FEhiHe2y"
db_host = "192.168.1.99"
db_port = 3306
db_name = "hytto_surfease"

engine = create_engine(f"mysql+pymysql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}", echo=True)
sql_database = SQLDatabase(engine)

table_names = sql_database.get_usable_table_names()
PROMPT_STR = """\
    Give me a summary of the table with the following format.
    
    - table_summary: Describe what the table is about in short. Columns: [col1(type), col2(type), ...]
    
    Table:
    {table_str}
     """
table_infos = {}
for table in table_names:
    table_info = sql_database.get_single_table_info(table)
    table_infos[table] = table_info

print(table_infos)