perf: Assistant token

Author: fit2cloud-chenyw

backend/apps/system/api/assistant.py

@@ -1,25 +1,29 @@
 from datetime import timedelta
 from fastapi import APIRouter, FastAPI, Request
 from sqlmodel import Session, select
-from apps.system.crud.user import get_user_by_account
+from apps.system.crud.assistant import get_assistant_info
 from apps.system.models.system_model import AssistantModel
+from apps.system.schemas.auth import CacheName, CacheNamespace
 from apps.system.schemas.system_schema import AssistantBase, AssistantDTO, AssistantValidator
 from common.core.deps import SessionDep
 from common.core.security import create_access_token
+from common.core.sqlbot_cache import clear_cache
 from common.utils.time import get_timestamp
 from starlette.middleware.cors import CORSMiddleware
 from common.core.config import settings
 router = APIRouter(tags=["system/assistant"], prefix="/system/assistant")
 
 @router.get("/validator/{id}", response_model=AssistantValidator) 
 async def info(session: SessionDep, id: int):
-    db_model = session.get(AssistantModel, id)
+    db_model = get_assistant_info(session, id)
     if not db_model:
         return AssistantValidator()
     access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
-    user = get_user_by_account(session=session, account='admin')
+    assistantDict = {
+        "id": 1, "account": 'admin', "oid": 1, "assistant_id": id
+    }
     access_token = create_access_token(
-        user.to_dict(), expires_delta=access_token_expires
+        assistantDict, expires_delta=access_token_expires
     )
     return AssistantValidator(True, True, True, access_token)
 
@@ -39,6 +43,7 @@ async def add(request: Request, session: SessionDep, creator: AssistantBase):
 
 
 @router.put("")
+@clear_cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="editor.id")  
 async def update(request: Request, session: SessionDep, editor: AssistantDTO):
     id = editor.id
     db_model = session.get(AssistantModel, id)
@@ -52,12 +57,13 @@ async def update(request: Request, session: SessionDep, editor: AssistantDTO):
 
 @router.get("/{id}", response_model=AssistantModel)    
 async def get_one(session: SessionDep, id: int):
-    db_model = session.get(AssistantModel, id)
+    db_model = get_assistant_info(session, id)
     if not db_model:
         raise ValueError(f"AssistantModel with id {id} not found")
     return db_model
 
-@router.delete("/{id}")  
+@router.delete("/{id}")
+@clear_cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="id")  
 async def delete(request: Request, session: SessionDep, id: int):
     db_model = session.get(AssistantModel, id)
     if not db_model:

backend/apps/system/crud/assistant.py

@@ -0,0 +1,12 @@
+
+
+from sqlmodel import Session
+from apps.system.models.system_model import AssistantModel
+from apps.system.schemas.auth import CacheName, CacheNamespace
+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

backend/apps/system/middleware/auth.py

@@ -1,28 +1,88 @@
 
-from fastapi import Depends
+from typing import Optional
+from fastapi import Request
 from fastapi.responses import JSONResponse
+import jwt
+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.user import get_user_info
+from apps.system.schemas.system_schema import UserInfoDTO
+from common.core import security
 from common.core.config import settings
-# from common.core.deps import get_current_user
+from common.core.schemas import TokenPayload
 from common.utils.whitelist import whiteUtils
-
+from fastapi.security.utils import get_authorization_scheme_param
 class TokenMiddleware(BaseHTTPMiddleware):
 
+    
+    
     def __init__(self, app):
         super().__init__(app)
 
     async def dispatch(self, request, call_next):
-        tokenkey = settings.TOKEN_KEY
+        
         if self.is_options(request) or whiteUtils.is_whitelisted(request.url.path):
             return await call_next(request)
+        assistantTokenKey = settings.ASSISTANT_TOKEN_KEY
+        assistantToken = request.headers.get(assistantTokenKey)
+        #if assistantToken and assistantToken.lower().startswith("assistant "):
+        if assistantToken:
+            validate_pass, data, assistant = await self.validateAssistant(assistantToken)
+            if validate_pass:
+                request.state.current_user = data
+                request.state.assistant = assistant
+                return await call_next(request)
+            return JSONResponse({"error": f"Unauthorized:[{data}]"}, status_code=401)
+        #validate pass
+        tokenkey = settings.TOKEN_KEY
         token = request.headers.get(tokenkey)
-        if not token or not token.startswith("Bearer "):
-            return JSONResponse({"error": "Unauthorized"}, status_code=401)
-        """ user = await get_current_user()
-        request.state.user = user """
-        return await call_next(request)
+        validate_pass, data = await self.validateToken(token)
+        if validate_pass:
+            request.state.current_user = data
+            return await call_next(request)
+        return JSONResponse({"error": f"Unauthorized:[{data}]"}, status_code=401)
 
-    def is_options(self, request):
+    def is_options(self, request: Request):
         return request.method == "OPTIONS"
 
-    
+    async def validateToken(self, token: Optional[str]):
+        if not token:
+            return False, f"Miss Token[{settings.TOKEN_KEY}]!"
+        schema, param = get_authorization_scheme_param(token)
+        if schema.lower() != "bearer":
+            return False, f"Token schema error!"
+        payload = jwt.decode(
+            param, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
+        )
+        token_data = TokenPayload(**payload)
+        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)
+                return True, session_user
+        except Exception as e:
+            return False, e
+            
+    
+    async def validateAssistant(self, assistantToken: Optional[str]):
+        if not assistantToken:
+            return False, f"Miss Token[{settings.TOKEN_KEY}]!"
+        schema, param = get_authorization_scheme_param(assistantToken)
+        if schema.lower() != "assistant":
+            return False, f"Token schema error!"
+        payload = jwt.decode(
+            param, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
+        )
+        token_data = TokenPayload(**payload)
+        if not payload['assistant_id']:
+            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)
+                assistant_info = get_assistant_info(session, payload['assistant_id'])
+                return True, session_user, assistant_info
+        except Exception as e:
+            return False, e

backend/apps/system/schemas/auth.py

@@ -8,10 +8,11 @@ class LocalLoginSchema(BaseModel):
 
 class CacheNamespace(Enum):
     AUTH_INFO = "sqlbot:auth"
+    EMBEDDED_INFO = "sqlbot:embedded"
     def __str__(self):
         return self.value
 class CacheName(Enum):
     USER_INFO = "user:info"
-
+    ASSISTANT_INFO = "assistant:info"
     def __str__(self):
         return self.value

backend/common/core/config.py

@@ -57,8 +57,9 @@ def all_cors_origins(self) -> list[str]:
     POSTGRES_PASSWORD: str = ""
     POSTGRES_DB: str = ""
 
-    TOKEN_KEY: str
-    DEFAULT_PWD: str
+    TOKEN_KEY: str =  "X-SQLBOT-TOKEN"
+    DEFAULT_PWD: str = "SQLBot@123456"
+    ASSISTANT_TOKEN_KEY: str = "X-SQLBOT-ASSISTANT-TOKEN"
 
     @computed_field  # type: ignore[prop-decorator]
     @property

backend/common/core/deps.py

@@ -1,51 +1,21 @@
 from typing import Annotated
 
-import jwt
-from fastapi import Depends, HTTPException, Request, status
-# from fastapi.security import OAuth2PasswordBearer
-from jwt.exceptions import InvalidTokenError
-from pydantic import ValidationError
+from fastapi import Depends, Request
 from sqlmodel import Session
-from apps.system.crud.user import get_user_info
 from apps.system.schemas.system_schema import UserInfoDTO
-from common.core.schemas import TokenPayload, XOAuth2PasswordBearer
-from common.core import security
-from common.core.config import settings
 from common.core.db import get_session
 from common.utils.locale import I18n
-reusable_oauth2 = XOAuth2PasswordBearer(
-    tokenUrl=f"{settings.API_V1_STR}/login/access-token"
-)
 
 
-    
-
 SessionDep = Annotated[Session, Depends(get_session)]
-TokenDep = Annotated[str, Depends(reusable_oauth2)]
 i18n = I18n()
 async def get_i18n(request: Request):
     return i18n(request)
 
 Trans = Annotated[I18n, Depends(get_i18n)]
-async def get_current_user(session: SessionDep, token: TokenDep) -> UserInfoDTO:
-    try:
-        payload = jwt.decode(
-            token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
-        )
-        token_data = TokenPayload(**payload)
-    except (InvalidTokenError, ValidationError):
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail="Could not validate credentials",
-        )
-    session_user = await get_user_info(session = session, user_id = token_data.id)
-    session_user = UserInfoDTO.model_validate(session_user)
-    if not session_user:
-        raise HTTPException(status_code=404, detail="User not found")
-    
-    if session_user.status != 1:
-        raise HTTPException(status_code=400, detail="Inactive user")
-    return session_user
+async def get_current_user(request: Request) -> UserInfoDTO:
+    return request.state.current_user
+
 CurrentUser = Annotated[UserInfoDTO, Depends(get_current_user)]
 
 

backend/common/core/schemas.py

@@ -18,8 +18,11 @@ class Token(SQLModel):
 class XOAuth2PasswordBearer(OAuth2PasswordBearer):
     async def __call__(self, request: Request) -> Optional[str]:
         authorization = request.headers.get(settings.TOKEN_KEY)
+        if request.headers.get(settings.ASSISTANT_TOKEN_KEY):
+            authorization = request.headers.get(settings.ASSISTANT_TOKEN_KEY)
         scheme, param = get_authorization_scheme_param(authorization)
-        if not authorization or scheme.lower() != "bearer":
+        
+        if not authorization or scheme.lower() not in  ["bearer", "assistant"]:
             if self.auto_error:
                 raise HTTPException(
                     status_code=HTTP_401_UNAUTHORIZED,

frontend/src/stores/assistant.ts

@@ -0,0 +1,31 @@
+import { defineStore } from 'pinia'
+import { store } from './index'
+
+interface AssistantState {
+  token: string
+}
+
+export const AssistantStore = defineStore('assistant', {
+  state: (): AssistantState => {
+    return {
+      token: '',
+    }
+  },
+  getters: {
+    getToken(): string {
+      return this.token
+    },
+  },
+  actions: {
+    setToken(token: string) {
+      this.token = token
+    },
+    clear() {
+      this.$reset()
+    },
+  },
+})
+
+export const useAssistantStore = () => {
+  return AssistantStore(store)
+}

frontend/src/utils/request.ts

@@ -10,6 +10,8 @@ import axios, {
 
 import { useCache } from '@/utils/useCache'
 import { getLocale } from './utils'
+import { useAssistantStore } from '@/stores/assistant'
+const assistantStore = useAssistantStore()
 const { wsCache } = useCache()
 // Response data structure
 export interface ApiResponse<T = unknown> {
@@ -70,6 +72,10 @@ class HttpService {
         if (token && config.headers) {
           config.headers['X-SQLBOT-TOKEN'] = `Bearer ${token}`
         }
+        if (assistantStore.getToken) {
+          config.headers['X-SQLBOT-ASSISTANT-TOKEN'] = `Assistant ${assistantStore.getToken}`
+          if (config.headers['X-SQLBOT-TOKEN']) config.headers.delete('X-SQLBOT-TOKEN')
+        }
         const locale = getLocale()
         if (locale) {
           /* const mapping = {

frontend/src/views/embedded/assistant.vue

@@ -62,8 +62,9 @@ import IconOpeEdit from '@/assets/svg/operate/ope-edit.svg'
 import IconOpeDelete from '@/assets/svg/operate/ope-delete.svg'
 import { useRoute } from 'vue-router'
 import { assistantApi } from '@/api/assistant'
-import { useUserStore } from '@/stores/user'
-const userStore = useUserStore()
+import { useAssistantStore } from '@/stores/assistant'
+
+const assistantStore = useAssistantStore()
 const route = useRoute()
 
 const chatRef = ref()
@@ -96,7 +97,7 @@ const loading = ref(true)
 onBeforeMount(async () => {
   const assistantId = route.params.id
   validator.value = await assistantApi.validate(assistantId)
-  userStore.setToken(validator.value.token)
+  assistantStore.setToken(validator.value.token)
   loading.value = false
 })
 </script>