|
| 1 | +import { publisher } from '../../constants' |
| 2 | + |
| 3 | +import type { |
| 4 | + AgentStepContext, |
| 5 | + StepText, |
| 6 | + ToolCall, |
| 7 | +} from '../../types/agent-definition' |
| 8 | +import type { SecretAgentDefinition } from '../../types/secret-agent-definition' |
| 9 | + |
| 10 | +/** |
| 11 | + * Creates a multi-prompt editor agent that spawns one implementor per prompt. |
| 12 | + * Each prompt specifies a slightly different implementation strategy/approach. |
| 13 | + */ |
| 14 | +export function createMultiPromptEditor(): Omit<SecretAgentDefinition, 'id'> { |
| 15 | + return { |
| 16 | + publisher, |
| 17 | + model: 'anthropic/claude-opus-4.5', |
| 18 | + displayName: 'Multi-Prompt Editor', |
| 19 | + spawnerPrompt: |
| 20 | + 'Edits code by spawning multiple implementor agents with different strategy prompts, selects the best implementation, and applies the changes. Pass an array of short prompts specifying different implementation approaches. Make sure to read any files intended to be edited before spawning this agent.', |
| 21 | + |
| 22 | + includeMessageHistory: true, |
| 23 | + inheritParentSystemPrompt: true, |
| 24 | + |
| 25 | + toolNames: [ |
| 26 | + 'spawn_agents', |
| 27 | + 'str_replace', |
| 28 | + 'write_file', |
| 29 | + 'set_messages', |
| 30 | + 'set_output', |
| 31 | + ], |
| 32 | + spawnableAgents: ['best-of-n-selector-opus', 'editor-implementor-opus'], |
| 33 | + |
| 34 | + inputSchema: { |
| 35 | + params: { |
| 36 | + type: 'object', |
| 37 | + properties: { |
| 38 | + prompts: { |
| 39 | + type: 'array', |
| 40 | + items: { type: 'string' }, |
| 41 | + description: |
| 42 | + 'Array of short prompts, each specifying a slightly different implementation strategy or approach. Example: ["use a cache for the data", "don\t cache anything", "make the minimal possible changes", "modularize your solution by creating new files"]', |
| 43 | + }, |
| 44 | + }, |
| 45 | + required: ['prompts'], |
| 46 | + }, |
| 47 | + }, |
| 48 | + outputMode: 'structured_output', |
| 49 | + |
| 50 | + handleSteps: handleStepsMultiPrompt, |
| 51 | + } |
| 52 | +} |
| 53 | + |
| 54 | +function* handleStepsMultiPrompt({ |
| 55 | + agentState, |
| 56 | + params, |
| 57 | + logger, |
| 58 | +}: AgentStepContext): ReturnType< |
| 59 | + NonNullable<SecretAgentDefinition['handleSteps']> |
| 60 | +> { |
| 61 | + const prompts = (params?.prompts as string[] | undefined) ?? [] |
| 62 | + |
| 63 | + if (prompts.length === 0) { |
| 64 | + yield { |
| 65 | + toolName: 'set_output', |
| 66 | + input: { |
| 67 | + error: 'No prompts provided. Please pass an array of strategy prompts.', |
| 68 | + }, |
| 69 | + } satisfies ToolCall<'set_output'> |
| 70 | + return |
| 71 | + } |
| 72 | + |
| 73 | + // Only keep messages up to just before the last user role message (skips input prompt, instructions prompt). |
| 74 | + const { messageHistory: initialMessageHistory } = agentState |
| 75 | + let userMessageIndex = initialMessageHistory.length |
| 76 | + |
| 77 | + while (userMessageIndex > 0) { |
| 78 | + const message = initialMessageHistory[userMessageIndex - 1] |
| 79 | + if (message.role === 'user') { |
| 80 | + userMessageIndex-- |
| 81 | + } else { |
| 82 | + break |
| 83 | + } |
| 84 | + } |
| 85 | + const updatedMessageHistory = initialMessageHistory.slice(0, userMessageIndex) |
| 86 | + yield { |
| 87 | + toolName: 'set_messages', |
| 88 | + input: { |
| 89 | + messages: updatedMessageHistory, |
| 90 | + }, |
| 91 | + includeToolCall: false, |
| 92 | + } satisfies ToolCall<'set_messages'> |
| 93 | + |
| 94 | + // Spawn one opus implementor per prompt |
| 95 | + const implementorAgents = prompts.map((prompt) => ({ |
| 96 | + agent_type: 'editor-implementor-opus', |
| 97 | + prompt: `Strategy: ${prompt}`, |
| 98 | + })) |
| 99 | + |
| 100 | + // Spawn all implementor agents |
| 101 | + const { toolResult: implementorResults } = yield { |
| 102 | + toolName: 'spawn_agents', |
| 103 | + input: { |
| 104 | + agents: implementorAgents, |
| 105 | + }, |
| 106 | + includeToolCall: false, |
| 107 | + } satisfies ToolCall<'spawn_agents'> |
| 108 | + |
| 109 | + // Extract spawn results |
| 110 | + const spawnedImplementations = extractSpawnResults( |
| 111 | + implementorResults, |
| 112 | + ) as any[] |
| 113 | + |
| 114 | + logger.info( |
| 115 | + { implementorResults, spawnedImplementations, prompts }, |
| 116 | + 'spawnedImplementations', |
| 117 | + ) |
| 118 | + |
| 119 | + // Extract all the implementations from the results |
| 120 | + const letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' |
| 121 | + const implementations = spawnedImplementations.map((result, index) => ({ |
| 122 | + id: letters[index], |
| 123 | + strategy: prompts[index], |
| 124 | + content: |
| 125 | + 'errorMessage' in result |
| 126 | + ? `Error: ${result.errorMessage}` |
| 127 | + : extractLastMessageText(result) ?? '', |
| 128 | + })) |
| 129 | + |
| 130 | + // Spawn selector with implementations as params |
| 131 | + const { toolResult: selectorResult, agentState: selectorAgentState } = yield { |
| 132 | + toolName: 'spawn_agents', |
| 133 | + input: { |
| 134 | + agents: [ |
| 135 | + { |
| 136 | + agent_type: 'best-of-n-selector-opus', |
| 137 | + params: { implementations }, |
| 138 | + }, |
| 139 | + ], |
| 140 | + }, |
| 141 | + includeToolCall: false, |
| 142 | + } satisfies ToolCall<'spawn_agents'> |
| 143 | + |
| 144 | + const selectorOutput = extractSpawnResults<{ |
| 145 | + value: { |
| 146 | + implementationId: string |
| 147 | + reasoning: string |
| 148 | + } |
| 149 | + }>(selectorResult)[0] |
| 150 | + |
| 151 | + if ('errorMessage' in selectorOutput) { |
| 152 | + yield { |
| 153 | + toolName: 'set_output', |
| 154 | + input: { error: selectorOutput.errorMessage }, |
| 155 | + } satisfies ToolCall<'set_output'> |
| 156 | + return |
| 157 | + } |
| 158 | + const { implementationId } = selectorOutput.value |
| 159 | + const chosenImplementation = implementations.find( |
| 160 | + (implementation) => implementation.id === implementationId, |
| 161 | + ) |
| 162 | + if (!chosenImplementation) { |
| 163 | + yield { |
| 164 | + toolName: 'set_output', |
| 165 | + input: { error: 'Failed to find chosen implementation.' }, |
| 166 | + } satisfies ToolCall<'set_output'> |
| 167 | + return |
| 168 | + } |
| 169 | + |
| 170 | + const numMessagesBeforeStepText = selectorAgentState.messageHistory.length |
| 171 | + |
| 172 | + const { agentState: postEditsAgentState } = yield { |
| 173 | + type: 'STEP_TEXT', |
| 174 | + text: chosenImplementation.content, |
| 175 | + } as StepText |
| 176 | + const { messageHistory } = postEditsAgentState |
| 177 | + |
| 178 | + // Set output with the messages from running the step text of the chosen implementation |
| 179 | + yield { |
| 180 | + toolName: 'set_output', |
| 181 | + input: { |
| 182 | + chosenStrategy: chosenImplementation.strategy, |
| 183 | + messages: messageHistory.slice(numMessagesBeforeStepText), |
| 184 | + }, |
| 185 | + includeToolCall: false, |
| 186 | + } satisfies ToolCall<'set_output'> |
| 187 | + |
| 188 | + /** |
| 189 | + * Extracts the array of subagent results from spawn_agents tool output. |
| 190 | + */ |
| 191 | + function extractSpawnResults<T>(results: any[] | undefined): T[] { |
| 192 | + if (!results || results.length === 0) return [] |
| 193 | + |
| 194 | + const jsonResult = results.find((r) => r.type === 'json') |
| 195 | + if (!jsonResult?.value) return [] |
| 196 | + |
| 197 | + const spawnedResults = Array.isArray(jsonResult.value) |
| 198 | + ? jsonResult.value |
| 199 | + : [jsonResult.value] |
| 200 | + |
| 201 | + return spawnedResults.map((result: any) => result?.value).filter(Boolean) |
| 202 | + } |
| 203 | + |
| 204 | + /** |
| 205 | + * Extracts the text content from a 'lastMessage' AgentOutput. |
| 206 | + */ |
| 207 | + function extractLastMessageText(agentOutput: any): string | null { |
| 208 | + if (!agentOutput) return null |
| 209 | + |
| 210 | + if ( |
| 211 | + agentOutput.type === 'lastMessage' && |
| 212 | + Array.isArray(agentOutput.value) |
| 213 | + ) { |
| 214 | + for (let i = agentOutput.value.length - 1; i >= 0; i--) { |
| 215 | + const message = agentOutput.value[i] |
| 216 | + if (message.role === 'assistant' && Array.isArray(message.content)) { |
| 217 | + for (const part of message.content) { |
| 218 | + if (part.type === 'text' && typeof part.text === 'string') { |
| 219 | + return part.text |
| 220 | + } |
| 221 | + } |
| 222 | + } |
| 223 | + } |
| 224 | + } |
| 225 | + return null |
| 226 | + } |
| 227 | +} |
| 228 | + |
| 229 | +const definition = { |
| 230 | + ...createMultiPromptEditor(), |
| 231 | + id: 'editor-multi-prompt', |
| 232 | +} |
| 233 | +export default definition |
0 commit comments