From 934b457365312d6bc0996d38eec2de1f78e4668b Mon Sep 17 00:00:00 2001 From: Ravi Kiran Date: Wed, 3 Dec 2025 15:09:48 +0530 Subject: [PATCH 1/2] Improve tool selection with additinal tool descriptions --- .../api/src/app/ai/chat/prompts.service.ts | 17 ++++++++++++- .../ai/mcp/external-tool-descriptions-data.ts | 7 ++++++ .../app/ai/mcp/external-tool-descriptions.ts | 25 +++++++++++++++++++ .../api/src/app/ai/mcp/llm-query-router.ts | 17 ++++++++++++- 4 files changed, 64 insertions(+), 2 deletions(-) create mode 100644 packages/server/api/src/app/ai/mcp/external-tool-descriptions-data.ts create mode 100644 packages/server/api/src/app/ai/mcp/external-tool-descriptions.ts diff --git a/packages/server/api/src/app/ai/chat/prompts.service.ts b/packages/server/api/src/app/ai/chat/prompts.service.ts index 07a5ca78c1..6f33308337 100644 --- a/packages/server/api/src/app/ai/chat/prompts.service.ts +++ b/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 ( diff --git a/packages/server/api/src/app/ai/mcp/external-tool-descriptions-data.ts b/packages/server/api/src/app/ai/mcp/external-tool-descriptions-data.ts new file mode 100644 index 0000000000..c9a5fdf4bf --- /dev/null +++ b/packages/server/api/src/app/ai/mcp/external-tool-descriptions-data.ts @@ -0,0 +1,7 @@ +export const MCP_TOOL_ADDITIONAL_DESCRIPTIONS: Record = { + '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.`, +}; diff --git a/packages/server/api/src/app/ai/mcp/external-tool-descriptions.ts b/packages/server/api/src/app/ai/mcp/external-tool-descriptions.ts new file mode 100644 index 0000000000..21f80ed074 --- /dev/null +++ b/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 | 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; +} diff --git a/packages/server/api/src/app/ai/mcp/llm-query-router.ts b/packages/server/api/src/app/ai/mcp/llm-query-router.ts index 7b907f7d4f..00cd6422d9 100644 --- a/packages/server/api/src/app/ai/mcp/llm-query-router.ts +++ b/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}` ); }; From 096d656cccae2a755ea9abd88cdc4ee40c222e4a Mon Sep 17 00:00:00 2001 From: Ravi Kiran Date: Thu, 4 Dec 2025 13:27:21 +0530 Subject: [PATCH 2/2] Improve utility functions and type --- .../api/src/app/ai/chat/prompts.service.ts | 8 ++----- .../ai/mcp/external-tool-descriptions-data.ts | 20 +++++++++++----- .../app/ai/mcp/external-tool-descriptions.ts | 23 ++++++++----------- .../api/src/app/ai/mcp/llm-query-router.ts | 14 ++++------- 4 files changed, 31 insertions(+), 34 deletions(-) diff --git a/packages/server/api/src/app/ai/chat/prompts.service.ts b/packages/server/api/src/app/ai/chat/prompts.service.ts index 6f33308337..8b762eb5db 100644 --- a/packages/server/api/src/app/ai/chat/prompts.service.ts +++ b/packages/server/api/src/app/ai/chat/prompts.service.ts @@ -3,17 +3,13 @@ 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 { 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 { - if (!selectedTools) { - return ''; - } - - const descriptions = getAdditionalDescriptionsForSelectedTools(selectedTools); + const descriptions = getAdditionalToolDescriptions(selectedTools); return descriptions.length > 0 ? descriptions.join('\n\n') : ''; } diff --git a/packages/server/api/src/app/ai/mcp/external-tool-descriptions-data.ts b/packages/server/api/src/app/ai/mcp/external-tool-descriptions-data.ts index c9a5fdf4bf..aad25fbd5d 100644 --- a/packages/server/api/src/app/ai/mcp/external-tool-descriptions-data.ts +++ b/packages/server/api/src/app/ai/mcp/external-tool-descriptions-data.ts @@ -1,7 +1,15 @@ -export const MCP_TOOL_ADDITIONAL_DESCRIPTIONS: Record = { - '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.`, +export type ToolDescription = { + note: string; + mcpServer?: string; }; + +export const MCP_TOOL_ADDITIONAL_DESCRIPTIONS: Record = + { + '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', + }, + }; diff --git a/packages/server/api/src/app/ai/mcp/external-tool-descriptions.ts b/packages/server/api/src/app/ai/mcp/external-tool-descriptions.ts index 21f80ed074..826f0b733f 100644 --- a/packages/server/api/src/app/ai/mcp/external-tool-descriptions.ts +++ b/packages/server/api/src/app/ai/mcp/external-tool-descriptions.ts @@ -1,23 +1,20 @@ +import { ToolSet } from 'ai'; 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 | undefined, +export function getAdditionalToolDescriptions( + tools: ToolSet | string[] | undefined, ): string[] { - if (!selectedTools) { + if (!tools) { return []; } + const toolNames = Array.isArray(tools) ? tools : Object.keys(tools); const descriptions: string[] = []; - for (const toolName of Object.keys(selectedTools)) { - const description = getToolAdditionalDescription(toolName); - if (description) { - descriptions.push(description); + + for (const toolName of toolNames) { + const description = MCP_TOOL_ADDITIONAL_DESCRIPTIONS[toolName]; + if (description?.note) { + descriptions.push(description.note); } } diff --git a/packages/server/api/src/app/ai/mcp/llm-query-router.ts b/packages/server/api/src/app/ai/mcp/llm-query-router.ts index 00cd6422d9..6853578952 100644 --- a/packages/server/api/src/app/ai/mcp/llm-query-router.ts +++ b/packages/server/api/src/app/ai/mcp/llm-query-router.ts @@ -7,7 +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 { getAdditionalToolDescriptions } from './external-tool-descriptions'; import { sanitizeMessages } from './tool-utils'; import { QueryClassification } from './types'; @@ -203,17 +203,13 @@ 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 additionalToolNotes = getAdditionalToolDescriptions( + toolList.map((t) => t.name), + ); const additionalToolNotesSection = additionalToolNotes.length > 0 - ? `\n\n## IMPORTANT TOOL USAGE NOTES:\n\n${additionalToolNotes.join( + ? `\n\n### IMPORTANT TOOL USAGE NOTES:\n\n${additionalToolNotes.join( '\n\n', )}\n` : '';