Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: allow passing context to tool calls #14751

Merged
merged 1 commit into from
Jan 22, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions packages/ai-chat/src/browser/ai-chat-frontend-module.ts
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,14 @@ import { aiChatPreferences } from './ai-chat-preferences';
import { AICustomAgentsFrontendApplicationContribution } from './custom-agent-frontend-application-contribution';
import { FrontendChatServiceImpl } from './frontend-chat-service';
import { CustomAgentFactory } from './custom-agent-factory';
import { ChatToolRequestService } from '../common/chat-tool-request-service';

export default new ContainerModule(bind => {
bindContributionProvider(bind, Agent);
bindContributionProvider(bind, ChatAgent);

bind(ChatToolRequestService).toSelf().inSingletonScope();

bind(ChatAgentServiceImpl).toSelf().inSingletonScope();
bind(ChatAgentService).toService(ChatAgentServiceImpl);
bind(DefaultChatAgentId).toConstantValue({ id: OrchestratorChatAgentId });
Expand Down
27 changes: 11 additions & 16 deletions packages/ai-chat/src/common/chat-agents.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,6 @@ import { inject, injectable, named, postConstruct, unmanaged } from '@theia/core
import { ChatAgentService } from './chat-agent-service';
import {
ChatModel,
ChatRequestModel,
ChatRequestModelImpl,
ChatResponseContent,
ErrorChatResponseContentImpl,
Expand All @@ -53,6 +52,7 @@ import {
import { findFirstMatch, parseContents } from './parse-contents';
import { DefaultResponseContentFactory, ResponseContentMatcher, ResponseContentMatcherProvider } from './response-content-matcher';
import { ChatHistoryEntry } from './chat-history-entry';
import { ChatToolRequestService } from './chat-tool-request-service';

/**
* A conversation consists of a sequence of ChatMessages.
Expand Down Expand Up @@ -123,10 +123,12 @@ export abstract class AbstractChatAgent {
@inject(LanguageModelRegistry) protected languageModelRegistry: LanguageModelRegistry;
@inject(ILogger) protected logger: ILogger;
@inject(CommunicationRecordingService) protected recordingService: CommunicationRecordingService;
@inject(ChatToolRequestService) protected chatToolRequestService: ChatToolRequestService;
@inject(PromptService) protected promptService: PromptService;

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

@inject(DefaultResponseContentFactory)
Expand Down Expand Up @@ -171,7 +173,6 @@ export abstract class AbstractChatAgent {
);
}

const tools: Map<string, ToolRequest> = new Map();
if (systemMessageDescription) {
const systemMsg: ChatMessage = {
actor: 'system',
Expand All @@ -180,16 +181,19 @@ export abstract class AbstractChatAgent {
};
// insert system message at the beginning of the request messages
messages.unshift(systemMsg);
systemMessageDescription.functionDescriptions?.forEach((tool, id) => {
tools.set(id, tool);
});
}
this.getTools(request)?.forEach(tool => tools.set(tool.id, tool));

const systemMessageToolRequests = systemMessageDescription?.functionDescriptions?.values();
const tools = [
...this.chatToolRequestService.getChatToolRequests(request),
...this.chatToolRequestService.toChatToolRequests(systemMessageToolRequests ? Array.from(systemMessageToolRequests) : [], request),
...this.chatToolRequestService.toChatToolRequests(this.additionalToolRequests, request)
];

const languageModelResponse = await this.callLlm(
languageModel,
messages,
tools.size > 0 ? Array.from(tools.values()) : undefined,
tools.length > 0 ? tools : undefined,
request.response.cancellationToken
);
await this.addContentsToResponse(languageModelResponse, request);
Expand Down Expand Up @@ -258,15 +262,6 @@ export abstract class AbstractChatAgent {
return requestMessages;
}

/**
* @returns the list of tools used by this agent, or undefined if none is needed.
*/
protected getTools(request: ChatRequestModel): ToolRequest[] | undefined {
return request.message.toolRequests.size > 0
? [...request.message.toolRequests.values()]
: undefined;
}

protected async callLlm(
languageModel: LanguageModel,
messages: ChatMessage[],
Expand Down
15 changes: 15 additions & 0 deletions packages/ai-chat/src/common/chat-model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,21 @@ export interface ChatRequestModel {
readonly data?: { [key: string]: unknown };
}

export namespace ChatRequestModel {
export function is(request: unknown): request is ChatRequestModel {
return !!(
request &&
typeof request === 'object' &&
'id' in request &&
typeof (request as { id: unknown }).id === 'string' &&
'session' in request &&
'request' in request &&
'response' in request &&
'message' in request
);
}
}

export interface ChatProgressMessage {
kind: 'progressMessage';
id: string;
Expand Down
59 changes: 59 additions & 0 deletions packages/ai-chat/src/common/chat-tool-request-service.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
// *****************************************************************************
// Copyright (C) 2025 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 { ToolRequest } from '@theia/ai-core';
import { injectable } from '@theia/core/shared/inversify';
import { ChatRequestModelImpl } from './chat-model';

export interface ChatToolRequest extends ToolRequest {
handler: (
arg_string: string,
context: ChatRequestModelImpl,
) => Promise<unknown>;
}

/**
* Wraps tool requests in a chat context.
*
* This service extracts tool requests from a given chat request model and wraps their
* handler functions to provide additional context, such as the chat request model.
*/
@injectable()
export class ChatToolRequestService {

getChatToolRequests(request: ChatRequestModelImpl): ChatToolRequest[] {
const toolRequests = request.message.toolRequests.size > 0 ? [...request.message.toolRequests.values()] : undefined;
if (!toolRequests) {
return [];
}
return this.toChatToolRequests(toolRequests, request);
}

toChatToolRequests(toolRequests: ToolRequest[] | undefined, request: ChatRequestModelImpl): ChatToolRequest[] {
if (!toolRequests) {
return [];
}
return toolRequests.map(toolRequest => this.toChatToolRequest(toolRequest, request));
}

protected toChatToolRequest(toolRequest: ToolRequest, request: ChatRequestModelImpl): ChatToolRequest {
return {
...toolRequest,
handler: async (arg_string: string) => toolRequest.handler(arg_string, request)
};
}

}
2 changes: 1 addition & 1 deletion packages/ai-core/src/common/language-model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ export interface ToolRequest {
name: string;
parameters?: ToolRequestParameters
description?: string;
handler: (arg_string: string) => Promise<unknown>;
handler: (arg_string: string, ctx?: unknown) => Promise<unknown>;
providerName?: string;
}

Expand Down
Loading