Merge pull request #1755 from benjidotsh/ai/bedrock-max-retries-config

feat(ai): make AWS Bedrock maxRetries configurable
RoadieHQ · Dec 11, 2024 · 2344125 · 2344125
2 parents 5417267 + 4e08917
commit 2344125
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 6 deletions.
diff --git a/.changeset/serious-insects-melt.md b/.changeset/serious-insects-melt.md
@@ -0,0 +1,5 @@
+---
+'@roadiehq/rag-ai-backend-embeddings-aws': minor
+---
+
+Made `maxRetries` configurable
diff --git a/plugins/backend/rag-ai-backend-embeddings-aws/README.md b/plugins/backend/rag-ai-backend-embeddings-aws/README.md
@@ -31,14 +31,16 @@ The module expects 2 pieces of configuration:
 
 ### AWS Bedrock configuration
 
-The module expects the name of the embeddings generative AI model to be configured via app-config.
+The module expects the name of the embeddings generative AI model to be configured via app-config. Optionally, you can also provide the maximum number of retries for the AWS SDK client.
 
 ```yaml
 ai:
   embeddings:
-    awsBedrock:
+    bedrock:
       # Name of the Bedrock model to use to create Embeddings
       modelName: 'amazon.titan-embed-text-v1'
+      # Maximum number of retries for the AWS SDK client
+      maxRetries: 3
 ```
 
 ### AWS Credentials Configuration

diff --git a/plugins/backend/rag-ai-backend-embeddings-aws/config.d.ts b/plugins/backend/rag-ai-backend-embeddings-aws/config.d.ts
@@ -21,11 +21,15 @@ export interface Config {
    */
   ai: {
     embeddings: {
-      awsBedrock: {
+      bedrock: {
         /**
          * Name of the Bedrock model to use to create Embeddings
          */
         modelName: string;
+        /**
+         * Maximum number of retries for the AWS SDK client. Defaults to 3
+         */
+        maxRetries?: string;
       };
     };
   };

diff --git a/plugins/backend/rag-ai-backend-embeddings-aws/src/RoadieBedrockAugmenter.ts b/plugins/backend/rag-ai-backend-embeddings-aws/src/RoadieBedrockAugmenter.ts
@@ -23,6 +23,7 @@ import { BedrockCohereEmbeddings } from './BedrockCohereEmbeddings';
 
 export type BedrockConfig = {
   modelName: string;
+  maxRetries?: number;
 };
 
 export class RoadieBedrockAugmenter extends DefaultVectorAugmentationIndexer {
@@ -41,8 +42,14 @@ export class RoadieBedrockAugmenter extends DefaultVectorAugmentationIndexer {
       model: config.bedrockConfig.modelName,
     };
     const embeddings = config.bedrockConfig.modelName.includes('cohere')
-      ? new BedrockCohereEmbeddings({ ...embeddingsConfig, maxRetries: 3 })
-      : new BedrockEmbeddings({ ...embeddingsConfig, maxRetries: 3 });
+      ? new BedrockCohereEmbeddings({
+          ...embeddingsConfig,
+          maxRetries: config.bedrockConfig.maxRetries ?? 3,
+        })
+      : new BedrockEmbeddings({
+          ...embeddingsConfig,
+          maxRetries: config.bedrockConfig.maxRetries ?? 3,
+        });
 
     super({ ...config, embeddings });
   }

diff --git a/plugins/backend/rag-ai-backend/README.md b/plugins/backend/rag-ai-backend/README.md
@@ -61,7 +61,7 @@ ai:
     concurrencyLimit: 10
 
     # AWS Bedrock Embeddings configuration
-    awsBedrock:
+    bedrock:
       # (Required) Name of the Bedrock model to use to create Embeddings.
       modelName: 'amazon.titan-embed-text-v1'