Update app.py

app.py

@@ -2,90 +2,77 @@ import gradio as gr
 from huggingface_hub import InferenceClient
 import re
 
-# 初始化 - 建议使用 72B 级别模型以获得更好的审计能力
-client = InferenceClient("Qwen/Qwen2.5-Coder-32B-Instruct")
+# 备选模型列表，增加稳定性
+MODELS = {
+    "Qwen 2.5 Coder (专业审计)": "Qwen/Qwen2.5-Coder-32B-Instruct",
+    "Llama 3.1 (逻辑推理)": "meta-llama/Llama-3.1-70B-Instruct",
+    "Mistral (快速响应)": "mistralai/Mistral-7B-Instruct-v0.3"
+}
 
-# --- [A] 核心审计逻辑 ---
-def run_security_scan(code):
-    if not code or len(code) < 5:
-        return "⚠️ 等待输入代码..."
-    findings = []
-    if re.search(r"exec\(|eval\(", code): findings.append("🔴 发现注入风险: eval/exec")
-    if re.search(r"select.*from.*where", code, re.I): findings.append("🟠 发现 SQL 风险")
-    if re.search(r"api_key|password|secret", code, re.I): findings.append("🟡 发现硬编码泄露")
-    return " \n".join(findings) if findings else "✅ 未发现基础漏洞"
-
-# --- [C] 文本处理逻辑 ---
-def process_file(file):
-    if file is None: return ""
-    try:
-        with open(file.name, "r", encoding="utf-8") as f:
-            return f"\n[文件上下文]:\n{f.read()[:800]}"
-    except:
-        return ""
+def ghost_chat(message, history, system_message, model_name, file_obj):
+    # 自动选择模型
+    repo_id = MODELS.get(model_name, MODELS["Qwen 2.5 Coder (专业审计)"])
+    client = InferenceClient(repo_id)
+    
+    # 文件处理
+    file_context = ""
+    if file_obj is not None:
+        try:
+            with open(file_obj.name, "r", encoding="utf-8") as f:
+                file_context = f"\n[File Context]:\n{f.read()[:1000]}"
+        except: pass
 
-# --- 核心交互引擎 ---
-def ghost_chat(message, history, system_message, file_obj):
-    file_context = process_file(file_obj)
     messages = [{"role": "system", "content": f"{system_message}\n{file_context}"}]
-    
     for user_msg, assistant_msg in history:
         messages.append({"role": "user", "content": user_msg})
         messages.append({"role": "assistant", "content": assistant_msg})
-    
     messages.append({"role": "user", "content": message})
-    
+
     response = ""
-    # 模拟连接感
-    yield history + [[message, "📡 连接节点..."]]
+    yield history + [[message, "🛰️ 正在链接卫星链路..."]]
     
-    for msg_chunk in client.chat_completion(messages, stream=True, max_tokens=2048):
-        token = msg_chunk.choices[0].delta.content
-        if token:
-            response += token
-            yield history + [[message, response]]
+    try:
+        for msg_chunk in client.chat_completion(messages, stream=True, max_tokens=2048):
+            token = msg_chunk.choices[0].delta.content
+            if token:
+                response += token
+                yield history + [[message, response]]
+    except Exception as e:
+        yield history + [[message, f"❌ 系统崩溃: {str(e)}"]]
 
-# --- [B] 样式定义 ---
-terminal_css = """
-.gradio-container { background: #020202 !important; color: #00ff41 !important; }
-#title-text { text-align: center; color: #00ff41; text-shadow: 0 0 5px #00ff41; }
-.message.user { border-left: 3px solid #00ff41 !important; }
+# --- 极简黑盒样式 ---
+style = """
+.gradio-container { background: #000 !important; color: #0f0 !important; }
+.message.user { border: 1px solid #0f0 !important; }
+.tabs { border: none !important; }
 """
 
-# --- 构建 UI ---
-with gr.Blocks(fill_height=True) as demo:
-    gr.Markdown("# 💀 GHOST-PROTOCOL v7.0", elem_id="title-text")
-
+with gr.Blocks(fill_height=True, css=style) as demo:
+    gr.Markdown("## 🌌 INTERSTELLAR CORE v8.0", textAlign="center")
+    
     with gr.Row():
-        # 工具栏
-        with gr.Column(scale=1, min_width=250):
-            with gr.Accordion("🛡️ 审计器", open=True):
-                code_box = gr.Textbox(placeholder="粘贴代码...", lines=4, label=None)
-                scan_btn = gr.Button("RUN", variant="primary")
-                scan_out = gr.Textbox(label="Result", interactive=False)
-                scan_btn.click(run_security_scan, [code_box], scan_out)
+        # 左侧控制面板
+        with gr.Column(scale=1):
+            with gr.Group():
+                model_dd = gr.Dropdown(list(MODELS.keys()), value="Qwen 2.5 Coder (专业审计)", label="选择核心")
+                sys_msg = gr.Textbox(value="你是一个顶尖的 AI 助手。", label="系统架构")
             
-            with gr.Accordion("📂 注入器", open=True):
-                file_up = gr.File(label=None)
+            with gr.Accordion("🛠️ 战术附件", open=False):
+                file_up = gr.File(label="注入知识库")
+                scan_input = gr.Textbox(label="扫描器输入")
+                scan_btn = gr.Button("RUN ANALYSIS")
+                scan_res = gr.Textbox(label="分析结果")
 
-        # 交互栏
+        # 右侧主控制台
         with gr.Column(scale=3):
-            # 彻底精简 Chatbot 参数，移除所有 show_... 类型的参数
-            chatbot = gr.Chatbot(height=600) 
+            chatbot = gr.Chatbot(height=650)
             with gr.Row():
-                msg = gr.Textbox(placeholder="root@ghost:~$", scale=9, container=False)
-                btn = gr.Button("EXE", scale=1)
-
-            with gr.Accordion("SYSTEM", open=False):
-                sys_msg = gr.Textbox(value="你是一个顶尖红队专家。", label=None)
+                msg = gr.Textbox(placeholder="输入指令...", scale=8, container=False)
+                send = gr.Button("EXE", scale=2)
 
-    # 绑定
-    msg.submit(ghost_chat, [msg, chatbot, sys_msg, file_up], [chatbot])
-    btn.click(ghost_chat, [msg, chatbot, sys_msg, file_up], [chatbot])
+    # 绑定逻辑
+    msg.submit(ghost_chat, [msg, chatbot, sys_msg, model_dd, file_up], [chatbot])
+    send.click(ghost_chat, [msg, chatbot, sys_msg, model_dd, file_up], [chatbot])
 
-# --- 启动 ---
 if __name__ == "__main__":
-    demo.launch(
-        css=terminal_css,
-        theme=gr.themes.Monochrome()
-    )
+    demo.launch()