mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
perf: Assistant dynamic certificate
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -151,8 +151,6 @@ class AiModelQuestion(BaseModel):
|
||||
class ChatQuestion(AiModelQuestion):
|
||||
question: str = ''
|
||||
chat_id: int = 0
|
||||
assistant_certificate: Optional[str] = None
|
||||
|
||||
|
||||
class ChatMcp(ChatQuestion):
|
||||
token: str = ''
|
||||
|
||||
@@ -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 @@ class LLMService:
|
||||
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)
|
||||
|
||||
@@ -10,7 +10,7 @@ from apps.datasource.models.datasource import CoreDatasource, DatasourceConf
|
||||
|
||||
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 @@ class AssistantOutDs:
|
||||
|
||||
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
|
||||
|
||||
@@ -9,7 +9,7 @@ from apps.system.models.system_model import AssistantModel
|
||||
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 @@ class TokenMiddleware(BaseHTTPMiddleware):
|
||||
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
|
||||
@@ -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
|
||||
|
||||
@@ -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)]
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user