From d26b96d12e7acd399bd08738edc17147ab6ee120 Mon Sep 17 00:00:00 2001
From: Alain Krok <alkrok@amazon.com>
Date: Tue, 7 May 2024 15:56:56 -0400
Subject: [PATCH] chore(model): add support for titan premier (#435)

chore(model): add support for titan premier
---
 apidocs/classes/bedrock.BedrockFoundationModel.md          | 7 +++++++
 .../question_answering/src/adapters/bedrock/titan.py       | 1 +
 .../question_answering/src/llms/types.py                   | 2 ++
 src/cdk-lib/bedrock/models.ts                              | 6 +++++-
 4 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/apidocs/classes/bedrock.BedrockFoundationModel.md b/apidocs/classes/bedrock.BedrockFoundationModel.md
index 19fbb4a0..b63711db 100644
--- a/apidocs/classes/bedrock.BedrockFoundationModel.md
+++ b/apidocs/classes/bedrock.BedrockFoundationModel.md
@@ -21,6 +21,7 @@ can instantiate a `BedrockFoundationModel` object, e.g: `new BedrockFoundationMo
 - [supportsAgents](bedrock.BedrockFoundationModel.md#supportsagents)
 - [supportsKnowledgeBase](bedrock.BedrockFoundationModel.md#supportsknowledgebase)
 - [vectorDimensions](bedrock.BedrockFoundationModel.md#vectordimensions)
+- [AMAZON\_TITAN\_PREMIER\_V1\_0](bedrock.BedrockFoundationModel.md#amazon_titan_premier_v1_0)
 - [AMAZON\_TITAN\_TEXT\_EXPRESS\_V1](bedrock.BedrockFoundationModel.md#amazon_titan_text_express_v1)
 - [ANTHROPIC\_CLAUDE\_HAIKU\_V1\_0](bedrock.BedrockFoundationModel.md#anthropic_claude_haiku_v1_0)
 - [ANTHROPIC\_CLAUDE\_INSTANT\_V1\_2](bedrock.BedrockFoundationModel.md#anthropic_claude_instant_v1_2)
@@ -79,6 +80,12 @@ ___
 
 ___
 
+### AMAZON\_TITAN\_PREMIER\_V1\_0
+
+▪ `Static` `Readonly` **AMAZON\_TITAN\_PREMIER\_V1\_0**: [`BedrockFoundationModel`](bedrock.BedrockFoundationModel.md)
+
+___
+
 ### AMAZON\_TITAN\_TEXT\_EXPRESS\_V1
 
 ▪ `Static` `Readonly` **AMAZON\_TITAN\_TEXT\_EXPRESS\_V1**: [`BedrockFoundationModel`](bedrock.BedrockFoundationModel.md)
diff --git a/lambda/aws-qa-appsync-opensearch/question_answering/src/adapters/bedrock/titan.py b/lambda/aws-qa-appsync-opensearch/question_answering/src/adapters/bedrock/titan.py
index 36b1e21d..feee2cb1 100644
--- a/lambda/aws-qa-appsync-opensearch/question_answering/src/adapters/bedrock/titan.py
+++ b/lambda/aws-qa-appsync-opensearch/question_answering/src/adapters/bedrock/titan.py
@@ -75,3 +75,4 @@ def get_prompt(self):
 # Register the adapter
 registry.register(r"^Bedrock.amazon.titan-t*", BedrockTitanAdapter)
 registry.register(r"^Bedrock.amazon.titan-e*", BedrockTitanAdapter)
+registry.register(r"^Bedrock.amazon.titan-p*", BedrockTitanAdapter)
diff --git a/lambda/aws-qa-appsync-opensearch/question_answering/src/llms/types.py b/lambda/aws-qa-appsync-opensearch/question_answering/src/llms/types.py
index 107d0c5c..a27daa4b 100644
--- a/lambda/aws-qa-appsync-opensearch/question_answering/src/llms/types.py
+++ b/lambda/aws-qa-appsync-opensearch/question_answering/src/llms/types.py
@@ -21,6 +21,7 @@ class BedrockModel(StrEnum):
     AMAZON_TITAN_EMBED_IMAGE_V1 = 'amazon.titan-embed-image-v1',
     AMAZON_TITAN_EMBED_TEXT_V1 = 'amazon.titan-embed-text-v1',
     AMAZON_TITAN_IMAGE_GENERATOR_V1 = 'amazon.titan-image-generator-v1',
+    AMAZON_TITAN_TEXT_PREMIER_V1 = 'amazon.titan-text-premier-v1:0'
 
 MAX_TOKENS_MAP = {
     Provider.BEDROCK+'.'+BedrockModel.ANTHROPIC_CLAUDE_V2_1 : 200000,
@@ -33,6 +34,7 @@ class BedrockModel(StrEnum):
     Provider.BEDROCK+'.'+BedrockModel.AMAZON_TITAN_EMBED_TEXT_V1: 8000,
     Provider.BEDROCK+'.'+BedrockModel.AMAZON_TITAN_EMBED_IMAGE_V1: 128,
     Provider.BEDROCK+'.'+BedrockModel.AMAZON_TITAN_IMAGE_GENERATOR_V1: 77,
+    Provider.BEDROCK+'.'+BedrockModel.AMAZON_TITAN_TEXT_PREMIER_V1: 32000,
 }
 
 
diff --git a/src/cdk-lib/bedrock/models.ts b/src/cdk-lib/bedrock/models.ts
index a39cb323..23d04ed1 100644
--- a/src/cdk-lib/bedrock/models.ts
+++ b/src/cdk-lib/bedrock/models.ts
@@ -65,6 +65,10 @@ export class BedrockFoundationModel {
     'anthropic.claude-3-haiku-20240307-v1:0',
     { supportsAgents: true },
   );
+  public static readonly AMAZON_TITAN_PREMIER_V1_0 = new BedrockFoundationModel(
+    'amazon.titan-text-premier-v1:0',
+    { supportsAgents: true },
+  );
 
   public static readonly TITAN_EMBED_TEXT_V1 = new BedrockFoundationModel(
     'amazon.titan-embed-text-v1',
@@ -98,4 +102,4 @@ export class BedrockFoundationModel {
     const region = cdk.Stack.of(construct).region;
     return `arn:aws:bedrock:${region}::foundation-model/${this.modelId}`;
   }
-}
+}
\ No newline at end of file