klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 01:06:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' of https://github.com/dataease/SQLBot

Browse Source
This commit is contained in:
dataeaseShu
2025-07-23 09:28:33 +08:00
parent 2bdbd2b8fd aeee30e5a9
commit 0069f767e2
14 changed files with 213 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/alembic/versions/033_chat_origin_ddl.py
+30
View File
@@ -0,0 +1,30 @@
"""033_chat_origin_ddl
Revision ID: 3cb5d6a54f2e
Revises: 6549e47f9adc
Create Date: 2025-07-22 22:00:48.599729
"""
from alembic import op
import sqlalchemy as sa
import sqlmodel.sql.sqltypes
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision = '3cb5d6a54f2e'
down_revision = '6549e47f9adc'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.add_column('chat', sa.Column('origin', sa.Integer(), nullable=True, default=0)) # 0: default, 1: mcp, 2: assistant
op.execute('UPDATE chat SET origin = 0 WHERE origin IS NULL')
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.drop_column('chat', 'origin')
# ### end Alembic commands ###
backend/apps/chat/api/chat.py
+3 -3
View File
@@ -18,9 +18,9 @@ async def chats(session: SessionDep, current_user: CurrentUser):
@router.get("/get/{chart_id}")
async def get_chat(session: SessionDep, current_user: CurrentUser, chart_id: int):
async def get_chat(session: SessionDep, current_user: CurrentUser, chart_id: int, current_assistant: CurrentAssistant):
try:
return get_chat_with_records(chart_id=chart_id, session=session, current_user=current_user)
return get_chat_with_records(chart_id=chart_id, session=session, current_user=current_user, current_assistant=current_assistant)
except Exception as e:
raise HTTPException(
status_code=500,
@@ -85,7 +85,7 @@ async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat
@router.post("/assistant/start")
async def start_chat(session: SessionDep, current_user: CurrentUser):
try:
return create_chat(session, current_user, CreateChat(), False)
return create_chat(session, current_user, CreateChat(origin=2), False)
except Exception as e:
raise HTTPException(
status_code=500,
backend/apps/chat/curd/chat.py
+14 -7
View File
@@ -8,7 +8,8 @@ from sqlalchemy.orm import load_only
from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat, ChatQuestion
from apps.datasource.models.datasource import CoreDatasource
from common.core.deps import SessionDep, CurrentUser
from apps.system.crud.assistant import AssistantOutDsFactory
from common.core.deps import CurrentAssistant, SessionDep, CurrentUser
from common.utils.utils import extract_nested_json
@@ -20,7 +21,7 @@ def list_chats(session: SessionDep, current_user: CurrentUser) -> List[Chat]:
def rename_chat(session: SessionDep, rename_object: RenameChat) -> str:
chat = session.query(Chat).filter(Chat.id == rename_object.id).first()
chat = session.get(Chat, rename_object.id)
if not chat:
raise Exception(f"Chat with id {rename_object.id} not found")
@@ -65,14 +66,19 @@ def get_chat_predict_data(session: SessionDep, chart_record_id: int):
return ''
def get_chat_with_records(session: SessionDep, chart_id: int, current_user: CurrentUser) -> ChatInfo:
chat = session.query(Chat).filter(Chat.id == chart_id).first()
def get_chat_with_records(session: SessionDep, chart_id: int, current_user: CurrentUser, current_assistant: CurrentAssistant) -> ChatInfo:
chat = session.get(Chat, chart_id)
if not chat:
raise Exception(f"Chat with id {chart_id} not found")
chat_info = ChatInfo(**chat.model_dump())
ds = session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first()
if current_assistant and current_assistant.type == 1:
out_ds_instance = AssistantOutDsFactory.get_instance(current_assistant)
ds = out_ds_instance.get_ds(chat.datasource)
else:
ds = session.get(CoreDatasource, chat.datasource) if chat.datasource else None
if not ds:
chat_info.datasource_exists = False
chat_info.datasource_name = 'Datasource not exist'
@@ -157,11 +163,12 @@ def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj:
chat = Chat(create_time=datetime.datetime.now(),
create_by=current_user.id,
oid=current_user.oid if current_user.oid is not None else 1,
brief=create_chat_obj.question.strip()[:20])
brief=create_chat_obj.question.strip()[:20],
origin=create_chat_obj.origin if create_chat_obj.origin is not None else 0)
ds: CoreDatasource | None = None
if create_chat_obj.datasource:
chat.datasource = create_chat_obj.datasource
ds = session.query(CoreDatasource).filter(CoreDatasource.id == create_chat_obj.datasource).first()
ds = session.get(CoreDatasource, create_chat_obj.datasource)
if not ds:
raise Exception(f"Datasource with id {create_chat_obj.datasource} not found")
backend/apps/chat/models/chat_model.py
+3 -1
View File
@@ -2,7 +2,7 @@ from datetime import datetime
from typing import List, Optional
from pydantic import BaseModel
from sqlalchemy import Column, Text, BigInteger, DateTime, Identity, Boolean
from sqlalchemy import Column, Integer, Text, BigInteger, DateTime, Identity, Boolean
from sqlmodel import SQLModel, Field
from apps.template.generate_analysis.generator import get_analysis_template
@@ -24,6 +24,7 @@ class Chat(SQLModel, table=True):
chat_type: str = Field(max_length=20, default="chat") # chat, datasource
datasource: int = Field(sa_column=Column(BigInteger, nullable=True))
engine_type: str = Field(max_length=64)
origin: Optional[int] = Field(sa_column=Column(Integer, nullable=False, default=0)) # 0: default, 1: mcp, 2: assistant
class ChatRecord(SQLModel, table=True):
@@ -72,6 +73,7 @@ class CreateChat(BaseModel):
id: int = None
question: str = None
datasource: int = None
origin: Optional[int] = 0
class RenameChat(BaseModel):
backend/apps/chat/task/llm.py
+18 -13
View File
@@ -31,7 +31,8 @@ from apps.datasource.crud.datasource import get_table_schema
from apps.datasource.crud.row_permission import transFilterTree
from apps.datasource.models.datasource import CoreDatasource, CoreTable
from apps.db.db import exec_sql
from apps.system.crud.assistant import AssistantOutDs, get_assistant_ds
from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, get_assistant_ds
from apps.system.schemas.system_schema import AssistantOutDsSchema
from common.core.config import settings
from common.core.deps import CurrentAssistant, SessionDep, CurrentUser
from common.utils.utils import extract_nested_json
@@ -67,15 +68,23 @@ class LLMService:
chat: Chat = self.session.get(Chat, chat_id)
if not chat:
raise Exception(f"Chat with id {chat_id} not found")
ds: CoreDatasource | None = None
ds: CoreDatasource | AssistantOutDsSchema | None = None
if chat.datasource:
# Get available datasource
# ds = self.session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first()
ds = self.session.get(CoreDatasource, chat.datasource)
if not ds:
raise Exception("No available datasource configuration found")
chat_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL'
if current_assistant and current_assistant.type == 1:
self.out_ds_instance = AssistantOutDsFactory.get_instance(current_assistant)
ds = self.out_ds_instance.get_ds(chat.datasource)
if not ds:
raise Exception("No available datasource configuration found")
chat_question.engine = ds.type
chat_question.db_schema = self.out_ds_instance.get_db_schema(ds.id)
else:
ds = self.session.get(CoreDatasource, chat.datasource)
if not ds:
raise Exception("No available datasource configuration found")
chat_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL'
chat_question.db_schema = get_table_schema(session=self.session, current_user=current_user, ds=ds)
history_records: List[ChatRecord] = list(
map(lambda x: ChatRecord(**x.model_dump()), filter(lambda r: True if r.first_chat != True else False,
@@ -84,13 +93,9 @@ class LLMService:
chart_id=chat_id))))
self.change_title = len(history_records) == 0
# get schema
if ds:
chat_question.db_schema = get_table_schema(session=self.session, current_user=current_user, ds=ds)
chat_question.lang = current_user.language
self.ds = CoreDatasource(**ds.model_dump()) if ds else None
self.ds = (ds if isinstance(ds, AssistantOutDsSchema) else CoreDatasource(**ds.model_dump())) if ds else None
self.chat_question = chat_question
self.config = get_default_config()
self.chat_question.ai_modal_id = self.config.model_id
@@ -298,7 +303,7 @@ class LLMService:
# get schema
if self.ds and not self.chat_question.db_schema:
self.chat_question.db_schema = get_table_schema(session=self.session, current_user=self.current_user, ds=self.ds)
self.chat_question.db_schema = self.out_ds_instance.get_db_schema(self.ds.id) if self.out_ds_instance else get_table_schema(session=self.session, current_user=self.current_user, ds=self.ds)
guess_msg: List[Union[BaseMessage, dict[str, Any]]] = []
guess_msg.append(SystemMessage(content=self.chat_question.guess_sys_question()))
backend/apps/mcp/mcp.py
+1 -1
View File
@@ -71,7 +71,7 @@ async def mcp_start(session: SessionDep, chat: ChatStart):
t = Token(access_token=create_access_token(
user_dict, expires_delta=access_token_expires
))
c = create_chat(session, user, CreateChat(), False)
c = create_chat(session, user, CreateChat(origin=1), False)
return {"access_token": t.access_token, "chat_id": c.id}
backend/apps/system/crud/assistant.py
+12 -3
View File
@@ -15,6 +15,7 @@ from common.core.sqlbot_cache import cache
from common.core.db import engine
from starlette.middleware.cors import CORSMiddleware
from common.core.config import settings
from common.utils.utils import string_to_numeric_hash
@cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="assistant_id")
async def get_assistant_info(*, session: Session, assistant_id: int) -> AssistantModel | None:
@@ -100,7 +101,7 @@ class AssistantOutDs:
header = {}
cookies = {}
for item in certificateList:
if item['target'] == 'head':
if item['target'] == 'header':
header[item['key']] = item['value']
if item['target'] == 'cookie':
cookies[item['key']] = item['value']
@@ -111,8 +112,8 @@ class AssistantOutDs:
if result_json.get('code') == 0:
temp_list = result_json.get('data', [])
self.ds_list = [
AssistantOutDsSchema(**{**item, "id": idx})
for idx, item in enumerate(temp_list, start=1)
self.convert2schema(item)
for item in temp_list
]
return self.ds_list
@@ -151,6 +152,14 @@ class AssistantOutDs:
raise Exception("Datasource list is not found.")
raise Exception(f"Datasource with id {ds_id} not found.")
def convert2schema(self, ds_dict: dict) -> AssistantOutDsSchema:
id_marker: str = ''
attr_list = ['name', 'type', 'host', 'port', 'user', 'dataBase', 'schema']
for attr in attr_list:
if attr in ds_dict:
id_marker += str(ds_dict.get(attr, '')) + '--sqlbot--'
id = string_to_numeric_hash(id_marker)
return AssistantOutDsSchema(**{**ds_dict, "id": id})
class AssistantOutDsFactory:
@staticmethod
backend/common/utils/utils.py
+6
View File
@@ -1,5 +1,7 @@
import hashlib
import logging
from datetime import datetime, timedelta, timezone
from typing import Optional
import jwt
import orjson
@@ -70,3 +72,7 @@ def extract_nested_json(text):
if len(results) > 0 and results[0]:
return results[0]
return None
def string_to_numeric_hash(text: str, bits: Optional[int] = 64) -> int:
hash_bytes = hashlib.sha256(text.encode()).digest()
return int.from_bytes(hash_bytes[:bits//8], byteorder='big')
frontend/src/views/chat/RecommendQuestion.vue
+10 -1
View File
@@ -10,12 +10,14 @@ const props = withDefaults(
currentChat?: ChatInfo
questions?: string
firstChat?: boolean
disabled?: boolean
}>(),
{
recordId: undefined,
currentChat: () => new ChatInfo(),
questions: '[]',
firstChat: false,
disabled: false,
}
)
@@ -47,7 +49,9 @@ const computedQuestions = computed<string>(() => {
const { t } = useI18n()
function clickQuestion(question: string): void {
emits('clickQuestion', question)
if (!props.disabled) {
emits('clickQuestion', question)
}
}
const stopFlag = ref(false)
@@ -160,6 +164,7 @@ defineExpose({ getRecommendQuestions, id: () => props.recordId, stop })
v-for="(question, index) in computedQuestions"
:key="index"
class="question"
:class="{ disabled: disabled }"
@click="clickQuestion(question)"
>
{{ question }}
@@ -199,6 +204,10 @@ defineExpose({ getRecommendQuestions, id: () => props.recordId, stop })
&:hover {
background: rgba(31, 35, 41, 0.1);
}
&.disabled {
cursor: not-allowed;
background: rgba(245, 246, 247, 1);
}
}
}
</style>
frontend/src/views/chat/component/BaseG2Chart.ts
+1
View File
@@ -9,6 +9,7 @@ export abstract class BaseG2Chart extends BaseChart {
this.chart = new Chart({
container: id,
autoFit: true,
padding: 'auto',
})
this.chart.theme({
frontend/src/views/chat/component/charts/Bar.ts
+52 -4
View File
@@ -31,9 +31,56 @@ export class Bar extends BaseG2Chart {
y: y[0].value,
color: series.length > 0 ? series[0].value : undefined,
},
style: {
radiusTopLeft: (d: ChartData) => {
if (d[y[0].value] && d[y[0].value] > 0) {
return 4
}
return 0
},
radiusTopRight: (d: ChartData) => {
if (d[y[0].value] && d[y[0].value] > 0) {
return 4
}
return 0
},
radiusBottomLeft: (d: ChartData) => {
if (d[y[0].value] && d[y[0].value] < 0) {
return 4
}
return 0
},
radiusBottomRight: (d: ChartData) => {
if (d[y[0].value] && d[y[0].value] < 0) {
return 4
}
return 0
},
},
axis: {
x: { title: x[0].name },
y: { title: y[0].name },
x: {
title: x[0].name,
labelFontSize: 12,
labelAutoHide: {
type: 'hide',
keepHeader: true,
keepTail: true,
},
labelAutoRotate: false,
labelAutoWrap: true,
labelAutoEllipsis: true,
},
y: {
title: y[0].name,
labelFontSize: 12,
labelAutoHide: {
type: 'hide',
keepHeader: true,
},
labelAutoRotate: false,
labelAutoWrap: true,
labelAutoEllipsis: true,
},
},
scale: {
x: {
@@ -44,7 +91,8 @@ export class Bar extends BaseG2Chart {
},
},
interaction: {
elementHighlight: { background: true },
elementHighlight: { background: true, region: true },
tooltip: { series: series.length > 0, shared: true },
},
tooltip: (data) => {
if (series.length > 0) {
@@ -60,8 +108,8 @@ export class Bar extends BaseG2Chart {
{
text: (data: any) => `${data[y[0].value]}${_data.isPercent ? '%' : ''}`,
transform: [
{ type: 'overlapDodgeY' },
{ type: 'contrastReverse' },
{ type: 'exceedAdjust' },
{ type: 'overlapHide' },
],
},
frontend/src/views/chat/component/charts/Column.ts
+41 -4
View File
@@ -30,8 +30,45 @@ export class Column extends BaseG2Chart {
y: y[0].value,
color: series.length > 0 ? series[0].value : undefined,
},
style: {
radiusTopLeft: (d: ChartData) => {
if (d[y[0].value] && d[y[0].value] > 0) {
return 4
}
return 0
},
radiusTopRight: (d: ChartData) => {
if (d[y[0].value] && d[y[0].value] > 0) {
return 4
}
return 0
},
radiusBottomLeft: (d: ChartData) => {
if (d[y[0].value] && d[y[0].value] < 0) {
return 4
}
return 0
},
radiusBottomRight: (d: ChartData) => {
if (d[y[0].value] && d[y[0].value] < 0) {
return 4
}
return 0
},
},
axis: {
x: { title: x[0].name },
x: {
title: x[0].name,
labelFontSize: 12,
labelAutoHide: {
type: 'hide',
keepHeader: true,
keepTail: true,
},
labelAutoRotate: false,
labelAutoWrap: true,
labelAutoEllipsis: true,
},
y: { title: y[0].name },
},
scale: {
@@ -43,7 +80,8 @@ export class Column extends BaseG2Chart {
},
},
interaction: {
elementHighlight: { background: true },
elementHighlight: { background: true, region: true },
tooltip: { series: series.length > 0, shared: true },
},
tooltip: (data) => {
if (series.length > 0) {
@@ -59,10 +97,9 @@ export class Column extends BaseG2Chart {
{
text: (data: any) => `${data[y[0].value]}${_data.isPercent ? '%' : ''}`,
position: 'top',
dy: -25,
transform: [
{ type: 'overlapDodgeY' },
{ type: 'contrastReverse' },
{ type: 'exceedAdjust' },
{ type: 'overlapHide' },
],
},
frontend/src/views/chat/component/charts/Line.ts
+20 -2
View File
@@ -31,7 +31,18 @@ export class Line extends BaseG2Chart {
color: series.length > 0 ? series[0].value : undefined,
},
axis: {
x: { title: x[0].name },
x: {
title: x[0].name,
labelFontSize: 12,
labelAutoHide: {
type: 'hide',
keepHeader: true,
keepTail: true,
},
labelAutoRotate: false,
labelAutoWrap: true,
labelAutoEllipsis: true,
},
y: { title: y[0].name },
},
scale: {
@@ -42,9 +53,16 @@ export class Line extends BaseG2Chart {
nice: true,
},
},
interaction: {
tooltip: { series: series.length > 0, shared: true },
},
children: [
{
type: 'line',
encode: {
shape: 'smooth',
},
labels: [
{
text: (data: any) => `${data[y[0].value]}${_data.isPercent ? '%' : ''}`,
@@ -53,7 +71,7 @@ export class Line extends BaseG2Chart {
dy: -12,
},
transform: [
{ type: 'overlapDodgeY' },
{ type: 'contrastReverse' },
{ type: 'exceedAdjust' },
{ type: 'overlapHide' },
],
frontend/src/views/chat/index.vue
+2
View File
@@ -97,6 +97,7 @@
:current-chat="currentChat"
:record-id="message.record?.id"
:questions="message.recommended_question"
:disabled="isTyping"
:first-chat="message.first_chat"
@click-question="quickAsk"
@stop="onChatStop"
@@ -191,6 +192,7 @@
:record-id="message.record?.id"
:questions="message.recommended_question"
:first-chat="message.first_chat"
:disabled="isTyping"
@click-question="quickAsk"
@stop="onChatStop"
/>