refactor the creation of messages for vertexAI to reduce repetition a…

…nd increase readability
FormulaMonks · May 5, 2024 · 9232142 · 9232142
1 parent 4bce785
commit 9232142
Showing 1 changed file with 22 additions and 27 deletions.
diff --git a/src/KurtVertexAI.ts b/src/KurtVertexAI.ts
@@ -39,35 +39,22 @@ export class KurtVertexAI implements Kurt {
   generateNaturalLanguage(
     options: KurtGenerateNaturalLanguageOptions
   ): KurtResult {
-    const systemPrompt = options.systemPrompt ?? this.options.systemPrompt
-    const prompt = options.prompt
-    const extraMessages = options.extraMessages ?? []
-
     const llm = this.options.vertexAI.getGenerativeModel({
       model: this.options.model,
     }) as VertexAIGenerativeModel
 
     return this.handleStream(
       undefined,
       llm.generateContentStreamPATCHED({
-        contents: [
-          ...(systemPrompt
-            ? [{ role: "system", parts: [{ text: systemPrompt }] }]
-            : []),
-          { role: "user", parts: [{ text: prompt }] },
-          ...toVertexAIMessages(extraMessages),
-        ],
+        contents: this.toVertexAIMessages(options),
       })
     )
   }
 
   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 ?? []
 
     const llm = this.options.vertexAI.getGenerativeModel({
       model: this.options.model,
@@ -76,13 +63,7 @@ export class KurtVertexAI implements Kurt {
     return this.handleStream(
       schema as KurtSchemaMaybe<T>,
       llm.generateContentStreamPATCHED({
-        contents: [
-          ...(systemPrompt
-            ? [{ role: "system", parts: [{ text: systemPrompt }] }]
-            : []),
-          { role: "user", parts: [{ text: prompt }] },
-          ...toVertexAIMessages(extraMessages),
-        ],
+        contents: this.toVertexAIMessages(options),
         tool_config: { function_calling_config: { mode: "ANY" } },
         tools: [
           {
@@ -142,15 +123,29 @@ export class KurtVertexAI implements Kurt {
 
     return new KurtResult<T>(generator())
   }
-}
 
-function toVertexAIMessages(messages: KurtMessage[]): VertexAIMessage[] {
-  return messages.map((message) => {
-    const { role, text } = message
-    return { role, parts: [{ text }] }
-  })
+  private toVertexAIMessages = ({
+    prompt,
+    systemPrompt = this.options.systemPrompt,
+    extraMessages = [],
+  }: KurtGenerateNaturalLanguageOptions): VertexAIMessage[] => {
+    const systemMessage: VertexAIMessage[] = systemPrompt
+      ? [toVertexAIMessage({ role: "system", text: systemPrompt })]
+      : []
+
+    const userMessage = toVertexAIMessage({ role: "user", text: prompt })
+
+    const extras = extraMessages.map(toVertexAIMessage)
+
+    return systemMessage.concat(userMessage, extras)
+  }
 }
 
+const toVertexAIMessage = ({ role, text }: KurtMessage): VertexAIMessage => ({
+  role,
+  parts: [{ text }],
+})
+
 function jsonSchemaForVertexAI<T extends KurtSchemaInner>(
   zodSchema: KurtSchema<T>
 ) {