From cd68760c75221935ef956aaf35ffd9429f037561 Mon Sep 17 00:00:00 2001
From: "Junyan Qin (Chin)" <rockchinq@gmail.com>
Date: Mon, 12 May 2025 18:21:08 +0800
Subject: [PATCH] Fix/runner bugs (#1388)

* fix: bugs in runners

* fix: model entity of exists conversation is None after changed runner
---
 pkg/pipeline/preproc/preproc.py    | 21 +++++++++++----------
 pkg/provider/runners/difysvapi.py  |  6 +++---
 pkg/provider/session/sessionmgr.py |  3 ---
 3 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/pkg/pipeline/preproc/preproc.py b/pkg/pipeline/preproc/preproc.py
index 2190a12d..6dd909be 100644
--- a/pkg/pipeline/preproc/preproc.py
+++ b/pkg/pipeline/preproc/preproc.py
@@ -30,12 +30,14 @@ class PreProcessor(stage.PipelineStage):
         stage_inst_name: str,
     ) -> entities.StageProcessResult:
         """处理"""
+        selected_runner = query.pipeline_config['ai']['runner']['runner']
+
         session = await self.ap.sess_mgr.get_session(query)
 
         # 非 local-agent 时，llm_model 为 None
         llm_model = (
             await self.ap.model_mgr.get_model_by_uuid(query.pipeline_config['ai']['local-agent']['model'])
-            if query.pipeline_config['ai']['runner'] == 'local-agent'
+            if selected_runner == 'local-agent'
             else None
         )
 
@@ -43,17 +45,18 @@ class PreProcessor(stage.PipelineStage):
             query,
             session,
             query.pipeline_config['ai']['local-agent']['prompt'],
-            llm_model,
         )
 
+        conversation.use_llm_model = llm_model
+
         # 设置query
         query.session = session
         query.prompt = conversation.prompt.copy()
         query.messages = conversation.messages.copy()
 
-        query.use_llm_model = conversation.use_llm_model
+        query.use_llm_model = llm_model
 
-        if query.pipeline_config['ai']['runner'] == 'local-agent':
+        if selected_runner == 'local-agent':
             query.use_funcs = (
                 conversation.use_funcs if query.use_llm_model.model_entity.abilities.__contains__('tool_call') else None
             )
@@ -68,9 +71,7 @@ class PreProcessor(stage.PipelineStage):
 
         # Check if this model supports vision, if not, remove all images
         # TODO this checking should be performed in runner, and in this stage, the image should be reserved
-        if query.pipeline_config['ai']['runner'][
-            'runner'
-        ] == 'local-agent' and not query.use_llm_model.model_entity.abilities.__contains__('vision'):
+        if selected_runner == 'local-agent' and not query.use_llm_model.model_entity.abilities.__contains__('vision'):
             for msg in query.messages:
                 if isinstance(msg.content, list):
                     for me in msg.content:
@@ -86,9 +87,9 @@ class PreProcessor(stage.PipelineStage):
                 content_list.append(llm_entities.ContentElement.from_text(me.text))
                 plain_text += me.text
             elif isinstance(me, platform_message.Image):
-                if query.pipeline_config['ai']['runner'][
-                    'runner'
-                ] != 'local-agent' or query.use_llm_model.model_entity.abilities.__contains__('vision'):
+                if selected_runner != 'local-agent' or query.use_llm_model.model_entity.abilities.__contains__(
+                    'vision'
+                ):
                     if me.base64 is not None:
                         content_list.append(llm_entities.ContentElement.from_image_base64(me.base64))
 
diff --git a/pkg/provider/runners/difysvapi.py b/pkg/provider/runners/difysvapi.py
index 26556851..35181d3c 100644
--- a/pkg/provider/runners/difysvapi.py
+++ b/pkg/provider/runners/difysvapi.py
@@ -121,7 +121,7 @@ class DifyServiceAPIRunner(runner.RequestRunner):
             user=f'{query.session.launcher_type.value}_{query.session.launcher_id}',
             conversation_id=cov_id,
             files=files,
-            timeout=self.pipeline_config['ai']['dify-service-api']['timeout'],
+            timeout=120,
         ):
             self.ap.logger.debug('dify-chat-chunk: ' + str(chunk))
 
@@ -184,7 +184,7 @@ class DifyServiceAPIRunner(runner.RequestRunner):
             response_mode='streaming',
             conversation_id=cov_id,
             files=files,
-            timeout=self.pipeline_config['ai']['dify-service-api']['timeout'],
+            timeout=120,
         ):
             self.ap.logger.debug('dify-agent-chunk: ' + str(chunk))
 
@@ -276,7 +276,7 @@ class DifyServiceAPIRunner(runner.RequestRunner):
             inputs=inputs,
             user=f'{query.session.launcher_type.value}_{query.session.launcher_id}',
             files=files,
-            timeout=self.pipeline_config['ai']['dify-service-api']['timeout'],
+            timeout=120,
         ):
             self.ap.logger.debug('dify-workflow-chunk: ' + str(chunk))
             if chunk['event'] in ignored_events:
diff --git a/pkg/provider/session/sessionmgr.py b/pkg/provider/session/sessionmgr.py
index e467bdac..8966be79 100644
--- a/pkg/provider/session/sessionmgr.py
+++ b/pkg/provider/session/sessionmgr.py
@@ -4,7 +4,6 @@ import asyncio
 
 from ...core import app, entities as core_entities
 from ...provider import entities as provider_entities
-from ...provider.modelmgr import entities as model_entities
 
 
 class SessionManager:
@@ -42,7 +41,6 @@ class SessionManager:
         query: core_entities.Query,
         session: core_entities.Session,
         prompt_config: list[dict],
-        llm_model: model_entities.LLMModelInfo,
     ) -> core_entities.Conversation:
         """获取对话或创建对话"""
 
@@ -64,7 +62,6 @@ class SessionManager:
             conversation = core_entities.Conversation(
                 prompt=prompt,
                 messages=[],
-                use_llm_model=llm_model,
                 use_funcs=await self.ap.tool_mgr.get_all_functions(
                     plugin_enabled=True,
                 ),