update

Author: namchuai

engine/common/download_task.h

@@ -4,16 +4,13 @@
 #include <filesystem>
 #include <sstream>
 #include <string>
-#include <unordered_map>
 
 enum class DownloadType { Model, Engine, Miscellaneous, CudaToolkit, Cortex };
 
 struct DownloadItem {
 
   std::string id;
 
-  std::optional<std::unordered_map<std::string, std::string>> headers;
-
   std::string downloadUrl;
 
   /**

engine/services/download_service.cc

@@ -8,6 +8,7 @@
 #include <ostream>
 #include <utility>
 #include "download_service.h"
+#include "utils/curl_utils.h"
 #include "utils/format_utils.h"
 #include "utils/logging_utils.h"
 #include "utils/result.hpp"
@@ -53,10 +54,7 @@ cpp::result<bool, std::string> DownloadService::AddDownloadTask(
 }
 
 cpp::result<uint64_t, std::string> DownloadService::GetFileSize(
-    const std::string& url,
-    const std::optional<
-        std::reference_wrapper<std::unordered_map<std::string, std::string>>>&
-        headers) const noexcept {
+    const std::string& url) const noexcept {
 
   auto curl = curl_easy_init();
   if (!curl) {
@@ -67,10 +65,11 @@ cpp::result<uint64_t, std::string> DownloadService::GetFileSize(
   curl_easy_setopt(curl, CURLOPT_NOBODY, 1L);
   curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
 
+  auto headers = curl_utils::GetHeaders(url);
   if (headers.has_value()) {
     curl_slist* curl_headers = nullptr;
 
-    for (const auto& [key, value] : headers->get()) {
+    for (const auto& [key, value] : headers.value()) {
       auto header = key + ": " + value;
       curl_headers = curl_slist_append(curl_headers, header.c_str());
     }
@@ -151,10 +150,11 @@ cpp::result<bool, std::string> DownloadService::Download(
   }
 
   curl_easy_setopt(curl, CURLOPT_URL, download_item.downloadUrl.c_str());
-  if (download_item.headers.has_value()) {
+  auto headers = curl_utils::GetHeaders(download_item.downloadUrl);
+  if (headers.has_value()) {
     curl_slist* curl_headers = nullptr;
 
-    for (const auto& [key, value] : download_item.headers.value()) {
+    for (const auto& [key, value] : headers.value()) {
       auto header = key + ": " + value;
       curl_headers = curl_slist_append(curl_headers, header.c_str());
     }
@@ -248,10 +248,11 @@ void DownloadService::ProcessTask(DownloadTask& task) {
     });
     downloading_data_map_.insert(std::make_pair(item.id, dl_data_ptr));
 
-    if (item.headers.has_value()) {
+    auto headers = curl_utils::GetHeaders(item.downloadUrl);
+    if (headers.has_value()) {
       curl_slist* curl_headers = nullptr;
 
-      for (const auto& [key, value] : item.headers.value()) {
+      for (const auto& [key, value] : headers.value()) {
         auto header = key + ": " + value;
         curl_headers = curl_slist_append(curl_headers, header.c_str());
       }

engine/services/download_service.h

@@ -73,10 +73,7 @@ class DownloadService {
    * @param url - url to get file size
    */
   cpp::result<uint64_t, std::string> GetFileSize(
-      const std::string& url,
-      const std::optional<
-          std::reference_wrapper<std::unordered_map<std::string, std::string>>>&
-          headers = std::nullopt) const noexcept;
+      const std::string& url) const noexcept;
 
   cpp::result<std::string, std::string> StopTask(const std::string& task_id);
 

engine/services/engine_service.cc

@@ -55,11 +55,6 @@ std::string GetEnginePath(std::string_view e) {
   }
   return kLlamaLibPath;
 };
-
-constexpr const int k200OK = 200;
-constexpr const int k400BadRequest = 400;
-constexpr const int k409Conflict = 409;
-constexpr const int k500InternalServerError = 500;
 }  // namespace
 
 // cpp::result<EngineInfo, std::string> EngineService::GetEngineInfo(

engine/services/model_service.cc

@@ -43,7 +43,12 @@ void ParseGguf(const DownloadItem& ggufDownloadItem,
   }
 
   auto url_obj = url_parser::FromUrlString(ggufDownloadItem.downloadUrl);
-  auto branch = url_obj.pathParams[3];
+  if (url_obj.has_error()) {
+    CTL_WRN("Error parsing url: " << ggufDownloadItem.downloadUrl);
+    return;
+  }
+
+  auto branch = url_obj->pathParams[3];
   CTL_INF("Adding model to modellist with branch: " << branch);
 
   auto rel = file_manager_utils::ToRelativeCortexDataPath(yaml_name);
@@ -66,11 +71,10 @@ cpp::result<DownloadTask, std::string> GetDownloadTask(
     const std::string& modelId, const std::string& branch = "main") {
   url_parser::Url url = {
       .protocol = "https",
-      .host = ModelService::kHuggingFaceHost,
+      .host = kHuggingFaceHost,
       .pathParams = {"api", "models", "cortexso", modelId, "tree", branch}};
 
   httplib::Client cli(url.GetProtocolAndHost());
-  // TODO: namh add header here
   auto result = curl_utils::SimpleGetJson(url.ToFullPath());
   if (result.has_error()) {
     return cpp::fail("Model " + modelId + " not found");
@@ -89,7 +93,7 @@ cpp::result<DownloadTask, std::string> GetDownloadTask(
     }
     url_parser::Url download_url = {
         .protocol = "https",
-        .host = ModelService::kHuggingFaceHost,
+        .host = kHuggingFaceHost,
         .pathParams = {"cortexso", modelId, "resolve", branch, path}};
 
     auto local_path = model_container_path / path;
@@ -221,21 +225,24 @@ std::optional<config::ModelConfig> ModelService::GetDownloadedModel(
 cpp::result<DownloadTask, std::string> ModelService::HandleDownloadUrlAsync(
     const std::string& url, std::optional<std::string> temp_model_id) {
   auto url_obj = url_parser::FromUrlString(url);
+  if (url_obj.has_error()) {
+    return cpp::fail("Invalid url: " + url);
+  }
 
-  if (url_obj.host == kHuggingFaceHost) {
-    if (url_obj.pathParams[2] == "blob") {
-      url_obj.pathParams[2] = "resolve";
+  if (url_obj->host == kHuggingFaceHost) {
+    if (url_obj->pathParams[2] == "blob") {
+      url_obj->pathParams[2] = "resolve";
     }
   }
-  auto author{url_obj.pathParams[0]};
-  auto model_id{url_obj.pathParams[1]};
-  auto file_name{url_obj.pathParams.back()};
+  auto author{url_obj->pathParams[0]};
+  auto model_id{url_obj->pathParams[1]};
+  auto file_name{url_obj->pathParams.back()};
 
   if (author == "cortexso") {
     return DownloadModelFromCortexsoAsync(model_id);
   }
 
-  if (url_obj.pathParams.size() < 5) {
+  if (url_obj->pathParams.size() < 5) {
     return cpp::fail("Invalid url: " + url);
   }
 
@@ -255,7 +262,7 @@ cpp::result<DownloadTask, std::string> ModelService::HandleDownloadUrlAsync(
   }
 
   auto local_path{file_manager_utils::GetModelsContainerPath() /
-                  "huggingface.co" / author / model_id / file_name};
+                  kHuggingFaceHost / author / model_id / file_name};
 
   try {
     std::filesystem::create_directories(local_path.parent_path());
@@ -265,7 +272,7 @@ cpp::result<DownloadTask, std::string> ModelService::HandleDownloadUrlAsync(
     std::filesystem::create_directories(local_path.parent_path());
   }
 
-  auto download_url = url_parser::FromUrl(url_obj);
+  auto download_url = url_parser::FromUrl(url_obj.value());
   // this assume that the model being downloaded is a single gguf file
   auto downloadTask{DownloadTask{.id = model_id,
                                  .type = DownloadType::Model,
@@ -287,22 +294,25 @@ cpp::result<DownloadTask, std::string> ModelService::HandleDownloadUrlAsync(
 cpp::result<std::string, std::string> ModelService::HandleUrl(
     const std::string& url) {
   auto url_obj = url_parser::FromUrlString(url);
+  if (url_obj.has_error()) {
+    return cpp::fail("Invalid url: " + url);
+  }
 
-  if (url_obj.host == kHuggingFaceHost) {
-    if (url_obj.pathParams[2] == "blob") {
-      url_obj.pathParams[2] = "resolve";
+  if (url_obj->host == kHuggingFaceHost) {
+    if (url_obj->pathParams[2] == "blob") {
+      url_obj->pathParams[2] = "resolve";
     }
   }
-  auto author{url_obj.pathParams[0]};
-  auto model_id{url_obj.pathParams[1]};
-  auto file_name{url_obj.pathParams.back()};
+  auto author{url_obj->pathParams[0]};
+  auto model_id{url_obj->pathParams[1]};
+  auto file_name{url_obj->pathParams.back()};
 
   if (author == "cortexso") {
     return DownloadModelFromCortexso(model_id);
   }
 
-  if (url_obj.pathParams.size() < 5) {
-    if (url_obj.pathParams.size() < 2) {
+  if (url_obj->pathParams.size() < 5) {
+    if (url_obj->pathParams.size() < 2) {
       return cpp::fail("Invalid url: " + url);
     }
     return DownloadHuggingFaceGgufModel(author, model_id, std::nullopt);
@@ -320,7 +330,7 @@ cpp::result<std::string, std::string> ModelService::HandleUrl(
   }
 
   auto local_path{file_manager_utils::GetModelsContainerPath() /
-                  "huggingface.co" / author / model_id / file_name};
+                  kHuggingFaceHost / author / model_id / file_name};
 
   try {
     std::filesystem::create_directories(local_path.parent_path());
@@ -330,7 +340,7 @@ cpp::result<std::string, std::string> ModelService::HandleUrl(
     std::filesystem::create_directories(local_path.parent_path());
   }
 
-  auto download_url = url_parser::FromUrl(url_obj);
+  auto download_url = url_parser::FromUrl(url_obj.value());
   // this assume that the model being downloaded is a single gguf file
   auto downloadTask{DownloadTask{.id = model_id,
                                  .type = DownloadType::Model,

engine/services/model_service.h

@@ -8,17 +8,16 @@
 #include "services/inference_service.h"
 
 struct StartParameterOverride {
-std::optional<bool> cache_enabled;
-std::optional<int> ngl;
-std::optional<int> n_parallel;
-std::optional<int> ctx_len;
-std::optional<std::string> custom_prompt_template;
-std::optional<std::string> cache_type;
+  std::optional<bool> cache_enabled;
+  std::optional<int> ngl;
+  std::optional<int> n_parallel;
+  std::optional<int> ctx_len;
+  std::optional<std::string> custom_prompt_template;
+  std::optional<std::string> cache_type;
 };
+
 class ModelService {
  public:
-  constexpr auto static kHuggingFaceHost = "huggingface.co";
-
   explicit ModelService(std::shared_ptr<DownloadService> download_service)
       : download_service_{download_service} {};
 

engine/test/components/test_url_parser.cc

@@ -9,9 +9,9 @@ class UrlParserTestSuite : public ::testing::Test {
 TEST_F(UrlParserTestSuite, TestParseUrlCorrectly) {
   auto url = url_parser::FromUrlString(kValidUrlWithOnlyPaths);
 
-  EXPECT_EQ(url.host, "jan.ai");
-  EXPECT_EQ(url.protocol, "https");
-  EXPECT_EQ(url.pathParams.size(), 2);
+  EXPECT_EQ(url->host, "jan.ai");
+  EXPECT_EQ(url->protocol, "https");
+  EXPECT_EQ(url->pathParams.size(), 2);
 }
 
 TEST_F(UrlParserTestSuite, ConstructUrlCorrectly) {

engine/utils/curl_utils.h

@@ -6,7 +6,10 @@
 #include <yaml-cpp/node/node.h>
 #include <yaml-cpp/node/parse.h>
 #include <string>
+#include "utils/engine_constants.h"
+#include "utils/file_manager_utils.h"
 #include "utils/result.hpp"
+#include "utils/url_parser.h"
 
 namespace curl_utils {
 namespace {
@@ -18,12 +21,10 @@ size_t WriteCallback(void* contents, size_t size, size_t nmemb,
 }
 }  // namespace
 
-inline cpp::result<std::string, std::string> SimpleGet(
-    const std::string& url,
-    const std::optional<std::reference_wrapper<
-        const std::unordered_map<std::string, std::string>>>& headers =
-        std::nullopt) {
+inline std::optional<std::unordered_map<std::string, std::string>> GetHeaders(
+    const std::string& url);
 
+inline cpp::result<std::string, std::string> SimpleGet(const std::string& url) {
   // Initialize libcurl
   curl_global_init(CURL_GLOBAL_DEFAULT);
   auto curl = curl_easy_init();
@@ -32,10 +33,11 @@ inline cpp::result<std::string, std::string> SimpleGet(
     return cpp::fail("Failed to init CURL");
   }
 
+  auto headers = GetHeaders(url);
   if (headers.has_value()) {
     curl_slist* curl_headers = nullptr;
 
-    for (const auto& [key, value] : headers->get()) {
+    for (const auto& [key, value] : headers.value()) {
       auto header = key + ": " + value;
       curl_headers = curl_slist_append(curl_headers, header.c_str());
     }
@@ -62,11 +64,8 @@ inline cpp::result<std::string, std::string> SimpleGet(
 }
 
 inline cpp::result<YAML::Node, std::string> ReadRemoteYaml(
-    const std::string& url,
-    const std::optional<std::reference_wrapper<
-        const std::unordered_map<std::string, std::string>>>& headers =
-        std::nullopt) {
-  auto result = SimpleGet(url, headers);
+    const std::string& url) {
+  auto result = SimpleGet(url);
   if (result.has_error()) {
     return cpp::fail(result.error());
   }
@@ -80,23 +79,48 @@ inline cpp::result<YAML::Node, std::string> ReadRemoteYaml(
 }
 
 inline cpp::result<Json::Value, std::string> SimpleGetJson(
-    const std::string& url,
-    const std::optional<std::reference_wrapper<
-        const std::unordered_map<std::string, std::string>>>& headers =
-        std::nullopt) {
-
-  auto result = SimpleGet(url, headers);
+    const std::string& url) {
+  auto result = SimpleGet(url);
   if (result.has_error()) {
     return cpp::fail(result.error());
   }
+
   Json::Value root;
   Json::Reader reader;
-
   if (!reader.parse(result.value(), root)) {
     return cpp::fail("JSON from " + url +
                      " parsing error: " + reader.getFormattedErrorMessages());
   }
 
   return root;
 }
+
+inline std::optional<std::unordered_map<std::string, std::string>> GetHeaders(
+    const std::string& url) {
+  auto url_obj = url_parser::FromUrlString(url);
+  if (url_obj.has_error()) {
+    return std::nullopt;
+  }
+
+  if (url_obj->host == kHuggingFaceHost) {
+    std::unordered_map<std::string, std::string> headers{};
+    headers["Content-Type"] = "application/json";
+    auto const& token = file_manager_utils::GetCortexConfig().huggingFaceToken;
+    if (!token.empty()) {
+      headers["Authorization"] = "Bearer " + token;
+    }
+    // TODO: namh printout last 6 characters of the token
+
+    return headers;
+  }
+
+  if (url_obj->host == kGitHubHost) {
+    std::unordered_map<std::string, std::string> headers{};
+    headers["Accept"] = "application/vnd.github.v3+json";
+    headers["User-Agent"] = "cortex";
+    return headers;
+  }
+
+  return std::nullopt;
+}
 }  // namespace curl_utils

engine/utils/engine_constants.h

@@ -12,4 +12,8 @@ constexpr const auto kPythonRuntimeRepo = "cortex.python";
 constexpr const auto kLlamaLibPath = "/engines/cortex.llamacpp";
 constexpr const auto kPythonRuntimeLibPath = "/engines/cortex.python";
 constexpr const auto kOnnxLibPath = "/engines/cortex.onnx";
-constexpr const auto kTensorrtLlmPath = "/engines/cortex.tensorrt-llm";
+constexpr const auto kTensorrtLlmPath = "/engines/cortex.tensorrt-llm";
+
+// other constants
+constexpr auto static kHuggingFaceHost = "huggingface.co";
+constexpr auto static kGitHubHost = "api.github.com";