From 1dab6fcf4317989655e4b5545ae95fdaca7a1180 Mon Sep 17 00:00:00 2001
From: Joyce Er <joyce.er@microsoft.com>
Date: Thu, 6 Jun 2024 16:21:57 -0700
Subject: [PATCH] fix: correctly implement rendered function messages as User
 messages with `LanguageModelChatMessageFunctionResultPart` content (#57)

---
 package.json                          |  2 +-
 src/base/index.ts                     |  5 +++
 src/base/tokenizer/tokenizer.ts       |  2 +-
 src/base/vscode.proposed.lmTools.d.ts | 58 +++++++++++++++++++++++++++
 src/base/vscodeTypes.d.ts             | 11 -----
 5 files changed, 65 insertions(+), 13 deletions(-)
 create mode 100644 src/base/vscode.proposed.lmTools.d.ts

diff --git a/package.json b/package.json
index 0ae7e8a..1a03c61 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
 	"name": "@vscode/prompt-tsx",
-	"version": "0.2.0-alpha",
+	"version": "0.2.1-alpha",
 	"description": "Declare LLM prompts with TSX",
 	"main": "./dist/base/index.js",
 	"types": "./dist/base/index.d.ts",
diff --git a/src/base/index.ts b/src/base/index.ts
index 726be23..1d9f9fc 100644
--- a/src/base/index.ts
+++ b/src/base/index.ts
@@ -102,6 +102,11 @@ export function toVsCodeChatMessages(messages: ChatMessage[]) {
 				return vscode.LanguageModelChatMessage.Assistant(m.content, m.name);
 			case ChatRole.User:
 				return vscode.LanguageModelChatMessage.User(m.content, m.name);
+			case ChatRole.Function: {
+				const message = vscode.LanguageModelChatMessage.User('');
+				message.content2 = new vscode.LanguageModelChatMessageFunctionResultPart(m.name, m.content);
+				return message;
+			}
 			default:
 				throw new Error(`Converting chat message with role ${m.role} to VS Code chat message is not supported.`);
 		}
diff --git a/src/base/tokenizer/tokenizer.ts b/src/base/tokenizer/tokenizer.ts
index 8f134e9..fba8841 100644
--- a/src/base/tokenizer/tokenizer.ts
+++ b/src/base/tokenizer/tokenizer.ts
@@ -46,7 +46,7 @@ export class AnyTokenizer implements ITokenizer {
 			case ChatRole.User: return 1;
 			case ChatRole.Assistant: return 2;
 			case ChatRole.System: return 1;
-			case ChatRole.Function: return 5;
+			case ChatRole.Function: return 1;
 		}
 	}
 }
diff --git a/src/base/vscode.proposed.lmTools.d.ts b/src/base/vscode.proposed.lmTools.d.ts
new file mode 100644
index 0000000..6aa0ac3
--- /dev/null
+++ b/src/base/vscode.proposed.lmTools.d.ts
@@ -0,0 +1,58 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+declare module 'vscode' {
+
+	// TODO@API capabilities
+
+	export type JSONSchema = object;
+
+	// API -> LM: an tool/function that is available to the language model
+	export interface LanguageModelChatFunction {
+		name: string;
+		description: string;
+		parametersSchema: JSONSchema;
+	}
+
+	// API -> LM: add tools as request option
+	export interface LanguageModelChatRequestOptions {
+		// TODO@API this will a heterogeneous array of different types of tools
+		tools?: LanguageModelChatFunction[];
+	}
+
+	// LM -> USER: function that should be used
+	export class LanguageModelChatResponseFunctionUsePart {
+		name: string;
+		parameters: any;
+
+		constructor(name: string, parameters: any);
+	}
+
+	// LM -> USER: text chunk
+	export class LanguageModelChatResponseTextPart {
+		value: string;
+
+		constructor(value: string);
+	}
+
+	export interface LanguageModelChatResponse {
+
+		stream: AsyncIterable<LanguageModelChatResponseTextPart | LanguageModelChatResponseFunctionUsePart>;
+	}
+
+
+	// USER -> LM: the result of a function call
+	export class LanguageModelChatMessageFunctionResultPart {
+		name: string;
+		content: string;
+		isError: boolean;
+
+		constructor(name: string, content: string, isError?: boolean);
+	}
+
+	export interface LanguageModelChatMessage {
+		content2: string | LanguageModelChatMessageFunctionResultPart;
+	}
+}
diff --git a/src/base/vscodeTypes.d.ts b/src/base/vscodeTypes.d.ts
index 46a2613..8910c2a 100644
--- a/src/base/vscodeTypes.d.ts
+++ b/src/base/vscodeTypes.d.ts
@@ -179,17 +179,6 @@ export enum LanguageModelChatMessageRole {
 	 * The assistant role, e.g. the language model generating responses.
 	 */
 	Assistant = 2,
-
-	/**
-	 * The tool role, e.g. the language model using a tool.
-	 */
-	Tool = 4,
-
-	/**
-	 * The function role, e.g. the language model calling a function.
-	 */
-	Function = 5,
-
 }
 
 /**