基于 openai assistant api 实现的文档问答助手

document_bot.py

import time

import openai
import streamlit as st


assistant_id = 'asst_XNwQHWKtsPojN9x29UK4wbrX'
client = openai.OpenAI()
# 创建 assistant 执行线程
thread = openai.beta.threads.create()
file_id = None


def upload_file(st_file_obj):
    """上传文档到知识库"""
    file_obj = client.files.create(
        file=st_file_obj.read(),
        purpose='assistants'
    )
    # 更新 assistants 和 文档的关联关系
    client.beta.assistants.update(assistant_id, file_ids=[file_obj.id])
    return file_obj


def delete_file(file_id):
    """删除文档"""
    client.files.delete(file_id=file_id)


def query(question):
    """提问"""
    # 创建提问消息
    client.beta.threads.messages.create(
        thread_id=thread.id,
        role='user',
        content=question
    )
    # 执行 assistants
    run = client.beta.threads.runs.create(
        thread_id=thread.id,
        assistant_id=assistant_id
    )

    # 等待服务器执行完成
    while True:
        run = client.beta.threads.runs.retrieve(
            thread_id=thread.id,
            run_id=run.id
        )
        if run.completed_at:
            break
        time.sleep(1)

    # 获取消息列表
    messages = client.beta.threads.messages.list(thread_id=thread.id)
    last_message = messages.data[0]
    return last_message.content[0].text.value


def upload():
    global file_id
    file_obj = st.session_state.upload_file
    if file_obj:
        # 如果之前关联过文档，就将之前的文档删除
        if file_id:
            delete_file(file_id)
        file_id = upload_file(file_obj)


def question_on_change():
    question = st.session_state.question
    if question:
        res = query(question)
    else:
        res = '请输入要提问的内容'

    st.session_state.result = res


st.title('文档问答助手')
st.file_uploader('请上传你要提问的文档', key='upload_file', on_change=upload)
st.text_input('请输入你的问题：', key='question', on_change=question_on_change)
st.text_area('回答：', '', key='result')

小菜鸟，试一试