Skip to content

Commit

Permalink
chore: streamline AI agent code (#14859)
Browse files Browse the repository at this point in the history 
Fix #14818
  • Loading branch information
@eneufeld
eneufeld authored Feb 13, 2025
1 parent 9a20727 commit ee0e526
Show file tree
Hide file tree
Showing 9 changed files with 148 additions and 226 deletions.
36 changes: 13 additions & 23 deletions examples/api-samples/src/browser/chat/ask-and-continue-chat-agent-contribution.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ import {
ChatRequestModelImpl,
lastProgressMessage,
QuestionResponseContentImpl,
SystemMessageDescription,
unansweredQuestions
} from '@theia/ai-chat';
import { Agent, PromptTemplate } from '@theia/ai-core';
Expand Down Expand Up @@ -109,14 +108,19 @@ Do not ask further questions once the text contains 5 or more "Question/Answer"
* This is a very simple example agent that asks questions and continues the conversation based on the user's answers.
*/
@injectable()
export class AskAndContinueChatAgent extends AbstractStreamParsingChatAgent implements ChatAgent {
override id = 'AskAndContinue';
readonly name = 'AskAndContinue';
override defaultLanguageModelPurpose = 'chat';
readonly description = 'This chat will ask questions related to the input and continues after that.';
readonly variables = [];
readonly agentSpecificVariables = [];
readonly functions = [];
export class AskAndContinueChatAgent extends AbstractStreamParsingChatAgent {
id = 'AskAndContinue';
name = 'AskAndContinue';
override description = 'This chat will ask questions related to the input and continues after that.';
protected defaultLanguageModelPurpose = 'chat';
override languageModelRequirements = [
{
purpose: 'chat',
identifier: 'openai/gpt-4o',
}
];
override promptTemplates = [systemPrompt];
protected override systemPromptId: string | undefined = systemPrompt.id;

@postConstruct()
addContentMatchers(): void {
Expand All @@ -133,20 +137,6 @@ export class AskAndContinueChatAgent extends AbstractStreamParsingChatAgent impl
});
}

override languageModelRequirements = [
{
purpose: 'chat',
identifier: 'openai/gpt-4o',
}
];

readonly promptTemplates = [systemPrompt];

protected override async getSystemMessageDescription(): Promise<SystemMessageDescription | undefined> {
const resolvedPrompt = await this.promptService.getPrompt(systemPrompt.id);
return resolvedPrompt ? SystemMessageDescription.fromResolvedPromptTemplate(resolvedPrompt) : undefined;
}

protected override async onResponseComplete(request: ChatRequestModelImpl): Promise<void> {
const unansweredQs = unansweredQuestions(request);
if (unansweredQs.length < 1) {
Expand Down
16 changes: 6 additions & 10 deletions examples/api-samples/src/browser/chat/change-set-chat-agent-contribution.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import {
SystemMessageDescription
} from '@theia/ai-chat';
import { ChangeSetFileElementFactory } from '@theia/ai-chat/lib/browser/change-set-file-element';
import { Agent, PromptTemplate } from '@theia/ai-core';
import { Agent, LanguageModelRequirement } from '@theia/ai-core';
import { URI } from '@theia/core';
import { inject, injectable, interfaces } from '@theia/core/shared/inversify';
import { FileService } from '@theia/filesystem/lib/browser/file-service';
Expand All @@ -39,16 +39,12 @@ export function bindChangeSetChatAgentContribution(bind: interfaces.Bind): void
* This is a test agent demonstrating how to create change sets in AI chats.
*/
@injectable()
export class ChangeSetChatAgent extends AbstractStreamParsingChatAgent implements ChatAgent {
override id = 'ChangeSet';
export class ChangeSetChatAgent extends AbstractStreamParsingChatAgent {
readonly id = 'ChangeSet';
readonly name = 'ChangeSet';
override defaultLanguageModelPurpose = 'chat';
readonly description = 'This chat will create and modify a change set.';
readonly variables = [];
readonly agentSpecificVariables = [];
readonly functions = [];
override languageModelRequirements = [];
promptTemplates: PromptTemplate[] = [];
readonly defaultLanguageModelPurpose = 'chat';
override readonly description = 'This chat will create and modify a change set.';
override languageModelRequirements: LanguageModelRequirement[] = [];

@inject(WorkspaceService)
protected readonly workspaceService: WorkspaceService;
Expand Down
41 changes: 27 additions & 14 deletions packages/ai-chat/src/common/chat-agents.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,15 @@
// Partially copied from https://github.com/microsoft/vscode/blob/a2cab7255c0df424027be05d58e1b7b941f4ea60/src/vs/workbench/contrib/chat/common/chatAgents.ts

import {
AgentSpecificVariables,
CommunicationRecordingService,
getTextOfResponse,
LanguageModel,
LanguageModelRequirement,
LanguageModelResponse,
LanguageModelStreamResponse,
PromptService,
PromptTemplate,
ResolvedPromptTemplate,
ToolRequest,
} from '@theia/ai-core';
Expand All @@ -39,7 +41,7 @@ import {
MessageActor,
} from '@theia/ai-core/lib/common';
import { CancellationToken, ContributionProvider, ILogger, isArray } from '@theia/core';
import { inject, injectable, named, postConstruct, unmanaged } from '@theia/core/shared/inversify';
import { inject, injectable, named, postConstruct } from '@theia/core/shared/inversify';
import { ChatAgentService } from './chat-agent-service';
import {
ChatModel,
Expand Down Expand Up @@ -119,7 +121,7 @@ export interface ChatAgent extends Agent {
}

@injectable()
export abstract class AbstractChatAgent {
export abstract class AbstractChatAgent implements ChatAgent {
@inject(LanguageModelRegistry) protected languageModelRegistry: LanguageModelRegistry;
@inject(ILogger) protected logger: ILogger;
@inject(CommunicationRecordingService) protected recordingService: CommunicationRecordingService;
Expand All @@ -128,21 +130,26 @@ export abstract class AbstractChatAgent {

@inject(ContributionProvider) @named(ResponseContentMatcherProvider)
protected contentMatcherProviders: ContributionProvider<ResponseContentMatcherProvider>;
protected additionalToolRequests: ToolRequest[] = [];
protected contentMatchers: ResponseContentMatcher[] = [];

@inject(DefaultResponseContentFactory)
protected defaultContentFactory: DefaultResponseContentFactory;

constructor(
@unmanaged() public id: string,
@unmanaged() public languageModelRequirements: LanguageModelRequirement[],
@unmanaged() protected defaultLanguageModelPurpose: string,
@unmanaged() public iconClass: string = 'codicon codicon-copilot',
@unmanaged() public locations: ChatAgentLocation[] = ChatAgentLocation.ALL,
@unmanaged() public tags: string[] = ['Chat'],
@unmanaged() public defaultLogging: boolean = true) {
}
readonly abstract id: string;
readonly abstract name: string;
readonly abstract languageModelRequirements: LanguageModelRequirement[];
iconClass: string = 'codicon codicon-copilot';
locations: ChatAgentLocation[] = ChatAgentLocation.ALL;
tags: string[] = ['Chat'];
description: string = '';
variables: string[] = [];
promptTemplates: PromptTemplate[] = [];
agentSpecificVariables: AgentSpecificVariables[] = [];
functions: string[] = [];
protected readonly abstract defaultLanguageModelPurpose: string;
protected defaultLogging: boolean = true;
protected systemPromptId: string | undefined = undefined;
protected additionalToolRequests: ToolRequest[] = [];
protected contentMatchers: ResponseContentMatcher[] = [];

@postConstruct()
init(): void {
Expand Down Expand Up @@ -236,7 +243,13 @@ export abstract class AbstractChatAgent {
return languageModel;
}

protected abstract getSystemMessageDescription(): Promise<SystemMessageDescription | undefined>;
protected async getSystemMessageDescription(): Promise<SystemMessageDescription | undefined> {
if (this.systemPromptId === undefined) {
return undefined;
}
const resolvedPrompt = await this.promptService.getPrompt(this.systemPromptId);
return resolvedPrompt ? SystemMessageDescription.fromResolvedPromptTemplate(resolvedPrompt) : undefined;
}

protected async getMessages(
model: ChatModel, includeResponseInProgress = false
Expand Down
28 changes: 8 additions & 20 deletions packages/ai-chat/src/common/custom-chat-agent.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,29 +14,17 @@
// 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 { LanguageModelRequirement } from '@theia/ai-core';
import { AbstractStreamParsingChatAgent } 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;
}
export class CustomChatAgent extends AbstractStreamParsingChatAgent {
id: string = 'CustomChatAgent';
name: string = 'CustomChatAgent';
languageModelRequirements: LanguageModelRequirement[] = [{ purpose: 'chat' }];
protected defaultLanguageModelPurpose: string = 'chat';
protected override systemPromptId: string = `${this.name}_prompt`;

set prompt(prompt: string) {
this.promptTemplates.push({ id: `${this.name}_prompt`, template: prompt });
Expand Down
40 changes: 14 additions & 26 deletions packages/ai-ide-agents/src/browser/coder-agent.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,38 +13,26 @@
//
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
// *****************************************************************************
import { AbstractStreamParsingChatAgent, ChatAgent, SystemMessageDescription } from '@theia/ai-chat/lib/common';
import { AgentSpecificVariables, PromptTemplate } from '@theia/ai-core';
import { AbstractStreamParsingChatAgent } from '@theia/ai-chat/lib/common';
import { injectable } from '@theia/core/shared/inversify';
import { FILE_CONTENT_FUNCTION_ID, GET_WORKSPACE_FILE_LIST_FUNCTION_ID, GET_WORKSPACE_DIRECTORY_STRUCTURE_FUNCTION_ID } from '../common/workspace-functions';
import { CODER_REPLACE_PROMPT_TEMPLATE_ID, getCoderReplacePromptTemplate } from '../common/coder-replace-prompt-template';
import { WriteChangeToFileProvider } from './file-changeset-functions';
import { LanguageModelRequirement } from '@theia/ai-core';

@injectable()
export class CoderAgent extends AbstractStreamParsingChatAgent implements ChatAgent {
name: string;
description: string;
promptTemplates: PromptTemplate[];
variables: never[];
readonly agentSpecificVariables: AgentSpecificVariables[];
readonly functions: string[];
export class CoderAgent extends AbstractStreamParsingChatAgent {
id: string = 'Coder';
name = 'Coder';
languageModelRequirements: LanguageModelRequirement[] = [{
purpose: 'chat',
identifier: 'openai/gpt-4o',
}];
protected defaultLanguageModelPurpose: string = 'chat';

constructor() {
super('Coder', [{
purpose: 'chat',
identifier: 'openai/gpt-4o',
}], 'chat');
this.name = 'Coder';
this.description = 'An AI assistant integrated into Theia IDE, designed to assist software developers with code tasks.';
this.promptTemplates = [getCoderReplacePromptTemplate(true), getCoderReplacePromptTemplate(false)];
this.variables = [];
this.agentSpecificVariables = [];
this.functions = [GET_WORKSPACE_DIRECTORY_STRUCTURE_FUNCTION_ID, GET_WORKSPACE_FILE_LIST_FUNCTION_ID, FILE_CONTENT_FUNCTION_ID, WriteChangeToFileProvider.ID];
}

protected override async getSystemMessageDescription(): Promise<SystemMessageDescription | undefined> {
const resolvedPrompt = await this.promptService.getPrompt(CODER_REPLACE_PROMPT_TEMPLATE_ID);
return resolvedPrompt ? SystemMessageDescription.fromResolvedPromptTemplate(resolvedPrompt) : undefined;
}
override description = 'An AI assistant integrated into Theia IDE, designed to assist software developers with code tasks.';
override promptTemplates = [getCoderReplacePromptTemplate(true), getCoderReplacePromptTemplate(false)];
override functions = [GET_WORKSPACE_DIRECTORY_STRUCTURE_FUNCTION_ID, GET_WORKSPACE_FILE_LIST_FUNCTION_ID, FILE_CONTENT_FUNCTION_ID, WriteChangeToFileProvider.ID];
protected override systemPromptId: string | undefined = CODER_REPLACE_PROMPT_TEMPLATE_ID;

}
44 changes: 15 additions & 29 deletions packages/ai-ide-agents/src/browser/workspace-agent.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,41 +13,27 @@
//
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
// *****************************************************************************
import { AbstractStreamParsingChatAgent, ChatAgent, SystemMessageDescription } from '@theia/ai-chat/lib/common';
import { AgentSpecificVariables, PromptTemplate, ToolInvocationRegistry } from '@theia/ai-core';
import { inject, injectable } from '@theia/core/shared/inversify';
import { AbstractStreamParsingChatAgent } from '@theia/ai-chat/lib/common';
import { LanguageModelRequirement } from '@theia/ai-core';
import { injectable } from '@theia/core/shared/inversify';
import { workspacePromptTemplate } from '../common/workspace-prompt-template';
import { FILE_CONTENT_FUNCTION_ID, GET_WORKSPACE_FILE_LIST_FUNCTION_ID } from '../common/workspace-functions';

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

@inject(ToolInvocationRegistry)
protected toolInvocationRegistry: ToolInvocationRegistry;
name = 'Workspace';
id = 'Workspace';
languageModelRequirements: LanguageModelRequirement[] = [{
purpose: 'chat',
identifier: 'openai/gpt-4o',
}];
protected defaultLanguageModelPurpose: string = 'chat';

constructor() {
super('Workspace', [{
purpose: 'chat',
identifier: 'openai/gpt-4o',
}], 'chat');
this.name = 'Workspace';
this.description = 'This agent can access the users workspace, it can get a list of all available files and retrieve their content. \
override description = 'This agent can access the users workspace, it can get a list of all available files and retrieve their content. \
It can therefore answer questions about the current project, project files and source code in the workspace, such as how to build the project, \
where to put source code, where to find specific code or configurations, etc.';
this.promptTemplates = [workspacePromptTemplate];
this.variables = [];
this.agentSpecificVariables = [];
this.functions = [GET_WORKSPACE_FILE_LIST_FUNCTION_ID, FILE_CONTENT_FUNCTION_ID];
}

protected override async getSystemMessageDescription(): Promise<SystemMessageDescription | undefined> {
const resolvedPrompt = await this.promptService.getPrompt(workspacePromptTemplate.id);
return resolvedPrompt ? SystemMessageDescription.fromResolvedPromptTemplate(resolvedPrompt) : undefined;
}
override promptTemplates = [workspacePromptTemplate];
override functions = [GET_WORKSPACE_FILE_LIST_FUNCTION_ID, FILE_CONTENT_FUNCTION_ID];
protected override systemPromptId: string | undefined = workspacePromptTemplate.id;
}
53 changes: 21 additions & 32 deletions packages/ai-ide-agents/src/common/command-chat-agents.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,8 @@
// *****************************************************************************

import { inject, injectable } from '@theia/core/shared/inversify';
import { AbstractTextToModelParsingChatAgent, ChatAgent, SystemMessageDescription } from '@theia/ai-chat/lib/common/chat-agents';
import {
PromptTemplate,
AgentSpecificVariables
} from '@theia/ai-core';
import { AbstractTextToModelParsingChatAgent, SystemMessageDescription } from '@theia/ai-chat/lib/common/chat-agents';
import { LanguageModelRequirement, PromptTemplate } from '@theia/ai-core';
import {
ChatRequestModelImpl,
ChatResponseContent,
Expand Down Expand Up @@ -250,38 +247,30 @@ interface ParsedCommand {
}

@injectable()
export class CommandChatAgent extends AbstractTextToModelParsingChatAgent<ParsedCommand> implements ChatAgent {
export class CommandChatAgent extends AbstractTextToModelParsingChatAgent<ParsedCommand> {
@inject(CommandRegistry)
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', [{
purpose: 'command',
identifier: 'openai/gpt-4o',
}], 'command');
this.name = 'Command';
this.description = 'This agent is aware of all commands that the user can execute within the Theia IDE, the tool that the user is currently working with. \
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> {
id: string = 'Command';
name = 'Command';
languageModelRequirements: LanguageModelRequirement[] = [{
purpose: 'command',
identifier: 'openai/gpt-4o',
}];
protected defaultLanguageModelPurpose: string = 'command';

override description = 'This agent is aware of all commands that the user can execute within the Theia IDE, the tool that the user is currently working with. \
Based on the user request, it can find the right command and then let the user execute it.';
override promptTemplates = [commandTemplate];
override agentSpecificVariables = [{
name: 'command-ids',
description: 'The list of available commands in Theia.',
usedInPrompt: true
}];

protected override async getSystemMessageDescription(): Promise<SystemMessageDescription | undefined> {
const knownCommands: string[] = [];
for (const command of this.commandRegistry.getAllCommands()) {
knownCommands.push(`${command.id}: ${command.label}`);
Expand Down
Loading

0 comments on commit ee0e526

Please sign in to comment.