perf: Simple assistant

Author: fit2cloud-chenyw

backend/apps/chat/api/chat.py

@@ -7,7 +7,7 @@
     delete_chat, get_chat_chart_data, get_chat_predict_data
 from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, ChatQuestion
 from apps.chat.task.llm import LLMService, run_task, run_analysis_or_predict_task, run_recommend_questions_task
-from common.core.deps import SessionDep, CurrentUser
+from common.core.deps import CurrentAssistant, SessionDep, CurrentUser
 
 router = APIRouter(tags=["Data Q&A"], prefix="/chat")
 
@@ -94,7 +94,7 @@ async def start_chat(session: SessionDep, current_user: CurrentUser):
 
 
 @router.post("/recommend_questions/{chat_record_id}")
-async def recommend_questions(session: SessionDep, current_user: CurrentUser, chat_record_id: int):
+async def recommend_questions(session: SessionDep, current_user: CurrentUser, chat_record_id: int, current_assistant: CurrentAssistant):
     try:
         record = session.query(ChatRecord).get(chat_record_id)
         if not record:
@@ -104,7 +104,7 @@ async def recommend_questions(session: SessionDep, current_user: CurrentUser, ch
             )
         request_question = ChatQuestion(chat_id=record.chat_id, question=record.question if record.question else '')
 
-        llm_service = LLMService(session, current_user, request_question)
+        llm_service = LLMService(session, current_user, request_question, current_assistant)
         llm_service.set_record(record)
     except Exception as e:
         traceback.print_exc()
@@ -117,7 +117,7 @@ async def recommend_questions(session: SessionDep, current_user: CurrentUser, ch
 
 
 @router.post("/question")
-async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion):
+async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion, current_assistant: CurrentAssistant):
     """Stream SQL analysis results
     
     Args:
@@ -130,7 +130,7 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
     """
 
     try:
-        llm_service = LLMService(session, current_user, request_question)
+        llm_service = LLMService(session, current_user, request_question, current_assistant)
         llm_service.init_record()
     except Exception as e:
         traceback.print_exc()
@@ -143,7 +143,7 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
 
 
 @router.post("/record/{chat_record_id}/{action_type}")
-async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chat_record_id: int, action_type: str):
+async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chat_record_id: int, action_type: str, current_assistant: CurrentAssistant):
     if action_type != 'analysis' and action_type != 'predict':
         raise HTTPException(
             status_code=404,
@@ -166,7 +166,7 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch
     request_question = ChatQuestion(chat_id=record.chat_id, question='')
 
     try:
-        llm_service = LLMService(session, current_user, request_question)
+        llm_service = LLMService(session, current_user, request_question, current_assistant)
     except Exception as e:
         traceback.print_exc()
         raise HTTPException(

backend/apps/chat/task/llm.py

@@ -353,11 +353,21 @@ def select_datasource(self):
         if self.current_assistant:
             _ds_list = get_assistant_ds(llm_service=self)
         else:
-            _ds_list = self.session.exec(select(CoreDatasource).options(
-                load_only(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description))).all()
+            oid: str = self.current_user.oid
+            stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(CoreDatasource.oid == oid)
+            _ds_list = [
+                {
+                    "id": ds.id,
+                    "name": ds.name,
+                    "description": ds.description
+                }
+                for ds in self.session.exec(stmt)
+            ]
+            """ _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:
-            _ds_list_dict.append({'id': _ds[0].id, 'name': _ds[0].name, 'description': _ds[0].description})
+            _ds_list_dict.append(_ds)
         datasource_msg.append(
             HumanMessage(self.chat_question.datasource_user_question(orjson.dumps(_ds_list_dict).decode())))
 

backend/apps/system/crud/assistant.py

@@ -33,12 +33,22 @@ def get_assistant_ds(llm_service) -> list[dict]:
             config: dict[any] = json.loads(configuration)
             oid: str = config['oid']
             stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(CoreDatasource.oid == oid)
-            private_list:list[int] = config['private_list']
+            private_list:list[int] = config.get('private_list') or None
             if private_list:
                 stmt.where(~CoreDatasource.id.in_(private_list))
-        db_ds_list = session.exec(stmt).all()
+        db_ds_list = session.exec(stmt)
+        
+        result_list = [
+            {
+                "id": ds.id,
+                "name": ds.name,
+                "description": ds.description
+            }
+            for ds in db_ds_list
+        ]
+    
         # filter private ds if offline
-        return db_ds_list
+        return result_list
     out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(assistant, llm_service.assistant_certificate)
     llm_service.out_ds_instance = out_ds_instance
     dslist = out_ds_instance.get_simple_ds_list()