Skip to content

Commit 070f556

Browse files
perf: Assistant
1 parent bf9bec3 commit 070f556

File tree

6 files changed

+57
-28
lines changed

6 files changed

+57
-28
lines changed

backend/apps/system/api/assistant.py

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from datetime import timedelta
2-
from fastapi import APIRouter, FastAPI, Request
2+
from typing import Optional
3+
from fastapi import APIRouter, FastAPI, Query, Request, Response
34
from sqlmodel import Session, select
45
from apps.system.crud.assistant import get_assistant_info
56
from apps.system.models.system_model import AssistantModel
@@ -13,15 +14,28 @@
1314
from common.core.config import settings
1415
router = APIRouter(tags=["system/assistant"], prefix="/system/assistant")
1516

16-
@router.get("/validator/{id}", response_model=AssistantValidator)
17-
async def info(session: SessionDep, id: str):
18-
as_id, flag = id.split('-')
19-
db_model = get_assistant_info(session, as_id)
17+
@router.get("/info/{id}")
18+
async def info(request: Request, response: Response, session: SessionDep, id: int) -> dict:
19+
db_model = await get_assistant_info(session=session, assistant_id=id)
20+
if not db_model:
21+
raise RuntimeError(f"assistant application not exist")
22+
db_model = AssistantModel.model_validate(db_model)
23+
response.headers["Access-Control-Allow-Origin"] = db_model.domain
24+
origin = request.headers.get("origin") or request.headers.get("referer")
25+
origin = origin.rstrip('/')
26+
""" if origin != db_model.domain:
27+
raise RuntimeError("invalid domain [{origin}]") """
28+
return db_model.model_dump()
29+
30+
@router.get("/validator", response_model=AssistantValidator)
31+
async def info(session: SessionDep, id: str, virtual: Optional[int] = Query(None), online: Optional[bool] = Query(default=False)):
32+
db_model = await get_assistant_info(session=session, assistant_id=id)
2033
if not db_model:
2134
return AssistantValidator()
35+
db_model = AssistantModel.model_validate(db_model)
2236
access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
2337
assistantDict = {
24-
"id": flag, "account": 'sqlbot-inner-assistant', "oid": 1, "assistant_id": id
38+
"id": virtual, "account": 'sqlbot-inner-assistant', "oid": 1, "assistant_id": id
2539
}
2640
access_token = create_access_token(
2741
assistantDict, expires_delta=access_token_expires
@@ -58,9 +72,10 @@ async def update(request: Request, session: SessionDep, editor: AssistantDTO):
5872

5973
@router.get("/{id}", response_model=AssistantModel)
6074
async def get_one(session: SessionDep, id: int):
61-
db_model = get_assistant_info(session, id)
75+
db_model = await get_assistant_info(session=session, assistant_id=id)
6276
if not db_model:
6377
raise ValueError(f"AssistantModel with id {id} not found")
78+
db_model = AssistantModel.model_validate(db_model)
6479
return db_model
6580

6681
@router.delete("/{id}")

backend/apps/system/middleware/auth.py

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import jwt
66
from sqlmodel import Session
77
from starlette.middleware.base import BaseHTTPMiddleware
8+
from apps.system.models.system_model import AssistantModel
89
from common.core.db import engine
910
from apps.system.crud.assistant import get_assistant_info, get_assistant_user
1011
from apps.system.crud.user import get_user_info
@@ -29,12 +30,12 @@ async def dispatch(self, request, call_next):
2930
assistantToken = request.headers.get(assistantTokenKey)
3031
#if assistantToken and assistantToken.lower().startswith("assistant "):
3132
if assistantToken:
32-
validate_pass, data, assistant = await self.validateAssistant(assistantToken)
33-
if validate_pass:
34-
request.state.current_user = data
35-
request.state.assistant = assistant
33+
validator: tuple[any] = await self.validateAssistant(assistantToken)
34+
if validator[0]:
35+
request.state.current_user = validator[1]
36+
request.state.assistant = validator[2]
3637
return await call_next(request)
37-
return JSONResponse({"error": f"Unauthorized:[{data}]"}, status_code=401)
38+
return JSONResponse({"error": f"Unauthorized:[{validator[1]}]"}, status_code=401)
3839
#validate pass
3940
tokenkey = settings.TOKEN_KEY
4041
token = request.headers.get(tokenkey)
@@ -66,7 +67,7 @@ async def validateToken(self, token: Optional[str]):
6667
return False, e
6768

6869

69-
async def validateAssistant(self, assistantToken: Optional[str]):
70+
async def validateAssistant(self, assistantToken: Optional[str]) -> tuple[any]:
7071
if not assistantToken:
7172
return False, f"Miss Token[{settings.TOKEN_KEY}]!"
7273
schema, param = get_authorization_scheme_param(assistantToken)
@@ -83,7 +84,8 @@ async def validateAssistant(self, assistantToken: Optional[str]):
8384
""" session_user = await get_user_info(session = session, user_id = token_data.id)
8485
session_user = UserInfoDTO.model_validate(session_user) """
8586
session_user = get_assistant_user(id = token_data.id)
86-
assistant_info = get_assistant_info(session, payload['assistant_id'])
87+
assistant_info = await get_assistant_info(session=session, assistant_id=payload['assistant_id'])
88+
assistant_info = AssistantModel.model_validate(assistant_info)
8789
return True, session_user, assistant_info
8890
except Exception as e:
8991
return False, e

backend/common/utils/whitelist.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@
2626
"/system/config/key",
2727
"/images/*",
2828
"/sse",
29-
"/system/assistant/validator/*"
29+
"/system/assistant/validator*",
30+
"/system/assistant/info/*"
3031
]
3132

3233
class WhitelistChecker:

frontend/public/assistant.js

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
id: '1',
44
show_guide: false,
55
float_icon: '',
6-
domain_url: 'http://localhost:5173/',
6+
domain_url: 'http://localhost:5173',
77
header_font_color: 'rgb(100, 106, 115)',
88
x_type: 'right',
99
y_type: 'bottom',
@@ -40,7 +40,7 @@
4040
const getChatContainerHtml = (data) => {
4141
return `
4242
<div id="sqlbot-assistant-chat-container">
43-
<iframe id="sqlbot-assistant-chat" allow="microphone" src="${data.domain_url}#/assistant?id=${data.id}"></iframe>
43+
<iframe id="sqlbot-assistant-chat" allow="microphone" src="${data.domain_url}/#/assistant?id=${data.id}"></iframe>
4444
<div class="sqlbot-assistant-operate">
4545
<div class="sqlbot-assistant-closeviewport sqlbot-assistant-viewportnone">
4646
<svg xmlns="http://www.w3.org/2000/svg" width="20" height="20" viewBox="0 0 20 20" fill="none">
@@ -337,19 +337,26 @@
337337

338338
function loadScript(src, id) {
339339
const domain_url = getDomain(src)
340-
// const url = `${domain_url}api/v1/system/assistant/validator/${id}`
341-
/* const url = `http://localhost:8000/api/v1/system/assistant/validator/${id}`
340+
let url = `${domain_url}/api/v1/system/assistant/info/${id}`
341+
if (domain_url.includes('5173')) {
342+
url = url.replace('5173', '8000')
343+
}
342344
fetch(url)
343345
.then((response) => response.json())
344-
.then((data) => {
345-
console.log(data)
346-
const tempData = Object.assign(defaultData, data)
346+
.then((res) => {
347+
const data = res.data
348+
const config_json = data.configuration
349+
let tempData = Object.assign(defaultData, { id, domain_url })
350+
if (config_json) {
351+
const config = JSON.parse(config_json)
352+
Object.assign(tempData, config)
353+
tempData = Object.assign(tempData, config)
354+
}
347355
initsqlbot_assistant(tempData)
348-
}) */
349-
initsqlbot_assistant(defaultData)
356+
})
350357
}
351358
function getDomain(src) {
352-
return src.substring(0, src.indexOf('assistant.js'))
359+
return src.substring(0, src.indexOf('/assistant.js'))
353360
}
354361
function init() {
355362
const sqlbotScripts = document.querySelectorAll(`script[id^="${script_id_prefix}"]`)

frontend/src/api/assistant.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,5 +7,5 @@ export const assistantApi = {
77
edit: (data: any) => request.put('/system/assistant', data),
88
delete: (id: number) => request.delete(`/system/assistant/${id}`),
99
query: (id: number) => request.get(`/system/assistant/${id}`),
10-
validate: (id: any) => request.get(`/system/assistant/validator/${id}`),
10+
validate: (data: any) => request.get('/system/assistant/validator', { params: data }),
1111
}

frontend/src/views/embedded/assistant.vue

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,12 @@ onBeforeMount(async () => {
100100
console.log(online)
101101
const now = Date.now()
102102
assistantStore.setFlag(now)
103-
const id = `${assistantId}-${assistantStore.getFlag}`
104-
validator.value = await assistantApi.validate(id)
103+
const param = {
104+
id: assistantId,
105+
virtual: assistantStore.getFlag,
106+
online,
107+
}
108+
validator.value = await assistantApi.validate(param)
105109
assistantStore.setToken(validator.value.token)
106110
assistantStore.setAssistant(true)
107111
loading.value = false

0 commit comments

Comments
 (0)