Merge remote-tracking branch 'origin/main'

Author: ziyujiahao

backend/alembic/versions/016_modify_chat.py

@@ -0,0 +1,71 @@
+"""016_modify_chat
+
+Revision ID: 031148da1d81
+Revises: 02d84523a979
+Create Date: 2025-06-26 17:00:07.054531
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '031148da1d81'
+down_revision = '02d84523a979'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('chat', 'datasource',
+               existing_type=sa.INTEGER(),
+               nullable=True)
+    op.add_column('chat_record', sa.Column('ai_modal_id', sa.Integer(), nullable=True))
+    op.add_column('chat_record', sa.Column('first_chat', sa.Boolean(), nullable=True))
+    op.add_column('chat_record', sa.Column('recommended_question_answer', sa.Text(), nullable=True))
+    op.add_column('chat_record', sa.Column('recommended_question', sa.Text(), nullable=True))
+    op.add_column('chat_record', sa.Column('datasource_select_answer', sa.Text(), nullable=True))
+    op.add_column('chat_record', sa.Column('token_sql', sa.Integer(), nullable=True))
+    op.add_column('chat_record', sa.Column('token_chart', sa.Integer(), nullable=True))
+    op.add_column('chat_record', sa.Column('token_analysis', sa.Integer(), nullable=True))
+    op.add_column('chat_record', sa.Column('token_predict', sa.Integer(), nullable=True))
+    op.add_column('chat_record', sa.Column('full_recommended_question_message', sa.Text(), nullable=True))
+    op.add_column('chat_record', sa.Column('token_recommended_question', sa.Integer(), nullable=True))
+    op.add_column('chat_record', sa.Column('full_select_datasource_message', sa.Text(), nullable=True))
+    op.add_column('chat_record', sa.Column('token_select_datasource_question', sa.Integer(), nullable=True))
+    op.alter_column('chat_record', 'chat_id',
+               existing_type=sa.INTEGER(),
+               nullable=False)
+    op.alter_column('chat_record', 'datasource',
+               existing_type=sa.INTEGER(),
+               nullable=True)
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('chat_record', 'datasource',
+               existing_type=sa.INTEGER(),
+               nullable=False)
+    op.alter_column('chat_record', 'chat_id',
+               existing_type=sa.INTEGER(),
+               nullable=True)
+    op.drop_column('chat_record', 'token_select_datasource_question')
+    op.drop_column('chat_record', 'full_select_datasource_message')
+    op.drop_column('chat_record', 'token_recommended_question')
+    op.drop_column('chat_record', 'full_recommended_question_message')
+    op.drop_column('chat_record', 'token_predict')
+    op.drop_column('chat_record', 'token_analysis')
+    op.drop_column('chat_record', 'token_chart')
+    op.drop_column('chat_record', 'token_sql')
+    op.drop_column('chat_record', 'datasource_select_answer')
+    op.drop_column('chat_record', 'recommended_question')
+    op.drop_column('chat_record', 'recommended_question_answer')
+    op.drop_column('chat_record', 'first_chat')
+    op.drop_column('chat_record', 'ai_modal_id')
+    op.alter_column('chat', 'datasource',
+               existing_type=sa.INTEGER(),
+               nullable=False)
+    # ### end Alembic commands ###

backend/alembic/versions/017_rsa_ddl.py

@@ -0,0 +1,35 @@
+"""017_rsa_ddl
+
+Revision ID: a0ba8268868d
+Revises: 031148da1d81
+Create Date: 2025-06-27 15:05:38.676825
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+
+
+# revision identifiers, used by Alembic.
+revision = 'a0ba8268868d'
+down_revision = '031148da1d81'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_table(
+        'rsa',
+        sa.Column('id', sa.BigInteger(), primary_key=True, nullable=False),
+        sa.Column('private_key', sa.Text(), default="", nullable=False),
+        sa.Column('public_key', sa.Text(), default="", nullable=False),
+        sa.Column('salt', sa.Text(), default="", nullable=False),
+        sa.Column('create_time', sa.BigInteger(), default=0, nullable=False),
+        sa.Column('update_time', sa.BigInteger(), default=0, nullable=False)
+    )
+    op.create_index(op.f('ix_rsa_id'), 'rsa', ['id'], unique=False)
+
+
+def downgrade():
+    op.drop_table('rsa')
+    op.drop_index(op.f('ix_rsa_id'), table_name='rsa')

backend/apps/chat/api/chat.py

@@ -8,7 +8,7 @@
 
 from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat, rename_chat, \
     delete_chat, list_records
-from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, Chat, ChatQuestion
+from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, Chat, ChatQuestion, ChatMcp
 from apps.chat.task.llm import LLMService
 from apps.datasource.crud.datasource import get_table_schema
 from apps.datasource.models.datasource import CoreDatasource
@@ -57,6 +57,18 @@ async def delete(session: SessionDep, chart_id: int):
         )
 
 
+@router.post("/mcp_start", operation_id="mcp_start")
+async def mcp_start(session: SessionDep, chat: ChatMcp):
+    user = await get_current_user(session, chat.token)
+    return create_chat(session, user, CreateChat(), False)
+
+
+@router.post("/mcp_question", operation_id="mcp_question")
+async def mcp_question(session: SessionDep, chat: ChatMcp):
+    user = await get_current_user(session, chat.token)
+    return await stream_sql(session, user, chat)
+
+
 @router.post("/start")
 async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat):
     try:
@@ -68,25 +80,6 @@ async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat
         )
 
 
-@router.post("/mcp_question", operation_id="mcp_question")
-async def mcp_question(session: SessionDep, token: str, request_question: ChatQuestion):
-    user = await get_current_user(session, token)
-    # return await stream_sql(session, user, request_question)
-    return {"content":"""步骤1: 确定需要查询的字段。
-我们需要统计上海的订单总数，因此需要从"城市"字段中筛选出值为"上海"的记录，并使用COUNT函数计算这些记录的数量。
-
-步骤2: 确定筛选条件。
-问题要求统计上海的订单总数，所以我们需要在SQL语句中添加WHERE "城市" = '上海'来筛选出符合条件的记录。
-
-步骤3: 避免关键字冲突。
-因为这个Excel/CSV数据库是 PostgreSQL 类型，所以在schema、表名、字段名和别名外层加双引号。
-
-最终答案:
-```json
-{"success":true,"sql":"SELECT COUNT(*) AS \"TotalOrders\" FROM \"public\".\"Sheet1_c27345b66e\" WHERE \"城市\" = '上海';"}
-```"""}
-
-
 @router.post("/question", operation_id="question")
 async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion):
     """Stream SQL analysis results
@@ -106,16 +99,17 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
             status_code=400,
             detail=f"Chat with id {request_question.chat_id} not found"
         )
-
-    # Get available datasource
-    ds = session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first()
-    if not ds:
-        raise HTTPException(
-            status_code=500,
-            detail="No available datasource configuration found"
-        )
-
-    request_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL'
+    ds: CoreDatasource | None = None
+    if chat.datasource:
+        # Get available datasource
+        ds = session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first()
+        if not ds:
+            raise HTTPException(
+                status_code=500,
+                detail="No available datasource configuration found"
+            )
+
+        request_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL'
 
     # Get available AI model
     aimodel = session.exec(select(AiModelDetail).where(
@@ -128,14 +122,18 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
             detail="No available AI model configuration found"
         )
 
-    history_records: List[ChatRecord] = list_records(session=session, current_user=current_user,
-                                                     chart_id=request_question.chat_id)
+    history_records: List[ChatRecord] = list(filter(lambda r: True if r.first_chat != True else False,
+                                                    list_records(session=session, current_user=current_user,
+                                                                 chart_id=request_question.chat_id)))
     # get schema
-    request_question.db_schema = get_table_schema(session=session, ds=ds)
+    if ds:
+        request_question.db_schema = get_table_schema(session=session, ds=ds)
+
     db_user = get_user_info(session=session, user_id=current_user.id)
     request_question.lang = db_user.language
 
-    llm_service = LLMService(request_question, aimodel, history_records, CoreDatasource(**ds.model_dump()))
+    llm_service = LLMService(request_question, aimodel, history_records,
+                             CoreDatasource(**ds.model_dump()) if ds else None)
 
     llm_service.init_record(session=session, current_user=current_user)
 
@@ -144,6 +142,16 @@ def run_task():
             # return id
             yield orjson.dumps({'type': 'id', 'id': llm_service.get_record().id}).decode() + '\n\n'
 
+            # select datasource if datasource is none
+            if not ds:
+                ds_res = llm_service.select_datasource(session=session)
+                for chunk in ds_res:
+                    yield orjson.dumps({'content': chunk, 'type': 'datasource-result'}).decode() + '\n\n'
+                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'
+
+                llm_service.chat_question.db_schema = get_table_schema(session=session, ds=llm_service.ds)
+
             # generate sql
             sql_res = llm_service.generate_sql(session=session)
             full_sql_text = ''

backend/apps/chat/curd/chat.py

@@ -60,6 +60,8 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr
         load_only(ChatRecord.id, ChatRecord.chat_id, ChatRecord.create_time, ChatRecord.finish_time,
                   ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql, ChatRecord.data,
                   ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.analysis, ChatRecord.predict,
+                  ChatRecord.datasource_select_answer, ChatRecord.recommended_question_answer,
+                  ChatRecord.recommended_question,
                   ChatRecord.predict_data, ChatRecord.finish, ChatRecord.error, ChatRecord.run_time)).filter(
         and_(Chat.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(ChatRecord.create_time).all()
 
@@ -74,7 +76,8 @@ def list_records(session: SessionDep, chart_id: int, current_user: CurrentUser)
     return record_list
 
 
-def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat, require_datasource: bool = True) -> ChatInfo:
+def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat,
+                require_datasource: bool = True) -> ChatInfo:
     if not create_chat_obj.datasource and require_datasource:
         raise Exception("Datasource cannot be None")
 
@@ -84,7 +87,7 @@ def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj:
     chat = Chat(create_time=datetime.datetime.now(),
                 create_by=current_user.id,
                 brief=create_chat_obj.question.strip()[:20])
-    ds: CoreDatasource = None
+    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()
@@ -93,6 +96,8 @@ def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj:
             raise Exception(f"Datasource with id {create_chat_obj.datasource} not found")
 
         chat.engine_type = ds.type_name
+    else:
+        chat.engine_type = ''
 
     chat_info = ChatInfo(**chat.model_dump())
 
@@ -102,15 +107,9 @@ def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj:
     chat_info.id = chat.id
     session.commit()
 
-    if not create_chat_obj.datasource:
-        # use AI to get ds
-
-        if not ds:
-            raise Exception(f"Datasource with id {create_chat_obj.datasource} not found")
-
-
-    chat_info.datasource_exists = True
-    chat_info.datasource_name = ds.name
+    if ds:
+        chat_info.datasource_exists = True
+        chat_info.datasource_name = ds.name
 
     return chat_info
 
@@ -205,6 +204,30 @@ def save_full_predict_message_and_answer(session: SessionDep, record_id: int, an
     return result
 
 
+def save_full_select_datasource_message_and_answer(session: SessionDep, record_id: int, answer: str,
+                                                   full_message: str, datasource: int = None,
+                                                   engine_type: str = None) -> ChatRecord:
+    if not record_id:
+        raise Exception("Record id cannot be None")
+    record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first()
+    record.full_select_datasource_message = full_message
+    record.datasource_select_answer = answer
+
+    if datasource:
+        record.datasource = datasource
+        record.engine_type = engine_type
+
+    result = ChatRecord(**record.model_dump())
+
+    session.add(record)
+    session.flush()
+    session.refresh(record)
+
+    session.commit()
+
+    return result
+
+
 def save_sql(session: SessionDep, record_id: int, sql: str) -> ChatRecord:
     if not record_id:
         raise Exception("Record id cannot be None")