Pending changes exported from your codespace

DeepSeekSelfTool.py

@@ -2,6 +2,7 @@ import sys
 import os
 import json
 import requests
+import re
 from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                              QTextEdit, QPushButton, QLabel, QHBoxLayout,
                              QSplitter, QScrollArea, QTabWidget, QFrame,QCheckBox ,QSizePolicy,QComboBox,QFileDialog,QProgressBar)
@@ -9,53 +10,94 @@ from PyQt5.QtCore import Qt, QSize, QThread, pyqtSignal
 from PyQt5.QtGui import QFont, QColor, QPalette, QLinearGradient
 import config
 import glob
+from openai import OpenAI
 from config import THEMES
 
+
 os.environ["QT_IM_MODULE"] = "none"
 
+
 class APIAdapter:
     def __init__(self):
-        self.api_type = config.API_TYPE
+        self.api_type = config.API_TYPE # Supported values: "deepseek", "qwen", "openai", "ollama"
-        if self.api_type == "deepseek":
+
-            self.api_key = config.DEEPSEEK_API_KEY
+        self.api_endpoints = {
-            self.api_endpoint = "https://api.deepseek.com/v1/chat/completions"
+            "deepseek": "https://api.deepseek.com/v1",
-        else:  # ollama
+            "qwen": "https://dashscope.aliyuncs.com/compatible-mode/v1",
+            "openai": getattr(config, "OPENAI_BASE_URL", "https://api.openai.com/v1"),
+        }
+
+        self.api_keys = {
+            key: getattr(config, f"{key.upper()}_API_KEY", None)
+            for key in self.api_endpoints
+        }
+        self.api_keys["ollama"] = "OLLAMA" # required, but not used
+
+        self.models = {
+            "ollama": getattr(config, "OLLAMA_MODEL", "deepseek-coder"),
+            "deepseek": getattr(config, "DEEPSEEK_MODEL", "deepseek-chat"),
+            "qwen": getattr(config, "QWEN_MODEL", "qwen-turbo"),
+            "openai": getattr(config, "OPENAI_MODEL", "gpt-4o-mini"),
+        }
+
+        if self.api_type in self.api_endpoints:
+            self.api_endpoint = self.api_endpoints[self.api_type]
+            self.api_key = self.api_keys.get(self.api_type)
+            self.model = self.models[self.api_type]
+        elif self.api_type == "ollama":
             self.api_endpoint = config.OLLAMA_API_URL
-            self.model = config.OLLAMA_MODEL
+            self.model = getattr(config, "OLLAMA_MODEL", "deepseek-coder")
+        else:
+            raise ValueError(f"Unsupported API type: {self.api_type}")
+        print(f"API类型: {self.api_type}")
+        print(f"API地址: {self.api_endpoint}")
+        print(f"模型: {self.model}")
+
 
     def chat_completion(self, prompt, temperature=0.3):
         try:
-            if self.api_type == "deepseek":
+            return self._send_request(self.model, prompt, temperature)
-                headers = {
+        except Exception as e:
-                    "Authorization": f"Bearer {self.api_key}",
+            raise ValueError(f"API请求失败: {str(e)}")
-                    "Content-Type": "application/json"
+
-                }
+    def _send_request(self, model, prompt, temperature):
-                payload = {
+        if self.api_type in ["deepseek", "qwen", "openai"]:
-                    "model": "deepseek-chat",
+            client = OpenAI(
-                    "messages": [{"role": "user", "content": prompt}],
+                api_key=self.api_key,
-                    "temperature": temperature
+                base_url=self.api_endpoint
-                }
+            )
-                response = requests.post(self.api_endpoint, headers=headers, json=payload, timeout=60)
+            chat_completion = client.chat.completions.create(
-            else:  # ollama
+                messages=[
-                payload = {
+                    {
-                    "model": self.model,
+                        "role": "user",
-                    "messages": [{"role": "user", "content": prompt}],
+                        "content": prompt
-                    "stream": False
+                    }
-                }
+                ],
-                response = requests.post(self.api_endpoint, json=payload)
+                model=model,
+                temperature=temperature,
+                stream=False,
+            )
+            return chat_completion.choices[0].message.content
+        
+        elif self.api_type == "ollama":
+            headers = {"Content-Type": "application/json"}
+            payload = {
+                "model": self.model,
+                "messages": [{"role": "user", "content": prompt}],
+                "stream": False
+            }
+            response = requests.post(self.api_endpoint, headers=headers, json=payload, timeout=60)
             
             response.raise_for_status()
-            if self.api_type == "deepseek":
+            data = response.json()
-                return response.json()["choices"][0]["message"]["content"]
+            return self._clean_ollama_response(data["message"]["content"])
-            else:
+        else:
-                content = response.json()["message"]["content"]
+            raise ValueError(f"Unsupported API type: {self.api_type}")
-                # 使用正则表达式移除<think></think>标签及其内容
+        
-                import re
+
-                content = re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL)
+    def _clean_ollama_response(self, content):
-                return content.strip()
+        return re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL).strip()
 
-        except Exception as e:
-            raise Exception(f"API请求错误: {str(e)}")
 
 class AnalysisThread(QThread):
     analysis_complete = pyqtSignal(str, bool)
@@ -1115,6 +1157,7 @@ class CyberSecurityApp(QMainWindow):
         </html>
         """)
         self.show_status("源码审计完成", "#2ed573")
+
 if __name__ == '__main__':
     if os.name == 'nt':
         print("当前系统是 Windows")

config.py

@@ -1,12 +1,31 @@
 # API配置
-API_TYPE="ollama"  # 可选值: "deepseek" 或 "ollama"
+API_TYPE="qwen"  # Supported values: "deepseek", "qwen", "openai", "ollama"
 
 # DeepSeek API配置
 DEEPSEEK_API_KEY=""
 
-# Ollama API配置
+# OPENAI
-OLLAMA_API_URL="http://localhost:11434/api/chat"  # Ollama API地址
+OPENAI_API_KEY=""
-OLLAMA_MODEL="qwen2.5-coder:14b"  # Ollama模型名称
+
+# QWEN
+QWEN_API_KEY=""
+
+# OLLAMA
+OLLAMA_API_KEY=""
+OLLAMA_API_URL="http://localhost:11434/v1"  # Ollama API地址
+
+
+## 可选参数
+
+## OPENAI 第三方API地址
+# OPENAI_API_URL="https://api.openai.com/v1/"
+
+## MODEL
+# DEEPSEEK_MODEL=""
+# QWEN_MODEL=""
+# OPENAI_MODEL=""
+# OLLAMA_MODEL=""
+
 
 # 主题配色方案
 THEMES = {

config.py.example

@@ -1,12 +1,29 @@
 # API配置
-API_TYPE="deepseek"  # 可选值: "deepseek" 或 "ollama"
+API_TYPE="qwen"  # Supported values: "deepseek", "qwen", "openai", "ollama"
 
 # DeepSeek API配置
 DEEPSEEK_API_KEY=""
 
-# Ollama API配置
+# OPENAI
-OLLAMA_API_URL="http://localhost:11434/api/chat"  # Ollama API地址
+OPENAI_API_KEY=""
-OLLAMA_MODEL="deepseek-coder"  # Ollama模型名称
+
+# QWEN
+QWEN_API_KEY=""
+
+# OLLAMA
+OLLAMA_API_KEY=""
+OLLAMA_API_URL="http://localhost:11434/v1"  # Ollama API地址
+
+
+## 可选参数
+## OPENAI 第三方API地址
+# OPENAI_API_URL="https://api.openai.com/v1/"
+## MODEL
+# DEEPSEEK_MODEL=""
+# QWEN_MODEL=""
+# OPENAI_MODEL=""
+# OLLAMA_MODEL=""
+
 
 # 主题配色方案
 THEMES = {

requirements.txt

@@ -1,3 +1,4 @@
 PyQt5==5.15.11
 PyQt5_sip==12.15.0
 Requests==2.32.3
+openai==1.16.1