From d8db76cc7ba83cf4f095fa12831467f60a2b084a Mon Sep 17 00:00:00 2001 From: feng <1304903146@qq.com> Date: Sat, 8 Feb 2025 14:53:32 +0800 Subject: [PATCH] perf: DeepSeek --- apps/jumpserver/conf.py | 4 +++- apps/jumpserver/settings/custom.py | 2 ++ apps/settings/api/chat.py | 7 ++++++- apps/settings/const.py | 18 ++++++++++++++++++ apps/settings/models.py | 9 ++++++++- apps/settings/serializers/feature.py | 21 ++++++++++++++------- apps/settings/serializers/public.py | 1 + apps/terminal/models/component/terminal.py | 3 ++- apps/users/api/profile.py | 1 - 9 files changed, 54 insertions(+), 12 deletions(-) diff --git a/apps/jumpserver/conf.py b/apps/jumpserver/conf.py index 01af7ec96..452b809d9 100644 --- a/apps/jumpserver/conf.py +++ b/apps/jumpserver/conf.py @@ -681,10 +681,12 @@ class Config(dict): # Chat AI 'CHAT_AI_ENABLED': False, + 'CHAT_AI_TYPE': 'gpt', 'GPT_API_KEY': '', 'GPT_BASE_URL': '', 'GPT_PROXY': '', - 'GPT_MODEL': 'gpt-3.5-turbo', + 'GPT_MODEL': 'gpt-4o-mini', + 'DEEPSEEK_MODEL': 'deepseek-chat', 'VIRTUAL_APP_ENABLED': False, 'FILE_UPLOAD_SIZE_LIMIT_MB': 200, diff --git a/apps/jumpserver/settings/custom.py b/apps/jumpserver/settings/custom.py index 952c9bdd7..f02eb1cb6 100644 --- a/apps/jumpserver/settings/custom.py +++ b/apps/jumpserver/settings/custom.py @@ -226,6 +226,8 @@ GPT_API_KEY = CONFIG.GPT_API_KEY GPT_BASE_URL = CONFIG.GPT_BASE_URL GPT_PROXY = CONFIG.GPT_PROXY GPT_MODEL = CONFIG.GPT_MODEL +CHAT_AI_TYPE = CONFIG.CHAT_AI_TYPE +DEEPSEEK_MODEL = CONFIG.DEEPSEEK_MODEL VIRTUAL_APP_ENABLED = CONFIG.VIRTUAL_APP_ENABLED diff --git a/apps/settings/api/chat.py b/apps/settings/api/chat.py index 14e21d68f..0d5135999 100644 --- a/apps/settings/api/chat.py +++ b/apps/settings/api/chat.py @@ -9,6 +9,7 @@ from rest_framework.response import Response from common.api import JMSModelViewSet from common.permissions import IsValidUser, OnlySuperUser from .. import serializers +from ..const import ChatAITypeChoices from ..models import ChatPrompt from ..prompt import DefaultChatPrompt @@ -41,7 +42,11 @@ class ChatAITestingAPI(GenericAPIView): ) proxy = config['GPT_PROXY'] - model = config['GPT_MODEL'] + tp = config['CHAT_AI_TYPE'] + if tp == ChatAITypeChoices.gpt: + model = config['GPT_MODEL'] + else: + model = config['DEEPSEEK_MODEL'] kwargs = { 'base_url': config['GPT_BASE_URL'] or None, diff --git a/apps/settings/const.py b/apps/settings/const.py index 4e4028d4f..d64b6ee2f 100644 --- a/apps/settings/const.py +++ b/apps/settings/const.py @@ -5,3 +5,21 @@ class ImportStatus(TextChoices): ok = 'ok', 'Ok' pending = 'pending', 'Pending' error = 'error', 'Error' + + +class ChatAITypeChoices(TextChoices): + gpt = 'gpt', 'GPT' + deep_seek = 'deep-seek', 'DeepSeek' + + +class GPTModelChoices(TextChoices): + gpt_4o_mini = 'gpt-4o-mini', 'gpt-4o-mini' + gpt_4o = 'gpt-4o', 'gpt-4o' + o3_mini = 'o3-mini', 'o3-mini' + o1_mini = 'o1-mini', 'o1-mini' + o1 = 'o1', 'o1' + + +class DeepSeekModelChoices(TextChoices): + deepseek_chat = 'deepseek-chat', 'DeepSeek-V3' + deepseek_reasoner = 'deepseek-reasoner', 'DeepSeek-R1' diff --git a/apps/settings/models.py b/apps/settings/models.py index 0ef3b4a5f..9c0a4053c 100644 --- a/apps/settings/models.py +++ b/apps/settings/models.py @@ -10,8 +10,9 @@ from django.utils.translation import gettext_lazy as _ from rest_framework.utils.encoders import JSONEncoder from common.db.models import JMSBaseModel -from common.utils import get_logger from common.db.utils import Encryptor +from common.utils import get_logger +from .const import ChatAITypeChoices from .signals import setting_changed logger = get_logger(__name__) @@ -191,3 +192,9 @@ class ChatPrompt(JMSBaseModel): def __str__(self): return self.name + + +def get_chatai_model(): + if settings.CHAT_AI_TYPE == ChatAITypeChoices.gpt: + return settings.GPT_MODEL + return settings.DEEPSEEK_MODEL diff --git a/apps/settings/serializers/feature.py b/apps/settings/serializers/feature.py index 89df9ea85..9fb994ab1 100644 --- a/apps/settings/serializers/feature.py +++ b/apps/settings/serializers/feature.py @@ -4,7 +4,6 @@ from django.utils import timezone from django.utils.translation import gettext_lazy as _ from rest_framework import serializers -from assets.const import Protocol from common.serializers.fields import EncryptedField from common.utils import date_expired_default @@ -14,6 +13,8 @@ __all__ = [ 'ChatAISettingSerializer', 'VirtualAppSerializer', 'AmazonSMSerializer', ] +from settings.const import ChatAITypeChoices, GPTModelChoices, DeepSeekModelChoices + class AnnouncementSerializer(serializers.Serializer): ID = serializers.CharField(required=False, allow_blank=True, allow_null=True) @@ -119,16 +120,17 @@ class AmazonSMSerializer(serializers.Serializer): class ChatAISettingSerializer(serializers.Serializer): PREFIX_TITLE = _('Chat AI') - API_MODEL = Protocol.gpt_protocols()[Protocol.chatgpt]['setting']['api_mode'] - GPT_MODEL_CHOICES = API_MODEL['choices'] - GPT_MODEL_DEFAULT = API_MODEL['default'] CHAT_AI_ENABLED = serializers.BooleanField( required=False, label=_('Chat AI') ) + CHAT_AI_TYPE = serializers.ChoiceField( + default=ChatAITypeChoices.gpt, choices=ChatAITypeChoices.choices, + label=_("Chat AI type"), required=False, + ) GPT_BASE_URL = serializers.CharField( - allow_blank=True, required=False, label=_('GPT Base URL'), - help_text=_('The base URL of the GPT service. For example: https://api.openai.com/v1') + allow_blank=True, required=False, label=_('Base URL'), + help_text=_('The base URL of the Chat service.') ) GPT_API_KEY = EncryptedField( allow_blank=True, required=False, label=_('API Key'), @@ -138,7 +140,12 @@ class ChatAISettingSerializer(serializers.Serializer): help_text=_('The proxy server address of the GPT service. For example: http://ip:port') ) GPT_MODEL = serializers.ChoiceField( - default=GPT_MODEL_DEFAULT, choices=GPT_MODEL_CHOICES, label=_("GPT Model"), required=False, + default=GPTModelChoices.gpt_4o_mini, choices=GPTModelChoices.choices, + label=_("GPT Model"), required=False, + ) + DEEPSEEK_MODEL = serializers.ChoiceField( + default=DeepSeekModelChoices.deepseek_chat, choices=DeepSeekModelChoices.choices, + label=_("DeepSeek Model"), required=False, ) diff --git a/apps/settings/serializers/public.py b/apps/settings/serializers/public.py index 3c7b11a6b..d1f428b86 100644 --- a/apps/settings/serializers/public.py +++ b/apps/settings/serializers/public.py @@ -63,6 +63,7 @@ class PrivateSettingSerializer(PublicSettingSerializer): VAULT_ENABLED = serializers.BooleanField() VIRTUAL_APP_ENABLED = serializers.BooleanField() CHAT_AI_ENABLED = serializers.BooleanField() + CHAT_AI_TYPE = serializers.CharField() GPT_MODEL = serializers.CharField() FILE_UPLOAD_SIZE_LIMIT_MB = serializers.IntegerField() FTP_FILE_MAX_STORE = serializers.IntegerField() diff --git a/apps/terminal/models/component/terminal.py b/apps/terminal/models/component/terminal.py index 15ee0524d..01ad3f8d8 100644 --- a/apps/terminal/models/component/terminal.py +++ b/apps/terminal/models/component/terminal.py @@ -9,6 +9,7 @@ from common.const.signals import OP_LOG_SKIP_SIGNAL from common.db.models import JMSBaseModel from common.utils import get_logger, lazyproperty from orgs.utils import tmp_to_root_org +from settings.models import get_chatai_model from terminal.const import TerminalType as TypeChoices from users.models import User from .status import Status @@ -124,7 +125,7 @@ class Terminal(StorageMixin, TerminalStatusMixin, JMSBaseModel): 'GPT_BASE_URL': settings.GPT_BASE_URL, 'GPT_API_KEY': settings.GPT_API_KEY, 'GPT_PROXY': settings.GPT_PROXY, - 'GPT_MODEL': settings.GPT_MODEL, + 'GPT_MODEL': get_chatai_model(), } @staticmethod diff --git a/apps/users/api/profile.py b/apps/users/api/profile.py index 7e37da52c..0dee24b8e 100644 --- a/apps/users/api/profile.py +++ b/apps/users/api/profile.py @@ -10,7 +10,6 @@ from common.utils import get_object_or_none from orgs.utils import tmp_to_root_org from users.notifications import ( ResetPasswordMsg, ResetPasswordSuccessMsg, ResetSSHKeyMsg, - ResetPublicKeySuccessMsg, ) from .mixins import UserQuerysetMixin from .. import serializers