fix(bedrock): simplify ChatMessage syntax ChatMessage.userMessage() -…

…> ChatMessage.user()
awslabs · Jan 6, 2025 · 9b1c22c · 9b1c22c
1 parent bac7f58
commit 9b1c22c
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 22 deletions.
diff --git 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
@@ -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
@@ -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'],