shield-chatbot / src /streamlit_app.py
DeepLearning101's picture
Update src/streamlit_app.py
1df899f verified
raw
history blame
3.7 kB
import streamlit as st
import google.generativeai as genai
import requests
import os
st.set_page_config(page_title="AI 新知小助手", page_icon="📚", layout="wide")
# 改從環境變數中取得 API Key
api_key = os.environ.get("GEMINI_API_KEY")
if not api_key:
st.error("請確認已經在 Space 的 Settings -> Variables and secrets 中設定了 GEMINI_API_KEY")
st.stop()
genai.configure(api_key=api_key)
# 1. 定義你要鎖定的 GitHub 檔案清單
# 請把 YOUR_ACCOUNT/YOUR_REPO 換成你真實的 GitHub 資訊
BASE_URL = "https://raw.githubusercontent.com/Deep-Learning-101/deep-learning-101.github.io/main/"
TARGET_FILES = {
"大型語言模型 (LLM)": f"{BASE_URL}Large-Language-Model.md",
"自然語言處理 (NLP)": f"{BASE_URL}Natural-Language-Processing.md",
"語音處理 (Speech)": f"{BASE_URL}Speech-Processing.md",
"電腦視覺 (CV)": f"{BASE_URL}Computer-Vision.md"
}
# 2. 批量抓取並格式化內容
def fetch_all_knowledge():
combined_knowledge = ""
with st.spinner("正在從 GitHub 同步 4 大領域最新資訊..."):
for category, url in TARGET_FILES.items():
try:
response = requests.get(url)
response.raise_for_status()
# 在每個檔案內容前加上明確的標題,幫助 AI 分類記憶
combined_knowledge += f"\n\n## 【領域:{category}】\n"
combined_knowledge += response.text
except Exception as e:
st.warning(f"無法同步 {category} 的資料:{e}")
return combined_knowledge
# 初始化 Session State
if "knowledge" not in st.session_state:
st.session_state.knowledge = fetch_all_knowledge()
if "messages" not in st.session_state:
st.session_state.messages = []
# 側邊欄設計
with st.sidebar:
st.title("⚙️ 知識庫狀態")
st.write("目前收錄以下領域:")
for category in TARGET_FILES.keys():
st.markdown(f"- {category}")
st.markdown("---")
if st.button("🔄 手動更新知識庫"):
st.session_state.knowledge = fetch_all_knowledge()
st.success("所有領域資料已重新抓取!")
# 主介面
st.title("📚 AI 演算法與論文社群助手")
st.caption("知識庫涵蓋 LLM、NLP、Speech、CV。歡迎直接提問!")
def get_gemini_response(user_input):
system_instruction = f"""
你是一位專業的 AI 技術分析專家。
以下是我從 GitHub 同步的四個領域 (LLM, NLP, Speech, CV) 的最新技術簡介與重點:
---
{st.session_state.knowledge}
---
請嚴格基於上述提供的資訊來回答使用者的問題。
如果使用者詢問了上述資訊中沒有涵蓋的細節,請明確告知:「目前這份懶人包中尚未收錄關於此問題的詳細資訊」。
"""
model = genai.GenerativeModel(
model_name="gemini-1.5-pro",
system_instruction=system_instruction
)
chat = model.start_chat(history=[])
response = chat.send_message(user_input)
return response.text
# 顯示對話紀錄
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.markdown(message["content"])
# 接收輸入
if prompt := st.chat_input("例如:請幫我比較一下 CV 跟 NLP 目前常用的技術?"):
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
with st.chat_message("assistant"):
response_text = get_gemini_response(prompt)
st.markdown(response_text)
st.session_state.messages.append({"role": "assistant", "content": response_text})