cleanaup bedrock provider

Author: narengogi

src/providers/bedrock/chatComplete.ts

@@ -1,8 +1,8 @@
 import {
   BEDROCK,
   documentMimeTypes,
-  fileExtensionMimeTypeMap,
   imagesMimeTypes,
+  fileExtensionMimeTypeMap,
 } from '../../globals';
 import {
   Message,
@@ -62,7 +62,7 @@ export interface BedrockChatCompletionsParams extends Params {
 }
 
 export interface BedrockConverseAnthropicChatCompletionsParams
-  extends Omit<BedrockChatCompletionsParams, 'anthropic_beta'> {
+  extends BedrockChatCompletionsParams {
   anthropic_version?: string;
   user?: string;
   thinking?: {
@@ -477,9 +477,6 @@ export const BedrockChatCompleteResponseTransform: (
   }
 
   if ('output' in response) {
-    const cacheReadInputTokens = response.usage?.cacheReadInputTokens || 0;
-    const cacheWriteInputTokens = response.usage?.cacheWriteInputTokens || 0;
-
     let content: string = '';
     content = response.output.message.content
       .filter((item) => item.text)
@@ -489,6 +486,9 @@ export const BedrockChatCompleteResponseTransform: (
       ? transformContentBlocks(response.output.message.content)
       : undefined;
 
+    const cacheReadInputTokens = response.usage?.cacheReadInputTokens || 0;
+    const cacheWriteInputTokens = response.usage?.cacheWriteInputTokens || 0;
+
     const responseObj: ChatCompletionResponse = {
       id: Date.now().toString(),
       object: 'chat.completion',
@@ -571,7 +571,6 @@ export const BedrockChatCompleteStreamChunkTransform: (
     streamState.currentToolCallIndex = -1;
   }
 
-  // final chunk
   if (parsedChunk.usage) {
     const cacheReadInputTokens = parsedChunk.usage?.cacheReadInputTokens || 0;
     const cacheWriteInputTokens = parsedChunk.usage?.cacheWriteInputTokens || 0;
@@ -605,9 +604,8 @@ export const BedrockChatCompleteStreamChunkTransform: (
           },
           // we only want to be sending this for anthropic models and this is not openai compliant
           ...((cacheReadInputTokens > 0 || cacheWriteInputTokens > 0) && {
-            cache_read_input_tokens: parsedChunk.usage.cacheReadInputTokens,
-            cache_creation_input_tokens:
-              parsedChunk.usage.cacheWriteInputTokens,
+            cache_read_input_tokens: cacheReadInputTokens,
+            cache_creation_input_tokens: cacheWriteInputTokens,
           }),
         },
       })}\n\n`,

src/providers/bedrock/constants.ts

@@ -1,3 +1,15 @@
+export const BEDROCK_STABILITY_V1_MODELS = [
+  'stable-diffusion-xl-v0',
+  'stable-diffusion-xl-v1',
+];
+
+export const bedrockInvokeModels = [
+  'cohere.command-light-text-v14',
+  'cohere.command-text-v14',
+  'ai21.j2-mid-v1',
+  'ai21.j2-ultra-v1',
+];
+
 export const LLAMA_2_SPECIAL_TOKENS = {
   BEGINNING_OF_SENTENCE: '<s>',
   END_OF_SENTENCE: '</s>',
@@ -34,15 +46,3 @@ export const MISTRAL_CONTROL_TOKENS = {
   MIDDLE: '[MIDDLE]',
   SUFFIX: '[SUFFIX]',
 };
-
-export const BEDROCK_STABILITY_V1_MODELS = [
-  'stable-diffusion-xl-v0',
-  'stable-diffusion-xl-v1',
-];
-
-export const bedrockInvokeModels = [
-  'cohere.command-light-text-v14',
-  'cohere.command-text-v14',
-  'ai21.j2-mid-v1',
-  'ai21.j2-ultra-v1',
-];

src/providers/bedrock/createFinetune.ts

@@ -46,7 +46,7 @@ export const BedrockCreateFinetuneConfig: ProviderConfig = {
         return undefined;
       }
       return {
-        s3Uri: decodeURIComponent(value.validation_file),
+        s3Uri: decodeURIComponent(value.validation_file ?? ''),
       };
     },
   },

src/providers/bedrock/embed.ts

@@ -60,6 +60,12 @@ export const BedrockCohereEmbedConfig: ProviderConfig = {
   },
 };
 
+const g1EmbedModels = [
+  'amazon.titan-embed-g1-text-02',
+  'amazon.titan-embed-text-v1',
+  'amazon.titan-embed-image-v1',
+];
+
 export const BedrockTitanEmbedConfig: ProviderConfig = {
   input: [
     {
@@ -117,6 +123,8 @@ export const BedrockTitanEmbedConfig: ProviderConfig = {
     param: 'embeddingTypes',
     required: false,
     transform: (params: any): string[] | undefined => {
+      const model = params.foundationModel || params.model || '';
+      if (g1EmbedModels.includes(model)) return undefined;
       if (Array.isArray(params.encoding_format)) return params.encoding_format;
       else if (typeof params.encoding_format === 'string')
         return [params.encoding_format];

src/providers/bedrock/getBatchOutput.ts

@@ -5,6 +5,7 @@ import { BedrockGetBatchResponse } from './types';
 import { getOctetStreamToOctetStreamTransformer } from '../../handlers/streamHandlerUtils';
 import { BedrockUploadFileResponseTransforms } from './uploadFileUtils';
 import { BEDROCK } from '../../globals';
+import { getAwsEndpointDomain } from './utils';
 
 const getModelProvider = (modelId: string) => {
   let provider = '';
@@ -89,7 +90,7 @@ export const BedrockGetBatchOutputRequestHandler = async ({
     const awsS3ObjectKey = `${primaryKey}${jobId}/${inputS3URIParts[inputS3URIParts.length - 1]}.out`;
     const awsModelProvider = batchDetails.modelId;
 
-    const s3FileURL = `https://${awsS3Bucket}.s3.${awsRegion}.amazonaws.com/${awsS3ObjectKey}`;
+    const s3FileURL = `https://${awsS3Bucket}.s3.${awsRegion}.${getAwsEndpointDomain(c)}/${awsS3ObjectKey}`;
     const s3FileHeaders = await BedrockAPIConfig.headers({
       c,
       providerOptions,

src/providers/bedrock/utils.ts

@@ -14,6 +14,9 @@ import { BedrockFinetuneRecord, BedrockInferenceProfile } from './types';
 import { FinetuneRequest } from '../types';
 import { BEDROCK } from '../../globals';
 
+export const getAwsEndpointDomain = (c: Context) =>
+  env(c).AWS_ENDPOINT_DOMAIN || 'amazonaws.com';
+
 export const generateAWSHeaders = async (
   body: Record<string, any> | string | undefined,
   headers: Record<string, string>,

src/types/requestBody.ts

@@ -422,7 +422,6 @@ export interface Params {
   // Google Vertex AI specific
   safety_settings?: any;
   // Anthropic specific
-  anthropic_beta?: string;
   anthropic_version?: string;
   thinking?: {
     type?: string;

src/utils/env.ts

@@ -0,0 +1,3 @@
+export interface Env {
+  AWS_ENDPOINT_DOMAIN: string;
+}