fix: Assistant ds

backend/apps/chat/task/llm.py

@@ -403,22 +403,24 @@ class LLMService:
 
             if data['id'] and data['id'] != 0:
                 _datasource = data['id']
-                if self.current_assistant.type == 1:
+                _chat = self.session.get(Chat, self.record.chat_id)
+                _chat.datasource = _datasource
+                if self.current_assistant and self.current_assistant.type == 1:
                     _ds = self.out_ds_instance.get_ds(data['id'])
                     self.ds = _ds
+                    self.chat_question.engine = _ds.type
+                    _engine_type = self.chat_question.engine
+                    _chat.engine_type =  _ds.type
                 else:
                     _ds = self.session.get(CoreDatasource, _datasource)
                     if not _ds:
                         _datasource = None
                         raise Exception(f"Datasource configuration with id {_datasource} not found")
                     self.ds = CoreDatasource(**_ds.model_dump())
-                self.chat_question.engine = _ds.type_name if _ds.type != 'excel' else 'PostgreSQL'
-                _engine_type = self.chat_question.engine
+                    self.chat_question.engine = _ds.type_name if _ds.type != 'excel' else 'PostgreSQL'
+                    _engine_type = self.chat_question.engine
+                    _chat.engine_type = _ds.type_name
                 # save chat
-                _chat = self.session.get(Chat, self.record.chat_id)
-                _chat.datasource = _datasource
-                _chat.engine_type = _ds.type_name
-
                 self.session.add(_chat)
                 self.session.flush()
                 self.session.refresh(_chat)
@@ -734,7 +736,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
                          'type': 'datasource-result'}).decode() + '\n\n'
             if in_chat:
                 yield orjson.dumps({'id': llm_service.ds.id, 'datasource_name': llm_service.ds.name,
-                                    'engine_type': llm_service.ds.type_name, 'type': 'datasource'}).decode() + '\n\n'
+                                    'engine_type': llm_service.ds.type_name or llm_service.ds.type, 'type': 'datasource'}).decode() + '\n\n'
 
             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)
 

backend/apps/system/crud/assistant.py

@@ -8,7 +8,7 @@ from sqlmodel import Session, select
 from apps.datasource.models.datasource import CoreDatasource
 from apps.system.models.system_model import AssistantModel
 from apps.system.schemas.auth import CacheName, CacheNamespace
-from apps.system.schemas.system_schema import UserInfoDTO
+from apps.system.schemas.system_schema import AssistantOutDsSchema, UserInfoDTO
 from common.core.sqlbot_cache import cache
 from common.core.db import engine
 from starlette.middleware.cors import CORSMiddleware
@@ -73,7 +73,7 @@ def init_dynamic_cors(app: FastAPI):
 
 class AssistantOutDs:
     assistant: AssistantModel
-    ds_list: Optional[list[dict]] = None
+    ds_list: Optional[list[AssistantOutDsSchema]] = None
     certificate: Optional[str] = None
     def __init__(self, assistant: AssistantModel, certificate: Optional[str] = None):
         self.assistant = assistant
@@ -99,9 +99,11 @@ class AssistantOutDs:
             result_json: dict[any] = json.loads(res.json())
             if result_json.get('code') == 0:
                 temp_list = result_json.get('data', [])
-                for idx, item in enumerate(temp_list, start=1):
-                    item["id"] = idx
-                self.ds_list = temp_list
+                self.ds_list = [
+                    AssistantOutDsSchema(**{**item, "id": idx})
+                    for idx, item in enumerate(temp_list, start=1)
+                ]
+                
                 return self.ds_list
             else:
                 raise Exception(f"Failed to get datasource list from {endpoint}, error: {result_json.get('message')}")
@@ -110,12 +112,25 @@ class AssistantOutDs:
     
     def get_simple_ds_list(self):
         if self.ds_list:
-            return [{'id': ds['id'], 'name': ds['name'], 'description': ds['comment']} for ds in self.ds_list]
+            return [{'id': ds.id, 'name': ds.name, 'description': ds.comment} for ds in self.ds_list]
         else:
             raise Exception("Datasource list is not found.")
        
-    def get_db_schema(self, ds_id: int):
-        return None
+    def get_db_schema(self, ds_id: int) -> str:
+        ds = self.get_ds(ds_id)
+        schema_str = ""
+        db_name = ds.schema
+        schema_str += f"【DB_ID】 {db_name}\n【Schema】\n"
+        for table in ds.tables:
+            schema_str += f"# Table: {db_name}.{table.name}"
+            schema_str += f", {table.comment}\n[\n"
+            field_list = []
+            for field in table.fields:
+                field_list.append(f"({field.name}:{field.type}, {field.comment})")
+            schema_str += ",\n".join(field_list)
+            schema_str += '\n]\n'
+        return schema_str
+    
     def get_ds(self, ds_id: int):
         if self.ds_list:
             for ds in self.ds_list:

backend/apps/system/schemas/system_schema.py

@@ -101,4 +101,53 @@ class UserWs(BaseCreatorDTO):
     name: str
     
 class UserWsOption(UserWs):
-    account: str
+    account: str
+    
+    
+class AssistantFieldSchema(BaseModel):
+    id: Optional[int] = None
+    name: Optional[str] = None
+    type: Optional[str] = None
+    comment: Optional[str] = None
+class AssistantTableSchema(BaseModel):
+    id: Optional[int] = None
+    name: Optional[str] = None
+    comment: Optional[str] = None
+    rule: Optional[str] = None
+    sql: Optional[str] = None
+    fields: Optional[list[AssistantFieldSchema]] = None
+
+class AssistantOutDsBase(BaseModel):
+    id: Optional[int] = None
+    name: str
+    type: Optional[str] = None
+    type_name: Optional[str] = None
+    comment: Optional[str] = None
+    description: Optional[str] = None
+    
+    def __init__(self, id: Optional[int] = None, name: str = '', type: Optional[str] = None,
+                 type_name: Optional[str] = None, comment: Optional[str] = None):
+        super().__init__(id=id, name=name, type=type, type_name=type_name, comment=comment)
+        
+class AssistantOutDsSchema(AssistantOutDsBase):
+    host: Optional[str] = None
+    port: Optional[int] = None
+    user: Optional[str] = None
+    password: Optional[str] = None
+    schema: Optional[str] = None
+    tables: Optional[list[AssistantTableSchema]] = None
+
+    def __init__(self, id: int, name: str, comment: Optional[str] = None, type: Optional[str] = None,
+                 type_name: Optional[str] = None, host: Optional[str] = None, port: Optional[int] = None,
+                 username: Optional[str] = None, password: Optional[str] = None, schema: Optional[str] = None):
+        self.id = id
+        self.name = name
+        self.comment = comment
+        self.type = type
+        self.type_name = type_name
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+        self.schema = schema
+        self.description = comment