feat: sqlite database implementation (#1336)

* feat: sqlite * chore: unit tests * fix: unit tests * refactor: db * fix: remove mutex * fix: rm file * fix: test * fix: test * refactor: LoadModelList * refactor: more * fix: transaction * fix: format * fix: make alias unique * fix: remove ModelStatus * fix: models
janhq · Sep 27, 2024 · 08b8841 · 08b8841
1 parent 6502797
commit 08b8841
Show file tree

Hide file tree

Showing 22 changed files with 657 additions and 525 deletions.
diff --git a/engine/CMakeLists.txt b/engine/CMakeLists.txt
@@ -77,11 +77,11 @@ find_package(unofficial-minizip CONFIG REQUIRED)
 find_package(LibArchive REQUIRED)
 find_package(tabulate CONFIG REQUIRED)
 find_package(CURL REQUIRED)
+find_package(SQLiteCpp REQUIRED)
 
 add_executable(${TARGET_NAME} main.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/utils/cpuid/cpu_info.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/utils/file_logger.cc
-    ${CMAKE_CURRENT_SOURCE_DIR}/utils/modellist_utils.cc
   )
 
 target_link_libraries(${TARGET_NAME} PRIVATE httplib::httplib)
@@ -93,6 +93,7 @@ target_link_libraries(${TARGET_NAME} PRIVATE tabulate::tabulate)
 target_link_libraries(${TARGET_NAME} PRIVATE CURL::libcurl)
 target_link_libraries(${TARGET_NAME} PRIVATE JsonCpp::JsonCpp Drogon::Drogon OpenSSL::SSL OpenSSL::Crypto yaml-cpp::yaml-cpp
   ${CMAKE_THREAD_LIBS_INIT})
+target_link_libraries(${TARGET_NAME} PRIVATE SQLiteCpp)
 
 # ##############################################################################
 
@@ -114,7 +115,8 @@ aux_source_directory(models MODEL_SRC)
 aux_source_directory(cortex-common CORTEX_COMMON)
 aux_source_directory(config CONFIG_SRC)
 aux_source_directory(commands COMMANDS_SRC)
-
+aux_source_directory(database DB_SRC)
+
 target_include_directories(${TARGET_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} )
 
-target_sources(${TARGET_NAME} PRIVATE ${COMMANDS_SRC} ${CONFIG_SRC} ${CTL_SRC} ${COMMON_SRC} ${SERVICES_SRC})
+target_sources(${TARGET_NAME} PRIVATE ${COMMANDS_SRC} ${CONFIG_SRC} ${CTL_SRC} ${COMMON_SRC} ${SERVICES_SRC} ${DB_SRC})
diff --git a/engine/commands/chat_cmd.cc b/engine/commands/chat_cmd.cc
@@ -2,11 +2,11 @@
 #include "httplib.h"
 
 #include "cortex_upd_cmd.h"
+#include "database/models.h"
 #include "model_status_cmd.h"
 #include "server_start_cmd.h"
 #include "trantor/utils/Logger.h"
 #include "utils/logging_utils.h"
-#include "utils/modellist_utils.h"
 
 namespace commands {
 namespace {
@@ -39,11 +39,15 @@ struct ChunkParser {
 
 void ChatCmd::Exec(const std::string& host, int port,
                    const std::string& model_handle, std::string msg) {
-  modellist_utils::ModelListUtils modellist_handler;
+  cortex::db::Models modellist_handler;
   config::YamlHandler yaml_handler;
   try {
     auto model_entry = modellist_handler.GetModelInfo(model_handle);
-    yaml_handler.ModelConfigFromFile(model_entry.path_to_model_yaml);
+    if (model_entry.has_error()) {
+      CLI_LOG("Error: " + model_entry.error());
+      return;
+    }
+    yaml_handler.ModelConfigFromFile(model_entry.value().path_to_model_yaml);
     auto mc = yaml_handler.GetModelConfig();
     Exec(host, port, mc, std::move(msg));
   } catch (const std::exception& e) {

diff --git a/engine/commands/model_alias_cmd.cc b/engine/commands/model_alias_cmd.cc
@@ -1,11 +1,11 @@
 #include "model_alias_cmd.h"
-#include "utils/modellist_utils.h"
+#include "database/models.h"
 
 namespace commands {
 
 void ModelAliasCmd::Exec(const std::string& model_handle,
                          const std::string& model_alias) {
-  modellist_utils::ModelListUtils modellist_handler;
+  cortex::db::Models modellist_handler;
   try {
     if (modellist_handler.UpdateModelAlias(model_handle, model_alias)) {
       CLI_LOG("Successfully set model alias '" + model_alias +

diff --git a/engine/commands/model_get_cmd.cc b/engine/commands/model_get_cmd.cc
@@ -5,18 +5,22 @@
 #include <vector>
 #include "cmd_info.h"
 #include "config/yaml_config.h"
+#include "database/models.h"
 #include "utils/file_manager_utils.h"
 #include "utils/logging_utils.h"
-#include "utils/modellist_utils.h"
 
 namespace commands {
 
 void ModelGetCmd::Exec(const std::string& model_handle) {
-  modellist_utils::ModelListUtils modellist_handler;
+  cortex::db::Models modellist_handler;
   config::YamlHandler yaml_handler;
   try {
     auto model_entry = modellist_handler.GetModelInfo(model_handle);
-    yaml_handler.ModelConfigFromFile(model_entry.path_to_model_yaml);
+    if (model_entry.has_error()) {
+      CLI_LOG("Error: " + model_entry.error());
+      return;
+    }
+    yaml_handler.ModelConfigFromFile(model_entry.value().path_to_model_yaml);
     auto model_config = yaml_handler.GetModelConfig();
 
     std::cout << model_config.ToString() << std::endl;

diff --git a/engine/commands/model_import_cmd.cc b/engine/commands/model_import_cmd.cc
@@ -3,9 +3,9 @@
 #include <vector>
 #include "config/gguf_parser.h"
 #include "config/yaml_config.h"
+#include "database/models.h"
 #include "utils/file_manager_utils.h"
 #include "utils/logging_utils.h"
-#include "utils/modellist_utils.h"
 
 namespace commands {
 
@@ -16,15 +16,15 @@ ModelImportCmd::ModelImportCmd(std::string model_handle, std::string model_path)
 void ModelImportCmd::Exec() {
   config::GGUFHandler gguf_handler;
   config::YamlHandler yaml_handler;
-  modellist_utils::ModelListUtils modellist_utils_obj;
+  cortex::db::Models modellist_utils_obj;
 
   std::string model_yaml_path = (file_manager_utils::GetModelsContainerPath() /
                                  std::filesystem::path("imported") /
                                  std::filesystem::path(model_handle_ + ".yml"))
                                     .string();
-  modellist_utils::ModelEntry model_entry{
+  cortex::db::ModelEntry model_entry{
       model_handle_,   "local",       "imported",
-      model_yaml_path, model_handle_, modellist_utils::ModelStatus::READY};
+      model_yaml_path, model_handle_};
   try {
     std::filesystem::create_directories(
         std::filesystem::path(model_yaml_path).parent_path());
@@ -34,7 +34,7 @@ void ModelImportCmd::Exec() {
     model_config.model = model_handle_;
     yaml_handler.UpdateModelConfig(model_config);
 
-    if (modellist_utils_obj.AddModelEntry(model_entry)) {
+    if (modellist_utils_obj.AddModelEntry(model_entry).value()) {
       yaml_handler.WriteYamlFile(model_yaml_path);
       CLI_LOG("Model is imported successfully!");
     } else {

diff --git a/engine/commands/model_list_cmd.cc b/engine/commands/model_list_cmd.cc
@@ -3,15 +3,15 @@
 #include <tabulate/table.hpp>
 #include <vector>
 #include "config/yaml_config.h"
+#include "database/models.h"
 #include "utils/file_manager_utils.h"
 #include "utils/logging_utils.h"
-#include "utils/modellist_utils.h"
 
 namespace commands {
 
 void ModelListCmd::Exec() {
   auto models_path = file_manager_utils::GetModelsContainerPath();
-  modellist_utils::ModelListUtils modellist_handler;
+  cortex::db::Models modellist_handler;
   config::YamlHandler yaml_handler;
   tabulate::Table table;
 
@@ -20,24 +20,24 @@ void ModelListCmd::Exec() {
   int count = 0;
   // Iterate through directory
 
-  try {
-    auto list_entry = modellist_handler.LoadModelList();
-    for (const auto& model_entry : list_entry) {
-      // auto model_entry = modellist_handler.GetModelInfo(model_handle);
-      try {
-        count += 1;
-        yaml_handler.ModelConfigFromFile(model_entry.path_to_model_yaml);
-        auto model_config = yaml_handler.GetModelConfig();
-        table.add_row({std::to_string(count), model_entry.model_id,
-                       model_entry.model_alias, model_config.engine,
-                       model_config.version});
-        yaml_handler.Reset();
-      } catch (const std::exception& e) {
-        CTL_ERR("Fail to get list model information: " + std::string(e.what()));
-      }
+  auto list_entry = modellist_handler.LoadModelList();
+  if (list_entry.has_error()) {
+    CTL_ERR("Fail to get list model information: " << list_entry.error());
+    return;
+  }
+  for (const auto& model_entry : list_entry.value()) {
+    // auto model_entry = modellist_handler.GetModelInfo(model_handle);
+    try {
+      count += 1;
+      yaml_handler.ModelConfigFromFile(model_entry.path_to_model_yaml);
+      auto model_config = yaml_handler.GetModelConfig();
+      table.add_row({std::to_string(count), model_entry.model_id,
+                     model_entry.model_alias, model_config.engine,
+                     model_config.version});
+      yaml_handler.Reset();
+    } catch (const std::exception& e) {
+      CTL_ERR("Fail to get list model information: " + std::string(e.what()));
     }
-  } catch (const std::exception& e) {
-    CTL_ERR("Fail to get list model information: " + std::string(e.what()));
   }
 
   for (int i = 0; i < 5; i++) {

diff --git a/engine/commands/model_start_cmd.cc b/engine/commands/model_start_cmd.cc
@@ -1,23 +1,27 @@
 #include "model_start_cmd.h"
 #include "cortex_upd_cmd.h"
+#include "database/models.h"
 #include "httplib.h"
 #include "model_status_cmd.h"
 #include "nlohmann/json.hpp"
 #include "server_start_cmd.h"
 #include "trantor/utils/Logger.h"
 #include "utils/file_manager_utils.h"
 #include "utils/logging_utils.h"
-#include "utils/modellist_utils.h"
 
 namespace commands {
 bool ModelStartCmd::Exec(const std::string& host, int port,
                          const std::string& model_handle) {
 
-  modellist_utils::ModelListUtils modellist_handler;
+  cortex::db::Models modellist_handler;
   config::YamlHandler yaml_handler;
   try {
     auto model_entry = modellist_handler.GetModelInfo(model_handle);
-    yaml_handler.ModelConfigFromFile(model_entry.path_to_model_yaml);
+    if (model_entry.has_error()) {
+      CLI_LOG("Error: " + model_entry.error());
+      return false;
+    }
+    yaml_handler.ModelConfigFromFile(model_entry.value().path_to_model_yaml);
     auto mc = yaml_handler.GetModelConfig();
     return Exec(host, port, mc);
   } catch (const std::exception& e) {

diff --git a/engine/commands/model_status_cmd.cc b/engine/commands/model_status_cmd.cc
@@ -1,18 +1,22 @@
 #include "model_status_cmd.h"
 #include "config/yaml_config.h"
+#include "database/models.h"
 #include "httplib.h"
 #include "nlohmann/json.hpp"
 #include "utils/logging_utils.h"
-#include "utils/modellist_utils.h"
 
 namespace commands {
 bool ModelStatusCmd::IsLoaded(const std::string& host, int port,
                               const std::string& model_handle) {
-  modellist_utils::ModelListUtils modellist_handler;
+  cortex::db::Models modellist_handler;
   config::YamlHandler yaml_handler;
   try {
     auto model_entry = modellist_handler.GetModelInfo(model_handle);
-    yaml_handler.ModelConfigFromFile(model_entry.path_to_model_yaml);
+    if (model_entry.has_error()) {
+      CLI_LOG("Error: " + model_entry.error());
+      return false;
+    }
+    yaml_handler.ModelConfigFromFile(model_entry.value().path_to_model_yaml);
     auto mc = yaml_handler.GetModelConfig();
     return IsLoaded(host, port, mc);
   } catch (const std::exception& e) {

diff --git a/engine/commands/model_upd_cmd.cc b/engine/commands/model_upd_cmd.cc
@@ -11,7 +11,11 @@ void ModelUpdCmd::Exec(
     const std::unordered_map<std::string, std::string>& options) {
   try {
     auto model_entry = model_list_utils_.GetModelInfo(model_handle_);
-    yaml_handler_.ModelConfigFromFile(model_entry.path_to_model_yaml);
+    if (model_entry.has_error()) {
+      CLI_LOG("Error: " + model_entry.error());
+      return;
+    }
+    yaml_handler_.ModelConfigFromFile(model_entry.value().path_to_model_yaml);
     model_config_ = yaml_handler_.GetModelConfig();
 
     for (const auto& [key, value] : options) {
@@ -21,7 +25,7 @@ void ModelUpdCmd::Exec(
     }
 
     yaml_handler_.UpdateModelConfig(model_config_);
-    yaml_handler_.WriteYamlFile(model_entry.path_to_model_yaml);
+    yaml_handler_.WriteYamlFile(model_entry.value().path_to_model_yaml);
     CLI_LOG("Successfully updated model ID '" + model_handle_ + "'!");
   } catch (const std::exception& e) {
     CLI_LOG("Failed to update model with model ID '" + model_handle_ +

diff --git a/engine/commands/model_upd_cmd.h b/engine/commands/model_upd_cmd.h
@@ -5,8 +5,8 @@
 #include <unordered_map>
 #include <vector>
 #include "config/model_config.h"
-#include "utils/modellist_utils.h"
 #include "config/yaml_config.h"
+#include "database/models.h"
 namespace commands {
 class ModelUpdCmd {
  public:
@@ -17,7 +17,7 @@ class ModelUpdCmd {
   std::string model_handle_;
   config::ModelConfig model_config_;
   config::YamlHandler yaml_handler_;
-  modellist_utils::ModelListUtils model_list_utils_;
+  cortex::db::Models model_list_utils_;
 
   void UpdateConfig(const std::string& key, const std::string& value);
   void UpdateVectorField(const std::string& key, const std::string& value);

diff --git a/engine/commands/run_cmd.cc b/engine/commands/run_cmd.cc
@@ -1,18 +1,18 @@
 #include "run_cmd.h"
 #include "chat_cmd.h"
 #include "config/yaml_config.h"
+#include "database/models.h"
 #include "model_start_cmd.h"
 #include "model_status_cmd.h"
 #include "server_start_cmd.h"
 #include "utils/logging_utils.h"
-#include "utils/modellist_utils.h"
 
 namespace commands {
 
 void RunCmd::Exec() {
   std::optional<std::string> model_id = model_handle_;
 
-  modellist_utils::ModelListUtils modellist_handler;
+  cortex::db::Models modellist_handler;
   config::YamlHandler yaml_handler;
   auto address = host_ + ":" + std::to_string(port_);
 
@@ -31,7 +31,11 @@ void RunCmd::Exec() {
 
   try {
     auto model_entry = modellist_handler.GetModelInfo(*model_id);
-    yaml_handler.ModelConfigFromFile(model_entry.path_to_model_yaml);
+    if (model_entry.has_error()) {
+      CLI_LOG("Error: " + model_entry.error());
+      return;
+    }
+    yaml_handler.ModelConfigFromFile(model_entry.value().path_to_model_yaml);
     auto mc = yaml_handler.GetModelConfig();
 
     // Check if engine existed. If not, download it