from langchain_community.chat_models.tongyi import ChatTongyi
from langchain_postgres import PostgresChatMessageHistory
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
import uuid
import psycopg
from utils.tools import all_tools
# ------------------ 配置 PostgreSQL 聊天记录 ------------------
conn = sync_connection=psycopg.connect("postgresql://postgres:123456@10.10.10.9/ktor2")
database_name = "ai_chat_history"
PostgresChatMessageHistory.create_tables(conn, database_name)

history = PostgresChatMessageHistory(
    database_name,  
    str(uuid.uuid4()),          # session_id
    sync_connection=conn
)
# ------------------ 配置 LLM ------------------
prompt = PromptTemplate(
    input_variables=["question"],
    template="""
        请基于上下文，保持中文简洁回答用户：
        上下文：{history},
        用户：{userStr}
    """
)
llm = ChatTongyi(streaming=False, api_key="sk-fb46eefb6b404382a0a5325202e923a6") 
llm_with_tools = llm.bind_tools(all_tools)
# 建立链
chain = prompt | llm_with_tools
# chain = prompt | llm

# ------------------ 循环聊天 ------------------
while True:
    userStr = input("用户: ")
    history.add_user_message(userStr)
    answer = chain.invoke({"history":history.messages,"userStr": userStr})
    print("AI:", answer.content)
    history.add_ai_message(answer)