Support prompt variants

fixed #14485 Signed-off-by: Jonas Helming <[email protected]>
eclipse-theia · Nov 19, 2024 · f3a3bf8 · f3a3bf8
1 parent ef33e14
commit f3a3bf8
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 0 deletions.
diff --git a/packages/ai-core/src/common/prompt-service.spec.ts b/packages/ai-core/src/common/prompt-service.spec.ts
@@ -95,4 +95,61 @@ describe('PromptService', () => {
             expect(prompt?.text).to.equal('Hello, John!');
         }
     });
+    it('should return all variant IDs of a given prompt', () => {
+        promptService.storePromptTemplate({ id: 'main', template: 'Main template' });
+
+        promptService.storePromptTemplate({
+            id: 'variant1',
+            template: 'Variant 1',
+            mainTemplateId: 'main'
+        });
+        promptService.storePromptTemplate({
+            id: 'variant2',
+            template: 'Variant 2',
+            mainTemplateId: 'main'
+        });
+        promptService.storePromptTemplate({
+            id: 'variant3',
+            template: 'Variant 3',
+            mainTemplateId: 'main'
+        });
+
+        const variantIds = promptService.getVariantIds('main');
+        expect(variantIds).to.deep.equal(['variant1', 'variant2', 'variant3']);
+    });
+
+    it('should return an empty array if no variants exist for a given prompt', () => {
+        promptService.storePromptTemplate({ id: 'main', template: 'Main template' });
+
+        const variantIds = promptService.getVariantIds('main');
+        expect(variantIds).to.deep.equal([]);
+    });
+
+    it('should return an empty array if the main prompt ID does not exist', () => {
+        const variantIds = promptService.getVariantIds('nonExistent');
+        expect(variantIds).to.deep.equal([]);
+    });
+
+    it('should not influence prompts without variants when other prompts have variants', () => {
+        promptService.storePromptTemplate({ id: 'mainWithVariants', template: 'Main template with variants' });
+        promptService.storePromptTemplate({ id: 'mainWithoutVariants', template: 'Main template without variants' });
+
+        promptService.storePromptTemplate({
+            id: 'variant1',
+            template: 'Variant 1',
+            mainTemplateId: 'mainWithVariants'
+        });
+        promptService.storePromptTemplate({
+            id: 'variant2',
+            template: 'Variant 2',
+            mainTemplateId: 'mainWithVariants'
+        });
+
+        const variantsForMainWithVariants = promptService.getVariantIds('mainWithVariants');
+        const variantsForMainWithoutVariants = promptService.getVariantIds('mainWithoutVariants');
+
+        expect(variantsForMainWithVariants).to.deep.equal(['variant1', 'variant2']);
+        expect(variantsForMainWithoutVariants).to.deep.equal([]);
+    });
+
 });
diff --git a/packages/ai-core/src/common/prompt-service.ts b/packages/ai-core/src/common/prompt-service.ts
@@ -25,6 +25,11 @@ import { PROMPT_VARIABLE_REGEX, PROMPT_FUNCTION_REGEX } from './prompt-service-u
 export interface PromptTemplate {
     id: string;
     template: string;
+    /**
+     * (Optional) The ID of the main template for which this template is a variant.
+     * If present, this indicates that the current template represents an alternative version of the specified main template.
+     */
+    mainTemplateId?: string;
 }
 
 export interface PromptMap { [id: string]: PromptTemplate }
@@ -63,6 +68,11 @@ export interface PromptService {
      * @param prompt the prompt template to store
      */
     storePrompt(id: string, prompt: string): void;
+    /**
+     * Adds a {@link PromptTemplate} to the list of prompts.
+     * @param promptTemplate the prompt template to store
+     */
+    storePromptTemplate(promptTemplate: PromptTemplate): void;
     /**
      * Removes a prompt from the list of prompts.
      * @param id the id of the prompt
@@ -72,6 +82,12 @@ export interface PromptService {
      * Return all known prompts as a {@link PromptMap map}.
      */
     getAllPrompts(): PromptMap;
+    /**
+     * Retrieve all variant IDs of a given {@link PromptTemplate}.
+     * @param id the id of the main {@link PromptTemplate}
+     * @returns an array of string IDs representing the variants of the given template
+     */
+    getVariantIds(id: string): string[];
 }
 
 export interface CustomAgentDescription {
@@ -258,4 +274,12 @@ export class PromptServiceImpl implements PromptService {
     removePrompt(id: string): void {
         delete this._prompts[id];
     }
+    getVariantIds(id: string): string[] {
+        return Object.values(this._prompts)
+            .filter(prompt => prompt.mainTemplateId === id)
+            .map(variant => variant.id);
+    }
+    storePromptTemplate(promptTemplate: PromptTemplate): void {
+        this._prompts[promptTemplate.id] = promptTemplate;
+    }
 }
diff --git a/packages/ai-core/src/common/settings-service.ts b/packages/ai-core/src/common/settings-service.ts
@@ -30,4 +30,9 @@ export type AISettings = Record<string, AgentSettings>;
 export interface AgentSettings {
     languageModelRequirements?: LanguageModelRequirement[];
     enable?: boolean;
+    /**
+     * A mapping of main template IDs to their selected variant IDs.
+     * If a main template is not present in this mapping, it means the main template is used.
+     */
+    selectedVariants?: Record<string, string>;
 }