From 663b39c7be290aef7a81e35427a82e7472018504 Mon Sep 17 00:00:00 2001
From: Anders Hejlsberg <andersh@microsoft.com>
Date: Thu, 31 Aug 2023 10:41:53 -0700
Subject: [PATCH] Add support for prompt roles and prompt preambles

---
 src/model.ts    | 26 ++++++++++++++++++++++----
 src/typechat.ts | 14 +++++++++-----
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/model.ts b/src/model.ts
index dc0dafcb..6567a064 100644
--- a/src/model.ts
+++ b/src/model.ts
@@ -1,6 +1,22 @@
 import axios from "axios";
 import { Result, success, error } from "./result";
 
+/**
+ * Represents a section of an LLM prompt with an associated role. TypeChat uses the "user" role for
+ * prompts it generates and the "assistant" role for previous LLM responses (which will be part of
+ * the prompt in repair attempts). TypeChat currently doesn't use the "system" role.
+ */
+export interface PromptSection {
+    /**
+     * Specifies the role of this section.
+     */
+    role: "system" | "user" | "assistant";
+    /**
+     * Specifies the content of this section.
+     */
+    content: string;
+}
+
 /**
  * Represents a AI language model that can complete prompts. TypeChat uses an implementation of this
  * interface to communicate with an AI service that can translate natural language requests to JSON
@@ -18,9 +34,10 @@ export interface TypeChatLanguageModel {
     retryPauseMs?: number;
     /**
      * Obtains a completion from the language model for the given prompt.
-     * @param prompt The prompt string.
+     * @param prompt A prompt string or an array of prompt sections. If a string is specified,
+     *   it is converted into a single "user" role prompt section.
      */
-    complete(prompt: string): Promise<Result<string>>;
+    complete(prompt: string | PromptSection[]): Promise<Result<string>>;
 }
 
 /**
@@ -93,14 +110,15 @@ function createAxiosLanguageModel(url: string, config: object, defaultParams: Re
     };
     return model;
 
-    async function complete(prompt: string) {
+    async function complete(prompt: string | PromptSection[]) {
         let retryCount = 0;
         const retryMaxAttempts = model.retryMaxAttempts ?? 3;
         const retryPauseMs = model.retryPauseMs ?? 1000;
+        const messages = typeof prompt === "string" ? [{ role: "user", content: prompt }] : prompt;
         while (true) {
             const params = {
                 ...defaultParams,
-                messages: [{ role: "user", content: prompt }],
+                messages,
                 temperature: 0,
                 n: 1
             };
diff --git a/src/typechat.ts b/src/typechat.ts
index 3ab774f6..dcf2b8ca 100644
--- a/src/typechat.ts
+++ b/src/typechat.ts
@@ -1,5 +1,5 @@
 import { Result, error } from "./result";
-import { TypeChatLanguageModel } from "./model";
+import { TypeChatLanguageModel, PromptSection } from "./model";
 import { TypeChatJsonValidator, createJsonValidator } from "./validate";
 
 /**
@@ -48,9 +48,11 @@ export interface TypeChatJsonTranslator<T extends object> {
      * attempt to translate the request will be made. The prompt for the second attempt will include the
      * diagnostics produced for the first attempt. This often helps produce a valid instance.
      * @param request The natural language request.
+     * @param promptPreamble An optional string or array of prompt sections to prepend to the generated
+     *   prompt. If a string is specified, it is converted into a single "user" role prompt section.
      * @returns A promise for the resulting object.
      */
-    translate(request: string): Promise<Result<T>>;
+    translate(request: string, promptPreamble?: string | PromptSection[]): Promise<Result<T>>;
 }
 
 /**
@@ -89,8 +91,9 @@ export function createJsonTranslator<T extends object>(model: TypeChatLanguageMo
             `The following is a revised JSON object:\n`;
     }
 
-    async function translate(request: string) {
-        let prompt = typeChat.createRequestPrompt(request);
+    async function translate(request: string, promptPreamble?: string | PromptSection[]) {
+        const preamble: PromptSection[] = typeof promptPreamble === "string" ? [{ role: "user", content: promptPreamble }] : promptPreamble ?? [];
+        let prompt: PromptSection[] = [...preamble, { role: "user", content: typeChat.createRequestPrompt(request) }];
         let attemptRepair = typeChat.attemptRepair;
         while (true) {
             const response = await model.complete(prompt);
@@ -111,7 +114,8 @@ export function createJsonTranslator<T extends object>(model: TypeChatLanguageMo
             if (!attemptRepair) {
                 return error(`JSON validation failed: ${validation.message}\n${jsonText}`);
             }
-            prompt += `${responseText}\n${typeChat.createRepairPrompt(validation.message)}`;
+            prompt.push({ role: "assistant", content: responseText });
+            prompt.push({ role: "user", content: typeChat.createRepairPrompt(validation.message) });
             attemptRepair = false;
         }
     }