BBIT-Kai
259 lines
7.6 KiB
Python
259 lines
7.6 KiB
Python
from typing import List
|
|
from typing import TypedDict
|
|
|
|
from langchain_core.prompts import PromptTemplate
|
|
from langgraph.graph import StateGraph, END
|
|
|
|
import db.milvus as milvus
|
|
import db.postgres as pgdb
|
|
from config.llm import llm
|
|
from config.llm import llmThink
|
|
|
|
|
|
# -------- 定义状态 --------
|
|
class State(TypedDict):
|
|
path: str # 开始聊天选择的路径
|
|
|
|
memory: str # 记忆
|
|
knowledge: str # 知识库内容
|
|
history: str # 聊天历史
|
|
|
|
ai_id: str # AI id
|
|
ai_name: str # AI 名称
|
|
ai_service: str # AI 角色 业务
|
|
ai_role: str # AI 角色 性格特点
|
|
kn_bases: List[str] # AI 所使用的知识库
|
|
|
|
userInput: str # 用户输入
|
|
reply: str # 最终回复
|
|
|
|
|
|
# -------- 定义节点 --------
|
|
# ------------------------------------------------------------------------ 向量数据库查询 --------
|
|
|
|
gen_sql_prompt = PromptTemplate(
|
|
input_variables=["userInput"],
|
|
template="""你的任务是对用户输入进行意图分析，并将其分解成方便进行知识向量数据库搜索的关键词。
|
|
以下是用户的输入：
|
|
<用户输入>
|
|
{userInput}
|
|
</用户输入>
|
|
在提取关键词时，请遵循以下方法和要求：
|
|
1. 去除输入中的停用词（如“的”“是”“在”等）、语气词和无实际意义的符号。
|
|
2. 识别输入中的核心概念、实体和关键动作。
|
|
3. 尽量使用简洁、通用的词汇作为关键词。
|
|
4. 确保关键词之间相互独立，不包含其他关键词。
|
|
关键词之间用空格分隔。
|
|
你的回答是：
|
|
""",
|
|
)
|
|
sqlChain = gen_sql_prompt | llm
|
|
|
|
|
|
def db_search(state: State):
|
|
key_words = sqlChain.invoke(
|
|
{
|
|
"userInput": state["userInput"],
|
|
}
|
|
).content
|
|
print("关键词是：", key_words)
|
|
knowledge = milvus.get_knowledge_by_key_words(key_words, state["kn_bases"])
|
|
print("知识库内容是：", knowledge)
|
|
state["knowledge"] = knowledge
|
|
ai_ids = [state["ai_id"]]
|
|
memory = milvus.get_memory_by_key_words(key_words, ai_ids)
|
|
print("记忆是：", memory)
|
|
state["memory"] = memory
|
|
return state
|
|
|
|
|
|
# ------------------------------------------------------------------------ 意图分析 --------
|
|
|
|
pathSelectPrompt = PromptTemplate(
|
|
input_variables=["userInput", "ai_service", "history"],
|
|
template="""
|
|
你是一个意图分类器，负责判断用户提问是否与你的工作相关，进而确定是否需要去查知识库。
|
|
以下是你负责的工作内容：
|
|
<ai_service>
|
|
{ai_service}
|
|
</ai_service>
|
|
这是你们的对话历史：
|
|
<history>
|
|
{history}
|
|
</history>
|
|
用户最新回复是：
|
|
<userInput>
|
|
{userInput}
|
|
</userInput>
|
|
判断规则如下：
|
|
如果用户最新回复与你的负责工作相关，需要去查知识库，输出“kn”；如果不相关，则输出“chat”，不要包含任何标点符号以及空格。
|
|
你生成的结果：
|
|
""",
|
|
)
|
|
pathSelectChain = pathSelectPrompt | llmThink
|
|
|
|
|
|
def decide_source(state: State, max_retry=3):
|
|
"""根据用户输入选择数据来源"""
|
|
for _ in range(max_retry):
|
|
choice = (
|
|
pathSelectChain.invoke(
|
|
{
|
|
"userInput": state["userInput"],
|
|
"ai_service": state["ai_service"],
|
|
"history": state["history"],
|
|
}
|
|
)
|
|
.content.strip()
|
|
.lower()
|
|
)
|
|
print("根据用户输入选择数据来源，路径是：", choice)
|
|
if choice in ["kn", "chat"]:
|
|
state["path"] = choice
|
|
break
|
|
else:
|
|
# 如果连续 max_retry 次都不合法，默认走 chat
|
|
state["path"] = "chat"
|
|
return state
|
|
|
|
|
|
# ------------------------------------------------------------------------ !普通聊天 --------
|
|
noChatPrompt = PromptTemplate(
|
|
input_variables=["ai_name", "ai_service", "ai_role", "history"],
|
|
template="""
|
|
你的名字是：{ai_name}，你负责的业务是{ai_service}，你具有{ai_role}的性格特点。
|
|
|
|
这是你和用户的对话历史
|
|
<history>
|
|
{history}
|
|
</history>
|
|
在回复用户时，请遵循以下指南:
|
|
1. 回复要与AI角色业务相关，体现AI的专业能力。
|
|
2. 回复内容的语气和风格要符合AI角色性格特点。
|
|
3. 参考聊天历史，使回复具有连贯性和针对性。
|
|
4. 回复要简洁明了，避免冗长和复杂的表述。
|
|
|
|
你的回答：
|
|
""",
|
|
)
|
|
|
|
noChatChain = noChatPrompt | llm
|
|
|
|
|
|
def chat(state: State):
|
|
state["reply"] = noChatChain.invoke(
|
|
{
|
|
"ai_name": state["ai_name"],
|
|
"ai_service": state["ai_service"],
|
|
"ai_role": state["ai_role"],
|
|
"history": state["history"],
|
|
"userStr": state["userInput"],
|
|
}
|
|
).content
|
|
print("直接回复")
|
|
return state
|
|
|
|
|
|
# ------------------------------------------------------------------------ 整理结果 --------
|
|
|
|
summarizePrompt = PromptTemplate(
|
|
input_variables=["ai_name", "ai_service", "ai_role", "history", "knowledge"],
|
|
template="""
|
|
你的任务是基于给定的AI名称、AI角色业务、AI角色性格特点和聊天历史来回复用户。请仔细阅读以下信息，并按照指示进行回复。
|
|
你的名字是：{ai_name}，你负责的业务是{ai_service}，你具有{ai_role}的性格特点。
|
|
|
|
这是你和用户的对话历史
|
|
<history>
|
|
{history}
|
|
</history>
|
|
这是给你参考的知识库：
|
|
<knowledge>
|
|
{knowledge}
|
|
</knowledge>
|
|
{memory}
|
|
在回复时，请遵循以下指南：
|
|
1. 回复内容要与你负责的业务相关。
|
|
2. 回复的语气要结合你的性格特点。
|
|
3. 确保回复内容清晰、简洁、有针对性。
|
|
请生成你的回复：
|
|
""",
|
|
)
|
|
summarizeChain = summarizePrompt | llm
|
|
|
|
|
|
def summarize_ai(state: State):
|
|
"""AI 总结输出"""
|
|
mem = state["memory"]
|
|
if mem != "":
|
|
memStr = (
|
|
"""
|
|
这是给你参考的相关历史记忆：
|
|
<memory>
|
|
%s
|
|
</memory>
|
|
"""
|
|
% mem
|
|
) # 这里用 % 把 mem 填进去
|
|
else:
|
|
memStr = "没有记忆内容"
|
|
print("历史记录是：", state["history"])
|
|
state["reply"] = summarizeChain.invoke(
|
|
{
|
|
"ai_role": state["ai_role"],
|
|
"ai_name": state["ai_name"],
|
|
"history": state["history"],
|
|
"ai_service": state["ai_service"],
|
|
"knowledge": state["knowledge"],
|
|
"memory": memStr,
|
|
}
|
|
).content
|
|
return state
|
|
|
|
|
|
# ------------------------------------------------------------------------ 构建有向图 --------
|
|
workflow = StateGraph(State)
|
|
workflow.add_node("decide", decide_source)
|
|
workflow.add_node("db_search", db_search)
|
|
workflow.add_node("chat", chat)
|
|
workflow.add_node("summarize", summarize_ai)
|
|
workflow.set_entry_point("decide")
|
|
# 条件边：根据 path 决定走向
|
|
workflow.add_conditional_edges(
|
|
"decide",
|
|
lambda state: state["path"], # 返回 state["path"] 的值
|
|
{
|
|
"kn": "db_search",
|
|
"chat": "chat",
|
|
},
|
|
)
|
|
workflow.add_edge("db_search", "summarize")
|
|
workflow.add_edge("summarize", END)
|
|
workflow.add_edge("chat", END)
|
|
graph = workflow.compile()
|
|
|
|
|
|
# 执行函数
|
|
def get_service_agent_reply(
|
|
aiId: str, userInput: str, history: str, kn_bases: List[str]
|
|
):
|
|
json = pgdb.get_ai_personality(aiId)
|
|
ai_service = json["业务"]
|
|
ai_role = json["性格"]
|
|
ai_name = json["名字"]
|
|
print("AI Name:", ai_name)
|
|
print("AI Service:", ai_service)
|
|
|
|
final_state = graph.invoke(
|
|
{
|
|
"ai_service": ai_service,
|
|
"ai_role": ai_role,
|
|
"ai_name": ai_name,
|
|
"history": history,
|
|
"kn_bases": kn_bases,
|
|
"table_info": pgdb.get_available_tables_str(aiId),
|
|
"userInput": userInput,
|
|
"ai_id": aiId,
|
|
}
|
|
)
|
|
return final_state["reply"]
|