From bbe95f1eaa8f5351c58e0866ba938415db7891e4 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Wed, 23 Jul 2025 16:11:23 -0400 Subject: feat(commands): Implement argument handling for custom commands via a prompt pipeline (#4702) --- .../prompt-processors/argumentProcessor.test.ts | 99 ++++++++++++++++++++++ .../prompt-processors/argumentProcessor.ts | 34 ++++++++ .../cli/src/services/prompt-processors/types.ts | 37 ++++++++ 3 files changed, 170 insertions(+) create mode 100644 packages/cli/src/services/prompt-processors/argumentProcessor.test.ts create mode 100644 packages/cli/src/services/prompt-processors/argumentProcessor.ts create mode 100644 packages/cli/src/services/prompt-processors/types.ts (limited to 'packages/cli/src/services/prompt-processors') diff --git a/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts b/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts new file mode 100644 index 00000000..6af578a9 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts @@ -0,0 +1,99 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + DefaultArgumentProcessor, + ShorthandArgumentProcessor, +} from './argumentProcessor.js'; +import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; + +describe('Argument Processors', () => { + describe('ShorthandArgumentProcessor', () => { + const processor = new ShorthandArgumentProcessor(); + + it('should replace a single {{args}} instance', async () => { + const prompt = 'Refactor the following code: {{args}}'; + const context = createMockCommandContext({ + invocation: { + raw: '/refactor make it faster', + name: 'refactor', + args: 'make it faster', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('Refactor the following code: make it faster'); + }); + + it('should replace multiple {{args}} instances', async () => { + const prompt = 'User said: {{args}}. I repeat: {{args}}!'; + const context = createMockCommandContext({ + invocation: { + raw: '/repeat hello world', + name: 'repeat', + args: 'hello world', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('User said: hello world. I repeat: hello world!'); + }); + + it('should handle an empty args string', async () => { + const prompt = 'The user provided no input: {{args}}.'; + const context = createMockCommandContext({ + invocation: { + raw: '/input', + name: 'input', + args: '', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('The user provided no input: .'); + }); + + it('should not change the prompt if {{args}} is not present', async () => { + const prompt = 'This is a static prompt.'; + const context = createMockCommandContext({ + invocation: { + raw: '/static some arguments', + name: 'static', + args: 'some arguments', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('This is a static prompt.'); + }); + }); + + describe('DefaultArgumentProcessor', () => { + const processor = new DefaultArgumentProcessor(); + + it('should append the full command if args are provided', async () => { + const prompt = 'Parse the command.'; + const context = createMockCommandContext({ + invocation: { + raw: '/mycommand arg1 "arg two"', + name: 'mycommand', + args: 'arg1 "arg two"', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('Parse the command.\n\n/mycommand arg1 "arg two"'); + }); + + it('should NOT append the full command if no args are provided', async () => { + const prompt = 'Parse the command.'; + const context = createMockCommandContext({ + invocation: { + raw: '/mycommand', + name: 'mycommand', + args: '', + }, + }); + const result = await processor.process(prompt, context); + expect(result).toBe('Parse the command.'); + }); + }); +}); diff --git a/packages/cli/src/services/prompt-processors/argumentProcessor.ts b/packages/cli/src/services/prompt-processors/argumentProcessor.ts new file mode 100644 index 00000000..a7efeea9 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/argumentProcessor.ts @@ -0,0 +1,34 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { IPromptProcessor, SHORTHAND_ARGS_PLACEHOLDER } from './types.js'; +import { CommandContext } from '../../ui/commands/types.js'; + +/** + * Replaces all instances of `{{args}}` in a prompt with the user-provided + * argument string. + */ +export class ShorthandArgumentProcessor implements IPromptProcessor { + async process(prompt: string, context: CommandContext): Promise { + return prompt.replaceAll( + SHORTHAND_ARGS_PLACEHOLDER, + context.invocation!.args, + ); + } +} + +/** + * Appends the user's full command invocation to the prompt if arguments are + * provided, allowing the model to perform its own argument parsing. + */ +export class DefaultArgumentProcessor implements IPromptProcessor { + async process(prompt: string, context: CommandContext): Promise { + if (context.invocation!.args) { + return `${prompt}\n\n${context.invocation!.raw}`; + } + return prompt; + } +} diff --git a/packages/cli/src/services/prompt-processors/types.ts b/packages/cli/src/services/prompt-processors/types.ts new file mode 100644 index 00000000..2ca61062 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/types.ts @@ -0,0 +1,37 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { CommandContext } from '../../ui/commands/types.js'; + +/** + * Defines the interface for a prompt processor, a module that can transform + * a prompt string before it is sent to the model. Processors are chained + * together to create a processing pipeline. + */ +export interface IPromptProcessor { + /** + * Processes a prompt string, applying a specific transformation as part of a pipeline. + * + * Each processor in a command's pipeline receives the output of the previous + * processor. This method provides the full command context, allowing for + * complex transformations that may require access to invocation details, + * application services, or UI state. + * + * @param prompt The current state of the prompt string. This may have been + * modified by previous processors in the pipeline. + * @param context The full command context, providing access to invocation + * details (like `context.invocation.raw` and `context.invocation.args`), + * application services, and UI handlers. + * @returns A promise that resolves to the transformed prompt string, which + * will be passed to the next processor or, if it's the last one, sent to the model. + */ + process(prompt: string, context: CommandContext): Promise; +} + +/** + * The placeholder string for shorthand argument injection in custom commands. + */ +export const SHORTHAND_ARGS_PLACEHOLDER = '{{args}}'; -- cgit v1.2.3