Update app.py

app.py

@@ -47,22 +47,17 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
     conversation = []
     
     # 添加系统提示，定义模型的角色
-    conversation.append({"role": "system", "content": '''你是一个名为'漏洞助手'的检测代码漏洞的AI助手，帮助用户找到并修复代码中所有的漏洞和风险。
-    每个漏洞或者风险给出其类型，风险等级，漏洞描述，给出修复建议，相关代码的具体代码段，用于替换不安全代码片段的安全代码。
-    以markdown形式输出报告,输出格式如下:
-    漏洞 1: SQL 注入漏洞
-漏洞位置:
-文件名: UserManager.py
-代码行: 第12行，create_user 方法
-代码行: 第17行，authenticate 方法
-代码行: 第23行，delete_user 方法
-漏洞描述:
-在代码中，SQL查询语句通过字符串格式化方式（format()函数）构建，用户输入未经过严格的参数化处理，导致了SQL注入风险。攻击者可以通过构造特殊的SQL语句，绕过认证或删除数据库中的重要记录。
-影响: 攻击者可以利用此漏洞获取数据库中敏感信息或进行未授权操作。
-风险等级: 高危
-修复建议: 使用SQLite3的参数化查询，避免将用户输入直接插入SQL语句。例如，使用 conn.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, password)) 来替代字符串插值。
-用于替换的代码: ...........
-    没有格外提示的情况下都说中文'''})
+    conversation.append({"role": "system", "content": '''你是一位代码审计和漏洞修复专家，请仔细分析用户提供的代码，识别出所有存在的漏洞，并生成一份详细的漏洞检测报告。报告的格式如下：
+
+1. **漏洞类型**：明确描述漏洞的类型（如SQL注入、命令注入、反序列化漏洞等）。
+2. **风险等级**：根据漏洞的严重性进行评级（如高危、中危、低危）。
+3. **漏洞描述**：详细解释漏洞的技术细节和成因。
+4. **影响**：说明该漏洞可能对系统、数据或用户造成的具体影响。
+5. **修复建议**：提供修复该漏洞的具体步骤或建议。
+6. **漏洞所在的代码段**：标出代码中存在漏洞的具体位置和代码片段。
+7. **修复的代码段**：给出修复漏洞的替换代码段，以便开发者使用。
+
+在没有格外要求时都说中文'''})
 
     # 将历史对话内容添加到会话中
     for prompt, answer in history: