Repository navigation
feat(cli): add /enhance command to improve user prompts #25135
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 1 commit
04f0545
48d7a72
03da89e
30b2af3
8a8590c
fcb859d
b6f9f07
652c1e2
7c3b3ce
de8c2d8
2128019
c741ce3
385fbbd
bbdfa50
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,11 +1,16 @@ | ||
| { | ||
| "experimental": { | ||
| "plan": true, | ||
| "extensionReloading": true, | ||
| "modelSteering": true, | ||
| "memoryManager": true | ||
| }, | ||
| "general": { | ||
| "devtools": true | ||
| "devtools": true, | ||
| "plan": { | ||
| "enabled": true | ||
| } | ||
| }, | ||
| "agents": { | ||
| "overrides": {} | ||
| } | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,175 @@ | ||
| /** | ||
| * @license | ||
| * Copyright 2026 Google LLC | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| import { describe, it, expect, beforeEach, vi, type Mock } from 'vitest'; | ||
| import { enhanceCommand } from './enhanceCommand.js'; | ||
| import { type CommandContext } from './types.js'; | ||
| import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; | ||
| import { MessageType } from '../types.js'; | ||
| import { LlmRole } from '@google/gemini-cli-core'; | ||
|
|
||
| describe('enhanceCommand', () => { | ||
| let mockContext: CommandContext; | ||
| let mockGenerateContent: Mock; | ||
|
|
||
| beforeEach(() => { | ||
| mockGenerateContent = vi.fn(); | ||
| mockContext = createMockCommandContext({ | ||
| services: { | ||
| agentContext: { | ||
| promptId: 'test-prompt-id', | ||
| geminiClient: { | ||
| getHistory: vi.fn().mockReturnValue([ | ||
| { role: 'user', parts: [{ text: 'previous user msg' }] }, | ||
| { role: 'model', parts: [{ text: 'previous model msg' }] }, | ||
| ]), | ||
| }, | ||
| config: { | ||
| getModel: vi.fn().mockReturnValue('test-model'), | ||
| getContentGenerator: vi.fn().mockReturnValue({ | ||
| generateContent: mockGenerateContent, | ||
| }), | ||
| getGemini31LaunchedSync: vi.fn().mockReturnValue(true), | ||
| getHasAccessToPreviewModel: vi.fn().mockReturnValue(true), | ||
| }, | ||
| }, | ||
| }, | ||
| ui: { | ||
| addItem: vi.fn(), | ||
| setDebugMessage: vi.fn(), | ||
| }, | ||
| } as unknown as CommandContext); | ||
| }); | ||
|
|
||
| it('should have the correct name and description', () => { | ||
| expect(enhanceCommand.name).toBe('enhance'); | ||
| expect(enhanceCommand.description).toBe( | ||
| 'Enhance a prompt with additional context and rephrasing', | ||
| ); | ||
| }); | ||
|
|
||
| it('should show error if no prompt is provided', async () => { | ||
| if (!enhanceCommand.action) throw new Error('Action must be defined'); | ||
|
|
||
| await enhanceCommand.action(mockContext, ''); | ||
|
|
||
| expect(mockContext.ui.addItem).toHaveBeenCalledWith( | ||
| expect.objectContaining({ | ||
| type: MessageType.ERROR, | ||
| text: expect.stringContaining('Please provide a prompt'), | ||
| }), | ||
| ); | ||
| }); | ||
|
|
||
| it('should call generateContent with correct parameters and show enhanced prompt', async () => { | ||
| if (!enhanceCommand.action) throw new Error('Action must be defined'); | ||
|
|
||
| mockGenerateContent.mockResolvedValue({ | ||
| candidates: [ | ||
| { | ||
| content: { | ||
| parts: [{ text: 'Enhanced: do something' }], | ||
| }, | ||
| }, | ||
| ], | ||
| }); | ||
|
|
||
| await enhanceCommand.action(mockContext, 'do something'); | ||
|
|
||
| expect(mockGenerateContent).toHaveBeenCalledWith( | ||
| expect.objectContaining({ | ||
| model: 'test-model', | ||
| contents: [ | ||
| { role: 'user', parts: [{ text: 'previous user msg' }] }, | ||
| { role: 'model', parts: [{ text: 'previous model msg' }] }, | ||
| { role: 'user', parts: [{ text: 'do something' }] }, | ||
| ], | ||
| config: { | ||
| systemInstruction: { | ||
| role: 'system', | ||
| parts: [ | ||
| { | ||
| text: expect.stringContaining( | ||
| "Generate an enhanced version of the user's prompt", | ||
| ), | ||
| }, | ||
| ], | ||
| }, | ||
| }, | ||
| }), | ||
| 'test-prompt-id', | ||
| LlmRole.UTILITY_TOOL, | ||
| ); | ||
|
|
||
| expect(mockContext.ui.addItem).toHaveBeenCalledWith( | ||
| expect.objectContaining({ | ||
| type: MessageType.INFO, | ||
| text: expect.stringContaining( | ||
| 'Enhanced prompt:\n\nEnhanced: do something', | ||
| ), | ||
| }), | ||
| ); | ||
| expect(mockContext.ui.setInput).toHaveBeenCalledWith( | ||
| 'Enhanced: do something', | ||
| ); | ||
| }); | ||
|
|
||
| it('should clean the response from markdown and quotes', async () => { | ||
| if (!enhanceCommand.action) throw new Error('Action must be defined'); | ||
|
|
||
| mockGenerateContent.mockResolvedValue({ | ||
| candidates: [ | ||
| { | ||
| content: { | ||
| parts: [{ text: '```markdown\n"Clean me"\n```' }], | ||
| }, | ||
| }, | ||
| ], | ||
| }); | ||
|
|
||
| await enhanceCommand.action(mockContext, 'dirty prompt'); | ||
|
|
||
| expect(mockContext.ui.addItem).toHaveBeenCalledWith( | ||
| expect.objectContaining({ | ||
| type: MessageType.INFO, | ||
| text: expect.stringContaining('Enhanced prompt:\n\nClean me'), | ||
| }), | ||
| ); | ||
| expect(mockContext.ui.setInput).toHaveBeenCalledWith('Clean me'); | ||
| }); | ||
|
|
||
| it('should handle API errors gracefully', async () => { | ||
| if (!enhanceCommand.action) throw new Error('Action must be defined'); | ||
|
|
||
| mockGenerateContent.mockRejectedValue(new Error('API Error')); | ||
|
|
||
| await enhanceCommand.action(mockContext, 'test prompt'); | ||
|
|
||
| expect(mockContext.ui.addItem).toHaveBeenCalledWith( | ||
| expect.objectContaining({ | ||
| type: MessageType.ERROR, | ||
| text: expect.stringContaining('Failed to enhance prompt: API Error'), | ||
| }), | ||
| ); | ||
| }); | ||
|
|
||
| it('should handle empty response from model', async () => { | ||
| if (!enhanceCommand.action) throw new Error('Action must be defined'); | ||
|
|
||
| mockGenerateContent.mockResolvedValue({ | ||
| candidates: [], | ||
| }); | ||
|
|
||
| await enhanceCommand.action(mockContext, 'test prompt'); | ||
|
|
||
| expect(mockContext.ui.addItem).toHaveBeenCalledWith( | ||
| expect.objectContaining({ | ||
| type: MessageType.ERROR, | ||
| text: expect.stringContaining('Empty response from model'), | ||
| }), | ||
| ); | ||
| }); | ||
| }); |
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
| @@ -0,0 +1,104 @@ | ||||||
| /** | ||||||
| * @license | ||||||
| * Copyright 2026 Google LLC | ||||||
| * SPDX-License-Identifier: Apache-2.0 | ||||||
| */ | ||||||
|
|
||||||
| import { CommandKind, type SlashCommand } from './types.js'; | ||||||
| import { MessageType } from '../types.js'; | ||||||
| import { LlmRole, resolveModel } from '@google/gemini-cli-core'; | ||||||
| import type { Content } from '@google/genai'; | ||||||
|
|
||||||
| const INSTRUCTION = | ||||||
| "Generate an enhanced version of the user's prompt, using the preceding conversation as context. Reply with ONLY the enhanced prompt - no conversation, explanations, lead-in, bullet points, placeholders, or surrounding quotes."; | ||||||
|
|
||||||
| function clean(text: string) { | ||||||
| const stripped = text.replace(/^```\w*\n?|```$/g, '').trim(); | ||||||
| return stripped.replace(/^(['"])([\s\S]*)\1$/, '$2').trim(); | ||||||
| } | ||||||
|
|
||||||
| export const enhanceCommand: SlashCommand = { | ||||||
| name: 'enhance', | ||||||
| description: 'Enhance a prompt with additional context and rephrasing', | ||||||
| kind: CommandKind.BUILT_IN, | ||||||
| autoExecute: true, | ||||||
| action: async (context, args) => { | ||||||
| const draft = args.trim(); | ||||||
| if (!draft) { | ||||||
| context.ui.addItem({ | ||||||
| type: MessageType.ERROR, | ||||||
| text: 'Please provide a prompt to enhance. Usage: /enhance <prompt>', | ||||||
| }); | ||||||
| return; | ||||||
| } | ||||||
|
|
||||||
| const agentContext = context.services.agentContext; | ||||||
| if (!agentContext) { | ||||||
| context.ui.addItem({ | ||||||
| type: MessageType.ERROR, | ||||||
| text: 'Agent context not available.', | ||||||
| }); | ||||||
| return; | ||||||
| } | ||||||
|
|
||||||
| const config = agentContext.config; | ||||||
| const contentGenerator = config.getContentGenerator(); | ||||||
| const promptId = agentContext.promptId; | ||||||
|
|
||||||
| const model = resolveModel( | ||||||
| config.getModel(), | ||||||
| config.getGemini31LaunchedSync?.() ?? false, | ||||||
| false, | ||||||
| config.getHasAccessToPreviewModel?.() ?? true, | ||||||
| config, | ||||||
| ); | ||||||
|
|
||||||
| context.ui.setDebugMessage('Enhancing prompt...'); | ||||||
|
|
||||||
| try { | ||||||
| const history = agentContext.geminiClient?.getHistory() ?? []; | ||||||
| const contents: Content[] = [ | ||||||
| ...history, | ||||||
| { role: 'user', parts: [{ text: draft }] }, | ||||||
| ]; | ||||||
|
|
||||||
| const response = await contentGenerator.generateContent( | ||||||
| { | ||||||
| model, | ||||||
| contents, | ||||||
| config: { | ||||||
| systemInstruction: { | ||||||
| role: 'system', | ||||||
| parts: [{ text: INSTRUCTION }], | ||||||
| }, | ||||||
| }, | ||||||
| }, | ||||||
| promptId, | ||||||
| LlmRole.UTILITY_TOOL, | ||||||
| ); | ||||||
|
|
||||||
| const enhancedText = response.candidates?.[0]?.content?.parts?.[0]?.text; | ||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The current implementation assumes the enhanced prompt is always in the first part of the first candidate. Gemini 2.0 models often include a 'thought' part before the actual response text. It is safer to find the first text part that isn't a thought. Additionally, to prevent prompt injection, sanitize the output by removing newlines and context-breaking characters like ']' before it is used in further prompts, as per repository rules.
Suggested change
References
|
||||||
|
|
||||||
| if (enhancedText) { | ||||||
| const cleanedText = clean(enhancedText); | ||||||
| context.ui.addItem({ | ||||||
| type: MessageType.INFO, | ||||||
| text: `Enhanced prompt:\n\n${cleanedText}`, | ||||||
| }); | ||||||
| context.ui.setInput(cleanedText); | ||||||
| } else { | ||||||
| context.ui.addItem({ | ||||||
| type: MessageType.ERROR, | ||||||
| text: 'Failed to enhance prompt: Empty response from model.', | ||||||
| }); | ||||||
| } | ||||||
| } catch (error) { | ||||||
| context.ui.addItem({ | ||||||
| type: MessageType.ERROR, | ||||||
| text: `Failed to enhance prompt: ${error instanceof Error ? error.message : String(error)}`, | ||||||
| }); | ||||||
| } finally { | ||||||
| context.ui.setDebugMessage(''); | ||||||
| } | ||||||
| }, | ||||||
| }; | ||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Since the
/enhancecommand requires a prompt argument to function correctly,autoExecuteshould be set tofalse. WhenautoExecuteistrue, selecting the command from the autocomplete list (e.g., by pressing Enter) will trigger it immediately with empty arguments, resulting in an error message. Setting it tofalseallows the user to complete the command name and then type the prompt they wish to enhance.