feat: modify generate chat

2026-11-08 01:06:59 +08:00 · 2025-07-01 18:46:40 +08:00
parent d03e8fcee2
commit 87241b6cda
3 changed files with 195 additions and 186 deletions
@@ -1,20 +1,14 @@
 import traceback
-from typing import List

 import orjson
 from fastapi import APIRouter, HTTPException
 from fastapi.responses import StreamingResponse
-from sqlmodel import select

 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, ChatMcp
-from apps.chat.task.llm import LLMService
-from apps.datasource.crud.datasource import get_table_schema
-from apps.datasource.models.datasource import CoreDatasource
-from apps.system.crud.user import get_user_info
-from apps.system.models.system_model import AiModelDetail
-from common.core.deps import SessionDep, CurrentUser, get_current_user
+    delete_chat
+from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, ChatQuestion
+from apps.chat.task.llm import LLMService, run_task
+from common.core.deps import SessionDep, CurrentUser

 router = APIRouter(tags=["Data Q&A"], prefix="/chat")

@@ -81,107 +75,17 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
        Streaming response with analysis results
    """

-    chat = session.query(Chat).filter(Chat.id == request_question.chat_id).first()
-    if not chat:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Chat with id {request_question.chat_id} not found"
-        )
-    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(
-        AiModelDetail.status == True,
-        AiModelDetail.api_key.is_not(None)
-    )).first()
-    if not aimodel:
+    try:
+        llm_service = LLMService(session, current_user, request_question)
+        llm_service.init_record()
+    except Exception as e:
+        traceback.print_exc()
        raise HTTPException(
            status_code=500,
-            detail="No available AI model configuration found"
+            detail=str(e)
        )

-    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
-    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()) if ds else None)
-
-    llm_service.init_record(session=session, current_user=current_user)
-
-    def run_task():
-        try:
-            # 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 = ''
-            for chunk in sql_res:
-                full_sql_text += chunk
-                yield orjson.dumps({'content': chunk, 'type': 'sql-result'}).decode() + '\n\n'
-            yield orjson.dumps({'type': 'info', 'msg': 'sql generated'}).decode() + '\n\n'
-
-            # filter sql
-            print(full_sql_text)
-            sql = llm_service.check_save_sql(session=session, res=full_sql_text)
-            print(sql)
-            yield orjson.dumps({'content': sql, 'type': 'sql'}).decode() + '\n\n'
-
-            # execute sql
-            result = llm_service.execute_sql(sql=sql)
-            llm_service.save_sql_data(session=session, data_obj=result)
-            yield orjson.dumps({'content': orjson.dumps(result).decode(), 'type': 'sql-data'}).decode() + '\n\n'
-
-            # generate chart
-            chart_res = llm_service.generate_chart(session=session)
-            full_chart_text = ''
-            for chunk in chart_res:
-                full_chart_text += chunk
-                yield orjson.dumps({'content': chunk, 'type': 'chart-result'}).decode() + '\n\n'
-            yield orjson.dumps({'type': 'info', 'msg': 'chart generated'}).decode() + '\n\n'
-
-            # filter chart
-            print(full_chart_text)
-            chart = llm_service.check_save_chart(session=session, res=full_chart_text)
-            print(chart)
-            yield orjson.dumps({'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n'
-
-            llm_service.finish(session=session)
-            yield orjson.dumps({'type': 'finish'}).decode() + '\n\n'
-
-        except Exception as e:
-            traceback.print_exc()
-            llm_service.save_error(session=session, message=str(e))
-            yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
-
-    return StreamingResponse(run_task(), media_type="text/event-stream")
+    return StreamingResponse(run_task(llm_service, session), media_type="text/event-stream")


@router.post("/record/{chart_record_id}/{action_type}")
@@ -205,35 +109,9 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch
            detail=f"Chat record with id {chart_record_id} has not generated chart, do not support to analyze it"
        )

-    chat = session.query(Chat).filter(Chat.id == record.chat_id).first()
-    if not chat:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Chat with id {record.chart_id} not found"
-        )
+    request_question = ChatQuestion(chat_id=record.chat_id, question='')

-    if chat.create_by != current_user.id:
-        raise HTTPException(
-            status_code=401,
-            detail=f"You cannot use the chat with id {record.chart_id}"
-        )
-
-    # Get available AI model
-    aimodel = session.exec(select(AiModelDetail).where(
-        AiModelDetail.status == True,
-        AiModelDetail.api_key.is_not(None)
-    )).first()
-    if not aimodel:
-        raise HTTPException(
-            status_code=500,
-            detail="No available AI model configuration found"
-        )
-
-    request_question = ChatQuestion(chat_id=chat.id, question='')
-    db_user = get_user_info(session=session, user_id=current_user.id)
-    request_question.lang = db_user.language
-
-    llm_service = LLMService(request_question, aimodel)
+    llm_service = LLMService(session, current_user, request_question)
    llm_service.set_record(record)

    def run_task():
@@ -249,14 +127,14 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch

            elif action_type == 'predict':
                # generate predict
-                analysis_res = llm_service.generate_predict(session=session)
+                analysis_res = llm_service.generate_predict()
                full_text = ''
                for chunk in analysis_res:
                    yield orjson.dumps({'content': chunk, 'type': 'predict-result'}).decode() + '\n\n'
                    full_text += chunk
                yield orjson.dumps({'type': 'info', 'msg': 'predict generated'}).decode() + '\n\n'

-                _data = llm_service.check_save_predict_data(session=session, res=full_text)
+                _data = llm_service.check_save_predict_data(res=full_text)
                yield orjson.dumps({'type': 'predict', 'content': _data}).decode() + '\n\n'

                yield orjson.dumps({'type': 'predict_finish'}).decode() + '\n\n'
@@ -3,6 +3,7 @@ import warnings
 from typing import Any, List, Union, Dict

 import orjson
+import pandas as pd
 from langchain_community.utilities import SQLDatabase
 from langchain_core.language_models import BaseLLM
 from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, AIMessageChunk
@@ -13,10 +14,12 @@ from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_llm_config
 from apps.chat.curd.chat import save_question, save_full_sql_message, save_full_sql_message_and_answer, save_sql, \
    save_error_message, save_sql_exec_data, save_full_chart_message, save_full_chart_message_and_answer, save_chart, \
    finish_record, save_full_analysis_message_and_answer, save_full_predict_message_and_answer, save_predict_data, \
-    save_full_select_datasource_message_and_answer
+    save_full_select_datasource_message_and_answer, list_records
 from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat
+from apps.datasource.crud.datasource import get_table_schema
 from apps.datasource.models.datasource import CoreDatasource
 from apps.db.db import exec_sql
+from apps.system.crud.user import get_user_info
 from apps.system.models.system_model import AiModelDetail
 from common.core.deps import SessionDep, CurrentUser

@@ -34,13 +37,48 @@ class LLMService:
    sql_message: List[Union[BaseMessage, dict[str, Any]]] = []
    chart_message: List[Union[BaseMessage, dict[str, Any]]] = []
    history_records: List[ChatRecord] = []
+    session: SessionDep
+    _current_user: CurrentUser

-    def __init__(self, chat_question: ChatQuestion, aimodel: AiModelDetail, history_records: List[ChatRecord] = [],
-                 ds: CoreDatasource = None):
-        self.ds = ds
+    def __init__(self, session: SessionDep, current_user: CurrentUser, chat_question: ChatQuestion):
+
+        self.session = session
+        self.current_user = current_user
+
+        chat = self.session.query(Chat).filter(Chat.id == chat_question.chat_id).first()
+        if not chat:
+            raise Exception(f"Chat with id {chat_question.chat_id} not found")
+        ds: CoreDatasource | None = None
+        if chat.datasource:
+            # Get available datasource
+            ds = self.session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first()
+            if not ds:
+                raise Exception("No available datasource configuration found")
+
+            chat_question.engine = ds.type_name if ds.type != 'excel' else 'PostgreSQL'
+
+        # Get available AI model
+        aimodel = self.session.exec(select(AiModelDetail).where(
+            AiModelDetail.status == True,
+            AiModelDetail.api_key.is_not(None)
+        )).first()
+        if not aimodel and aimodel[0]:
+            raise Exception("No available AI model configuration found")
+
+        history_records: List[ChatRecord] = list(filter(lambda r: True if r.first_chat != True else False,
+                                                        list_records(session=self.session, current_user=current_user,
+                                                                     chart_id=chat_question.chat_id)))
+        # get schema
+        if ds:
+            chat_question.db_schema = get_table_schema(session=self.session, ds=ds)
+
+        db_user = get_user_info(session=self.session, user_id=current_user.id)
+        chat_question.lang = db_user.language
+
+        self.ds = CoreDatasource(**ds.model_dump()) if ds else None
        self.chat_question = chat_question
-        self.chat_question.ai_modal_id = aimodel.id
-        self.config = get_llm_config(aimodel)
+        self.chat_question.ai_modal_id = aimodel[0].id
+        self.config = get_llm_config(aimodel[0])

        # Create LLM instance through factory
        llm_instance = LLMFactory.create_llm(self.config)
@@ -112,8 +150,8 @@ class LLMService:
                    _msg = AIMessage(content=last_chart_message['content'])
                    self.chart_message.append(_msg)

-    def init_record(self, session: SessionDep, current_user: CurrentUser) -> ChatRecord:
-        self.record = save_question(session=session, current_user=current_user, question=self.chat_question)
+    def init_record(self) -> ChatRecord:
+        self.record = save_question(session=self.session, current_user=self.current_user, question=self.chat_question)
        return self.record

    def get_record(self):
@@ -183,7 +221,7 @@ class LLMService:
                                                                                       in
                                                                                       analysis_msg]).decode())

-    def generate_predict(self, session: SessionDep):
+    def generate_predict(self):
        fields = self.get_fields_from_chart()

        self.chat_question.fields = orjson.dumps(fields).decode()
@@ -196,7 +234,7 @@ class LLMService:
        if self.record.full_predict_message and self.record.full_predict_message.strip() != '':
            history_msg = orjson.loads(self.record.full_predict_message)

-        self.record = save_full_predict_message_and_answer(session=session, record_id=self.record.id, answer='',
+        self.record = save_full_predict_message_and_answer(session=self.session, record_id=self.record.id, answer='',
                                                           data='',
                                                           full_message=orjson.dumps(history_msg +
                                                                                     [{'type': msg.type,
@@ -218,7 +256,7 @@ class LLMService:
                continue

        predict_msg.append(AIMessage(full_predict_text))
-        self.record = save_full_predict_message_and_answer(session=session, record_id=self.record.id,
+        self.record = save_full_predict_message_and_answer(session=self.session, record_id=self.record.id,
                                                           answer=full_predict_text, data='',
                                                           full_message=orjson.dumps(history_msg +
                                                                                     [{'type': msg.type,
@@ -226,10 +264,10 @@ class LLMService:
                                                                                      in
                                                                                      predict_msg]).decode())

-    def select_datasource(self, session: SessionDep):
+    def select_datasource(self):
        datasource_msg: List[Union[BaseMessage, dict[str, Any]]] = []
        datasource_msg.append(SystemMessage(self.chat_question.datasource_sys_question()))
-        _ds_list = session.exec(select(CoreDatasource).options(
+        _ds_list = self.session.exec(select(CoreDatasource).options(
            load_only(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description))).all()
        _ds_list_dict = []
        for _ds in _ds_list:
@@ -241,7 +279,7 @@ class LLMService:
        if self.record.full_select_datasource_message and self.record.full_select_datasource_message.strip() != '':
            history_msg = orjson.loads(self.record.full_select_datasource_message)

-        self.record = save_full_select_datasource_message_and_answer(session=session, record_id=self.record.id,
+        self.record = save_full_select_datasource_message_and_answer(session=self.session, record_id=self.record.id,
                                                                     answer='',
                                                                     full_message=orjson.dumps(history_msg +
                                                                                               [{'type': msg.type,
@@ -273,7 +311,7 @@ class LLMService:

            if data['id'] and data['id'] != 0:
                _datasource = data['id']
-                _ds = session.query(CoreDatasource).filter(CoreDatasource.id == _datasource).first()
+                _ds = self.session.query(CoreDatasource).filter(CoreDatasource.id == _datasource).first()
                if not _ds:
                    _datasource = None
                    raise Exception(f"Datasource configuration with id {_datasource} not found")
@@ -281,14 +319,14 @@ class LLMService:
                self.chat_question.engine = _ds.type_name if _ds.type != 'excel' else 'PostgreSQL'
                _engine_type = self.chat_question.engine
                # save chat
-                _chat = session.query(Chat).filter(Chat.id == self.record.chat_id).first()
+                _chat = self.session.query(Chat).filter(Chat.id == self.record.chat_id).first()
                _chat.datasource = _datasource
                _chat.engine_type = _ds.type_name

-                session.add(_chat)
-                session.flush()
-                session.refresh(_chat)
-                session.commit()
+                self.session.add(_chat)
+                self.session.flush()
+                self.session.refresh(_chat)
+                self.session.commit()

            elif data['fail']:
                raise Exception(data['fail'])
@@ -298,7 +336,7 @@ class LLMService:
        except Exception as e:
            _error = e

-        self.record = save_full_select_datasource_message_and_answer(session=session, record_id=self.record.id,
+        self.record = save_full_select_datasource_message_and_answer(session=self.session, record_id=self.record.id,
                                                                     answer=full_text, datasource=_datasource,
                                                                     engine_type=_engine_type,
                                                                     full_message=orjson.dumps(history_msg +
@@ -312,10 +350,10 @@ class LLMService:
        if _error:
            raise _error

-    def generate_sql(self, session: SessionDep):
+    def generate_sql(self):
        # append current question
        self.sql_message.append(HumanMessage(self.chat_question.sql_user_question()))
-        self.record = save_full_sql_message(session=session, record_id=self.record.id,
+        self.record = save_full_sql_message(session=self.session, record_id=self.record.id,
                                            full_message=orjson.dumps(
                                                [{'type': msg.type, 'content': msg.content} for msg in
                                                 self.sql_message]).decode())
@@ -333,16 +371,16 @@ class LLMService:
                continue

        self.sql_message.append(AIMessage(full_sql_text))
-        self.record = save_full_sql_message_and_answer(session=session, record_id=self.record.id,
+        self.record = save_full_sql_message_and_answer(session=self.session, record_id=self.record.id,
                                                       answer=full_sql_text,
                                                       full_message=orjson.dumps(
                                                           [{'type': msg.type, 'content': msg.content} for msg in
                                                            self.sql_message]).decode())

-    def generate_chart(self, session: SessionDep):
+    def generate_chart(self):
        # append current question
        self.chart_message.append(HumanMessage(self.chat_question.chart_user_question()))
-        self.record = save_full_chart_message(session=session, record_id=self.record.id,
+        self.record = save_full_chart_message(session=self.session, record_id=self.record.id,
                                              full_message=orjson.dumps(
                                                  [{'type': msg.type, 'content': msg.content} for msg in
                                                   self.chart_message]).decode())
@@ -359,13 +397,13 @@ class LLMService:
                continue

        self.chart_message.append(AIMessage(full_chart_text))
-        self.record = save_full_chart_message_and_answer(session=session, record_id=self.record.id,
+        self.record = save_full_chart_message_and_answer(session=self.session, record_id=self.record.id,
                                                         answer=full_chart_text,
                                                         full_message=orjson.dumps(
                                                             [{'type': msg.type, 'content': msg.content} for msg in
                                                              self.chart_message]).decode())

-    def check_save_sql(self, session: SessionDep, res: str) -> str:
+    def check_save_sql(self, res: str) -> str:

        json_str = extract_nested_json(res)
        data = orjson.loads(json_str)
@@ -385,13 +423,13 @@ class LLMService:
        if sql.strip() == '':
            raise Exception("SQL query is empty")

-        save_sql(session=session, sql=sql, record_id=self.record.id)
+        save_sql(session=self.session, sql=sql, record_id=self.record.id)

        self.chat_question.sql = sql

        return sql

-    def check_save_chart(self, session: SessionDep, res: str) -> Dict[str, Any]:
+    def check_save_chart(self, res: str) -> Dict[str, Any]:

        json_str = extract_nested_json(res)
        data = orjson.loads(json_str)
@@ -413,30 +451,30 @@ class LLMService:
        if error:
            raise Exception(message)

-        save_chart(session=session, chart=orjson.dumps(chart).decode(), record_id=self.record.id)
+        save_chart(session=self.session, chart=orjson.dumps(chart).decode(), record_id=self.record.id)

        return chart

-    def check_save_predict_data(self, session: SessionDep, res: str) -> Dict[str, Any]:
+    def check_save_predict_data(self, res: str) -> Dict[str, Any]:

        json_str = extract_nested_json(res)

        if not json_str:
            json_str = ''

-        save_predict_data(session=session, record_id=self.record.id, data=json_str)
+        save_predict_data(session=self.session, record_id=self.record.id, data=json_str)

        return json_str

-    def save_error(self, session: SessionDep, message: str):
-        return save_error_message(session=session, record_id=self.record.id, message=message)
+    def save_error(self, message: str):
+        return save_error_message(session=self.session, record_id=self.record.id, message=message)

-    def save_sql_data(self, session: SessionDep, data_obj: Dict[str, Any]):
-        return save_sql_exec_data(session=session, record_id=self.record.id,
+    def save_sql_data(self, data_obj: Dict[str, Any]):
+        return save_sql_exec_data(session=self.session, record_id=self.record.id,
                                  data=orjson.dumps(data_obj).decode())

-    def finish(self, session: SessionDep):
-        return finish_record(session=session, record_id=self.record.id)
+    def finish(self):
+        return finish_record(session=self.session, record_id=self.record.id)

    def execute_sql(self, sql: str):
        """Execute SQL query
@@ -503,3 +541,87 @@ def execute_sql_with_db(db: SQLDatabase, sql: str) -> str:
        error_msg = f"SQL execution failed: {str(e)}"
        logging.error(error_msg)
        raise RuntimeError(error_msg)
+
+
+def run_task(llm_service: LLMService, session: SessionDep, stream: bool = True):
+    try:
+        # return id
+        if stream:
+            yield orjson.dumps({'type': 'id', 'id': llm_service.get_record().id}).decode() + '\n\n'
+
+        # select datasource if datasource is none
+        if not llm_service.ds:
+            ds_res = llm_service.select_datasource()
+            if stream:
+                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()
+        full_sql_text = ''
+        for chunk in sql_res:
+            full_sql_text += chunk
+            if stream:
+                yield orjson.dumps({'content': chunk, 'type': 'sql-result'}).decode() + '\n\n'
+        if stream:
+            yield orjson.dumps({'type': 'info', 'msg': 'sql generated'}).decode() + '\n\n'
+
+        # filter sql
+        print(full_sql_text)
+        sql = llm_service.check_save_sql(res=full_sql_text)
+        print(sql)
+        if stream:
+            yield orjson.dumps({'content': sql, 'type': 'sql'}).decode() + '\n\n'
+
+        # execute sql
+        result = llm_service.execute_sql(sql=sql)
+        llm_service.save_sql_data(data_obj=result)
+        if stream:
+            yield orjson.dumps({'content': orjson.dumps(result).decode(), 'type': 'sql-data'}).decode() + '\n\n'
+
+        # generate chart
+        chart_res = llm_service.generate_chart()
+        full_chart_text = ''
+        for chunk in chart_res:
+            full_chart_text += chunk
+            if stream:
+                yield orjson.dumps({'content': chunk, 'type': 'chart-result'}).decode() + '\n\n'
+        if stream:
+            yield orjson.dumps({'type': 'info', 'msg': 'chart generated'}).decode() + '\n\n'
+
+        # filter chart
+        print(full_chart_text)
+        chart = llm_service.check_save_chart(res=full_chart_text)
+        print(chart)
+        if stream:
+            yield orjson.dumps({'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n'
+
+        record = llm_service.finish()
+        if stream:
+            yield orjson.dumps({'type': 'finish'}).decode() + '\n\n'
+        else:
+            md_str = f'```sql\n{sql}\n```\n\n'
+            # todo generate picture
+            if chart['type'] == 'table':
+                data = {}
+                for _data in result['data']:
+                    for field in result['fields']:
+                        if not data[field]:
+                            data[field] = []
+                        data[field].append(_data[field])
+                df = pd.DataFrame(data, columns=result['fields'])
+                markdown_table = df.to_markdown(index=False)
+                md_str += markdown_table
+            else:
+                md_str += ''
+
+    except Exception as e:
+        llm_service.save_error(message=str(e))
+        if stream:
+            yield orjson.dumps({'content': str(e), 'type': 'error'}).decode() + '\n\n'
+        else:
+            raise e
@@ -1,19 +1,22 @@
 # Author: Junjun
 # Date: 2025/7/1

+from datetime import timedelta
 from typing import Annotated
+
 from fastapi import APIRouter, Depends, HTTPException
 from fastapi.security import OAuth2PasswordRequestForm
-from common.core.deps import SessionDep, get_current_user
-from apps.system.crud.user import authenticate
-from common.core.security import create_access_token
-from datetime import timedelta
-from common.core.config import settings
-from common.core.schemas import Token
-from apps.datasource.crud.datasource import get_datasource_list
-from apps.system.models.system_model import AiModelDetail
+
+from apps.chat.api.chat import create_chat
 from apps.chat.models.chat_model import ChatMcp, CreateChat
-from apps.chat.api.chat import create_chat, stream_sql
+from apps.chat.task.llm import LLMService, run_task
+from apps.datasource.crud.datasource import get_datasource_list
+from apps.system.crud.user import authenticate
+from apps.system.models.system_model import AiModelDetail
+from common.core.config import settings
+from common.core.deps import SessionDep, get_current_user
+from common.core.schemas import Token
+from common.core.security import create_access_token

 router = APIRouter(tags=["mcp"], prefix="/mcp")

@@ -52,6 +55,12 @@ async def mcp_start(session: SessionDep, chat: ChatMcp):
@router.post("/mcp_question", operation_id="mcp_question")
 async def mcp_question(session: SessionDep, chat: ChatMcp):
    user = await get_current_user(session, chat.token)
+
+    llm_service = LLMService(session, user, chat)
+    llm_service.init_record()
+
+    run_task(llm_service, session)
+
    # return await stream_sql(session, user, chat)
    return {"content": """这是一段写死的测试内容：

@@ -68,4 +77,4 @@ async def mcp_question(session: SessionDep, chat: ChatMcp):
    ```json
    {"success":true,"sql":"SELECT COUNT(*) AS \"TotalOrders\" FROM \"public\".\"Sheet1_c27345b66e\" WHERE \"城市\" = '上海';"}
    ```
-    <img src="https://sqlbot.fit2cloud.cn/images/111.png">"""}
+    <img src="https://sqlbot.fit2cloud.cn/images/111.png">"""}