perf: Assistant chat history

Author: fit2cloud-chenyw

backend/apps/chat/curd/chat.py

@@ -4,6 +4,7 @@
 import orjson
 from sqlalchemy import and_
 from sqlalchemy.orm import load_only
+from sqlmodel import select
 
 from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat, ChatQuestion
 from apps.datasource.models.datasource import CoreDatasource
@@ -12,8 +13,7 @@
 
 
 def list_chats(session: SessionDep, current_user: CurrentUser) -> List[Chat]:
-    chart_list = session.query(Chat).filter(Chat.create_by == current_user.id).order_by(
-        Chat.create_time.desc()).all()
+    chart_list = session.exec(select(Chat).where(Chat.create_by == current_user.id).order_by(Chat.create_time.desc())).all()
     return chart_list
 
 

backend/apps/system/api/assistant.py

@@ -14,13 +14,14 @@
 router = APIRouter(tags=["system/assistant"], prefix="/system/assistant")
 
 @router.get("/validator/{id}", response_model=AssistantValidator) 
-async def info(session: SessionDep, id: int):
-    db_model = get_assistant_info(session, id)
+async def info(session: SessionDep, id: str):
+    as_id, flag = id.split('-')
+    db_model = get_assistant_info(session, as_id)
     if not db_model:
         return AssistantValidator()
     access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
     assistantDict = {
-        "id": 1, "account": 'admin', "oid": 1, "assistant_id": id
+        "id": flag, "account": 'sqlbot-inner-assistant', "oid": 1, "assistant_id": id
     }
     access_token = create_access_token(
         assistantDict, expires_delta=access_token_expires

backend/apps/system/crud/assistant.py

@@ -3,10 +3,14 @@
 from sqlmodel import Session
 from apps.system.models.system_model import AssistantModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
+from apps.system.schemas.system_schema import UserInfoDTO
 from common.core.sqlbot_cache import cache
 
 
 @cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="assistant_id")
 async def get_assistant_info(*, session: Session, assistant_id: int) -> AssistantModel | None:
     db_model = session.get(AssistantModel, assistant_id)
-    return db_model
+    return db_model
+
+def get_assistant_user(*, id: int):
+    return UserInfoDTO(id=id, account="sqlbot-inner-assistant", oid=1, name="sqlbot-inner-assistant", email="[email protected]")

backend/apps/system/middleware/auth.py

@@ -6,7 +6,7 @@
 from sqlmodel import Session
 from starlette.middleware.base import BaseHTTPMiddleware
 from common.core.db import engine 
-from apps.system.crud.assistant import get_assistant_info
+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 common.core import security
@@ -80,8 +80,9 @@ async def validateAssistant(self, assistantToken: Optional[str]):
             return False, f"Miss assistant payload error!"
         try: 
             with Session(engine) as session:
-                session_user = await get_user_info(session = session, user_id = token_data.id)
-                session_user = UserInfoDTO.model_validate(session_user)
+                """ session_user = await get_user_info(session = session, user_id = token_data.id)
+                session_user = UserInfoDTO.model_validate(session_user) """
+                session_user = get_assistant_user(id = token_data.id)
                 assistant_info = get_assistant_info(session, payload['assistant_id'])
                 return True, session_user, assistant_info
         except Exception as e:

backend/common/core/deps.py

@@ -2,6 +2,7 @@
 
 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 common.core.db import get_session
 from common.utils.locale import I18n
@@ -18,5 +19,10 @@ async def get_current_user(request: Request) -> UserInfoDTO:
 
 CurrentUser = Annotated[UserInfoDTO, Depends(get_current_user)]
 
+async def get_current_assistant(request: Request) -> AssistantModel:
+    return request.state.assistant
+
+CurrentAssistant = Annotated[AssistantModel, Depends(get_current_assistant)]
+
 
 

frontend/src/stores/assistant.ts

@@ -1,17 +1,22 @@
 import { defineStore } from 'pinia'
 import { store } from './index'
 import { chatApi, ChatInfo } from '@/api/chat'
+import { useCache } from '@/utils/useCache'
 
+const { wsCache } = useCache()
+const flagKey = 'sqlbit-assistant-flag'
 interface AssistantState {
   token: string
   assistant: boolean
+  flag: number
 }
 
 export const AssistantStore = defineStore('assistant', {
   state: (): AssistantState => {
     return {
       token: '',
       assistant: false,
+      flag: 0,
     }
   },
   getters: {
@@ -21,6 +26,9 @@ export const AssistantStore = defineStore('assistant', {
     getAssistant(): boolean {
       return this.assistant
     },
+    getFlag(): number {
+      return this.flag
+    },
   },
   actions: {
     setToken(token: string) {
@@ -29,6 +37,14 @@ export const AssistantStore = defineStore('assistant', {
     setAssistant(assistant: boolean) {
       this.assistant = assistant
     },
+    setFlag(flag: number) {
+      if (wsCache.get(flagKey)) {
+        this.flag = wsCache.get(flagKey)
+      } else {
+        this.flag = flag
+        wsCache.set(flagKey, flag)
+      }
+    },
     async setChat() {
       if (!this.assistant) {
         return null
@@ -38,6 +54,7 @@ export const AssistantStore = defineStore('assistant', {
       return chat
     },
     clear() {
+      wsCache.delete(flagKey)
       this.$reset()
     },
   },

frontend/src/utils/request.ts

@@ -238,7 +238,7 @@ class HttpService {
     }
     if (assistantStore.getToken) {
       heads['X-SQLBOT-ASSISTANT-TOKEN'] = `Assistant ${assistantStore.getToken}`
-      if (heads['X-SQLBOT-TOKEN']) heads.delete('X-SQLBOT-TOKEN')
+      if (heads['X-SQLBOT-TOKEN']) delete heads['X-SQLBOT-TOKEN']
     }
     // eslint-disable-next-line @typescript-eslint/ban-ts-comment
     // @ts-ignore

frontend/src/views/embedded/assistant.vue

@@ -96,7 +96,10 @@ const validator = ref({
 const loading = ref(true)
 onBeforeMount(async () => {
   const assistantId = route.params.id
-  validator.value = await assistantApi.validate(assistantId)
+  const now = Date.now()
+  assistantStore.setFlag(now)
+  const id = `${assistantId}-${assistantStore.getFlag}`
+  validator.value = await assistantApi.validate(id)
   assistantStore.setToken(validator.value.token)
   assistantStore.setAssistant(true)
   loading.value = false