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 16, 2024
1 parent 1fea534 commit a431a3a
Show file tree
Hide file tree
Showing 11 changed files with 169 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,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 {

Expand All @@ -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()
Expand All @@ -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());
Expand Down Expand Up @@ -91,6 +114,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 }}>{agent.name}</div>
<div style={{ paddingBottom: 10 }}>{agent.description}</div>
Expand All @@ -100,16 +128,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 @@ -125,9 +143,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);
}
Expand Down
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
14 changes: 13 additions & 1 deletion packages/ai-core/src/common/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -37,12 +43,18 @@ 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. */
readonly promptTemplates: PromptTemplate[];

/** Required language models. This includes the purpose and optional language model selector arguments. See #47. */
readonly languageModelRequirements: LanguageModelRequirement[];

/** 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[];
}
7 changes: 6 additions & 1 deletion packages/ai-core/src/common/prompt-service.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 { 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';
Expand Down Expand Up @@ -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()
Expand Down
7 changes: 7 additions & 0 deletions packages/ai-terminal/src/browser/ai-terminal-agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
Loading

0 comments on commit a431a3a

Please sign in to comment.