perf: Assistant dynamic certificate

Author: fit2cloud-chenyw

backend/apps/chat/api/chat.py

@@ -96,7 +96,7 @@ async def start_chat(session: SessionDep, current_user: CurrentUser):
 @router.post("/recommend_questions/{chat_record_id}")
 async def recommend_questions(session: SessionDep, current_user: CurrentUser, chat_record_id: int, current_assistant: CurrentAssistant):
     try:
-        record = session.query(ChatRecord).get(chat_record_id)
+        record = session.get(ChatRecord, chat_record_id)
         if not record:
             raise HTTPException(
                 status_code=400,

backend/apps/chat/models/chat_model.py

@@ -151,8 +151,6 @@ def filter_user_question(self):
 class ChatQuestion(AiModelQuestion):
     question: str = ''
     chat_id: int = 0
-    assistant_certificate: Optional[str] = None
-
 
 class ChatMcp(ChatQuestion):
     token: str = ''

backend/apps/chat/task/llm.py

@@ -53,7 +53,6 @@ class LLMService:
     session: SessionDep
     current_user: CurrentUser
     current_assistant: Optional[CurrentAssistant] = None
-    assistant_certificate: str
     out_ds_instance: Optional[AssistantOutDs] = None
     change_title: bool = False
 
@@ -63,8 +62,6 @@ def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question
         self.session = session
         self.current_user = current_user
         self.current_assistant = current_assistant
-        if chat_question.assistant_certificate:
-            self.assistant_certificate = chat_question.assistant_certificate
         # chat = self.session.query(Chat).filter(Chat.id == chat_question.chat_id).first()
         chat_id = chat_question.chat_id
         chat: Chat = self.session.get(Chat, chat_id)

backend/apps/system/crud/assistant.py

@@ -10,7 +10,7 @@
 
 from apps.system.models.system_model import AssistantModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
-from apps.system.schemas.system_schema import AssistantOutDsSchema, UserInfoDTO
+from apps.system.schemas.system_schema import AssistantHeader, AssistantOutDsSchema, UserInfoDTO
 from common.core.sqlbot_cache import cache
 from common.core.db import engine
 from starlette.middleware.cors import CORSMiddleware
@@ -24,9 +24,8 @@ async def get_assistant_info(*, session: Session, assistant_id: int) -> Assistan
 def get_assistant_user(*, id: int):
     return UserInfoDTO(id=id, account="sqlbot-inner-assistant", oid=1, name="sqlbot-inner-assistant", email="[email protected]")
 
-# def get_assistant_ds(*, session: Session, assistant: AssistantModel):
 def get_assistant_ds(llm_service) -> list[dict]:
-    assistant: AssistantModel = llm_service.current_assistant
+    assistant: AssistantHeader = llm_service.current_assistant
     session: Session = llm_service.session
     type = assistant.type
     if type == 0:
@@ -51,7 +50,7 @@ def get_assistant_ds(llm_service) -> list[dict]:
 
         # filter private ds if offline
         return result_list
-    out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(assistant, llm_service.assistant_certificate)
+    out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(assistant)
     llm_service.out_ds_instance = out_ds_instance
     dslist = out_ds_instance.get_simple_ds_list()
     # format?
@@ -84,20 +83,20 @@ def init_dynamic_cors(app: FastAPI):
 
 
 class AssistantOutDs:
-    assistant: AssistantModel
+    assistant: AssistantHeader
     ds_list: Optional[list[AssistantOutDsSchema]] = None
     certificate: Optional[str] = None
-    def __init__(self, assistant: AssistantModel, certificate: Optional[str] = None):
+    def __init__(self, assistant: AssistantHeader):
         self.assistant = assistant
         self.ds_list = None
-        self.certificate = certificate
-        self.get_ds_from_api(certificate)
+        self.certificate = assistant.certificate
+        self.get_ds_from_api()
 
     #@cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_DS, keyExpression="current_user.id")    
-    def get_ds_from_api(self, certificate: Optional[str] = None):
+    def get_ds_from_api(self):
         config: dict[any] = json.loads(self.assistant.configuration)
         endpoint: str = config['endpoint']
-        certificateList: list[any] = json.loads(certificate)
+        certificateList: list[any] = json.loads(self.certificate)
         header = {}
         cookies = {}
         for item in certificateList:
@@ -155,8 +154,8 @@ def get_ds(self, ds_id: int):
 
 class AssistantOutDsFactory:
     @staticmethod
-    def get_instance(assistant: AssistantModel, certificate: Optional[str] = None) -> AssistantOutDs:
-        return AssistantOutDs(assistant, certificate)
+    def get_instance(assistant: AssistantHeader) -> AssistantOutDs:
+        return AssistantOutDs(assistant)
 
 def get_ds_engine(ds: AssistantOutDsSchema) -> Engine:
     timeout: int = 30

backend/apps/system/middleware/auth.py

@@ -9,7 +9,7 @@
 from common.core.db import engine 
 from apps.system.crud.assistant import get_assistant_info, get_assistant_user
 from apps.system.crud.user import get_user_info
-from apps.system.schemas.system_schema import UserInfoDTO
+from apps.system.schemas.system_schema import AssistantHeader, UserInfoDTO
 from common.core import security
 from common.core.config import settings
 from common.core.schemas import TokenPayload
@@ -92,6 +92,7 @@ async def validateAssistant(self, assistantToken: Optional[str]) -> tuple[any]:
                 session_user = get_assistant_user(id = token_data.id)
                 assistant_info = await get_assistant_info(session=session, assistant_id=payload['assistant_id'])
                 assistant_info = AssistantModel.model_validate(assistant_info)
+                assistant_info = AssistantHeader.model_validate(assistant_info.model_dump(exclude_unset=True))
                 return True, session_user, assistant_info
         except Exception as e:
             return False, e

backend/apps/system/schemas/system_schema.py

@@ -71,6 +71,12 @@ class AssistantBase(BaseModel):
 class AssistantDTO(AssistantBase, BaseCreatorDTO):
     pass
 
+class AssistantHeader(AssistantDTO):
+    unique: Optional[str] = None
+    certificate: Optional[str] = None
+    online: bool = False
+    
+
 class AssistantValidator(BaseModel):
     valid: bool = False
     id_match: bool = False

backend/common/core/deps.py

@@ -1,9 +1,10 @@
+import base64
 from typing import Annotated
+from urllib.parse import unquote
 
 from fastapi import Depends, Request
 from sqlmodel import Session
-from apps.system.models.system_model import AssistantModel
-from apps.system.schemas.system_schema import UserInfoDTO
+from apps.system.schemas.system_schema import AssistantHeader, UserInfoDTO
 from common.core.db import get_session
 from common.utils.locale import I18n
 
@@ -19,10 +20,14 @@ async def get_current_user(request: Request) -> UserInfoDTO:
 
 CurrentUser = Annotated[UserInfoDTO, Depends(get_current_user)]
 
-async def get_current_assistant(request: Request) -> AssistantModel | None:
-    return request.state.assistant if hasattr(request.state, "assistant") else None
+async def get_current_assistant(request: Request) -> AssistantHeader | None:
+    base_assistant = request.state.assistant if hasattr(request.state, "assistant") else None
+    if request.headers.get("X-SQLBOT-ASSISTANT-CERTIFICATE"):
+        entry_certificate = request.headers['X-SQLBOT-ASSISTANT-CERTIFICATE']
+        base_assistant.certificate = unquote(base64.b64decode(entry_certificate).decode('utf-8'))
+    return base_assistant
 
-CurrentAssistant = Annotated[AssistantModel, Depends(get_current_assistant)]
+CurrentAssistant = Annotated[AssistantHeader, Depends(get_current_assistant)]
 
 
 

frontend/src/utils/request.ts

@@ -75,6 +75,11 @@ class HttpService {
         if (assistantStore.getToken) {
           config.headers['X-SQLBOT-ASSISTANT-TOKEN'] = `Assistant ${assistantStore.getToken}`
           if (config.headers['X-SQLBOT-TOKEN']) config.headers.delete('X-SQLBOT-TOKEN')
+          if (assistantStore.getType && assistantStore.getCertificate) {
+            config.headers['X-SQLBOT-ASSISTANT-CERTIFICATE'] = btoa(
+              encodeURIComponent(assistantStore.getCertificate)
+            )
+          }
         }
         const locale = getLocale()
         if (locale) {
@@ -239,6 +244,11 @@ class HttpService {
     if (assistantStore.getToken) {
       heads['X-SQLBOT-ASSISTANT-TOKEN'] = `Assistant ${assistantStore.getToken}`
       if (heads['X-SQLBOT-TOKEN']) delete heads['X-SQLBOT-TOKEN']
+      if (assistantStore.getType && assistantStore.getCertificate) {
+        heads['X-SQLBOT-ASSISTANT-CERTIFICATE'] = btoa(
+          encodeURIComponent(assistantStore.getCertificate)
+        )
+      }
     }
     // eslint-disable-next-line @typescript-eslint/ban-ts-comment
     // @ts-ignore

frontend/src/views/chat/answer/ChartAnswer.vue

@@ -1,9 +1,7 @@
 <script setup lang="ts">
 import BaseAnswer from './BaseAnswer.vue'
 import { Chat, chatApi, ChatInfo, type ChatMessage, ChatRecord, questionApi } from '@/api/chat.ts'
-import { useAssistantStore } from '@/stores/assistant'
 import { computed, nextTick, onBeforeUnmount, ref } from 'vue'
-const assistantStore = useAssistantStore()
 const props = withDefaults(
   defineProps<{
     chatList?: Array<ChatInfo>
@@ -102,10 +100,7 @@ const sendMessage = async () => {
     const param = {
       question: currentRecord.question,
       chat_id: _currentChatId.value,
-      assistant_certificate: assistantStore.getCertificate,
     }
-    console.log(assistantStore.getCertificate)
-    console.log(param)
     const response = await questionApi.add(param, controller)
     const reader = response.body.getReader()
     const decoder = new TextDecoder()

frontend/src/views/embedded/assistant.vue

@@ -109,7 +109,6 @@ const communicationCb = async (event: any) => {
     }
     if (event.data?.busi == 'certificate') {
       const certificate = event.data['certificate']
-      console.log(certificate)
       assistantStore.setType(1)
       assistantStore.setCertificate(certificate)
       // store certificate to pinia