Skip to content

Commit

Permalink
feat: show variables and function on ai agent configuration
Browse files Browse the repository at this point in the history
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
  • Loading branch information
eneufeld committed Sep 19, 2024
1 parent 4a0228b commit 28cf1f2
Show file tree
Hide file tree
Showing 11 changed files with 222 additions and 22 deletions.
11 changes: 10 additions & 1 deletion packages/ai-chat/src/common/command-chat-agents.ts
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
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
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
11 changes: 9 additions & 2 deletions packages/ai-code-completion/src/common/code-completion-agent.ts
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 { CancellationToken, generateUuid, ILogger } from '@theia/core';
Expand All @@ -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;
Expand Down
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>;
};
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 { 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';
Expand Down Expand Up @@ -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 => {
Expand Down Expand Up @@ -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);
}
}

Expand Down
6 changes: 4 additions & 2 deletions packages/ai-core/src/browser/style/index.css
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading

0 comments on commit 28cf1f2

Please sign in to comment.