Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,7 +6,6 @@ import requests
 import os
 import sys
 import subprocess
-import time
 from openai import OpenAI
 from rank_bm25 import BM25Okapi
 from sklearn.metrics.pairwise import cosine_similarity
@@ -28,10 +27,10 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
-# --- 注入自定义 CSS (核心美化部分) ---
+# --- 注入自定义 CSS (保持之前的审美) ---
 st.markdown("""
 <style>
-    /* 1. 整体背景设定 - 深空黑 */
+    /* 1. 整体背景 - 深空黑 */
     .stApp {
         background-color: #050505;
         background-image: radial-gradient(circle at 50% 0%, #1a1f35 0%, #050505 60%);
@@ -44,7 +43,7 @@ st.markdown("""
     footer {visibility: hidden;}
     header {visibility: hidden;}
 
-    /* 3. 聊天气泡 - 毛玻璃特效 */
+    /* 3. 聊天气泡 */
     [data-testid="stChatMessage"] {
         background: rgba(255, 255, 255, 0.03);
         border: 1px solid rgba(255, 255, 255, 0.08);
@@ -52,13 +51,9 @@ st.markdown("""
         backdrop-filter: blur(12px);
         box-shadow: 0 4px 20px rgba(0,0,0,0.2);
         padding: 1.2rem;
-        transition: transform 0.2s;
-    }
-    [data-testid="stChatMessage"]:hover {
-        border-color: rgba(41, 181, 232, 0.3);
     }
     
-    /* 用户气泡特别样式 */
+    /* 用户气泡 */
     [data-testid="stChatMessage"][data-testid="user"] {
         background: rgba(41, 181, 232, 0.1);
         border-color: rgba(41, 181, 232, 0.2);
@@ -82,7 +77,7 @@ st.markdown("""
         letter-spacing: -1px;
     }
     
-    /* 5. 快捷按钮 (Starter Pills) */
+    /* 5. 快捷按钮 */
     div.stButton > button {
         background: rgba(255,255,255,0.05);
         color: #aaa;
@@ -100,7 +95,7 @@ st.markdown("""
         transform: translateY(-2px);
     }
 
-    /* 6. 输入框美化 */
+    /* 6. 输入框 */
     .stChatInputContainer textarea {
         background-color: #0f1115 !important;
         border: 1px solid #333 !important;
@@ -108,29 +103,23 @@ st.markdown("""
         border-radius: 12px !important;
     }
     
-    /* 7. 证据栏 Expander 美化 */
+    /* 7. Expander */
     .streamlit-expanderHeader {
         background-color: rgba(255,255,255,0.02);
         border: 1px solid rgba(255,255,255,0.05);
         border-radius: 8px;
         color: #bbb;
     }
-    .evidence-highlight {
-        color: #29B5E8;
-        font-weight: bold;
-        font-family: monospace;
-    }
 </style>
 """, unsafe_allow_html=True)
 
-# ================= 2. 核心逻辑（保留原本的鲁棒设计） =================
+# ================= 2. 核心逻辑（数据与RAG） =================
 
 if not API_KEY:
     st.error("⚠️ 未检测到 API Key。请在 Settings -> Secrets 中配置 `SILICONFLOW_API_KEY`。")
     st.stop()
 
 def download_with_curl(url, output_path):
-    """使用系统 curl 命令下载，模拟浏览器行为绕过 403"""
     try:
         cmd = [
             "curl", "-L", 
@@ -140,8 +129,7 @@ def download_with_curl(url, output_path):
             url
         ]
         result = subprocess.run(cmd, capture_output=True, text=True)
-        if result.returncode != 0:
-            raise Exception(f"Curl failed: {result.stderr}")
+        if result.returncode != 0: raise Exception(f"Curl failed: {result.stderr}")
         return True
     except Exception as e:
         print(f"Curl download error: {e}")
@@ -155,15 +143,11 @@ def get_data_file_path():
         os.path.join("..", DATA_FILENAME), "/tmp/" + DATA_FILENAME
     ]
     for path in possible_paths:
-        if os.path.exists(path):
-            return path
+        if os.path.exists(path): return path
             
-    # 下载逻辑
     download_target = "/app/" + DATA_FILENAME
-    try:
-        os.makedirs(os.path.dirname(download_target), exist_ok=True)
-    except:
-        download_target = "/tmp/" + DATA_FILENAME # 回退到 tmp
+    try: os.makedirs(os.path.dirname(download_target), exist_ok=True)
+    except: download_target = "/tmp/" + DATA_FILENAME
 
     status_container = st.empty()
     status_container.info("📡 正在接入神经元网络... (下载核心数据中)")
@@ -172,84 +156,62 @@ def get_data_file_path():
         status_container.empty()
         return download_target
         
-    # 备用下载
     try:
         headers = {'User-Agent': 'Mozilla/5.0'}
         r = requests.get(DATA_URL, headers=headers, stream=True)
         r.raise_for_status()
         with open(download_target, 'wb') as f:
-            for chunk in r.iter_content(chunk_size=8192):
-                f.write(chunk)
+            for chunk in r.iter_content(chunk_size=8192): f.write(chunk)
         status_container.empty()
         return download_target
     except Exception as e:
         st.error(f"❌ 数据链路中断。Error: {e}")
         st.stop()
 
-class RerankClient:
-    def __init__(self, api_base, api_key, model):
-        self.api_url = f"{api_base}/rerank" 
-        self.headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
-        self.model = model
-
-    def rerank(self, query: str, documents: list, top_n: int = 5):
-        if not documents: return []
-        payload = {"model": self.model, "query": query, "documents": documents, "top_n": top_n}
-        try:
-            response = requests.post(self.api_url, headers=self.headers, json=payload, timeout=10)
-            if response.status_code == 200:
-                return response.json().get('results', [])
-            return [{"index": i, "relevance_score": 0.0} for i in range(len(documents))]
-        except:
-            return [{"index": i, "relevance_score": 0.0} for i in range(len(documents))]
-
 class FullRetriever:
     def __init__(self, parquet_path):
-        try:
-            self.df = pd.read_parquet(parquet_path)
-        except Exception as e:
-            st.error(f"Memory Matrix Load Failed: {e}")
-            st.stop()
-            
+        try: self.df = pd.read_parquet(parquet_path)
+        except Exception as e: st.error(f"Memory Matrix Load Failed: {e}"); st.stop()
         self.documents = self.df['content'].tolist()
         self.embeddings = np.stack(self.df['embedding'].values)
         self.bm25 = BM25Okapi([jieba.lcut(str(d).lower()) for d in self.documents])
         self.client = OpenAI(base_url=API_BASE, api_key=API_KEY)
-        self.reranker = RerankClient(API_BASE, API_KEY, RERANK_MODEL)
+        # Reranker 初始化移到这里，减少重复调用
+        self.rerank_headers = {"Content-Type": "application/json", "Authorization": f"Bearer {API_KEY}"}
+        self.rerank_url = f"{API_BASE}/rerank"
 
     def _get_emb(self, q):
-        try:
-            return self.client.embeddings.create(model=EMBEDDING_MODEL, input=[q]).data[0].embedding
-        except:
-            return [0.0] * 1024
+        try: return self.client.embeddings.create(model=EMBEDDING_MODEL, input=[q]).data[0].embedding
+        except: return [0.0] * 1024
 
     def hybrid_search(self, query: str, top_k=5):
         # 1. Vector
         q_emb = self._get_emb(query)
         vec_scores = cosine_similarity([q_emb], self.embeddings)[0]
         vec_idx = np.argsort(vec_scores)[-100:][::-1]
-        
         # 2. Keyword
         kw_idx = np.argsort(self.bm25.get_scores(jieba.lcut(query.lower())))[-100:][::-1]
-        
         # 3. RRF Fusion
         fused = {}
         for r, i in enumerate(vec_idx): fused[i] = fused.get(i, 0) + 1/(60+r+1)
         for r, i in enumerate(kw_idx): fused[i] = fused.get(i, 0) + 1/(60+r+1)
-        
         c_idxs = [x[0] for x in sorted(fused.items(), key=lambda x:x[1], reverse=True)[:50]]
         c_docs = [self.documents[i] for i in c_idxs]
         
         # 4. Rerank
-        results = self.reranker.rerank(query, c_docs, top_n=top_k)
-        
+        try:
+            payload = {"model": RERANK_MODEL, "query": query, "documents": c_docs, "top_n": top_k}
+            resp = requests.post(self.rerank_url, headers=self.rerank_headers, json=payload, timeout=10)
+            results = resp.json().get('results', [])
+        except:
+            results = [{"index": i, "relevance_score": 0.0} for i in range(len(c_docs))][:top_k]
+            
         final_res = []
         context = ""
         for i, item in enumerate(results):
             orig_idx = c_idxs[item['index']]
             row = self.df.iloc[orig_idx]
             final_res.append({
-                "rank": i+1,
                 "score": item['relevance_score'],
                 "filename": row['filename'],
                 "content": row['content']
@@ -262,17 +224,16 @@ def load_engine():
     real_path = get_data_file_path()
     return FullRetriever(real_path)
 
-# ================= 3. UI 主逻辑 =================
+# ================= 3. UI 主程序 =================
 
 def main():
-    # 自定义 Header 区域
     st.markdown("""
     <div class="custom-header">
         <div style="font-size: 3rem;">🌌</div>
         <div>
             <div class="glitch-text">COMSOL DARK EXPERT</div>
             <div style="color: #666; font-size: 0.9rem; letter-spacing: 1px;">
-                NEURAL SIMULATION ASSISTANT <span style="color:#29B5E8">V4.0</span>
+                NEURAL SIMULATION ASSISTANT <span style="color:#29B5E8">V4.1 Fixed</span>
             </div>
         </div>
     </div>
@@ -280,78 +241,75 @@ def main():
 
     retriever = load_engine()
 
-    # --- 侧边栏 ---
     with st.sidebar:
         st.markdown("### ⚙️ 控制台")
-        top_k = st.slider("Retrieval Depth (检索深度)", 1, 10, 4)
-        temp = st.slider("Temperature (发散度)", 0.0, 1.0, 0.3)
+        top_k = st.slider("检索深度", 1, 10, 4)
+        temp = st.slider("发散度", 0.0, 1.0, 0.3)
         st.markdown("---")
-        
-        if st.button("🗑️ 格式化内存 (Clear History)", use_container_width=True):
+        if st.button("🗑️ 清空记忆 (Clear)", use_container_width=True):
             st.session_state.messages = []
             st.session_state.current_refs = []
             st.rerun()
-            
-        st.markdown("""
-        <div style="margin-top: 2rem; padding: 1rem; background: rgba(255,255,255,0.05); border-radius: 8px; font-size: 0.8rem; color: #888;">
-        <strong>System Status:</strong><br>
-        ✅ Vector Engine: Online<br>
-        ✅ Reranker: Active<br>
-        ✅ Neural Link: Stable
-        </div>
-        """, unsafe_allow_html=True)
 
-    # --- 初始化 Session ---
     if "messages" not in st.session_state: st.session_state.messages = []
     if "current_refs" not in st.session_state: st.session_state.current_refs = []
 
-    # --- 布局：左聊天，右证据 ---
     col_chat, col_evidence = st.columns([0.65, 0.35], gap="large")
 
+    # ------------------ 处理输入源 ------------------
+    # 我们定义一个变量 user_input，不管它来自按钮还是输入框
+    user_input = None
+
     with col_chat:
-        # 1. 如果没有历史消息，显示快捷提问胶囊 (Starter Pills)
+        # 1. 如果历史为空，显示快捷按钮
         if not st.session_state.messages:
             st.markdown("##### 💡 初始化提问序列 (Starter Sequence)")
             c1, c2, c3 = st.columns(3)
+            # 点击按钮直接赋值给 user_input
             if c1.button("🌊 流固耦合接口设置"):
-                st.session_state.messages.append({"role": "user", "content": "怎么设置流固耦合接口？"})
-                st.rerun()
-            if c2.button("⚡ 低频电磁场网格"):
-                st.session_state.messages.append({"role": "user", "content": "低频电磁场网格划分有哪些技巧？"})
-                st.rerun()
-            if c3.button("📉 求解器不收敛"):
-                st.session_state.messages.append({"role": "user", "content": "求解器不收敛通常怎么解决？"})
-                st.rerun()
-
+                user_input = "怎么设置流固耦合接口？"
+            elif c2.button("⚡ 低频电磁场网格"):
+                user_input = "低频电磁场网格划分有哪些技巧？"
+            elif c3.button("📉 求解器不收敛"):
+                user_input = "求解器不收敛通常怎么解决？"
+        
         # 2. 渲染历史消息
         for msg in st.session_state.messages:
             with st.chat_message(msg["role"]):
                 st.markdown(msg["content"])
 
-        # 3. 输入处理
-        if prompt := st.chat_input("输入指令或物理参数问题..."):
-            st.session_state.messages.append({"role": "user", "content": prompt})
-            with st.chat_message("user"): st.markdown(prompt)
+        # 3. 处理底部输入框 (如果有按钮输入，这里会被跳过，因为 user_input 已经有值了)
+        if not user_input:
+            user_input = st.chat_input("输入指令或物理参数问题...")
+
+    # ------------------ 统一处理消息追加 ------------------
+    if user_input:
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        # 强制刷新以立即在 UI 上显示用户的提问（对于按钮点击尤为重要）
+        st.rerun()
 
-            # RAG 检索
+    # ------------------ 统一触发生成 (修复的核心) ------------------
+    # 检查：如果有消息，且最后一条是 User 发的，说明需要 Assistant 回答
+    if st.session_state.messages and st.session_state.messages[-1]["role"] == "user":
+        
+        # 获取最后一条用户消息
+        last_query = st.session_state.messages[-1]["content"]
+        
+        with col_chat: # 确保在聊天栏显示
             with st.spinner("🔍 正在扫描向量空间..."):
-                refs, context = retriever.hybrid_search(prompt, top_k=top_k)
+                refs, context = retriever.hybrid_search(last_query, top_k=top_k)
                 st.session_state.current_refs = refs
             
-            # 增强版 System Prompt
             system_prompt = f"""你是一个COMSOL高级仿真专家。请基于提供的文档回答问题。
-            
             要求：
             1. 语气专业、客观，逻辑严密。
             2. 涉及物理公式时，**必须**使用 LaTeX 格式（例如 $E = mc^2$）。
             3. 涉及步骤或参数对比时，优先使用 Markdown 列表或表格。
-            4. 严禁编造文档中不存在的参数。
             
             参考文档：
             {context}
             """
             
-            # 生成回答
             with st.chat_message("assistant"):
                 resp_cont = st.empty()
                 full_resp = ""
@@ -360,7 +318,7 @@ def main():
                 try:
                     stream = client.chat.completions.create(
                         model=GEN_MODEL_NAME,
-                        messages=[{"role": "system", "content": system_prompt}] + st.session_state.messages[-6:],
+                        messages=[{"role": "system", "content": system_prompt}] + st.session_state.messages[-6:], # 除去当前的System
                         temperature=temp,
                         stream=True
                     )
@@ -368,16 +326,13 @@ def main():
                         txt = chunk.choices[0].delta.content
                         if txt:
                             full_resp += txt
-                            # 打字机效果 + 光标
                             resp_cont.markdown(full_resp + " ▌")
-                    resp_cont.markdown(full_resp) # 移除光标
+                    resp_cont.markdown(full_resp)
                     st.session_state.messages.append({"role": "assistant", "content": full_resp})
-                    # 强制刷新以更新右侧证据栏（虽然Streamlit通常会自动更新，但显式rerun更稳）
-                    # st.rerun() -> 这里不需要rerun，因为session_state已更新，下一轮循环会渲染
                 except Exception as e:
                     st.error(f"Neural Generation Failed: {e}")
 
-    # --- 右侧证据栏美化 ---
+    # ------------------ 渲染右侧证据栏 ------------------
     with col_evidence:
         st.markdown("### 📚 神经记忆 (Evidence)")
         if st.session_state.current_refs:
@@ -385,7 +340,6 @@ def main():
                 score = ref['score']
                 score_color = "#00ff41" if score > 0.6 else "#ffb700" if score > 0.4 else "#ff003c"
                 
-                # 使用 Expander 折叠详细内容
                 with st.expander(f"📄 Doc {i+1}: {ref['filename'][:20]}...", expanded=(i==0)):
                     st.markdown(f"""
                     <div style="margin-bottom:5px;">