Skip to content

Commit

Permalink
feat: add support for custom agents (#14301)
Browse files Browse the repository at this point in the history
This adds support for custom agents which are basically a custom
system prompt with additional metadata (id, name and description).
This features allows to add very specific agents without coding.
All features, like variable and functions are supported.
  • Loading branch information
eneufeld authored Oct 18, 2024
1 parent fc1b88e commit aaaa73f
Show file tree
Hide file tree
Showing 13 changed files with 329 additions and 9 deletions.
29 changes: 25 additions & 4 deletions packages/ai-chat/src/browser/ai-chat-frontend-module.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,9 @@
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
// *****************************************************************************

import { Agent, AIVariableContribution } from '@theia/ai-core/lib/common';
import { Agent, AgentService, AIVariableContribution } from '@theia/ai-core/lib/common';
import { bindContributionProvider } from '@theia/core';
import { PreferenceContribution } from '@theia/core/lib/browser';
import { FrontendApplicationContribution, PreferenceContribution } from '@theia/core/lib/browser';
import { ContainerModule } from '@theia/core/shared/inversify';
import {
ChatAgent,
Expand All @@ -27,13 +27,16 @@ import {
ChatService,
DefaultChatAgentId
} from '../common';
import { ChatAgentsVariableContribution } from '../common/chat-agents-variable-contribution';
import { CommandChatAgent } from '../common/command-chat-agents';
import { CustomChatAgent } from '../common/custom-chat-agent';
import { OrchestratorChatAgent, OrchestratorChatAgentId } from '../common/orchestrator-chat-agent';
import { DefaultResponseContentFactory, DefaultResponseContentMatcherProvider, ResponseContentMatcherProvider } from '../common/response-content-matcher';
import { UniversalChatAgent } from '../common/universal-chat-agent';
import { aiChatPreferences } from './ai-chat-preferences';
import { ChatAgentsVariableContribution } from '../common/chat-agents-variable-contribution';
import { AICustomAgentsFrontendApplicationContribution } from './custom-agent-frontend-application-contribution';
import { FrontendChatServiceImpl } from './frontend-chat-service';
import { DefaultResponseContentMatcherProvider, DefaultResponseContentFactory, ResponseContentMatcherProvider } from '../common/response-content-matcher';
import { CustomAgentFactory } from './custom-agent-factory';

export default new ContainerModule(bind => {
bindContributionProvider(bind, Agent);
Expand Down Expand Up @@ -69,4 +72,22 @@ export default new ContainerModule(bind => {
bind(ChatAgent).toService(CommandChatAgent);

bind(PreferenceContribution).toConstantValue({ schema: aiChatPreferences });

bind(CustomChatAgent).toSelf();
bind(CustomAgentFactory).toFactory<CustomChatAgent, [string, string, string, string, string]>(
ctx => (id: string, name: string, description: string, prompt: string, defaultLLM: string) => {
const agent = ctx.container.get<CustomChatAgent>(CustomChatAgent);
agent.id = id;
agent.name = name;
agent.description = description;
agent.prompt = prompt;
agent.languageModelRequirements = [{
purpose: 'chat',
identifier: defaultLLM,
}];
ctx.container.get<ChatAgentService>(ChatAgentService).registerChatAgent(agent);
ctx.container.get<AgentService>(AgentService).registerAgent(agent);
return agent;
});
bind(FrontendApplicationContribution).to(AICustomAgentsFrontendApplicationContribution).inSingletonScope();
});
20 changes: 20 additions & 0 deletions packages/ai-chat/src/browser/custom-agent-factory.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
// *****************************************************************************
// Copyright (C) 2024 EclipseSource GmbH.
//
// This program and the accompanying materials are made available under the
// terms of the Eclipse Public License v. 2.0 which is available at
// http://www.eclipse.org/legal/epl-2.0.
//
// This Source Code may also be made available under the following Secondary
// Licenses when the conditions for such availability set forth in the Eclipse
// Public License v. 2.0 are satisfied: GNU General Public License, version 2
// with the GNU Classpath Exception which is available at
// https://www.gnu.org/software/classpath/license.html.
//
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
// *****************************************************************************

import { CustomChatAgent } from '../common';

export const CustomAgentFactory = Symbol('CustomAgentFactory');
export type CustomAgentFactory = (id: string, name: string, description: string, prompt: string, defaultLLM: string) => CustomChatAgent;
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
// *****************************************************************************
// Copyright (C) 2024 EclipseSource GmbH.
//
// This program and the accompanying materials are made available under the
// terms of the Eclipse Public License v. 2.0 which is available at
// http://www.eclipse.org/legal/epl-2.0.
//
// This Source Code may also be made available under the following Secondary
// Licenses when the conditions for such availability set forth in the Eclipse
// Public License v. 2.0 are satisfied: GNU General Public License, version 2
// with the GNU Classpath Exception which is available at
// https://www.gnu.org/software/classpath/license.html.
//
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
// *****************************************************************************

import { AgentService, CustomAgentDescription, PromptCustomizationService } from '@theia/ai-core';
import { FrontendApplicationContribution } from '@theia/core/lib/browser';
import { inject, injectable } from '@theia/core/shared/inversify';
import { ChatAgentService } from '../common';
import { CustomAgentFactory } from './custom-agent-factory';

@injectable()
export class AICustomAgentsFrontendApplicationContribution implements FrontendApplicationContribution {
@inject(CustomAgentFactory)
protected readonly customAgentFactory: CustomAgentFactory;

@inject(PromptCustomizationService)
protected readonly customizationService: PromptCustomizationService;

@inject(AgentService)
private readonly agentService: AgentService;

@inject(ChatAgentService)
private readonly chatAgentService: ChatAgentService;

private knownCustomAgents: Map<string, CustomAgentDescription> = new Map();
onStart(): void {
this.customizationService?.getCustomAgents().then(customAgents => {
customAgents.forEach(agent => {
this.customAgentFactory(agent.id, agent.name, agent.description, agent.prompt, agent.defaultLLM);
this.knownCustomAgents.set(agent.id, agent);
});
}).catch(e => {
console.error('Failed to load custom agents', e);
});
this.customizationService?.onDidChangeCustomAgents(() => {
this.customizationService?.getCustomAgents().then(customAgents => {
const customAgentsToAdd = customAgents.filter(agent =>
!this.knownCustomAgents.has(agent.id) || !CustomAgentDescription.equals(this.knownCustomAgents.get(agent.id)!, agent));
const customAgentIdsToRemove = [...this.knownCustomAgents.values()].filter(agent =>
!customAgents.find(a => CustomAgentDescription.equals(a, agent))).map(a => a.id);

// delete first so we don't have to deal with the case where we add and remove the same agentId
customAgentIdsToRemove.forEach(id => {
this.chatAgentService.unregisterChatAgent(id);
this.agentService.unregisterAgent(id);
this.knownCustomAgents.delete(id);
});
customAgentsToAdd
.forEach(agent => {
this.customAgentFactory(agent.id, agent.name, agent.description, agent.prompt, agent.defaultLLM);
this.knownCustomAgents.set(agent.id, agent);
});
}).catch(e => {
console.error('Failed to load custom agents', e);
});
});
}

onStop(): void {
}
}
15 changes: 15 additions & 0 deletions packages/ai-chat/src/common/chat-agent-service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,18 @@ export interface ChatAgentService {
* Returns all agents, including disabled ones.
*/
getAllAgents(): ChatAgent[];

/**
* Allows to register a chat agent programmatically.
* @param agent the agent to register
*/
registerChatAgent(agent: ChatAgent): void;

/**
* Allows to unregister a chat agent programmatically.
* @param agentId the agent id to unregister
*/
unregisterChatAgent(agentId: string): void;
}
@injectable()
export class ChatAgentServiceImpl implements ChatAgentService {
Expand All @@ -65,6 +77,9 @@ export class ChatAgentServiceImpl implements ChatAgentService {
registerChatAgent(agent: ChatAgent): void {
this._agents.push(agent);
}
unregisterChatAgent(agentId: string): void {
this._agents = this._agents.filter(a => a.id !== agentId);
}

getAgent(id: string): ChatAgent | undefined {
if (!this._agentIsEnabled(id)) {
Expand Down
44 changes: 44 additions & 0 deletions packages/ai-chat/src/common/custom-chat-agent.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
// *****************************************************************************
// Copyright (C) 2024 EclipseSource GmbH.
//
// This program and the accompanying materials are made available under the
// terms of the Eclipse Public License v. 2.0 which is available at
// http://www.eclipse.org/legal/epl-2.0.
//
// This Source Code may also be made available under the following Secondary
// Licenses when the conditions for such availability set forth in the Eclipse
// Public License v. 2.0 are satisfied: GNU General Public License, version 2
// with the GNU Classpath Exception which is available at
// https://www.gnu.org/software/classpath/license.html.
//
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
// *****************************************************************************

import { AgentSpecificVariables, PromptTemplate } from '@theia/ai-core';
import { AbstractStreamParsingChatAgent, ChatAgent, SystemMessageDescription } from './chat-agents';
import { injectable } from '@theia/core/shared/inversify';

@injectable()
export class CustomChatAgent
extends AbstractStreamParsingChatAgent
implements ChatAgent {
name: string;
description: string;
readonly variables: string[] = [];
readonly functions: string[] = [];
readonly promptTemplates: PromptTemplate[] = [];
readonly agentSpecificVariables: AgentSpecificVariables[] = [];

constructor(
) {
super('CustomChatAgent', [{ purpose: 'chat' }], 'chat');
}
protected override async getSystemMessageDescription(): Promise<SystemMessageDescription | undefined> {
const resolvedPrompt = await this.promptService.getPrompt(`${this.name}_prompt`);
return resolvedPrompt ? SystemMessageDescription.fromResolvedPromptTemplate(resolvedPrompt) : undefined;
}

set prompt(prompt: string) {
this.promptTemplates.push({ id: `${this.name}_prompt`, template: prompt });
}
}
7 changes: 4 additions & 3 deletions packages/ai-chat/src/common/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,13 @@
//
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
// *****************************************************************************
export * from './chat-agent-service';
export * from './chat-agents';
export * from './chat-agent-service';
export * from './chat-model';
export * from './parsed-chat-request';
export * from './chat-request-parser';
export * from './chat-service';
export * from './command-chat-agents';
export * from './universal-chat-agent';
export * from './custom-chat-agent';
export * from './parsed-chat-request';
export * from './orchestrator-chat-agent';
export * from './universal-chat-agent';
2 changes: 2 additions & 0 deletions packages/ai-core/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
"@theia/output": "1.54.0",
"@theia/variable-resolver": "1.54.0",
"@theia/workspace": "1.54.0",
"@types/js-yaml": "^4.0.9",
"js-yaml": "^4.1.0",
"minimatch": "^5.1.0",
"tslib": "^2.6.2"
},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,7 @@ export class AIAgentConfigurationWidget extends ReactWidget {

this.aiSettingsService.onDidChange(() => this.update());
this.aiConfigurationSelectionService.onDidAgentChange(() => this.update());
this.agentService.onDidChangeAgents(() => this.update());
this.update();
}

Expand All @@ -100,6 +101,9 @@ export class AIAgentConfigurationWidget extends ReactWidget {
</li>
)}
</ul>
<div className='configuration-agents-add'>
<button style={{ marginLeft: 0 }} className='theia-button main' onClick={() => this.addCustomAgent()}>Add Custom Agent</button>
</div>
</div>
<div className='configuration-agent-panel preferences-editor-widget'>
{this.renderAgentDetails()}
Expand Down Expand Up @@ -205,6 +209,10 @@ export class AIAgentConfigurationWidget extends ReactWidget {
this.aiConfigurationSelectionService.selectConfigurationTab(AIVariableConfigurationWidget.ID);
}

protected addCustomAgent(): void {
this.promptCustomizationService.openCustomAgentYaml();
}

protected setActiveAgent(agent: Agent): void {
this.aiConfigurationSelectionService.setActiveAgent(agent);
this.update();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,22 @@
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';
import { PromptCustomizationService, PromptTemplate, 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';

const templateEntry = {
id: 'my_agent',
name: 'My Agent',
description: 'This is an example agent. Please adapt the properties to fit your needs.',
prompt: 'You are an example agent. Be nice and helpful to the user.',
defaultLLM: 'openai/gpt-4o'
};

@injectable()
export class FrontendPromptCustomizationServiceImpl implements PromptCustomizationService {
Expand Down Expand Up @@ -51,6 +60,9 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati
private readonly onDidChangePromptEmitter = new Emitter<string>();
readonly onDidChangePrompt: Event<string> = this.onDidChangePromptEmitter.event;

private readonly onDidChangeCustomAgentsEmitter = new Emitter<void>();
readonly onDidChangeCustomAgents = this.onDidChangeCustomAgentsEmitter.event;

@postConstruct()
protected init(): void {
this.preferences.onPreferenceChanged(event => {
Expand All @@ -72,6 +84,9 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati

this.toDispose.push(this.fileService.watch(templateURI, { recursive: true, excludes: [] }));
this.toDispose.push(this.fileService.onDidFilesChange(async (event: FileChangesEvent) => {
if (event.changes.some(change => change.resource.toString().endsWith('customAgents.yml'))) {
this.onDidChangeCustomAgentsEmitter.fire();
}
// check deleted templates
for (const deletedFile of event.getDeleted()) {
if (this.trackedTemplateURIs.has(deletedFile.resource.toString())) {
Expand Down Expand Up @@ -103,6 +118,7 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati

}));

this.onDidChangeCustomAgentsEmitter.fire();
const stat = await this.fileService.resolve(templateURI);
if (stat.children === undefined) {
return;
Expand Down Expand Up @@ -194,4 +210,47 @@ export class FrontendPromptCustomizationServiceImpl implements PromptCustomizati
return undefined;
}

async getCustomAgents(): Promise<CustomAgentDescription[]> {
const customAgentYamlUri = (await this.getTemplatesDirectoryURI()).resolve('customAgents.yml');
const yamlExists = await this.fileService.exists(customAgentYamlUri);
if (!yamlExists) {
return [];
}
const filecontent = await this.fileService.read(customAgentYamlUri, { encoding: 'utf-8' });
try {
const doc = load(filecontent.value);
if (!Array.isArray(doc) || !doc.every(entry => CustomAgentDescription.is(entry))) {
console.debug('Invalid customAgents.yml file content');
return [];
}
const readAgents = doc as CustomAgentDescription[];
// make sure all agents are unique (id and name)
const uniqueAgentIds = new Set<string>();
const uniqueAgens: CustomAgentDescription[] = [];
readAgents.forEach(agent => {
if (uniqueAgentIds.has(agent.id)) {
return;
}
uniqueAgentIds.add(agent.id);
uniqueAgens.push(agent);
});
return uniqueAgens;
} catch (e) {
console.debug(e.message, e);
return [];
}
}

async openCustomAgentYaml(): Promise<void> {
const customAgentYamlUri = (await this.getTemplatesDirectoryURI()).resolve('customAgents.yml');
const content = dump([templateEntry]);
if (! await this.fileService.exists(customAgentYamlUri)) {
await this.fileService.createFile(customAgentYamlUri, BinaryBuffer.fromString(content));
} else {
const fileContent = (await this.fileService.readFile(customAgentYamlUri)).value;
await this.fileService.writeFile(customAgentYamlUri, BinaryBuffer.concat([fileContent, BinaryBuffer.fromString(content)]));
}
const openHandler = await this.openerService.getOpener(customAgentYamlUri);
openHandler.open(customAgentYamlUri);
}
}
5 changes: 5 additions & 0 deletions packages/ai-core/src/browser/style/index.css
Original file line number Diff line number Diff line change
Expand Up @@ -88,3 +88,8 @@
border-radius: calc(var(--theia-ui-padding) * 2 / 3);
background: hsla(0, 0%, 68%, 0.31);
}

.configuration-agents-add {
margin-top: 3em;
margin-left: 0;
}
Loading

0 comments on commit aaaa73f

Please sign in to comment.