From 10eb41f77517eb0601b63844aec354c189f56f7c Mon Sep 17 00:00:00 2001 From: fit2cloud-chenyw Date: Mon, 21 Jul 2025 17:31:03 +0800 Subject: [PATCH] fix: Assistant ds --- backend/apps/chat/task/llm.py | 18 ++++--- backend/apps/system/crud/assistant.py | 31 +++++++++--- backend/apps/system/schemas/system_schema.py | 51 +++++++++++++++++++- 3 files changed, 83 insertions(+), 17 deletions(-) diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 2392058c..d8e23569 100644 --- a/backend/apps/chat/task/llm.py +++ b/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) diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py index b830e104..5aadc8f0 100644 --- a/backend/apps/system/crud/assistant.py +++ b/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: diff --git a/backend/apps/system/schemas/system_schema.py b/backend/apps/system/schemas/system_schema.py index 30592c31..d94bffd8 100644 --- a/backend/apps/system/schemas/system_schema.py +++ b/backend/apps/system/schemas/system_schema.py @@ -101,4 +101,53 @@ class UserWs(BaseCreatorDTO): name: str class UserWsOption(UserWs): - account: str \ No newline at end of file + 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 \ No newline at end of file