perf: Assistant dynamic ds id

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,

backend/apps/chat/curd/chat.py

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

backend/apps/chat/task/llm.py

@@ -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/system/crud/assistant.py

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

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