From c3e39e822cb1fa47e36eb04c623b800afc6f332a Mon Sep 17 00:00:00 2001
From: RYDE-WORK <rydewu@primedigitaltech.com>
Date: Tue, 27 Aug 2024 20:34:05 +0800
Subject: [PATCH] Initial commit

---
 file_chat.py | 65 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)
 create mode 100644 file_chat.py

diff --git a/file_chat.py b/file_chat.py
new file mode 100644
index 0000000..03504bf
--- /dev/null
+++ b/file_chat.py
@@ -0,0 +1,65 @@
+import os
+import json
+import streamlit as st
+from zhipuai import ZhipuAI
+
+
+def main():
+    st.set_page_config(page_title="文件对话demo", page_icon="📖", layout="wide")  # 页面设置
+    st.header("📖文件对话demo")  # 标题
+    
+    with st.sidebar:  # 侧边栏
+        st.markdown("调用智谱AI文件对话接口进行问答")
+        api_key_input = st.text_input(
+            "智谱API Key",
+            type="password",
+            value=os.environ.get("API_KEY", None)
+            or st.session_state.get("API_KEY", ""),
+            key="API_KEY",
+        )
+
+    api_key = st.session_state.get("API_KEY")
+    if not api_key:
+        st.warning(
+            "请在侧边栏输入API_KEY。可从"
+            " https://bigmodel.cn/usercenter/apikeys 获取"
+        )
+    client = ZhipuAI(api_key=api_key)
+    uploaded_file = st.file_uploader(
+        "上传文件",
+        type=["pdf", "docx", "txt"],
+        help="文件大小不超过50M，图片大小不超过5M",
+        on_change=lambda: st.session_state.pop("messages", None),
+    )
+
+    if not uploaded_file:
+        st.stop()
+    else:
+        if "messages" not in st.session_state:
+            file_object = client.files.create(file=uploaded_file, purpose="file-extract")
+            file_content = json.loads(client.files.content(file_id=file_object.id).content)["content"]
+            message_content = f"请对\n{file_content}\n的内容进行分析，并撰写一份论文摘要。"
+            response = client.chat.completions.create(
+                model="glm-4-long",
+                messages=[
+                    {"role": "user", "content": message_content}
+                ],
+            )
+            msg = response.choices[0].message.content
+            st.session_state["messages"] = [{"role": "assistant", "content": msg}]  # 每次提问时不带上全文信息
+            # st.session_state["messages"] = [{"role": "user", "content": message_content}, {"role": "assistant", "content": msg}]  # 每次提问时带上全文信息
+
+        for msg in st.session_state.get("messages", []):
+            st.chat_message(msg["role"]).write(msg["content"])
+
+        if prompt := st.chat_input():
+            st.session_state.messages.append({"role": "user", "content": prompt})
+            st.chat_message("user").write(prompt)
+            response = client.chat.completions.create(model="glm-4-long", messages=st.session_state.messages)
+            msg = response.choices[0].message.content
+            st.session_state.messages.append({"role": "assistant", "content": msg})
+            st.chat_message("assistant").write(msg)
+
+
+if __name__ == "__main__":
+    main()