diff --git a/packages/core/src/prompts/promptTemplating.test.ts b/packages/core/src/prompts/promptTemplating.test.ts index 6d09f604c3..da482ff103 100644 --- a/packages/core/src/prompts/promptTemplating.test.ts +++ b/packages/core/src/prompts/promptTemplating.test.ts @@ -5,7 +5,7 @@ */ import { describe, expect, it } from 'vitest'; -import { promptComponent, toPrompt, enabledWhen, section, xmlSection, each, switchOn } from './promptTemplating.js'; +import { promptComponent, promptTemplate, renderSnippet, renderTemplate, enabledWhen, section, xmlSection, each, switchOn, type Snippet, type PromptTemplateBase } from './promptTemplating.js'; // Sample prompt components. type MainPromptOptions = { isInteractive: boolean }; @@ -37,60 +37,113 @@ const emptyMarkdownPrompt = section('Rules', ''); const customHeaderMarkdownPrompt = section('Rules', 'Be helpful', { headerLevel: 3 }); -const eachPrompt = each<{ items: string[] }>('items', (item) => `Item: ${item}`); +const eachPrompt = each<{ items: string[] }, string>('items', (item: string) => `Item: ${item}`); -const switchOnPrompt = switchOn<{ mode: string }>('mode', { +const eachCustomSeparatorPrompt = each<{ items: string[] }, string>('items', (item: string) => item, ', '); + +const switchOnPrompt = switchOn<{ mode: string | number }>('mode', { fast: 'Speed is key', safe: 'Safety first', + 1: 'Mode one', +}); + +// A high-level template. +interface SystemPromptOptions { name: string } +interface SystemPromptTemplate extends PromptTemplateBase { + identity: Snippet; + instructions: Snippet; +} + +const mainTemplate = promptTemplate({ + identity: (opt: SystemPromptOptions) => `Your name is ${opt.name}`, + instructions: 'Be helpful', }); describe('promptTemplating', () => { it('should take variadic arguments', () => { - expect(toPrompt({}, variadicPrompt)).toBe('A,B,C'); + expect(renderSnippet({}, variadicPrompt)).toBe('A,B,C'); }); it('should handle nested arrays (variadic with arrays)', () => { - expect(toPrompt({}, nestedVariadicPrompt)).toBe('A,B,C,D'); + expect(renderSnippet({}, nestedVariadicPrompt)).toBe('A,B,C,D'); }); it('should handle enabledWhen', () => { - expect(toPrompt({ show: true }, enabledWhenPrompt)).toBe('Always,Sometimes'); - expect(toPrompt({ show: false }, enabledWhenPrompt)).toBe('Always,'); + expect(renderSnippet({ show: true }, enabledWhenPrompt)).toBe('Always,Sometimes'); + expect(renderSnippet({ show: false }, enabledWhenPrompt)).toBe('Always,'); }); it('should handle xmlSection', () => { - expect(toPrompt({}, xmlPrompt)).toBe('\nBe helpful\n'); + expect(renderSnippet({}, xmlPrompt)).toBe('\nBe helpful\n'); }); it('should return empty string for xmlSection if content is empty', () => { - expect(toPrompt({}, emptyXmlPrompt)).toBe(''); + expect(renderSnippet({}, emptyXmlPrompt)).toBe(''); }); it('should handle markdown section with default header', () => { - expect(toPrompt({}, markdownPrompt)).toBe('# Rules\n\nBe helpful'); + expect(renderSnippet({}, markdownPrompt)).toBe('# Rules\n\nBe helpful'); }); it('should return empty string for markdown section if content is empty', () => { - expect(toPrompt({}, emptyMarkdownPrompt)).toBe(''); + expect(renderSnippet({}, emptyMarkdownPrompt)).toBe(''); }); it('should handle markdown section with custom header level', () => { - expect(toPrompt({}, customHeaderMarkdownPrompt)).toBe('### Rules\n\nBe helpful'); + expect(renderSnippet({}, customHeaderMarkdownPrompt)).toBe('### Rules\n\nBe helpful'); }); it('should handle each', () => { const options = { items: ['A', 'B'] }; - expect(toPrompt(options, eachPrompt)).toBe('Item: A\nItem: B'); + expect(renderSnippet(options, eachPrompt)).toBe('Item: A\nItem: B'); }); it('should handle switchOn', () => { - expect(toPrompt({ mode: 'fast' }, switchOnPrompt)).toBe('Speed is key'); - expect(toPrompt({ mode: 'safe' }, switchOnPrompt)).toBe('Safety first'); - expect(toPrompt({ mode: 'unknown' }, switchOnPrompt)).toBe(''); + expect(renderSnippet({ mode: 'fast' }, switchOnPrompt)).toBe('Speed is key'); + expect(renderSnippet({ mode: 'safe' }, switchOnPrompt)).toBe('Safety first'); + expect(renderSnippet({ mode: 'unknown' }, switchOnPrompt)).toBe(''); + }); + + it('should handle switchOn with numeric values', () => { + expect(renderSnippet({ mode: 1 }, switchOnPrompt)).toBe('Mode one'); + }); + + it('should handle each with custom separator', () => { + const options = { items: ['A', 'B'] }; + expect(renderSnippet(options, eachCustomSeparatorPrompt)).toBe('A, B'); + }); + + it('should handle each with non-array value', () => { + const options = { items: 'not an array' as unknown as string[] }; + expect(renderSnippet(options, eachPrompt)).toBe(''); + }); + + it('should handle section with different header levels', () => { + const h2Section = section('Level 2', 'Content', { headerLevel: 2 }); + const h6Section = section('Level 6', 'Content', { headerLevel: 6 }); + expect(renderSnippet({}, h2Section)).toBe('## Level 2\n\nContent'); + expect(renderSnippet({}, h6Section)).toBe('###### Level 6\n\nContent'); + }); + + it('should handle toPrompt with basic types', () => { + expect(renderSnippet({}, 'just a string')).toBe('just a string'); + expect(renderSnippet({ name: 'test' }, (opt: { name: string }) => opt.name)).toBe('test'); + }); + + it('should handle toPrompt with empty array or object', () => { + expect(renderSnippet({}, [])).toBe(''); + expect(renderSnippet({}, {})).toBe(''); + }); + + it('should handle renderTemplate with a record of snippets (implements required members)', () => { + const options = { + name: 'Gemini CLI', + }; + expect(renderTemplate(options, mainTemplate)).toBe('Your name is Gemini CLI,Be helpful'); }); it('should output correctly for mainPrompt', () => { - const text = toPrompt({ isInteractive: true }, mainPrompt); + const text = renderSnippet({ isInteractive: true }, mainPrompt); expect(text).toContain('Your name is Gemini CLI and you are a coding agent'); expect(text).toContain('You are operating in an interactive session'); }); diff --git a/packages/core/src/prompts/promptTemplating.ts b/packages/core/src/prompts/promptTemplating.ts index 6f43a174bb..16fa62ef08 100644 --- a/packages/core/src/prompts/promptTemplating.ts +++ b/packages/core/src/prompts/promptTemplating.ts @@ -6,19 +6,35 @@ export type DynamicSnippet = (options: TOption) => string; -export type Snippet = string | DynamicSnippet | Array>; +export type Snippet = + | string + | DynamicSnippet + | Array> + | PromptTemplateBase; + +export interface PromptTemplateBase { + [key: string]: Snippet; +} + +export type PromptTemplate> = TTemplate; + +export function promptTemplate>( + template: TTemplate +): PromptTemplate { + return template; +} export function promptComponent(...snippets: Array>): Snippet { return snippets; } export function enabledWhen(optionName: keyof TOption, snippet: Snippet): Snippet { - return (option: TOption) => option[optionName] ? toPrompt(option, snippet) : ''; + return (option: TOption) => (option[optionName] ? renderSnippet(option, snippet) : ''); } export function xmlSection(name: string, snippet: Snippet): Snippet { return (options: TOption) => { - const content = toPrompt(options, snippet); + const content = renderSnippet(options, snippet); return content ? `<${name}>\n${content}\n` : ''; }; } @@ -27,24 +43,28 @@ export interface SectionOptions { headerLevel?: number; } -export function section(name: string, snippet: Snippet, sectionOptions?: SectionOptions): Snippet { +export function section( + name: string, + snippet: Snippet, + sectionOptions?: SectionOptions +): Snippet { return (options: TOption) => { - const content = toPrompt(options, snippet); + const content = renderSnippet(options, snippet); const level = sectionOptions?.headerLevel ?? 1; const hashes = '#'.repeat(level); return content ? `${hashes} ${name}\n\n${content}` : ''; }; } -export function each( +export function each( optionName: keyof TOption, - snippet: Snippet, + snippet: Snippet, separator = '\n' ): Snippet { return (options: TOption) => { const items = options[optionName]; if (Array.isArray(items)) { - return items.map((item: any) => toPrompt(item, snippet)).join(separator); + return items.map((item: TItem) => renderSnippet(item, snippet)).join(separator); } return ''; }; @@ -57,16 +77,29 @@ export function switchOn( return (options: TOption) => { const value = String(options[optionName]); const caseSnippet = cases[value]; - return caseSnippet ? toPrompt(options, caseSnippet) : ''; + return caseSnippet ? renderSnippet(options, caseSnippet) : ''; }; } -export function toPrompt(options: TOption, snippet: Snippet): string { +export function renderTemplate>( + options: TOption, + implementation: TTemplate +): string { + return Object.values(implementation) + .map(eachSnippet => renderSnippet(options, eachSnippet)) + .join(); +} + +export function renderSnippet(options: TOption, snippet: Snippet): string { if (typeof snippet === 'string') { return snippet; } else if (Array.isArray(snippet)) { - return snippet.map(eachSnippet => toPrompt(options, eachSnippet)).join(); - } else { + return snippet.map(eachSnippet => renderSnippet(options, eachSnippet)).join(); + } else if (typeof snippet === 'function') { return snippet(options); + } else { + return Object.values(snippet) + .map(eachSnippet => renderSnippet(options, eachSnippet)) + .join(); } }