@theia/ai-core 1.59.0 → 1.60.0-next.47

package/src/common/prompt-service.spec.ts

@@ -20,6 +20,10 @@ import { expect } from 'chai';
 import { Container } from 'inversify';
 import { PromptService, PromptServiceImpl } from './prompt-service';
 import { DefaultAIVariableService, AIVariableService } from './variable-service';
+import { ToolInvocationRegistry } from './tool-invocation-registry';
+import { ToolRequest } from './language-model';
+import { Logger } from '@theia/core';
+import * as sinon from 'sinon';
 
 describe('PromptService', () => {
     let promptService: PromptService;
@@ -27,8 +31,9 @@ describe('PromptService', () => {
     beforeEach(() => {
         const container = new Container();
         container.bind<PromptService>(PromptService).to(PromptServiceImpl).inSingletonScope();
+        const logger = sinon.createStubInstance(Logger);
 
-        const variableService = new DefaultAIVariableService({ getContributions: () => [] });
+        const variableService = new DefaultAIVariableService({ getContributions: () => [] }, logger);
         const nameVariable = { id: 'test', name: 'name', description: 'Test name ' };
         variableService.registerResolver(nameVariable, {
             canResolve: () => 100,
@@ -298,4 +303,66 @@ describe('PromptService', () => {
         expect(variantsForMainWithVariants).to.deep.equal(['variant1', 'variant2']);
         expect(variantsForMainWithoutVariants).to.deep.equal([]);
     });
+
+    it('should resolve function references within resolved variable replacements', async () => {
+        // Mock the tool invocation registry
+        const toolInvocationRegistry = {
+            getFunction: sinon.stub()
+        };
+
+        // Create a test tool request that will be returned by the registry
+        const testFunction: ToolRequest = {
+            id: 'testFunction',
+            name: 'Test Function',
+            description: 'A test function',
+            parameters: {
+                type: 'object',
+                properties: {
+                    param1: {
+                        type: 'string',
+                        description: 'Test parameter'
+                    }
+                }
+            },
+            providerName: 'test-provider',
+            handler: sinon.stub()
+        };
+        toolInvocationRegistry.getFunction.withArgs('testFunction').returns(testFunction);
+
+        // Create a container with our mocked registry
+        const container = new Container();
+        container.bind<PromptService>(PromptService).to(PromptServiceImpl).inSingletonScope();
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        container.bind<ToolInvocationRegistry>(ToolInvocationRegistry).toConstantValue(toolInvocationRegistry as any);
+
+        // Set up a variable service that returns a fragment with a function reference
+        const variableService = new DefaultAIVariableService({ getContributions: () => [] }, sinon.createStubInstance(Logger));
+        const fragmentVariable = { id: 'test', name: 'fragment', description: 'Test fragment with function' };
+        variableService.registerResolver(fragmentVariable, {
+            canResolve: () => 100,
+            resolve: async () => ({
+                variable: fragmentVariable,
+                value: 'This fragment contains a function reference: ~{testFunction}'
+            })
+        });
+        container.bind<AIVariableService>(AIVariableService).toConstantValue(variableService);
+
+        const testPromptService = container.get<PromptService>(PromptService);
+        testPromptService.storePromptTemplate({ id: 'testPrompt', template: 'Template with fragment: {{fragment}}' });
+
+        // Get the resolved prompt
+        const resolvedPrompt = await testPromptService.getPrompt('testPrompt');
+
+        // Verify that the function was resolved
+        expect(resolvedPrompt).to.not.be.undefined;
+        expect(resolvedPrompt?.text).to.include('This fragment contains a function reference:');
+        expect(resolvedPrompt?.text).to.not.include('~{testFunction}');
+
+        // Verify that the function description was added to functionDescriptions
+        expect(resolvedPrompt?.functionDescriptions?.size).to.equal(1);
+        expect(resolvedPrompt?.functionDescriptions?.get('testFunction')).to.deep.equal(testFunction);
+
+        // Verify that the tool invocation registry was called
+        expect(toolInvocationRegistry.getFunction.calledWith('testFunction')).to.be.true;
+    });
 });

package/src/common/prompt-service.ts

@@ -16,7 +16,7 @@
 
 import { URI, Event } from '@theia/core';
 import { inject, injectable, optional } from '@theia/core/shared/inversify';
-import { AIVariableContext, AIVariableService } from './variable-service';
+import { AIVariableArg, AIVariableContext, AIVariableService, createAIResolveVariableCache, ResolvedAIVariable } from './variable-service';
 import { ToolInvocationRegistry } from './tool-invocation-registry';
 import { toolRequestToPromptText } from './language-model-util';
 import { ToolRequest } from './language-model';
@@ -41,6 +41,8 @@ export interface ResolvedPromptTemplate {
     text: string;
     /** All functions referenced in the prompt template. */
     functionDescriptions?: Map<string, ToolRequest>;
+    /** All variables resolved in the prompt template */
+    variables?: ResolvedAIVariable[];
 }
 
 export const PromptService = Symbol('PromptService');
@@ -64,10 +66,34 @@ export interface PromptService {
      * Allows to directly replace placeholders in the prompt. The supported format is 'Hi {{name}}!'.
      * The placeholder is then searched inside the args object and replaced.
      * Function references are also supported via format '~{functionId}'.
+     *
+     * All placeholders are replaced before function references are resolved.
+     * This allows to resolve function references contained in placeholders.
+     *
      * @param id the id of the prompt
      * @param args the object with placeholders, mapping the placeholder key to the value
      */
     getPrompt(id: string, args?: { [key: string]: unknown }, context?: AIVariableContext): Promise<ResolvedPromptTemplate | undefined>;
+
+    /**
+     * Allows to directly replace placeholders in the prompt. The supported format is 'Hi {{name}}!'.
+     * The placeholder is then searched inside the args object and replaced.
+     *
+     * In contrast to {@link getPrompt}, this method does not resolve function references but leaves them as is.
+     * This allows resolving them later as part of the prompt or chat message containing the fragment.
+     *
+     * @param id the id of the prompt
+     * @param args the object with placeholders, mapping the placeholder key to the value
+     * @param context the {@link AIVariableContext} to use during variable resolvement
+     * @param resolveVariable the variable resolving method. Fall back to using the {@link AIVariableService} if not given.
+     */
+    getPromptFragment(
+        id: string,
+        args?: { [key: string]: unknown },
+        context?: AIVariableContext,
+        resolveVariable?: (variable: AIVariableArg) => Promise<ResolvedAIVariable | undefined>
+    ): Promise<Omit<ResolvedPromptTemplate, 'functionDescriptions'> | undefined>;
+
     /**
      * Adds a {@link PromptTemplate} to the list of prompts.
      * @param promptTemplate the prompt template to store
@@ -246,27 +272,14 @@ export class PromptServiceImpl implements PromptService {
             return undefined;
         }
 
-        const matches = matchVariablesRegEx(prompt.template);
-        const variableAndArgReplacements = await Promise.all(matches.map(async match => {
-            const completeText = match[0];
-            const variableAndArg = match[1];
-            let variableName = variableAndArg;
-            let argument: string | undefined;
-            const parts = variableAndArg.split(':', 2);
-            if (parts.length > 1) {
-                variableName = parts[0];
-                argument = parts[1];
-            }
-            return {
-                placeholder: completeText,
-                value: String(args?.[variableAndArg] ?? (await this.variableService?.resolveVariable({
-                    variable: variableName,
-                    arg: argument
-                }, context ?? {}))?.value ?? completeText)
-            };
-        }));
+        // First resolve variables and arguments
+        let resolvedTemplate = prompt.template;
+        const variableAndArgReplacements = await this.getVariableAndArgReplacements(prompt.template, args, context);
+        variableAndArgReplacements.replacements.forEach(replacement => resolvedTemplate = resolvedTemplate.replace(replacement.placeholder, replacement.value));
 
-        const functionMatches = matchFunctionsRegEx(prompt.template);
+        // Then resolve function references with already resolved variables and arguments
+        // This allows to resolve function references contained in resolved variables (e.g. prompt fragments)
+        const functionMatches = matchFunctionsRegEx(resolvedTemplate);
         const functions = new Map<string, ToolRequest>();
         const functionReplacements = functionMatches.map(match => {
             const completeText = match[0];
@@ -280,16 +293,88 @@ export class PromptServiceImpl implements PromptService {
                 value: toolRequest ? toolRequestToPromptText(toolRequest) : completeText
             };
         });
+        functionReplacements.forEach(replacement => resolvedTemplate = resolvedTemplate.replace(replacement.placeholder, replacement.value));
+
+        return {
+            id,
+            text: resolvedTemplate,
+            functionDescriptions: functions.size > 0 ? functions : undefined,
+            variables: variableAndArgReplacements.resolvedVariables
+        };
+    }
 
+    async getPromptFragment(
+        id: string,
+        args?: { [key: string]: unknown },
+        context?: AIVariableContext,
+        resolveVariable?: (variable: AIVariableArg) => Promise<ResolvedAIVariable | undefined>
+    ): Promise<Omit<ResolvedPromptTemplate, 'functionDescriptions'> | undefined> {
+        const variantId = await this.getVariantId(id);
+        const prompt = this.getUnresolvedPrompt(variantId);
+        if (prompt === undefined) {
+            return undefined;
+        }
+
+        const replacements = await this.getVariableAndArgReplacements(prompt.template, args, context, resolveVariable);
         let resolvedTemplate = prompt.template;
-        const replacements = [...variableAndArgReplacements, ...functionReplacements];
-        replacements.forEach(replacement => resolvedTemplate = resolvedTemplate.replace(replacement.placeholder, replacement.value));
+        replacements.replacements.forEach(replacement => resolvedTemplate = resolvedTemplate.replace(replacement.placeholder, replacement.value));
+
         return {
             id,
             text: resolvedTemplate,
-            functionDescriptions: functions.size > 0 ? functions : undefined
+            variables: replacements.resolvedVariables
         };
     }
+
+    /**
+     * Calculates all variable and argument replacements for an unresolved template.
+     *
+     * @param template the unresolved template text
+     * @param args the object with placeholders, mapping the placeholder key to the value
+     * @param context the {@link AIVariableContext} to use during variable resolvement
+     * @param resolveVariable the variable resolving method. Fall back to using the {@link AIVariableService} if not given.
+     */
+    protected async getVariableAndArgReplacements(
+        template: string,
+        args?: { [key: string]: unknown },
+        context?: AIVariableContext,
+        resolveVariable?: (variable: AIVariableArg) => Promise<ResolvedAIVariable | undefined>
+    ): Promise<{ replacements: { placeholder: string; value: string }[], resolvedVariables: ResolvedAIVariable[] }> {
+        const matches = matchVariablesRegEx(template);
+        const variableCache = createAIResolveVariableCache();
+        const variableAndArgReplacements: { placeholder: string; value: string }[] = [];
+        const resolvedVariables: Set<ResolvedAIVariable> = new Set();
+        for (const match of matches) {
+            const completeText = match[0];
+            const variableAndArg = match[1];
+            let variableName = variableAndArg;
+            let argument: string | undefined;
+            const parts = variableAndArg.split(':', 2);
+            if (parts.length > 1) {
+                variableName = parts[0];
+                argument = parts[1];
+            }
+            let value: string;
+            if (args && args[variableAndArg] !== undefined) {
+                value = String(args[variableAndArg]);
+            } else {
+                const toResolve = { variable: variableName, arg: argument };
+                const resolved = resolveVariable
+                    ? await resolveVariable(toResolve)
+                    : await this.variableService?.resolveVariable(toResolve, context ?? {}, variableCache);
+                // Track resolved variable and its dependencies in all resolved variables
+                if (resolved) {
+                    resolvedVariables.add(resolved);
+                    resolved.allResolvedDependencies?.forEach(v => resolvedVariables.add(v));
+                }
+                value = String(resolved?.value ?? completeText);
+            }
+            variableAndArgReplacements.push({ placeholder: completeText, value });
+        }
+
+        return { replacements: variableAndArgReplacements, resolvedVariables: Array.from(resolvedVariables) };
+    }
+
     getAllPrompts(): PromptMap {
         if (this.customizationService !== undefined) {
             const myCustomization = this.customizationService;

package/src/common/prompt-variable-contribution.ts

@@ -0,0 +1,138 @@
+// *****************************************************************************
+// Copyright (C) 2025 EclipseSource GmbH.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License v. 2.0 which is available at
+// http://www.eclipse.org/legal/epl-2.0.
+//
+// This Source Code may also be made available under the following Secondary
+// Licenses when the conditions for such availability set forth in the Eclipse
+// Public License v. 2.0 are satisfied: GNU General Public License, version 2
+// with the GNU Classpath Exception which is available at
+// https://www.gnu.org/software/classpath/license.html.
+//
+// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
+// *****************************************************************************
+import { CommandService, nls } from '@theia/core';
+import { injectable, inject, optional } from '@theia/core/shared/inversify';
+import * as monaco from '@theia/monaco-editor-core';
+import {
+    AIVariable,
+    AIVariableContribution,
+    AIVariableService,
+    AIVariableResolutionRequest,
+    AIVariableContext,
+    ResolvedAIVariable,
+    AIVariableResolverWithVariableDependencies,
+    AIVariableArg
+} from './variable-service';
+import { PromptCustomizationService, PromptService } from './prompt-service';
+import { PromptText } from './prompt-text';
+
+export const PROMPT_VARIABLE: AIVariable = {
+    id: 'prompt-provider',
+    description: nls.localize('theia/ai/core/promptVariable/description', 'Resolves prompt templates via the prompt service'),
+    name: 'prompt',
+    args: [
+        { name: 'id', description: nls.localize('theia/ai/core/promptVariable/argDescription', 'The prompt template id to resolve') }
+    ]
+};
+
+@injectable()
+export class PromptVariableContribution implements AIVariableContribution, AIVariableResolverWithVariableDependencies {
+
+    @inject(CommandService)
+    protected readonly commandService: CommandService;
+
+    @inject(PromptService)
+    protected readonly promptService: PromptService;
+
+    @inject(PromptCustomizationService) @optional()
+    protected readonly promptCustomizationService: PromptCustomizationService;
+
+    registerVariables(service: AIVariableService): void {
+        service.registerResolver(PROMPT_VARIABLE, this);
+        service.registerArgumentPicker(PROMPT_VARIABLE, this.triggerArgumentPicker.bind(this));
+        service.registerArgumentCompletionProvider(PROMPT_VARIABLE, this.provideArgumentCompletionItems.bind(this));
+    }
+
+    canResolve(request: AIVariableResolutionRequest, context: AIVariableContext): number {
+        if (request.variable.name === PROMPT_VARIABLE.name) {
+            return 1;
+        }
+        return -1;
+    }
+
+    async resolve(
+        request: AIVariableResolutionRequest,
+        context: AIVariableContext,
+        resolveDependency?: (variable: AIVariableArg) => Promise<ResolvedAIVariable | undefined>
+    ): Promise<ResolvedAIVariable | undefined> {
+        if (request.variable.name === PROMPT_VARIABLE.name) {
+            const promptId = request.arg?.trim();
+            if (promptId) {
+                const resolvedPrompt = await this.promptService.getPromptFragment(promptId, undefined, context, resolveDependency);
+                if (resolvedPrompt) {
+                    return {
+                        variable: request.variable,
+                        value: resolvedPrompt.text,
+                        allResolvedDependencies: resolvedPrompt.variables
+                    };
+                }
+            }
+        }
+        return undefined;
+    }
+
+    protected async triggerArgumentPicker(): Promise<string | undefined> {
+        // Trigger the suggestion command to show argument completions
+        this.commandService.executeCommand('editor.action.triggerSuggest');
+        // Return undefined because we don't actually pick the argument here.
+        // The argument is selected and inserted by the monaco editor's completion mechanism.
+        return undefined;
+    }
+
+    protected async provideArgumentCompletionItems(
+        model: monaco.editor.ITextModel,
+        position: monaco.Position
+    ): Promise<monaco.languages.CompletionItem[] | undefined> {
+        const lineContent = model.getLineContent(position.lineNumber);
+
+        // Only provide completions once the variable argument separator is typed
+        const triggerCharIndex = lineContent.lastIndexOf(PromptText.VARIABLE_SEPARATOR_CHAR, position.column - 1);
+        if (triggerCharIndex === -1) {
+            return undefined;
+        }
+
+        // Check if the text immediately before the trigger is the prompt variable, i.e #prompt
+        const requiredVariable = `${PromptText.VARIABLE_CHAR}${PROMPT_VARIABLE.name}`;
+        if (triggerCharIndex < requiredVariable.length ||
+            lineContent.substring(triggerCharIndex - requiredVariable.length, triggerCharIndex) !== requiredVariable) {
+            return undefined;
+        }
+
+        const range = new monaco.Range(position.lineNumber, triggerCharIndex + 2, position.lineNumber, position.column);
+
+        const customPromptIds = this.promptCustomizationService?.getCustomPromptTemplateIDs() ?? [];
+        const builtinPromptIds = Object.keys(this.promptService.getAllPrompts());
+
+        const customPromptCompletions = customPromptIds.map(promptId => ({
+            label: promptId,
+            kind: monaco.languages.CompletionItemKind.Enum,
+            insertText: promptId,
+            range,
+            detail: nls.localize('theia/ai/core/promptVariable/completions/detail/custom', 'Custom prompt template'),
+            sortText: `AAA${promptId}` // Sort before everything else including all built-in prompts
+        }));
+        const builtinPromptCompletions = builtinPromptIds.map(promptId => ({
+            label: promptId,
+            kind: monaco.languages.CompletionItemKind.Variable,
+            insertText: promptId,
+            range,
+            detail: nls.localize('theia/ai/core/promptVariable/completions/detail/builtin', 'Built-in prompt template'),
+            sortText: `AAB${promptId}` // Sort after all custom prompts but before others
+        }));
+
+        return [...customPromptCompletions, ...builtinPromptCompletions];
+    }
+}