@theia/ai-core 1.46.0-next.241

package/src/common/language-model.ts

@@ -0,0 +1,242 @@
+// *****************************************************************************
+// Copyright (C) 2024 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 { ContributionProvider, ILogger, isFunction, isObject, Event, Emitter, CancellationToken } from '@theia/core';
+import { inject, injectable, named, postConstruct } from '@theia/core/shared/inversify';
+
+export type MessageActor = 'user' | 'ai' | 'system';
+
+export interface LanguageModelRequestMessage {
+    actor: MessageActor;
+    type: 'text';
+    query: string;
+}
+export const isLanguageModelRequestMessage = (obj: unknown): obj is LanguageModelRequestMessage =>
+    !!(obj && typeof obj === 'object' &&
+        'type' in obj &&
+        typeof (obj as { type: unknown }).type === 'string' &&
+        (obj as { type: unknown }).type === 'text' &&
+        'query' in obj &&
+        typeof (obj as { query: unknown }).query === 'string'
+    );
+export interface ToolRequest {
+    id: string;
+    name: string;
+    parameters?: { type?: 'object', properties: Record<string, { type: string, [key: string]: unknown }> };
+    description?: string;
+    handler: (arg_string: string) => Promise<unknown>;
+}
+export interface LanguageModelRequest {
+    messages: LanguageModelRequestMessage[],
+    tools?: ToolRequest[];
+    response_format?: { type: 'text' } | { type: 'json_object' } | ResponseFormatJsonSchema;
+    settings?: { [key: string]: unknown };
+}
+export interface ResponseFormatJsonSchema {
+    type: 'json_schema';
+    json_schema: {
+        name: string,
+        description?: string,
+        schema?: Record<string, unknown>,
+        strict?: boolean | null
+    };
+}
+
+export interface LanguageModelTextResponse {
+    text: string;
+}
+export const isLanguageModelTextResponse = (obj: unknown): obj is LanguageModelTextResponse =>
+    !!(obj && typeof obj === 'object' && 'text' in obj && typeof (obj as { text: unknown }).text === 'string');
+
+export interface LanguageModelStreamResponsePart {
+    content?: string | null;
+    tool_calls?: ToolCall[];
+}
+
+export interface ToolCall {
+    id?: string;
+    function?: {
+        arguments?: string;
+        name?: string;
+    },
+    finished?: boolean;
+    result?: string;
+}
+
+export interface LanguageModelStreamResponse {
+    stream: AsyncIterable<LanguageModelStreamResponsePart>;
+}
+export const isLanguageModelStreamResponse = (obj: unknown): obj is LanguageModelStreamResponse =>
+    !!(obj && typeof obj === 'object' && 'stream' in obj);
+
+export interface LanguageModelParsedResponse {
+    parsed: unknown;
+    content: string;
+}
+export const isLanguageModelParsedResponse = (obj: unknown): obj is LanguageModelParsedResponse =>
+    !!(obj && typeof obj === 'object' && 'parsed' in obj && 'content' in obj);
+
+export type LanguageModelResponse = LanguageModelTextResponse | LanguageModelStreamResponse | LanguageModelParsedResponse;
+
+///////////////////////////////////////////
+// Language Model Provider
+///////////////////////////////////////////
+
+export const LanguageModelProvider = Symbol('LanguageModelProvider');
+export type LanguageModelProvider = () => Promise<LanguageModel[]>;
+
+// See also VS Code `ILanguageModelChatMetadata`
+export interface LanguageModelMetaData {
+    readonly id: string;
+    readonly name?: string;
+    readonly vendor?: string;
+    readonly version?: string;
+    readonly family?: string;
+    readonly maxInputTokens?: number;
+    readonly maxOutputTokens?: number;
+    /**
+     * Default request settings for the language model. These settings can be set by a user preferences.
+     * Settings in a request will override these default settings.
+     */
+    readonly defaultRequestSettings?: { [key: string]: unknown };
+}
+
+export namespace LanguageModelMetaData {
+    export function is(arg: unknown): arg is LanguageModelMetaData {
+        return isObject(arg) && 'id' in arg;
+    }
+}
+
+export interface LanguageModel extends LanguageModelMetaData {
+    request(request: LanguageModelRequest, cancellationToken?: CancellationToken): Promise<LanguageModelResponse>;
+}
+
+export namespace LanguageModel {
+    export function is(arg: unknown): arg is LanguageModel {
+        return isObject(arg) && 'id' in arg && isFunction(arg.request);
+    }
+}
+
+// See also VS Code `ILanguageModelChatSelector`
+interface VsCodeLanguageModelSelector {
+    readonly identifier?: string;
+    readonly name?: string;
+    readonly vendor?: string;
+    readonly version?: string;
+    readonly family?: string;
+    readonly tokens?: number;
+}
+
+export interface LanguageModelSelector extends VsCodeLanguageModelSelector {
+    readonly agent: string;
+    readonly purpose: string;
+}
+
+export type LanguageModelRequirement = Omit<LanguageModelSelector, 'agent'>;
+
+export const LanguageModelRegistry = Symbol('LanguageModelRegistry');
+export interface LanguageModelRegistry {
+    onChange: Event<{ models: LanguageModel[] }>;
+    addLanguageModels(models: LanguageModel[]): void;
+    getLanguageModels(): Promise<LanguageModel[]>;
+    getLanguageModel(id: string): Promise<LanguageModel | undefined>;
+    removeLanguageModels(id: string[]): void;
+    selectLanguageModel(request: LanguageModelSelector): Promise<LanguageModel | undefined>;
+    selectLanguageModels(request: LanguageModelSelector): Promise<LanguageModel[]>;
+}
+
+@injectable()
+export class DefaultLanguageModelRegistryImpl implements LanguageModelRegistry {
+    @inject(ILogger)
+    protected logger: ILogger;
+    @inject(ContributionProvider) @named(LanguageModelProvider)
+    protected readonly languageModelContributions: ContributionProvider<LanguageModelProvider>;
+
+    protected languageModels: LanguageModel[] = [];
+
+    protected markInitialized: () => void;
+    protected initialized: Promise<void> = new Promise(resolve => { this.markInitialized = resolve; });
+
+    protected changeEmitter = new Emitter<{ models: LanguageModel[] }>();
+    onChange = this.changeEmitter.event;
+
+    @postConstruct()
+    protected init(): void {
+        const contributions = this.languageModelContributions.getContributions();
+        const promises = contributions.map(provider => provider());
+        Promise.allSettled(promises).then(results => {
+            for (const result of results) {
+                if (result.status === 'fulfilled') {
+                    this.languageModels.push(...result.value);
+                } else {
+                    this.logger.error('Failed to add some language models:', result.reason);
+                }
+            }
+            this.markInitialized();
+        });
+    }
+
+    addLanguageModels(models: LanguageModel[]): void {
+        models.forEach(model => {
+            if (this.languageModels.find(lm => lm.id === model.id)) {
+                console.warn(`Tried to add already existing language model with id ${model.id}. The new model will be ignored.`);
+                return;
+            }
+            this.languageModels.push(model);
+            this.changeEmitter.fire({ models: this.languageModels });
+        });
+    }
+
+    async getLanguageModels(): Promise<LanguageModel[]> {
+        await this.initialized;
+        return this.languageModels;
+    }
+
+    async getLanguageModel(id: string): Promise<LanguageModel | undefined> {
+        await this.initialized;
+        return this.languageModels.find(model => model.id === id);
+    }
+
+    removeLanguageModels(ids: string[]): void {
+        ids.forEach(id => {
+            const index = this.languageModels.findIndex(model => model.id === id);
+            if (index !== -1) {
+                this.languageModels.splice(index, 1);
+                this.changeEmitter.fire({ models: this.languageModels });
+            } else {
+                console.warn(`Language model with id ${id} was requested to be removed, however it does not exist`);
+            }
+        });
+    }
+
+    async selectLanguageModels(request: LanguageModelSelector): Promise<LanguageModel[]> {
+        await this.initialized;
+        // TODO check for actor and purpose against settings
+        return this.languageModels.filter(model => isModelMatching(request, model));
+    }
+
+    async selectLanguageModel(request: LanguageModelSelector): Promise<LanguageModel | undefined> {
+        return (await this.selectLanguageModels(request))[0];
+    }
+}
+
+export function isModelMatching(request: LanguageModelSelector, model: LanguageModel): boolean {
+    return (!request.identifier || model.id === request.identifier) &&
+        (!request.name || model.name === request.name) &&
+        (!request.vendor || model.vendor === request.vendor) &&
+        (!request.version || model.version === request.version) &&
+        (!request.family || model.family === request.family);
+}

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

@@ -0,0 +1,31 @@
+// *****************************************************************************
+// Copyright (C) 2024 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
+// *****************************************************************************
+
+/** 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;
+
+export function matchFunctionsRegEx(template: string): RegExpMatchArray[] {
+    return [...template.matchAll(PROMPT_FUNCTION_REGEX)];
+}

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

@@ -0,0 +1,301 @@
+// *****************************************************************************
+// Copyright (C) 2024 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 'reflect-metadata';
+
+import { expect } from 'chai';
+import { Container } from 'inversify';
+import { PromptService, PromptServiceImpl } from './prompt-service';
+import { DefaultAIVariableService, AIVariableService } from './variable-service';
+
+describe('PromptService', () => {
+    let promptService: PromptService;
+
+    beforeEach(() => {
+        const container = new Container();
+        container.bind<PromptService>(PromptService).to(PromptServiceImpl).inSingletonScope();
+
+        const variableService = new DefaultAIVariableService({ getContributions: () => [] });
+        const nameVariable = { id: 'test', name: 'name', description: 'Test name ' };
+        variableService.registerResolver(nameVariable, {
+            canResolve: () => 100,
+            resolve: async () => ({ variable: nameVariable, value: 'Jane' })
+        });
+        container.bind<AIVariableService>(AIVariableService).toConstantValue(variableService);
+
+        promptService = container.get<PromptService>(PromptService);
+        promptService.storePromptTemplate({ id: '1', template: 'Hello, {{name}}!' });
+        promptService.storePromptTemplate({ id: '2', template: 'Goodbye, {{name}}!' });
+        promptService.storePromptTemplate({ id: '3', template: 'Ciao, {{invalid}}!' });
+        promptService.storePromptTemplate({ id: '8', template: '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', () => {
+        const rawPrompt = promptService.getRawPrompt('1');
+        expect(rawPrompt?.template).to.equal('Hello, {{name}}!');
+    });
+
+    it('should format prompt with provided arguments', async () => {
+        const formattedPrompt = await promptService.getPrompt('1', { name: 'John' });
+        expect(formattedPrompt?.text).to.equal('Hello, John!');
+    });
+
+    it('should store a new prompt', () => {
+        promptService.storePromptTemplate({ id: '3', template: 'Welcome, {{name}}!' });
+        const newPrompt = promptService.getRawPrompt('3');
+        expect(newPrompt?.template).to.equal('Welcome, {{name}}!');
+    });
+
+    it('should replace placeholders with provided arguments', async () => {
+        const prompt = await promptService.getPrompt('1', { name: 'John' });
+        expect(prompt?.text).to.equal('Hello, John!');
+    });
+
+    it('should use variable service to resolve placeholders if argument value is not provided', async () => {
+        const prompt = await promptService.getPrompt('1');
+        expect(prompt?.text).to.equal('Hello, Jane!');
+    });
+
+    it('should return the prompt even if there are no replacements', async () => {
+        const prompt = await promptService.getPrompt('3');
+        expect(prompt?.text).to.equal('Ciao, {{invalid}}!');
+    });
+
+    it('should return undefined if the prompt id is not found', async () => {
+        const prompt = await promptService.getPrompt('4');
+        expect(prompt).to.be.undefined;
+    });
+
+    it('should ignore whitespace in variables', async () => {
+        promptService.storePromptTemplate({ id: '4', template: 'Hello, {{name }}!' });
+        promptService.storePromptTemplate({ id: '5', template: 'Hello, {{ name}}!' });
+        promptService.storePromptTemplate({ id: '6', template: 'Hello, {{ name }}!' });
+        promptService.storePromptTemplate({ id: '7', template: 'Hello, {{       name           }}!' });
+        for (let i = 4; i <= 7; i++) {
+            const prompt = await promptService.getPrompt(`${i}`, { name: 'John' });
+            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.storePromptTemplate({ id: '9', template: 'Hello, {{{name }}}' });
+        promptService.storePromptTemplate({ id: '10', template: 'Hello, {{{ name}}}' });
+        promptService.storePromptTemplate({ id: '11', template: 'Hello, {{{ name }}}' });
+        promptService.storePromptTemplate({ id: '12', template: '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.storePromptTemplate({ id: '9', template: 'Hello, {{name' });
+        promptService.storePromptTemplate({ id: '10', template: 'Hello, {{{name' });
+        promptService.storePromptTemplate({ id: '11', template: '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.storePromptTemplate({ id: '12', template: 'Hi, {{name}}}' });            // (invalid)
+        promptService.storePromptTemplate({ id: '13', template: 'Hello, {{{name}}' });         // (invalid)
+        promptService.storePromptTemplate({ id: '14', template: 'Greetings, {{{name}}}}' });   // (invalid)
+        promptService.storePromptTemplate({ id: '15', template: 'Bye, {{{{name}}}' });         // (invalid)
+        promptService.storePromptTemplate({ id: '16', template: 'Ciao, {{{{name}}}}' });       // (invalid)
+        promptService.storePromptTemplate({ id: '17', template: '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');
+    });
+
+    it('should strip single-line comments at the start of the template', () => {
+        promptService.storePromptTemplate({ id: 'comment-basic', template: '{{!-- Comment --}}Hello, {{name}}!' });
+        const prompt = promptService.getUnresolvedPrompt('comment-basic');
+        expect(prompt?.template).to.equal('Hello, {{name}}!');
+    });
+
+    it('should remove line break after first-line comment', () => {
+        promptService.storePromptTemplate({ id: 'comment-line-break', template: '{{!-- Comment --}}\nHello, {{name}}!' });
+        const prompt = promptService.getUnresolvedPrompt('comment-line-break');
+        expect(prompt?.template).to.equal('Hello, {{name}}!');
+    });
+
+    it('should strip multiline comments at the start of the template', () => {
+        promptService.storePromptTemplate({ id: 'comment-multiline', template: '{{!--\nMultiline comment\n--}}\nGoodbye, {{name}}!' });
+        const prompt = promptService.getUnresolvedPrompt('comment-multiline');
+        expect(prompt?.template).to.equal('Goodbye, {{name}}!');
+    });
+
+    it('should not strip comments not in the first line', () => {
+        promptService.storePromptTemplate({ id: 'comment-second-line', template: 'Hello, {{name}}!\n{{!-- Comment --}}' });
+        const prompt = promptService.getUnresolvedPrompt('comment-second-line');
+        expect(prompt?.template).to.equal('Hello, {{name}}!\n{{!-- Comment --}}');
+    });
+
+    it('should treat unclosed comments as regular text', () => {
+        promptService.storePromptTemplate({ id: 'comment-unclosed', template: '{{!-- Unclosed comment' });
+        const prompt = promptService.getUnresolvedPrompt('comment-unclosed');
+        expect(prompt?.template).to.equal('{{!-- Unclosed comment');
+    });
+
+    it('should treat standalone closing delimiters as regular text', () => {
+        promptService.storePromptTemplate({ id: 'comment-standalone', template: '--}} Hello, {{name}}!' });
+        const prompt = promptService.getUnresolvedPrompt('comment-standalone');
+        expect(prompt?.template).to.equal('--}} Hello, {{name}}!');
+    });
+
+    it('should handle nested comments and stop at the first closing tag', () => {
+        promptService.storePromptTemplate({ id: 'nested-comment', template: '{{!-- {{!-- Nested comment --}} --}}text' });
+        const prompt = promptService.getUnresolvedPrompt('nested-comment');
+        expect(prompt?.template).to.equal('--}}text');
+    });
+
+    it('should handle templates with only comments', () => {
+        promptService.storePromptTemplate({ id: 'comment-only', template: '{{!-- Only comments --}}' });
+        const prompt = promptService.getUnresolvedPrompt('comment-only');
+        expect(prompt?.template).to.equal('');
+    });
+
+    it('should handle mixed delimiters on the same line', () => {
+        promptService.storePromptTemplate({ id: 'comment-mixed', template: '{{!-- Unclosed comment --}}' });
+        const prompt = promptService.getUnresolvedPrompt('comment-mixed');
+        expect(prompt?.template).to.equal('');
+    });
+
+    it('should resolve variables after stripping single-line comments', async () => {
+        promptService.storePromptTemplate({ id: 'comment-resolve', template: '{{!-- Comment --}}Hello, {{name}}!' });
+        const prompt = await promptService.getPrompt('comment-resolve', { name: 'John' });
+        expect(prompt?.text).to.equal('Hello, John!');
+    });
+
+    it('should resolve variables in multiline templates with comments', async () => {
+        promptService.storePromptTemplate({ id: 'comment-multiline-vars', template: '{{!--\nMultiline comment\n--}}\nHello, {{name}}!' });
+        const prompt = await promptService.getPrompt('comment-multiline-vars', { name: 'John' });
+        expect(prompt?.text).to.equal('Hello, John!');
+    });
+
+    it('should resolve variables with standalone closing delimiters', async () => {
+        promptService.storePromptTemplate({ id: 'comment-standalone-vars', template: '--}} Hello, {{name}}!' });
+        const prompt = await promptService.getPrompt('comment-standalone-vars', { name: 'John' });
+        expect(prompt?.text).to.equal('--}} Hello, John!');
+    });
+
+    it('should treat unclosed comments as text and resolve variables', async () => {
+        promptService.storePromptTemplate({ id: 'comment-unclosed-vars', template: '{{!-- Unclosed comment\nHello, {{name}}!' });
+        const prompt = await promptService.getPrompt('comment-unclosed-vars', { name: 'John' });
+        expect(prompt?.text).to.equal('{{!-- Unclosed comment\nHello, John!');
+    });
+
+    it('should handle templates with mixed comments and variables', async () => {
+        promptService.storePromptTemplate({ id: 'comment-mixed-vars', template: '{{!-- Comment --}}Hi, {{name}}! {{!-- Another comment --}}' });
+        const prompt = await promptService.getPrompt('comment-mixed-vars', { name: 'John' });
+        expect(prompt?.text).to.equal('Hi, John! {{!-- Another comment --}}');
+    });
+
+    it('should return all variant IDs of a given prompt', () => {
+        promptService.storePromptTemplate({ id: 'main', template: 'Main template' });
+
+        promptService.storePromptTemplate({
+            id: 'variant1',
+            template: 'Variant 1',
+            variantOf: 'main'
+        });
+        promptService.storePromptTemplate({
+            id: 'variant2',
+            template: 'Variant 2',
+            variantOf: 'main'
+        });
+        promptService.storePromptTemplate({
+            id: 'variant3',
+            template: 'Variant 3',
+            variantOf: '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',
+            variantOf: 'mainWithVariants'
+        });
+        promptService.storePromptTemplate({
+            id: 'variant2',
+            template: 'Variant 2',
+            variantOf: 'mainWithVariants'
+        });
+
+        const variantsForMainWithVariants = promptService.getVariantIds('mainWithVariants');
+        const variantsForMainWithoutVariants = promptService.getVariantIds('mainWithoutVariants');
+
+        expect(variantsForMainWithVariants).to.deep.equal(['variant1', 'variant2']);
+        expect(variantsForMainWithoutVariants).to.deep.equal([]);
+    });
+});