feat: add chat completion (#8)

* feat: add chat completion * chore: self mutation Signed-off-by: github-actions <[email protected]> --------- Signed-off-by: github-actions <[email protected]> Co-authored-by: github-actions <[email protected]>
WiseGPT · Mar 19, 2023 · da95564 · da95564
1 parent 37325e8
commit da95564
Show file tree

Hide file tree

Showing 12 changed files with 308 additions and 52 deletions.
diff --git a/src/conversation-prompt/conversation-prompt-service.dto.ts b/src/conversation-prompt/conversation-prompt-service.dto.ts
@@ -17,6 +17,21 @@ export type ModelConfiguration = Pick<
   | "frequency_penalty"
 >;
 
+export type ConversationChatCompleteInput = {
+  // prompt generation related details
+  prompt: {
+    aiPersona: AIPersona;
+    conversation: Conversation;
+  };
+  modelConfiguration: ModelConfiguration;
+};
+
+export type ConversationChatCompleteOutput = {
+  text: string;
+  usage: APIUsageInfo;
+  headers: APIResponseHeaders;
+};
+
 export type ConversationCompleteInput = {
   // prompt generation related details
   prompt: {

diff --git a/src/conversation-prompt/conversation-prompt.service.ts b/src/conversation-prompt/conversation-prompt.service.ts
@@ -1,15 +1,20 @@
 import {
+  ChatCompletionRequestMessage,
+  CreateChatCompletionResponse,
   CreateCompletionResponse,
   CreateCompletionResponseUsage,
   OpenAIApi,
 } from "openai";
 import {
+  ConversationChatCompleteInput,
+  ConversationChatCompleteOutput,
   ConversationCompleteInput,
   ConversationCompleteOutput,
   ConversationSummaryInput,
   ConversationSummaryOutput,
   ModelConfiguration,
 } from "./conversation-prompt-service.dto";
+import { createConversationChatCompletionPrompt } from "./prompts/create-conversation-chat-completion-prompt";
 import { createConversationCompletionPrompt } from "./prompts/create-conversation-completion-prompt";
 import { createConversationSummaryPrompt } from "./prompts/create-conversation-summary-prompt";
 import { STATEMENT_SEPARATOR_TOKEN } from "./prompts/prompts.constants";
@@ -63,6 +68,23 @@ export class ConversationPromptService {
 
   constructor(private readonly openAIApi: OpenAIApi) {}
 
+  async chatCompletion(
+    input: ConversationChatCompleteInput
+  ): Promise<ConversationChatCompleteOutput> {
+    const messages = createConversationChatCompletionPrompt(input.prompt);
+
+    const {
+      firstChoice: text,
+      usage,
+      headers,
+    } = await this.makeChatCompletionRequest({
+      messages,
+      modelConfiguration: input.modelConfiguration,
+    });
+
+    return { text, usage, headers };
+  }
+
   async completion(
     input: ConversationCompleteInput
   ): Promise<ConversationCompleteOutput> {
@@ -96,6 +118,44 @@ export class ConversationPromptService {
     return { summary, usage, headers };
   }
 
+  // TODO: remove duplication
+  private async makeChatCompletionRequest(input: {
+    messages: Array<ChatCompletionRequestMessage>;
+    modelConfiguration: ModelConfiguration;
+  }) {
+    try {
+      const { data, headers }: APIResponse<CreateChatCompletionResponse> =
+        await this.openAIApi.createChatCompletion({
+          ...input.modelConfiguration,
+          max_tokens: input.modelConfiguration.max_tokens ?? undefined,
+          messages: input.messages,
+          n: 1,
+        });
+
+      const firstChoice = data.choices?.[0]?.message?.content?.trim()!;
+
+      return {
+        firstChoice,
+        usage: ConversationPromptService.extractAPIUsageInfo(data),
+        headers: ConversationPromptService.extractAPIResponseHeaders(headers),
+      };
+    } catch (err: unknown) {
+      if (ConversationPromptService.isErrorResponse(err)) {
+        const { message, type } = err.response.data.error;
+
+        throw new ConversationPromptServiceError(
+          err.response.status,
+          { message, type },
+          ConversationPromptService.extractAPIResponseHeaders(
+            err.response.headers
+          )
+        );
+      }
+
+      throw err;
+    }
+  }
+
   private async makeCompletionRequest(input: {
     prompt: string;
     modelConfiguration: ModelConfiguration;

diff --git a/src/conversation-prompt/mention.ts b/src/conversation-prompt/mention.ts
@@ -1,11 +1,11 @@
 import { Author } from "../types";
 
-export const BOT_MENTION = "<@bot>";
+export const ASSISTANT_MENTION = "<@assistant>";
 
 export const buildMention = (author: Author): string => {
   switch (author.type) {
     case "BOT":
-      return BOT_MENTION;
+      return ASSISTANT_MENTION;
     case "USER":
       return `<@${author.id}>`;
     default:

diff --git a/src/conversation-prompt/prompts/create-conversation-chat-completion-prompt.ts b/src/conversation-prompt/prompts/create-conversation-chat-completion-prompt.ts
@@ -0,0 +1,27 @@
+import { ChatCompletionRequestMessage } from "openai";
+import { renderAIPersona } from "./render-ai-persona";
+import { renderConversationForChat } from "./render-conversation-for-chat";
+import { AIPersona, Conversation } from "../../types";
+
+const CONVERSATION_SUMMARY_SUFFIX = `The conversations starts with a detailed summary of a previous conversation. While answering questions, take this summary into account. Summary:`;
+
+export type CreateConversationChatCompletionPromptInput = {
+  aiPersona: AIPersona;
+  conversation: Conversation;
+};
+
+export function createConversationChatCompletionPrompt({
+  aiPersona,
+  conversation,
+}: CreateConversationChatCompletionPromptInput): Array<ChatCompletionRequestMessage> {
+  const systemMessage = {
+    role: "system" as const,
+    content: renderAIPersona(aiPersona),
+  };
+
+  if (conversation.summary) {
+    systemMessage.content += `\n\n${CONVERSATION_SUMMARY_SUFFIX} ${conversation.summary}`;
+  }
+
+  return [systemMessage, ...renderConversationForChat(conversation)];
+}
diff --git a/src/conversation-prompt/prompts/create-conversation-completion-prompt.ts b/src/conversation-prompt/prompts/create-conversation-completion-prompt.ts
@@ -3,7 +3,7 @@ import { renderAIPersona } from "./render-ai-persona";
 import { renderConversation } from "./render-conversation";
 import { renderFormatAndExamples } from "./render-format-and-examples";
 import { AIPersona, Conversation } from "../../types";
-import { BOT_MENTION } from "../mention";
+import { ASSISTANT_MENTION } from "../mention";
 
 const CURRENT_CONVERSATION_PROMPT = `Continue the conversation, paying very close attention to things entities told you; such as their name, and personal details. Never say "${STATEMENT_SEPARATOR_TOKEN}". Current conversation:`;
 
@@ -22,11 +22,11 @@ export function createConversationCompletionPrompt({
 
   return (
     renderAIPersona(aiPersona) +
-    `\n${renderFormatAndExamples({
+    `\n\n${renderFormatAndExamples({
       hasSummary,
       exampleConversations,
     })}` +
     `\n\n${CURRENT_CONVERSATION_PROMPT}\n\n` +
-    (renderConversation(conversation) + `${BOT_MENTION}:`)
+    (renderConversation(conversation) + `${ASSISTANT_MENTION}:`)
   );
 }
diff --git a/src/conversation-prompt/prompts/create-conversation-summary-prompt.ts b/src/conversation-prompt/prompts/create-conversation-summary-prompt.ts
@@ -53,7 +53,7 @@ export const createConversationSummaryPrompt = ({
 
   return (
     renderAIPersona(aiPersona) +
-    `\n${renderFormatAndExamples({
+    `\n\n${renderFormatAndExamples({
       hasSummary,
     })}` +
     `\n\n${prompt}\n\n` +

diff --git a/src/conversation-prompt/prompts/render-ai-persona.ts b/src/conversation-prompt/prompts/render-ai-persona.ts
@@ -1,10 +1,9 @@
 import { AIPersona } from "../../types";
-import { BOT_MENTION } from "../mention";
+import { ASSISTANT_MENTION } from "../mention";
 
 export const renderAIPersona = (aiPersona: AIPersona): string =>
-  `Instructions for ${BOT_MENTION}, this is how you should behave in a conversation, but this is not your personality:\n` +
-  `Your name is "${aiPersona.name}". You are referenced in conversations as "${BOT_MENTION}".\n` +
+  `Instructions for ${ASSISTANT_MENTION}, this is how you should behave in a conversation, but this is not your personality:\n` +
+  `Your name is "${aiPersona.name}". You are referenced in conversations as "${ASSISTANT_MENTION}".\n` +
   aiPersona.instructions +
   `\n\nThis is your personality:\n` +
-  aiPersona.personality +
-  "\n";
+  aiPersona.personality;
diff --git a/src/conversation-prompt/prompts/render-conversation-for-chat.ts b/src/conversation-prompt/prompts/render-conversation-for-chat.ts
@@ -0,0 +1,17 @@
+import { ChatCompletionRequestMessage } from "openai";
+import { Conversation } from "../../types";
+import { buildMention } from "../mention";
+
+export const renderConversationForChat = ({
+  messages,
+}: Conversation): Array<ChatCompletionRequestMessage> => {
+  return messages.map((message): ChatCompletionRequestMessage => {
+    const author = message.author;
+
+    return {
+      role: author.type === "BOT" ? "assistant" : "user",
+      name: buildMention(author).replace(/[<>@]/g, ""),
+      content: message.text,
+    };
+  });
+};
diff --git a/src/conversation-prompt/prompts/render-format-and-examples.ts b/src/conversation-prompt/prompts/render-format-and-examples.ts
@@ -1,7 +1,7 @@
 import { STATEMENT_SEPARATOR_TOKEN } from "./prompts.constants";
 import { renderConversation } from "./render-conversation";
 import { Conversation } from "../../types";
-import { BOT_MENTION, buildMention } from "../mention";
+import { ASSISTANT_MENTION, buildMention } from "../mention";
 
 const CONVERSATION_FORMAT = `The conversations are in this format, there can be an arbitrary amount of newlines between chat entries. "<@id>" format is used to reference entities in the conversation, where "id" is replaced with message author's unique id. The text "${STATEMENT_SEPARATOR_TOKEN}" is used to separate chat entries and make it easier for you to understand the context:`;
 const CONVERSATION_FORMAT_WITH_EXAMPLE = `The conversations are in this format, there can be an arbitrary amount of newlines between chat entries. "<@id>" format is used to reference entities in the conversation, where "id" is replaced with message author's unique id. The text "${STATEMENT_SEPARATOR_TOKEN}" is used to separate chat entries and make it easier for you to understand the context. The conversations start with a "Summary:" which includes a detailed summary of messages in the same conversation. Summary ends with "${STATEMENT_SEPARATOR_TOKEN}":`;
@@ -22,7 +22,7 @@ const BASIC_EXAMPLE_CONVERSATIONS: Conversation[] = [
   {
     messages: [
       {
-        text: `hello ${BOT_MENTION}`,
+        text: `hello ${ASSISTANT_MENTION}`,
         author: { type: "USER", id: "U02" },
       },
       {

diff --git a/test/conversation-prompt/prompts/create-conversation-chat-completion-prompt.test.ts b/test/conversation-prompt/prompts/create-conversation-chat-completion-prompt.test.ts
@@ -0,0 +1,128 @@
+import {
+  AIPersona,
+  Author,
+  ASSISTANT_MENTION,
+  buildMention,
+} from "../../../src";
+import { createConversationChatCompletionPrompt } from "../../../src/conversation-prompt/prompts/create-conversation-chat-completion-prompt";
+
+describe("createConversationChatCompletionPrompt", () => {
+  const aiPersona: AIPersona = {
+    name: "Lenard",
+    instructions: `When providing code examples, use triple backticks.`,
+    personality: `You are a software engineer.`,
+  };
+  const authors: Record<string, Author> = {
+    bot: { type: "BOT" },
+    exampleUser1: { type: "USER", id: "EU01" },
+    exampleUser2: { type: "USER", id: "EU02" },
+    user1: { type: "USER", id: "U01" },
+  };
+
+  it("should work without summary", () => {
+    expect(
+      createConversationChatCompletionPrompt({
+        aiPersona,
+        conversation: {
+          messages: [
+            {
+              author: authors.user1,
+              text: "hello!",
+            },
+            {
+              author: authors.bot,
+              text: "hello! how can I help you?",
+            },
+            {
+              author: authors.user1,
+              text: "can you write me fibonacci function in Typescript?",
+            },
+          ],
+        },
+      })
+    ).toMatchInlineSnapshot(`
+      [
+        {
+          "content": "Instructions for <@assistant>, this is how you should behave in a conversation, but this is not your personality:
+      Your name is "Lenard". You are referenced in conversations as "<@assistant>".
+      When providing code examples, use triple backticks.
+
+      This is your personality:
+      You are a software engineer.",
+          "role": "system",
+        },
+        {
+          "content": "hello!",
+          "name": "U01",
+          "role": "user",
+        },
+        {
+          "content": "hello! how can I help you?",
+          "name": "assistant",
+          "role": "assistant",
+        },
+        {
+          "content": "can you write me fibonacci function in Typescript?",
+          "name": "U01",
+          "role": "user",
+        },
+      ]
+    `);
+  });
+
+  it("should work with summary", () => {
+    expect(
+      createConversationChatCompletionPrompt({
+        aiPersona,
+        conversation: {
+          summary: `${buildMention(
+            authors.user1
+          )} asked ${ASSISTANT_MENTION} whether it knows Typescript.`,
+          messages: [
+            {
+              author: authors.user1,
+              text: "hello!",
+            },
+            {
+              author: authors.bot,
+              text: "hello! how can I help you?",
+            },
+            {
+              author: authors.user1,
+              text: "can you write me fibonacci function in Typescript?",
+            },
+          ],
+        },
+      })
+    ).toMatchInlineSnapshot(`
+      [
+        {
+          "content": "Instructions for <@assistant>, this is how you should behave in a conversation, but this is not your personality:
+      Your name is "Lenard". You are referenced in conversations as "<@assistant>".
+      When providing code examples, use triple backticks.
+
+      This is your personality:
+      You are a software engineer.
+
+      The conversations starts with a detailed summary of a previous conversation. While answering questions, take this summary into account. Summary: <@U01> asked <@assistant> whether it knows Typescript.",
+          "role": "system",
+        },
+        {
+          "content": "hello!",
+          "name": "U01",
+          "role": "user",
+        },
+        {
+          "content": "hello! how can I help you?",
+          "name": "assistant",
+          "role": "assistant",
+        },
+        {
+          "content": "can you write me fibonacci function in Typescript?",
+          "name": "U01",
+          "role": "user",
+        },
+      ]
+    `);
+  });
+});