PDF问答工具(基于openai API和streamlit)

utils.py

python 复制代码
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS

from langchain_openai import ChatOpenAI
from langchain.chains import ConversationalRetrievalChain

def rag_tool(api_key,memory,uploaded_file,question):
    # 将文件内容写入到本地,才能有文件路径(用户上传的文件直接储存在内存里,无路径)
    file_content = uploaded_file.read() #返回文件的二进制数据
    temp_file_path = "temp.pdf"
    with open(temp_file_path,"wb") as fwb: #以二进制写入方式打开文件,并创建文件对象(与文件进行交互的接口)
        fwb.write(file_content) #写入上面读取到的文件内容

    # 加载
    loader = PyPDFLoader(temp_file_path) 
    docs = loader.load() 
    # 分割
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=50,
        separators=["\n\n","\n","。","!","?",",","","、"," ",""]
    ) 
    texts = text_splitter.split_documents(docs)
    # 嵌入模型
    embeddings_model = OpenAIEmbeddings(api_key=api_key,
                                        base_url="https://api.gptsapi.net/v1") 
    # 嵌入并储存
    db = FAISS.from_documents(texts,embeddings_model)
    # 检索
    retriever = db.as_retriever() 


    # 模型
    model = ChatOpenAI(model="gpt-3.5-turbo",
                    api_key=api_key,
                    base_url="https://api.gptsapi.net/v1")

    # 记忆

    # 创建带记忆的检索增强对话链(有了检索器、模型、记忆)
    chain = ConversationalRetrievalChain.from_llm(
        llm = model,
        retriever = retriever,
        memory = memory,
        return_source_documents=True
    )

    result = chain.invoke(
        {
            "chat_history":memory,
            "question":question
        }
    )
    return result

main.py

python 复制代码
import streamlit as st
from utils import rag_tool

from langchain.memory import ConversationBufferMemory

st.title("智能PDF问答工具")

with st.sidebar:
    api_key = st.text_input("请输入你的OpenAI API密钥:",type="password")
    st.markdown("[获取OpenAI API密钥](https://2233.ai/api)")

    # 上传文件
    uploaded_file = st.file_uploader("请上传你的PDF文件:",type="pdf")


# 初始化会话状态
if "memory" not in st.session_state:
    st.session_state.memory = ConversationBufferMemory(
        return_messages=True,
        memory_key="chat_history",
        output_key="answer"
    )
    st.session_state.messages = [
        {
            "role":"ai","content":"你好,我是PDF分析小助手,上传文件向我提问吧!"
        }
    ]
    st.session_state.documents = []

# 显示历史对话消息(初始显示)和历史资料
num = 0
for message in st.session_state.messages:
    st.chat_message(message["role"]).write(message["content"])
    if message["role"]=="ai":
        if message["content"]!="你好,我是PDF分析小助手,上传文件向我提问吧!": #注意超范围问题,第一句ai消息没有相应的相关资料
            with st.expander("相关资料"):
                st.write(st.session_state.documents[num][0].page_content) 
            num += 1

# 获取用户输入
question = st.chat_input("对PDF的内容进行提问")
if question:
    if not api_key:
        st.info("请输入你的OpenAI API密钥")
        st.stop()
    if not uploaded_file:
        st.info("请先上传文件!")
        st.stop()

    # 合法后就显示
    st.session_state.messages.append(
        {"role":"human","content":question}
    )
    st.chat_message("human").write(question)

    # 获取AI的回复
    with st.spinner("AI正在思考中,请稍等···"):
        result = rag_tool(api_key=api_key,memory=st.session_state.memory,uploaded_file=uploaded_file,question=question)

    answer = result["answer"]
    st.session_state.messages.append(
        {"role":"ai","content":answer}
    )
    st.chat_message("ai").write(answer)

    relavant_docs = result["source_documents"]
    st.session_state.documents.append(relavant_docs)
    with st.expander("相关资料"):
        st.write(relavant_docs[0].page_content) #Document(id='1abe1f48-fd94-445f-b3ff-02b1b5f12f29', metadata={'source': './京剧介绍.txt'},page_content='***') 这种里面的元素是对象的属性,不是键值
相关推荐
在猴站学算法24 分钟前
机器学习(西瓜书) 第二章 模型评估与选择
人工智能·机器学习
科技宅说1 小时前
36氪专访丨乐橙CEO谢运:AI科技下的业务创新与长期主义下的品牌坚守
人工智能·科技
学术小八2 小时前
2025年人工智能、虚拟现实与交互设计国际学术会议
人工智能·交互·vr
仗剑_走天涯3 小时前
基于pytorch.nn模块实现线性模型
人工智能·pytorch·python·深度学习
cnbestec4 小时前
协作机器人UR7e与UR12e:轻量化设计与高负载能力助力“小而美”智造升级
人工智能·机器人·协作机器人·ur协作机器人·ur7e·ur12e
zskj_zhyl4 小时前
毫米波雷达守护银发安全:七彩喜跌倒检测仪重构居家养老防线
人工智能·安全·重构
gaosushexiangji5 小时前
利用sCMOS科学相机测量激光散射强度
大数据·人工智能·数码相机·计算机视觉
ai小鬼头7 小时前
AIStarter新版重磅来袭!永久订阅限时福利抢先看
人工智能·开源·github
说私域7 小时前
从品牌附庸到自我表达:定制开发开源AI智能名片S2B2C商城小程序赋能下的营销变革
人工智能·小程序
飞哥数智坊8 小时前
新版定价不够用,Cursor如何退回旧版定价
人工智能·cursor