Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,342 +1,667 @@
 import streamlit as st
+
 import os
+
 import io
+
 import numpy as np
+
 import faiss
+
 import uuid
+
 import time
-import google.generativeai as genai
-from google.generativeai.types import GenerationConfig, HarmCategory, HarmBlockThreshold # 引入必要的型別
+
+import google.generativeai as genai  # <--- 新增 Google SDK
+
+
 
 # === RAG 相關套件 ===
-import torch
+
+# 這裡保留 Torch 和 HuggingFaceEmbeddings 是為了向量化 (Embedding)，這部分吃資源很少
+
+import torch 
+
 from langchain_community.embeddings import HuggingFaceEmbeddings
+
 from langchain_core.documents import Document
+
 from langchain_community.vectorstores import FAISS
+
 from langchain_community.vectorstores.utils import DistanceStrategy
+
 from langchain_community.docstore.in_memory import InMemoryDocstore
 
+
+
 # 嘗試匯入 pypdf
+
 try:
+
     import pypdf
+
 except ImportError:
+
     pypdf = None
 
+
+
 # --- 頁面設定 ---
+
 st.set_page_config(page_title="Cybersecurity AI Assistant (Gemini RAG)", page_icon="🛡️", layout="wide")
-st.title("🛡️ Gemini-2.5-Flash RAG 資安分析助理")
+
+st.title("🛡️ Gemini-1.5-Flash with FAISS RAG")
+
 st.markdown("已啟用：**IndexFlatIP** + **L2 正規化** + **Google Gemini API**")
 
+
+
 # --- 側邊欄設定 ---
+
 with st.sidebar:
+
     st.header("⚙️ 設定")
+
     
-    # Google API Key 輸入
+
+    # 改為 Google API Key
+
     default_key = os.getenv("GOOGLE_API_KEY", "")
+
     google_api_key = st.text_input("Google API Key", value=default_key, type="password")
+
     
+
     if not google_api_key:
+
         st.warning("請輸入 Google API Key 以繼續。")
+
     
+
     st.divider()
+
     st.subheader("📂 上傳分析檔案 (建立 RAG 庫)")
+
     uploaded_file = st.file_uploader("上傳 Logs/PDF/Code", type=['txt', 'py', 'log', 'csv', 'md', 'json', 'pdf'])
+
     
+
     st.divider()
+
     st.subheader("🔍 RAG 檢索設定")
+
     similarity_threshold = st.slider(
+
         "📐 Cosine Similarity 門檻", 
+
         0.0, 1.0, 0.4, 0.01,
+
         help="數值越大越相似。一般建議 0.4~0.7"
+
     )
+
     
+
     st.divider()
+
     st.subheader("模型參數")
-    # 調整 System Prompt 預設值，鼓勵模型提供結構化資安分析
-    system_prompt = st.text_area(
-        "System Prompt", 
-        value="You are a Tier 3 Senior Security Analyst. Use the retrieved context to answer the user's question. Specifically, follow the strict analysis framework provided by the user (Ransomware Kill Chain, Timeline Reconstruction) and respond in Traditional Chinese. If no malicious indicators are found, state clearly.", 
-        height=100
-    )
-    # 預設 Max Output Tokens 調整到 4096，以避免中斷
-    max_output_tokens = st.slider("Max Output Tokens", 128, 8192, 4096, 128, help="調高此值可避免回應被截斷 (MAX_TOKENS 錯誤)。")
-    temperature = st.slider("Temperature", 0.0, 2.0, 0.1, 0.1, help="資安分析建議使用極低的 Temperature (0.1-0.3)。")
+
+    system_prompt = st.text_area("System Prompt", value="You are a Senior Security Analyst. Use the retrieved context to answer the user's question. Every claim you make MUST be supported by a specific Event Record ID from the retrieved context.", height=100)
+
+    # Gemini 不需要 max_new_tokens 來限制記憶體，但可以設定輸出上限
+
+    max_output_tokens = st.slider("Max Output Tokens", 128, 8192, 2048, 128)
+
+    temperature = st.slider("Temperature", 0.0, 2.0, 0.1, 0.1)
+
     
+
     st.divider()
+
     if st.button("🗑️ 清除對話紀錄"):
+
         st.session_state.messages = []
+
         st.rerun()
 
+
+
 # --- 初始化 Gemini ---
+
 genai_model = None
+
 if google_api_key:
+
     try:
+
         genai.configure(api_key=google_api_key)
+
         # 使用 Flash 模型，速度快且便宜，適合 RAG 大量文本閱讀
-        genai_model = genai.GenerativeModel('gemini-2.5-flash')
+
+        genai_model = genai.GenerativeModel('gemini-2.5-pro')
+
     except Exception as e:
+
         st.error(f"Gemini 設定失敗: {e}")
 
+
+
 # === Embedding 模型 (保留原本的 Jina 或其他 HF 模型) ===
+
+# Embedding 還是建議用專門的模型，不一定要換成 Google 的 Embedding
+
 @st.cache_resource
+
 def load_embedding_model():
+
     model_kwargs = {
-        'device': 'cpu', 
+
+        'device': 'cpu',  # Embedding 通常 CPU 夠用，若有 GPU 也可改 cuda
+
         'trust_remote_code': True 
+
     }
+
     encode_kwargs = {
+
         'normalize_embeddings': False 
+
     }
+
     return HuggingFaceEmbeddings(
+
         model_name="jinaai/jina-embeddings-v2-base-code",
+
         model_kwargs=model_kwargs,
+
         encode_kwargs=encode_kwargs
+
     )
 
+
+
 with st.spinner("正在載入 Embedding 模型..."):
+
     embedding_model = load_embedding_model()
 
+
+
 # === 建立向量庫 (Strict Cosine) - 邏輯維持不變 ===
+
 def process_file_to_faiss(uploaded_file):
+
     text_content = ""
+
     try:
+
         if uploaded_file.type == "application/pdf":
+
             if pypdf:
+
                 pdf_reader = pypdf.PdfReader(uploaded_file)
+
                 for page in pdf_reader.pages:
+
                     text_content += page.extract_text() + "\n"
+
             else:
+
                 return None, "PDF library missing"
+
         else:
+
             stringio = io.StringIO(uploaded_file.getvalue().decode("utf-8"))
+
             text_content = stringio.read()
+
         
+
         if not text_content.strip():
+
             return None, "File is empty"
 
+
+
         # 簡單切分
+
         events = [e + "</Event>" for e in text_content.split("</Event>") if e.strip()]
+
         if len(events) <= 1:
+
               events = [line for line in text_content.split("\n") if line.strip()]
+
         
+
         docs = [Document(page_content=e) for e in events]
+
         
+
         if not docs:
+
             return None, "No documents created"
 
+
+
         embeddings = embedding_model.embed_documents([d.page_content for d in docs])
+
         embeddings_np = np.array(embeddings).astype("float32")
+
         faiss.normalize_L2(embeddings_np) 
+
         
+
         dimension = embeddings_np.shape[1]
+
         index = faiss.IndexFlatIP(dimension)
+
         index.add(embeddings_np)
+
         
+
         doc_ids = [str(uuid.uuid4()) for _ in range(len(docs))]
+
         docstore = InMemoryDocstore({_id: doc for _id, doc in zip(doc_ids, docs)})
+
         index_to_docstore_id = {i: _id for i, _id in enumerate(doc_ids)}
+
         
+
         vector_store = FAISS(
+
             embedding_function=embedding_model,
+
             index=index,
+
             docstore=docstore,
+
             index_to_docstore_id=index_to_docstore_id,
+
             distance_strategy=DistanceStrategy.COSINE
+
         )
+
         
-        return vector_store, f"建立了 {len(docs)} 個日誌片段。"
+
+        return vector_store, f"{len(docs)} chunks created."
+
     except Exception as e:
-        return None, f"錯誤: {str(e)}"
+
+        return None, f"Error: {str(e)}"
+
+
 
 # === 檔案處理邏輯 ===
+
 if uploaded_file:
+
     file_key = f"vs_{uploaded_file.name}_{uploaded_file.size}"
+
     
+
     if "current_file_key" not in st.session_state or st.session_state.current_file_key != file_key:
+
         with st.spinner("偵測到新檔案，正在更新知識庫..."):
+
             vs, msg = process_file_to_faiss(uploaded_file)
+
             if vs:
+
                 st.session_state.vector_store = vs
+
                 st.session_state.current_file_key = file_key
+
                 st.toast(f"知識庫已更新！{msg}", icon="✅")
+
             else:
+
                 st.error(msg)
+
 else:
+
     if "vector_store" in st.session_state:
+
         del st.session_state.vector_store
+
         st.info("檔案已移除，已清除知識庫，回到一般模式。")
+
     if "current_file_key" in st.session_state:
-        if 'vector_store' not in st.session_state: # 避免重複刪除
-            del st.session_state.current_file_key
+
+        del st.session_state.current_file_key
+
+
 
 # === 顯示對話歷史 ===
+
 if "messages" not in st.session_state:
+
     st.session_state.messages = []
 
+
+
 for idx, message in enumerate(st.session_state.messages):
+
     with st.chat_message(message["role"]):
+
         st.markdown(message["content"])
+
         if message.get("context"):
-            # 確保只顯示最近一次對話的 expander
-            if idx == len(st.session_state.messages) - 1:
-                is_expanded = True
-            else:
-                is_expanded = False
-                
-            with st.expander("查看參考片段", expanded=is_expanded):
-                st.code(message["context"], language="log")
+
+            with st.expander(f"查看參考片段 (Turn {idx})"):
+
+                st.code(message["context"])
+
                 st.download_button(
+
                     label="📥 下載此參考內容 (.txt)",
+
                     data=message["context"],
+
                     file_name=f"rag_context_{idx}.txt",
+
                     mime="text/plain",
+
                     key=f"dl_btn_{idx}"
+
                 )
 
+
+
 # === Search 函數 ===
+
 def faiss_cosine_search_all(vector_store, query, threshold):
+
     q_emb = embedding_model.embed_query(query)
+
     q_emb = np.array([q_emb]).astype("float32")
+
     faiss.normalize_L2(q_emb)
+
     
+
     index = vector_store.index
+
     D, I = index.search(q_emb, k=index.ntotal)
+
     
+
     selected = []
-    # 這裡只取相似度高於門檻的片段
+
     for score, idx in zip(D[0], I[0]):
+
         if idx == -1: continue
+
         if score >= threshold:
+
             doc_id = vector_store.index_to_docstore_id[idx]
+
             doc = vector_store.docstore.search(doc_id)
+
             selected.append((doc, score))
+
             
+
     selected.sort(key=lambda x: x[1], reverse=True)
+
     return selected
 
+
+
 # === Gemini 產生回答 ===
+
 def generate_rag_response_gemini(prompt, history, sys_prompt, vector_store=None, threshold=0.5):
+
     context_text = ""
+
+    top_k_selected = []
+
     
+
     # 1. 檢索
+
     if vector_store:
-        # 為了資安分析，我們需要擷取所有相關的 Log，所以將 threshold 作為篩選標準
+
         selected = faiss_cosine_search_all(vector_store, prompt, threshold)
-        
+
         if selected:
-            # 取前 50 個片段，以利用 Gemini-Flash 的大上下文視窗
-            top_k_selected = selected[:50]
+
+            top_k_selected = selected
+
+            # 取前 30 個或更多 (Gemini Context Window 很大，可以塞多一點)
+
             retrieved_contents = [
+
                 f"--- Chunk (sim={score:.3f}) ---\n{doc.page_content}"
-                for (doc, score) in top_k_selected
+
+                for i, (doc, score) in enumerate(top_k_selected[:30]) 
+
             ]
+
             context_text = "\n".join(retrieved_contents)
 
+
+
     # 2. 構建 Prompt
-    # 我們將系統指令與上下文合併，作為使用者訊息的第一部分，以確保指令被嚴格遵循
+
     if context_text:
+
         full_user_input = f"""
-System Instruction (CRITICAL: Adhere to the following framework and respond in Traditional Chinese):
-{sys_prompt}
+
+System Instruction: {sys_prompt}
+
+
 
 === RETRIEVED CONTEXT (Cosine ≥ {threshold}) ===
+
 {context_text}
+
 === END CONTEXT ===
 
+
+
 Question: {prompt}
-Analyze the question based strictly on the context and output using the required 4-part Chinese structure.
+
+Answer the question strictly based on the provided context.
+
 """
+
     else:
+
         full_user_input = f"""
+
 System Instruction: {sys_prompt}
 
+
+
 Question: {prompt}
+
 """
 
+
+
     # 3. 轉換歷史訊息格式 (Streamlit -> Gemini)
+
+    # Gemini 格式: [{'role': 'user', 'parts': [...]}, {'role': 'model', 'parts': [...]}]
+
     gemini_history = []
+
     for msg in history:
-        # 由於我們在 full_user_input 塞入了 System Prompt 和 Context，這裡只傳遞純對話以避免上下文重複
+
         role = "user" if msg["role"] == "user" else "model"
-        gemini_history.append({"role": role, "parts": [{"text": msg["content"]}]})
+
+        # 濾除非文字內容 (簡單處理)
+
+        content_text = msg["content"]
+
+        # 這裡不把之前的 context 重複塞入歷史，避免 context window 爆炸或混淆，僅傳遞純對話
+
+        gemini_history.append({"role": role, "parts": [content_text]})
+
+
 
     # 4. 呼叫 Gemini
+
     try:
+
+        chat = genai_model.start_chat(history=gemini_history)
+
+        
+
         # 設定生成參數
-        generation_config = GenerationConfig(
+
+        generation_config = genai.types.GenerationConfig(
+
             candidate_count=1,
+
             max_output_tokens=max_output_tokens,
+
             temperature=temperature,
+
         )
+
         
-        # 安全設定 (設為 BLOCK_NONE，這是解決資安敏感內容被阻擋的關鍵)
+
+        # 安全設定 (設為 BLOCK_NONE 以避免資安 Log 被誤判為有害內容)
+
         safety_settings = [
-            {"category": HarmCategory.HARM_CATEGORY_HARASSMENT, "threshold": HarmBlockThreshold.BLOCK_NONE},
-            {"category": HarmCategory.HARM_CATEGORY_HATE_SPEECH, "threshold": HarmBlockThreshold.BLOCK_NONE},
-            {"category": HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, "threshold": HarmBlockThreshold.BLOCK_NONE},
-            {"category": HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, "threshold": HarmBlockThreshold.BLOCK_NONE},
+
+            {
+
+                "category": "HARM_CATEGORY_HARASSMENT",
+
+                "threshold": "BLOCK_NONE",
+
+            },
+
+            {
+
+                "category": "HARM_CATEGORY_HATE_SPEECH",
+
+                "threshold": "BLOCK_NONE",
+
+            },
+
+            {
+
+                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+
+                "threshold": "BLOCK_NONE",
+
+            },
+
+            {
+
+                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+
+                "threshold": "BLOCK_NONE",
+
+            },
+
         ]
 
-        chat = genai_model.start_chat(history=gemini_history)
-        
+
+
         response = chat.send_message(
+
             full_user_input, 
+
             generation_config=generation_config,
+
             safety_settings=safety_settings
+
         )
-        
-        # 檢查是否有安全阻擋或錯誤
-        if response.prompt_feedback.block_reason or not response.candidates:
-             # 安全性阻擋的錯誤處理
-            reason = response.prompt_feedback.block_reason.name if response.prompt_feedback.block_reason else "Unknown"
-            return f"Gemini API 錯誤: 由於安全原因，回應被阻擋。原因: {reason}", context_text
-        
+
         return response.text, context_text
 
+
+
     except Exception as e:
-        return f"Gemini API 錯誤: {str(e)}", context_text
+
+        return f"Gemini API Error: {str(e)}", context_text
+
+
 
 # === 處理使用者輸入 ===
+
 if prompt := st.chat_input("請輸入問題..."):
-    if not google_api_key:
+
+    if not genai_model:
+
         st.error("請先輸入有效的 Google API Key")
-    elif not genai_model:
-        st.error("Gemini 模型初始化失敗，請檢查 API Key")
+
     else:
+
         vs = st.session_state.get("vector_store", None)
+
         display_prompt = prompt
+
+        if vs:
+
+            display_prompt = f"🔍 **[RAG]** {prompt}"
+
         
-        st.chat_message("user").markdown(f"🔍 **[RAG]** {prompt}" if vs else prompt)
+
+        st.chat_message("user").markdown(display_prompt)
+
         
+
         with st.chat_message("assistant"):
+
             msg_placeholder = st.empty()
+
             
-            with st.spinner("Gemini Thinking... 正在進行日誌分析與 RAG 檢索"):
+
+            with st.spinner("Gemini Thinking..."):
+
                 response, retrieved_ctx = generate_rag_response_gemini(
+
                     prompt, 
+
                     st.session_state.messages, 
+
                     system_prompt, 
+
                     vector_store=vs, 
+
                     threshold=similarity_threshold,
+
                 )
+
                 
+
                 msg_placeholder.markdown(response)
+
                 
+
                 if retrieved_ctx:
-                    # 再次顯示擴展器，確保當前回合的參考資料可見
+
                     with st.expander("查看檢索到的參考片段"):
-                        st.code(retrieved_ctx, language="log")
+
+                        st.code(retrieved_ctx)
+
                         st.download_button(
+
                             label="📥 下載此參考內容 (.txt)",
+
                             data=retrieved_ctx,
+
                             file_name=f"rag_context_current.txt",
+
                             mime="text/plain"
+
                         )
 
-            # 更新歷史 (將原始 prompt 和回應存入 session state)
-            st.session_state.messages.append({"role": "user", "content": f"🔍 **[RAG]** {prompt}" if vs else prompt})
+
+
+            # 更新歷史
+
+            st.session_state.messages.append({"role": "user", "content": display_prompt})
+
             st.session_state.messages.append({
+
                 "role": "assistant", 
+
                 "content": response,
+
                 "context": retrieved_ctx 
+
             })