这段代码实现了一个基于LangChain框架的聊天客服机器人,专门用于“易速鲜花”业务。该机器人能够处理PDF、Word和纯文本文件,并利用OpenAI的嵌入模型和语言模型来生成回答。代码通过Gradio库创建了一个用户界面,使得用户可以与机器人进行交互。
-
DoubaoEmbeddings类 :
class DoubaoEmbeddings(BaseModel, Embeddings):
client: Ark = None
api_key: str = ""
model: str
def __init__(self, **data: Any):
super().__init__(**data)
if self.api_key == "":
self.api_key = os.environ["OPENAI_API_KEY"]
self.client = Ark(
base_url=os.environ["OPENAI_BASE_URL"],
api_key=self.api_key
)
def embed_query(self, text: str) -> List[float]:
embeddings = self.client.embeddings.create(model=self.model, input=text)
return embeddings.data[0].embedding
def embed_documents(self, texts: List[str]) -> List[List[float]]:
return [self.embed_query(text) for text in texts]
class Config:
arbitrary_types_allowed = True
该类用于生成文本的嵌入向量。它通过OpenAI的API生成嵌入,并支持单个文本和多个文本的嵌入生成。
### ChatbotWithRetrieval 类详细介绍
ChatbotWithRetrieval 类是整个聊天客服机器人的核心实现,它负责加载文档、处理文本、生成嵌入向量、管理对话记忆,并最终通过检索链生成回答。以下是对该类的详细介绍:
1. 初始化方法 __init__
def init(self, dir):
__init__ 方法在创建 ChatbotWithRetrieval 类的实例时被调用,它接受一个参数 dir,表示文档的存放目录。该方法负责初始化聊天机器人的各个组件。
1.1 加载文档
base_dir = dir
documents = []
for file in os.listdir(base_dir):
file_path = os.path.join(base_dir, file)
if file.endswith(".pdf"):
loader = PyPDFLoader(file_path)
documents.extend(loader.load())
elif file.endswith(".docx") or file.endswith(".doc"):
loader = Docx2txtLoader(file_path)
documents.extend(loader.load())
elif file.endswith(".txt"):
loader = TextLoader(file_path)
documents.extend(loader.load())
这部分代码遍历指定目录中的所有文件,根据文件扩展名使用相应的加载器(PyPDFLoader, Docx2txtLoader, TextLoader)加载文档内容,并将加载的文档添加到 documents 列表中。
1.2 文本分割
text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=0)
all_splits = text_splitter.split_documents(documents)
使用 RecursiveCharacterTextSplitter 将加载的文档分割成大小为 200 个字符的块,且块之间没有重叠。分割后的文档块存储在 all_splits 列表中。
1.3 向量数据库
self.vectorstore = Qdrant.from_documents(
documents=all_splits,
embedding=DoubaoEmbeddings(
model=os.environ["EMBEDDING_MODELEND"],
),
location=":memory:",
collection_name="my_documents",
)
使用 Qdrant 向量数据库存储分割后的文档块。DoubaoEmbeddings 类用于生成文档块的嵌入向量。location=":memory:" 表示向量数据库存储在内存中,collection_name="my_documents" 指定了集合名称。
1.4 初始化语言模型
self.llm = ChatOpenAI(
model=os.environ["LLM_MODELEND"],
temperature=0,
)
使用 ChatOpenAI 初始化语言模型,model 参数指定了使用的模型名称,temperature=0 表示生成回答时采用确定性较高的模式。
1.5 初始化对话记忆
self.memory = ConversationSummaryMemory(
llm=self.llm, memory_key="chat_history", return_messages=True
)
使用 ConversationSummaryMemory 初始化对话记忆,llm 参数指定了用于生成对话摘要的语言模型,memory_key="chat_history" 指定了记忆的键名,return_messages=True 表示返回对话消息。
1.6 初始化对话历史
self.conversation_history = ""
初始化对话历史为空字符串。
1.7 设置检索链
retriever = self.vectorstore.as_retriever()
self.qa = ConversationalRetrievalChain.from_llm(
self.llm, retriever=retriever, memory=self.memory
)
使用 Qdrant 向量数据库的检索器初始化 ConversationalRetrievalChain,该链结合了语言模型、检索器和对话记忆,用于生成回答。
2. 获取回答方法 get_response
def get_response(self, user_input):
response = self.qa(user_input)
self.conversation_history += (
f"你: {user_input}\nChatbot: {response['answer']}\n"
)
return self.conversation_history
get_response 方法用于处理用户输入并生成回答。它接受 user_input 参数,表示用户的输入。通过调用 self.qa(user_input) 生成回答,并将用户输入和生成的回答更新到对话历史中。最后返回对话历史。