🐛 Bugfix: Enhance prompt generation with knowledge base display names part2

backend/apps/prompt_app.py

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

backend/consts/model.py

@@ -304,6 +304,8 @@ class GeneratePromptRequest(BaseModel):
         None, description="Optional: tool IDs from frontend (takes precedence over database query)")
     sub_agent_ids: Optional[List[int]] = Field(
         None, description="Optional: sub-agent IDs from frontend (takes precedence over database query)")
+    knowledge_base_display_names: Optional[List[str]] = Field(
+        None, description="Optional: knowledge base display names from frontend (takes precedence over database query)")
 
 
 class GenerateTitleRequest(BaseModel):

backend/prompts/utils/prompt_generate_en.yaml

@@ -249,6 +249,13 @@ USER_PROMPT: |-
   You have no available assistants
   {% endif %}
 
+  {% if knowledge_base_names %}
+  ### Knowledge Base Configuration Note:
+  When generating few-shot examples, if using the knowledge_base_search tool, you MUST use the following actual configured knowledge base names:
+  {{knowledge_base_names}}
+  Please use these names directly in examples, e.g.: knowledge_base_search(query="xxx", index_names=[{{knowledge_base_names}}])
+  {% endif %}
+
 
 AGENT_NAME_REGENERATE_SYSTEM_PROMPT: |-
   ### You are an [Agent Variable Name Refinement Expert]

backend/prompts/utils/prompt_generate_zh.yaml

@@ -245,6 +245,13 @@ USER_PROMPT: |-
   你没有可用的助手
   {% endif %}
 
+  {% if knowledge_base_names %}
+  ### 知识库配置说明：
+  在生成 few-shot 示例时，如果使用 knowledge_base_search 工具，必须使用以下实际配置的知识库名称：
+  {{knowledge_base_names}}
+  请将这些名称直接用于示例中，例如：knowledge_base_search(query="xxx", index_names=[{{knowledge_base_names}}])
+  {% endif %}
+
 
 AGENT_NAME_REGENERATE_SYSTEM_PROMPT: |-
   ### 你是【Agent变量名调整专家】

backend/services/prompt_service.py

@@ -12,7 +12,8 @@
 from consts.exceptions import AppException
 from database.agent_db import search_agent_info_by_agent_id, query_all_agent_info_by_tenant_id, \
     query_sub_agents_id_list
-from database.tool_db import query_tools_by_ids
+from database.knowledge_db import get_knowledge_name_map_by_index_names
+from database.tool_db import query_tools_by_ids, query_tool_instances_by_id
 from services.agent_service import (
     get_enable_tool_id_by_agent_id,
     _check_agent_name_duplicate,
@@ -29,7 +30,7 @@
 logger = logging.getLogger("prompt_service")
 
 
-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):
+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, knowledge_base_display_names: Optional[List[str]] = None):
     try:
         for system_prompt in generate_and_save_system_prompt_impl(
             agent_id=agent_id,
@@ -39,7 +40,8 @@
             tenant_id=tenant_id,
             language=language,
             tool_ids=tool_ids,
-            sub_agent_ids=sub_agent_ids
+            sub_agent_ids=sub_agent_ids,
+            knowledge_base_display_names=knowledge_base_display_names
         ):
             # SSE format, each message ends with \n\n
             yield f"data: {json.dumps({'success': True, 'data': system_prompt}, ensure_ascii=False)}\n\n"
@@ -63,7 +65,8 @@
                                          tenant_id: str,
                                          language: str,
                                          tool_ids: Optional[List[int]] = None,
-                                         sub_agent_ids: Optional[List[int]] = None):
+                                         sub_agent_ids: Optional[List[int]] = None,
+                                         knowledge_base_display_names: Optional[List[str]] = None):
     # Get description of tool and agent from frontend-provided IDs
     # Frontend always provides tool_ids and sub_agent_ids (could be empty arrays)
 
@@ -77,6 +80,18 @@
         tool_info_list = get_enabled_tool_description_for_generate_prompt(
             tenant_id=tenant_id, agent_id=agent_id)
 
+    # Get knowledge base display names for few-shot examples
+    # Priority: frontend-provided > database query
+    if knowledge_base_display_names:
+        logger.debug(f"Using frontend-provided knowledge base display names: {knowledge_base_display_names}")
+    else:
+        knowledge_base_display_names = get_knowledge_base_display_names(
+            tool_info_list=tool_info_list,
+            agent_id=agent_id,
+            tenant_id=tenant_id
+        )
+        logger.debug(f"Using database query for knowledge base display names: {knowledge_base_display_names}")
+
     # Handle sub-agent IDs
     if sub_agent_ids and len(sub_agent_ids) > 0:
         sub_agent_info_list = []
@@ -114,7 +129,7 @@
 
     # Collect results and yield non-name fields immediately, but hold name fields for duplicate checking
     for result_data in generate_system_prompt(sub_agent_info_list, task_description, tool_info_list, tenant_id,
-                                              model_id, language):
+                                              model_id, language, knowledge_base_display_names):
         result_type = result_data["type"]
         final_results[result_type] = result_data["content"]
 
@@ -223,7 +238,7 @@
         raise Exception("Failed to generate prompt content.")
 
 
-def generate_system_prompt(sub_agent_info_list, task_description, tool_info_list, tenant_id: str, model_id: int, language: str = LANGUAGE["ZH"]):
+def generate_system_prompt(sub_agent_info_list, task_description, tool_info_list, tenant_id: str, model_id: int, language: str = LANGUAGE["ZH"], knowledge_base_display_names: Optional[List[str]] = None):
     """Main function for generating system prompts"""
     prompt_for_generate = get_prompt_generate_prompt_template(language)
 
@@ -233,7 +248,8 @@
         sub_agent_info_list=sub_agent_info_list,
         task_description=task_description,
         tool_info_list=tool_info_list,
-        language=language
+        language=language,
+        knowledge_base_display_names=knowledge_base_display_names
     )
 
     # Initialize state
@@ -352,7 +368,7 @@
             last_results[tag] = latest[tag]
 
 
-def join_info_for_generate_system_prompt(prompt_for_generate, sub_agent_info_list, task_description, tool_info_list, language: str = LANGUAGE["ZH"]):
+def join_info_for_generate_system_prompt(prompt_for_generate, sub_agent_info_list, task_description, tool_info_list, language: str = LANGUAGE["ZH"], knowledge_base_display_names: Optional[List[str]] = None):
     input_label = "Inputs" if language == 'en' else "接受输入"
     output_label = "Output type" if language == 'en' else "返回输出类型"
 
@@ -361,12 +377,21 @@
          for tool in tool_info_list])
     assistant_description = "\n".join(
         [f"- {sub_agent_info['name']}: {sub_agent_info['description']}" for sub_agent_info in sub_agent_info_list])
-    # Generate content using template
-    content = Template(prompt_for_generate["USER_PROMPT"], undefined=StrictUndefined).render({
+
+    # Build template context
+    template_context = {
         "task_description": task_description,
         "tool_description": tool_description,
         "assistant_description": assistant_description
-    })
+    }
+
+    # Add knowledge base display names for few-shot examples if available
+    if knowledge_base_display_names:
+        kb_names_str = ", ".join(f'"{name}"' for name in knowledge_base_display_names)
+        template_context["knowledge_base_names"] = kb_names_str
+
+    # Generate content using template
+    content = Template(prompt_for_generate["USER_PROMPT"], undefined=StrictUndefined).render(template_context)
     return content
 
 
@@ -379,6 +404,68 @@
     return tool_info_list
 
 
+def get_knowledge_base_display_names(tool_info_list: List[dict], agent_id: int, tenant_id: str) -> Optional[List[str]]:
+    """
+    Extract knowledge base display names from tool configurations.
+    This is used to ensure few-shot examples use actual configured knowledge base names.
+
+    Args:
+        tool_info_list: List of tool info dictionaries
+        agent_id: Agent ID for querying tool instances
+        tenant_id: Tenant ID for database queries
+
+    Returns:
+        List of knowledge base display names if knowledge_base_search tool is configured, None otherwise
+    """
+    # Check if knowledge_base_search tool is in the list
+    kb_tool_ids = [tool['tool_id'] for tool in tool_info_list if tool.get('name') == 'knowledge_base_search']
+    if not kb_tool_ids:
+        logger.debug("No knowledge_base_search tool found in tool list")
+        return None
+
+    # Get the index_names from ToolInstance for knowledge_base_search tool
+    all_index_names = []
+    for kb_tool_id in kb_tool_ids:
+        try:
+            tool_instance = query_tool_instances_by_id(
+                agent_id=agent_id,
+                tool_id=kb_tool_id,
+                tenant_id=tenant_id
+            )
+            if tool_instance and tool_instance.get('params', {}).get('index_names'):
+                index_names = tool_instance['params']['index_names']
+                if isinstance(index_names, list):
+                    all_index_names.extend(index_names)
+                elif isinstance(index_names, str):
+                    # Handle JSON string format
+                    try:
+                        all_index_names.extend(json.loads(index_names))
+                    except json.JSONDecodeError:
+                        logger.warning(f"Failed to parse index_names JSON: {index_names}")
+        except Exception as e:
+            logger.warning(f"Failed to get tool instance for tool_id {kb_tool_id}: {e}")
+
+    if not all_index_names:
+        logger.debug("No index_names configured for knowledge_base_search tool")
+        return None
+
+    # Remove duplicates while preserving order
+    unique_index_names = list(dict.fromkeys(all_index_names))
+
+    # Convert to display names
+    knowledge_name_map = get_knowledge_name_map_by_index_names(unique_index_names)
+
+    # Return list of display names (knowledge_name) for each configured index_name
+    display_names = []
+    for index_name in unique_index_names:
+        display_name = knowledge_name_map.get(index_name, index_name)
+        if display_name and display_name not in display_names:
+            display_names.append(display_name)
+
+    logger.debug(f"Converted index_names {unique_index_names} to display_names: {display_names}")
+    return display_names if display_names else None
+
+
 def get_enabled_sub_agent_description_for_generate_prompt(agent_id: int, tenant_id: str):
     logger.info("Fetching sub-agents information")
 

frontend/app/[locale]/agents/components/agentConfig/tool/ToolConfigModal.tsx

@@ -966,8 +966,16 @@ export default function ToolConfigModal({
         {} as Record<string, any>
       );
 
-      // Update local state: Add tool to selected tools with updated params
-      const updatedTool = { ...toolToSave, initParams: currentParams };
+      // Update local state: Add tool to selected tools with updated params and display_names
+      // Include display_names for knowledge base tools to pass to prompt generation
+      const updatedTool: typeof toolToSave = {
+        ...toolToSave,
+        initParams: currentParams,
+        // Store knowledge base display names for prompt generation
+        ...(toolRequiresKbSelection && selectedKbDisplayNames.length > 0
+          ? { display_names: selectedKbDisplayNames }
+          : {})
+      };
       const currentTools = useAgentConfigStore.getState().editedAgent.tools;
 
       // Check if tool already exists, if so replace it, otherwise add it

frontend/app/[locale]/agents/components/agentInfo/AgentGenerateDetail.tsx

@@ -591,6 +591,18 @@ export default function AgentGenerateDetail({
       businessLogicModelName: businessInfo.businessLogicModelName,
     });
 
+
+    // Extract knowledge base display names from selected tools
+    // This allows the backend to use frontend-configured display names without database lookup
+    const knowledgeBaseDisplayNames: string[] = [];
+    if (Array.isArray(editedAgent.tools)) {
+      for (const tool of editedAgent.tools) {
+        if (typeof tool === "object" && tool.display_names && Array.isArray(tool.display_names)) {
+          knowledgeBaseDisplayNames.push(...tool.display_names);
+        }
+      }
+    }
+
     try {
       await generatePromptStream(
         {
@@ -605,6 +617,8 @@ export default function AgentGenerateDetail({
                 : tool
             )
             : [],
+          // Pass knowledge base display names from frontend-configured tools
+          knowledge_base_display_names: knowledgeBaseDisplayNames.length > 0 ? knowledgeBaseDisplayNames : undefined,
         },
         (data) => {
           // Track the agent this generation was for

frontend/types/agentConfig.ts

@@ -76,6 +76,12 @@ export interface Tool {
   usage?: string;
   inputs?: string;
   category?: string;
+  /**
+   * Knowledge base display names associated with this tool.
+   * This is populated when the tool (e.g., knowledge_base_search) has knowledge bases configured.
+   * Used to pass knowledge base names to prompt generation without requiring database lookup.
+   */
+  display_names?: string[];
 }
 
 export interface ToolParam {
@@ -403,6 +409,13 @@ export interface GeneratePromptParams {
   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)
+  /**
+   * Optional: Knowledge base display names for few-shot examples.
+   * If provided, the backend will use these instead of querying the database.
+   * This allows the frontend to pass the latest configured knowledge base names
+   * without waiting for tool config to be saved first.
+   */
+  knowledge_base_display_names?: string[];
 }
 
 /**