fix: properly initialize custom agents (#14354)

Prompt templates of newly added custom agents were not stored within the PromptService. As a consequence their prompt was ignored when invoked in Chat. This issue is now fixed by restructuring the prompt storing mechanism. Prompts are now implicitly stored and removed based on all currently registered agents.
eclipse-theia · Oct 28, 2024 · a5f26a3 · a5f26a3
1 parent c1e4279
commit a5f26a3
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 50 deletions.
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
@@ -25,7 +25,7 @@ export interface TemplateSettingProps {
 
 export const TemplateRenderer: React.FC<TemplateSettingProps> = ({ agentId, template, promptCustomizationService }) => {
     const openTemplate = React.useCallback(async () => {
-        promptCustomizationService.editTemplate(template.id);
+        promptCustomizationService.editTemplate(template.id, template.template);
     }, [template, promptCustomizationService]);
     const resetTemplate = React.useCallback(async () => {
         promptCustomizationService.resetTemplate(template.id);

diff --git a/packages/ai-core/src/browser/ai-core-frontend-application-contribution.ts b/packages/ai-core/src/browser/ai-core-frontend-application-contribution.ts
@@ -15,23 +15,22 @@
 // *****************************************************************************
 
 import { FrontendApplicationContribution } from '@theia/core/lib/browser';
-import { inject, injectable } from '@theia/core/shared/inversify';
-import { PromptService } from '../common';
+import { inject, injectable, named } from '@theia/core/shared/inversify';
+import { Agent } from '../common';
 import { AgentService } from '../common/agent-service';
+import { ContributionProvider } from '@theia/core/lib/common/contribution-provider';
 
 @injectable()
 export class AICoreFrontendApplicationContribution implements FrontendApplicationContribution {
     @inject(AgentService)
     private readonly agentService: AgentService;
 
-    @inject(PromptService)
-    private readonly promptService: PromptService;
+    @inject(ContributionProvider) @named(Agent)
+    protected readonly agentsProvider: ContributionProvider<Agent>;
 
     onStart(): void {
-        this.agentService.getAllAgents().forEach(a => {
-            a.promptTemplates.forEach(t => {
-                this.promptService.storePrompt(t.id, t.template);
-            });
+        this.agentsProvider.getContributions().forEach(agent => {
+            this.agentService.registerAgent(agent);
         });
     }
 

diff --git a/packages/ai-core/src/browser/frontend-prompt-customization-service.ts b/packages/ai-core/src/browser/frontend-prompt-customization-service.ts
@@ -17,12 +17,11 @@
 import { DisposableCollection, URI, Event, Emitter } from '@theia/core';
 import { OpenerService } from '@theia/core/lib/browser';
 import { inject, injectable, postConstruct } from '@theia/core/shared/inversify';
-import { PromptCustomizationService, PromptTemplate, CustomAgentDescription } from '../common';
+import { PromptCustomizationService, CustomAgentDescription } from '../common';
 import { BinaryBuffer } from '@theia/core/lib/common/buffer';
 import { FileService } from '@theia/filesystem/lib/browser/file-service';
 import { FileChangesEvent } from '@theia/filesystem/lib/common/files';
 import { AICorePreferences, PREFERENCE_NAME_PROMPT_TEMPLATES } from './ai-core-preferences';
-import { AgentService } from '../common/agent-service';
 import { EnvVariablesServer } from '@theia/core/lib/common/env-variables';
 import { load, dump } from 'js-yaml';
 
@@ -49,9 +48,6 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati
     @inject(OpenerService)
     protected readonly openerService: OpenerService;
 
-    @inject(AgentService)
-    protected readonly agentService: AgentService;
-
     protected readonly trackedTemplateURIs = new Set<string>();
     protected templates = new Map<string, string>();
 
@@ -168,17 +164,10 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati
         return this.templates.get(id);
     }
 
-    async editTemplate(id: string, content?: string): Promise<void> {
-        const template = this.getOriginalTemplate(id);
-        if (template === undefined) {
-            throw new Error(`Unable to edit template ${id}: template not found.`);
-        }
+    async editTemplate(id: string, defaultContent?: string): Promise<void> {
         const editorUri = await this.getTemplateURI(id);
         if (! await this.fileService.exists(editorUri)) {
-            await this.fileService.createFile(editorUri, BinaryBuffer.fromString(content ?? template.template));
-        } else if (content) {
-            // Write content to the file before opening it
-            await this.fileService.writeFile(editorUri, BinaryBuffer.fromString(content));
+            await this.fileService.createFile(editorUri, BinaryBuffer.fromString(defaultContent ?? ''));
         }
         const openHandler = await this.openerService.getOpener(editorUri);
         openHandler.open(editorUri);
@@ -191,17 +180,6 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati
         }
     }
 
-    getOriginalTemplate(id: string): PromptTemplate | undefined {
-        for (const agent of this.agentService.getAllAgents()) {
-            for (const template of agent.promptTemplates) {
-                if (template.id === id) {
-                    return template;
-                }
-            }
-        }
-        return undefined;
-    }
-
     getTemplateIDFromURI(uri: URI): string | undefined {
         const id = this.removePromptTemplateSuffix(uri.path.name);
         if (this.templates.has(id)) {

diff --git a/packages/ai-core/src/common/agent-service.ts b/packages/ai-core/src/common/agent-service.ts
@@ -13,10 +13,11 @@
 //
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
-import { inject, injectable, named, optional, postConstruct } from '@theia/core/shared/inversify';
-import { ContributionProvider, Emitter, Event } from '@theia/core';
+import { inject, injectable, optional, postConstruct } from '@theia/core/shared/inversify';
+import { Emitter, Event } from '@theia/core';
 import { Agent } from './agent';
 import { AISettingsService } from './settings-service';
+import { PromptService } from './prompt-service';
 
 export const AgentService = Symbol('AgentService');
 
@@ -71,12 +72,12 @@ export interface AgentService {
 @injectable()
 export class AgentServiceImpl implements AgentService {
 
-    @inject(ContributionProvider) @named(Agent)
-    protected readonly agentsProvider: ContributionProvider<Agent>;
-
     @inject(AISettingsService) @optional()
     protected readonly aiSettingsService: AISettingsService | undefined;
 
+    @inject(PromptService)
+    protected readonly promptService: PromptService;
+
     protected disabledAgents = new Set<string>();
 
     protected _agents: Agent[] = [];
@@ -95,28 +96,29 @@ export class AgentServiceImpl implements AgentService {
         });
     }
 
-    private get agents(): Agent[] {
-        // We can't collect the contributions at @postConstruct because this will lead to a circular dependency
-        // with agents reusing the chat agent service (e.g. orchestrator) which in turn injects the agent service
-        return [...this.agentsProvider.getContributions(), ...this._agents];
-    }
-
     registerAgent(agent: Agent): void {
         this._agents.push(agent);
+        agent.promptTemplates.forEach(
+            template => this.promptService.storePrompt(template.id, template.template)
+        );
         this.onDidChangeAgentsEmitter.fire();
     }
 
     unregisterAgent(agentId: string): void {
+        const agent = this._agents.find(a => a.id === agentId);
         this._agents = this._agents.filter(a => a.id !== agentId);
         this.onDidChangeAgentsEmitter.fire();
+        agent?.promptTemplates.forEach(
+            template => this.promptService.removePrompt(template.id)
+        );
     }
 
     getAgents(): Agent[] {
-        return this.agents.filter(agent => this.isEnabled(agent.id));
+        return this._agents.filter(agent => this.isEnabled(agent.id));
     }
 
     getAllAgents(): Agent[] {
-        return this.agents;
+        return this._agents;
     }
 
     enableAgent(agentId: string): void {

diff --git a/packages/ai-core/src/common/prompt-service.ts b/packages/ai-core/src/common/prompt-service.ts
@@ -58,11 +58,16 @@ export interface PromptService {
      */
     getPrompt(id: string, args?: { [key: string]: unknown }): Promise<ResolvedPromptTemplate | undefined>;
     /**
-     * Manually add a prompt to the list of prompts.
+     * Adds a prompt to the list of prompts.
      * @param id the id of the prompt
      * @param prompt the prompt template to store
      */
     storePrompt(id: string, prompt: string): void;
+    /**
+     * Removes a prompt from the list of prompts.
+     * @param id the id of the prompt
+     */
+    removePrompt(id: string): void;
     /**
      * Return all known prompts as a {@link PromptMap map}.
      */
@@ -113,9 +118,9 @@ export interface PromptCustomizationService {
      * on the implementation. Implementation may for example decide to
      * open an editor, or request more information from the user, ...
      * @param id the template id.
-     * @param content optional content to customize the template.
+     * @param content optional default content to initialize the template
      */
-    editTemplate(id: string, content?: string): void;
+    editTemplate(id: string, defaultContent?: string): void;
 
     /**
      * Reset the template to its default value.
@@ -250,4 +255,7 @@ export class PromptServiceImpl implements PromptService {
     storePrompt(id: string, prompt: string): void {
         this._prompts[id] = { id, template: prompt };
     }
+    removePrompt(id: string): void {
+        delete this._prompts[id];
+    }
 }