From 9b1c22c2abf6f745823945373fdb39bd2d2c785c Mon Sep 17 00:00:00 2001
From: Rafael Mosca <rafams@amazon.com>
Date: Mon, 6 Jan 2025 17:16:35 +0100
Subject: [PATCH] fix(bedrock): simplify ChatMessage syntax
 ChatMessage.userMessage() -> ChatMessage.user()

---
 src/cdk-lib/bedrock/prompts/prompt-variant.ts | 10 +++----
 test/cdk-lib/bedrock/prompts.test.ts          | 28 +++++++++----------
 test/integ/chatprompt.integ.ts                |  6 ++--
 3 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/src/cdk-lib/bedrock/prompts/prompt-variant.ts b/src/cdk-lib/bedrock/prompts/prompt-variant.ts
index 9157ea2f..02d92efc 100644
--- a/src/cdk-lib/bedrock/prompts/prompt-variant.ts
+++ b/src/cdk-lib/bedrock/prompts/prompt-variant.ts
@@ -92,10 +92,10 @@ export enum ChatMessageRole {
 }
 
 export class ChatMessage {
-  public static userMessage(text: string) {
+  public static user(text: string) {
     return new ChatMessage(ChatMessageRole.USER, text);
   }
-  public static assistantMessage(text: string) {
+  public static assistant(text: string) {
     return new ChatMessage(ChatMessageRole.ASSISTANT, text);
   }
   public readonly role: ChatMessageRole;
@@ -226,9 +226,9 @@ export abstract class PromptVariant {
           system: props.system ? [{ text: props.system }] : undefined,
           toolConfiguration: props.toolConfiguration
             ? {
-              toolChoice: props.toolConfiguration.toolChoice.__render(),
-              tools: props.toolConfiguration.tools,
-            }
+                toolChoice: props.toolConfiguration.toolChoice.__render(),
+                tools: props.toolConfiguration.tools,
+              }
             : undefined,
         },
       },
diff --git a/test/cdk-lib/bedrock/prompts.test.ts b/test/cdk-lib/bedrock/prompts.test.ts
index 1ffb8aec..7ad75f18 100644
--- a/test/cdk-lib/bedrock/prompts.test.ts
+++ b/test/cdk-lib/bedrock/prompts.test.ts
@@ -51,7 +51,7 @@ describe('Prompt', () => {
         Name: 'prompt1',
         Description: 'my cmk prompt',
         CustomerEncryptionKeyArn: cmk.keyArn,
-      }),
+      })
     );
     expect(prompt.promptName).toEqual('prompt1');
   });
@@ -62,7 +62,7 @@ describe('Prompt', () => {
     const variant1 = PromptVariant.text({
       variantName: 'variant1',
       model: BedrockFoundationModel.fromCdkFoundationModelId(
-        cdk_bedrock.FoundationModelIdentifier.ANTHROPIC_CLAUDE_3_SONNET_20240229_V1_0,
+        cdk_bedrock.FoundationModelIdentifier.ANTHROPIC_CLAUDE_3_SONNET_20240229_V1_0
       ),
       promptVariables: ['topic'],
       promptText: 'This is my first text prompt. Please summarize our conversation on {{topic}}.',
@@ -105,7 +105,7 @@ describe('Prompt', () => {
             TemplateType: 'TEXT',
           },
         ],
-      }),
+      })
     );
   });
 
@@ -155,7 +155,7 @@ describe('Prompt', () => {
             TemplateType: 'TEXT',
           },
         ],
-      }),
+      })
     );
   });
 
@@ -170,9 +170,9 @@ describe('Prompt', () => {
         maxTokens: 2000,
       },
       messages: [
-        ChatMessage.userMessage('From now on, you speak Japanese!'),
-        ChatMessage.assistantMessage('Konnichiwa!'),
-        ChatMessage.userMessage('From now on, you speak {{language}}!'),
+        ChatMessage.user('From now on, you speak Japanese!'),
+        ChatMessage.assistant('Konnichiwa!'),
+        ChatMessage.user('From now on, you speak {{language}}!'),
       ],
       system: 'You are a helpful assistant that only speaks the language you`re told.',
       promptVariables: ['language'],
@@ -294,7 +294,7 @@ describe('Prompt', () => {
             TemplateType: 'CHAT',
           },
         ],
-      }),
+      })
     );
   });
 
@@ -349,7 +349,7 @@ describe('Prompt', () => {
             TemplateType: 'TEXT',
           },
         ],
-      }),
+      })
     );
   });
 
@@ -368,12 +368,12 @@ describe('Prompt', () => {
       haveResource('AWS::Bedrock::Prompt', {
         Name: 'prompt1',
         Description: 'my versioned prompt',
-      }),
+      })
     );
     cdkExpect(stack).to(
       haveResource('AWS::Bedrock::PromptVersion', {
         Description: 'first version',
-      }),
+      })
     );
   });
 
@@ -395,7 +395,7 @@ describe('Prompt', () => {
       PromptVariant.text({
         variantName: `variant${id}`,
         model: BedrockFoundationModel.fromCdkFoundationModelId(
-          cdk_bedrock.FoundationModelIdentifier.ANTHROPIC_CLAUDE_3_SONNET_20240229_V1_0,
+          cdk_bedrock.FoundationModelIdentifier.ANTHROPIC_CLAUDE_3_SONNET_20240229_V1_0
         ),
         promptVariables: ['topic'],
         promptText: `This is my text prompt ${id}. Please summarize our conversation on {{topic}}.`,
@@ -404,7 +404,7 @@ describe('Prompt', () => {
           topP: 0.999,
           maxTokens: 2000,
         },
-      }),
+      })
     );
     new Prompt(stack, 'prompt1', {
       promptName: 'my-prompt',
@@ -413,7 +413,7 @@ describe('Prompt', () => {
     });
     // THEN
     expect(() => app.synth()).toThrow(
-      'Error: Too many variants specified. The maximum allowed is 3, but you have provided 4 variants.',
+      'Error: Too many variants specified. The maximum allowed is 3, but you have provided 4 variants.'
     );
   });
 });
diff --git a/test/integ/chatprompt.integ.ts b/test/integ/chatprompt.integ.ts
index 71f7e9ef..37eabcb4 100644
--- a/test/integ/chatprompt.integ.ts
+++ b/test/integ/chatprompt.integ.ts
@@ -24,9 +24,9 @@ const variantChat = PromptVariant.chat({
   variantName: 'variant1',
   model: BedrockFoundationModel.ANTHROPIC_CLAUDE_3_5_SONNET_V1_0,
   messages: [
-    ChatMessage.userMessage('From now on, you speak Japanese!'),
-    ChatMessage.assistantMessage('Konnichiwa!'),
-    ChatMessage.userMessage('From now on, you speak {{language}}!'),
+    ChatMessage.user('From now on, you speak Japanese!'),
+    ChatMessage.assistant('Konnichiwa!'),
+    ChatMessage.user('From now on, you speak {{language}}!'),
   ],
   system: 'You are a helpful assistant that only speaks the language you`re told.',
   promptVariables: ['language'],