Skip to content

Commit

Permalink
feat: allow passing context to tool calls
Browse files Browse the repository at this point in the history
With this change we enable passing optional context to tool calls. This
is then used for tool calls in the context of a chat to pass on the chat
request model. This way tool call handlers can retrieve additional
information as a context or update the chat model directly.
  • Loading branch information
planger committed Jan 22, 2025
1 parent 998bf47 commit d493ea9
Show file tree
Hide file tree
Showing 5 changed files with 89 additions and 17 deletions.
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

0 comments on commit d493ea9

Please sign in to comment.