From b1ac3fece55ed0019218f085a12e5bf07bbe0755 Mon Sep 17 00:00:00 2001
From: Santiago Kent <santiagokent@gmail.com>
Date: Sun, 5 May 2024 12:51:29 -0300
Subject: [PATCH] refactor the creation of messages for openAI to reduce
 repetition and increase readability

---
 src/KurtOpenAI.ts | 58 +++++++++++++++++++++++------------------------
 1 file changed, 28 insertions(+), 30 deletions(-)

diff --git a/src/KurtOpenAI.ts b/src/KurtOpenAI.ts
index 393b0e2..fd849f0 100644
--- a/src/KurtOpenAI.ts
+++ b/src/KurtOpenAI.ts
@@ -38,22 +38,12 @@ export class KurtOpenAI implements Kurt {
   generateNaturalLanguage(
     options: KurtGenerateNaturalLanguageOptions
   ): KurtResult {
-    const systemPrompt = options.systemPrompt ?? this.options.systemPrompt
-    const prompt = options.prompt
-    const extraMessages = options.extraMessages ?? []
-
     return this.handleStream(
       undefined,
       this.options.openAI.chat.completions.create({
         stream: true,
         model: this.options.model,
-        messages: [
-          ...(systemPrompt
-            ? [{ role: "system", content: systemPrompt } as const]
-            : []),
-          { role: "user", content: prompt },
-          ...toOpenAIMessages(extraMessages),
-        ],
+        messages: this.toOpenAIMessages(options),
       })
     )
   }
@@ -61,23 +51,14 @@ export class KurtOpenAI implements Kurt {
   generateStructuredData<T extends KurtSchemaInner>(
     options: KurtGenerateStructuredDataOptions<T>
   ): KurtResult<T> {
-    const systemPrompt = options.systemPrompt ?? this.options.systemPrompt
-    const prompt = options.prompt
     const schema = options.schema
-    const extraMessages = options.extraMessages ?? []
 
     return this.handleStream(
       schema as KurtSchemaMaybe<T>,
       this.options.openAI.chat.completions.create({
         stream: true,
         model: this.options.model,
-        messages: [
-          ...(systemPrompt
-            ? [{ role: "system", content: systemPrompt } as const]
-            : []),
-          { role: "user", content: prompt },
-          ...toOpenAIMessages(extraMessages),
-        ],
+        messages: this.toOpenAIMessages(options),
         tool_choice: {
           type: "function",
           function: { name: "structured_data" },
@@ -144,14 +125,31 @@ export class KurtOpenAI implements Kurt {
 
     return new KurtResult<T>(generator())
   }
-}
 
-function toOpenAIMessages(messages: KurtMessage[]): OpenAIMessage[] {
-  return messages.map((message) => {
-    const { role, text } = message
-    return {
-      role: role === "model" ? "assistant" : role,
-      content: text,
-    }
-  })
+  private toOpenAIMessages = ({
+    prompt,
+    systemPrompt = this.options.systemPrompt,
+    extraMessages = [],
+  }: KurtGenerateNaturalLanguageOptions): OpenAIMessage[] => {
+    const systemMessage: OpenAIMessage[] = systemPrompt
+      ? [toOpenAIMessage({ role: "system", text: systemPrompt })]
+      : []
+
+    const userMessage = toOpenAIMessage({ role: "user", text: prompt })
+
+    const extras = extraMessages.map(toOpenAIMessage)
+
+    return systemMessage.concat(userMessage, extras)
+  }
 }
+
+const toOpenAIMessage = ({ role, text }: KurtMessage): OpenAIMessage => ({
+  role: openAIRoleMapping[role],
+  content: text,
+})
+
+const openAIRoleMapping = {
+  model: "assistant",
+  system: "system",
+  user: "user",
+} as const satisfies Record<KurtMessage["role"], OpenAIMessage["role"]>