Ubuntu
155 lines
4.7 KiB
Python
155 lines
4.7 KiB
Python
"""
|
|
AI 客服核心逻辑 - 基于 LangGraph
|
|
"""
|
|
from typing import TypedDict, Annotated, Sequence
|
|
import operator
|
|
|
|
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.graph import StateGraph, END
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
|
|
from config import LLM_API_KEY, LLM_BASE_URL, LLM_MODEL, SYSTEM_PROMPT
|
|
from knowledge import get_knowledge_base
|
|
|
|
|
|
# 定义状态类型
|
|
class ChatState(TypedDict):
|
|
"""对话状态"""
|
|
messages: Annotated[Sequence[BaseMessage], operator.add] # 消息历史
|
|
context: str # 检索到的知识上下文
|
|
session_id: str # 会话ID
|
|
|
|
|
|
class AIChatbot:
|
|
"""AI 客服机器人"""
|
|
|
|
def __init__(self):
|
|
# 初始化 LLM
|
|
self.llm = ChatOpenAI(
|
|
model=LLM_MODEL,
|
|
api_key=LLM_API_KEY,
|
|
base_url=LLM_BASE_URL,
|
|
temperature=0.7,
|
|
max_tokens=1000
|
|
)
|
|
|
|
# 获取知识库
|
|
self.knowledge_base = get_knowledge_base()
|
|
|
|
# 构建工作流
|
|
self.workflow = self._build_workflow()
|
|
|
|
def _retrieve_knowledge(self, state: ChatState) -> ChatState:
|
|
"""检索相关知识"""
|
|
# 获取最后一条用户消息
|
|
last_message = state["messages"][-1]
|
|
if not isinstance(last_message, HumanMessage):
|
|
return state
|
|
|
|
query = last_message.content
|
|
|
|
# 从知识库检索
|
|
docs = self.knowledge_base.search(query)
|
|
|
|
# 构建上下文
|
|
if docs:
|
|
context = "\n\n".join([f"[知识库内容 {i+1}]\n{doc.page_content}"
|
|
for i, doc in enumerate(docs)])
|
|
else:
|
|
context = "【知识库中没有找到相关信息】"
|
|
|
|
return {"context": context, "messages": state["messages"], "session_id": state["session_id"]}
|
|
|
|
def _generate_response(self, state: ChatState) -> ChatState:
|
|
"""生成回答"""
|
|
# 构建系统提示
|
|
system_prompt = f"""{SYSTEM_PROMPT}
|
|
|
|
【知识库信息】
|
|
{state['context']}
|
|
|
|
请基于以上知识库信息回答用户问题。如果知识库中没有相关信息，请诚实地告诉用户你无法回答。
|
|
"""
|
|
|
|
# 构建消息列表
|
|
messages = [SystemMessage(content=system_prompt)] + list(state["messages"])
|
|
|
|
# 调用 LLM
|
|
response = self.llm.invoke(messages)
|
|
|
|
# 返回新的消息（会自动添加到历史）
|
|
return {"messages": [response], "context": state["context"], "session_id": state["session_id"]}
|
|
|
|
def _build_workflow(self) -> StateGraph:
|
|
"""构建 LangGraph 工作流"""
|
|
# 创建工作流
|
|
workflow = StateGraph(ChatState)
|
|
|
|
# 添加节点
|
|
workflow.add_node("retrieve", self._retrieve_knowledge)
|
|
workflow.add_node("generate", self._generate_response)
|
|
|
|
# 添加边
|
|
workflow.set_entry_point("retrieve")
|
|
workflow.add_edge("retrieve", "generate")
|
|
workflow.add_edge("generate", END)
|
|
|
|
# 编译（添加内存检查点）
|
|
memory = MemorySaver()
|
|
return workflow.compile(checkpointer=memory)
|
|
|
|
async def chat(self, message: str, session_id: str) -> str:
|
|
"""
|
|
处理用户消息
|
|
|
|
Args:
|
|
message: 用户消息
|
|
session_id: 会话ID（用于区分不同用户）
|
|
|
|
Returns:
|
|
AI 回复
|
|
"""
|
|
# 构建配置
|
|
config = {"configurable": {"thread_id": session_id}}
|
|
|
|
# 运行工作流
|
|
result = await self.workflow.ainvoke(
|
|
{
|
|
"messages": [HumanMessage(content=message)],
|
|
"context": "",
|
|
"session_id": session_id
|
|
},
|
|
config
|
|
)
|
|
|
|
# 返回最后一条 AI 消息
|
|
last_message = result["messages"][-1]
|
|
if isinstance(last_message, AIMessage):
|
|
return last_message.content
|
|
return str(last_message)
|
|
|
|
def get_history(self, session_id: str) -> list:
|
|
"""获取对话历史"""
|
|
config = {"configurable": {"thread_id": session_id}}
|
|
# 从 checkpointer 获取状态
|
|
state = self.workflow.get_state(config)
|
|
if state and "messages" in state.values:
|
|
return [
|
|
{"role": "user" if isinstance(m, HumanMessage) else "assistant",
|
|
"content": m.content}
|
|
for m in state.values["messages"]
|
|
]
|
|
return []
|
|
|
|
|
|
# 全局客服实例
|
|
_chatbot = None
|
|
|
|
def get_chatbot() -> AIChatbot:
|
|
"""获取客服单例"""
|
|
global _chatbot
|
|
if _chatbot is None:
|
|
_chatbot = AIChatbot()
|
|
return _chatbot
|