BBIT-Kai
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
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)
|
|
|
|
|