diff --git a/frontend/src/lib/components/flows/agentToolUtils.test.ts b/frontend/src/lib/components/flows/agentToolUtils.test.ts new file mode 100644 index 0000000000..a88eb2833a --- /dev/null +++ b/frontend/src/lib/components/flows/agentToolUtils.test.ts @@ -0,0 +1,68 @@ +import { describe, expect, it, vi } from 'vitest' + +vi.mock('../aiProviderStorage', () => ({ + loadStoredConfig: () => undefined +})) + +vi.mock('./flowInfers', () => ({ + AI_AGENT_SCHEMA: { properties: {} } +})) + +import type { FlowModule } from '$lib/gen' +import { removeAgentToolByIdDeep } from './agentToolUtils' + +function makeRawModule(id: string): FlowModule { + return { + id, + summary: id, + value: { type: 'rawscript', content: '', language: 'python3' } as any + } as FlowModule +} + +function makeAiAgent(id: string, tools: any[]): FlowModule { + return { + id, + summary: id, + value: { + type: 'aiagent', + tools, + input_transforms: {} + } as any + } as FlowModule +} + +function makeFlowModuleTool(module: FlowModule) { + return { + id: module.id, + summary: module.summary, + value: { + tool_type: 'flowmodule', + ...module.value + } + } +} + +describe('removeAgentToolByIdDeep', () => { + it('removes a direct tool from an ai agent', () => { + const tool = makeFlowModuleTool(makeRawModule('lookup_user')) + const agent = makeAiAgent('agent', [tool]) + const removed: string[] = [] + + expect(removeAgentToolByIdDeep([agent], 'lookup_user', (x) => removed.push(x.id))).toBe(true) + expect((agent.value as any).tools).toEqual([]) + expect(removed).toEqual(['lookup_user']) + }) + + it('removes a nested tool from a nested ai agent tool', () => { + const nestedTool = makeFlowModuleTool(makeRawModule('create_ticket')) + const nestedAgent = makeAiAgent('support_agent', [nestedTool]) + const rootAgent = makeAiAgent('root_agent', [makeFlowModuleTool(nestedAgent)]) + const removed: string[] = [] + + expect(removeAgentToolByIdDeep([rootAgent], 'create_ticket', (x) => removed.push(x.id))).toBe( + true + ) + expect((((rootAgent.value as any).tools as any[])[0].value as any).tools).toEqual([]) + expect(removed).toEqual(['create_ticket']) + }) +}) diff --git a/frontend/src/lib/components/flows/agentToolUtils.ts b/frontend/src/lib/components/flows/agentToolUtils.ts index 84e102ce56..104c36bafc 100644 --- a/frontend/src/lib/components/flows/agentToolUtils.ts +++ b/frontend/src/lib/components/flows/agentToolUtils.ts @@ -108,6 +108,66 @@ export function createWebsearchTool(id: string): WebsearchTool { } } +/** + * Remove an AI agent tool by id, recursively traversing nested modules and nested AI agents. + * Returns true when a matching tool was found and removed. + */ +export function removeAgentToolByIdDeep( + modules: FlowModule[], + id: string, + onRemove?: (tool: AgentTool) => void +): boolean { + for (const module of modules) { + if (module.value.type === 'forloopflow' || module.value.type === 'whileloopflow') { + if (removeAgentToolByIdDeep(module.value.modules, id, onRemove)) { + return true + } + continue + } + + if (module.value.type === 'branchall') { + for (const branch of module.value.branches) { + if (removeAgentToolByIdDeep(branch.modules, id, onRemove)) { + return true + } + } + continue + } + + if (module.value.type === 'branchone') { + if (removeAgentToolByIdDeep(module.value.default, id, onRemove)) { + return true + } + for (const branch of module.value.branches) { + if (removeAgentToolByIdDeep(branch.modules, id, onRemove)) { + return true + } + } + continue + } + + if (module.value.type !== 'aiagent') { + continue + } + + const toolIndex = module.value.tools.findIndex((tool) => tool.id === id) + if (toolIndex !== -1) { + const [removed] = module.value.tools.splice(toolIndex, 1) + onRemove?.(removed) + return true + } + + const nestedToolModules = module.value.tools + .filter(isFlowModuleTool) + .map((tool) => agentToolToFlowModule(tool)) + if (removeAgentToolByIdDeep(nestedToolModules, id, onRemove)) { + return true + } + } + + return false +} + /** * Convert a FlowModuleTool to a FlowModule for use with loadFlowModuleState etc. * Strips the extra `tool_type` field and maps AgentTool fields to FlowModule fields. diff --git a/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte b/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte index a3f2545805..7db4d0e39a 100644 --- a/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte +++ b/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte @@ -50,7 +50,8 @@ createWebsearchTool, createAiAgentTool, SPECIAL_TOOL_KINDS, - agentToolToFlowModule + agentToolToFlowModule, + removeAgentToolByIdDeep } from '../agentToolUtils' import { loadFlowModuleState } from '../flowStateUtils.svelte' import { getNoteEditorContext } from '$lib/components/graph/noteEditor.svelte' @@ -257,7 +258,7 @@ * Tools are always leaf nodes, so we just need to delete their state directly */ function removeAgentToolById(tools: AgentTool[], id: string): AgentTool[] { - const index = tools.findIndex((tool) => tool.id == id) + const index = tools.findIndex((tool) => tool.id === id) if (index != -1) { const [removed] = tools.splice(index, 1) deleteFlowStateById(removed.id, flowStateStore) @@ -383,6 +384,8 @@ } export function deleteMultiple(ids: string[]) { + const structureIds = ids.filter((id) => findInStructure(proxy.items, id)) + const toolIds = ids.filter((id) => !findInStructure(proxy.items, id)) const deletingSet = new Set(ids) const allDeps: Record = {} for (const id of ids) { @@ -395,18 +398,30 @@ } const opts = { displayState: groupDisplayState } - const { emptiedGroups, duplicateGroups, commit } = proxy.prepareMutation((tree) => { - for (const id of ids) { - const found = findInStructure(tree, id) - if (found) found.parentChildren.splice(found.index, 1) - } - }, opts) + const { emptiedGroups, duplicateGroups, commit } = + structureIds.length > 0 + ? proxy.prepareMutation((tree) => { + for (const id of structureIds) { + const found = findInStructure(tree, id) + if (found) found.parentChildren.splice(found.index, 1) + } + }, opts) + : { + emptiedGroups: [], + duplicateGroups: [], + commit: () => {} + } const affectedGroups = [...emptiedGroups, ...duplicateGroups] const cb = () => { push(history, flowStore.val) commit({ removeDuplicates: duplicateGroups.length > 0 }) + for (const id of toolIds) { + removeAgentToolByIdDeep(flowStore.val.value.modules, id, (removed) => { + deleteFlowStateById(removed.id, flowStateStore) + }) + } for (const id of ids) { delete flowStateStore.val[id] } @@ -678,6 +693,25 @@ return } + if (!findInStructure(proxy.items, id)) { + const cb = () => { + push(history, flowStore.val) + selectNextId(id) + const removed = removeAgentToolByIdDeep(flowStore.val.value.modules, id, (tool) => { + deleteFlowStateById(tool.id, flowStateStore) + }) + if (!removed) return + refreshStateStore(flowStore) + onDelete?.(id) + } + if (Object.keys(dependents).length > 0) { + deleteCallback = cb + } else { + cb() + } + return + } + const dsOpts = { displayState: groupDisplayState } const { emptiedGroups, duplicateGroups, commit } = proxy.prepareMutation((tree) => { const found = findInStructure(tree, id)