import pprint
from dotenv import load_dotenv
import logging
import os
from pandasai.llm import OpenAI

load_dotenv()
model = os.getenv('LLM_MODEL')
api_key = os.getenv('LLM_API_KEY')
api_base = os.getenv('LLM_API_BASE')
OpenAI._supported_chat_models.append(model)
llm = OpenAI(verbose=True, model=model, api_token=api_key, api_base=api_base)


def initLogger(level=logging.INFO):
    logging.basicConfig(
        level=level,  # 设置日志级别
        format='%(asctime)s [%(levelname)s] %(name)s - %(message)s',  # 设置日志格式
        datefmt='%Y-%m-%d %H:%M:%S',  # 设置日期格式
        handlers=[
            logging.StreamHandler(),  # 输出到控制台
            logging.FileHandler('viz_llm.log')  # 输出到文件
        ]
    )
    logging.info(f"App Start!")


if os.getenv('LOG_LEVEL') == 'DEBUG':
    initLogger(logging.DEBUG)
else:
    initLogger()


if __name__ == "__main__":
    res = llm.chat_completion("hello", None)
    pprint.pprint(res)