🐛 Bugfix: The intelligent generation agent cannot recognize the selected tool

Author: Jasonxia007

backend/apps/prompt_app.py

@@ -27,7 +27,9 @@ async def generate_and_save_system_prompt_api(
             task_description=prompt_request.task_description,
             user_id=user_id,
             tenant_id=tenant_id,
-            language=language
+            language=language,
+            tool_ids=prompt_request.tool_ids,
+            sub_agent_ids=prompt_request.sub_agent_ids
         ), media_type="text/event-stream")
     except Exception as e:
         logger.exception(f"Error occurred while generating system prompt: {e}")

backend/consts/model.py

@@ -193,6 +193,8 @@ class GeneratePromptRequest(BaseModel):
     task_description: str
     agent_id: int
     model_id: int
+    tool_ids: Optional[List[int]] = None  # Optional: tool IDs from frontend (takes precedence over database query)
+    sub_agent_ids: Optional[List[int]] = None  # Optional: sub-agent IDs from frontend (takes precedence over database query)
 
 
 class GenerateTitleRequest(BaseModel):

backend/services/prompt_service.py

@@ -2,17 +2,17 @@
 import logging
 import queue
 import threading
+from typing import Optional, List
 
 from jinja2 import StrictUndefined, Template
 from smolagents import OpenAIServerModel
 
-from consts.const import LANGUAGE, MODEL_CONFIG_MAPPING, MESSAGE_ROLE, THINK_END_PATTERN, THINK_START_PATTERN
+from consts.const import LANGUAGE, MESSAGE_ROLE, THINK_END_PATTERN, THINK_START_PATTERN
 from consts.model import AgentInfoRequest
-from database.agent_db import update_agent, query_sub_agents_id_list, search_agent_info_by_agent_id
+from database.agent_db import update_agent, search_agent_info_by_agent_id
 from database.model_management_db import get_model_by_model_id
 from database.tool_db import query_tools_by_ids
-from services.agent_service import get_enable_tool_id_by_agent_id
-from utils.config_utils import tenant_config_manager, get_model_name_from_config
+from utils.config_utils import get_model_name_from_config
 from utils.prompt_template_utils import get_prompt_generate_prompt_template
 
 # Configure logging
@@ -34,7 +34,7 @@ def _process_thinking_tokens(new_token: str, is_thinking: bool, token_join: list
     """
     # Handle thinking mode
     if is_thinking:
-        return not (THINK_END_PATTERN in new_token)
+        return THINK_END_PATTERN not in new_token
 
     # Handle start of thinking
     if THINK_START_PATTERN in new_token:
@@ -98,14 +98,16 @@ def call_llm_for_system_prompt(model_id: int, user_prompt: str, system_prompt: s
         raise e
 
 
-def gen_system_prompt_streamable(agent_id: int, model_id: int, task_description: str, user_id: str, tenant_id: str, language: str):
+def gen_system_prompt_streamable(agent_id: int, model_id: int, task_description: str, user_id: str, tenant_id: str, language: str, tool_ids: Optional[List[int]] = None, sub_agent_ids: Optional[List[int]] = None):
     for system_prompt in generate_and_save_system_prompt_impl(
         agent_id=agent_id,
         model_id=model_id,
         task_description=task_description,
         user_id=user_id,
         tenant_id=tenant_id,
-        language=language
+        language=language,
+        tool_ids=tool_ids,
+        sub_agent_ids=sub_agent_ids
     ):
         # SSE format, each message ends with \n\n
         yield f"data: {json.dumps({'success': True, 'data': system_prompt}, ensure_ascii=False)}\n\n"
@@ -116,17 +118,35 @@ def generate_and_save_system_prompt_impl(agent_id: int,
                                          task_description: str,
                                          user_id: str,
                                          tenant_id: str,
-                                         language: str):
-    # Get description of tool and agent
-    # In create mode (agent_id=0), return empty lists
-    if agent_id == 0:
+                                         language: str,
+                                         tool_ids: Optional[List[int]] = None,
+                                         sub_agent_ids: Optional[List[int]] = None):
+    # Get description of tool and agent from frontend-provided IDs
+    # Frontend always provides tool_ids and sub_agent_ids (could be empty arrays)
+
+    # Handle tool IDs
+    if tool_ids and len(tool_ids) > 0:
+        tool_info_list = query_tools_by_ids(tool_ids)
+        logger.debug(f"Using frontend-provided tool IDs: {tool_ids}")
+    else:
         tool_info_list = []
+        logger.debug("No tools selected (empty tool_ids list)")
+
+    # Handle sub-agent IDs
+    if sub_agent_ids and len(sub_agent_ids) > 0:
         sub_agent_info_list = []
+        for sub_agent_id in sub_agent_ids:
+            try:
+                sub_agent_info = search_agent_info_by_agent_id(
+                    agent_id=sub_agent_id, tenant_id=tenant_id)
+                sub_agent_info_list.append(sub_agent_info)
+            except Exception as e:
+                logger.warning(
+                    f"Failed to get sub-agent info for agent_id {sub_agent_id}: {str(e)}")
+        logger.debug(f"Using frontend-provided sub-agent IDs: {sub_agent_ids}")
     else:
-        tool_info_list = get_enabled_tool_description_for_generate_prompt(
-            tenant_id=tenant_id, agent_id=agent_id)
-        sub_agent_info_list = get_enabled_sub_agent_description_for_generate_prompt(
-            tenant_id=tenant_id, agent_id=agent_id)
+        sub_agent_info_list = []
+        logger.debug("No sub-agents selected (empty sub_agent_ids list)")
 
     # 1. Real-time streaming push
     final_results = {"duty": "", "constraint": "", "few_shots": "", "agent_var_name": "", "agent_display_name": "",
@@ -292,27 +312,3 @@ def join_info_for_generate_system_prompt(prompt_for_generate, sub_agent_info_lis
         "assistant_description": assistant_description
     })
     return content
-
-
-def get_enabled_tool_description_for_generate_prompt(agent_id: int, tenant_id: str):
-    # Get tool information
-    logger.info("Fetching tool instances")
-    tool_id_list = get_enable_tool_id_by_agent_id(
-        agent_id=agent_id, tenant_id=tenant_id)
-    tool_info_list = query_tools_by_ids(tool_id_list)
-    return tool_info_list
-
-
-def get_enabled_sub_agent_description_for_generate_prompt(agent_id: int, tenant_id: str):
-    logger.info("Fetching sub-agents information")
-
-    sub_agent_id_list = query_sub_agents_id_list(
-        main_agent_id=agent_id, tenant_id=tenant_id)
-
-    sub_agent_info_list = []
-    for sub_agent_id in sub_agent_id_list:
-        sub_agent_info = search_agent_info_by_agent_id(
-            agent_id=sub_agent_id, tenant_id=tenant_id)
-
-        sub_agent_info_list.append(sub_agent_info)
-    return sub_agent_info_list

frontend/app/[locale]/agents/AgentConfiguration.tsx

@@ -161,12 +161,25 @@ export default forwardRef<AgentConfigHandle, AgentConfigProps>(function AgentCon
       const currentAgentName = agentName;
       const currentAgentDisplayName = agentDisplayName;
 
+      // Extract tool IDs from selected tools (convert string IDs to numbers)
+      // Always pass tool_ids array (empty array means no tools selected, undefined means use database)
+      // In edit mode, we want to use current selection, so pass the array even if empty
+      const toolIds = selectedTools.map((tool) => Number(tool.id));
+      
+      // Get sub-agent IDs from enabledAgentIds
+      // Always pass sub_agent_ids array (empty array means no sub-agents selected, undefined means use database)
+      // In edit mode, we want to use current selection, so pass the array even if empty
+      const subAgentIds = [...enabledAgentIds];
+
       // Call backend API to generate agent prompt
+      // Pass tool_ids and sub_agent_ids to use frontend selection instead of database query
       await generatePromptStream(
         {
           agent_id: agentIdToUse,
           task_description: businessLogic,
           model_id: selectedModel?.id?.toString() || "",
+          tool_ids: toolIds,
+          sub_agent_ids: subAgentIds,
         },
         (data) => {
           // Process streaming response data

frontend/types/agentConfig.ts

@@ -323,6 +323,8 @@ export interface GeneratePromptParams {
   agent_id: number;
   task_description: string;
   model_id: string;
+  tool_ids?: number[]; // Optional: tool IDs selected in frontend (takes precedence over database query)
+  sub_agent_ids?: number[]; // Optional: sub-agent IDs selected in frontend (takes precedence over database query)
 }
 
 /**

test/backend/services/test_prompt_service.py

@@ -36,8 +36,6 @@
     call_llm_for_system_prompt,
     generate_and_save_system_prompt_impl,
     gen_system_prompt_streamable,
-    get_enabled_tool_description_for_generate_prompt,
-    get_enabled_sub_agent_description_for_generate_prompt,
     generate_system_prompt,
     join_info_for_generate_system_prompt,
     _process_thinking_tokens
@@ -91,21 +89,21 @@ def test_call_llm_for_system_prompt(self, mock_get_model_name, mock_openai, mock
         )
 
     @patch('backend.services.prompt_service.generate_system_prompt')
-    @patch('backend.services.prompt_service.get_enabled_sub_agent_description_for_generate_prompt')
-    @patch('backend.services.prompt_service.get_enabled_tool_description_for_generate_prompt')
+    @patch('backend.services.prompt_service.query_tools_by_ids')
+    @patch('backend.services.prompt_service.search_agent_info_by_agent_id')
     @patch('backend.services.prompt_service.update_agent')
-    def test_generate_and_save_system_prompt_impl(self, mock_update_agent, mock_get_tool_desc,
-                                                  mock_get_agent_desc, mock_generate_system_prompt):
+    def test_generate_and_save_system_prompt_impl(self, mock_update_agent, mock_search_agent_info,
+                                                  mock_query_tools, mock_generate_system_prompt):
         # Setup
         mock_tool1 = {"name": "tool1", "description": "Tool 1 desc",
                       "inputs": "input1", "output_type": "output1"}
         mock_tool2 = {"name": "tool2", "description": "Tool 2 desc",
                       "inputs": "input2", "output_type": "output2"}
-        mock_get_tool_desc.return_value = [mock_tool1, mock_tool2]
+        mock_query_tools.return_value = [mock_tool1, mock_tool2]
 
         mock_agent1 = {"name": "agent1", "description": "Agent 1 desc"}
         mock_agent2 = {"name": "agent2", "description": "Agent 2 desc"}
-        mock_get_agent_desc.return_value = [mock_agent1, mock_agent2]
+        mock_search_agent_info.side_effect = [mock_agent1, mock_agent2]
 
         # Mock the generator to return the expected data structure
         def mock_generator(*args, **kwargs):
@@ -121,33 +119,34 @@ def mock_generator(*args, **kwargs):
 
         mock_generate_system_prompt.side_effect = mock_generator
 
-        # Execute - test as a generator
+        # Execute - test as a generator with frontend-provided IDs
         result_gen = generate_and_save_system_prompt_impl(
             agent_id=123,
             model_id=self.test_model_id,
             task_description="Test task",
             user_id="user123",
             tenant_id="tenant456",
-            language="zh"
+            language="zh",
+            tool_ids=[1, 2],
+            sub_agent_ids=[10, 20]
         )
         result = list(result_gen)  # Convert generator to list for assertion
 
         # Assert
         self.assertGreater(len(result), 0)
 
-        mock_get_tool_desc.assert_called_once_with(
-            agent_id=123, tenant_id="tenant456")
-        mock_get_agent_desc.assert_called_once_with(
-            agent_id=123, tenant_id="tenant456")
+        # Verify tools and agents were queried using frontend-provided IDs
+        mock_query_tools.assert_called_once_with([1, 2])
+        self.assertEqual(mock_search_agent_info.call_count, 2)
+        mock_search_agent_info.assert_any_call(agent_id=10, tenant_id="tenant456")
+        mock_search_agent_info.assert_any_call(agent_id=20, tenant_id="tenant456")
 
-        mock_generate_system_prompt.assert_called_once_with(
-            mock_get_agent_desc.return_value,
-            "Test task",
-            mock_get_tool_desc.return_value,
-            "tenant456",
-            self.test_model_id,
-            "zh"
-        )
+        # Verify generate_system_prompt was called with correct parameters
+        mock_generate_system_prompt.assert_called_once()
+        call_args = mock_generate_system_prompt.call_args
+        self.assertEqual(call_args[0][0], [mock_agent1, mock_agent2])  # sub_agent_info_list
+        self.assertEqual(call_args[0][1], "Test task")  # task_description
+        self.assertEqual(call_args[0][2], [mock_tool1, mock_tool2])  # tool_info_list
 
         # Verify update_agent was called with the correct parameters
         mock_update_agent.assert_called_once()
@@ -162,11 +161,8 @@ def mock_generator(*args, **kwargs):
         self.assertEqual(agent_info.business_description, "Test task")
 
     @patch('backend.services.prompt_service.generate_system_prompt')
-    @patch('backend.services.prompt_service.get_enabled_sub_agent_description_for_generate_prompt')
-    @patch('backend.services.prompt_service.get_enabled_tool_description_for_generate_prompt')
     @patch('backend.services.prompt_service.update_agent')
-    def test_generate_and_save_system_prompt_impl_create_mode(self, mock_update_agent, mock_get_tool_desc,
-                                                               mock_get_agent_desc, mock_generate_system_prompt):
+    def test_generate_and_save_system_prompt_impl_create_mode(self, mock_update_agent, mock_generate_system_prompt):
         """Test generate_and_save_system_prompt_impl in create mode (agent_id=0)"""
         # Setup - Mock the generator to return the expected data structure
         def mock_generator(*args, **kwargs):
@@ -182,25 +178,22 @@ def mock_generator(*args, **kwargs):
 
         mock_generate_system_prompt.side_effect = mock_generator
 
-        # Execute - test as a generator with agent_id=0 (create mode)
+        # Execute - test as a generator with agent_id=0 (create mode) and empty tool/sub-agent IDs
         result_gen = generate_and_save_system_prompt_impl(
             agent_id=0,
             model_id=self.test_model_id,
             task_description="Test task",
             user_id="user123",
             tenant_id="tenant456",
-            language="zh"
+            language="zh",
+            tool_ids=[],
+            sub_agent_ids=[]
         )
         result = list(result_gen)  # Convert generator to list for assertion
 
         # Assert
         self.assertGreater(len(result), 0)
 
-        # In create mode, should NOT call get_enabled_tool_description_for_generate_prompt
-        # and get_enabled_sub_agent_description_for_generate_prompt
-        mock_get_tool_desc.assert_not_called()
-        mock_get_agent_desc.assert_not_called()
-
         # Should call generate_system_prompt with empty lists for tools and sub-agents
         mock_generate_system_prompt.assert_called_once_with(
             [],  # Empty sub_agent_info_list
@@ -246,7 +239,9 @@ def test_gen_system_prompt_streamable(self, mock_generate_impl):
             task_description="Test task",
             user_id="user123",
             tenant_id="tenant456",
-            language="zh"
+            language="zh",
+            tool_ids=None,
+            sub_agent_ids=None,
         )
 
         # Verify output format - should be SSE format
@@ -473,68 +468,6 @@ def test_join_info_for_generate_system_prompt(self, mock_template):
         self.assertEqual(
             template_vars["task_description"], mock_task_description)
 
-    @patch('backend.services.prompt_service.get_enable_tool_id_by_agent_id')
-    @patch('backend.services.prompt_service.query_tools_by_ids')
-    def test_get_enabled_tool_description_for_generate_prompt(self, mock_query_tools, mock_get_tool_ids):
-        # Setup
-        mock_get_tool_ids.return_value = [1, 2, 3]
-        mock_tools = [{"id": 1, "name": "tool1"}, {
-            "id": 2, "name": "tool2"}, {"id": 3, "name": "tool3"}]
-        mock_query_tools.return_value = mock_tools
-
-        # Execute
-        result = get_enabled_tool_description_for_generate_prompt(
-            agent_id=123,
-            tenant_id="tenant456"
-        )
-
-        # Assert
-        self.assertEqual(result, mock_tools)
-        mock_get_tool_ids.assert_called_once_with(
-            agent_id=123, tenant_id="tenant456")
-        mock_query_tools.assert_called_once_with([1, 2, 3])
-
-    @patch('backend.services.prompt_service.search_agent_info_by_agent_id')
-    @patch('backend.services.prompt_service.query_sub_agents_id_list')
-    def test_get_enabled_sub_agent_description_for_generate_prompt(self, mock_query_sub_agents_id_list, mock_search_agent_info):
-        # Setup
-        mock_query_sub_agents_id_list.return_value = [1, 2, 3]
-
-        # Mock search_agent_info_by_agent_id to return different agent info for each ID
-        def mock_search_agent_info_side_effect(agent_id, tenant_id):
-            agent_info_map = {
-                1: {"id": 1, "name": "agent1", "enabled": True},
-                2: {"id": 2, "name": "agent2", "enabled": False},
-                3: {"id": 3, "name": "agent3", "enabled": True}
-            }
-            return agent_info_map.get(agent_id, {})
-
-        mock_search_agent_info.side_effect = mock_search_agent_info_side_effect
-
-        # Execute
-        result = get_enabled_sub_agent_description_for_generate_prompt(
-            agent_id=123,
-            tenant_id="tenant456"
-        )
-
-        # Assert
-        expected_result = [
-            {"id": 1, "name": "agent1", "enabled": True},
-            {"id": 2, "name": "agent2", "enabled": False},
-            {"id": 3, "name": "agent3", "enabled": True}
-        ]
-        self.assertEqual(result, expected_result)
-        mock_query_sub_agents_id_list.assert_called_once_with(
-            main_agent_id=123, tenant_id="tenant456")
-
-        # Verify search_agent_info_by_agent_id was called for each sub agent ID
-        self.assertEqual(mock_search_agent_info.call_count, 3)
-        mock_search_agent_info.assert_any_call(
-            agent_id=1, tenant_id="tenant456")
-        mock_search_agent_info.assert_any_call(
-            agent_id=2, tenant_id="tenant456")
-        mock_search_agent_info.assert_any_call(
-            agent_id=3, tenant_id="tenant456")
 
     @patch('backend.services.prompt_service.get_model_by_model_id')
     @patch('backend.services.prompt_service.OpenAIServerModel')