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

Author: dataeaseShu

backend/alembic/versions/033_chat_origin_ddl.py

@@ -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

@@ -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

@@ -8,7 +8,8 @@
 
 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

@@ -2,7 +2,7 @@
 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

@@ -31,7 +31,8 @@
 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 @@ def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question
         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 @@ def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question
                                                                                  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 @@ def generate_recommend_questions_task(self):
 
         # 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

@@ -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

@@ -15,6 +15,7 @@
 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 @@ def get_ds_from_api(self):
         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 @@ def get_ds_from_api(self):
             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 @@ def get_ds(self, ds_id: int):
             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

@@ -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,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

@@ -9,6 +9,7 @@ export abstract class BaseG2Chart extends BaseChart {
     this.chart = new Chart({
       container: id,
       autoFit: true,
+      padding: 'auto',
     })
 
     this.chart.theme({