Merge pull request #360 from janhq/303-feat-nitro-chat-completion-wit…

…h-image-supporting-local-image-path feat: Support for LVM with openai compatible API with local image path
janhq · Jan 24, 2024 · cb6c3d8 · cb6c3d8
2 parents 92b5a5c + 236e3ba
commit cb6c3d8
Show file tree

Hide file tree

Showing 3 changed files with 108 additions and 6 deletions.
diff --git a/controllers/llamaCPP.cc b/controllers/llamaCPP.cc
@@ -222,18 +222,33 @@ void llamaCPP::chatCompletion(
           for (auto content_piece : message["content"]) {
             role = user_prompt;
 
+            json content_piece_image_data;
+            content_piece_image_data["data"] = "";
+
             auto content_piece_type = content_piece["type"].asString();
             if (content_piece_type == "text") {
               auto text = content_piece["text"].asString();
               formatted_output += text;
             } else if (content_piece_type == "image_url") {
               auto image_url = content_piece["image_url"]["url"].asString();
-              auto base64_image_data = nitro_utils::extractBase64(image_url);
-              LOG_INFO << base64_image_data;
-              formatted_output += "[img-" + std::to_string(no_images) + "]";
-
-              json content_piece_image_data;
+              std::string base64_image_data;
+              if (image_url.find("http") != std::string::npos) {
+                LOG_INFO << "Remote image detected but not supported yet";
+              } else if (image_url.find("data:image") != std::string::npos) {
+                LOG_INFO << "Base64 image detected";
+                base64_image_data = nitro_utils::extractBase64(image_url);
+                LOG_INFO << base64_image_data;
+              } else {
+                LOG_INFO << "Local image detected";
+                nitro_utils::processLocalImage(
+                    image_url, [&](const std::string &base64Image) {
+                      base64_image_data = base64Image;
+                    });
+                LOG_INFO << base64_image_data;
+              }
               content_piece_image_data["data"] = base64_image_data;
+
+              formatted_output += "[img-" + std::to_string(no_images) + "]";
               content_piece_image_data["id"] = no_images;
               data["image_data"].push_back(content_piece_image_data);
               no_images++;

diff --git a/llama.cpp b/llama.cpp
diff --git a/utils/nitro_utils.h b/utils/nitro_utils.h
@@ -3,10 +3,13 @@
 #include "random"
 #include "string"
 #include <algorithm>
+#include <drogon/HttpClient.h>
 #include <drogon/HttpResponse.h>
+#include <fstream>
 #include <iostream>
 #include <ostream>
 #include <regex>
+#include <vector>
 // Include platform-specific headers
 #ifdef _WIN32
 #include <winsock2.h>
@@ -32,6 +35,90 @@ inline std::string extractBase64(const std::string &input) {
   return "";
 }
 
+// Helper function to encode data to Base64
+inline std::string base64Encode(const std::vector<unsigned char> &data) {
+  static const char encodingTable[] =
+      "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
+  std::string encodedData;
+  int i = 0;
+  int j = 0;
+  unsigned char array3[3];
+  unsigned char array4[4];
+
+  for (unsigned char c : data) {
+    array3[i++] = c;
+    if (i == 3) {
+      array4[0] = (array3[0] & 0xfc) >> 2;
+      array4[1] = ((array3[0] & 0x03) << 4) + ((array3[1] & 0xf0) >> 4);
+      array4[2] = ((array3[1] & 0x0f) << 2) + ((array3[2] & 0xc0) >> 6);
+      array4[3] = array3[2] & 0x3f;
+
+      for (i = 0; i < 4; i++)
+        encodedData += encodingTable[array4[i]];
+      i = 0;
+    }
+  }
+
+  if (i) {
+    for (j = i; j < 3; j++)
+      array3[j] = '\0';
+
+    array4[0] = (array3[0] & 0xfc) >> 2;
+    array4[1] = ((array3[0] & 0x03) << 4) + ((array3[1] & 0xf0) >> 4);
+    array4[2] = ((array3[1] & 0x0f) << 2) + ((array3[2] & 0xc0) >> 6);
+
+    for (j = 0; j < i + 1; j++)
+      encodedData += encodingTable[array4[j]];
+
+    while (i++ < 3)
+      encodedData += '=';
+  }
+
+  return encodedData;
+}
+
+// Function to load an image and convert it to Base64
+inline std::string imageToBase64(const std::string &imagePath) {
+  std::ifstream imageFile(imagePath, std::ios::binary);
+  if (!imageFile.is_open()) {
+    throw std::runtime_error("Could not open the image file.");
+  }
+
+  std::vector<unsigned char> buffer(std::istreambuf_iterator<char>(imageFile),
+                                    {});
+  return base64Encode(buffer);
+}
+
+// Helper function to generate a unique filename
+inline std::string generateUniqueFilename(const std::string &prefix,
+                                          const std::string &extension) {
+  // Get current time as a timestamp
+  auto now = std::chrono::system_clock::now();
+  auto now_ms = std::chrono::time_point_cast<std::chrono::milliseconds>(now);
+  auto epoch = now_ms.time_since_epoch();
+  auto value = std::chrono::duration_cast<std::chrono::milliseconds>(epoch);
+
+  // Generate a random number
+  std::random_device rd;
+  std::mt19937 gen(rd());
+  std::uniform_int_distribution<> dis(1000, 9999);
+
+  std::stringstream ss;
+  ss << prefix << value.count() << "_" << dis(gen) << extension;
+  return ss.str();
+}
+
+inline void
+processLocalImage(const std::string &localPath,
+                  std::function<void(const std::string &)> callback) {
+  try {
+    std::string base64Image = imageToBase64(localPath);
+    callback(base64Image); // Invoke the callback with the Base64 string
+  } catch (const std::exception &e) {
+    std::cerr << "Error during processing: " << e.what() << std::endl;
+  }
+}
+
 inline std::vector<std::string> listFilesInDir(const std::string &path) {
   std::vector<std::string> files;