feat: 添加嵌入访问限制,白名单

apps/application/serializers/application_serializers.py

@@ -17,6 +17,8 @@ from django.core import cache
 from django.core import signing
 from django.db import transaction, models
 from django.db.models import QuerySet
+from django.http import HttpResponse
+from django.template import Template, Context
 from rest_framework import serializers
 
 from application.models import Application, ApplicationDatasetMapping
@@ -26,8 +28,10 @@ from common.constants.authentication_type import AuthenticationType
 from common.db.search import get_dynamics_model, native_search, native_page_search
 from common.db.sql_execute import select_list
 from common.exception.app_exception import AppApiException, NotFound404
+from common.util.common import getRestSeconds, set_embed_identity_cookie
 from common.util.field_message import ErrMessage
 from common.util.file_util import get_file_content
+from common.util.rsa_util import encrypt
 from dataset.models import DataSet, Document
 from dataset.serializers.common_serializers import list_paragraph
 from setting.models import AuthOperate
@@ -38,6 +42,7 @@ from smartdoc.conf import PROJECT_DIR
 from smartdoc.settings import JWT_AUTH
 
 token_cache = cache.caches['token_cache']
+chat_cache = cache.caches['chat_cache']
 
 
 class ModelDatasetAssociation(serializers.Serializer):
@@ -104,6 +109,31 @@ class ApplicationSerializer(serializers.Serializer):
         ModelDatasetAssociation(data={'user_id': user_id, 'model_id': self.data.get('model_id'),
                                       'dataset_id_list': self.data.get('dataset_id_list')}).is_valid()
 
+    class Embed(serializers.Serializer):
+        host = serializers.CharField(required=True, error_messages=ErrMessage.char("主机"))
+        protocol = serializers.CharField(required=True, error_messages=ErrMessage.char("协议"))
+        token = serializers.CharField(required=True, error_messages=ErrMessage.char("token"))
+
+        def get_embed(self, request, with_valid=True):
+            if with_valid:
+                self.is_valid(raise_exception=True)
+            index_path = os.path.join(PROJECT_DIR, 'apps', "application", 'template', 'embed.js')
+            file = open(index_path, "r", encoding='utf-8')
+            content = file.read()
+            file.close()
+            is_auth = 'true'
+            try:
+                ApplicationSerializer.Authentication(data={'access_token': self.data.get('token')}).auth()
+            except Exception as e:
+                is_auth = 'false'
+            t = Template(content)
+            s = t.render(
+                Context(
+                    {'is_auth': is_auth, 'protocol': 'http', 'host': 'localhost:8000', 'token': '0a8d892c755f1a75'}))
+            response = HttpResponse(s, status=200, headers={'Content-Type': 'text/javascript'})
+            set_embed_identity_cookie(request, response)
+            return response
+
     class AccessTokenSerializer(serializers.Serializer):
         application_id = serializers.UUIDField(required=True, error_messages=ErrMessage.boolean("应用id"))
 

apps/application/serializers/chat_message_serializers.py

@@ -24,6 +24,7 @@ from application.chat_pipeline.step.reset_problem_step.impl.base_reset_problem_s
 from application.chat_pipeline.step.search_dataset_step.impl.base_search_dataset_step import BaseSearchDatasetStep
 from application.models import ChatRecord, Chat, Application, ApplicationDatasetMapping
 from common.exception.app_exception import AppApiException
+from common.util.field_message import ErrMessage
 from common.util.rsa_util import decrypt
 from common.util.split_model import flat_map
 from dataset.models import Paragraph, Document
@@ -31,6 +32,7 @@ from setting.models import Model
 from setting.models_provider.constants.model_provider_constants import ModelProvideConstants
 
 chat_cache = caches['model_cache']
+chat_embed_identity_cache = caches['chat_cache']
 
 
 class ChatInfo:

apps/application/template/embed.js

@@ -0,0 +1,294 @@
+const guideHtml=`
+<div class="mask">
+  <div class="content"></div>
+</div>
+<div class="tips">
+  <div class="close">
+      <svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 20 20" fill="none">
+          <path d="M9.95317 8.73169L15.5511 3.13376C15.7138 2.97104 15.9776 2.97104 16.1403 3.13376L16.7296 3.72301C16.8923 3.88573 16.8923 4.14955 16.7296 4.31227L11.1317 9.9102L16.7296 15.5081C16.8923 15.6708 16.8923 15.9347 16.7296 16.0974L16.1403 16.6866C15.9776 16.8494 15.7138 16.8494 15.5511 16.6866L9.95317 11.0887L4.35524 16.6866C4.19252 16.8494 3.9287 16.8494 3.76598 16.6866L3.17673 16.0974C3.01401 15.9347 3.01401 15.6708 3.17673 15.5081L8.77465 9.9102L3.17673 4.31227C3.01401 4.14955 3.01401 3.88573 3.17673 3.72301L3.76598 3.13376C3.9287 2.97104 4.19252 2.97104 4.35524 3.13376L9.95317 8.73169Z" fill="#ffffff"></path>
+          </svg>
+  </div>
+ 
+  <div class="title"> 🌟 遇见问题，不再是障碍！</div>
+  <p>你好，我是你的智能小助手。<br/>
+      点我，开启高效解答模式，让问题变成过去式。</p>
+  <div class="button">
+      <button>我知道了</button>
+  </div>
+  <span class="arrow" ></span>
+</div>
+`
+const chatButtonHtml=
+`<div class="chat_button"><svg xmlns="http://www.w3.org/2000/svg" width="48" height="56" viewBox="0 0 48 56" fill="none">
+<g filter="url(#filter0_d_349_49711)">
+<path d="M8 24C8 12.9543 16.9543 4 28 4H48V44H28C16.9543 44 8 35.0457 8 24Z" fill="url(#paint0_linear_349_49711)"/>
+</g>
+<path d="M31.6667 15.6665H28.3333V18.1665H29.1667V19.8332H24.5833C23.6629 19.8332 22.9167 20.5794 22.9167 21.4998V30.6665C22.9167 31.587 23.6629 32.3332 24.5833 32.3332H35.4167C36.3371 32.3332 37.0833 31.587 37.0833 30.6665V21.4998C37.0833 20.5794 36.3371 19.8332 35.4167 19.8332H30.8333V18.1665H31.6667V15.6665ZM25.8333 24.8332H28.3333V27.3332H25.8333V24.8332ZM34.1667 24.8332V27.3332H31.6667V24.8332H34.1667Z" fill="white"/>
+<path d="M21.6667 23.9998H20V28.1665H21.6667V23.9998Z" fill="white"/>
+<path d="M38.3333 23.9998H40V28.1665H38.3333V23.9998Z" fill="white"/>
+<defs>
+<filter id="filter0_d_349_49711" x="0" y="0" width="56" height="56" filterUnits="userSpaceOnUse" color-interpolation-filters="sRGB">
+<feFlood flood-opacity="0" result="BackgroundImageFix"/>
+<feColorMatrix in="SourceAlpha" type="matrix" values="0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 127 0" result="hardAlpha"/>
+<feOffset dy="4"/>
+<feGaussianBlur stdDeviation="4"/>
+<feComposite in2="hardAlpha" operator="out"/>
+<feColorMatrix type="matrix" values="0 0 0 0 0.168627 0 0 0 0 0.372549 0 0 0 0 0.85098 0 0 0 0.24 0"/>
+<feBlend mode="normal" in2="BackgroundImageFix" result="effect1_dropShadow_349_49711"/>
+<feBlend mode="normal" in="SourceGraphic" in2="effect1_dropShadow_349_49711" result="shape"/>
+</filter>
+<linearGradient id="paint0_linear_349_49711" x1="48" y1="25.6667" x2="8" y2="25.6667" gradientUnits="userSpaceOnUse">
+<stop stop-color="#9258F7"/>
+<stop offset="1" stop-color="#3370FF"/>
+</linearGradient>
+</defs>
+</svg>
+</div>`
+
+ 
+
+const getChatContainerHtml=(protocol,host,token)=>{
+ return `<div id="chat_container">
+<iframe id="chat" src=${protocol}://${host}/ui/chat/${token}></iframe>
+<div class="closeviewport viewportnone"><svg t="1710214539671"  viewBox="0 0 1024 1024" version="1.1" xmlns="http://www.w3.org/2000/svg"  fill="rgb(100, 106, 115)" width="16" height="16"><path d="M85.333333 384c25.6 0 42.666667-17.066667 42.666667-42.666667V128h213.333333c25.6 0 42.666667-17.066667 42.666667-42.666667s-17.066667-42.666667-42.666667-42.666666H85.333333c-25.6 0-42.666667 17.066667-42.666666 42.666666v256c0 25.6 17.066667 42.666667 42.666666 42.666667zM938.666667 640c-25.6 0-42.666667 17.066667-42.666667 42.666667v213.333333h-213.333333c-25.6 0-42.666667 17.066667-42.666667 42.666667s17.066667 42.666667 42.666667 42.666666h256c25.6 0 42.666667-17.066667 42.666666-42.666666v-256c0-25.6-17.066667-42.666667-42.666666-42.666667zM601.6 401.066667c4.266667 8.533333 12.8 17.066667 21.333333 21.333333 4.266667 4.266667 12.8 4.266667 17.066667 4.266667h256c25.6 0 42.666667-17.066667 42.666667-42.666667s-17.066667-42.666667-42.666667-42.666667h-153.6l226.133333-226.133333c17.066667-17.066667 17.066667-42.666667 0-59.733333-8.533333-8.533333-17.066667-12.8-29.866666-12.8s-21.333333 4.266667-29.866667 12.8L682.666667 281.6V128c0-25.6-17.066667-42.666667-42.666667-42.666667s-42.666667 17.066667-42.666667 42.666667v256c0 4.266667 0 12.8 4.266667 17.066667zM115.2 968.533333L341.333333 742.4V896c0 25.6 17.066667 42.666667 42.666667 42.666667s42.666667-17.066667 42.666667-42.666667v-256c0-4.266667 0-12.8-4.266667-17.066667-4.266667-8.533333-12.8-17.066667-21.333333-21.333333-4.266667-4.266667-12.8-4.266667-17.066667-4.266667H128c-25.6 0-42.666667 17.066667-42.666667 42.666667s17.066667 42.666667 42.666667 42.666667h153.6l-226.133333 226.133333c-17.066667 17.066667-17.066667 42.666667 0 59.733333s42.666667 17.066667 59.733333 0z" p-id="10189"></path></svg></div>
+<div class="openviewport">
+            <svg t="1710150885892"   viewBox="0 0 1024 1024" version="1.1" xmlns="http://www.w3.org/2000/svg"  fill="rgb(100, 106, 115)" width="16" height="16" ><path d="M85.333333 384c25.6 0 42.666667-17.066667 42.666667-42.666667V128h213.333333c25.6 0 42.666667-17.066667 42.666667-42.666667s-17.066667-42.666667-42.666667-42.666666H85.333333c-25.6 0-42.666667 17.066667-42.666666 42.666666v256c0 25.6 17.066667 42.666667 42.666666 42.666667zM938.666667 640c-25.6 0-42.666667 17.066667-42.666667 42.666667v213.333333h-213.333333c-25.6 0-42.666667 17.066667-42.666667 42.666667s17.066667 42.666667 42.666667 42.666666h256c25.6 0 42.666667-17.066667 42.666666-42.666666v-256c0-25.6-17.066667-42.666667-42.666666-42.666667zM977.066667 68.266667c-4.266667-8.533333-12.8-17.066667-21.333334-21.333334-4.266667-4.266667-12.8-4.266667-17.066666-4.266666h-256c-25.6 0-42.666667 17.066667-42.666667 42.666666s17.066667 42.666667 42.666667 42.666667h153.6l-226.133334 226.133333c-17.066667 17.066667-17.066667 42.666667 0 59.733334 8.533333 8.533333 17.066667 12.8 29.866667 12.8s21.333333-4.266667 29.866667-12.8L896 187.733333V341.333333c0 25.6 17.066667 42.666667 42.666667 42.666667s42.666667-17.066667 42.666666-42.666667V85.333333c0-4.266667 0-12.8-4.266666-17.066666zM354.133333 610.133333L128 836.266667V682.666667c0-25.6-17.066667-42.666667-42.666667-42.666667s-42.666667 17.066667-42.666666 42.666667v256c0 4.266667 0 12.8 4.266666 17.066666 4.266667 8.533333 12.8 17.066667 21.333334 21.333334 4.266667 4.266667 12.8 4.266667 17.066666 4.266666h256c25.6 0 42.666667-17.066667 42.666667-42.666666s-17.066667-42.666667-42.666667-42.666667H187.733333l226.133334-226.133333c17.066667-17.066667 17.066667-42.666667 0-59.733334s-42.666667-17.066667-59.733334 0z" p-id="8645"></path></svg>
+</div>
+           <div class="close"><svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 20 20" fill="none">
+            <path d="M9.95317 8.73169L15.5511 3.13376C15.7138 2.97104 15.9776 2.97104 16.1403 3.13376L16.7296 3.72301C16.8923 3.88573 16.8923 4.14955 16.7296 4.31227L11.1317 9.9102L16.7296 15.5081C16.8923 15.6708 16.8923 15.9347 16.7296 16.0974L16.1403 16.6866C15.9776 16.8494 15.7138 16.8494 15.5511 16.6866L9.95317 11.0887L4.35524 16.6866C4.19252 16.8494 3.9287 16.8494 3.76598 16.6866L3.17673 16.0974C3.01401 15.9347 3.01401 15.6708 3.17673 15.5081L8.77465 9.9102L3.17673 4.31227C3.01401 4.14955 3.01401 3.88573 3.17673 3.72301L3.76598 3.13376C3.9287 2.97104 4.19252 2.97104 4.35524 3.13376L9.95317 8.73169Z" fill="#646A73"/>
+  </svg>
+ </div>`
+}
+/**
+ * 初始化引导
+ * @param {*} root  
+ */
+const initGuide=(root)=>{
+   root.insertAdjacentHTML("beforeend",guideHtml)
+   const button=root.querySelector(".button")
+   const close_icon=root.querySelector('.close')
+   const close_func=()=>{
+     root.removeChild(root.querySelector('.tips'))
+     root.removeChild(root.querySelector('.mask'))
+     localStorage.setItem('maxkbMaskTip',true)
+   }
+   button.onclick=close_func
+   close_icon.onclick=close_func
+}
+const initChat=(root)=>{
+  // 添加对话icon
+  root.insertAdjacentHTML("beforeend",chatButtonHtml)
+  // 添加对话框
+  root.insertAdjacentHTML('beforeend',getChatContainerHtml('{{protocol}}','{{host}}','{{token}}'))
+  // 按钮元素
+  const chat_button=root.querySelector('.chat_button')
+  //  对话框元素
+  const chat_container=root.querySelector('#chat_container')
+
+  const viewport=root.querySelector('.openviewport')
+  const closeviewport=root.querySelector('.closeviewport')
+  const close_func=()=>{
+    chat_container.style['display']=chat_container.style['display']=='block'?'none':'block'
+  }
+  close_icon=chat_container.querySelector('.close')
+  chat_button.onclick = close_func
+  close_icon.onclick=close_func
+  const viewport_func=()=>{
+    if(chat_container.classList.contains('enlarge')){
+      chat_container.classList.remove("enlarge");
+      viewport.classList.remove('viewportnone')
+      closeviewport.classList.add('viewportnone')
+    }else{
+      chat_container.classList.add("enlarge");
+      viewport.classList.add('viewportnone')
+      closeviewport.classList.remove('viewportnone')
+    }
+  }
+  viewport.onclick=viewport_func
+  closeviewport.onclick=viewport_func
+}
+/**
+ * 第一次进来的引导提示
+ */
+function initMaxkb(){
+  initMaxkbStyle()
+  const root=document.createElement('div')
+  root.id="maxkb"
+  document.body.appendChild(root)
+  const maxkbMaskTip=localStorage.getItem('maxkbMaskTip')
+  if(maxkbMaskTip==null){
+    initGuide(root)
+  }
+  initChat(root)
+}
+
+ 
+// 初始化全局样式
+function initMaxkbStyle(){
+  style=document.createElement('style')
+  style.type='text/css'
+  style.innerText=  `/* 放大 */
+  #maxkb .enlarge {
+      width: 50%!important;
+      height: 100%!important;
+      bottom: 0!important;
+      right: 0 !important;
+  }
+  @media only screen and (max-width: 768px){
+  #maxkb .enlarge {
+      width: 100%!important;
+      height: 100%!important;
+      right: 0 !important;
+      bottom: 0!important;
+  }
+  }
+  
+  /* 引导 */
+  
+  #maxkb .mask {
+      position: fixed;
+      z-index: 999;
+      background-color: transparent;
+      height: 100%;
+      width: 100%;
+      top: 0;
+      left: 0;
+  }
+  #maxkb .mask .content {
+      width: 45px;
+      height: 50px;
+      box-shadow: 1px 1px 1px 2000px rgba(0,0,0,.6);
+      border-radius: 50% 0 0 50%;
+      position: absolute;
+      right: 0;
+      bottom: 42px;
+      z-index: 1000;
+  }
+  #maxkb .tips {
+      position: absolute;
+      bottom: 30px;
+      right: 60px;
+      padding: 22px 24px 24px;
+      border-radius: 6px;
+      color: #ffffff;
+      font-size: 14px;
+      background: #3370FF;
+      z-index: 1000;
+  }
+  #maxkb .tips .arrow {
+      position: absolute;
+      background: #3370FF;
+      width: 10px;
+      height: 10px;
+      pointer-events: none;
+      transform: rotate(45deg);
+      box-sizing: border-box;
+      /* left  */
+      right: -5px;
+      bottom: 33px;
+      border-left-color: transparent;
+      border-bottom-color: transparent
+  }
+  #maxkb .tips .title {
+      font-size: 20px;
+      font-weight: 500;
+      margin-bottom: 8px;
+  }
+  #maxkb .tips .button {
+      text-align: right;
+      margin-top: 24px;
+  }
+  #maxkb .tips .button button {
+      border-radius: 4px;
+      background: #FFF;
+      padding: 3px 12px;
+      color: #3370FF;
+      cursor: pointer;
+      outline: none;
+      border: none;
+  }
+  #maxkb .tips .button button::after{
+      border: none;
+    }
+  #maxkb .tips .close {
+      position: absolute;
+      right: 20px;
+      top: 20px;
+      cursor: pointer;
+  
+  }
+  #chat_container {
+        width: 420px;
+        height: 600px;
+        display:none;
+      }
+  @media only screen and (max-width: 768px) {
+        #chat_container {
+          width: 100%;
+          height: 70%;
+          right: 0 !important;
+        }
+      }
+      
+      #maxkb .chat_button{
+        position: fixed;
+        bottom: 30px;
+        right: 0;
+        cursor: pointer;
+    }
+    #maxkb #chat_container{
+        z-index:10000;position: relative;
+              border-radius: 8px;
+              border: 1px solid var(--N300, #DEE0E3);
+              background: linear-gradient(188deg, rgba(235, 241, 255, 0.20) 39.6%, rgba(231, 249, 255, 0.20) 94.3%), #EFF0F1;
+              box-shadow: 0px 4px 8px 0px rgba(31, 35, 41, 0.10);
+              position: fixed;bottom: 20px;right: 45px;overflow: hidden;
+    }
+    #maxkb #chat_container .close{
+        position: absolute;
+            top: 15px;
+            right: 10px;
+            cursor: pointer;
+    }
+    #maxkb #chat_container .openviewport{
+           position: absolute;
+            top: 15px;
+            right: 50px;
+            cursor: pointer;
+    }
+    #maxkb #chat_container .closeviewport{
+      position: absolute;
+      top: 15px;
+      right: 50px;
+      cursor: pointer;
+    }
+    #maxkb #chat_container .viewportnone{
+      display:none;
+    }
+    #maxkb #chat_container #chat{
+     height:100%;
+     width:100%;
+     border: none;
+}
+    #maxkb #chat_container {
+                animation: appear .4s ease-in-out;
+              }
+              @keyframes appear {
+                from {
+                  height: 0;;
+                }
+        
+                to {
+                  height: 600px;
+                }
+              }`
+  document.head.appendChild(style)
+}
+
+function embedChatbot() {
+  if ({{is_auth}}) {
+    // 初始化maxkb智能小助手
+    initMaxkb()
+  } else console.error('invalid parameter')
+}
+window.onload = embedChatbot

apps/application/urls.py

@@ -6,6 +6,7 @@ app_name = "application"
 urlpatterns = [
     path('application', views.Application.as_view(), name="application"),
     path('application/profile', views.Application.Profile.as_view()),
+    path('application/embed', views.Application.Embed.as_view()),
     path('application/authentication', views.Application.Authentication.as_view()),
     path('application/<str:application_id>/model', views.Application.Model.as_view()),
     path('application/<str:application_id>/hit_test', views.Application.HitTest.as_view()),

apps/application/views/application_views.py

@@ -6,6 +6,7 @@
     @date：2023/10/27 14:56
     @desc:
 """
+
 from django.http import HttpResponse
 from drf_yasg.utils import swagger_auto_schema
 from rest_framework.decorators import action
@@ -20,13 +21,24 @@ from common.constants.permission_constants import CompareConstants, PermissionCo
 from common.exception.app_exception import AppAuthenticationFailed
 from common.response import result
 from common.swagger_api.common_api import CommonApi
-from common.util.common import query_params_to_single_dict
+from common.util.common import query_params_to_single_dict, set_embed_identity_cookie
 from dataset.serializers.dataset_serializers import DataSetSerializers
 
 
 class Application(APIView):
     authentication_classes = [TokenAuth]
 
+    class Embed(APIView):
+        @action(methods=["GET"], detail=False)
+        @swagger_auto_schema(operation_summary="获取嵌入js",
+                             operation_id="获取嵌入js",
+                             tags=["应用"],
+                             manual_parameters=ApplicationApi.ApiKey.get_request_params_api())
+        def get(self, request: Request):
+            return ApplicationSerializer.Embed(
+                data={'protocol': request.query_params.get('protocol'), 'token': request.query_params.get('token'),
+                      'host': request.query_params.get('host'), }).get_embed(request)
+
     class Model(APIView):
         authentication_classes = [TokenAuth]
 
@@ -185,7 +197,7 @@ class Application(APIView):
                          "Access-Control-Allow-Methods": "POST",
                          "Access-Control-Allow-Headers": "Origin,Content-Type,Cookie,Accept,Token"}
             )
-
+            set_embed_identity_cookie(request, response)
             return response
 
     @action(methods=['POST'], detail=False)

apps/application/views/chat_views.py

@@ -18,7 +18,7 @@ from common.auth import TokenAuth, has_permissions
 from common.constants.permission_constants import Permission, Group, Operate, \
     RoleConstants, ViewPermission, CompareConstants
 from common.response import result
-from common.util.common import query_params_to_single_dict
+from common.util.common import query_params_to_single_dict, set_embed_identity_cookie
 
 
 class ChatView(APIView):
@@ -71,9 +71,13 @@ class ChatView(APIView):
                                                            dynamic_tag=keywords.get('application_id'))])
         )
         def post(self, request: Request, chat_id: str):
-            return ChatMessageSerializer(data={'chat_id': chat_id}).chat(request.data.get('message'), request.data.get(
+            response = ChatMessageSerializer(data={'chat_id': chat_id}).chat(request.data.get('message'),
-                're_chat') if 're_chat' in request.data else False, request.data.get(
+                                                                             request.data.get(
-                'stream') if 'stream' in request.data else True)
+                                                                                 're_chat') if 're_chat' in request.data else False,
+                                                                             request.data.get(
+                                                                                 'stream') if 'stream' in request.data else True)
+            set_embed_identity_cookie(request, response)
+            return response
 
     @action(methods=['GET'], detail=False)
     @swagger_auto_schema(operation_summary="获取对话列表",

apps/common/auth/authenticate.py

@@ -6,21 +6,28 @@
     @date：2023/9/4 11:16
     @desc:  认证类
 """
+import datetime
+import traceback
+from urllib.parse import urlparse
 
 from django.core import cache
 from django.core import signing
 from django.db.models import QuerySet
+from ipware import get_client_ip
 from rest_framework.authentication import TokenAuthentication
 
 from application.models.api_key_model import ApplicationAccessToken, ApplicationApiKey
 from common.constants.authentication_type import AuthenticationType
 from common.constants.permission_constants import Auth, get_permission_list_by_role, RoleConstants, Permission, Group, \
     Operate
-from common.exception.app_exception import AppAuthenticationFailed
+from common.exception.app_exception import AppAuthenticationFailed, AppEmbedIdentityFailed, AppChatNumOutOfBoundsFailed
+from common.util.common import getRestSeconds
+from common.util.rsa_util import decrypt
 from smartdoc.settings import JWT_AUTH
 from users.models.user import User, get_user_dynamics_permission
 
 token_cache = cache.caches['token_cache']
+chat_cache = cache.caches['chat_cache']
 
 
 class AnonymousAuthentication(TokenAuthentication):
@@ -80,6 +87,35 @@ class TokenAuth(TokenAuthentication):
                     raise AppAuthenticationFailed(1002, "身份验证信息不正确")
                 if not application_access_token.access_token == auth_details.get('access_token'):
                     raise AppAuthenticationFailed(1002, "身份验证信息不正确")
+                if application_access_token.white_active:
+                    referer = request.META.get('HTTP_REFERER')
+                    if referer is not None:
+                        client_ip = urlparse(referer).hostname
+                    else:
+                        client_ip = get_client_ip(request)
+                    if not application_access_token.white_list.__contains__(client_ip):
+                        raise AppAuthenticationFailed(1002, "身份验证信息不正确")
+                if 'embed_identity' in request.COOKIES and request.path.__contains__('/api/application/chat_message/'):
+                    embed_identity = request.COOKIES['embed_identity']
+                    try:
+                        # 如果无法解密 说明embed_identity并非系统颁发
+                        value = decrypt(embed_identity)
+                    except Exception as e:
+                        raise AppEmbedIdentityFailed(1004, '嵌入cookie不正确')
+                    embed_identity_number = chat_cache.get(value)
+                    if embed_identity_number is not None:
+                        if application_access_token.access_num <= embed_identity_number:
+                            raise AppChatNumOutOfBoundsFailed(1003, '访问次数超过今日访问量')
+                    # 对话次数+1
+                    try:
+                        if not chat_cache.incr(value):
+                            # 如果修改失败则设置为1
+                            chat_cache.set(value, 1,
+                                           timeout=getRestSeconds())
+                    except Exception as e:
+                        # 如果修改失败则设置为1 证明 key不存在
+                        chat_cache.add(value, 1,
+                                       timeout=getRestSeconds())
                 return application_access_token.application.user, Auth(
                     role_list=[RoleConstants.APPLICATION_ACCESS_TOKEN],
                     permission_list=[
@@ -94,4 +130,7 @@ class TokenAuth(TokenAuthentication):
                 raise AppAuthenticationFailed(1002, "身份验证信息不正确！非法用户")
 
         except Exception as e:
+            traceback.format_exc()
+            if isinstance(e, AppEmbedIdentityFailed) or isinstance(e, AppChatNumOutOfBoundsFailed):
+                raise e
             raise AppAuthenticationFailed(1002, "身份验证信息不正确！非法用户")

apps/common/exception/app_exception.py

@@ -51,3 +51,25 @@ class AppUnauthorizedFailed(AppApiException):
     def __init__(self, code, message):
         self.code = code
         self.message = message
+
+
+class AppEmbedIdentityFailed(AppApiException):
+    """
+    嵌入cookie异常
+    """
+    status_code = 460
+
+    def __init__(self, code, message):
+        self.code = code
+        self.message = message
+
+
+class AppChatNumOutOfBoundsFailed(AppApiException):
+    """
+      访问次数超过今日访问量
+    """
+    status_code = 461
+
+    def __init__(self, code, message):
+        self.code = code
+        self.message = message

apps/common/util/common.py

@@ -6,9 +6,36 @@
     @date：2023/10/16 16:42
     @desc:
 """
+import datetime
 import importlib
+import uuid
 from functools import reduce
 from typing import Dict, List
+from django.core import cache
+
+from .rsa_util import encrypt
+
+chat_cache = cache.caches['chat_cache']
+
+
+def set_embed_identity_cookie(request, response):
+    if 'embed_identity' in request.COOKIES:
+        embed_identity = request.COOKIES['embed_identity']
+    else:
+        value = str(uuid.uuid1())
+        embed_identity = encrypt(value)
+        chat_cache.set(value, 0, timeout=getRestSeconds())
+    response.set_cookie("embed_identity", embed_identity, max_age=3600 * 24 * 100, samesite='None',
+                        secure=True)
+    return response
+
+
+def getRestSeconds():
+    now = datetime.datetime.now()
+    today_begin = datetime.datetime(now.year, now.month, now.day, 0, 0, 0)
+    tomorrow_begin = today_begin + datetime.timedelta(days=1)
+    rest_seconds = (tomorrow_begin - now).seconds
+    return rest_seconds
 
 
 def sub_array(array: List, item_num=10):

apps/smartdoc/settings/base.py

@@ -103,8 +103,7 @@ CACHES = {
         'LOCATION': os.path.join(PROJECT_DIR, 'data', 'cache', "token_cache")  # 文件夹路径
     },
     "chat_cache": {
-        'BACKEND': 'common.cache.file_cache.FileCache',
+        'BACKEND': 'django.core.cache.backends.locmem.LocMemCache',
-        'LOCATION': os.path.join(PROJECT_DIR, 'data', 'cache', "chat_cache")  # 文件夹路径
     }
 }
 

pyproject.toml

@@ -28,6 +28,7 @@ pycryptodome = "^3.19.0"
 beautifulsoup4 = "^4.12.2"
 html2text = "^2024.2.26"
 langchain-openai = "^0.0.8"
+django-ipware = "^6.0.4"
 
 [build-system]
 requires = ["poetry-core"]

ui/src/components/ai-chat/index.vue

@@ -399,10 +399,10 @@ const getWrite = (chat: any, reader: any, stream: boolean) => {
   }
   return stream ? write_stream : write_json
 }
-const errorWrite = (chat: any) => {
+const errorWrite = (chat: any, message?: string) => {
   ChatManagement.addChatRecord(chat, 50, loading)
   ChatManagement.write(chat.id)
-  ChatManagement.append(chat.id, '抱歉，当前正在维护，无法提供服务，请稍后再试！')
+  ChatManagement.append(chat.id, message || '抱歉，当前正在维护，无法提供服务，请稍后再试！')
   ChatManagement.close(chat.id)
 }
 function chatMessage(chat?: any, problem?: string) {
@@ -444,6 +444,10 @@ function chatMessage(chat?: any, problem?: string) {
             .catch((err) => {
               errorWrite(chat)
             })
+        } else if (response.status === 460) {
+          return Promise.reject('无法识别用户身份')
+        } else if (response.status === 461) {
+          return Promise.reject('抱歉，您的提问已达到最大限制，请明天再来吧！')
         } else {
           nextTick(() => {
             // 将滚动条滚动到最下面
@@ -468,7 +472,7 @@ function chatMessage(chat?: any, problem?: string) {
         ChatManagement.close(chat.id)
       })
       .catch((e: any) => {
-        MsgError(e)
+        errorWrite(chat, e + '')
       })
   }
 }

ui/src/request/index.ts

@@ -58,7 +58,10 @@ instance.interceptors.response.use(
       }
     }
     if (err.response?.status === 401) {
-      if (!err.response.config.url.includes('chat/open')) {
+      if (
+        !err.response.config.url.includes('chat/open') &&
+        !err.response.config.url.includes('application/profile')
+      ) {
         router.push({ name: 'login' })
       }
     }

ui/src/views/applicaiton-overview/component/EmbedDialog.vue

@@ -29,7 +29,7 @@
                 <AppIcon iconName="app-copy"></AppIcon>
               </el-button>
             </div>
-            <div class="mt-8">
+            <div class="mt-8 white-space">
               {{ source2 }}
             </div>
           </div>
@@ -68,13 +68,13 @@ frameborder="0"
 allow="microphone">
 </iframe>
 `
-  source2.value = `<script> window.maxkbChatConfig = { 
+
-  token: "${val}",
+  source2.value = `
-  host: "${window.location.host}",
+<script src="${
-  protocol:"${window.location.protocol}"
+    window.location.origin
- }
+  }/api/application/embed?protocol=${window.location.protocol.replace(':', '')}&host=${
- <\/script>
+    window.location.host
-<script src="${window.location.origin}/ui/embeb.js">
+  }&token=${val}">
 <\/script>
 `
   dialogVisible.value = true
@@ -95,6 +95,9 @@ defineExpose({ open })
     font-size: 13px;
     white-space: pre;
     height: 180px;
+    .white-space {
+      white-space: pre-wrap;
+    }
   }
 }
 </style>