diff --git a/backend/alembic/versions/031_modify_chat_record.py b/backend/alembic/versions/031_modify_chat_record.py new file mode 100644 index 00000000..aa391e26 --- /dev/null +++ b/backend/alembic/versions/031_modify_chat_record.py @@ -0,0 +1,33 @@ +"""031_modify_chat_record + +Revision ID: bd2ed188b5bd +Revises: c1d7ac00b3a8 +Create Date: 2025-07-21 17:27:55.985821 + +""" +from alembic import op +import sqlalchemy as sa +import sqlmodel.sql.sqltypes +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = 'bd2ed188b5bd' +down_revision = 'c1d7ac00b3a8' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column('chat_record', 'engine_type', + existing_type=sa.VARCHAR(length=64), + nullable=True) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column('chat_record', 'engine_type', + existing_type=sa.VARCHAR(length=64), + nullable=False) + # ### end Alembic commands ### diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index 199fae5a..a83c6b07 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -36,7 +36,7 @@ class ChatRecord(SQLModel, table=True): finish_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True)) create_by: int = Field(sa_column=Column(BigInteger, nullable=True)) datasource: int = Field(sa_column=Column(BigInteger, nullable=True)) - engine_type: str = Field(max_length=64) + engine_type: str = Field(max_length=64, nullable=True) question: str = Field(sa_column=Column(Text, nullable=True)) sql_answer: str = Field(sa_column=Column(Text, nullable=True)) sql: str = Field(sa_column=Column(Text, nullable=True)) 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/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index d60a173e..0168c420 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -230,19 +230,26 @@ def updateField(session: SessionDep, field: CoreField): def preview(session: SessionDep, id: int, data: TableObj): + if data.fields is None or len(data.fields) == 0: + return {"fields": [], "data": [], "sql": ''} + + fields = [f.field_name for f in data.fields if f.checked] + if fields is None or len(fields) == 0: + return {"fields": [], "data": [], "sql": ''} + ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first() conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config() sql: str = "" if ds.type == "mysql": - sql = f"""SELECT `{"`, `".join([f.field_name for f in data.fields if f.checked])}` FROM `{data.table.table_name}` LIMIT 100""" + sql = f"""SELECT `{"`, `".join(fields)}` FROM `{data.table.table_name}` LIMIT 100""" elif ds.type == "sqlServer": - sql = f"""SELECT [{"], [".join([f.field_name for f in data.fields if f.checked])}] FROM [{conf.dbSchema}].[{data.table.table_name}] + sql = f"""SELECT [{"], [".join(fields)}] FROM [{conf.dbSchema}].[{data.table.table_name}] ORDER BY [{data.fields[0].field_name}] OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY""" elif ds.type == "pg" or ds.type == "excel": - sql = f"""SELECT "{'", "'.join([f.field_name for f in data.fields if f.checked])}" FROM "{conf.dbSchema}"."{data.table.table_name}" LIMIT 100""" + sql = f"""SELECT "{'", "'.join(fields)}" FROM "{conf.dbSchema}"."{data.table.table_name}" LIMIT 100""" elif ds.type == "oracle": - sql = f"""SELECT "{'", "'.join([f.field_name for f in data.fields if f.checked])}" FROM "{conf.dbSchema}"."{data.table.table_name}" + sql = f"""SELECT "{'", "'.join(fields)}" FROM "{conf.dbSchema}"."{data.table.table_name}" ORDER BY "{data.fields[0].field_name}" OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY""" return exec_sql(ds, sql) 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 diff --git a/frontend/src/views/chat/RecommendQuestion.vue b/frontend/src/views/chat/RecommendQuestion.vue index b539e82b..6cf0a025 100644 --- a/frontend/src/views/chat/RecommendQuestion.vue +++ b/frontend/src/views/chat/RecommendQuestion.vue @@ -1,18 +1,37 @@