Improve tool selection with additional tool descriptions (#1718)

Fixes OPS-3138.

Author: ravikiranvm

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

@@ -3,10 +3,16 @@ import { ChatFlowContext, CODE_BLOCK_NAME, isNil } from '@openops/shared';
 import { ToolSet } from 'ai';
 import { readFile } from 'fs/promises';
 import { join } from 'path';
+import { getAdditionalToolDescriptions } 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 {
+  const descriptions = getAdditionalToolDescriptions(selectedTools);
+  return descriptions.length > 0 ? descriptions.join('\n\n') : '';
+}
+
 export const getMcpSystemPrompt = async ({
   queryClassification,
   selectedTools,
@@ -58,7 +64,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,15 @@
+export type ToolDescription = {
+  note: string;
+  mcpServer?: string;
+};
+
+export const MCP_TOOL_ADDITIONAL_DESCRIPTIONS: Record<string, ToolDescription> =
+  {
+    'session-sql': {
+      note: `### 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, table schema, or table operations, DO NOT use 'session-sql'. 
+- This tool is only for AWS billing and cost management related SQL queries.`,
+      mcpServer: 'aws-billing-cost-management',
+    },
+  };

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

@@ -0,0 +1,22 @@
+import { ToolSet } from 'ai';
+import { MCP_TOOL_ADDITIONAL_DESCRIPTIONS } from './external-tool-descriptions-data';
+
+export function getAdditionalToolDescriptions(
+  tools: ToolSet | string[] | undefined,
+): string[] {
+  if (!tools) {
+    return [];
+  }
+
+  const toolNames = Array.isArray(tools) ? tools : Object.keys(tools);
+  const descriptions: string[] = [];
+
+  for (const toolName of toolNames) {
+    const description = MCP_TOOL_ADDITIONAL_DESCRIPTIONS[toolName];
+    if (description?.note) {
+      descriptions.push(description.note);
+    }
+  }
+
+  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 { getAdditionalToolDescriptions } from './external-tool-descriptions';
 import { sanitizeMessages } from './tool-utils';
 import { QueryClassification } from './types';
 
@@ -202,6 +203,16 @@ const getSystemPrompt = async (
   const toolsMessage = toolList
     .map((t) => `- ${t.name}: ${t.description}`)
     .join('\n');
+  const additionalToolNotes = getAdditionalToolDescriptions(
+    toolList.map((t) => t.name),
+  );
+
+  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 +222,7 @@ const getSystemPrompt = async (
     'Include ALL relevant categories that apply. ' +
     `${
       uiContext ? `${await buildUIContextSection(uiContext)}\n` : ''
-    } Tools: ${toolsMessage}`
+    } Tools: ${toolsMessage}${additionalToolNotesSection}`
   );
 };