All three brackets for variables

fixed #14464
eclipse-theia · Nov 17, 2024 · 03aab6f · 03aab6f
1 parent be43deb
commit 03aab6f
Show file tree

Hide file tree

Showing 5 changed files with 126 additions and 12 deletions.
diff --git a/packages/ai-core/data/prompttemplate.tmLanguage.json b/packages/ai-core/data/prompttemplate.tmLanguage.json
@@ -2,17 +2,56 @@
   "scopeName": "source.prompttemplate",
   "patterns": [
     {
-      "name": "variable.other.prompttemplate",
-      "begin": "{{",
+      "name": "invalid.illegal.mismatched.prompttemplate",
+      "match": "\\{\\{\\{[^}]*\\}\\}(?!\\})",
+      "captures": {
+        "0": {
+          "name": "invalid.illegal.bracket.mismatch"
+        }
+      }
+    },
+    {
+      "name": "invalid.illegal.mismatched.prompttemplate",
+      "match": "\\{\\{[^}]*\\}\\}\\}(?!\\})",
+      "captures": {
+        "0": {
+          "name": "invalid.illegal.bracket.mismatch"
+        }
+      }
+    },
+    {
+      "name": "variable.other.prompttemplate.double",
+      "begin": "\\{\\{",
       "beginCaptures": {
         "0": {
-          "name": "punctuation.definition.brace.begin"
+          "name": "punctuation.definition.variable.begin"
         }
       },
-      "end": "}}",
+      "end": "\\}\\}(?!\\})",
       "endCaptures": {
         "0": {
-          "name": "punctuation.definition.brace.end"
+          "name": "punctuation.definition.variable.end"
+        }
+      },
+      "patterns": [
+        {
+          "name": "keyword.control",
+          "match": "[a-zA-Z_][a-zA-Z0-9_]*"
+        }
+      ]
+    },
+    {
+      "name": "variable.other.prompttemplate.triple",
+      "begin": "\\{\\{\\{",
+      "beginCaptures": {
+        "0": {
+          "name": "punctuation.definition.variable.begin"
+        }
+      },
+      "end": "\\}\\}\\}(?!\\})",
+      "endCaptures": {
+        "0": {
+          "name": "punctuation.definition.variable.end"
         }
       },
       "patterns": [
@@ -49,4 +88,4 @@
   "fileTypes": [
     ".prompttemplate"
   ]
-}
+}
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
@@ -23,8 +23,8 @@ import {
     AIVariableService,
     LanguageModel,
     LanguageModelRegistry,
+    matchVariablesRegEx,
     PROMPT_FUNCTION_REGEX,
-    PROMPT_VARIABLE_REGEX,
     PromptCustomizationService,
     PromptService,
 } from '../../common';
@@ -182,7 +182,7 @@ export class AIAgentConfigurationWidget extends ReactWidget {
         promptTemplates.forEach(template => {
             const storedPrompt = this.promptService.getRawPrompt(template.id);
             const prompt = storedPrompt?.template ?? template.template;
-            const variableMatches = [...prompt.matchAll(PROMPT_VARIABLE_REGEX)];
+            const variableMatches = matchVariablesRegEx(prompt);
 
             variableMatches.forEach(match => {
                 const variableId = match[1];

diff --git a/packages/ai-core/src/common/prompt-service-util.ts b/packages/ai-core/src/common/prompt-service-util.ts
@@ -14,8 +14,14 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 
-/** Should match the one from VariableResolverService. The format is `{{variableName:arg}}`. */
-export const PROMPT_VARIABLE_REGEX = /\{\{\s*(.*?)\s*\}\}/g;
+/** Should match the one from VariableResolverService. The format is `{{variableName:arg}}`. We allow {{}} and {{{}}} but no mixtures */
+export const PROMPT_VARIABLE_TWO_BRACES_REGEX = /(?<!\{)\{\{\s*([^{}]+?)\s*\}\}(?!\})/g;
+export const PROMPT_VARIABLE_THREE_BRACES_REGEX = /(?<!\{)\{\{\{\s*([^{}]+?)\s*\}\}\}(?!\})/g;
+export function matchVariablesRegEx(template: string): RegExpMatchArray[] {
+    const twoBraceMatches = [...template.matchAll(PROMPT_VARIABLE_TWO_BRACES_REGEX)];
+    const threeBraceMatches = [...template.matchAll(PROMPT_VARIABLE_THREE_BRACES_REGEX)];
+    return twoBraceMatches.concat(threeBraceMatches);
+}
 
 /** Match function/tool references in the prompt. The format is `~{functionId}`. */
 export const PROMPT_FUNCTION_REGEX = /\~\{\s*(.*?)\s*\}/g;
diff --git a/packages/ai-core/src/common/prompt-service.spec.ts b/packages/ai-core/src/common/prompt-service.spec.ts
@@ -40,13 +40,15 @@ describe('PromptService', () => {
         promptService.storePrompt('1', 'Hello, {{name}}!');
         promptService.storePrompt('2', 'Goodbye, {{name}}!');
         promptService.storePrompt('3', 'Ciao, {{invalid}}!');
+        promptService.storePrompt('8', 'Hello, {{{name}}}');
     });
 
     it('should initialize prompts from PromptCollectionService', () => {
         const allPrompts = promptService.getAllPrompts();
         expect(allPrompts['1'].template).to.equal('Hello, {{name}}!');
         expect(allPrompts['2'].template).to.equal('Goodbye, {{name}}!');
         expect(allPrompts['3'].template).to.equal('Ciao, {{invalid}}!');
+        expect(allPrompts['8'].template).to.equal('Hello, {{{name}}}');
     });
 
     it('should retrieve raw prompt by id', () => {
@@ -95,4 +97,66 @@ describe('PromptService', () => {
             expect(prompt?.text).to.equal('Hello, John!');
         }
     });
+
+    it('should retrieve raw prompt by id (three bracket)', () => {
+        const rawPrompt = promptService.getRawPrompt('8');
+        expect(rawPrompt?.template).to.equal('Hello, {{{name}}}');
+    });
+
+    it('should correctly replace variables (three brackets)', async () => {
+        const formattedPrompt = await promptService.getPrompt('8');
+        expect(formattedPrompt?.text).to.equal('Hello, Jane');
+    });
+
+    it('should ignore whitespace in variables (three bracket)', async () => {
+        promptService.storePrompt('9', 'Hello, {{{name }}}');
+        promptService.storePrompt('10', 'Hello, {{{ name}}}');
+        promptService.storePrompt('11', 'Hello, {{{ name }}}');
+        promptService.storePrompt('12', 'Hello, {{{       name           }}}');
+        for (let i = 9; i <= 12; i++) {
+            const prompt = await promptService.getPrompt(`${i}`, { name: 'John' });
+            expect(prompt?.text).to.equal('Hello, John');
+        }
+    });
+
+    it('should ignore invalid prompts with unmatched brackets', async () => {
+        promptService.storePrompt('9', 'Hello, {{name');
+        promptService.storePrompt('10', 'Hello, {{{name');
+        promptService.storePrompt('11', 'Hello, name}}}}');
+        const prompt1 = await promptService.getPrompt('9', { name: 'John' });
+        expect(prompt1?.text).to.equal('Hello, {{name'); // Not matching due to missing closing brackets
+
+        const prompt2 = await promptService.getPrompt('10', { name: 'John' });
+        expect(prompt2?.text).to.equal('Hello, {{{name'); // Matches pattern due to valid three-start-two-end brackets
+
+        const prompt3 = await promptService.getPrompt('11', { name: 'John' });
+        expect(prompt3?.text).to.equal('Hello, name}}}}'); // Extra closing bracket, does not match cleanly
+    });
+
+    it('should handle a mixture of two and three brackets correctly', async () => {
+        promptService.storePrompt('12', 'Hi, {{name}}}');            // (invalid)
+        promptService.storePrompt('13', 'Hello, {{{name}}');         // (invalid)
+        promptService.storePrompt('14', 'Greetings, {{{name}}}}');   // (invalid)
+        promptService.storePrompt('15', 'Bye, {{{{name}}}');         // (invalid)
+        promptService.storePrompt('16', 'Ciao, {{{{name}}}}');       // (invalid)
+        promptService.storePrompt('17', 'Hi, {{name}}! {{{name}}}'); // Mixed valid patterns
+
+        const prompt12 = await promptService.getPrompt('12', { name: 'John' });
+        expect(prompt12?.text).to.equal('Hi, {{name}}}');
+
+        const prompt13 = await promptService.getPrompt('13', { name: 'John' });
+        expect(prompt13?.text).to.equal('Hello, {{{name}}');
+
+        const prompt14 = await promptService.getPrompt('14', { name: 'John' });
+        expect(prompt14?.text).to.equal('Greetings, {{{name}}}}');
+
+        const prompt15 = await promptService.getPrompt('15', { name: 'John' });
+        expect(prompt15?.text).to.equal('Bye, {{{{name}}}');
+
+        const prompt16 = await promptService.getPrompt('16', { name: 'John' });
+        expect(prompt16?.text).to.equal('Ciao, {{{{name}}}}');
+
+        const prompt17 = await promptService.getPrompt('17', { name: 'John' });
+        expect(prompt17?.text).to.equal('Hi, John! John');
+    });
 });
diff --git a/packages/ai-core/src/common/prompt-service.ts b/packages/ai-core/src/common/prompt-service.ts
@@ -20,7 +20,7 @@ import { AIVariableService } from './variable-service';
 import { ToolInvocationRegistry } from './tool-invocation-registry';
 import { toolRequestToPromptText } from './language-model-util';
 import { ToolRequest } from './language-model';
-import { PROMPT_VARIABLE_REGEX, PROMPT_FUNCTION_REGEX } from './prompt-service-util';
+import { PROMPT_FUNCTION_REGEX, matchVariablesRegEx } from './prompt-service-util';
 
 export interface PromptTemplate {
     id: string;
@@ -181,13 +181,18 @@ export class PromptServiceImpl implements PromptService {
     getDefaultRawPrompt(id: string): PromptTemplate | undefined {
         return this._prompts[id];
     }
+
+    matchVariables(template: string): RegExpMatchArray[] {
+        return matchVariablesRegEx(template);
+    }
+
     async getPrompt(id: string, args?: { [key: string]: unknown }): Promise<ResolvedPromptTemplate | undefined> {
         const prompt = this.getRawPrompt(id);
         if (prompt === undefined) {
             return undefined;
         }
 
-        const matches = [...prompt.template.matchAll(PROMPT_VARIABLE_REGEX)];
+        const matches = this.matchVariables(prompt.template);
         const variableAndArgReplacements = await Promise.all(matches.map(async match => {
             const completeText = match[0];
             const variableAndArg = match[1];