Support prompt variants (#14487)

* Support prompt variants fixed #14485 Signed-off-by: Jonas Helming <[email protected]>
eclipse-theia · Nov 23, 2024 · 2a0232d · 2a0232d
1 parent 0556d2a
commit 2a0232d
Show file tree

Hide file tree

Showing 8 changed files with 302 additions and 74 deletions.
diff --git a/packages/ai-chat/src/common/universal-chat-agent.ts b/packages/ai-chat/src/common/universal-chat-agent.ts
@@ -78,6 +78,12 @@ simple solutions.
 `
 };
 
+export const universalTemplateVariant: PromptTemplate = {
+   id: 'universal-system-empty',
+   template: '',
+   variantOf: universalTemplate.id,
+};
+
 @injectable()
 export class UniversalChatAgent extends AbstractStreamParsingChatAgent implements ChatAgent {
    name: string;
@@ -98,7 +104,7 @@ export class UniversalChatAgent extends AbstractStreamParsingChatAgent implement
          + 'questions the user might ask. The universal agent currently does not have any context by default, i.e. it cannot '
          + 'access the current user context or the workspace.';
       this.variables = [];
-      this.promptTemplates = [universalTemplate];
+      this.promptTemplates = [universalTemplate, universalTemplateVariant];
       this.functions = [];
       this.agentSpecificVariables = [];
    }

diff --git a/packages/ai-core/src/browser/ai-configuration/agent-configuration-widget.tsx b/packages/ai-core/src/browser/ai-configuration/agent-configuration-widget.tsx
@@ -137,14 +137,31 @@ export class AIAgentConfigurationWidget extends ReactWidget {
                     Enable Agent
                 </label>
             </div>
-            <div className='ai-templates'>
-                {agent.promptTemplates?.map(template =>
-                    <TemplateRenderer
-                        key={agent?.id + '.' + template.id}
-                        agentId={agent.id}
-                        template={template}
-                        promptCustomizationService={this.promptCustomizationService} />)}
+            <div className="settings-section-subcategory-title ai-settings-section-subcategory-title">
+                Prompt Templates
             </div>
+            <div className="ai-templates">
+                {(() => {
+                    const defaultTemplates = agent.promptTemplates?.filter(template => !template.variantOf) || [];
+                    return defaultTemplates.length > 0 ? (
+                        defaultTemplates.map(template => (
+                            <div key={agent.id + '.' + template.id}>
+                                <TemplateRenderer
+                                    key={agent.id + '.' + template.id}
+                                    agentId={agent.id}
+                                    template={template}
+                                    promptService={this.promptService}
+                                    aiSettingsService={this.aiSettingsService}
+                                    promptCustomizationService={this.promptCustomizationService}
+                                />
+                            </div>
+                        ))
+                    ) : (
+                        <div>No default template available</div>
+                    );
+                })()}
+            </div>
+
             <div className='ai-lm-requirements'>
                 <LanguageModelRenderer
                     agent={agent}

diff --git a/packages/ai-core/src/browser/ai-configuration/template-settings-renderer.tsx b/packages/ai-core/src/browser/ai-configuration/template-settings-renderer.tsx
@@ -14,26 +14,103 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 import * as React from '@theia/core/shared/react';
-import { PromptCustomizationService } from '../../common/prompt-service';
-import { PromptTemplate } from '../../common';
+import { PromptCustomizationService, PromptService } from '../../common/prompt-service';
+import { AISettingsService, PromptTemplate } from '../../common';
 
-export interface TemplateSettingProps {
+const DEFAULT_VARIANT = 'default';
+
+export interface TemplateRendererProps {
     agentId: string;
     template: PromptTemplate;
     promptCustomizationService: PromptCustomizationService;
+    promptService: PromptService;
+    aiSettingsService: AISettingsService;
 }
 
-export const TemplateRenderer: React.FC<TemplateSettingProps> = ({ agentId, template, promptCustomizationService }) => {
+export const TemplateRenderer: React.FC<TemplateRendererProps> = ({
+    agentId,
+    template,
+    promptCustomizationService,
+    promptService,
+    aiSettingsService,
+}) => {
+    const [variantIds, setVariantIds] = React.useState<string[]>([]);
+    const [selectedVariant, setSelectedVariant] = React.useState<string>(DEFAULT_VARIANT);
+
+    React.useEffect(() => {
+        (async () => {
+            const variants = promptService.getVariantIds(template.id);
+            setVariantIds([DEFAULT_VARIANT, ...variants]);
+
+            const agentSettings = await aiSettingsService.getAgentSettings(agentId);
+            const currentVariant =
+                agentSettings?.selectedVariants?.[template.id] || DEFAULT_VARIANT;
+            setSelectedVariant(currentVariant);
+        })();
+    }, [template.id, promptService, aiSettingsService, agentId]);
+
+    const handleVariantChange = async (event: React.ChangeEvent<HTMLSelectElement>) => {
+        const newVariant = event.target.value;
+        setSelectedVariant(newVariant);
+
+        const agentSettings = await aiSettingsService.getAgentSettings(agentId);
+        const selectedVariants = agentSettings?.selectedVariants || {};
+
+        const updatedVariants = { ...selectedVariants };
+        if (newVariant === DEFAULT_VARIANT) {
+            delete updatedVariants[template.id];
+        } else {
+            updatedVariants[template.id] = newVariant;
+        }
+
+        await aiSettingsService.updateAgentSettings(agentId, {
+            selectedVariants: updatedVariants,
+        });
+    };
+
     const openTemplate = React.useCallback(async () => {
-        promptCustomizationService.editTemplate(template.id, template.template);
-    }, [template, promptCustomizationService]);
+        const templateId = selectedVariant === DEFAULT_VARIANT ? template.id : selectedVariant;
+        const selectedTemplate = promptService.getRawPrompt(templateId);
+        promptCustomizationService.editTemplate(templateId, selectedTemplate?.template || '');
+    }, [selectedVariant, template.id, promptService, promptCustomizationService]);
+
     const resetTemplate = React.useCallback(async () => {
-        promptCustomizationService.resetTemplate(template.id);
-    }, [promptCustomizationService, template]);
-
-    return <>
-        {template.id}
-        <button className='theia-button main' onClick={openTemplate}>Edit</button>
-        <button className='theia-button secondary' onClick={resetTemplate}>Reset</button>
-    </>;
+        const templateId = selectedVariant === DEFAULT_VARIANT ? template.id : selectedVariant;
+        promptCustomizationService.resetTemplate(templateId);
+    }, [selectedVariant, template.id, promptCustomizationService]);
+
+    return (
+        <div className="template-renderer">
+            <div className="settings-section-title template-header">
+                <strong>{template.id}</strong>
+            </div>
+            <div className="template-controls">
+                {variantIds.length > 1 && (
+                    <>
+                        <label htmlFor={`variant-selector-${template.id}`} className="template-select-label">
+                            Select Variant:
+                        </label>
+                        <select
+                            id={`variant-selector-${template.id}`}
+                            className="theia-select template-variant-selector"
+                            value={selectedVariant}
+                            onChange={handleVariantChange}
+                        >
+                            {variantIds.map(variantId => (
+                                <option key={variantId} value={variantId}>
+                                    {variantId}
+                                </option>
+                            ))}
+                        </select>
+                    </>
+                )}
+                <button className="theia-button main" onClick={openTemplate}>
+                    Edit
+                </button>
+                <button className="theia-button secondary" onClick={resetTemplate}>
+                    Reset
+                </button>
+            </div>
+        </div>
+    );
 };
diff --git a/packages/ai-core/src/browser/style/index.css b/packages/ai-core/src/browser/style/index.css
@@ -14,14 +14,42 @@
   margin-left: var(--theia-ui-padding);
 }
 
+.theia-settings-container .settings-section-subcategory-title.ai-settings-section-subcategory-title {
+  padding-left: 0;
+}
+
 .ai-templates {
-  display: grid;
-  /** Display content in 3 columns */
-  grid-template-columns: 1fr auto auto;
-  /** add a 3px gap between rows */
-  row-gap: 3px;
+  display: flex;
+  flex-direction: column;
+  gap: 5px;
+}
+
+.template-renderer {
+  display: flex;
+  flex-direction: column;
+  padding: 10px;
+}
+
+.template-header {
+  margin-bottom: 8px;
 }
 
+.template-controls {
+  display: flex;
+  align-items: center;
+  gap: 10px;
+}
+
+.template-select-label {
+  margin-right: 5px;
+}
+
+.template-variant-selector {
+  min-width: 120px;
+}
+
+
+
 #ai-variable-configuration-container-widget,
 #ai-agent-configuration-container-widget {
   margin-top: 5px;

diff --git a/packages/ai-core/src/common/agent-service.ts b/packages/ai-core/src/common/agent-service.ts
@@ -99,7 +99,7 @@ export class AgentServiceImpl implements AgentService {
     registerAgent(agent: Agent): void {
         this._agents.push(agent);
         agent.promptTemplates.forEach(
-            template => this.promptService.storePrompt(template.id, template.template)
+            template => this.promptService.storePromptTemplate(template)
         );
         this.onDidChangeAgentsEmitter.fire();
     }