Improve tool selection with additinal tool descriptions

Author: ravikiranvm

packages/server/api/src/app/ai/chat/prompts.service.ts

@@ -3,10 +3,20 @@ import { ChatFlowContext, CODE_BLOCK_NAME, isNil } from '@openops/shared';
 import { ToolSet } from 'ai';
 import { readFile } from 'fs/promises';
 import { join } from 'path';
+import { getAdditionalDescriptionsForSelectedTools } from '../mcp/external-tool-descriptions';
 import { hasToolProvider } from '../mcp/tool-utils';
 import { QueryClassification } from '../mcp/types';
 import { MCPChatContext } from './ai-chat.service';
 
+function buildAdditionalToolNotes(selectedTools: ToolSet | undefined): string {
+  if (!selectedTools) {
+    return '';
+  }
+
+  const descriptions = getAdditionalDescriptionsForSelectedTools(selectedTools);
+  return descriptions.length > 0 ? descriptions.join('\n\n') : '';
+}
+
 export const getMcpSystemPrompt = async ({
   queryClassification,
   selectedTools,
@@ -58,7 +68,12 @@ export const getMcpSystemPrompt = async ({
 
   const allPrompts = await Promise.all(promptPromises);
 
-  return allPrompts.join('\n\n');
+  const additionalToolNotes = buildAdditionalToolNotes(selectedTools);
+  const finalPrompts = additionalToolNotes
+    ? [...allPrompts, additionalToolNotes]
+    : allPrompts;
+
+  return finalPrompts.join('\n\n');
 };
 
 export const getBlockSystemPrompt = async (

packages/server/api/src/app/ai/mcp/external-tool-descriptions-data.ts

@@ -0,0 +1,7 @@
+export const MCP_TOOL_ADDITIONAL_DESCRIPTIONS: Record<string, string> = {
+  'session-sql': `## IMPORTANT: Tool Usage Note for 'session-sql'
+  
+  The 'session-sql' tool from AWS billing and cost management MCP server does NOT work with OpenOps tables.
+  
+  When the user asks about OpenOps tables, database schema, or table operations, DO NOT use 'session-sql'. This tool is only for AWS billing and cost management related SQL queries.`,
+};

packages/server/api/src/app/ai/mcp/external-tool-descriptions.ts

@@ -0,0 +1,25 @@
+import { MCP_TOOL_ADDITIONAL_DESCRIPTIONS } from './external-tool-descriptions-data';
+
+export function getToolAdditionalDescription(
+  toolName: string,
+): string | undefined {
+  return MCP_TOOL_ADDITIONAL_DESCRIPTIONS[toolName];
+}
+
+export function getAdditionalDescriptionsForSelectedTools(
+  selectedTools: Record<string, unknown> | undefined,
+): string[] {
+  if (!selectedTools) {
+    return [];
+  }
+
+  const descriptions: string[] = [];
+  for (const toolName of Object.keys(selectedTools)) {
+    const description = getToolAdditionalDescription(toolName);
+    if (description) {
+      descriptions.push(description);
+    }
+  }
+
+  return descriptions;
+}

packages/server/api/src/app/ai/mcp/llm-query-router.ts

@@ -7,6 +7,7 @@ import { projectService } from '../../project/project-service';
 import { getChatTools } from '../chat/ai-chat.service';
 import { buildUIContextSection } from '../chat/prompts.service';
 import { getAdditionalQueryClassificationDescriptions } from './extensions';
+import { getToolAdditionalDescription } from './external-tool-descriptions';
 import { sanitizeMessages } from './tool-utils';
 import { QueryClassification } from './types';
 
@@ -202,6 +203,20 @@ const getSystemPrompt = async (
   const toolsMessage = toolList
     .map((t) => `- ${t.name}: ${t.description}`)
     .join('\n');
+  const additionalToolNotes: string[] = [];
+  for (const tool of toolList) {
+    const additionalDescription = getToolAdditionalDescription(tool.name);
+    if (additionalDescription) {
+      additionalToolNotes.push(additionalDescription);
+    }
+  }
+
+  const additionalToolNotesSection =
+    additionalToolNotes.length > 0
+      ? `\n\n## IMPORTANT TOOL USAGE NOTES:\n\n${additionalToolNotes.join(
+          '\n\n',
+        )}\n`
+      : '';
   return (
     "Given the following conversation history and the list of available tools, select the tools that are most relevant to answer the user's request. " +
     `IMPORTANT: Tables tools should always be included in the output if the user asks a question involving those table names: ${openopsTablesNames.join(
@@ -211,7 +226,7 @@ const getSystemPrompt = async (
     'Include ALL relevant categories that apply. ' +
     `${
       uiContext ? `${await buildUIContextSection(uiContext)}\n` : ''
-    } Tools: ${toolsMessage}`
+    } Tools: ${toolsMessage}${additionalToolNotesSection}`
   );
 };