feat: show variables and function on ai agent configuration

The agent configuration now shows also the variables (global and agent specific) as well as the functions that the agents uses. If a variable or function is used in the prompt but is not declared by the agent then this is also marked in the view. All agents are updated to declare the variables and functions they use. fixes #14133
eclipse-theia · Sep 18, 2024 · 4a4c035 · 4a4c035
1 parent bdbe90e
commit 4a4c035
Show file tree

Hide file tree

Showing 11 changed files with 169 additions and 22 deletions.
diff --git a/packages/ai-chat/src/common/command-chat-agents.ts b/packages/ai-chat/src/common/command-chat-agents.ts
@@ -18,6 +18,7 @@ import { inject, injectable } from '@theia/core/shared/inversify';
 import { AbstractTextToModelParsingChatAgent, ChatAgent, SystemMessageDescription } from './chat-agents';
 import {
     PromptTemplate,
+    AgentSpecificVariables
 } from '@theia/ai-core';
 import {
     ChatRequestModelImpl,
@@ -252,11 +253,13 @@ export class CommandChatAgent extends AbstractTextToModelParsingChatAgent<Parsed
     protected commandRegistry: CommandRegistry;
     @inject(MessageService)
     protected messageService: MessageService;
-
     readonly name: string;
     readonly description: string;
     readonly variables: string[];
     readonly promptTemplates: PromptTemplate[];
+    readonly functions: string[];
+    readonly agentSpecificVariables: AgentSpecificVariables[];
+
     constructor(
     ) {
         super('Command', [{
@@ -268,6 +271,12 @@ export class CommandChatAgent extends AbstractTextToModelParsingChatAgent<Parsed
         Based on the user request, it can find the right command and then let the user execute it.';
         this.variables = [];
         this.promptTemplates = [commandTemplate];
+        this.functions = [];
+        this.agentSpecificVariables = [{
+            name: 'command-ids',
+            description: 'The list of available commands in Theia.',
+            usedInPrompt: true
+        }];
     }
 
     protected async getSystemMessageDescription(): Promise<SystemMessageDescription | undefined> {

diff --git a/packages/ai-chat/src/common/orchestrator-chat-agent.ts b/packages/ai-chat/src/common/orchestrator-chat-agent.ts
@@ -14,7 +14,7 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 
-import { getJsonOfResponse, LanguageModelResponse } from '@theia/ai-core';
+import { AgentSpecificVariables, getJsonOfResponse, LanguageModelResponse } from '@theia/ai-core';
 import {
     PromptTemplate
 } from '@theia/ai-core/lib/common';
@@ -64,9 +64,12 @@ export const OrchestratorChatAgentId = 'Orchestrator';
 export class OrchestratorChatAgent extends AbstractStreamParsingChatAgent implements ChatAgent {
     name: string;
     description: string;
-    variables: string[];
+    readonly variables: string[];
     promptTemplates: PromptTemplate[];
     fallBackChatAgentId: string;
+    readonly functions: string[] = [];
+    readonly agentSpecificVariables: AgentSpecificVariables[] = [];
+
     constructor() {
         super(OrchestratorChatAgentId, [{
             purpose: 'agent-selection',
@@ -78,6 +81,8 @@ export class OrchestratorChatAgent extends AbstractStreamParsingChatAgent implem
         this.variables = ['chatAgents'];
         this.promptTemplates = [orchestratorTemplate];
         this.fallBackChatAgentId = 'Universal';
+        this.functions = [];
+        this.agentSpecificVariables = [];
     }
 
     @inject(ChatAgentService)

diff --git a/packages/ai-chat/src/common/universal-chat-agent.ts b/packages/ai-chat/src/common/universal-chat-agent.ts
@@ -14,6 +14,7 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 
+import { AgentSpecificVariables } from '@theia/ai-core';
 import {
    PromptTemplate
 } from '@theia/ai-core/lib/common';
@@ -81,6 +82,8 @@ export class UniversalChatAgent extends AbstractStreamParsingChatAgent implement
    description: string;
    variables: string[];
    promptTemplates: PromptTemplate[];
+   readonly functions: string[];
+   readonly agentSpecificVariables: AgentSpecificVariables[];
 
    constructor() {
       super('Universal', [{
@@ -94,6 +97,8 @@ export class UniversalChatAgent extends AbstractStreamParsingChatAgent implement
          + 'access the current user context or the workspace.';
       this.variables = [];
       this.promptTemplates = [universalTemplate];
+      this.functions = [];
+      this.agentSpecificVariables = [];
    }
 
    protected override async getSystemMessageDescription(): Promise<SystemMessageDescription | undefined> {

diff --git a/packages/ai-code-completion/src/common/code-completion-agent.ts b/packages/ai-code-completion/src/common/code-completion-agent.ts
@@ -15,7 +15,7 @@
 // *****************************************************************************
 
 import {
-    Agent, CommunicationHistoryEntry, CommunicationRecordingService, getTextOfResponse,
+    Agent, AgentSpecificVariables, CommunicationHistoryEntry, CommunicationRecordingService, getTextOfResponse,
     LanguageModelRegistry, LanguageModelRequest, LanguageModelRequirement, PromptService, PromptTemplate
 } from '@theia/ai-core/lib/common';
 import { CancellationToken, generateUuid, ILogger } from '@theia/core';
@@ -32,7 +32,14 @@ export interface CodeCompletionAgent extends Agent {
 
 @injectable()
 export class CodeCompletionAgentImpl implements CodeCompletionAgent {
-    variables: string[] = [];
+    readonly variables: string[] = [];
+    readonly functions: string[] = [];
+    readonly agentSpecificVariables: AgentSpecificVariables[] = [
+        { name: 'file', usedInPrompt: true, description: 'The uri of the file being edited.' },
+        { name: 'language', usedInPrompt: true, description: 'The languageId of the file being edited.' },
+        { name: 'snippet', usedInPrompt: true, description: 'The code snippet to be completed.' },
+        { name: 'MARKER', usedInPrompt: true, description: 'The position where the completion should be inserted.' }
+    ];
 
     @inject(ILogger) @named('code-completion-agent')
     protected logger: ILogger;

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
@@ -17,14 +17,30 @@
 import { codicon, ReactWidget } from '@theia/core/lib/browser';
 import { inject, injectable, postConstruct } from '@theia/core/shared/inversify';
 import * as React from '@theia/core/shared/react';
-import { Agent, LanguageModel, LanguageModelRegistry, PromptCustomizationService } from '../../common';
+import {
+    Agent,
+    AIVariableService,
+    LanguageModel,
+    LanguageModelRegistry,
+    PROMPT_FUNCTION_REGEX,
+    PROMPT_VARIABLE_REGEX,
+    PromptCustomizationService,
+    PromptService,
+    PromptTemplate,
+} from '../../common';
 import { AISettingsService } from '../ai-settings-service';
 import { LanguageModelRenderer } from './language-model-renderer';
 import { TemplateRenderer } from './template-settings-renderer';
 import { AIConfigurationSelectionService } from './ai-configuration-service';
 import { AIVariableConfigurationWidget } from './variable-configuration-widget';
 import { AgentService } from '../../common/agent-service';
 
+interface ParsedPrompt {
+    functions: string[];
+    globalVariables: string[];
+    agentSpecificVariables: string[];
+};
+
 @injectable()
 export class AIAgentConfigurationWidget extends ReactWidget {
 
@@ -46,6 +62,12 @@ export class AIAgentConfigurationWidget extends ReactWidget {
     @inject(AIConfigurationSelectionService)
     protected readonly aiConfigurationSelectionService: AIConfigurationSelectionService;
 
+    @inject(AIVariableService)
+    protected readonly variableService: AIVariableService;
+
+    @inject(PromptService)
+    protected promptService: PromptService;
+
     protected languageModels: LanguageModel[] | undefined;
 
     @postConstruct()
@@ -62,6 +84,7 @@ export class AIAgentConfigurationWidget extends ReactWidget {
             this.languageModels = models;
             this.update();
         }));
+        this.toDispose.push(this.promptCustomizationService.onDidChangePrompt(() => this.update()));
 
         this.aiSettingsService.onDidChange(() => this.update());
         this.aiConfigurationSelectionService.onDidAgentChange(() => this.update());
@@ -98,6 +121,11 @@ export class AIAgentConfigurationWidget extends ReactWidget {
 
         const enabled = this.agentService.isEnabled(agent.id);
 
+        const parsedPromptParts = this.parsePromptTemplatesForVariableAndFunction(agent.promptTemplates);
+        const globalVariables = Array.from(new Set([...parsedPromptParts.globalVariables, ...agent.variables]));
+        const allOtherVariables = Array.from(new Set([...parsedPromptParts.agentSpecificVariables, ...(agent.agentSpecificVariables.map(v => v.name))]));
+        const functions = Array.from(new Set([...parsedPromptParts.functions, ...agent.functions]));
+
         return <div key={agent.id} style={{ display: 'flex', flexDirection: 'column', alignItems: 'flex-start' }}>
             <div className='settings-section-title settings-section-category-title' style={{ paddingLeft: 0, paddingBottom: 10 }}>{this.renderAgentName(agent)}</div>
             <div style={{ paddingBottom: 10 }}>{agent.description}</div>
@@ -107,16 +135,6 @@ export class AIAgentConfigurationWidget extends ReactWidget {
                     Enable Agent
                 </label>
             </div>
-            <div style={{ paddingBottom: 10 }}>
-                <span style={{ marginRight: '0.5rem' }}>Variables:</span>
-                <ul className='variable-references'>
-                    {agent.variables.map(variableId => <li key={variableId} className='theia-TreeNode theia-CompositeTreeNode theia-ExpandableTreeNode theia-mod-selected'>
-                        <div key={variableId} onClick={() => { this.showVariableConfigurationTab(); }} className='variable-reference'>
-                            <span>{variableId}</span>
-                            <i className={codicon('chevron-right')}></i>
-                        </div></li>)}
-                </ul>
-            </div>
             <div className='ai-templates'>
                 {agent.promptTemplates?.map(template =>
                     <TemplateRenderer
@@ -132,9 +150,76 @@ export class AIAgentConfigurationWidget extends ReactWidget {
                     aiSettingsService={this.aiSettingsService}
                     languageModelRegistry={this.languageModelRegistry} />
             </div>
+            <div style={{ paddingBottom: 10 }}>
+                <span style={{ marginRight: '0.5rem' }}>Used Global Variables:</span>
+                <ul className='variable-references'>
+                    {globalVariables.map(variableId => <li key={variableId} className='theia-TreeNode theia-CompositeTreeNode theia-ExpandableTreeNode theia-mod-selected'>
+                        <div key={variableId} onClick={() => { this.showVariableConfigurationTab(); }} className='variable-reference'>
+                            <span>{variableId}</span>
+                            <i className={codicon('chevron-right')}></i>
+                        </div></li>)}
+                </ul>
+            </div>
+            <div style={{ paddingBottom: 10 }}>
+                <span style={{ marginRight: '0.5rem' }}>Used agent-specific Variables:</span>
+                <ul className='variable-references'>
+                    {allOtherVariables.map(variableId => {
+                        const agentSpecificVariable = agent.agentSpecificVariables.find(asv => asv.name === variableId);
+                        const undeclared = agentSpecificVariable === undefined;
+                        const notUsed = !parsedPromptParts.agentSpecificVariables.includes(variableId) && agentSpecificVariable?.usedInPrompt === true;
+                        return <li key={variableId} className='theia-TreeNode theia-CompositeTreeNode theia-ExpandableTreeNode theia-mod-selected'>
+                            <div>
+                                <div><span>Name:</span> <span>{variableId}</span></div>
+                                {undeclared ? <div><span>Undeclared</span></div> :
+                                    (<React.Fragment>
+                                        <div><span>Description:</span> <span>{agentSpecificVariable.description}</span></div>
+                                        {notUsed && <div>Not used in prompt</div>}
+                                    </React.Fragment>)}
+                                <hr />
+                            </div>
+                        </li>;
+                    })}
+                </ul>
+            </div>
+            <div style={{ paddingBottom: 10 }}>
+                <span style={{ marginRight: '0.5rem' }}>Used Functions:</span>
+                <ul className='function-references'>
+                    {functions.map(functionId => <li key={functionId} className='theia-TreeNode theia-CompositeTreeNode theia-ExpandableTreeNode theia-mod-selected'>
+                        <div key={functionId} onClick={() => { this.showVariableConfigurationTab(); }} className='variable-reference'>
+                            <span>{functionId}</span>
+                            <i className={codicon('chevron-right')}></i>
+                        </div></li>)}
+                </ul>
+            </div>
         </div>;
     }
 
+    private parsePromptTemplatesForVariableAndFunction(promptTemplates: PromptTemplate[]): ParsedPrompt {
+        const result: ParsedPrompt = { functions: [], globalVariables: [], agentSpecificVariables: [] };
+        promptTemplates.forEach(template => {
+            const storedPrompt = this.promptService.getRawPrompt(template.id);
+            const prompt = storedPrompt?.template ?? template.template;
+            const variableMatches = [...prompt.matchAll(PROMPT_VARIABLE_REGEX)];
+
+            variableMatches.forEach(match => {
+                const variableId = match[1];
+                if (this.variableService.hasVariable(variableId)) {
+                    result.globalVariables.push(variableId);
+                } else {
+                    result.agentSpecificVariables.push(variableId);
+                }
+            });
+
+            const functionMatches = [...prompt.matchAll(PROMPT_FUNCTION_REGEX)];
+            functionMatches.forEach(match => {
+                const functionId = match[1];
+                result.functions.push(functionId);
+            });
+
+        });
+        return result;
+    }
+
     protected showVariableConfigurationTab(): void {
         this.aiConfigurationSelectionService.selectConfigurationTab(AIVariableConfigurationWidget.ID);
     }

diff --git a/packages/ai-core/src/browser/frontend-prompt-customization-service.ts b/packages/ai-core/src/browser/frontend-prompt-customization-service.ts
@@ -14,7 +14,7 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 
-import { DisposableCollection, URI } from '@theia/core';
+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 } from '../common';
@@ -48,6 +48,9 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati
 
     protected toDispose = new DisposableCollection();
 
+    private readonly onDidChangePromptEmitter = new Emitter<string>();
+    readonly onDidChangePrompt: Event<string> = this.onDidChangePromptEmitter.event;
+
     @postConstruct()
     protected init(): void {
         this.preferences.onPreferenceChanged(event => {
@@ -85,6 +88,8 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati
                             _templates.set(this.removePromptTemplateSuffix(updatedFile.resource.path.name), filecontent.value);
                         }
                     }
+                    const id = this.removePromptTemplateSuffix(new URI(child).path.name);
+                    this.onDidChangePromptEmitter.fire(id);
                 }
             }
 

diff --git a/packages/ai-core/src/browser/style/index.css b/packages/ai-core/src/browser/style/index.css
@@ -66,14 +66,16 @@
 }
 
 #ai-variable-configuration-container-widget .variable-references,
-#ai-agent-configuration-container-widget .variable-references {
+#ai-agent-configuration-container-widget .variable-references,
+#ai-agent-configuration-container-widget .function-references {
   margin-left: 0.5rem;
   padding: 0.5rem;
   border-left: solid 1px var(--theia-tree-indentGuidesStroke);
 }
 
 #ai-variable-configuration-container-widget .variable-reference,
-#ai-agent-configuration-container-widget .variable-reference {
+#ai-agent-configuration-container-widget .variable-reference,
+#ai-agent-configuration-container-widget .function-reference {
   display: flex;
   flex-direction: row;
   align-items: center;

diff --git a/packages/ai-core/src/common/agent.ts b/packages/ai-core/src/common/agent.ts
@@ -17,6 +17,12 @@
 import { LanguageModelRequirement } from './language-model';
 import { PromptTemplate } from './prompt-service';
 
+export interface AgentSpecificVariables {
+    name: string;
+    description: string;
+    usedInPrompt: boolean;
+}
+
 export const Agent = Symbol('Agent');
 /**
  * Agents represent the main functionality of the AI system. They are responsible for processing user input, collecting information from the environment,
@@ -44,7 +50,7 @@ export interface Agent {
     /** A markdown description of its functionality and its privacy-relevant requirements, including function call handlers that access some data autonomously. */
     readonly description: string;
 
-    /** The list of variable identifiers this agent needs to clarify its context requirements. See #39. */
+    /** The list of global variable identifiers this agent needs to clarify its context requirements. See #39. */
     readonly variables: string[];
 
     /** The prompt templates introduced and used by this agent. */
@@ -55,4 +61,10 @@ export interface Agent {
 
     /** A list of tags to filter agents and to display capabilities in the UI */
     readonly tags?: String[];
+
+    /** The list of local variable identifiers this agent needs to clarify its context requirements. */
+    readonly agentSpecificVariables: AgentSpecificVariables[];
+
+    /** The list of global function identifiers this agent needs to clarify its context requirements. */
+    readonly functions: string[];
 }
diff --git a/packages/ai-core/src/common/prompt-service.ts b/packages/ai-core/src/common/prompt-service.ts
@@ -14,7 +14,7 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 
-import { URI } from '@theia/core';
+import { URI, Event } from '@theia/core';
 import { inject, injectable, optional } from '@theia/core/shared/inversify';
 import { AIVariableService } from './variable-service';
 import { ToolInvocationRegistry } from './tool-invocation-registry';
@@ -104,6 +104,11 @@ export interface PromptCustomizationService {
      * @param uri the uri of the template file
      */
     getTemplateIDFromURI(uri: URI): string | undefined;
+
+    /**
+     * Event which is fired when the prompt template is changed.
+     */
+    readonly onDidChangePrompt: Event<string>;
 }
 
 @injectable()

diff --git a/packages/ai-terminal/src/browser/ai-terminal-agent.ts b/packages/ai-terminal/src/browser/ai-terminal-agent.ts
@@ -40,6 +40,13 @@ export class AiTerminalAgent implements Agent {
         Based on the user\'s request, it suggests commands and allows the user to directly paste and execute them in the terminal. \
         It accesses the current directory, environment and the recent terminal output of the terminal session to provide context-aware assistance';
     variables = [];
+    functions = [];
+    agentSpecificVariables = [
+        { name: 'userRequest', usedInPrompt: true, description: 'The user\'s question or request.' },
+        { name: 'shell', usedInPrompt: true, description: 'The shell being used, e.g., /usr/bin/zsh.' },
+        { name: 'cwd', usedInPrompt: true, description: 'The current working directory.' },
+        { name: 'recentTerminalContents', usedInPrompt: true, description: 'The last 0 to 50 recent lines visible in the terminal.' }
+    ];
     promptTemplates = [
         {
             id: 'terminal-system',