Bedrock converse integration

Portkey-AI · Oct 14, 2024 · ecc2def · ecc2def
1 parent b64b8a2
commit ecc2def
Show file tree

Hide file tree

Showing 6 changed files with 262 additions and 1,344 deletions.
diff --git a/src/handlers/streamHandler.ts b/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,

diff --git a/src/providers/bedrock/api.ts b/src/providers/bedrock/api.ts
@@ -31,8 +31,12 @@ const BedrockAPIConfig: ProviderAPIConfig = {
     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') {
+      endpoint = `/model/${model}/converse`;
+      streamEndpoint = `/model/${model}/converse-stream`;
+    }
     switch (mappedFn) {
       case 'chatComplete': {
         return endpoint;