import os import chromadb from dotenv import load_dotenv from openai import OpenAI from pdfminer.high_level import extract_pages from pdfminer.layout import LTTextContainer # 按照固定字符切割文档 def sliding_window_chunks(text, chunk_size, stride): return [text[i:i + chunk_size] for i in range(0, len(text), stride)] # 读取PDF 分割文档 def extract_text_from_pdf(filename, page_numbers=None, min_line_length=1): '''从 PDF 文件中（按指定页码）提取文字''' paragraphs = [] buffer = '' full_text = '' # 提取全部文本 for i, page_layout in enumerate(extract_pages(filename)): # print(i, page_layout) # 如果指定了页码范围，跳过范围外的页 if page_numbers is not None and i not in page_numbers: continue for element in page_layout: # print(element) # 检查element是不是文本 if isinstance(element, LTTextContainer): # print(element) # print(element.get_text()) # 将换行和空格去掉 full_text += element.get_text().replace("\n", "").replace(" ", "") # print(full_text) # print(full_text) if full_text: # 调用切割函数 text_chunks = sliding_window_chunks(full_text, 250, 100) for text in text_chunks: paragraphs.append(text) return paragraphs class MyVectorDBConnector: def __init__(self): # 创建数据库的链接 self.db = chromadb.Client() # 创建数据库 self.collection = self.db.get_or_create_collection(name="demo") def get_embeddings(self, texts, model="text-embedding-v2"): '''封装 qwen 的 Embedding 模型接口''' # print('texts', texts) data = client.embeddings.create(input=texts, model=model).data return [x.embedding for x in data] def add_documents(self, instructions): '''向 collection 中添加文档与向量''' # get_embeddings(instructions) # 将数据向量化 embeddings = self.get_embeddings(instructions) # 把向量化的数据和原文存入向量数据库 self.collection.add( embeddings=embeddings, # 每个文档的向量 documents=instructions, # 文档的原文 ids=[f"id{i}" for i in range(len(instructions))] # 每个文档的 id ) # print(self.collection.count()) def search(self, query, n_results): '''检索向量数据库''' # 把我们查询的问题向量化, 在chroma当中进行查询 results = self.collection.query( query_embeddings=self.get_embeddings([query]), n_results=n_results, ) return results class RAG_bot(): def __init__(self, vector_db, n_res): self.vector_db = vector_db self.n_res = n_res # llm模型 def get_completion(self, prompt, model="qwen-turbo"): '''封装 千问 接口''' messages = [{"role": "user", "content": prompt}] response = client.chat.completions.create( model=model, messages=messages, temperature=0, # 模型输出的随机性，0 表示随机性最小 ) return response.choices[0].message.content def chat(self, queyr): # 1.检索 search_Data = self.vector_db.search(queyr, self.n_res) print(search_Data['documents'][0]) # 2. 构建提示词 prompt = prompt_template.replace('__INFO__', '\n'.join(search_Data['documents'][0])).replace('__QUERY__', queyr) response = self.get_completion(prompt) print('AI回复的内容：', response) if __name__ == '__main__': load_dotenv() client = OpenAI(api_key=os.getenv("DASHSCOPE_API_KEY"), base_url="https://dashscope.aliyuncs.com/compatible-mode/v1") prompt_template = """ 你是一个问答机器人。 你的任务是根据下述给定的已知信息回答用户问题。 确保你的回复完全依据下述已知信息。不要编造答案。 如果下述已知信息不足以回答用户的问题，请直接回复"我无法回答您的问题"。 已知信息: __INFO__ 用户问： __QUERY__ 请用中文回答用户问题。 """ page_data = extract_text_from_pdf('财务管理文档.pdf', page_numbers=[0, 1, 2]) # print(page_data) # 创建向量数据库对象 vector_db = MyVectorDBConnector() # 添加向量 vector_db.add_documents(page_data) # 创建机器人对象 rag_bot = RAG_bot(vector_db, n_res=2) rag_bot.chat('财务管理权限划分')
