diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 14cfdbf4..459dd050 100644 --- a/backend/apps/chat/api/chat.py +++ b/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, diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index a83c6b07..db72e80d 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -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 = '' diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 02746c15..4194e75f 100644 --- a/backend/apps/chat/task/llm.py +++ b/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 @@ 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) diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py index 11934f94..45f451e1 100644 --- a/backend/apps/system/crud/assistant.py +++ b/backend/apps/system/crud/assistant.py @@ -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="sqlbot-inner-assistant@sqlbot.com") -# 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 diff --git a/backend/apps/system/middleware/auth.py b/backend/apps/system/middleware/auth.py index db5bc4f3..89c3b1f3 100644 --- a/backend/apps/system/middleware/auth.py +++ b/backend/apps/system/middleware/auth.py @@ -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 \ No newline at end of file diff --git a/backend/apps/system/schemas/system_schema.py b/backend/apps/system/schemas/system_schema.py index 2183db38..28e30505 100644 --- a/backend/apps/system/schemas/system_schema.py +++ b/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 diff --git a/backend/common/core/deps.py b/backend/common/core/deps.py index be8d1efb..afc2dc36 100644 --- a/backend/common/core/deps.py +++ b/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)] diff --git a/frontend/src/utils/request.ts b/frontend/src/utils/request.ts index 479c948e..2893d9f0 100644 --- a/frontend/src/utils/request.ts +++ b/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 diff --git a/frontend/src/views/chat/answer/ChartAnswer.vue b/frontend/src/views/chat/answer/ChartAnswer.vue index fd478e9d..805a05f6 100644 --- a/frontend/src/views/chat/answer/ChartAnswer.vue +++ b/frontend/src/views/chat/answer/ChartAnswer.vue @@ -1,9 +1,7 @@