feat: auto select datasource

Author: ulleo

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/apps/chat/api/chat.py

@@ -72,7 +72,7 @@ async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat
 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: 确定需要查询的字段。
+    return {"content": """步骤1: 确定需要查询的字段。
 我们需要统计上海的订单总数，因此需要从"城市"字段中筛选出值为"上海"的记录，并使用COUNT函数计算这些记录的数量。
 
 步骤2: 确定筛选条件。
@@ -106,16 +106,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 +129,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 +149,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()
@@ -102,15 +105,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 +202,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")

backend/apps/chat/models/chat_model.py

@@ -7,8 +7,10 @@
 
 from apps.template.generate_analysis.generator import get_analysis_template
 from apps.template.generate_chart.generator import get_chart_template
+from apps.template.generate_guess_question.generator import get_guess_question_template
 from apps.template.generate_predict.generator import get_predict_template
 from apps.template.generate_sql.generator import get_sql_template
+from apps.template.select_datasource.generator import get_datasource_template
 
 
 class Chat(SQLModel, table=True):
@@ -25,11 +27,13 @@ class Chat(SQLModel, table=True):
 class ChatRecord(SQLModel, table=True):
     __tablename__ = "chat_record"
     id: Optional[int] = Field(sa_column=Column(Integer, Identity(always=True), primary_key=True))
-    chat_id: int = Field(sa_column=Column(Integer))
+    chat_id: int = Field(sa_column=Column(Integer, nullable=False))
+    ai_modal_id: Optional[int] = Field(sa_column=Column(Integer))
+    first_chat: bool = Field(sa_column=Column(Boolean, nullable=True, default=False))
     create_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
     finish_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
     create_by: int = Field(sa_column=Column(BigInteger, nullable=True))
-    datasource: int = Field(sa_column=Column(Integer, nullable=False))
+    datasource: int = Field(sa_column=Column(Integer, nullable=True))
     engine_type: str = Field(max_length=64)
     question: str = Field(sa_column=Column(Text, nullable=True))
     sql_answer: str = Field(sa_column=Column(Text, nullable=True))
@@ -41,10 +45,21 @@ class ChatRecord(SQLModel, table=True):
     analysis: str = Field(sa_column=Column(Text, nullable=True))
     predict: str = Field(sa_column=Column(Text, nullable=True))
     predict_data: str = Field(sa_column=Column(Text, nullable=True))
+    recommended_question_answer: str = Field(sa_column=Column(Text, nullable=True))
+    recommended_question: str = Field(sa_column=Column(Text, nullable=True))
+    datasource_select_answer: str = Field(sa_column=Column(Text, nullable=True))
     full_sql_message: str = Field(sa_column=Column(Text, nullable=True))
+    token_sql: int = Field(default=0, nullable=True)
     full_chart_message: str = Field(sa_column=Column(Text, nullable=True))
+    token_chart: int = Field(default=0, nullable=True)
     full_analysis_message: str = Field(sa_column=Column(Text, nullable=True))
+    token_analysis: int = Field(default=0, nullable=True)
     full_predict_message: str = Field(sa_column=Column(Text, nullable=True))
+    token_predict: int = Field(default=0, nullable=True)
+    full_recommended_question_message: str = Field(sa_column=Column(Text, nullable=True))
+    token_recommended_question: int = Field(default=0, nullable=True)
+    full_select_datasource_message: str = Field(sa_column=Column(Text, nullable=True))
+    token_select_datasource_question: int = Field(default=0, nullable=True)
     finish: bool = Field(sa_column=Column(Boolean, nullable=True, default=False))
     error: str = Field(sa_column=Column(Text, nullable=True))
     run_time: float = Field(default=0)
@@ -67,7 +82,7 @@ class ChatInfo(BaseModel):
     create_by: int = None
     brief: str = ''
     chat_type: str = "chat"
-    datasource: int = None
+    datasource: Optional[int] = None
     engine_type: str = ''
     datasource_name: str = ''
     datasource_exists: bool = True
@@ -108,6 +123,19 @@ def predict_sys_question(self):
     def predict_user_question(self):
         return get_predict_template()['user'].format(fields=self.fields, data=self.data, lang=self.lang)
 
+    def datasource_sys_question(self):
+        return get_datasource_template()['system']
+
+    def datasource_user_question(self, datasource_list: str = "[]"):
+        return get_datasource_template()['user'].format(question=self.question, data=datasource_list, lang=self.lang)
+
+    def datasource_guess_sys_question(self):
+        return get_guess_question_template()['system']
+
+    def datasource_guess_user_question(self, old_questions: str = "[]"):
+        return get_guess_question_template()['user'].format(question=self.question, schema=self.db_schema,
+                                                            old_questions=old_questions, lang=self.lang)
+
 
 class ChatQuestion(AiModelQuestion):
     question: str