RAG工厂制度问答 - 源码

from dotenv import load_dotenv
from openai import OpenAI
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer
import chromadb
import os


# 按照固定字符切割文档
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):
    '''从 PDF 文件中（按指定页码）提取文字'''
    full_text = ''
    # print(extract_pages(filename))
    # # 提取全部文本
    for i, page_layout in enumerate(extract_pages(filename)):
        # print(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.get_text())
            #     # 将换行和空格去掉
                full_text += element.get_text().replace("\n", "").replace(" ", "")
        # print(full_text)
    text_chunks = sliding_window_chunks(full_text, 250, 100)

    return text_chunks


class MyVectorDBConnector:
    def __init__(self, collection_name):
        # 创建一个客户端
        chroma_client = chromadb.Client()
        # 默认是欧式距离
        self.collection = chroma_client.get_or_create_collection(
            name=collection_name,
            metadata={'hnsw:space': 'cosine'}
        )

    def get_embeddings(self, texts, model="text-embedding-v2", batch_size=20): # 改批量，支持批量
        '''封装 qwen 的 Embedding 模型接口'''
        # print('texts', texts)
        # data = client.embeddings.create(input=texts, model=model).data
        # return [x.embedding for x in data]
        all_embeddings = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            data = client.embeddings.create(input=batch, model=model).data
            all_embeddings.extend([x.embedding for x in data])
        return all_embeddings


    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
        )

    def search(self, query, n_results):
        '''检索向量数据库'''
        # 把我们查询的问题向量化, 在chroma当中进行查询
        # D:\software\miniconda\envs\py_ai\Lib\site-packages\chromadb\segment\impl\vector\local_hnsw.py  默认相似度匹配欧式距离
        results = self.collection.query(
            query_embeddings=self.get_embeddings([query]),
            n_results=n_results,   # topk
        )
        return results


class RAG_Bot:
    def __init__(self, vector_db, n_results=2):
        self.vector_db = vector_db
        self.n_results = n_results

        # llm模型
    def get_completion(self, prompt, model="qwen3.8-27b"):
        '''封装 千问 接口'''
        messages = [{"role": "user", "content": prompt}]
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=0,  # 模型输出的随机性，0 表示随机性最小
        )
        # print(response)
        return response.choices[0].message.content

    def chat(self, user_query):
        search_res = self.vector_db.search(user_query, self.n_results)
        print(search_res)
        prompt = prompt_template.replace('__INFO__', ','.join(search_res['documents'][0])).replace('__QUERY__', user_query)
        res = self.get_completion(prompt)
        print('用户问题:', user_query)
        print('大模型回复:', res)



if __name__ == '__main__':
    load_dotenv()
    client = OpenAI(api_key=os.getenv("DASHSCOPE_API_KEY"), base_url=os.getenv("DASHSCOPE_BASE_URL"))

    # 生成
    prompt_template = """
            你是一个问答机器人。
            你的任务是根据下述给定的已知信息回答用户问题。
            确保你的回复完全依据下述已知信息。不要编造答案。
            如果下述已知信息不足以回答用户的问题，请直接回复"我无法回答您的问题"。

            已知信息:
            __INFO__

            用户问：
            __QUERY__

            请用中文回答用户问题。
            """

    # 创建索引
    docx_file = '工厂员工手册.pdf'
    # paragraphs = extract_text_from_pdf(docx_file, page_numbers=[0, 1, 2,9,10]) # 固定提取的页码
    paragraphs = extract_text_from_pdf(docx_file)  # 固定提取的页码
    # print(paragraphs)
    vector_db = MyVectorDBConnector('test1')
    # print(paragraphs)
    vector_db.add_documents(paragraphs)

    # 检索
    bot = RAG_Bot(vector_db, 5)
    user_query = "录用员工必须签订合同么?"
    response = bot.chat(user_query)