Skip to content

Commit

Permalink
feat: show variables and function on ai agent configuration (#14177)
Browse files Browse the repository at this point in the history 
* 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

* fix missing emit on propmt template folder switch
  • Loading branch information
@eneufeld
eneufeld authored Sep 20, 2024
1 parent 0f1ebdf commit 3a339a7
Show file tree
Hide file tree
Showing 11 changed files with 243 additions and 42 deletions.
11 changes: 10 additions & 1 deletion packages/ai-chat/src/common/command-chat-agents.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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', [{
Expand All @@ -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> {
Expand Down
9 changes: 7 additions & 2 deletions packages/ai-chat/src/common/orchestrator-chat-agent.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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',
Expand All @@ -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)
Expand Down
5 changes: 5 additions & 0 deletions packages/ai-chat/src/common/universal-chat-agent.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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', [{
Expand All @@ -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> {
Expand Down
13 changes: 10 additions & 3 deletions packages/ai-code-completion/src/common/code-completion-agent.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 { generateUuid, ILogger } from '@theia/core';
Expand Down Expand Up @@ -116,8 +116,6 @@ export class CodeCompletionAgentImpl implements CodeCompletionAgent {
enableForwardStability: true,
};
}
tags?: String[] | undefined;
variables: string[] = [];

@inject(ILogger)
@named('code-completion-agent')
Expand Down Expand Up @@ -154,4 +152,13 @@ Only return the exact replacement for [[MARKER]] to complete the snippet.`,
identifier: 'openai/gpt-4o',
},
];
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: 'textUntilCurrentPosition', usedInPrompt: true, description: 'The code before the current position of the cursor.' },
{ name: 'textAfterCurrentPosition', usedInPrompt: true, description: 'The code after the current position of the cursor.' }
];
readonly tags?: String[] | undefined;
}
160 changes: 149 additions & 11 deletions packages/ai-core/src/browser/ai-configuration/agent-configuration-widget.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,29 @@
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,
} 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 {

Expand All @@ -46,6 +61,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()
Expand All @@ -62,6 +83,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());
Expand Down Expand Up @@ -98,6 +120,10 @@ export class AIAgentConfigurationWidget extends ReactWidget {

const enabled = this.agentService.isEnabled(agent.id);

const parsedPromptParts = this.parsePromptTemplatesForVariableAndFunction(agent);
const globalVariables = Array.from(new Set([...parsedPromptParts.globalVariables, ...agent.variables]));
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>
Expand All @@ -107,16 +133,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
Expand All @@ -132,9 +148,59 @@ export class AIAgentConfigurationWidget extends ReactWidget {
aiSettingsService={this.aiSettingsService}
languageModelRegistry={this.languageModelRegistry} />
</div>
<div>
<span>Used Global Variables:</span>
<ul className='variable-references'>
<AgentGlobalVariables variables={globalVariables} showVariableConfigurationTab={this.showVariableConfigurationTab.bind(this)} />
</ul>
</div>
<div>
<span>Used agent-specific Variables:</span>
<ul className='variable-references'>
<AgentSpecificVariables
promptVariables={parsedPromptParts.agentSpecificVariables}
agent={agent}
/>
</ul>
</div>
<div>
<span>Used Functions:</span>
<ul className='function-references'>
<AgentFunctions functions={functions} />
</ul>
</div>
</div>;
}

private parsePromptTemplatesForVariableAndFunction(agent: Agent): ParsedPrompt {
const promptTemplates = agent.promptTemplates;
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 the variable is part of the variable service and not part of the agent specific variables then it is a global variable
if (this.variableService.hasVariable(variableId) &&
agent.agentSpecificVariables.find(v => v.name === variableId) === undefined) {
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);
}
Expand All @@ -159,3 +225,75 @@ export class AIAgentConfigurationWidget extends ReactWidget {
};

}
interface AgentGlobalVariablesProps {
variables: string[];
showVariableConfigurationTab: () => void;
}
const AgentGlobalVariables = ({ variables: globalVariables, showVariableConfigurationTab }: AgentGlobalVariablesProps) => {
if (globalVariables.length === 0) {
return <>None</>;
}
return <>
{globalVariables.map(variableId => <li key={variableId} className='theia-TreeNode theia-CompositeTreeNode theia-ExpandableTreeNode theia-mod-selected'>
<div key={variableId} onClick={() => { showVariableConfigurationTab(); }} className='variable-reference'>
<span>{variableId}</span>
<i className={codicon('chevron-right')}></i>
</div></li>)}

</>;
};

interface AgentFunctionsProps {
functions: string[];
}
const AgentFunctions = ({ functions }: AgentFunctionsProps) => {
if (functions.length === 0) {
return <>None</>;
}
return <>
{functions.map(functionId => <li key={functionId} className='variable-reference'>
<span>{functionId}</span>
</li>)}
</>;
};

interface AgentSpecificVariablesProps {
promptVariables: string[];
agent: Agent;
}
const AgentSpecificVariables = ({ promptVariables, agent }: AgentSpecificVariablesProps) => {
const agentDefinedVariablesName = agent.agentSpecificVariables.map(v => v.name);
const variables = Array.from(new Set([...promptVariables, ...agentDefinedVariablesName]));
if (variables.length === 0) {
return <>None</>;
}
return <>
{variables.map(variableId =>
<AgentSpecifcVariable
key={variableId}
variableId={variableId}
agent={agent}
promptVariables={promptVariables} />

)}
</>;
};
interface AgentSpecifcVariableProps {
variableId: string;
agent: Agent;
promptVariables: string[];
}
const AgentSpecifcVariable = ({ variableId, agent, promptVariables }: AgentSpecifcVariableProps) => {
const agentDefinedVariable = agent.agentSpecificVariables.find(v => v.name === variableId);
const undeclared = agentDefinedVariable === undefined;
const notUsed = !promptVariables.includes(variableId) && agentDefinedVariable?.usedInPrompt === true;
return <li key={variableId}>
<div><span>Name:</span> <span>{variableId}</span></div>
{undeclared ? <div><span>Undeclared</span></div> :
(<>
<div><span>Description:</span> <span>{agentDefinedVariable.description}</span></div>
{notUsed && <div>Not used in prompt</div>}
</>)}
<hr />
</li>;
};
Loading

0 comments on commit 3a339a7

Please sign in to comment.