Skip to content

Commit 55dba79

Browse files
dataeaseShudataeaseShu
committed
Merge branch 'main' of https://github.com/dataease/SQLBot 
2 parents 6e97f7b + 9b2aead commit 55dba79

File tree

10 files changed

+178
-32
lines changed

10 files changed

+178
-32
lines changed

‎.typos.toml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
[default.extend-words]
2+
[default.extend-identifiers]
3+
maintain_column_froms = "maintain_column_froms"
4+
[files]
5+
extend-exclude = ["frontend/public", "backend/alembic/versions"]

‎backend/apps/chat/task/llm.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -296,7 +296,7 @@ def generate_recommend_questions_task(self):
296296

297297
# get schema
298298
if self.ds and not self.chat_question.db_schema:
299-
self.chat_question.db_schema = get_table_schema(session=self.session, ds=self.ds)
299+
self.chat_question.db_schema = get_table_schema(session=self.session, current_user=self.current_user, ds=self.ds)
300300

301301
guess_msg: List[Union[BaseMessage, dict[str, Any]]] = []
302302
guess_msg.append(SystemMessage(content=self.chat_question.guess_sys_question()))
@@ -726,7 +726,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
726726
yield orjson.dumps({'id': llm_service.ds.id, 'datasource_name': llm_service.ds.name,
727727
'engine_type': llm_service.ds.type_name, 'type': 'datasource'}).decode() + '\n\n'
728728

729-
llm_service.chat_question.db_schema = get_table_schema(session=llm_service.session, ds=llm_service.ds)
729+
llm_service.chat_question.db_schema = get_table_schema(session=llm_service.session, current_user=llm_service.current_user, ds=llm_service.ds)
730730

731731
# generate sql
732732
sql_res = llm_service.generate_sql()

‎backend/apps/datasource/api/datasource.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from apps.db.engine import create_table, get_data_engine, insert_data
1010
from common.core.deps import SessionDep, CurrentUser
1111
from ..crud.datasource import get_datasource_list, check_status, create_ds, update_ds, delete_ds, getTables, getFields, \
12-
execSql, update_table_and_fields, getTablesByDs, chooseTables, preview, updateTable, updateField, get_ds
12+
execSql, update_table_and_fields, getTablesByDs, chooseTables, preview, updateTable, updateField, get_ds, fieldEnum
1313
from ..crud.field import get_fields_by_table_id
1414
from ..crud.table import get_tables_by_ds_id
1515
from ..models.datasource import CoreDatasource, CreateDatasource, TableObj, CoreTable, CoreField
@@ -103,6 +103,11 @@ async def edit_local(session: SessionDep, id: int, data: TableObj):
103103
return preview(session, id, data)
104104

105105

106+
@router.post("/fieldEnum/{id}")
107+
async def field_enum(session: SessionDep, id: int):
108+
return fieldEnum(session, id)
109+
110+
106111
@router.post("/uploadExcel")
107112
async def upload_excel(session: SessionDep, file: UploadFile = File(...)):
108113
ALLOWED_EXTENSIONS = {"xlsx", "xls", "csv"}

‎backend/apps/datasource/crud/datasource.py‎

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,9 @@
99
from sqlmodel import select
1010

1111
from apps.datasource.utils.utils import aes_decrypt
12+
from apps.db.constant import DB
1213
from apps.db.db import get_engine, get_tables, get_fields, exec_sql
13-
from apps.db.engine import get_engine_config
14-
from apps.db.engine import get_engine_conn
14+
from apps.db.engine import get_engine_config, get_engine_conn
1515
from apps.db.type import db_type_relation
1616
from common.core.deps import SessionDep, CurrentUser
1717
from common.utils.utils import deepcopy_ignore_extra
@@ -245,6 +245,23 @@ def preview(session: SessionDep, id: int, data: TableObj):
245245
return exec_sql(ds, sql)
246246

247247

248+
def fieldEnum(session: SessionDep, id: int):
249+
field = session.query(CoreField).filter(CoreField.id == id).first()
250+
if field is None:
251+
return []
252+
table = session.query(CoreTable).filter(CoreTable.id == field.table_id).first()
253+
if table is None:
254+
return []
255+
ds = session.query(CoreDatasource).filter(CoreDatasource.id == table.ds_id).first()
256+
if ds is None:
257+
return []
258+
259+
db = DB.get_db(ds.type)
260+
sql = f"""SELECT DISTINCT {db.prefix}{field.field_name}{db.suffix} FROM {db.prefix}{table.table_name}{db.suffix}"""
261+
res = exec_sql(ds, sql)
262+
return [item.get(res.get('fields')[0]) for item in res.get('data')]
263+
264+
248265
def updateNum(session: SessionDep, ds: CoreDatasource):
249266
all_tables = get_tables(ds)
250267
selected_tables = get_tables_by_ds_id(session, ds.id)

‎backend/apps/db/constant.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
# Author: Junjun
2+
# Date: 2025/7/16
3+
4+
from enum import Enum
5+
6+
7+
class DB(Enum):
8+
mysql = ('mysql', '`', '`')
9+
sqlServer = ('sqlServer', '[', ']')
10+
pg = ('pg', '"', '"')
11+
excel = ('excel', '"', '"')
12+
oracle = ('oracle', '"', '"')
13+
14+
def __init__(self, type, prefix, suffix):
15+
self.type = type
16+
self.prefix = prefix
17+
self.suffix = suffix
18+
19+
@classmethod
20+
def get_db(cls, type):
21+
for db in cls:
22+
if db.type == type:
23+
return db
24+
raise ValueError(f"Invalid db type: {type}")

‎backend/apps/system/api/workspace.py‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from sqlmodel import exists, or_, select
44
from apps.system.models.system_model import UserWsModel, WorkspaceBase, WorkspaceEditor, WorkspaceModel
55
from apps.system.models.user import UserModel
6-
from apps.system.schemas.system_schema import UserWsBase, UserWsDTO, UserWsOption, WorkspaceUser
6+
from apps.system.schemas.system_schema import UserWsBase, UserWsDTO, UserWsEditor, UserWsOption, WorkspaceUser
77
from common.core.deps import CurrentUser, SessionDep, Trans
88
from common.core.pagination import Paginator
99
from common.core.schemas import PaginatedResponse, PaginationParams
@@ -42,6 +42,29 @@ async def option_pager(
4242
stmt=stmt,
4343
pagination=pagination,
4444
)
45+
46+
@router.get("/uws/option", response_model=UserWsOption)
47+
async def option_user(
48+
session: SessionDep,
49+
current_user: CurrentUser,
50+
keyword: str = Query(description="搜索关键字")
51+
):
52+
if (not current_user.isAdmin) and current_user.weight == 0:
53+
raise RuntimeError("no permission to execute this api")
54+
oid = current_user.oid
55+
56+
stmt = select(UserModel.id, UserModel.account, UserModel.name).where(
57+
~exists().where(UserWsModel.uid == UserModel.id, UserWsModel.oid == oid)
58+
)
59+
60+
if keyword:
61+
stmt = stmt.where(
62+
or_(
63+
UserModel.account == int(keyword),
64+
UserModel.name == keyword,
65+
)
66+
)
67+
return session.exec(stmt).first()
4568

4669
@router.get("/uws/pager/{pageNum}/{pageSize}", response_model=PaginatedResponse[WorkspaceUser])
4770
async def pager(
@@ -95,6 +118,20 @@ async def create(session: SessionDep, creator: UserWsDTO):
95118
session.add_all(db_model_list)
96119
session.commit()
97120

121+
@router.put("/uws")
122+
async def edit(session: SessionDep, editor: UserWsEditor):
123+
if not editor.oid or not editor.uid:
124+
raise RuntimeError("param [oid, uid] miss")
125+
db_model = session.exec(select(UserWsModel).where(UserWsModel.uid == editor.uid, UserWsModel.oid == editor.oid)).first()
126+
if not db_model:
127+
raise RuntimeError("uws not exist")
128+
if editor.weight == db_model.weight:
129+
return
130+
131+
db_model.weight = editor.weight
132+
session.add(db_model)
133+
session.commit()
134+
98135
@router.delete("/uws")
99136
async def delete(session: SessionDep, dto: UserWsBase):
100137
db_model_list: list[UserWsModel] = session.exec(select(UserWsModel).where(UserWsModel.uid.in_(dto.uid_list), UserWsModel.oid == dto.oid)).all()

‎backend/apps/system/schemas/system_schema.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,11 @@ class UserWsBase(BaseModel):
5050
oid: int
5151
class UserWsDTO(UserWsBase):
5252
weight: int = 0
53-
53+
54+
class UserWsEditor(BaseModel):
55+
uid: int
56+
oid: int
57+
weight: int = 0
5458

5559
class UserInfoDTO(UserEditor):
5660
language: str = "zh-CN"

‎frontend/src/assets/svg/icon_new_chat_outlined.svg‎

Lines changed: 1 addition & 4 deletions
Loading

‎frontend/src/views/chat/index.vue‎

Lines changed: 58 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -72,11 +72,14 @@
7272
v-if="currentChatId === undefined"
7373
size="large"
7474
type="primary"
75+
class="greeting-btn"
7576
@click="createNewChat"
7677
>
77-
<el-icon>
78-
<icon_new_chat_outlined />
79-
</el-icon>
78+
<span class="inner-icon">
79+
<el-icon>
80+
<icon_new_chat_outlined />
81+
</el-icon>
82+
</span>
8083
{{ t('qa.start_sqlbot') }}
8184
</el-button>
8285
</div>
@@ -292,7 +295,7 @@ const createNewChat = async () => {
292295
if (isAssistant.value) {
293296
const assistantChat = await assistantStore.setChat()
294297
if (assistantChat) {
295-
onChatCreated(assistantChat as any)
298+
onChatCreatedQuick(assistantChat as any)
296299
}
297300
return
298301
}
@@ -317,6 +320,26 @@ function onClickHistory(chat: Chat) {
317320
scrollToBottom()
318321
}
319322
323+
function toAssistantHistory(chat: Chat) {
324+
currentChat.value = new ChatInfo(chat)
325+
if (chat !== undefined && chat.id !== undefined && !loading.value) {
326+
currentChatId.value = chat.id
327+
loading.value = true
328+
chatApi
329+
.get(chat.id)
330+
.then((res) => {
331+
const info = chatApi.toChatInfo(res)
332+
if (info) {
333+
currentChat.value = info
334+
onClickHistory(info)
335+
}
336+
})
337+
.finally(() => {
338+
loading.value = false
339+
})
340+
}
341+
}
342+
320343
function onChatDeleted(id: number) {
321344
console.log('deleted', id)
322345
}
@@ -434,11 +457,6 @@ const sendMessage = async () => {
434457
loading.value = true
435458
isTyping.value = true
436459
437-
/* const assistantChat = await assistantStore.setChat()
438-
if (assistantChat) {
439-
onChatCreated(assistantChat as any)
440-
} */
441-
442460
const currentRecord = new ChatRecord()
443461
currentRecord.create_time = new Date()
444462
currentRecord.chat_id = currentChatId.value
@@ -833,10 +851,9 @@ const getCurrentChatId = () => {
833851
}
834852
defineExpose({
835853
getHistoryList,
836-
onClickHistory,
837-
onChatDeleted,
838-
onChatRenamed,
854+
toAssistantHistory,
839855
getCurrentChatId,
856+
createNewChat,
840857
})
841858
</script>
842859

@@ -879,7 +896,6 @@ defineExpose({
879896
880897
.chat-record-list {
881898
padding: 0 0 20px 0;
882-
background: rgba(255, 255, 255, 1);
883899
border-radius: 0 12px 12px 0;
884900
885901
&.hide-sidebar {
@@ -974,6 +990,35 @@ defineExpose({
974990
font-size: 16px;
975991
line-height: 24px;
976992
}
993+
994+
.greeting-btn {
995+
width: 100%;
996+
height: 88px;
997+
998+
border-style: dashed;
999+
1000+
.inner-icon {
1001+
display: flex;
1002+
flex-direction: row;
1003+
align-items: center;
1004+
1005+
margin-right: 6px;
1006+
}
1007+
1008+
font-size: 16px;
1009+
line-height: 24px;
1010+
font-weight: 500;
1011+
1012+
--ed-button-text-color: rgba(28, 186, 144, 1);
1013+
--ed-button-hover-text-color: rgba(28, 186, 144, 1);
1014+
--ed-button-active-text-color: rgba(28, 186, 144, 1);
1015+
--ed-button-bg-color: rgba(248, 249, 250, 1);
1016+
--ed-button-hover-bg-color: rgba(28, 186, 144, 0.1);
1017+
--ed-button-border-color: rgba(217, 220, 223, 1);
1018+
--ed-button-hover-border-color: rgba(28, 186, 144, 1);
1019+
--ed-button-active-bg-color: rgba(28, 186, 144, 0.2);
1020+
--ed-button-active-border-color: rgba(28, 186, 144, 1);
1021+
}
9771022
}
9781023
}
9791024
</style>

‎frontend/src/views/embedded/assistant.vue‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@
1111
<chat-component v-if="!loading" ref="chatRef" class="sqlbot-chat-container" />
1212
</div>
1313
<div class="sqlbot-top-btn">
14+
<el-icon style="cursor: pointer" @click="createChat">
15+
<icon_new_chat_outlined />
16+
</el-icon>
1417
<el-icon style="cursor: pointer" @click="openHistory">
1518
<history></history>
1619
</el-icon>
@@ -41,10 +44,10 @@
4144
@click="onClickHistory(chat)"
4245
>
4346
<span class="title">{{ chat.brief ?? 'Untitled' }}</span>
44-
<div class="history-operate">
47+
<!-- <div class="history-operate">
4548
<el-icon @click="EditPen(chat)"><IconOpeEdit /></el-icon>
4649
<el-icon @click="Delete(chat)"><IconOpeDelete /></el-icon>
47-
</div>
50+
</div> -->
4851
</div>
4952
</template>
5053
</el-scrollbar>
@@ -58,8 +61,9 @@ import { onBeforeMount, ref } from 'vue'
5861
import ChatComponent from '@/views/chat/index.vue'
5962
import AssistantGif from '@/assets/img/assistant.gif'
6063
import history from '@/assets/svg/chart/history.svg'
61-
import IconOpeEdit from '@/assets/svg/operate/ope-edit.svg'
62-
import IconOpeDelete from '@/assets/svg/operate/ope-delete.svg'
64+
import icon_new_chat_outlined from '@/assets/svg/icon_new_chat_outlined.svg'
65+
/* import IconOpeEdit from '@/assets/svg/operate/ope-edit.svg'
66+
import IconOpeDelete from '@/assets/svg/operate/ope-delete.svg' */
6367
import { useRoute } from 'vue-router'
6468
import { assistantApi } from '@/api/assistant'
6569
import { useAssistantStore } from '@/stores/assistant'
@@ -72,21 +76,24 @@ const chatList = ref<Array<any>>([])
7276
const drawer = ref(false)
7377
const currentId = ref()
7478
79+
const createChat = () => {
80+
chatRef.value?.createNewChat()
81+
}
7582
const openHistory = () => {
7683
chatList.value = chatRef.value?.getHistoryList()
7784
currentId.value = chatRef.value?.getCurrentChatId()
7885
drawer.value = true
7986
}
8087
const onClickHistory = (chat: any) => {
81-
chatRef.value?.onClickHistory(chat)
88+
chatRef.value?.toAssistantHistory(chat)
8289
}
8390
84-
const EditPen = (chat: any) => {
91+
/* const EditPen = (chat: any) => {
8592
chatRef.value?.onChatRenamed(chat)
8693
}
8794
const Delete = (chat: any) => {
8895
chatRef.value?.onChatDeleted(chat.id)
89-
}
96+
} */
9097
const validator = ref({
9198
id: '',
9299
valid: false,
@@ -153,10 +160,15 @@ onBeforeMount(async () => {
153160
.sqlbot-top-btn {
154161
right: 85px;
155162
z-index: 2009;
156-
position: absolute;
163+
position: fixed;
157164
top: 16px;
165+
display: flex;
166+
column-gap: 16px;
158167
// right: 16px;
159168
font-size: 22px;
169+
i:first-child {
170+
color: var(--primary-color);
171+
}
160172
}
161173
.sqlbot-history-container {
162174
width: 100%;

0 commit comments

Comments
 (0)