diff --git a/backend/apps/system/api/assistant.py b/backend/apps/system/api/assistant.py index 836461ee..6577c3b4 100644 --- a/backend/apps/system/api/assistant.py +++ b/backend/apps/system/api/assistant.py @@ -1,5 +1,6 @@ from datetime import timedelta -from fastapi import APIRouter, FastAPI, Request +from typing import Optional +from fastapi import APIRouter, FastAPI, Query, Request, Response from sqlmodel import Session, select from apps.system.crud.assistant import get_assistant_info from apps.system.models.system_model import AssistantModel @@ -13,15 +14,28 @@ 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: str): - as_id, flag = id.split('-') - db_model = get_assistant_info(session, as_id) +@router.get("/info/{id}") +async def info(request: Request, response: Response, session: SessionDep, id: int) -> dict: + db_model = await get_assistant_info(session=session, assistant_id=id) + if not db_model: + raise RuntimeError(f"assistant application not exist") + db_model = AssistantModel.model_validate(db_model) + response.headers["Access-Control-Allow-Origin"] = db_model.domain + origin = request.headers.get("origin") or request.headers.get("referer") + origin = origin.rstrip('/') + """ if origin != db_model.domain: + raise RuntimeError("invalid domain [{origin}]") """ + return db_model.model_dump() + +@router.get("/validator", response_model=AssistantValidator) +async def info(session: SessionDep, id: str, virtual: Optional[int] = Query(None), online: Optional[bool] = Query(default=False)): + db_model = await get_assistant_info(session=session, assistant_id=id) if not db_model: return AssistantValidator() + db_model = AssistantModel.model_validate(db_model) access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) assistantDict = { - "id": flag, "account": 'sqlbot-inner-assistant', "oid": 1, "assistant_id": id + "id": virtual, "account": 'sqlbot-inner-assistant', "oid": 1, "assistant_id": id } access_token = create_access_token( assistantDict, expires_delta=access_token_expires @@ -58,9 +72,10 @@ 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 = get_assistant_info(session, id) + db_model = await get_assistant_info(session=session, assistant_id=id) if not db_model: raise ValueError(f"AssistantModel with id {id} not found") + db_model = AssistantModel.model_validate(db_model) return db_model @router.delete("/{id}") diff --git a/backend/apps/system/middleware/auth.py b/backend/apps/system/middleware/auth.py index 809ba15d..f20fba96 100644 --- a/backend/apps/system/middleware/auth.py +++ b/backend/apps/system/middleware/auth.py @@ -5,6 +5,7 @@ from fastapi.responses import JSONResponse import jwt from sqlmodel import Session from starlette.middleware.base import BaseHTTPMiddleware +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 @@ -29,12 +30,12 @@ class TokenMiddleware(BaseHTTPMiddleware): 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 + validator: tuple[any] = await self.validateAssistant(assistantToken) + if validator[0]: + request.state.current_user = validator[1] + request.state.assistant = validator[2] return await call_next(request) - return JSONResponse({"error": f"Unauthorized:[{data}]"}, status_code=401) + return JSONResponse({"error": f"Unauthorized:[{validator[1]}]"}, status_code=401) #validate pass tokenkey = settings.TOKEN_KEY token = request.headers.get(tokenkey) @@ -66,7 +67,7 @@ class TokenMiddleware(BaseHTTPMiddleware): return False, e - async def validateAssistant(self, assistantToken: Optional[str]): + async def validateAssistant(self, assistantToken: Optional[str]) -> tuple[any]: if not assistantToken: return False, f"Miss Token[{settings.TOKEN_KEY}]!" schema, param = get_authorization_scheme_param(assistantToken) @@ -83,7 +84,8 @@ class TokenMiddleware(BaseHTTPMiddleware): """ 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']) + assistant_info = await get_assistant_info(session=session, assistant_id=payload['assistant_id']) + assistant_info = AssistantModel.model_validate(assistant_info) return True, session_user, assistant_info except Exception as e: return False, e \ No newline at end of file diff --git a/backend/common/utils/whitelist.py b/backend/common/utils/whitelist.py index 137104c9..10fa0c02 100644 --- a/backend/common/utils/whitelist.py +++ b/backend/common/utils/whitelist.py @@ -26,7 +26,8 @@ wlist = [ "/system/config/key", "/images/*", "/sse", - "/system/assistant/validator/*" + "/system/assistant/validator*", + "/system/assistant/info/*" ] class WhitelistChecker: diff --git a/frontend/public/assistant.js b/frontend/public/assistant.js index 5b3d0fda..eb2d51a6 100644 --- a/frontend/public/assistant.js +++ b/frontend/public/assistant.js @@ -3,7 +3,7 @@ id: '1', show_guide: false, float_icon: '', - domain_url: 'http://localhost:5173/', + domain_url: 'http://localhost:5173', header_font_color: 'rgb(100, 106, 115)', x_type: 'right', y_type: 'bottom', @@ -40,7 +40,7 @@ const getChatContainerHtml = (data) => { return `