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..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,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 ( 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..aad25fbd5d --- /dev/null +++ b/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 = + { + '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 new file mode 100644 index 0000000000..826f0b733f --- /dev/null +++ b/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; +} 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..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,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}` ); };