From bc3eeaf06d8993ce5c4f145a6cabb4a69fb26a4d Mon Sep 17 00:00:00 2001
From: Alex Ross <alros@microsoft.com>
Date: Tue, 22 Oct 2024 12:11:55 +0200
Subject: [PATCH] Work around User message utility issue (#6349)

* Work around User message utility issue

* Fix in search tool
---
 src/lm/participants.ts            | 10 +++++++---
 src/lm/tools/displayIssuesTool.ts | 10 +++++++---
 src/lm/tools/searchTools.ts       |  2 +-
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/lm/participants.ts b/src/lm/participants.ts
index 106bd159a3..e5f123cb9e 100644
--- a/src/lm/participants.ts
+++ b/src/lm/participants.ts
@@ -25,11 +25,15 @@ export class ChatParticipantState {
 		return [];
 	}
 
-	get firstUserMessage(): string | undefined {
+	get firstUserMessage(): vscode.LanguageModelTextPart | undefined {
 		for (let i = 0; i < this._messages.length; i++) {
 			const message = this._messages[i];
 			if (message.role === vscode.LanguageModelChatMessageRole.User && message.content) {
-				return message.content;
+				for (const part of message.content) {
+					if (part instanceof vscode.LanguageModelTextPart) {
+						return part;
+					}
+				}
 			}
 		}
 	}
@@ -96,7 +100,7 @@ export class ChatParticipant implements vscode.Disposable {
 			{ modelMaxPromptTokens: model.maxInputTokens },
 			model);
 
-		this.state.addMessages(messages);
+		this.state.addMessages(messages as any);
 
 		const toolReferences = [...request.toolReferences];
 		const options: vscode.LanguageModelChatRequestOptions = {
diff --git a/src/lm/tools/displayIssuesTool.ts b/src/lm/tools/displayIssuesTool.ts
index c00019941e..bb8c11e785 100644
--- a/src/lm/tools/displayIssuesTool.ts
+++ b/src/lm/tools/displayIssuesTool.ts
@@ -57,7 +57,11 @@ export class DisplayIssuesTool extends ToolBase<DisplayIssuesParameters> {
 		return finalColumns;
 	}
 
-	private async getImportantColumns(issueItemsInfo: string, issues: IssueSearchResultItem[], token: vscode.CancellationToken): Promise<IssueColumn[]> {
+	private async getImportantColumns(issueItemsInfo: vscode.LanguageModelTextPart | undefined, issues: IssueSearchResultItem[], token: vscode.CancellationToken): Promise<IssueColumn[]> {
+		if (!issueItemsInfo) {
+			return ['number', 'title', 'state'];
+		}
+
 		// Try to get the llm to tell us which columns are important based on information it has about the issues
 		const models = await vscode.lm.selectChatModels({
 			vendor: 'copilot',
@@ -68,7 +72,7 @@ export class DisplayIssuesTool extends ToolBase<DisplayIssuesParameters> {
 			justification: 'Answering user questions pertaining to GitHub.'
 		};
 		const messages = [vscode.LanguageModelChatMessage.Assistant(this.assistantPrompt(issues))];
-		messages.push(vscode.LanguageModelChatMessage.User(issueItemsInfo));
+		messages.push(new vscode.LanguageModelChatMessage(vscode.LanguageModelChatMessageRole.User, issueItemsInfo?.value));
 		const response = await model.sendRequest(messages, chatOptions, token);
 		const result = this.postProcess(await concatAsyncIterable(response.text), issues);
 		const indexOfUrl = result.indexOf('url');
@@ -119,7 +123,7 @@ export class DisplayIssuesTool extends ToolBase<DisplayIssuesParameters> {
 	}
 
 	async invoke(options: vscode.LanguageModelToolInvocationOptions<DisplayIssuesParameters>, token: vscode.CancellationToken): Promise<vscode.LanguageModelToolResult | undefined> {
-		let issueItemsInfo: string = this.chatParticipantState.firstUserMessage ?? '';
+		let issueItemsInfo: vscode.LanguageModelTextPart | undefined = this.chatParticipantState.firstUserMessage;
 		const issueItems: IssueSearchResultItem[] = options.parameters.arrayOfIssues;
 		if (issueItems.length === 0) {
 			return new vscode.LanguageModelToolResult([new vscode.LanguageModelTextPart(vscode.l10n.t('No issues found. Please try another query.'))]);
diff --git a/src/lm/tools/searchTools.ts b/src/lm/tools/searchTools.ts
index 222119b01c..5122d763ab 100644
--- a/src/lm/tools/searchTools.ts
+++ b/src/lm/tools/searchTools.ts
@@ -354,7 +354,7 @@ You are getting ready to make a GitHub search query. Given a natural language qu
 
 	async invoke(options: vscode.LanguageModelToolInvocationOptions<ConvertToQuerySyntaxParameters>, token: vscode.CancellationToken): Promise<vscode.LanguageModelToolResult | undefined> {
 		const { owner, name, folderManager } = this.getRepoInfo({ owner: options.parameters.repo?.owner, name: options.parameters.repo?.name });
-		const firstUserMessage = `${this.chatParticipantState.firstUserMessage}, ${options.parameters.naturalLanguageString}`;
+		const firstUserMessage = `${this.chatParticipantState.firstUserMessage?.value}, ${options.parameters.naturalLanguageString}`;
 
 		const labels = await folderManager.getLabels(undefined, { owner, repo: name });