cleanaup bedrock provider

add wrapper for environment variables with support for fetching from file path

bedrock cleanup

add conditional import

Author: narengogi

src/handlers/handlerUtils.ts

@@ -856,7 +856,6 @@ export function constructConfigFromRequestHeaders(
     azureApiVersion: requestHeaders[`x-${POWERED_BY}-azure-api-version`],
     azureEndpointName: requestHeaders[`x-${POWERED_BY}-azure-endpoint-name`],
     azureFoundryUrl: requestHeaders[`x-${POWERED_BY}-azure-foundry-url`],
-    azureExtraParams: requestHeaders[`x-${POWERED_BY}-azure-extra-params`],
     azureAdToken: requestHeaders[`x-${POWERED_BY}-azure-ad-token`],
     azureAuthMode: requestHeaders[`x-${POWERED_BY}-azure-auth-mode`],
     azureManagedClientId:
@@ -865,6 +864,7 @@ export function constructConfigFromRequestHeaders(
     azureEntraClientSecret:
       requestHeaders[`x-${POWERED_BY}-azure-entra-client-secret`],
     azureEntraTenantId: requestHeaders[`x-${POWERED_BY}-azure-entra-tenant-id`],
+    azureExtraParameters: requestHeaders[`x-${POWERED_BY}-azure-extra-params`],
   };
 
   const awsConfig = {

src/providers/azure-ai-inference/api.ts

@@ -39,14 +39,14 @@ const AzureAIInferenceAPI: ProviderAPIConfig = {
   headers: async ({ providerOptions, fn }) => {
     const {
       apiKey,
-      azureExtraParams,
+      azureExtraParameters,
       azureDeploymentName,
       azureAdToken,
       azureAuthMode,
     } = providerOptions;
 
     const headers: Record<string, string> = {
-      'extra-parameters': azureExtraParams ?? 'drop',
+      'extra-parameters': azureExtraParameters ?? 'drop',
       ...(azureDeploymentName && {
         'azureml-model-deployment': azureDeploymentName,
       }),

src/providers/bedrock/api.ts

@@ -1,10 +1,10 @@
 import { Context } from 'hono';
-import { Options } from '../../types/requestBody';
+import { Options, Params } from '../../types/requestBody';
 import { endpointStrings, ProviderAPIConfig } from '../types';
 import { bedrockInvokeModels } from './constants';
 import {
+  getAwsEndpointDomain,
   generateAWSHeaders,
-  getAssumedRoleCredentials,
   getFoundationModelFromInferenceProfile,
   providerAssumedRoleCredentials,
 } from './utils';
@@ -18,6 +18,7 @@ interface BedrockAPIConfigInterface extends Omit<ProviderAPIConfig, 'headers'> {
     transformedRequestBody: Record<string, any> | string;
     transformedRequestUrl: string;
     gatewayRequestBody?: Params;
+    headers?: Record<string, string>;
   }) => Promise<Record<string, any>> | Record<string, any>;
 }
 
@@ -66,7 +67,14 @@ const ENDPOINTS_TO_ROUTE_TO_S3 = [
   'initiateMultipartUpload',
 ];
 
-const getMethod = (fn: endpointStrings, transformedRequestUrl: string) => {
+const getMethod = (
+  fn: endpointStrings,
+  transformedRequestUrl: string,
+  c: Context
+) => {
+  if (fn === 'proxy') {
+    return c.req.method;
+  }
   if (fn === 'uploadFile') {
     const url = new URL(transformedRequestUrl);
     return url.searchParams.get('partNumber') ? 'PUT' : 'POST';
@@ -121,36 +129,47 @@ const BedrockAPIConfig: BedrockAPIConfigInterface = {
         gatewayRequestURL.split('/v1/files/')[1]
       );
       const bucketName = s3URL.replace('s3://', '').split('/')[0];
-      return `https://${bucketName}.s3.${providerOptions.awsRegion || 'us-east-1'}.amazonaws.com`;
+      return `https://${bucketName}.s3.${providerOptions.awsRegion || 'us-east-1'}.${getAwsEndpointDomain(c)}`;
     }
     if (fn === 'retrieveFileContent') {
       const s3URL = decodeURIComponent(
         gatewayRequestURL.split('/v1/files/')[1]
       );
       const bucketName = s3URL.replace('s3://', '').split('/')[0];
-      return `https://${bucketName}.s3.${providerOptions.awsRegion || 'us-east-1'}.amazonaws.com`;
+      return `https://${bucketName}.s3.${providerOptions.awsRegion || 'us-east-1'}.${getAwsEndpointDomain(c)}`;
     }
     if (fn === 'uploadFile')
-      return `https://${providerOptions.awsS3Bucket}.s3.${providerOptions.awsRegion || 'us-east-1'}.amazonaws.com`;
+      return `https://${providerOptions.awsS3Bucket}.s3.${providerOptions.awsRegion || 'us-east-1'}.${getAwsEndpointDomain(c)}`;
     const isAWSControlPlaneEndpoint =
       fn && AWS_CONTROL_PLANE_ENDPOINTS.includes(fn);
-    return `https://${isAWSControlPlaneEndpoint ? 'bedrock' : 'bedrock-runtime'}.${providerOptions.awsRegion || 'us-east-1'}.amazonaws.com`;
+    return `https://${isAWSControlPlaneEndpoint ? 'bedrock' : 'bedrock-runtime'}.${providerOptions.awsRegion || 'us-east-1'}.${getAwsEndpointDomain(c)}`;
   },
   headers: async ({
     c,
     fn,
     providerOptions,
     transformedRequestBody,
     transformedRequestUrl,
+    gatewayRequestBody, // for proxy use the passed body blindly
+    headers: requestHeaders,
   }) => {
-    const method = getMethod(fn as endpointStrings, transformedRequestUrl);
-    const service = getService(fn as endpointStrings);
+    const { awsService } = providerOptions;
+    const method =
+      c.get('method') || // method set specifically into context
+      getMethod(fn as endpointStrings, transformedRequestUrl, c); // method calculated
+    const service = awsService || getService(fn as endpointStrings);
 
-    const headers: Record<string, string> = {
-      'content-type': 'application/json',
-    };
+    let headers: Record<string, string> = {};
 
-    if (method === 'PUT' || method === 'GET') {
+    if (fn === 'proxy' && service !== 'bedrock') {
+      headers = { ...(requestHeaders ?? {}) };
+    } else {
+      headers = {
+        'content-type': 'application/json',
+      };
+    }
+
+    if ((method === 'PUT' || method === 'GET') && fn !== 'proxy') {
       delete headers['content-type'];
     }
 
@@ -160,7 +179,8 @@ const BedrockAPIConfig: BedrockAPIConfigInterface = {
       await providerAssumedRoleCredentials(c, providerOptions);
     }
 
-    let finalRequestBody = transformedRequestBody;
+    let finalRequestBody =
+      fn === 'proxy' ? gatewayRequestBody : transformedRequestBody;
 
     if (['cancelFinetune', 'cancelBatch'].includes(fn as endpointStrings)) {
       // Cancel doesn't require any body, but fetch is sending empty body, to match the signature this block is required.
@@ -183,7 +203,6 @@ const BedrockAPIConfig: BedrockAPIConfigInterface = {
     fn,
     gatewayRequestBodyJSON: gatewayRequestBody,
     gatewayRequestURL,
-    c,
   }) => {
     if (fn === 'retrieveFile') {
       const fileId = decodeURIComponent(

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/listBatches.ts

@@ -28,12 +28,8 @@ export const BedrockListBatchesResponseTransform = (
       output_file_id: encodeURIComponent(
         batch.outputDataConfig.s3OutputDataConfig.s3Uri
       ),
-      finalizing_at: batch.endTime
-        ? new Date(batch.endTime).getTime()
-        : undefined,
-      expires_at: batch.jobExpirationTime
-        ? new Date(batch.jobExpirationTime).getTime()
-        : undefined,
+      finalizing_at: new Date(batch.endTime).getTime(),
+      expires_at: new Date(batch.jobExpirationTime).getTime(),
     }));
 
     return {

src/providers/bedrock/listFinetunes.ts

@@ -10,6 +10,7 @@ export const BedrockListFinetuneResponseTransform: (
   if (responseStatus !== 200) {
     return BedrockErrorResponseTransform(response) || response;
   }
+
   const records =
     response?.modelCustomizationJobSummaries as BedrockFinetuneRecord[];
   const openaiRecords = records.map(bedrockFinetuneToOpenAI);