Merge pull request #678 from narengogi/feat/converse-api

Feature: Bedrock converse integration

Author: VisargD

src/globals.ts

@@ -139,3 +139,52 @@ export const MULTIPART_FORM_DATA_ENDPOINTS: endpointStrings[] = [
   'createTranscription',
   'createTranslation',
 ];
+
+export const fileExtensionMimeTypeMap = {
+  mp4: 'video/mp4',
+  jpeg: 'image/jpeg',
+  jpg: 'image/jpeg',
+  png: 'image/png',
+  bmp: 'image/bmp',
+  tiff: 'image/tiff',
+  webp: 'image/webp',
+  pdf: 'application/pdf',
+  csv: 'text/csv',
+  doc: 'application/msword',
+  docx: 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+  xls: 'application/vnd.ms-excel',
+  xlsx: 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
+  html: 'text/html',
+  md: 'text/markdown',
+  mp3: 'audio/mp3',
+  wav: 'audio/wav',
+  txt: 'text/plain',
+  mov: 'video/mov',
+  mpeg: 'video/mpeg',
+  mpg: 'video/mpg',
+  avi: 'video/avi',
+  wmv: 'video/wmv',
+  mpegps: 'video/mpegps',
+  flv: 'video/flv',
+};
+
+export const imagesMimeTypes = [
+  fileExtensionMimeTypeMap.jpeg,
+  fileExtensionMimeTypeMap.jpg,
+  fileExtensionMimeTypeMap.png,
+  fileExtensionMimeTypeMap.bmp,
+  fileExtensionMimeTypeMap.tiff,
+  fileExtensionMimeTypeMap.webp,
+];
+
+export const documentMimeTypes = [
+  fileExtensionMimeTypeMap.pdf,
+  fileExtensionMimeTypeMap.csv,
+  fileExtensionMimeTypeMap.doc,
+  fileExtensionMimeTypeMap.docx,
+  fileExtensionMimeTypeMap.xls,
+  fileExtensionMimeTypeMap.xlsx,
+  fileExtensionMimeTypeMap.html,
+  fileExtensionMimeTypeMap.md,
+  fileExtensionMimeTypeMap.txt,
+];

src/handlers/streamHandler.ts

@@ -33,7 +33,10 @@ function getPayloadFromAWSChunk(chunk: Uint8Array): string {
 
   const payloadLength = chunkLength - headersEnd - 4; // Subtracting 4 for the message crc
   const payload = chunk.slice(headersEnd, headersEnd + payloadLength);
-  return decoder.decode(payload);
+  const decodedJson = JSON.parse(decoder.decode(payload));
+  return decodedJson.bytes
+    ? Buffer.from(decodedJson.bytes, 'base64').toString()
+    : JSON.stringify(decodedJson);
 }
 
 function concatenateUint8Arrays(a: Uint8Array, b: Uint8Array): Uint8Array {
@@ -60,10 +63,7 @@ export async function* readAWSStream(
           const data = buffer.subarray(0, expectedLength);
           buffer = buffer.subarray(expectedLength);
           expectedLength = readUInt32BE(buffer, 0);
-          const payload = Buffer.from(
-            JSON.parse(getPayloadFromAWSChunk(data)).bytes,
-            'base64'
-          ).toString();
+          const payload = getPayloadFromAWSChunk(data);
           if (transformFunction) {
             const transformedChunk = transformFunction(
               payload,
@@ -96,11 +96,7 @@ export async function* readAWSStream(
       buffer = buffer.subarray(expectedLength);
 
       expectedLength = readUInt32BE(buffer, 0);
-      const payload = Buffer.from(
-        JSON.parse(getPayloadFromAWSChunk(data)).bytes,
-        'base64'
-      ).toString();
-
+      const payload = getPayloadFromAWSChunk(data);
       if (transformFunction) {
         const transformedChunk = transformFunction(
           payload,

src/providers/bedrock/api.ts

@@ -1,4 +1,6 @@
+import { GatewayError } from '../../errors/GatewayError';
 import { ProviderAPIConfig } from '../types';
+import { bedrockInvokeModels } from './constants';
 import { generateAWSHeaders } from './utils';
 
 const BedrockAPIConfig: ProviderAPIConfig = {
@@ -27,12 +29,20 @@ const BedrockAPIConfig: ProviderAPIConfig = {
   },
   getEndpoint: ({ fn, gatewayRequestBody }) => {
     const { model, stream } = gatewayRequestBody;
+    if (!model) throw new GatewayError('Model is required');
     let mappedFn = fn;
     if (stream) {
       mappedFn = `stream-${fn}`;
     }
-    const endpoint = `/model/${model}/invoke`;
-    const streamEndpoint = `/model/${model}/invoke-with-response-stream`;
+    let endpoint = `/model/${model}/invoke`;
+    let streamEndpoint = `/model/${model}/invoke-with-response-stream`;
+    if (
+      (mappedFn === 'chatComplete' || mappedFn === 'stream-chatComplete') &&
+      !bedrockInvokeModels.includes(model)
+    ) {
+      endpoint = `/model/${model}/converse`;
+      streamEndpoint = `/model/${model}/converse-stream`;
+    }
     switch (mappedFn) {
       case 'chatComplete': {
         return endpoint;