bedrock cleanup

Author: narengogi

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/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);

src/providers/bedrock/types.ts

@@ -81,6 +81,16 @@ export interface BedrockInferenceProfile {
   type: string;
 }
 
+// https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html#API_runtime_Converse_ResponseSyntax
+export enum BEDROCK_STOP_REASON {
+  end_turn = 'end_turn',
+  tool_use = 'tool_use',
+  max_tokens = 'max_tokens',
+  stop_sequence = 'stop_sequence',
+  guardrail_intervened = 'guardrail_intervened',
+  content_filtered = 'content_filtered',
+}
+
 export interface BedrockMessagesParams extends MessageCreateParamsBase {
   additionalModelRequestFields?: Record<string, any>;
   additional_model_request_fields?: Record<string, any>;

src/providers/bedrock/uploadFileUtils.ts

@@ -830,6 +830,10 @@ interface BedrockAnthropicChatCompleteResponse {
   stop_reason: string;
   model: string;
   stop_sequence: null | string;
+  usage: {
+    input_tokens: number;
+    output_tokens: number;
+  };
 }
 
 export const BedrockAnthropicChatCompleteResponseTransform: (
@@ -874,9 +878,10 @@ export const BedrockAnthropicChatCompleteResponseTransform: (
         },
       ],
       usage: {
-        prompt_tokens: 0,
-        completion_tokens: 0,
-        total_tokens: 0,
+        prompt_tokens: response.usage.input_tokens,
+        completion_tokens: response.usage.output_tokens,
+        total_tokens:
+          response.usage.input_tokens + response.usage.output_tokens,
       },
     };
   }
@@ -933,6 +938,7 @@ export const BedrockMistralChatCompleteResponseTransform: (
           finish_reason: response.outputs[0].stop_reason,
         },
       ],
+      // mistral not sending usage.
       usage: {
         prompt_tokens: 0,
         completion_tokens: 0,

src/providers/bedrock/utils.ts

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