Skip to content

Commit cb997d5

Browse files
committed
Merge remote-tracking branch 'origin/main'
2 parents 35e0a36 + 7c2b064 commit cb997d5

12 files changed

Lines changed: 262 additions & 99 deletions

File tree

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
"""031_modify_chat_record
2+
3+
Revision ID: bd2ed188b5bd
4+
Revises: c1d7ac00b3a8
5+
Create Date: 2025-07-21 17:27:55.985821
6+
7+
"""
8+
from alembic import op
9+
import sqlalchemy as sa
10+
import sqlmodel.sql.sqltypes
11+
from sqlalchemy.dialects import postgresql
12+
13+
# revision identifiers, used by Alembic.
14+
revision = 'bd2ed188b5bd'
15+
down_revision = 'c1d7ac00b3a8'
16+
branch_labels = None
17+
depends_on = None
18+
19+
20+
def upgrade():
21+
# ### commands auto generated by Alembic - please adjust! ###
22+
op.alter_column('chat_record', 'engine_type',
23+
existing_type=sa.VARCHAR(length=64),
24+
nullable=True)
25+
# ### end Alembic commands ###
26+
27+
28+
def downgrade():
29+
# ### commands auto generated by Alembic - please adjust! ###
30+
op.alter_column('chat_record', 'engine_type',
31+
existing_type=sa.VARCHAR(length=64),
32+
nullable=False)
33+
# ### end Alembic commands ###

backend/apps/chat/models/chat_model.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ class ChatRecord(SQLModel, table=True):
3636
finish_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True))
3737
create_by: int = Field(sa_column=Column(BigInteger, nullable=True))
3838
datasource: int = Field(sa_column=Column(BigInteger, nullable=True))
39-
engine_type: str = Field(max_length=64)
39+
engine_type: str = Field(max_length=64, nullable=True)
4040
question: str = Field(sa_column=Column(Text, nullable=True))
4141
sql_answer: str = Field(sa_column=Column(Text, nullable=True))
4242
sql: str = Field(sa_column=Column(Text, nullable=True))

backend/apps/chat/task/llm.py

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -403,22 +403,24 @@ def select_datasource(self):
403403

404404
if data['id'] and data['id'] != 0:
405405
_datasource = data['id']
406-
if self.current_assistant.type == 1:
406+
_chat = self.session.get(Chat, self.record.chat_id)
407+
_chat.datasource = _datasource
408+
if self.current_assistant and self.current_assistant.type == 1:
407409
_ds = self.out_ds_instance.get_ds(data['id'])
408410
self.ds = _ds
411+
self.chat_question.engine = _ds.type
412+
_engine_type = self.chat_question.engine
413+
_chat.engine_type = _ds.type
409414
else:
410415
_ds = self.session.get(CoreDatasource, _datasource)
411416
if not _ds:
412417
_datasource = None
413418
raise Exception(f"Datasource configuration with id {_datasource} not found")
414419
self.ds = CoreDatasource(**_ds.model_dump())
415-
self.chat_question.engine = _ds.type_name if _ds.type != 'excel' else 'PostgreSQL'
416-
_engine_type = self.chat_question.engine
420+
self.chat_question.engine = _ds.type_name if _ds.type != 'excel' else 'PostgreSQL'
421+
_engine_type = self.chat_question.engine
422+
_chat.engine_type = _ds.type_name
417423
# save chat
418-
_chat = self.session.get(Chat, self.record.chat_id)
419-
_chat.datasource = _datasource
420-
_chat.engine_type = _ds.type_name
421-
422424
self.session.add(_chat)
423425
self.session.flush()
424426
self.session.refresh(_chat)
@@ -734,7 +736,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
734736
'type': 'datasource-result'}).decode() + '\n\n'
735737
if in_chat:
736738
yield orjson.dumps({'id': llm_service.ds.id, 'datasource_name': llm_service.ds.name,
737-
'engine_type': llm_service.ds.type_name, 'type': 'datasource'}).decode() + '\n\n'
739+
'engine_type': llm_service.ds.type_name or llm_service.ds.type, 'type': 'datasource'}).decode() + '\n\n'
738740

739741
llm_service.chat_question.db_schema = llm_service.out_ds_instance.get_db_schema() if llm_service.out_ds_instance else get_table_schema(session=llm_service.session, current_user=llm_service.current_user, ds=llm_service.ds)
740742

backend/apps/datasource/crud/datasource.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -230,19 +230,26 @@ def updateField(session: SessionDep, field: CoreField):
230230

231231

232232
def preview(session: SessionDep, id: int, data: TableObj):
233+
if data.fields is None or len(data.fields) == 0:
234+
return {"fields": [], "data": [], "sql": ''}
235+
236+
fields = [f.field_name for f in data.fields if f.checked]
237+
if fields is None or len(fields) == 0:
238+
return {"fields": [], "data": [], "sql": ''}
239+
233240
ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
234241
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
235242
sql: str = ""
236243
if ds.type == "mysql":
237-
sql = f"""SELECT `{"`, `".join([f.field_name for f in data.fields if f.checked])}` FROM `{data.table.table_name}` LIMIT 100"""
244+
sql = f"""SELECT `{"`, `".join(fields)}` FROM `{data.table.table_name}` LIMIT 100"""
238245
elif ds.type == "sqlServer":
239-
sql = f"""SELECT [{"], [".join([f.field_name for f in data.fields if f.checked])}] FROM [{conf.dbSchema}].[{data.table.table_name}]
246+
sql = f"""SELECT [{"], [".join(fields)}] FROM [{conf.dbSchema}].[{data.table.table_name}]
240247
ORDER BY [{data.fields[0].field_name}]
241248
OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY"""
242249
elif ds.type == "pg" or ds.type == "excel":
243-
sql = f"""SELECT "{'", "'.join([f.field_name for f in data.fields if f.checked])}" FROM "{conf.dbSchema}"."{data.table.table_name}" LIMIT 100"""
250+
sql = f"""SELECT "{'", "'.join(fields)}" FROM "{conf.dbSchema}"."{data.table.table_name}" LIMIT 100"""
244251
elif ds.type == "oracle":
245-
sql = f"""SELECT "{'", "'.join([f.field_name for f in data.fields if f.checked])}" FROM "{conf.dbSchema}"."{data.table.table_name}"
252+
sql = f"""SELECT "{'", "'.join(fields)}" FROM "{conf.dbSchema}"."{data.table.table_name}"
246253
ORDER BY "{data.fields[0].field_name}"
247254
OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY"""
248255
return exec_sql(ds, sql)

backend/apps/system/crud/assistant.py

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from apps.datasource.models.datasource import CoreDatasource
99
from apps.system.models.system_model import AssistantModel
1010
from apps.system.schemas.auth import CacheName, CacheNamespace
11-
from apps.system.schemas.system_schema import UserInfoDTO
11+
from apps.system.schemas.system_schema import AssistantOutDsSchema, UserInfoDTO
1212
from common.core.sqlbot_cache import cache
1313
from common.core.db import engine
1414
from starlette.middleware.cors import CORSMiddleware
@@ -73,7 +73,7 @@ def init_dynamic_cors(app: FastAPI):
7373

7474
class AssistantOutDs:
7575
assistant: AssistantModel
76-
ds_list: Optional[list[dict]] = None
76+
ds_list: Optional[list[AssistantOutDsSchema]] = None
7777
certificate: Optional[str] = None
7878
def __init__(self, assistant: AssistantModel, certificate: Optional[str] = None):
7979
self.assistant = assistant
@@ -99,9 +99,11 @@ async def get_ds_from_api(self, certificate: Optional[str] = None):
9999
result_json: dict[any] = json.loads(res.json())
100100
if result_json.get('code') == 0:
101101
temp_list = result_json.get('data', [])
102-
for idx, item in enumerate(temp_list, start=1):
103-
item["id"] = idx
104-
self.ds_list = temp_list
102+
self.ds_list = [
103+
AssistantOutDsSchema(**{**item, "id": idx})
104+
for idx, item in enumerate(temp_list, start=1)
105+
]
106+
105107
return self.ds_list
106108
else:
107109
raise Exception(f"Failed to get datasource list from {endpoint}, error: {result_json.get('message')}")
@@ -110,12 +112,25 @@ async def get_ds_from_api(self, certificate: Optional[str] = None):
110112

111113
def get_simple_ds_list(self):
112114
if self.ds_list:
113-
return [{'id': ds['id'], 'name': ds['name'], 'description': ds['comment']} for ds in self.ds_list]
115+
return [{'id': ds.id, 'name': ds.name, 'description': ds.comment} for ds in self.ds_list]
114116
else:
115117
raise Exception("Datasource list is not found.")
116118

117-
def get_db_schema(self, ds_id: int):
118-
return None
119+
def get_db_schema(self, ds_id: int) -> str:
120+
ds = self.get_ds(ds_id)
121+
schema_str = ""
122+
db_name = ds.schema
123+
schema_str += f"【DB_ID】 {db_name}\n【Schema】\n"
124+
for table in ds.tables:
125+
schema_str += f"# Table: {db_name}.{table.name}"
126+
schema_str += f", {table.comment}\n[\n"
127+
field_list = []
128+
for field in table.fields:
129+
field_list.append(f"({field.name}:{field.type}, {field.comment})")
130+
schema_str += ",\n".join(field_list)
131+
schema_str += '\n]\n'
132+
return schema_str
133+
119134
def get_ds(self, ds_id: int):
120135
if self.ds_list:
121136
for ds in self.ds_list:

backend/apps/system/schemas/system_schema.py

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,4 +101,53 @@ class UserWs(BaseCreatorDTO):
101101
name: str
102102

103103
class UserWsOption(UserWs):
104-
account: str
104+
account: str
105+
106+
107+
class AssistantFieldSchema(BaseModel):
108+
id: Optional[int] = None
109+
name: Optional[str] = None
110+
type: Optional[str] = None
111+
comment: Optional[str] = None
112+
class AssistantTableSchema(BaseModel):
113+
id: Optional[int] = None
114+
name: Optional[str] = None
115+
comment: Optional[str] = None
116+
rule: Optional[str] = None
117+
sql: Optional[str] = None
118+
fields: Optional[list[AssistantFieldSchema]] = None
119+
120+
class AssistantOutDsBase(BaseModel):
121+
id: Optional[int] = None
122+
name: str
123+
type: Optional[str] = None
124+
type_name: Optional[str] = None
125+
comment: Optional[str] = None
126+
description: Optional[str] = None
127+
128+
def __init__(self, id: Optional[int] = None, name: str = '', type: Optional[str] = None,
129+
type_name: Optional[str] = None, comment: Optional[str] = None):
130+
super().__init__(id=id, name=name, type=type, type_name=type_name, comment=comment)
131+
132+
class AssistantOutDsSchema(AssistantOutDsBase):
133+
host: Optional[str] = None
134+
port: Optional[int] = None
135+
user: Optional[str] = None
136+
password: Optional[str] = None
137+
schema: Optional[str] = None
138+
tables: Optional[list[AssistantTableSchema]] = None
139+
140+
def __init__(self, id: int, name: str, comment: Optional[str] = None, type: Optional[str] = None,
141+
type_name: Optional[str] = None, host: Optional[str] = None, port: Optional[int] = None,
142+
username: Optional[str] = None, password: Optional[str] = None, schema: Optional[str] = None):
143+
self.id = id
144+
self.name = name
145+
self.comment = comment
146+
self.type = type
147+
self.type_name = type_name
148+
self.host = host
149+
self.port = port
150+
self.username = username
151+
self.password = password
152+
self.schema = schema
153+
self.description = comment

frontend/src/views/chat/RecommendQuestion.vue

Lines changed: 102 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,37 @@
11
<script setup lang="ts">
2-
import { computed } from 'vue'
2+
import { computed, nextTick, ref } from 'vue'
33
import { endsWith, startsWith } from 'lodash-es'
44
import { useI18n } from 'vue-i18n'
5+
import { chatApi, ChatInfo } from '@/api/chat.ts'
56
67
const props = withDefaults(
78
defineProps<{
9+
recordId?: number
10+
currentChat?: ChatInfo
811
questions?: string
912
firstChat?: boolean
1013
}>(),
1114
{
15+
recordId: undefined,
16+
currentChat: () => new ChatInfo(),
1217
questions: '[]',
1318
firstChat: false,
1419
}
1520
)
21+
22+
const emits = defineEmits(['clickQuestion', 'update:currentChat'])
23+
24+
const loading = ref(false)
25+
26+
const _currentChat = computed({
27+
get() {
28+
return props.currentChat
29+
},
30+
set(v) {
31+
emits('update:currentChat', v)
32+
},
33+
})
34+
1635
const computedQuestions = computed<string>(() => {
1736
if (
1837
props.questions &&
@@ -25,20 +44,98 @@ const computedQuestions = computed<string>(() => {
2544
return []
2645
})
2746
28-
const emits = defineEmits(['clickQuestion'])
29-
3047
const { t } = useI18n()
3148
3249
function clickQuestion(question: string): void {
3350
emits('clickQuestion', question)
3451
}
52+
53+
async function getRecommendQuestions() {
54+
loading.value = true
55+
try {
56+
const response = await chatApi.recommendQuestions(props.recordId)
57+
const reader = response.body.getReader()
58+
const decoder = new TextDecoder()
59+
60+
while (true) {
61+
const { done, value } = await reader.read()
62+
if (done) {
63+
break
64+
}
65+
66+
const chunk = decoder.decode(value)
67+
68+
let _list = [chunk]
69+
70+
const lines = chunk.trim().split('}\n\n{')
71+
if (lines.length > 1) {
72+
_list = []
73+
for (let line of lines) {
74+
if (!line.trim().startsWith('{')) {
75+
line = '{' + line.trim()
76+
}
77+
if (!line.trim().endsWith('}')) {
78+
line = line.trim() + '}'
79+
}
80+
_list.push(line)
81+
}
82+
}
83+
84+
for (const str of _list) {
85+
let data
86+
try {
87+
data = JSON.parse(str)
88+
} catch (err) {
89+
console.error('JSON string:', str)
90+
throw err
91+
}
92+
93+
if (data.code && data.code !== 200) {
94+
ElMessage({
95+
message: data.msg,
96+
type: 'error',
97+
showClose: true,
98+
})
99+
return
100+
}
101+
102+
switch (data.type) {
103+
case 'recommended_question':
104+
if (
105+
data.content &&
106+
data.content.length > 0 &&
107+
startsWith(data.content.trim(), '[') &&
108+
endsWith(data.content.trim(), ']')
109+
) {
110+
if (_currentChat.value?.records) {
111+
for (let record of _currentChat.value.records) {
112+
if (record.id === props.recordId) {
113+
record.recommended_question = data.content
114+
115+
await nextTick()
116+
}
117+
}
118+
}
119+
}
120+
}
121+
}
122+
}
123+
} finally {
124+
loading.value = false
125+
}
126+
}
127+
128+
defineExpose({ getRecommendQuestions, id: () => props.recordId })
35129
</script>
36130

37131
<template>
38-
<div v-if="computedQuestions.length > 0" class="recommend-questions">
132+
<div v-if="computedQuestions.length > 0 || loading" class="recommend-questions">
39133
<div v-if="firstChat">{{ t('qa.guess_u_ask') }}</div>
40134
<div v-else class="continue-ask">{{ t('qa.continue_to_ask') }}</div>
41-
<div class="question-grid">
135+
<div v-if="loading">
136+
<el-button style="min-width: unset" type="primary" link loading />
137+
</div>
138+
<div v-else class="question-grid">
42139
<div
43140
v-for="(question, index) in computedQuestions"
44141
:key="index"

frontend/src/views/chat/answer/AnalysisAnswer.vue

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,9 @@ import { computed, nextTick } from 'vue'
55
import MdComponent from '@/views/chat/component/MdComponent.vue'
66
const props = withDefaults(
77
defineProps<{
8-
chatList: Array<ChatInfo>
8+
chatList?: Array<ChatInfo>
99
currentChatId?: number
10-
currentChat: ChatInfo
10+
currentChat?: ChatInfo
1111
message?: ChatMessage
1212
loading?: boolean
1313
}>(),

frontend/src/views/chat/answer/ChartAnswer.vue

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@ import { Chat, chatApi, ChatInfo, type ChatMessage, ChatRecord, questionApi } fr
44
import { computed, nextTick } from 'vue'
55
const props = withDefaults(
66
defineProps<{
7-
chatList: Array<ChatInfo>
7+
chatList?: Array<ChatInfo>
88
currentChatId?: number
9-
currentChat: ChatInfo
9+
currentChat?: ChatInfo
1010
message?: ChatMessage
1111
loading?: boolean
1212
reasoningName: 'sql_answer' | 'chart_answer' | Array<'sql_answer' | 'chart_answer'>

0 commit comments

Comments
 (0)