Last active
December 23, 2023 05:21
-
-
Save liaokongVFX/1bd95f9a8dd6be3b837c05af88e0063b to your computer and use it in GitHub Desktop.
基于 openai assistant api 实现的文档问答助手
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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') |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
小菜鸟,试一试