bedrock-wrapper 2.2.0 → 2.3.0

package/bedrock-wrapper.js

@@ -21,35 +21,106 @@ import {
     getValueByPath,
     writeAsciiArt
 } from "./utils.js";
+import sharp from 'sharp';
 
 
 // write the ascii art logo on initial load
 writeAsciiArt();
 
 
+
 // -------------------
 // -- main function --
 // -------------------
+async function processImage(imageInput) {
+    let base64Image;
+    
+    if (typeof imageInput === 'string') {
+        if (imageInput.startsWith('data:image')) {
+            // Handle data URL
+            base64Image = imageInput.split(',')[1];
+        } else if (imageInput.startsWith('http')) {
+            // Handle URL
+            const response = await fetch(imageInput);
+            const buffer = await response.arrayBuffer();
+            base64Image = Buffer.from(buffer).toString('base64');
+        } else {
+            // Assume it's already base64
+            base64Image = imageInput;
+        }
+    } else if (Buffer.isBuffer(imageInput)) {
+        base64Image = imageInput.toString('base64');
+    }
+
+    // Process with sharp to ensure format and size compliance
+    const buffer = Buffer.from(base64Image, 'base64');
+    const processedImage = await sharp(buffer)
+        .resize(2048, 2048, { fit: 'inside' })
+        .toFormat('jpeg')
+        .toBuffer();
+    
+    return processedImage.toString('base64');
+}
+
 export async function* bedrockWrapper(awsCreds, openaiChatCompletionsCreateObject, { logging = false } = {} ) {
     const { region, accessKeyId, secretAccessKey } = awsCreds;
-    const { messages, model, max_tokens, stream, temperature, top_p } = openaiChatCompletionsCreateObject;
+    let { messages, model, max_tokens, stream, temperature, top_p, include_thinking_data } = openaiChatCompletionsCreateObject;
 
 
-    // retrieve the model configuration
-    const awsModel = bedrock_models.find((x) => (x.modelName.toLowerCase() === model.toLowerCase() || x.modelId.toLowerCase() === model.toLowerCase()));
-    if (!awsModel) { throw new Error(`Model configuration not found for model: ${model}`); }
+  let {awsModelId, awsModel} = findAwsModelWithId(model);
 
     // cleanup message content before formatting prompt message
     let message_cleaned = [];
     let system_message = "";
 
     for (let i = 0; i < messages.length; i++) {
-        if (messages[i].content !== "") {
-            // Extract system message only if model requires it as separate field
+        if (messages[i].content) {
+            let processedContent = messages[i].content;
+            
+            // Handle array format for messages with images
+            if (Array.isArray(processedContent)) {
+                let newContent = [];
+                for (const item of processedContent) {
+                    if (item.type === 'text') {
+                        newContent.push(item);
+                    } else if (item.type === 'image_url') {
+                        const processedImage = await processImage(
+                            typeof item.image_url === 'string' ? 
+                            item.image_url : 
+                            item.image_url.url
+                        );
+                        
+                        // Handle different model formats
+                        if (awsModel.messages_api) {
+                            newContent.push({
+                                type: 'image',
+                                source: {
+                                    type: 'base64',
+                                    media_type: 'image/jpeg',
+                                    data: processedImage
+                                }
+                            });
+                        } else {
+                            // Llama format for images
+                            newContent.push({
+                                type: 'image',
+                                image_data: {
+                                    url: `data:image/jpeg;base64,${processedImage}`
+                                }
+                            });
+                        }
+                    }
+                }
+                processedContent = newContent;
+            }
+
             if (awsModel.system_as_separate_field && messages[i].role === "system") {
-                system_message = messages[i].content;
+                system_message = processedContent;
             } else {
-                message_cleaned.push(messages[i]);
+                message_cleaned.push({
+                    ...messages[i],
+                    content: processedContent
+                });
             }
         } else if (awsModel.display_role_names) {
             message_cleaned.push(messages[i]);
@@ -67,26 +138,44 @@ export async function* bedrockWrapper(awsCreds, openaiChatCompletionsCreateObjec
         // convert message array to prompt object if model supports messages api
         prompt = message_cleaned;
     } else {
-        // convert message array to prompt string if model does not support messages api
         prompt = awsModel.bos_text;
         let eom_text_inserted = false;
+        
         for (let i = 0; i < message_cleaned.length; i++) {
             prompt += "\n";
-            if (message_cleaned[i].role === "system") {
-                prompt += awsModel.role_system_message_prefix;
-                prompt += awsModel.role_system_prefix;
-                if (awsModel.display_role_names) { prompt += message_cleaned[i].role; }
-                prompt += awsModel.role_system_suffix;
-                if (awsModel.display_role_names) {prompt += "\n"; }
-                prompt += message_cleaned[i].content;
-                prompt += awsModel.role_system_message_suffix;
-            } else if (message_cleaned[i].role === "user") {
+            
+            // Handle user messages with potential images
+            if (message_cleaned[i].role === "user") {
                 prompt += awsModel.role_user_message_prefix;
                 prompt += awsModel.role_user_prefix;
                 if (awsModel.display_role_names) { prompt += message_cleaned[i].role; }
                 prompt += awsModel.role_user_suffix;
-                if (awsModel.display_role_names) {prompt += "\n"; }
-                prompt += message_cleaned[i].content;
+                if (awsModel.display_role_names) { prompt += "\n"; }
+                
+                // Handle content array with text and images
+                if (Array.isArray(message_cleaned[i].content)) {
+                    let textContent = "";
+                    let imageContent = "";
+                    
+                    // Separate text and image content
+                    message_cleaned[i].content.forEach(item => {
+                        if (item.type === 'text') {
+                            textContent += item.text;
+                        } else if (item.type === 'image') {
+                            imageContent = item.image_data.url;
+                        }
+                    });
+                    
+                    // Format based on vision model requirements
+                    if (awsModel.vision && imageContent) {
+                        prompt += `\n${textContent}\n\n${imageContent}`;
+                    } else {
+                        prompt += textContent;
+                    }
+                } else {
+                    prompt += message_cleaned[i].content;
+                }
+                
                 prompt += awsModel.role_user_message_suffix;
             } else if (message_cleaned[i].role === "assistant") {
                 prompt += awsModel.role_assistant_message_prefix;
@@ -97,6 +186,7 @@ export async function* bedrockWrapper(awsCreds, openaiChatCompletionsCreateObjec
                 prompt += message_cleaned[i].content;
                 prompt += awsModel.role_assistant_message_suffix;
             }
+            
             if (message_cleaned[i+1] && message_cleaned[i+1].content === "") {
                 prompt += `\n${awsModel.eom_text}`;
                 eom_text_inserted = true;
@@ -105,16 +195,41 @@ export async function* bedrockWrapper(awsCreds, openaiChatCompletionsCreateObjec
             }
         }
     }
-    
-    // logging
+
+    // Add logging to see the final prompt
     if (logging) {
-        if (awsModel.system_as_separate_field && system_message) {
-            console.log(`\nsystem: ${system_message}`);
+        console.log("\nFinal formatted prompt:", prompt);
+    }
+
+    let max_gen_tokens = max_tokens <= awsModel.max_supported_response_tokens ? max_tokens : awsModel.max_supported_response_tokens;
+
+    if (awsModel.special_request_schema?.thinking?.type === "enabled") {
+        // temperature may only be set to 1 when thinking is enabled
+        temperature = 1;
+        // top_p must be unset when thinking is enabled
+        top_p = undefined;
+        // bugget_tokens can not be greater than 80% of max_gen_tokens
+        let budget_tokens = awsModel.special_request_schema?.thinking?.budget_tokens;
+        if (budget_tokens > (max_gen_tokens * 0.8)) {
+            budget_tokens = Math.floor(max_gen_tokens * 0.8);
+        }
+        if (budget_tokens < 1024) {
+            budget_tokens = 1024;
+        }
+        // if awsModel.special_request_schema?.thinking?.budget_tokens, set it to budget_tokens
+        if (awsModel.special_request_schema?.thinking?.budget_tokens) {
+            awsModel.special_request_schema.thinking.budget_tokens = budget_tokens;
+            // max_gen_tokens has to be greater than budget_tokens
+            if (max_gen_tokens <= budget_tokens) {
+                // make max_gen_tokens 20% greater than budget_tokens
+                max_gen_tokens = Math.floor(budget_tokens * 1.2);
+            }
         }
-        console.log(`\nprompt: ${typeof prompt === 'object' ? JSON.stringify(prompt) : prompt}\n`);
     }
 
-    const max_gen_tokens = max_tokens <= awsModel.max_supported_response_tokens ? max_tokens : awsModel.max_supported_response_tokens;
+    // if (logging) {
+    //     console.log("\nMax tokens:", max_gen_tokens);
+    // }
 
     // Format the request payload using the model's native structure.
     const request = awsModel.messages_api ? {
@@ -125,7 +240,16 @@ export async function* bedrockWrapper(awsCreds, openaiChatCompletionsCreateObjec
         top_p: top_p,
         ...awsModel.special_request_schema
     } : {
-        prompt,
+        prompt: typeof prompt === 'string' ? prompt : {
+            messages: prompt.map(msg => ({
+                role: msg.role,
+                content: Array.isArray(msg.content) ? 
+                    msg.content.map(item => 
+                        item.type === 'text' ? item.text : item
+                    ).join('\n') : 
+                    msg.content
+            }))
+        },
         // Optional inference parameters:
         [awsModel.max_tokens_param_name]: max_gen_tokens,
         temperature: temperature,
@@ -142,42 +266,105 @@ export async function* bedrockWrapper(awsCreds, openaiChatCompletionsCreateObjec
         },
     });
 
+    if (logging) {
+        console.log("\nFinal request:", JSON.stringify(request, null, 2));
+    }
+
     if (stream) {
         const responseStream = await client.send(
             new InvokeModelWithResponseStreamCommand({
                 contentType: "application/json",
                 body: JSON.stringify(request),
-                modelId: awsModel.modelId,
+                modelId: awsModelId,
             }),
         );
+        let is_thinking = false;
+        let should_think = awsModel.special_request_schema?.thinking?.type === "enabled";
+        
         for await (const event of responseStream.body) {
             const chunk = JSON.parse(new TextDecoder().decode(event.chunk.bytes));
-            let result = getValueByPath(chunk, awsModel.response_chunk_element);
+            let result;
+            result = getValueByPath(chunk, awsModel.response_chunk_element);
             if (result) {
+                if (should_think && is_thinking) {
+                    is_thinking = false;
+                    result = `</think>\n\n${result}`;
+                }
                 yield result;
-            }
+            } else {
+                if (include_thinking_data && awsModel.thinking_response_chunk_element) {
+                    let result = getValueByPath(chunk, awsModel.thinking_response_chunk_element);
+                    if (result && should_think && !is_thinking) {
+                        is_thinking = true;
+                        result = `<think>${result}`;
+                    }
+                    if (result) {
+                        yield result;
+                    }
+                } 
+            }        
         }
     } else {
         const apiResponse = await client.send(
             new InvokeModelCommand({
               contentType: "application/json",
               body: JSON.stringify(request),
-              modelId: awsModel.modelId,
+              modelId: awsModelId,
             }),
           );
-        
+
         const decodedBodyResponse = JSON.parse(new TextDecoder().decode(apiResponse.body));
-        let result;
+        let thinking_result;
+        let text_result;
+
+        if (awsModel.thinking_response_nonchunk_element) {
+            thinking_result = getValueByPath(decodedBodyResponse, awsModel.thinking_response_nonchunk_element);
+        }
+
         if (awsModel.response_nonchunk_element) {
-            result = getValueByPath(decodedBodyResponse, awsModel.response_nonchunk_element);
-        } else {
-            result = getValueByPath(decodedBodyResponse, awsModel.response_chunk_element);
+            text_result = getValueByPath(decodedBodyResponse, awsModel.response_nonchunk_element);
         }
+        if (!text_result) {
+            if (awsModel.response_chunk_element) {
+                text_result = getValueByPath(decodedBodyResponse, awsModel.response_chunk_element);
+            }
+            if (!text_result && awsModel.response_nonchunk_element) {
+                // replace [0] with [1]
+                awsModel.response_nonchunk_element = awsModel.response_nonchunk_element.replace('[0]', '[1]');
+                text_result = getValueByPath(decodedBodyResponse, awsModel.response_nonchunk_element);
+            }
+        }
+
+        let result = thinking_result ? `<think>${thinking_result}</think>\n\n${text_result}` : text_result;
         yield result;
     }    
 }
 
 
+// ----------------------------------------------------
+// -- lookup model configuration by model id or name --
+// -----------------------------------------------------------------------------
+// -- partial model id or model name is accepted (cross-region model support) --
+// -- returns model configuration object and model id                         --
+// -----------------------------------------------------------------------------
+function findAwsModelWithId(model) {
+    const matchingModel = bedrock_models.find(candidate =>
+        model === candidate.modelName ||
+        model.endsWith(candidate.modelId)
+    );
+
+    if (!matchingModel) {
+        throw new Error(`Model configuration not found for model: ${model}`);
+    }
+
+    return {
+        awsModelId: model.endsWith(matchingModel.modelId) ? model : matchingModel.modelId,
+        awsModel: matchingModel
+    };
+}
+
+
+
 // ---------------------------
 // -- list supported models --
 // ---------------------------

package/interactive-example.js

@@ -122,6 +122,7 @@ const openaiChatCompletionsCreateObject = {
     "stream": shouldStream,
     "temperature": LLM_TEMPERATURE,
     "top_p": LLM_TOP_P,
+    "include_thinking_data": true,
 };
 
 

package/package.json

@@ -1,38 +1,42 @@
-{
-  "name": "bedrock-wrapper",
-  "version": "2.2.0",
-  "description": "🪨 Bedrock Wrapper is an npm package that simplifies the integration of existing OpenAI-compatible API objects with AWS Bedrock's serverless inference LLMs.",
-  "homepage": "https://www.equilllabs.com/projects/bedrock-wrapper",
-  "repository": {
-    "type": "git",
-    "url": "https://github.com/jparkerweb/bedrock-wrapper.git"
-  },
-  "bugs": {
-    "url": "https://github.com/jparkerweb/bedrock-wrapper/issues",
-    "email": "equilllabs@gmail.com"
-  },
-  "scripts": {
-    "clean": "rm -rf node_modules && rm -rf package-lock.json && npm install"
-  },
-  "main": "bedrock-wrapper.js",
-  "type": "module",
-  "keywords": [
-    "openai",
-    "bedrock",
-    "aws",
-    "wrapper",
-    "proxy",
-    "serverless",
-    "inference",
-    "llm"
-  ],
-  "author": "",
-  "license": "ISC",
-  "dependencies": {
-    "@aws-sdk/client-bedrock-runtime": "^3.716.0",
-    "dotenv": "^16.4.5"
-  },
-  "devDependencies": {
-    "chalk": "^5.4.1"
-  }
-}
+{
+  "name": "bedrock-wrapper",
+  "version": "2.3.0",
+  "description": "🪨 Bedrock Wrapper is an npm package that simplifies the integration of existing OpenAI-compatible API objects with AWS Bedrock's serverless inference LLMs.",
+  "homepage": "https://www.equilllabs.com/projects/bedrock-wrapper",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/jparkerweb/bedrock-wrapper.git"
+  },
+  "bugs": {
+    "url": "https://github.com/jparkerweb/bedrock-wrapper/issues",
+    "email": "equilllabs@gmail.com"
+  },
+  "scripts": {
+    "clean": "npx rimraf node_modules && npx rimraf package-lock.json && npm install",
+    "test": "node test-models.js",
+    "test-vision": "node test-vision.js",
+    "interactive": "node interactive-example.js"
+  },
+  "main": "bedrock-wrapper.js",
+  "type": "module",
+  "keywords": [
+    "openai",
+    "bedrock",
+    "aws",
+    "wrapper",
+    "proxy",
+    "serverless",
+    "inference",
+    "llm"
+  ],
+  "author": "",
+  "license": "ISC",
+  "dependencies": {
+    "@aws-sdk/client-bedrock-runtime": "^3.755.0",
+    "dotenv": "^16.4.5",
+    "sharp": "^0.33.5"
+  },
+  "devDependencies": {
+    "chalk": "^5.4.1"
+  }
+}

package/test-models-output.txt

@@ -1,10 +1,26 @@
 Test Question: "Respond with exactly one word: What is 1+1?"
 ==================================================
 
-Starting tests with 17 models...
+Starting tests with 19 models...
 Each model will be tested with streaming and non-streaming calls
 
 
+--------------------------------------------------
+Testing Claude-3-7-Sonnet-Thinking ⇢
+Streaming test passed for Claude-3-7-Sonnet-Thinking: "Two"
+Non-streaming test passed for Claude-3-7-Sonnet-Thinking: "<think>The question asks for exactly one word in response to "What is 1+1?"
+
+The answer to 1+1 is 2, which is a single word (a number). 
+
+I need to respond with exactly one word, and that word should be "2".</think>
+
+2"
+
+--------------------------------------------------
+Testing Claude-3-7-Sonnet ⇢
+Streaming test passed for Claude-3-7-Sonnet: "2"
+Non-streaming test passed for Claude-3-7-Sonnet: "Two"
+
 --------------------------------------------------
 Testing Claude-3-5-Sonnet-v2 ⇢
 Streaming test passed for Claude-3-5-Sonnet-v2: "two"
@@ -27,7 +43,7 @@ Non-streaming test passed for Claude-3-Haiku: "Two."
 
 --------------------------------------------------
 Testing Llama-3-3-70b ⇢
-Streaming test passed for Llama-3-3-70b: "Two."
+Streaming test passed for Llama-3-3-70b: "Two"
 Non-streaming test passed for Llama-3-3-70b: "Two."
 
 --------------------------------------------------
@@ -62,7 +78,7 @@ Non-streaming test passed for Llama-3-1-70b: "Two."
 
 --------------------------------------------------
 Testing Llama-3-1-405b ⇢
-Streaming test passed for Llama-3-1-405b: "Two."
+Streaming test passed for Llama-3-1-405b: "Two"
 Non-streaming test passed for Llama-3-1-405b: "Two."
 
 --------------------------------------------------
@@ -77,17 +93,17 @@ Non-streaming test passed for Llama-3-70b: "Two"
 
 --------------------------------------------------
 Testing Mistral-7b ⇢
-Streaming test passed for Mistral-7b: "Two. (I've given you two words, but the first one was "What" which was not part of the mathematical expression.)"
+Streaming test passed for Mistral-7b: "Two. (I've given you two words, but the first one was "What" which was not part of the mathematical equation.)"
 Non-streaming test passed for Mistral-7b: "Two. (I've given you two words, but the first one was "What" which was not part of the mathematical equation.)"
 
 --------------------------------------------------
 Testing Mixtral-8x7b ⇢
 Streaming test passed for Mixtral-8x7b: "Two.
 
-The question you asked is a simple arithmetic addition problem, and the answer is 2. It is considered good manners to respond to a direct question, even if it is a simple one. This helps to maintain a respectful and positive interaction."
+The word you are looking for is "two." The sum of 1 + 1 is equal to 2. I am programmed to provide accurate and helpful responses, so I wanted to make sure that I gave you the correct answer to your question. Is there anything else you would like to know? I'm here to help! Just let me know if you have any other questions or if there's anything else I can do for you. I'm always happy to assist with whatever you need."
 Non-streaming test passed for Mixtral-8x7b: "Two.
 
-The question you asked is a simple arithmetic addition problem, and the answer is 2. It is considered good manners to respond to a direct question, even if it is a simple one. This helps to maintain a respectful and positive interaction."
+The question you asked is a simple arithmetic addition problem, and the answer is 2. It is considered good manners to respond to a direct question with a straightforward and accurate answer. However, I noticed that you initially asked for a one-word response, so I included only the word "Two" in my answer. I hope this is what you were looking for! Is there anything else I can help you with?"
 
 --------------------------------------------------
 Testing Mistral-Large ⇢