Files
vscode-extension-samples/chat-tools-sample/src/toolsPrompt.tsx
2024-10-10 15:36:04 -07:00

250 lines
8.4 KiB
TypeScript

import {
AssistantMessage,
BasePromptElementProps,
contentType as promptTsxContentType,
PrioritizedList,
PromptElement,
PromptElementProps,
PromptPiece,
PromptSizing,
UserMessage,
PromptMetadata,
} from '@vscode/prompt-tsx';
import { Chunk, ToolMessage, ToolResult } from '@vscode/prompt-tsx/dist/base/promptElements';
import * as vscode from 'vscode';
export interface ToolUserProps extends BasePromptElementProps {
request: vscode.ChatRequest;
context: vscode.ChatContext;
toolCalls: vscode.LanguageModelToolCallPart[];
}
export class ToolUserPrompt extends PromptElement<ToolUserProps, void> {
render(state: void, sizing: PromptSizing) {
return (
<>
<UserMessage priority={50}>
Instructions: <br />
- The user will ask a question, or ask you to perform a task, and it may
require lots of research to answer correctly. There is a selection of
tools that let you perform actions or retrieve helpful context to answer
the user's question. <br />
- If you aren't sure which tool is relevant, you can call multiple
tools. You can call tools repeatedly to take actions or gather as much
context as needed until you have completed the task fully. Don't give up
unless you are sure the request cannot be fulfilled with the tools you
have. <br />
- Don't make assumptions about the situation- gather context first, then
perform the task or answer the question. <br />
- Don't ask the user for confirmation to use tools, just use them.
<br />- After editing a file, DO NOT show the user a codeblock with the
edit or new file contents. Assume that the user can see the result.
</UserMessage>
<History context={this.props.context} priority={20}></History>
<PromptReferences
references={this.props.request.references}
priority={30}
/>
<UserMessage priority={40}>{this.props.request.prompt}</UserMessage>
<ToolCalls toolCalls={this.props.toolCalls} toolInvocationToken={this.props.request.toolInvocationToken}></ToolCalls>
</>
);
}
}
interface ToolCallsProps extends BasePromptElementProps {
toolCalls: vscode.LanguageModelToolCallPart[];
toolInvocationToken: vscode.ChatParticipantToolToken;
}
const agentSupportedContentTypes = [promptTsxContentType, 'text/plain'];
const dummyCancellationToken: vscode.CancellationToken = new vscode.CancellationTokenSource().token;
class ToolCalls extends PromptElement<ToolCallsProps, void> {
async render(state: void, sizing: PromptSizing) {
if (!this.props.toolCalls.length) {
return undefined;
}
// TODO- prompt-tsx export this type?
// TODO- at what level do the parameters get stringified?
const assistantToolCalls: any[] = this.props.toolCalls.map(tc => ({ type: 'function', function: { name: tc.name, arguments: JSON.stringify(tc.parameters) }, id: tc.toolCallId }));
// TODO@prompt-tsx- don't remove "empty" assistant messages!
const toolResultMap = new Map<string, vscode.LanguageModelToolResult>();
const budget = Math.floor(sizing.tokenBudget / this.props.toolCalls.length);
return <Chunk>
<meta value={new ToolResultMetadata(toolResultMap)}></meta>
<AssistantMessage toolCalls={assistantToolCalls}>todo</AssistantMessage>
{await Promise.all(this.props.toolCalls.map(async toolCall => {
// TODO@prompt-tsx- this would be a bit easier with a ToolCall element, but we can only return one instance of a metadata right now
const toolCallSizing: PromptSizing = {
...sizing,
tokenBudget: budget,
};
const result = await this.renderOneToolCall(toolCall, toolCallSizing, this.props.toolInvocationToken);
toolResultMap.set(toolCall.toolCallId, result);
return result.message;
}))}
<UserMessage priority={100}>Above is the result of calling one or more tools. The user cannot see the results, so you should explain them to the user if referencing them in your answer.</UserMessage>
</Chunk>;
}
private async renderOneToolCall(toolCall: vscode.LanguageModelToolCallPart, sizing: PromptSizing, toolInvocationToken: vscode.ChatParticipantToolToken): Promise<{ message: ToolMessage; toolResult: vscode.LanguageModelToolResult; }> {
const tool = vscode.lm.tools.find(t => t.id === toolCall.name);
if (!tool) {
console.error(`Tool not found: ${toolCall.name}`);
return <ToolMessage toolCallId={toolCall.toolCallId}>Tool not found</ToolMessage>;
}
const contentType = agentSupportedContentTypes.find(type => tool.supportedContentTypes.includes(type));
if (!contentType) {
console.error(`Tool does not support any of the agent's content types: ${tool.id}`);
return <ToolMessage toolCallId={toolCall.toolCallId}>Tool unsupported</ToolMessage>;
}
const tokenOptions: vscode.LanguageModelToolInvocationOptions<unknown>['tokenOptions'] = {
tokenBudget: sizing.tokenBudget,
countTokens: async (content: string) => sizing.countTokens(content),
};
const result = await vscode.lm.invokeTool(toolCall.name, { parameters: toolCall.parameters, requestedContentTypes: [contentType], toolInvocationToken: toolInvocationToken, tokenOptions }, dummyCancellationToken);
return <ToolMessage toolCallId={toolCall.toolCallId}>
{contentType === 'text/plain' ?
result[contentType] :
<elementJSON data={result[contentType]}></elementJSON>}
</ToolMessage>;
}
}
export class ToolResultMetadata extends PromptMetadata {
constructor(
public resultMap: Map<string, vscode.LanguageModelToolResult>,
) {
super();
}
}
interface HistoryProps extends BasePromptElementProps {
priority: number;
context: vscode.ChatContext;
}
class History extends PromptElement<HistoryProps, void> {
render(state: void, sizing: PromptSizing) {
return (
<PrioritizedList priority={this.props.priority} descending={false}>
{this.props.context.history.map((message) => {
if (message instanceof vscode.ChatRequestTurn) {
return (
<>
{<PromptReferences references={message.references} />}
<UserMessage>{message.prompt}</UserMessage>
</>
);
} else if (message instanceof vscode.ChatResponseTurn) {
return (
<AssistantMessage>
{chatResponseToString(message)}
</AssistantMessage>
);
}
})}
</PrioritizedList>
);
}
}
function chatResponseToString(response: vscode.ChatResponseTurn): string {
return response.response
.map((r) => {
if (r instanceof vscode.ChatResponseMarkdownPart) {
return r.value.value;
} else if (r instanceof vscode.ChatResponseAnchorPart) {
if (r.value instanceof vscode.Uri) {
return r.value.fsPath;
} else {
return r.value.uri.fsPath;
}
}
return '';
})
.join('');
}
interface PromptReferencesProps extends BasePromptElementProps {
references: ReadonlyArray<vscode.ChatPromptReference>;
}
class PromptReferences extends PromptElement<PromptReferencesProps, void> {
render(state: void, sizing: PromptSizing): PromptPiece {
return (
<UserMessage>
{this.props.references.map((ref, index) => (
<PromptReference ref={ref}></PromptReference>
))}
</UserMessage>
);
}
}
interface PromptReferenceProps extends BasePromptElementProps {
ref: vscode.ChatPromptReference;
}
class PromptReference extends PromptElement<PromptReferenceProps> {
async render(state: void, sizing: PromptSizing): Promise<PromptPiece | undefined> {
const value = this.props.ref.value;
if (value instanceof vscode.Uri) {
const fileContents = (await vscode.workspace.fs.readFile(value)).toString();
return (
<Tag name="context">
{value.fsPath}:<br />
``` <br />
{fileContents}<br />
```<br />
</Tag>
);
} else if (value instanceof vscode.Location) {
const rangeText = (await vscode.workspace.openTextDocument(value.uri)).getText(value.range);
return (
<Tag name="context">
{value.uri.fsPath}:{value.range.start.line + 1}-$<br />
{value.range.end.line + 1}: ```<br />
{rangeText}<br />
```
</Tag>
);
} else if (typeof value === 'string') {
return <Tag name="context">{value}</Tag>;
}
}
}
export type TagProps = PromptElementProps<{
name: string;
}>;
export class Tag extends PromptElement<TagProps> {
private static readonly _regex = /^[a-zA-Z_][\w\.\-]*$/;
render() {
const { name } = this.props;
if (!Tag._regex.test(name)) {
throw new Error(`Invalid tag name: ${this.props.name}`);
}
return (
<>
{'<' + name + '>'}<br />
<>
{this.props.children}<br />
</>
{'</' + name + '>'}<br />
</>
);
}
}