konnostool-sanitizer.ts53 lines · main
| 1 | import type { ToolUIPart, UIMessage } from 'ai' |
| 2 | |
| 3 | import type { ToolName } from '../tool-filter' |
| 4 | import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' |
| 5 | |
| 6 | interface ToolSanitizer { |
| 7 | toolName: ToolName |
| 8 | sanitize: <Tool extends ToolUIPart>(tool: Tool, optInLevel: AiOptInLevel) => Tool |
| 9 | } |
| 10 | |
| 11 | export const NO_DATA_PERMISSIONS = |
| 12 | 'The query was executed and the user has viewed the results but decided not to share in the conversation due to permission levels. Continue with your plan unless instructed to interpret the result.' |
| 13 | |
| 14 | const executeSqlSanitizer: ToolSanitizer = { |
| 15 | toolName: 'execute_sql', |
| 16 | sanitize: (tool, optInLevel) => { |
| 17 | const output = tool.output |
| 18 | let sanitizedOutput: unknown |
| 19 | |
| 20 | if (optInLevel !== 'schema_and_log_and_data') { |
| 21 | if (Array.isArray(output)) { |
| 22 | sanitizedOutput = NO_DATA_PERMISSIONS |
| 23 | } |
| 24 | } else { |
| 25 | sanitizedOutput = output |
| 26 | } |
| 27 | |
| 28 | return { |
| 29 | ...tool, |
| 30 | output: sanitizedOutput, |
| 31 | } |
| 32 | }, |
| 33 | } |
| 34 | |
| 35 | export const ALL_TOOL_SANITIZERS = { |
| 36 | [executeSqlSanitizer.toolName]: executeSqlSanitizer, |
| 37 | } |
| 38 | |
| 39 | export function sanitizeMessagePart( |
| 40 | part: UIMessage['parts'][number], |
| 41 | optInLevel: AiOptInLevel |
| 42 | ): UIMessage['parts'][number] { |
| 43 | if (part.type.startsWith('tool-')) { |
| 44 | const toolPart = part as ToolUIPart |
| 45 | const toolName = toolPart.type.slice('tool-'.length) |
| 46 | const sanitizer = ALL_TOOL_SANITIZERS[toolName] |
| 47 | if (sanitizer) { |
| 48 | return sanitizer.sanitize(toolPart, optInLevel) |
| 49 | } |
| 50 | } |
| 51 | |
| 52 | return part |
| 53 | } |