From 98e12ae9a03a513054abc6770a7c4f29c00c3f8a Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Wed, 25 Sep 2024 15:01:01 +0700 Subject: [PATCH 01/15] feat: sqlite --- engine/CMakeLists.txt | 2 + engine/utils/modellist_utils.cc | 247 +++++++++++++++++--------------- engine/utils/modellist_utils.h | 12 +- engine/vcpkg.json | 3 +- 4 files changed, 139 insertions(+), 125 deletions(-) diff --git a/engine/CMakeLists.txt b/engine/CMakeLists.txt index 76cdcf303..454dc8f2c 100644 --- a/engine/CMakeLists.txt +++ b/engine/CMakeLists.txt @@ -77,6 +77,7 @@ 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 @@ -93,6 +94,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) # ############################################################################## diff --git a/engine/utils/modellist_utils.cc b/engine/utils/modellist_utils.cc index d577519f3..84fff1d91 100644 --- a/engine/utils/modellist_utils.cc +++ b/engine/utils/modellist_utils.cc @@ -12,75 +12,84 @@ const std::string ModelListUtils::kModelListPath = (file_manager_utils::GetModelsContainerPath() / std::filesystem::path("model.list")) .string(); +namespace { +const std::string kDefaultDbPath = + file_manager_utils::GetCortexDataPath().string() + "/cortex.db"; +} +ModelListUtils::ModelListUtils() + : db_(kDefaultDbPath, SQLite::OPEN_READWRITE | SQLite::OPEN_CREATE) { + db_.exec( + "CREATE TABLE IF NOT EXISTS models (" + "model_id TEXT PRIMARY KEY," + "author_repo_id TEXT," + "branch_name TEXT," + "path_to_model_yaml TEXT," + "model_alias TEXT," + "status TEXT);"); +} +ModelListUtils::~ModelListUtils() {} std::vector ModelListUtils::LoadModelList() const { std::vector entries; - std::filesystem::path file_path(kModelListPath); - - // Check if the file exists, if not, create it - if (!std::filesystem::exists(file_path)) { - std::ofstream create_file(kModelListPath); - if (!create_file) { - throw std::runtime_error("Unable to create model.list file: " + - kModelListPath); - } - create_file.close(); - return entries; // Return empty vector for newly created file - } - - std::ifstream file(kModelListPath); - if (!file.is_open()) { - throw std::runtime_error("Unable to open model.list file: " + - kModelListPath); - } + SQLite::Statement query( + db_, + "SELECT model_id, author_repo_id, branch_name, " + "path_to_model_yaml, model_alias, status FROM models"); - std::string line; - while (std::getline(file, line)) { - std::istringstream iss(line); + while (query.executeStep()) { ModelEntry entry; - std::string status_str; - if (!(iss >> entry.model_id >> entry.author_repo_id >> entry.branch_name >> - entry.path_to_model_yaml >> entry.model_alias >> status_str)) { - LOG_WARN << "Invalid entry in model.list: " << line; - } else { - entry.status = - (status_str == "RUNNING") ? ModelStatus::RUNNING : ModelStatus::READY; - entries.push_back(entry); - } + entry.model_id = query.getColumn(0).getString(); + entry.author_repo_id = query.getColumn(1).getString(); + entry.branch_name = query.getColumn(2).getString(); + entry.path_to_model_yaml = query.getColumn(3).getString(); + entry.model_alias = query.getColumn(4).getString(); + std::string status_str = query.getColumn(5).getString(); + entry.status = + (status_str == "RUNNING") ? ModelStatus::RUNNING : ModelStatus::READY; + entries.push_back(entry); } return entries; } -bool ModelListUtils::IsUnique(const std::vector& entries, - const std::string& model_id, +bool ModelListUtils::IsUnique(const std::string& model_id, const std::string& model_alias) const { - return std::none_of( - entries.begin(), entries.end(), [&](const ModelEntry& entry) { - return entry.model_id == model_id || entry.model_alias == model_id || - entry.model_id == model_alias || - entry.model_alias == model_alias; - }); + SQLite::Statement query(db_, + "SELECT COUNT(*) FROM models WHERE model_id = ? OR " + "model_id = ? OR model_alias = ? OR model_alias = ?"); + query.bind(1, model_id); + query.bind(2, model_alias); + query.bind(3, model_id); + query.bind(4, model_alias); + if (query.executeStep()) { + return query.getColumn(0).getInt() == 0; + } + return false; } void ModelListUtils::SaveModelList( const std::vector& entries) const { std::ofstream file(kModelListPath); - if (!file.is_open()) { - throw std::runtime_error("Unable to open model.list file for writing: " + - kModelListPath); - } - + db_.exec("BEGIN TRANSACTION;"); for (const auto& entry : entries) { - file << entry.model_id << " " << entry.author_repo_id << " " - << entry.branch_name << " " << entry.path_to_model_yaml << " " - << entry.model_alias << " " - << (entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY") - << std::endl; + SQLite::Statement insert( + db_, + "INSERT OR REPLACE INTO models (model_id, author_repo_id, " + "branch_name, path_to_model_yaml, model_alias, status) VALUES (?, ?, " + "?, ?, ?, ?)"); + insert.bind(1, entry.model_id); + insert.bind(2, entry.author_repo_id); + insert.bind(3, entry.branch_name); + insert.bind(4, entry.path_to_model_yaml); + insert.bind(5, entry.model_alias); + insert.bind(6, + (entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY")); + insert.exec(); } + db_.exec("COMMIT;"); } std::string ModelListUtils::GenerateShortenedAlias( - const std::string& model_id, const std::vector& entries) const { + const std::string& model_id) const { std::vector parts; std::istringstream iss(model_id); std::string part; @@ -123,7 +132,7 @@ std::string ModelListUtils::GenerateShortenedAlias( // Find the first unique candidate for (const auto& candidate : candidates) { - if (IsUnique(entries, model_id, candidate)) { + if (IsUnique(model_id, candidate)) { return candidate; } } @@ -132,7 +141,7 @@ std::string ModelListUtils::GenerateShortenedAlias( std::string base_candidate = candidates.back(); int suffix = 1; std::string unique_candidate = base_candidate; - while (!IsUnique(entries, model_id, unique_candidate)) { + while (!IsUnique(model_id, unique_candidate)) { unique_candidate = base_candidate + "-" + std::to_string(suffix++); } @@ -141,14 +150,24 @@ std::string ModelListUtils::GenerateShortenedAlias( ModelEntry ModelListUtils::GetModelInfo(const std::string& identifier) const { std::lock_guard lock(mutex_); - auto entries = LoadModelList(); - auto it = std::find_if( - entries.begin(), entries.end(), [&identifier](const ModelEntry& entry) { - return entry.model_id == identifier || entry.model_alias == identifier; - }); - - if (it != entries.end()) { - return *it; + SQLite::Statement query(db_, + "SELECT model_id, author_repo_id, branch_name, " + "path_to_model_yaml, model_alias, status FROM models " + "WHERE model_id = ? OR model_alias = ?"); + + query.bind(1, identifier); + query.bind(2, identifier); + if (query.executeStep()) { + ModelEntry entry; + entry.model_id = query.getColumn(0).getString(); + entry.author_repo_id = query.getColumn(1).getString(); + entry.branch_name = query.getColumn(2).getString(); + entry.path_to_model_yaml = query.getColumn(3).getString(); + entry.model_alias = query.getColumn(4).getString(); + std::string status_str = query.getColumn(5).getString(); + entry.status = + (status_str == "RUNNING") ? ModelStatus::RUNNING : ModelStatus::READY; + return entry; } else { throw std::runtime_error("Model not found: " + identifier); } @@ -166,16 +185,27 @@ void ModelListUtils::PrintModelInfo(const ModelEntry& entry) const { bool ModelListUtils::AddModelEntry(ModelEntry new_entry, bool use_short_alias) { std::lock_guard lock(mutex_); - auto entries = LoadModelList(); - if (IsUnique(entries, new_entry.model_id, new_entry.model_alias)) { + if (IsUnique(new_entry.model_id, new_entry.model_alias)) { if (use_short_alias) { - new_entry.model_alias = - GenerateShortenedAlias(new_entry.model_id, entries); + new_entry.model_alias = GenerateShortenedAlias(new_entry.model_id); } new_entry.status = ModelStatus::READY; // Set default status to READY - entries.push_back(std::move(new_entry)); - SaveModelList(entries); + + SQLite::Statement insert( + db_, + "INSERT INTO models (model_id, author_repo_id, " + "branch_name, path_to_model_yaml, model_alias, status) VALUES (?, ?, " + "?, ?, ?, ?)"); + insert.bind(1, new_entry.model_id); + insert.bind(2, new_entry.author_repo_id); + insert.bind(3, new_entry.branch_name); + insert.bind(4, new_entry.path_to_model_yaml); + insert.bind(5, new_entry.model_alias); + insert.bind( + 6, (new_entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY")); + insert.exec(); + return true; } return false; // Entry not added due to non-uniqueness @@ -184,73 +214,52 @@ bool ModelListUtils::AddModelEntry(ModelEntry new_entry, bool use_short_alias) { bool ModelListUtils::UpdateModelEntry(const std::string& identifier, const ModelEntry& updated_entry) { std::lock_guard lock(mutex_); - auto entries = LoadModelList(); - auto it = std::find_if( - entries.begin(), entries.end(), [&identifier](const ModelEntry& entry) { - return entry.model_id == identifier || entry.model_alias == identifier; - }); - - if (it != entries.end()) { - *it = updated_entry; - SaveModelList(entries); - return true; - } - return false; // Entry not found + SQLite::Statement upd(db_, + "UPDATE models " + "SET author_repo_id = ?, branch_name = ?, " + "path_to_model_yaml = ?, status = ? " + "WHERE model_id = ? OR model_alias = ?"); + upd.bind(1, updated_entry.author_repo_id); + upd.bind(2, updated_entry.branch_name); + upd.bind(3, updated_entry.path_to_model_yaml); + upd.bind( + 4, (updated_entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY")); + upd.bind(5, identifier); + upd.bind(6, identifier); + return upd.exec() == 1; } bool ModelListUtils::UpdateModelAlias(const std::string& model_id, const std::string& new_model_alias) { std::lock_guard lock(mutex_); - auto entries = LoadModelList(); - auto it = std::find_if( - entries.begin(), entries.end(), [&model_id](const ModelEntry& entry) { - return entry.model_id == model_id || entry.model_alias == model_id; - }); - bool check_alias_unique = std::none_of( - entries.begin(), entries.end(), [&](const ModelEntry& entry) { - return (entry.model_id == new_model_alias && - entry.model_id != model_id) || - entry.model_alias == new_model_alias; - }); - if (it != entries.end() && check_alias_unique) { - - (*it).model_alias = new_model_alias; - SaveModelList(entries); - return true; - } - return false; // Entry not found + SQLite::Statement upd(db_, + "UPDATE models " + "SET model_alias = ? " + "WHERE model_id = ? OR model_alias = ?"); + upd.bind(1, new_model_alias); + upd.bind(2, model_id); + upd.bind(3, model_id); + return upd.exec() == 1; } bool ModelListUtils::DeleteModelEntry(const std::string& identifier) { std::lock_guard lock(mutex_); - auto entries = LoadModelList(); - auto it = std::find_if(entries.begin(), entries.end(), - [&identifier](const ModelEntry& entry) { - return (entry.model_id == identifier || - entry.model_alias == identifier) && - entry.status == ModelStatus::READY; - }); - - if (it != entries.end()) { - entries.erase(it); - SaveModelList(entries); - return true; - } - return false; // Entry not found or not in READY state + SQLite::Statement del( + db_, "DELETE from models WHERE model_id = ? OR model_alias = ?"); + del.bind(1, identifier); + del.bind(2, identifier); + return del.exec() == 1; } bool ModelListUtils::HasModel(const std::string& identifier) const { std::lock_guard lock(mutex_); - auto entries = LoadModelList(); - auto it = std::find_if( - entries.begin(), entries.end(), [&identifier](const ModelEntry& entry) { - return entry.model_id == identifier || entry.model_alias == identifier; - }); - - if (it != entries.end()) { - return true; - } else { - return false; + SQLite::Statement query( + db_, "SELECT COUNT(*) FROM models WHERE model_id = ? OR model_alias = ?"); + query.bind(1, identifier); + query.bind(2, identifier); + if (query.executeStep()) { + return query.getColumn(0).getInt() > 0; } + return false; } } // namespace modellist_utils diff --git a/engine/utils/modellist_utils.h b/engine/utils/modellist_utils.h index 113591f25..610f07cb4 100644 --- a/engine/utils/modellist_utils.h +++ b/engine/utils/modellist_utils.h @@ -5,6 +5,8 @@ #include #include +#include "sqlitecpp/SQLiteCpp.h" + namespace modellist_utils { enum class ModelStatus { READY, RUNNING }; @@ -21,20 +23,20 @@ struct ModelEntry { class ModelListUtils { private: + mutable SQLite::Database db_; mutable std::mutex mutex_; // For thread safety - bool IsUnique(const std::vector& entries, - const std::string& model_id, + bool IsUnique(const std::string& model_id, const std::string& model_alias) const; void SaveModelList(const std::vector& entries) const; public: static const std::string kModelListPath; std::vector LoadModelList() const; - ModelListUtils() = default; + ModelListUtils(); + ~ModelListUtils(); std::string GenerateShortenedAlias( - const std::string& model_id, - const std::vector& entries) const; + const std::string& model_id) const; ModelEntry GetModelInfo(const std::string& identifier) const; void PrintModelInfo(const ModelEntry& entry) const; bool AddModelEntry(ModelEntry new_entry, bool use_short_alias = false); diff --git a/engine/vcpkg.json b/engine/vcpkg.json index 40abc186e..25b5f3de4 100644 --- a/engine/vcpkg.json +++ b/engine/vcpkg.json @@ -15,6 +15,7 @@ "nlohmann-json", "yaml-cpp", "libarchive", - "tabulate" + "tabulate", + "sqlitecpp" ] } From 2275b0cb81707197d6107081bf50f7f6c8fe52e1 Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 06:58:47 +0700 Subject: [PATCH 02/15] chore: unit tests --- engine/test/components/CMakeLists.txt | 2 + .../test/components/test_modellist_utils.cc | 108 ++++++++++-------- engine/utils/modellist_utils.cc | 33 ++++-- engine/utils/modellist_utils.h | 3 +- 4 files changed, 90 insertions(+), 56 deletions(-) diff --git a/engine/test/components/CMakeLists.txt b/engine/test/components/CMakeLists.txt index 13bb9c526..ec25787cb 100644 --- a/engine/test/components/CMakeLists.txt +++ b/engine/test/components/CMakeLists.txt @@ -20,6 +20,7 @@ find_package(httplib CONFIG REQUIRED) find_package(unofficial-minizip CONFIG REQUIRED) find_package(LibArchive REQUIRED) find_package(CURL REQUIRED) +find_package(SQLiteCpp REQUIRED) target_link_libraries(${PROJECT_NAME} PRIVATE Drogon::Drogon GTest::gtest GTest::gtest_main yaml-cpp::yaml-cpp ${CMAKE_THREAD_LIBS_INIT}) @@ -28,6 +29,7 @@ target_link_libraries(${PROJECT_NAME} PRIVATE httplib::httplib) target_link_libraries(${PROJECT_NAME} PRIVATE unofficial::minizip::minizip) target_link_libraries(${PROJECT_NAME} PRIVATE LibArchive::LibArchive) target_link_libraries(${PROJECT_NAME} PRIVATE CURL::libcurl) +target_link_libraries(${PROJECT_NAME} PRIVATE SQLiteCpp) target_include_directories(${PROJECT_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../) add_test(NAME ${PROJECT_NAME} diff --git a/engine/test/components/test_modellist_utils.cc b/engine/test/components/test_modellist_utils.cc index d1dbf91e3..3ad0dded9 100644 --- a/engine/test/components/test_modellist_utils.cc +++ b/engine/test/components/test_modellist_utils.cc @@ -1,38 +1,48 @@ #include #include #include "gtest/gtest.h" -#include "utils/modellist_utils.h" #include "utils/file_manager_utils.h" +#include "utils/modellist_utils.h" + +namespace { +constexpr const auto kTestDb = "./test.db"; +} class ModelListUtilsTestSuite : public ::testing::Test { + public: + ModelListUtilsTestSuite() : model_list_(kTestDb) {} + protected: modellist_utils::ModelListUtils model_list_; const modellist_utils::ModelEntry kTestModel{ - "test_model_id", "test_author", - "main", "/path/to/model.yaml", - "test_alias", modellist_utils::ModelStatus::READY}; -}; - void SetUp() { - // Create a temporary directory for tests - file_manager_utils::CreateConfigFileIfNotExist(); + kTestModel.model_id, + "test_author", + "main", + "/path/to/model.yaml", + "test_alias", + modellist_utils::ModelStatus::READY}; + + void SetUp() { } void TearDown() { // Clean up the temporary directory - std::remove((file_manager_utils::GetModelsContainerPath() / "model.list").string().c_str()); + std::remove(kTestDb); } +}; + TEST_F(ModelListUtilsTestSuite, TestAddModelEntry) { EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); - auto retrieved_model = model_list_.GetModelInfo("test_model_id"); + auto retrieved_model = model_list_.GetModelInfo(kTestModel.model_id); EXPECT_EQ(retrieved_model.model_id, kTestModel.model_id); EXPECT_EQ(retrieved_model.author_repo_id, kTestModel.author_repo_id); } TEST_F(ModelListUtilsTestSuite, TestGetModelInfo) { - model_list_.AddModelEntry(kTestModel); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); - auto model_by_id = model_list_.GetModelInfo("test_model_id"); + auto model_by_id = model_list_.GetModelInfo(kTestModel.model_id); EXPECT_EQ(model_by_id.model_id, kTestModel.model_id); auto model_by_alias = model_list_.GetModelInfo("test_alias"); @@ -43,92 +53,96 @@ TEST_F(ModelListUtilsTestSuite, TestGetModelInfo) { } TEST_F(ModelListUtilsTestSuite, TestUpdateModelEntry) { - model_list_.AddModelEntry(kTestModel); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); modellist_utils::ModelEntry updated_model = kTestModel; updated_model.status = modellist_utils::ModelStatus::RUNNING; - EXPECT_TRUE(model_list_.UpdateModelEntry("test_model_id", updated_model)); + EXPECT_TRUE(model_list_.UpdateModelEntry(kTestModel.model_id, updated_model)); - auto retrieved_model = model_list_.GetModelInfo("test_model_id"); + auto retrieved_model = model_list_.GetModelInfo(kTestModel.model_id); EXPECT_EQ(retrieved_model.status, modellist_utils::ModelStatus::RUNNING); updated_model.status = modellist_utils::ModelStatus::READY; - model_list_.UpdateModelEntry("test_model_id", updated_model); + EXPECT_TRUE(model_list_.UpdateModelEntry(kTestModel.model_id, updated_model)); } TEST_F(ModelListUtilsTestSuite, TestDeleteModelEntry) { - model_list_.AddModelEntry(kTestModel); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); - EXPECT_TRUE(model_list_.DeleteModelEntry("test_model_id")); - EXPECT_THROW(model_list_.GetModelInfo("test_model_id"), std::runtime_error); + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); + EXPECT_THROW(model_list_.GetModelInfo(kTestModel.model_id), + std::runtime_error); } TEST_F(ModelListUtilsTestSuite, TestGenerateShortenedAlias) { + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); auto alias = model_list_.GenerateShortenedAlias( - "huggingface.co/bartowski/llama3.1-7b-gguf/Model_ID_Xxx.gguf", {}); + "huggingface.co/bartowski/llama3.1-7b-gguf/Model_ID_Xxx.gguf"); EXPECT_EQ(alias, "model_id_xxx"); + EXPECT_TRUE(model_list_.UpdateModelAlias(kTestModel.model_id, alias)); // Test with existing entries to force longer alias - modellist_utils::ModelEntry existing_model = kTestModel; - existing_model.model_alias = "model_id_xxx"; - std::vector existing_entries = {existing_model}; - alias = model_list_.GenerateShortenedAlias( - "huggingface.co/bartowski/llama3.1-7b-gguf/Model_ID_Xxx.gguf", - existing_entries); + "huggingface.co/bartowski/llama3.1-7b-gguf/Model_ID_Xxx.gguf"); EXPECT_EQ(alias, "llama3.1-7b-gguf:model_id_xxx"); } TEST_F(ModelListUtilsTestSuite, TestPersistence) { - model_list_.AddModelEntry(kTestModel); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); // Create a new ModelListUtils instance to test if it loads from file - modellist_utils::ModelListUtils new_model_list; - auto retrieved_model = new_model_list.GetModelInfo("test_model_id"); + modellist_utils::ModelListUtils new_model_list(kTestDb); + auto retrieved_model = new_model_list.GetModelInfo(kTestModel.model_id); EXPECT_EQ(retrieved_model.model_id, kTestModel.model_id); EXPECT_EQ(retrieved_model.author_repo_id, kTestModel.author_repo_id); - model_list_.DeleteModelEntry("test_model_id"); + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } TEST_F(ModelListUtilsTestSuite, TestUpdateModelAlias) { + constexpr const auto kNewTestAlias = "new_test_alias"; + constexpr const auto kNonExistentModel = "non_existent_model"; + constexpr const auto kAnotherAlias = "another_alias"; + constexpr const auto kFinalTestAlias = "final_test_alias"; + constexpr const auto kAnotherModelId = "another_model_id"; // Add the test model ASSERT_TRUE(model_list_.AddModelEntry(kTestModel)); // Test successful update - EXPECT_TRUE(model_list_.UpdateModelAlias("test_model_id", "new_test_alias")); - auto updated_model = model_list_.GetModelInfo("new_test_alias"); - EXPECT_EQ(updated_model.model_alias, "new_test_alias"); - EXPECT_EQ(updated_model.model_id, "test_model_id"); + EXPECT_TRUE(model_list_.UpdateModelAlias(kTestModel.model_id, kNewTestAlias)); + auto updated_model = model_list_.GetModelInfo(kNewTestAlias); + EXPECT_EQ(updated_model.model_alias, kNewTestAlias); + EXPECT_EQ(updated_model.model_id, kTestModel.model_id); // Test update with non-existent model - EXPECT_FALSE(model_list_.UpdateModelAlias("non_existent_model", "another_alias")); + EXPECT_FALSE(model_list_.UpdateModelAlias(kNonExistentModel, kAnotherAlias)); // Test update with non-unique alias modellist_utils::ModelEntry another_model = kTestModel; - another_model.model_id = "another_model_id"; - another_model.model_alias = "another_alias"; + another_model.model_id = kAnotherModelId; + another_model.model_alias = kAnotherAlias; ASSERT_TRUE(model_list_.AddModelEntry(another_model)); - EXPECT_FALSE(model_list_.UpdateModelAlias("test_model_id", "another_alias")); + EXPECT_FALSE( + model_list_.UpdateModelAlias(kTestModel.model_id, kAnotherAlias)); // Test update using model alias instead of model ID - EXPECT_TRUE(model_list_.UpdateModelAlias("new_test_alias", "final_test_alias")); - updated_model = model_list_.GetModelInfo("final_test_alias"); - EXPECT_EQ(updated_model.model_alias, "final_test_alias"); - EXPECT_EQ(updated_model.model_id, "test_model_id"); + EXPECT_TRUE(model_list_.UpdateModelAlias(kNewTestAlias, kFinalTestAlias)); + updated_model = model_list_.GetModelInfo(kFinalTestAlias); + EXPECT_EQ(updated_model.model_alias, kFinalTestAlias); + EXPECT_EQ(updated_model.model_id, kTestModel.model_id); // Clean up - model_list_.DeleteModelEntry("test_model_id"); - model_list_.DeleteModelEntry("another_model_id"); + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); + EXPECT_TRUE(model_list_.DeleteModelEntry(kAnotherModelId)); } TEST_F(ModelListUtilsTestSuite, TestHasModel) { - model_list_.AddModelEntry(kTestModel); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); - EXPECT_TRUE(model_list_.HasModel("test_model_id")); + EXPECT_TRUE(model_list_.HasModel(kTestModel.model_id)); EXPECT_TRUE(model_list_.HasModel("test_alias")); EXPECT_FALSE(model_list_.HasModel("non_existent_model")); // Clean up - model_list_.DeleteModelEntry("test_model_id"); + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } \ No newline at end of file diff --git a/engine/utils/modellist_utils.cc b/engine/utils/modellist_utils.cc index 84fff1d91..24cc80559 100644 --- a/engine/utils/modellist_utils.cc +++ b/engine/utils/modellist_utils.cc @@ -27,6 +27,19 @@ ModelListUtils::ModelListUtils() "model_alias TEXT," "status TEXT);"); } + +ModelListUtils::ModelListUtils(const std::string& db_path) + : db_(db_path, SQLite::OPEN_READWRITE | SQLite::OPEN_CREATE) { + db_.exec( + "CREATE TABLE IF NOT EXISTS models (" + "model_id TEXT PRIMARY KEY," + "author_repo_id TEXT," + "branch_name TEXT," + "path_to_model_yaml TEXT," + "model_alias TEXT," + "status TEXT);"); +} + ModelListUtils::~ModelListUtils() {} std::vector ModelListUtils::LoadModelList() const { @@ -232,14 +245,18 @@ bool ModelListUtils::UpdateModelEntry(const std::string& identifier, bool ModelListUtils::UpdateModelAlias(const std::string& model_id, const std::string& new_model_alias) { std::lock_guard lock(mutex_); - SQLite::Statement upd(db_, - "UPDATE models " - "SET model_alias = ? " - "WHERE model_id = ? OR model_alias = ?"); - upd.bind(1, new_model_alias); - upd.bind(2, model_id); - upd.bind(3, model_id); - return upd.exec() == 1; + // Check new_model_alias is unique + if (IsUnique(new_model_alias, new_model_alias)) { + SQLite::Statement upd(db_, + "UPDATE models " + "SET model_alias = ? " + "WHERE model_id = ? OR model_alias = ?"); + upd.bind(1, new_model_alias); + upd.bind(2, model_id); + upd.bind(3, model_id); + return upd.exec() == 1; + } + return false; } bool ModelListUtils::DeleteModelEntry(const std::string& identifier) { diff --git a/engine/utils/modellist_utils.h b/engine/utils/modellist_utils.h index 610f07cb4..be849a891 100644 --- a/engine/utils/modellist_utils.h +++ b/engine/utils/modellist_utils.h @@ -5,7 +5,7 @@ #include #include -#include "sqlitecpp/SQLiteCpp.h" +#include "SQLiteCpp/SQLiteCpp.h" namespace modellist_utils { @@ -34,6 +34,7 @@ class ModelListUtils { static const std::string kModelListPath; std::vector LoadModelList() const; ModelListUtils(); + ModelListUtils(const std::string& db_path); ~ModelListUtils(); std::string GenerateShortenedAlias( const std::string& model_id) const; From a4bf17b9a80d4e4551870cbd1b63e569d605c78d Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 08:06:48 +0700 Subject: [PATCH 03/15] fix: unit tests --- .../test/components/test_modellist_utils.cc | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/engine/test/components/test_modellist_utils.cc b/engine/test/components/test_modellist_utils.cc index 3ad0dded9..a3f975ee0 100644 --- a/engine/test/components/test_modellist_utils.cc +++ b/engine/test/components/test_modellist_utils.cc @@ -21,14 +21,6 @@ class ModelListUtilsTestSuite : public ::testing::Test { "/path/to/model.yaml", "test_alias", modellist_utils::ModelStatus::READY}; - - void SetUp() { - } - - void TearDown() { - // Clean up the temporary directory - std::remove(kTestDb); - } }; TEST_F(ModelListUtilsTestSuite, TestAddModelEntry) { @@ -37,6 +29,9 @@ TEST_F(ModelListUtilsTestSuite, TestAddModelEntry) { auto retrieved_model = model_list_.GetModelInfo(kTestModel.model_id); EXPECT_EQ(retrieved_model.model_id, kTestModel.model_id); EXPECT_EQ(retrieved_model.author_repo_id, kTestModel.author_repo_id); + + // Clean up + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } TEST_F(ModelListUtilsTestSuite, TestGetModelInfo) { @@ -50,6 +45,9 @@ TEST_F(ModelListUtilsTestSuite, TestGetModelInfo) { EXPECT_THROW(model_list_.GetModelInfo("non_existent_model"), std::runtime_error); + + // Clean up + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } TEST_F(ModelListUtilsTestSuite, TestUpdateModelEntry) { @@ -64,6 +62,9 @@ TEST_F(ModelListUtilsTestSuite, TestUpdateModelEntry) { EXPECT_EQ(retrieved_model.status, modellist_utils::ModelStatus::RUNNING); updated_model.status = modellist_utils::ModelStatus::READY; EXPECT_TRUE(model_list_.UpdateModelEntry(kTestModel.model_id, updated_model)); + + // Clean up + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } TEST_F(ModelListUtilsTestSuite, TestDeleteModelEntry) { @@ -85,6 +86,9 @@ TEST_F(ModelListUtilsTestSuite, TestGenerateShortenedAlias) { alias = model_list_.GenerateShortenedAlias( "huggingface.co/bartowski/llama3.1-7b-gguf/Model_ID_Xxx.gguf"); EXPECT_EQ(alias, "llama3.1-7b-gguf:model_id_xxx"); + + // Clean up + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } TEST_F(ModelListUtilsTestSuite, TestPersistence) { From 77d28d2efb4af7ead2276cb0cc4ca78d6abdd63c Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 09:18:12 +0700 Subject: [PATCH 04/15] refactor: db --- engine/CMakeLists.txt | 6 +- engine/commands/chat_cmd.cc | 4 +- engine/commands/model_alias_cmd.cc | 4 +- engine/commands/model_del_cmd.cc | 4 +- engine/commands/model_get_cmd.cc | 4 +- engine/commands/model_import_cmd.cc | 8 +- engine/commands/model_list_cmd.cc | 4 +- engine/commands/model_start_cmd.cc | 4 +- engine/commands/model_status_cmd.cc | 4 +- engine/commands/model_upd_cmd.h | 4 +- engine/commands/run_cmd.cc | 4 +- engine/controllers/models.cc | 16 ++-- engine/database/database.h | 28 ++++++ .../modellist_utils.cc => database/models.cc} | 94 +++++++++---------- .../modellist_utils.h => database/models.h} | 18 ++-- engine/services/model_service.cc | 14 +-- engine/test/components/CMakeLists.txt | 2 +- ...t_modellist_utils.cc => test_models_db.cc} | 45 +++++---- engine/utils/model_callback_utils.h | 14 +-- 19 files changed, 151 insertions(+), 130 deletions(-) create mode 100644 engine/database/database.h rename engine/{utils/modellist_utils.cc => database/models.cc} (75%) rename engine/{utils/modellist_utils.h => database/models.h} (76%) rename engine/test/components/{test_modellist_utils.cc => test_models_db.cc} (80%) diff --git a/engine/CMakeLists.txt b/engine/CMakeLists.txt index 454dc8f2c..6279813a6 100644 --- a/engine/CMakeLists.txt +++ b/engine/CMakeLists.txt @@ -82,7 +82,6 @@ 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) @@ -116,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 index 922dc32ed..3b516e794 100644 --- a/engine/commands/chat_cmd.cc +++ b/engine/commands/chat_cmd.cc @@ -6,7 +6,7 @@ #include "server_start_cmd.h" #include "trantor/utils/Logger.h" #include "utils/logging_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.h" namespace commands { namespace { @@ -39,7 +39,7 @@ 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); diff --git a/engine/commands/model_alias_cmd.cc b/engine/commands/model_alias_cmd.cc index 2123d06cf..4a4ef98af 100644 --- 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_del_cmd.cc b/engine/commands/model_del_cmd.cc index 7f6b6d32a..399f2de2a 100644 --- a/engine/commands/model_del_cmd.cc +++ b/engine/commands/model_del_cmd.cc @@ -2,11 +2,11 @@ #include "cmd_info.h" #include "config/yaml_config.h" #include "utils/file_manager_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.h" namespace commands { bool ModelDelCmd::Exec(const std::string& model_handle) { - modellist_utils::ModelListUtils modellist_handler; + cortex::db::Models modellist_handler; config::YamlHandler yaml_handler; try { diff --git a/engine/commands/model_get_cmd.cc b/engine/commands/model_get_cmd.cc index 715728c1f..6a488f0fc 100644 --- a/engine/commands/model_get_cmd.cc +++ b/engine/commands/model_get_cmd.cc @@ -7,12 +7,12 @@ #include "config/yaml_config.h" #include "utils/file_manager_utils.h" #include "utils/logging_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.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); diff --git a/engine/commands/model_import_cmd.cc b/engine/commands/model_import_cmd.cc index 3fb047a9d..96281ddd3 100644 --- a/engine/commands/model_import_cmd.cc +++ b/engine/commands/model_import_cmd.cc @@ -5,7 +5,7 @@ #include "config/yaml_config.h" #include "utils/file_manager_utils.h" #include "utils/logging_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.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_, cortex::db::ModelStatus::READY}; try { std::filesystem::create_directories( std::filesystem::path(model_yaml_path).parent_path()); diff --git a/engine/commands/model_list_cmd.cc b/engine/commands/model_list_cmd.cc index 6e3990eb6..9ca2db619 100644 --- a/engine/commands/model_list_cmd.cc +++ b/engine/commands/model_list_cmd.cc @@ -6,13 +6,13 @@ #include "config/yaml_config.h" #include "utils/file_manager_utils.h" #include "utils/logging_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.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; diff --git a/engine/commands/model_start_cmd.cc b/engine/commands/model_start_cmd.cc index 1340614d9..4f8a7ca61 100644 --- a/engine/commands/model_start_cmd.cc +++ b/engine/commands/model_start_cmd.cc @@ -7,13 +7,13 @@ #include "trantor/utils/Logger.h" #include "utils/file_manager_utils.h" #include "utils/logging_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.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); diff --git a/engine/commands/model_status_cmd.cc b/engine/commands/model_status_cmd.cc index e6ba9bbe0..959006bec 100644 --- a/engine/commands/model_status_cmd.cc +++ b/engine/commands/model_status_cmd.cc @@ -3,12 +3,12 @@ #include "httplib.h" #include "nlohmann/json.hpp" #include "utils/logging_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.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); diff --git a/engine/commands/model_upd_cmd.h b/engine/commands/model_upd_cmd.h index 51f5a88d3..c2bffc81a 100644 --- a/engine/commands/model_upd_cmd.h +++ b/engine/commands/model_upd_cmd.h @@ -5,7 +5,7 @@ #include #include #include "config/model_config.h" -#include "utils/modellist_utils.h" +#include "database/models.h" #include "config/yaml_config.h" namespace commands { class ModelUpdCmd { @@ -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 index ea551987f..94f596bb3 100644 --- a/engine/commands/run_cmd.cc +++ b/engine/commands/run_cmd.cc @@ -5,14 +5,14 @@ #include "model_status_cmd.h" #include "server_start_cmd.h" #include "utils/logging_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.h" namespace commands { void RunCmd::Exec() { std::optional 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_); diff --git a/engine/controllers/models.cc b/engine/controllers/models.cc index 4660b50e5..c7e946886 100644 --- a/engine/controllers/models.cc +++ b/engine/controllers/models.cc @@ -5,7 +5,7 @@ #include "utils/cortex_utils.h" #include "utils/file_manager_utils.h" #include "utils/model_callback_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.h" void Models::PullModel( const HttpRequestPtr& req, @@ -56,7 +56,7 @@ void Models::ListModel( // Iterate through directory try { - modellist_utils::ModelListUtils modellist_handler; + cortex::db::Models modellist_handler; config::YamlHandler yaml_handler; auto list_entry = modellist_handler.LoadModelList(); @@ -107,7 +107,7 @@ void Models::GetModel( Json::Value data(Json::arrayValue); try { - modellist_utils::ModelListUtils modellist_handler; + cortex::db::Models modellist_handler; config::YamlHandler yaml_handler; auto model_entry = modellist_handler.GetModelInfo(model_handle); yaml_handler.ModelConfigFromFile(model_entry.path_to_model_yaml); @@ -164,7 +164,7 @@ void Models::UpdateModel( auto model_id = (*(req->getJsonObject())).get("modelId", "").asString(); auto json_body = *(req->getJsonObject()); try { - modellist_utils::ModelListUtils model_list_utils; + cortex::db::Models model_list_utils; auto model_entry = model_list_utils.GetModelInfo(model_id); config::YamlHandler yaml_handler; yaml_handler.ModelConfigFromFile(model_entry.path_to_model_yaml); @@ -209,15 +209,15 @@ void Models::ImportModel( auto modelPath = (*(req->getJsonObject())).get("modelPath", "").asString(); 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(modelHandle + ".yml")) .string(); - modellist_utils::ModelEntry model_entry{ + cortex::db::ModelEntry model_entry{ modelHandle, "local", "imported", - model_yaml_path, modelHandle, modellist_utils::ModelStatus::READY}; + model_yaml_path, modelHandle, cortex::db::ModelStatus::READY}; try { std::filesystem::create_directories( std::filesystem::path(model_yaml_path).parent_path()); @@ -281,7 +281,7 @@ void Models::SetModelAlias( LOG_DEBUG << "GetModel, Model handle: " << model_handle << ", Model alias: " << model_alias; - modellist_utils::ModelListUtils modellist_handler; + cortex::db::Models modellist_handler; try { if (modellist_handler.UpdateModelAlias(model_handle, model_alias)) { std::string message = "Successfully set model alias '" + model_alias + diff --git a/engine/database/database.h b/engine/database/database.h new file mode 100644 index 000000000..239337b18 --- /dev/null +++ b/engine/database/database.h @@ -0,0 +1,28 @@ +#pragma once +#include +#include +#include "SQLiteCpp/SQLiteCpp.h" +#include "utils/file_manager_utils.h" + +namespace cortex::db { +const std::string kDefaultDbPath = + file_manager_utils::GetCortexDataPath().string() + "/cortex.db"; +class Database { + public: + Database(Database const&) = delete; + Database& operator=(Database const&) = delete; + ~Database() {} + + static Database& GetInstance() { + static Database db; + return db; + } + + SQLite::Database& db() { return db_; } + + private: + Database() + : db_(kDefaultDbPath, SQLite::OPEN_READWRITE | SQLite::OPEN_CREATE) {} + SQLite::Database db_; +}; +} // namespace cortex::db \ No newline at end of file diff --git a/engine/utils/modellist_utils.cc b/engine/database/models.cc similarity index 75% rename from engine/utils/modellist_utils.cc rename to engine/database/models.cc index 24cc80559..5de93830f 100644 --- a/engine/utils/modellist_utils.cc +++ b/engine/database/models.cc @@ -1,23 +1,16 @@ -#include "modellist_utils.h" +#include "models.h" #include #include #include #include #include #include -#include "file_manager_utils.h" +#include "database.h" +#include "utils/file_manager_utils.h" -namespace modellist_utils { -const std::string ModelListUtils::kModelListPath = - (file_manager_utils::GetModelsContainerPath() / - std::filesystem::path("model.list")) - .string(); -namespace { -const std::string kDefaultDbPath = - file_manager_utils::GetCortexDataPath().string() + "/cortex.db"; -} -ModelListUtils::ModelListUtils() - : db_(kDefaultDbPath, SQLite::OPEN_READWRITE | SQLite::OPEN_CREATE) { +namespace cortex::db { + +Models::Models() : db_(cortex::db::Database::GetInstance().db()) { db_.exec( "CREATE TABLE IF NOT EXISTS models (" "model_id TEXT PRIMARY KEY," @@ -28,8 +21,7 @@ ModelListUtils::ModelListUtils() "status TEXT);"); } -ModelListUtils::ModelListUtils(const std::string& db_path) - : db_(db_path, SQLite::OPEN_READWRITE | SQLite::OPEN_CREATE) { +Models::Models(SQLite::Database& db) : db_(db) { db_.exec( "CREATE TABLE IF NOT EXISTS models (" "model_id TEXT PRIMARY KEY," @@ -40,9 +32,9 @@ ModelListUtils::ModelListUtils(const std::string& db_path) "status TEXT);"); } -ModelListUtils::~ModelListUtils() {} +Models::~Models() {} -std::vector ModelListUtils::LoadModelList() const { +std::vector Models::LoadModelList() const { std::vector entries; SQLite::Statement query( db_, @@ -64,8 +56,8 @@ std::vector ModelListUtils::LoadModelList() const { return entries; } -bool ModelListUtils::IsUnique(const std::string& model_id, - const std::string& model_alias) const { +bool Models::IsUnique(const std::string& model_id, + const std::string& model_alias) const { SQLite::Statement query(db_, "SELECT COUNT(*) FROM models WHERE model_id = ? OR " "model_id = ? OR model_alias = ? OR model_alias = ?"); @@ -79,30 +71,28 @@ bool ModelListUtils::IsUnique(const std::string& model_id, return false; } -void ModelListUtils::SaveModelList( - const std::vector& entries) const { - std::ofstream file(kModelListPath); - db_.exec("BEGIN TRANSACTION;"); - for (const auto& entry : entries) { - SQLite::Statement insert( - db_, - "INSERT OR REPLACE INTO models (model_id, author_repo_id, " - "branch_name, path_to_model_yaml, model_alias, status) VALUES (?, ?, " - "?, ?, ?, ?)"); - insert.bind(1, entry.model_id); - insert.bind(2, entry.author_repo_id); - insert.bind(3, entry.branch_name); - insert.bind(4, entry.path_to_model_yaml); - insert.bind(5, entry.model_alias); - insert.bind(6, - (entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY")); - insert.exec(); - } - db_.exec("COMMIT;"); -} +// void Models::SaveModelList(const std::vector& entries) const { +// std::ofstream file(kModelListPath); +// db_.exec("BEGIN TRANSACTION;"); +// for (const auto& entry : entries) { +// SQLite::Statement insert( +// db_, +// "INSERT OR REPLACE INTO models (model_id, author_repo_id, " +// "branch_name, path_to_model_yaml, model_alias, status) VALUES (?, ?, " +// "?, ?, ?, ?)"); +// insert.bind(1, entry.model_id); +// insert.bind(2, entry.author_repo_id); +// insert.bind(3, entry.branch_name); +// insert.bind(4, entry.path_to_model_yaml); +// insert.bind(5, entry.model_alias); +// insert.bind(6, +// (entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY")); +// insert.exec(); +// } +// db_.exec("COMMIT;"); +// } -std::string ModelListUtils::GenerateShortenedAlias( - const std::string& model_id) const { +std::string Models::GenerateShortenedAlias(const std::string& model_id) const { std::vector parts; std::istringstream iss(model_id); std::string part; @@ -161,7 +151,7 @@ std::string ModelListUtils::GenerateShortenedAlias( return unique_candidate; } -ModelEntry ModelListUtils::GetModelInfo(const std::string& identifier) const { +ModelEntry Models::GetModelInfo(const std::string& identifier) const { std::lock_guard lock(mutex_); SQLite::Statement query(db_, "SELECT model_id, author_repo_id, branch_name, " @@ -186,7 +176,7 @@ ModelEntry ModelListUtils::GetModelInfo(const std::string& identifier) const { } } -void ModelListUtils::PrintModelInfo(const ModelEntry& entry) const { +void Models::PrintModelInfo(const ModelEntry& entry) const { LOG_INFO << "Model ID: " << entry.model_id; LOG_INFO << "Author/Repo ID: " << entry.author_repo_id; LOG_INFO << "Branch Name: " << entry.branch_name; @@ -196,7 +186,7 @@ void ModelListUtils::PrintModelInfo(const ModelEntry& entry) const { << (entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY"); } -bool ModelListUtils::AddModelEntry(ModelEntry new_entry, bool use_short_alias) { +bool Models::AddModelEntry(ModelEntry new_entry, bool use_short_alias) { std::lock_guard lock(mutex_); if (IsUnique(new_entry.model_id, new_entry.model_alias)) { @@ -224,8 +214,8 @@ bool ModelListUtils::AddModelEntry(ModelEntry new_entry, bool use_short_alias) { return false; // Entry not added due to non-uniqueness } -bool ModelListUtils::UpdateModelEntry(const std::string& identifier, - const ModelEntry& updated_entry) { +bool Models::UpdateModelEntry(const std::string& identifier, + const ModelEntry& updated_entry) { std::lock_guard lock(mutex_); SQLite::Statement upd(db_, "UPDATE models " @@ -242,8 +232,8 @@ bool ModelListUtils::UpdateModelEntry(const std::string& identifier, return upd.exec() == 1; } -bool ModelListUtils::UpdateModelAlias(const std::string& model_id, - const std::string& new_model_alias) { +bool Models::UpdateModelAlias(const std::string& model_id, + const std::string& new_model_alias) { std::lock_guard lock(mutex_); // Check new_model_alias is unique if (IsUnique(new_model_alias, new_model_alias)) { @@ -259,7 +249,7 @@ bool ModelListUtils::UpdateModelAlias(const std::string& model_id, return false; } -bool ModelListUtils::DeleteModelEntry(const std::string& identifier) { +bool Models::DeleteModelEntry(const std::string& identifier) { std::lock_guard lock(mutex_); SQLite::Statement del( db_, "DELETE from models WHERE model_id = ? OR model_alias = ?"); @@ -268,7 +258,7 @@ bool ModelListUtils::DeleteModelEntry(const std::string& identifier) { return del.exec() == 1; } -bool ModelListUtils::HasModel(const std::string& identifier) const { +bool Models::HasModel(const std::string& identifier) const { std::lock_guard lock(mutex_); SQLite::Statement query( db_, "SELECT COUNT(*) FROM models WHERE model_id = ? OR model_alias = ?"); @@ -279,4 +269,4 @@ bool ModelListUtils::HasModel(const std::string& identifier) const { } return false; } -} // namespace modellist_utils +} // namespace cortex::db diff --git a/engine/utils/modellist_utils.h b/engine/database/models.h similarity index 76% rename from engine/utils/modellist_utils.h rename to engine/database/models.h index be849a891..12cb662e6 100644 --- a/engine/utils/modellist_utils.h +++ b/engine/database/models.h @@ -7,7 +7,7 @@ #include "SQLiteCpp/SQLiteCpp.h" -namespace modellist_utils { +namespace cortex::db { enum class ModelStatus { READY, RUNNING }; @@ -20,24 +20,22 @@ struct ModelEntry { ModelStatus status; }; -class ModelListUtils { +class Models { private: - mutable SQLite::Database db_; + SQLite::Database& db_; mutable std::mutex mutex_; // For thread safety bool IsUnique(const std::string& model_id, const std::string& model_alias) const; - void SaveModelList(const std::vector& entries) const; public: static const std::string kModelListPath; std::vector LoadModelList() const; - ModelListUtils(); - ModelListUtils(const std::string& db_path); - ~ModelListUtils(); - std::string GenerateShortenedAlias( - const std::string& model_id) const; + Models(); + Models(SQLite::Database& db); + ~Models(); + std::string GenerateShortenedAlias(const std::string& model_id) const; ModelEntry GetModelInfo(const std::string& identifier) const; void PrintModelInfo(const ModelEntry& entry) const; bool AddModelEntry(ModelEntry new_entry, bool use_short_alias = false); @@ -48,4 +46,4 @@ class ModelListUtils { const std::string& model_alias); bool HasModel(const std::string& identifier) const; }; -} // namespace modellist_utils +} // namespace cortex::db diff --git a/engine/services/model_service.cc b/engine/services/model_service.cc index 36a5e7aa5..1693168e2 100644 --- a/engine/services/model_service.cc +++ b/engine/services/model_service.cc @@ -9,7 +9,7 @@ #include "utils/file_manager_utils.h" #include "utils/huggingface_utils.h" #include "utils/logging_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.h" #include "utils/string_utils.h" std::optional ModelService::DownloadModel( @@ -172,14 +172,14 @@ std::optional ModelService::DownloadModelFromCortexso( model_yml_item->localPath.string()); auto mc = yaml_handler.GetModelConfig(); - modellist_utils::ModelListUtils modellist_utils_obj; - modellist_utils::ModelEntry model_entry{ + cortex::db::Models modellist_utils_obj; + cortex::db::ModelEntry model_entry{ .model_id = model_id, .author_repo_id = "cortexso", .branch_name = branch, .path_to_model_yaml = model_yml_item->localPath.string(), .model_alias = model_id, - .status = modellist_utils::ModelStatus::READY}; + .status = cortex::db::ModelStatus::READY}; modellist_utils_obj.AddModelEntry(model_entry); } }); @@ -250,13 +250,13 @@ void ModelService::ParseGguf(const DownloadItem& ggufDownloadItem, CTL_INF("Adding model to modellist with branch: " << branch); auto author_id = author.has_value() ? author.value() : "cortexso"; - modellist_utils::ModelListUtils modellist_utils_obj; - modellist_utils::ModelEntry model_entry{ + cortex::db::Models modellist_utils_obj; + cortex::db::ModelEntry model_entry{ .model_id = ggufDownloadItem.id, .author_repo_id = author_id, .branch_name = branch, .path_to_model_yaml = yaml_name.string(), .model_alias = ggufDownloadItem.id, - .status = modellist_utils::ModelStatus::READY}; + .status = cortex::db::ModelStatus::READY}; modellist_utils_obj.AddModelEntry(model_entry, true); } diff --git a/engine/test/components/CMakeLists.txt b/engine/test/components/CMakeLists.txt index ec25787cb..a321b1821 100644 --- a/engine/test/components/CMakeLists.txt +++ b/engine/test/components/CMakeLists.txt @@ -5,12 +5,12 @@ enable_testing() add_executable(${PROJECT_NAME} ${SRCS} - ${CMAKE_CURRENT_SOURCE_DIR}/../../utils/modellist_utils.cc ${CMAKE_CURRENT_SOURCE_DIR}/../../config/yaml_config.cc ${CMAKE_CURRENT_SOURCE_DIR}/../../config/gguf_parser.cc ${CMAKE_CURRENT_SOURCE_DIR}/../../commands/cortex_upd_cmd.cc ${CMAKE_CURRENT_SOURCE_DIR}/../../commands/server_stop_cmd.cc ${CMAKE_CURRENT_SOURCE_DIR}/../../services/download_service.cc + ${CMAKE_CURRENT_SOURCE_DIR}/../../database/models.cc ) find_package(Drogon CONFIG REQUIRED) diff --git a/engine/test/components/test_modellist_utils.cc b/engine/test/components/test_models_db.cc similarity index 80% rename from engine/test/components/test_modellist_utils.cc rename to engine/test/components/test_models_db.cc index a3f975ee0..58f10b163 100644 --- a/engine/test/components/test_modellist_utils.cc +++ b/engine/test/components/test_models_db.cc @@ -2,28 +2,32 @@ #include #include "gtest/gtest.h" #include "utils/file_manager_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.h" +namespace cortex::db { namespace { constexpr const auto kTestDb = "./test.db"; } -class ModelListUtilsTestSuite : public ::testing::Test { +class ModelsTestSuite : public ::testing::Test { public: - ModelListUtilsTestSuite() : model_list_(kTestDb) {} + ModelsTestSuite() + : db_(kTestDb, SQLite::OPEN_READWRITE | SQLite::OPEN_CREATE), + model_list_(db_) {} protected: - modellist_utils::ModelListUtils model_list_; + SQLite::Database db_; + cortex::db::Models model_list_; - const modellist_utils::ModelEntry kTestModel{ + const cortex::db::ModelEntry kTestModel{ kTestModel.model_id, "test_author", "main", "/path/to/model.yaml", "test_alias", - modellist_utils::ModelStatus::READY}; + cortex::db::ModelStatus::READY}; }; -TEST_F(ModelListUtilsTestSuite, TestAddModelEntry) { +TEST_F(ModelsTestSuite, TestAddModelEntry) { EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); auto retrieved_model = model_list_.GetModelInfo(kTestModel.model_id); @@ -34,7 +38,7 @@ TEST_F(ModelListUtilsTestSuite, TestAddModelEntry) { EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } -TEST_F(ModelListUtilsTestSuite, TestGetModelInfo) { +TEST_F(ModelsTestSuite, TestGetModelInfo) { EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); auto model_by_id = model_list_.GetModelInfo(kTestModel.model_id); @@ -50,24 +54,24 @@ TEST_F(ModelListUtilsTestSuite, TestGetModelInfo) { EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } -TEST_F(ModelListUtilsTestSuite, TestUpdateModelEntry) { +TEST_F(ModelsTestSuite, TestUpdateModelEntry) { EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); - modellist_utils::ModelEntry updated_model = kTestModel; - updated_model.status = modellist_utils::ModelStatus::RUNNING; + cortex::db::ModelEntry updated_model = kTestModel; + updated_model.status = cortex::db::ModelStatus::RUNNING; EXPECT_TRUE(model_list_.UpdateModelEntry(kTestModel.model_id, updated_model)); auto retrieved_model = model_list_.GetModelInfo(kTestModel.model_id); - EXPECT_EQ(retrieved_model.status, modellist_utils::ModelStatus::RUNNING); - updated_model.status = modellist_utils::ModelStatus::READY; + EXPECT_EQ(retrieved_model.status, cortex::db::ModelStatus::RUNNING); + updated_model.status = cortex::db::ModelStatus::READY; EXPECT_TRUE(model_list_.UpdateModelEntry(kTestModel.model_id, updated_model)); // Clean up EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } -TEST_F(ModelListUtilsTestSuite, TestDeleteModelEntry) { +TEST_F(ModelsTestSuite, TestDeleteModelEntry) { EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); @@ -75,7 +79,7 @@ TEST_F(ModelListUtilsTestSuite, TestDeleteModelEntry) { std::runtime_error); } -TEST_F(ModelListUtilsTestSuite, TestGenerateShortenedAlias) { +TEST_F(ModelsTestSuite, TestGenerateShortenedAlias) { EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); auto alias = model_list_.GenerateShortenedAlias( "huggingface.co/bartowski/llama3.1-7b-gguf/Model_ID_Xxx.gguf"); @@ -91,11 +95,11 @@ TEST_F(ModelListUtilsTestSuite, TestGenerateShortenedAlias) { EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } -TEST_F(ModelListUtilsTestSuite, TestPersistence) { +TEST_F(ModelsTestSuite, TestPersistence) { EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); // Create a new ModelListUtils instance to test if it loads from file - modellist_utils::ModelListUtils new_model_list(kTestDb); + cortex::db::Models new_model_list(db_); auto retrieved_model = new_model_list.GetModelInfo(kTestModel.model_id); EXPECT_EQ(retrieved_model.model_id, kTestModel.model_id); @@ -103,7 +107,7 @@ TEST_F(ModelListUtilsTestSuite, TestPersistence) { EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } -TEST_F(ModelListUtilsTestSuite, TestUpdateModelAlias) { +TEST_F(ModelsTestSuite, TestUpdateModelAlias) { constexpr const auto kNewTestAlias = "new_test_alias"; constexpr const auto kNonExistentModel = "non_existent_model"; constexpr const auto kAnotherAlias = "another_alias"; @@ -122,7 +126,7 @@ TEST_F(ModelListUtilsTestSuite, TestUpdateModelAlias) { EXPECT_FALSE(model_list_.UpdateModelAlias(kNonExistentModel, kAnotherAlias)); // Test update with non-unique alias - modellist_utils::ModelEntry another_model = kTestModel; + cortex::db::ModelEntry another_model = kTestModel; another_model.model_id = kAnotherModelId; another_model.model_alias = kAnotherAlias; ASSERT_TRUE(model_list_.AddModelEntry(another_model)); @@ -141,7 +145,7 @@ TEST_F(ModelListUtilsTestSuite, TestUpdateModelAlias) { EXPECT_TRUE(model_list_.DeleteModelEntry(kAnotherModelId)); } -TEST_F(ModelListUtilsTestSuite, TestHasModel) { +TEST_F(ModelsTestSuite, TestHasModel) { EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); EXPECT_TRUE(model_list_.HasModel(kTestModel.model_id)); @@ -149,4 +153,5 @@ TEST_F(ModelListUtilsTestSuite, TestHasModel) { EXPECT_FALSE(model_list_.HasModel("non_existent_model")); // Clean up EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); +} } \ No newline at end of file diff --git a/engine/utils/model_callback_utils.h b/engine/utils/model_callback_utils.h index c6e98dd48..9ef6947cc 100644 --- a/engine/utils/model_callback_utils.h +++ b/engine/utils/model_callback_utils.h @@ -8,7 +8,7 @@ #include "services/download_service.h" #include "utils/huggingface_utils.h" #include "utils/logging_utils.h" -#include "utils/modellist_utils.h" +#include "database/models.h" namespace model_callback_utils { @@ -35,14 +35,14 @@ inline void ParseGguf(const DownloadItem& ggufDownloadItem, CTL_INF("Adding model to modellist with branch: " << branch); auto author_id = author.has_value() ? author.value() : "cortexso"; - modellist_utils::ModelListUtils modellist_utils_obj; - modellist_utils::ModelEntry model_entry{ + cortex::db::Models modellist_utils_obj; + cortex::db::ModelEntry model_entry{ .model_id = model_config.id, .author_repo_id = author_id, .branch_name = branch, .path_to_model_yaml = yaml_name.string(), .model_alias = model_config.id, - .status = modellist_utils::ModelStatus::READY}; + .status = cortex::db::ModelStatus::READY}; modellist_utils_obj.AddModelEntry(model_entry); } @@ -76,14 +76,14 @@ inline void DownloadModelCb(const DownloadTask& finishedTask) { yaml_handler.ModelConfigFromFile(model_yml_di->localPath.string()); auto mc = yaml_handler.GetModelConfig(); - modellist_utils::ModelListUtils modellist_utils_obj; - modellist_utils::ModelEntry model_entry{ + cortex::db::Models modellist_utils_obj; + cortex::db::ModelEntry model_entry{ .model_id = mc.name, .author_repo_id = "cortexso", .branch_name = branch, .path_to_model_yaml = model_yml_di->localPath.string(), .model_alias = mc.name, - .status = modellist_utils::ModelStatus::READY}; + .status = cortex::db::ModelStatus::READY}; modellist_utils_obj.AddModelEntry(model_entry); } } From 4833e6356543569c6c44381ca40f9cebf298ea39 Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 09:27:36 +0700 Subject: [PATCH 05/15] fix: remove mutex --- engine/database/models.cc | 33 +++++---------------------------- engine/database/models.h | 2 -- 2 files changed, 5 insertions(+), 30 deletions(-) diff --git a/engine/database/models.cc b/engine/database/models.cc index 5de93830f..3f6689be4 100644 --- a/engine/database/models.cc +++ b/engine/database/models.cc @@ -7,6 +7,7 @@ #include #include "database.h" #include "utils/file_manager_utils.h" +#include "utils/scope_exit.h" namespace cortex::db { @@ -71,27 +72,6 @@ bool Models::IsUnique(const std::string& model_id, return false; } -// void Models::SaveModelList(const std::vector& entries) const { -// std::ofstream file(kModelListPath); -// db_.exec("BEGIN TRANSACTION;"); -// for (const auto& entry : entries) { -// SQLite::Statement insert( -// db_, -// "INSERT OR REPLACE INTO models (model_id, author_repo_id, " -// "branch_name, path_to_model_yaml, model_alias, status) VALUES (?, ?, " -// "?, ?, ?, ?)"); -// insert.bind(1, entry.model_id); -// insert.bind(2, entry.author_repo_id); -// insert.bind(3, entry.branch_name); -// insert.bind(4, entry.path_to_model_yaml); -// insert.bind(5, entry.model_alias); -// insert.bind(6, -// (entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY")); -// insert.exec(); -// } -// db_.exec("COMMIT;"); -// } - std::string Models::GenerateShortenedAlias(const std::string& model_id) const { std::vector parts; std::istringstream iss(model_id); @@ -152,7 +132,6 @@ std::string Models::GenerateShortenedAlias(const std::string& model_id) const { } ModelEntry Models::GetModelInfo(const std::string& identifier) const { - std::lock_guard lock(mutex_); SQLite::Statement query(db_, "SELECT model_id, author_repo_id, branch_name, " "path_to_model_yaml, model_alias, status FROM models " @@ -187,8 +166,8 @@ void Models::PrintModelInfo(const ModelEntry& entry) const { } bool Models::AddModelEntry(ModelEntry new_entry, bool use_short_alias) { - std::lock_guard lock(mutex_); - + db_.exec("BEGIN TRANSACTION;"); + utils::ScopeExit se([this] { db_.exec("COMMIT;"); }); if (IsUnique(new_entry.model_id, new_entry.model_alias)) { if (use_short_alias) { new_entry.model_alias = GenerateShortenedAlias(new_entry.model_id); @@ -216,7 +195,6 @@ bool Models::AddModelEntry(ModelEntry new_entry, bool use_short_alias) { bool Models::UpdateModelEntry(const std::string& identifier, const ModelEntry& updated_entry) { - std::lock_guard lock(mutex_); SQLite::Statement upd(db_, "UPDATE models " "SET author_repo_id = ?, branch_name = ?, " @@ -234,7 +212,8 @@ bool Models::UpdateModelEntry(const std::string& identifier, bool Models::UpdateModelAlias(const std::string& model_id, const std::string& new_model_alias) { - std::lock_guard lock(mutex_); + db_.exec("BEGIN TRANSACTION;"); + utils::ScopeExit se([this] { db_.exec("COMMIT;"); }); // Check new_model_alias is unique if (IsUnique(new_model_alias, new_model_alias)) { SQLite::Statement upd(db_, @@ -250,7 +229,6 @@ bool Models::UpdateModelAlias(const std::string& model_id, } bool Models::DeleteModelEntry(const std::string& identifier) { - std::lock_guard lock(mutex_); SQLite::Statement del( db_, "DELETE from models WHERE model_id = ? OR model_alias = ?"); del.bind(1, identifier); @@ -259,7 +237,6 @@ bool Models::DeleteModelEntry(const std::string& identifier) { } bool Models::HasModel(const std::string& identifier) const { - std::lock_guard lock(mutex_); SQLite::Statement query( db_, "SELECT COUNT(*) FROM models WHERE model_id = ? OR model_alias = ?"); query.bind(1, identifier); diff --git a/engine/database/models.h b/engine/database/models.h index 12cb662e6..a775cacc1 100644 --- a/engine/database/models.h +++ b/engine/database/models.h @@ -1,7 +1,6 @@ #pragma once #include -#include #include #include @@ -24,7 +23,6 @@ class Models { private: SQLite::Database& db_; - mutable std::mutex mutex_; // For thread safety bool IsUnique(const std::string& model_id, const std::string& model_alias) const; From f0f1441242d9bacf7f65d5df3a757dd1864e54d7 Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 09:37:46 +0700 Subject: [PATCH 06/15] fix: rm file --- engine/test/components/test_models_db.cc | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/engine/test/components/test_models_db.cc b/engine/test/components/test_models_db.cc index 58f10b163..6823461fb 100644 --- a/engine/test/components/test_models_db.cc +++ b/engine/test/components/test_models_db.cc @@ -1,8 +1,8 @@ #include #include +#include "database/models.h" #include "gtest/gtest.h" #include "utils/file_manager_utils.h" -#include "database/models.h" namespace cortex::db { namespace { @@ -13,18 +13,21 @@ class ModelsTestSuite : public ::testing::Test { ModelsTestSuite() : db_(kTestDb, SQLite::OPEN_READWRITE | SQLite::OPEN_CREATE), model_list_(db_) {} + ~ModelsTestSuite() { + try { + std::filesystem::remove(kTestDb); + } catch (const std::exception& e) { + + } + } protected: SQLite::Database db_; cortex::db::Models model_list_; const cortex::db::ModelEntry kTestModel{ - kTestModel.model_id, - "test_author", - "main", - "/path/to/model.yaml", - "test_alias", - cortex::db::ModelStatus::READY}; + kTestModel.model_id, "test_author", "main", + "/path/to/model.yaml", "test_alias", cortex::db::ModelStatus::READY}; }; TEST_F(ModelsTestSuite, TestAddModelEntry) { @@ -154,4 +157,4 @@ TEST_F(ModelsTestSuite, TestHasModel) { // Clean up EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); } -} \ No newline at end of file +} // namespace cortex::db \ No newline at end of file From 7bd1c5dd5decd3741a9810b266603b4fcad2d014 Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 10:24:50 +0700 Subject: [PATCH 07/15] fix: test --- engine/test/components/test_models_db.cc | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/engine/test/components/test_models_db.cc b/engine/test/components/test_models_db.cc index 6823461fb..d2fcffec3 100644 --- a/engine/test/components/test_models_db.cc +++ b/engine/test/components/test_models_db.cc @@ -13,12 +13,8 @@ class ModelsTestSuite : public ::testing::Test { ModelsTestSuite() : db_(kTestDb, SQLite::OPEN_READWRITE | SQLite::OPEN_CREATE), model_list_(db_) {} - ~ModelsTestSuite() { - try { - std::filesystem::remove(kTestDb); - } catch (const std::exception& e) { - - } + void SetUp() { + db_.exec("DELETE FROM models"); } protected: From 201c9886f51c3d7c0bbc8423e0f30832296bae9c Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 10:55:06 +0700 Subject: [PATCH 08/15] fix: test --- engine/test/components/test_models_db.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/engine/test/components/test_models_db.cc b/engine/test/components/test_models_db.cc index d2fcffec3..4abe66d1c 100644 --- a/engine/test/components/test_models_db.cc +++ b/engine/test/components/test_models_db.cc @@ -14,7 +14,9 @@ class ModelsTestSuite : public ::testing::Test { : db_(kTestDb, SQLite::OPEN_READWRITE | SQLite::OPEN_CREATE), model_list_(db_) {} void SetUp() { - db_.exec("DELETE FROM models"); + try { + db_.exec("DELETE FROM models"); + } catch (const std::exception& e) {} } protected: @@ -22,7 +24,7 @@ class ModelsTestSuite : public ::testing::Test { cortex::db::Models model_list_; const cortex::db::ModelEntry kTestModel{ - kTestModel.model_id, "test_author", "main", + "test_model_id", "test_author", "main", "/path/to/model.yaml", "test_alias", cortex::db::ModelStatus::READY}; }; From 26581758be690213064b77ab8ad311cec62d01bb Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 11:11:28 +0700 Subject: [PATCH 09/15] refactor: LoadModelList --- engine/commands/model_list_cmd.cc | 36 ++++++++++++------------- engine/controllers/models.cc | 17 ++++++------ engine/database/models.cc | 44 ++++++++++++++++++------------- engine/database/models.h | 4 +-- 4 files changed, 54 insertions(+), 47 deletions(-) diff --git a/engine/commands/model_list_cmd.cc b/engine/commands/model_list_cmd.cc index 4448b7ed0..3fe0b4700 100644 --- a/engine/commands/model_list_cmd.cc +++ b/engine/commands/model_list_cmd.cc @@ -3,9 +3,9 @@ #include #include #include "config/yaml_config.h" +#include "database/models.h" #include "utils/file_manager_utils.h" #include "utils/logging_utils.h" -#include "database/models.h" namespace commands { @@ -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/controllers/models.cc b/engine/controllers/models.cc index 60029281e..c010a9fc1 100644 --- a/engine/controllers/models.cc +++ b/engine/controllers/models.cc @@ -65,13 +65,12 @@ void Models::ListModel( // Iterate through directory - try { - cortex::db::Models modellist_handler; - config::YamlHandler yaml_handler; - - auto list_entry = modellist_handler.LoadModelList(); + cortex::db::Models modellist_handler; + config::YamlHandler yaml_handler; - for (const auto& model_entry : list_entry) { + auto list_entry = modellist_handler.LoadModelList(); + if (list_entry) { + for (const auto& model_entry : list_entry.value()) { // auto model_entry = modellist_handler.GetModelInfo(model_handle); try { @@ -91,9 +90,9 @@ void Models::ListModel( auto resp = cortex_utils::CreateCortexHttpJsonResponse(ret); resp->setStatusCode(k200OK); callback(resp); - } catch (const std::exception& e) { - std::string message = - "Fail to get list model information: " + std::string(e.what()); + } else { + std::string message = "Fail to get list model information: " + + std::string(list_entry.error()); LOG_ERROR << message; ret["data"] = data; ret["result"] = "Fail to get list model information"; diff --git a/engine/database/models.cc b/engine/database/models.cc index 3f6689be4..a03f82db2 100644 --- a/engine/database/models.cc +++ b/engine/database/models.cc @@ -7,6 +7,7 @@ #include #include "database.h" #include "utils/file_manager_utils.h" +#include "utils/result.hpp" #include "utils/scope_exit.h" namespace cortex::db { @@ -35,26 +36,33 @@ Models::Models(SQLite::Database& db) : db_(db) { Models::~Models() {} -std::vector Models::LoadModelList() const { - std::vector entries; - SQLite::Statement query( - db_, - "SELECT model_id, author_repo_id, branch_name, " - "path_to_model_yaml, model_alias, status FROM models"); +cpp::result, std::string> Models::LoadModelList() + const { + try { + std::vector entries; + db_.exec("BEGIN TRANSACTION;"); + utils::ScopeExit se([this] { db_.exec("COMMIT;"); }); + SQLite::Statement query( + db_, + "SELECT model_id, author_repo_id, branch_name, " + "path_to_model_yaml, model_alias, status FROM models"); - while (query.executeStep()) { - ModelEntry entry; - entry.model_id = query.getColumn(0).getString(); - entry.author_repo_id = query.getColumn(1).getString(); - entry.branch_name = query.getColumn(2).getString(); - entry.path_to_model_yaml = query.getColumn(3).getString(); - entry.model_alias = query.getColumn(4).getString(); - std::string status_str = query.getColumn(5).getString(); - entry.status = - (status_str == "RUNNING") ? ModelStatus::RUNNING : ModelStatus::READY; - entries.push_back(entry); + while (query.executeStep()) { + ModelEntry entry; + entry.model_id = query.getColumn(0).getString(); + entry.author_repo_id = query.getColumn(1).getString(); + entry.branch_name = query.getColumn(2).getString(); + entry.path_to_model_yaml = query.getColumn(3).getString(); + entry.model_alias = query.getColumn(4).getString(); + std::string status_str = query.getColumn(5).getString(); + entry.status = + (status_str == "RUNNING") ? ModelStatus::RUNNING : ModelStatus::READY; + entries.push_back(entry); + } + return entries; + } catch (const std::exception& e) { + return cpp::fail(e.what()); } - return entries; } bool Models::IsUnique(const std::string& model_id, diff --git a/engine/database/models.h b/engine/database/models.h index a775cacc1..79ca27908 100644 --- a/engine/database/models.h +++ b/engine/database/models.h @@ -3,8 +3,8 @@ #include #include #include - #include "SQLiteCpp/SQLiteCpp.h" +#include "utils/result.hpp" namespace cortex::db { @@ -29,7 +29,7 @@ class Models { public: static const std::string kModelListPath; - std::vector LoadModelList() const; + cpp::result, std::string> LoadModelList() const; Models(); Models(SQLite::Database& db); ~Models(); From 383c6b0313153087f65f25c077ee1605cae4f690 Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 12:15:46 +0700 Subject: [PATCH 10/15] refactor: more --- engine/commands/chat_cmd.cc | 6 +- engine/commands/model_get_cmd.cc | 6 +- engine/commands/model_import_cmd.cc | 4 +- engine/commands/model_start_cmd.cc | 8 +- engine/commands/model_status_cmd.cc | 6 +- engine/commands/model_upd_cmd.cc | 8 +- engine/commands/run_cmd.cc | 8 +- engine/controllers/models.cc | 12 +- engine/database/models.cc | 240 +++++++++++++---------- engine/database/models.h | 18 +- engine/services/model_service.cc | 10 +- engine/test/components/test_models_db.cc | 98 +++++---- 12 files changed, 256 insertions(+), 168 deletions(-) diff --git a/engine/commands/chat_cmd.cc b/engine/commands/chat_cmd.cc index 3b516e794..b1e56773f 100644 --- a/engine/commands/chat_cmd.cc +++ b/engine/commands/chat_cmd.cc @@ -43,7 +43,11 @@ void ChatCmd::Exec(const std::string& host, int port, 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_get_cmd.cc b/engine/commands/model_get_cmd.cc index 6a488f0fc..df11f982d 100644 --- a/engine/commands/model_get_cmd.cc +++ b/engine/commands/model_get_cmd.cc @@ -16,7 +16,11 @@ void ModelGetCmd::Exec(const std::string& model_handle) { 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 index 96281ddd3..688753511 100644 --- a/engine/commands/model_import_cmd.cc +++ b/engine/commands/model_import_cmd.cc @@ -3,9 +3,9 @@ #include #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 "database/models.h" namespace commands { @@ -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_start_cmd.cc b/engine/commands/model_start_cmd.cc index 4f8a7ca61..2b0c8f2b9 100644 --- a/engine/commands/model_start_cmd.cc +++ b/engine/commands/model_start_cmd.cc @@ -1,5 +1,6 @@ #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" @@ -7,7 +8,6 @@ #include "trantor/utils/Logger.h" #include "utils/file_manager_utils.h" #include "utils/logging_utils.h" -#include "database/models.h" namespace commands { bool ModelStartCmd::Exec(const std::string& host, int port, @@ -17,7 +17,11 @@ bool ModelStartCmd::Exec(const std::string& host, int port, 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 index 959006bec..12b939f13 100644 --- a/engine/commands/model_status_cmd.cc +++ b/engine/commands/model_status_cmd.cc @@ -12,7 +12,11 @@ bool ModelStatusCmd::IsLoaded(const std::string& host, int port, 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 index 65883def3..e2c52ab46 100644 --- 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& 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/run_cmd.cc b/engine/commands/run_cmd.cc index 38b6c1f70..12182ec55 100644 --- a/engine/commands/run_cmd.cc +++ b/engine/commands/run_cmd.cc @@ -1,11 +1,11 @@ #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 "database/models.h" namespace commands { @@ -30,7 +30,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 diff --git a/engine/controllers/models.cc b/engine/controllers/models.cc index c010a9fc1..af7a9fd85 100644 --- a/engine/controllers/models.cc +++ b/engine/controllers/models.cc @@ -119,7 +119,11 @@ void Models::GetModel( cortex::db::Models modellist_handler; config::YamlHandler yaml_handler; 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(); Json::Value obj = model_config.ToJson(); @@ -174,11 +178,11 @@ void Models::UpdateModel( cortex::db::Models model_list_utils; auto model_entry = model_list_utils.GetModelInfo(model_id); config::YamlHandler yaml_handler; - yaml_handler.ModelConfigFromFile(model_entry.path_to_model_yaml); + yaml_handler.ModelConfigFromFile(model_entry.value().path_to_model_yaml); config::ModelConfig model_config = yaml_handler.GetModelConfig(); model_config.FromJson(json_body); yaml_handler.UpdateModelConfig(model_config); - yaml_handler.WriteYamlFile(model_entry.path_to_model_yaml); + yaml_handler.WriteYamlFile(model_entry.value().path_to_model_yaml); std::string message = "Successfully update model ID '" + model_id + "': " + json_body.toStyledString(); LOG_INFO << message; @@ -234,7 +238,7 @@ void Models::ImportModel( model_config.name = modelHandle; 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); std::string success_message = "Model is imported successfully!"; LOG_INFO << success_message; diff --git a/engine/database/models.cc b/engine/database/models.cc index a03f82db2..aac233c6a 100644 --- a/engine/database/models.cc +++ b/engine/database/models.cc @@ -40,8 +40,6 @@ cpp::result, std::string> Models::LoadModelList() const { try { std::vector entries; - db_.exec("BEGIN TRANSACTION;"); - utils::ScopeExit se([this] { db_.exec("COMMIT;"); }); SQLite::Statement query( db_, "SELECT model_id, author_repo_id, branch_name, " @@ -61,26 +59,24 @@ cpp::result, std::string> Models::LoadModelList() } return entries; } catch (const std::exception& e) { + CTL_WRN(e.what()); return cpp::fail(e.what()); } } -bool Models::IsUnique(const std::string& model_id, +bool Models::IsUnique(const std::vector& entries, + const std::string& model_id, const std::string& model_alias) const { - SQLite::Statement query(db_, - "SELECT COUNT(*) FROM models WHERE model_id = ? OR " - "model_id = ? OR model_alias = ? OR model_alias = ?"); - query.bind(1, model_id); - query.bind(2, model_alias); - query.bind(3, model_id); - query.bind(4, model_alias); - if (query.executeStep()) { - return query.getColumn(0).getInt() == 0; - } - return false; + return std::none_of( + entries.begin(), entries.end(), [&](const ModelEntry& entry) { + return entry.model_id == model_id || entry.model_alias == model_id || + entry.model_id == model_alias || + entry.model_alias == model_alias; + }); } -std::string Models::GenerateShortenedAlias(const std::string& model_id) const { +std::string Models::GenerateShortenedAlias( + const std::string& model_id, const std::vector& entries) const { std::vector parts; std::istringstream iss(model_id); std::string part; @@ -123,7 +119,7 @@ std::string Models::GenerateShortenedAlias(const std::string& model_id) const { // Find the first unique candidate for (const auto& candidate : candidates) { - if (IsUnique(model_id, candidate)) { + if (IsUnique(entries, model_id, candidate)) { return candidate; } } @@ -132,34 +128,40 @@ std::string Models::GenerateShortenedAlias(const std::string& model_id) const { std::string base_candidate = candidates.back(); int suffix = 1; std::string unique_candidate = base_candidate; - while (!IsUnique(model_id, unique_candidate)) { + while (!IsUnique(entries, model_id, unique_candidate)) { unique_candidate = base_candidate + "-" + std::to_string(suffix++); } return unique_candidate; } -ModelEntry Models::GetModelInfo(const std::string& identifier) const { - SQLite::Statement query(db_, - "SELECT model_id, author_repo_id, branch_name, " - "path_to_model_yaml, model_alias, status FROM models " - "WHERE model_id = ? OR model_alias = ?"); +cpp::result Models::GetModelInfo( + const std::string& identifier) const { + try { + SQLite::Statement query( + db_, + "SELECT model_id, author_repo_id, branch_name, " + "path_to_model_yaml, model_alias, status FROM models " + "WHERE model_id = ? OR model_alias = ?"); - query.bind(1, identifier); - query.bind(2, identifier); - if (query.executeStep()) { - ModelEntry entry; - entry.model_id = query.getColumn(0).getString(); - entry.author_repo_id = query.getColumn(1).getString(); - entry.branch_name = query.getColumn(2).getString(); - entry.path_to_model_yaml = query.getColumn(3).getString(); - entry.model_alias = query.getColumn(4).getString(); - std::string status_str = query.getColumn(5).getString(); - entry.status = - (status_str == "RUNNING") ? ModelStatus::RUNNING : ModelStatus::READY; - return entry; - } else { - throw std::runtime_error("Model not found: " + identifier); + query.bind(1, identifier); + query.bind(2, identifier); + if (query.executeStep()) { + ModelEntry entry; + entry.model_id = query.getColumn(0).getString(); + entry.author_repo_id = query.getColumn(1).getString(); + entry.branch_name = query.getColumn(2).getString(); + entry.path_to_model_yaml = query.getColumn(3).getString(); + entry.model_alias = query.getColumn(4).getString(); + std::string status_str = query.getColumn(5).getString(); + entry.status = + (status_str == "RUNNING") ? ModelStatus::RUNNING : ModelStatus::READY; + return entry; + } else { + return cpp::fail("Model not found: " + identifier); + } + } catch (const std::exception& e) { + return cpp::fail(e.what()); } } @@ -173,85 +175,123 @@ void Models::PrintModelInfo(const ModelEntry& entry) const { << (entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY"); } -bool Models::AddModelEntry(ModelEntry new_entry, bool use_short_alias) { - db_.exec("BEGIN TRANSACTION;"); - utils::ScopeExit se([this] { db_.exec("COMMIT;"); }); - if (IsUnique(new_entry.model_id, new_entry.model_alias)) { - if (use_short_alias) { - new_entry.model_alias = GenerateShortenedAlias(new_entry.model_id); +cpp::result Models::AddModelEntry(ModelEntry new_entry, + bool use_short_alias) { + try { + db_.exec("BEGIN TRANSACTION;"); + utils::ScopeExit se([this] { db_.exec("COMMIT;"); }); + auto model_list = LoadModelList(); + if (model_list.has_error()) { + CTL_WRN(model_list.error()); + std::cout << "Test: " << model_list.error(); + return cpp::fail(model_list.error()); } - new_entry.status = ModelStatus::READY; // Set default status to READY + if (IsUnique(model_list.value(), new_entry.model_id, + new_entry.model_alias)) { + if (use_short_alias) { + new_entry.model_alias = + GenerateShortenedAlias(new_entry.model_id, model_list.value()); + } + new_entry.status = ModelStatus::READY; // Set default status to READY - SQLite::Statement insert( - db_, - "INSERT INTO models (model_id, author_repo_id, " - "branch_name, path_to_model_yaml, model_alias, status) VALUES (?, ?, " - "?, ?, ?, ?)"); - insert.bind(1, new_entry.model_id); - insert.bind(2, new_entry.author_repo_id); - insert.bind(3, new_entry.branch_name); - insert.bind(4, new_entry.path_to_model_yaml); - insert.bind(5, new_entry.model_alias); - insert.bind( - 6, (new_entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY")); - insert.exec(); + SQLite::Statement insert( + db_, + "INSERT INTO models (model_id, author_repo_id, " + "branch_name, path_to_model_yaml, model_alias, status) VALUES (?, ?, " + "?, ?, ?, ?)"); + insert.bind(1, new_entry.model_id); + insert.bind(2, new_entry.author_repo_id); + insert.bind(3, new_entry.branch_name); + insert.bind(4, new_entry.path_to_model_yaml); + insert.bind(5, new_entry.model_alias); + insert.bind( + 6, (new_entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY")); + insert.exec(); - return true; + return true; + } + return false; // Entry not added due to non-uniqueness + } catch (const std::exception& e) { + CTL_WRN(e.what()); + return cpp::fail(e.what()); } - return false; // Entry not added due to non-uniqueness -} - -bool Models::UpdateModelEntry(const std::string& identifier, - const ModelEntry& updated_entry) { - SQLite::Statement upd(db_, - "UPDATE models " - "SET author_repo_id = ?, branch_name = ?, " - "path_to_model_yaml = ?, status = ? " - "WHERE model_id = ? OR model_alias = ?"); - upd.bind(1, updated_entry.author_repo_id); - upd.bind(2, updated_entry.branch_name); - upd.bind(3, updated_entry.path_to_model_yaml); - upd.bind( - 4, (updated_entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY")); - upd.bind(5, identifier); - upd.bind(6, identifier); - return upd.exec() == 1; } -bool Models::UpdateModelAlias(const std::string& model_id, - const std::string& new_model_alias) { - db_.exec("BEGIN TRANSACTION;"); - utils::ScopeExit se([this] { db_.exec("COMMIT;"); }); - // Check new_model_alias is unique - if (IsUnique(new_model_alias, new_model_alias)) { +cpp::result Models::UpdateModelEntry( + const std::string& identifier, const ModelEntry& updated_entry) { + try { SQLite::Statement upd(db_, "UPDATE models " - "SET model_alias = ? " + "SET author_repo_id = ?, branch_name = ?, " + "path_to_model_yaml = ?, status = ? " "WHERE model_id = ? OR model_alias = ?"); - upd.bind(1, new_model_alias); - upd.bind(2, model_id); - upd.bind(3, model_id); + upd.bind(1, updated_entry.author_repo_id); + upd.bind(2, updated_entry.branch_name); + upd.bind(3, updated_entry.path_to_model_yaml); + upd.bind(4, (updated_entry.status == ModelStatus::RUNNING ? "RUNNING" + : "READY")); + upd.bind(5, identifier); + upd.bind(6, identifier); return upd.exec() == 1; + } catch (const std::exception& e) { + return cpp::fail(e.what()); } - return false; } -bool Models::DeleteModelEntry(const std::string& identifier) { - SQLite::Statement del( - db_, "DELETE from models WHERE model_id = ? OR model_alias = ?"); - del.bind(1, identifier); - del.bind(2, identifier); - return del.exec() == 1; +cpp::result Models::UpdateModelAlias( + const std::string& model_id, const std::string& new_model_alias) { + try { + db_.exec("BEGIN TRANSACTION;"); + utils::ScopeExit se([this] { db_.exec("COMMIT;"); }); + auto model_list = LoadModelList(); + if (model_list.has_error()) { + CTL_WRN(model_list.error()); + return cpp::fail(model_list.error()); + } + // Check new_model_alias is unique + if (IsUnique(model_list.value(), new_model_alias, new_model_alias)) { + SQLite::Statement upd(db_, + "UPDATE models " + "SET model_alias = ? " + "WHERE model_id = ? OR model_alias = ?"); + upd.bind(1, new_model_alias); + upd.bind(2, model_id); + upd.bind(3, model_id); + return upd.exec() == 1; + } + return false; + } catch (const std::exception& e) { + return cpp::fail(e.what()); + } +} + +cpp::result Models::DeleteModelEntry( + const std::string& identifier) { + try { + SQLite::Statement del( + db_, "DELETE from models WHERE model_id = ? OR model_alias = ?"); + del.bind(1, identifier); + del.bind(2, identifier); + return del.exec() == 1; + } catch (const std::exception& e) { + return cpp::fail(e.what()); + } } bool Models::HasModel(const std::string& identifier) const { - SQLite::Statement query( - db_, "SELECT COUNT(*) FROM models WHERE model_id = ? OR model_alias = ?"); - query.bind(1, identifier); - query.bind(2, identifier); - if (query.executeStep()) { - return query.getColumn(0).getInt() > 0; + try { + SQLite::Statement query( + db_, + "SELECT COUNT(*) FROM models WHERE model_id = ? OR model_alias = ?"); + query.bind(1, identifier); + query.bind(2, identifier); + if (query.executeStep()) { + return query.getColumn(0).getInt() > 0; + } + return false; + } catch (const std::exception& e) { + CTL_WRN(e.what()); + return false; } - return false; } } // namespace cortex::db diff --git a/engine/database/models.h b/engine/database/models.h index 79ca27908..fca823368 100644 --- a/engine/database/models.h +++ b/engine/database/models.h @@ -24,7 +24,8 @@ class Models { private: SQLite::Database& db_; - bool IsUnique(const std::string& model_id, + bool IsUnique(const std::vector& entries, + const std::string& model_id, const std::string& model_alias) const; public: @@ -33,14 +34,17 @@ class Models { Models(); Models(SQLite::Database& db); ~Models(); - std::string GenerateShortenedAlias(const std::string& model_id) const; - ModelEntry GetModelInfo(const std::string& identifier) const; + std::string GenerateShortenedAlias( + const std::string& model_id, + const std::vector& entries) const; + cpp::result GetModelInfo(const std::string& identifier) const; void PrintModelInfo(const ModelEntry& entry) const; - bool AddModelEntry(ModelEntry new_entry, bool use_short_alias = false); - bool UpdateModelEntry(const std::string& identifier, + cpp::result AddModelEntry(ModelEntry new_entry, + bool use_short_alias = false); + cpp::result UpdateModelEntry(const std::string& identifier, const ModelEntry& updated_entry); - bool DeleteModelEntry(const std::string& identifier); - bool UpdateModelAlias(const std::string& model_id, + cpp::result DeleteModelEntry(const std::string& identifier); + cpp::result UpdateModelAlias(const std::string& model_id, const std::string& model_alias); bool HasModel(const std::string& identifier) const; }; diff --git a/engine/services/model_service.cc b/engine/services/model_service.cc index e369f6fd2..4612c0948 100644 --- a/engine/services/model_service.cc +++ b/engine/services/model_service.cc @@ -342,12 +342,16 @@ cpp::result ModelService::DeleteModel( 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 cpp::fail(model_entry.error()); + } + yaml_handler.ModelConfigFromFile(model_entry.value().path_to_model_yaml); auto mc = yaml_handler.GetModelConfig(); // Remove yaml file - std::filesystem::remove(model_entry.path_to_model_yaml); + std::filesystem::remove(model_entry.value().path_to_model_yaml); // Remove model files if they are not imported locally - if (model_entry.branch_name != "imported") { + if (model_entry.value().branch_name != "imported") { if (mc.files.size() > 0) { if (mc.engine == "cortex.llamacpp") { for (auto& file : mc.files) { diff --git a/engine/test/components/test_models_db.cc b/engine/test/components/test_models_db.cc index 4abe66d1c..67f1bd0be 100644 --- a/engine/test/components/test_models_db.cc +++ b/engine/test/components/test_models_db.cc @@ -29,83 +29,91 @@ class ModelsTestSuite : public ::testing::Test { }; TEST_F(ModelsTestSuite, TestAddModelEntry) { - EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel).value()); auto retrieved_model = model_list_.GetModelInfo(kTestModel.model_id); - EXPECT_EQ(retrieved_model.model_id, kTestModel.model_id); - EXPECT_EQ(retrieved_model.author_repo_id, kTestModel.author_repo_id); + EXPECT_TRUE(retrieved_model); + EXPECT_EQ(retrieved_model.value().model_id, kTestModel.model_id); + EXPECT_EQ(retrieved_model.value().author_repo_id, kTestModel.author_repo_id); - // Clean up - EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); + // // Clean up + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id).value()); } TEST_F(ModelsTestSuite, TestGetModelInfo) { - EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel).value()); auto model_by_id = model_list_.GetModelInfo(kTestModel.model_id); - EXPECT_EQ(model_by_id.model_id, kTestModel.model_id); + EXPECT_TRUE(model_by_id); + EXPECT_EQ(model_by_id.value().model_id, kTestModel.model_id); auto model_by_alias = model_list_.GetModelInfo("test_alias"); - EXPECT_EQ(model_by_alias.model_id, kTestModel.model_id); + EXPECT_TRUE(model_by_alias); + EXPECT_EQ(model_by_alias.value().model_id, kTestModel.model_id); - EXPECT_THROW(model_list_.GetModelInfo("non_existent_model"), - std::runtime_error); + EXPECT_TRUE(model_list_.GetModelInfo("non_existent_model").has_error()); // Clean up - EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id).value()); } TEST_F(ModelsTestSuite, TestUpdateModelEntry) { - EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel).value()); cortex::db::ModelEntry updated_model = kTestModel; updated_model.status = cortex::db::ModelStatus::RUNNING; - EXPECT_TRUE(model_list_.UpdateModelEntry(kTestModel.model_id, updated_model)); + EXPECT_TRUE( + model_list_.UpdateModelEntry(kTestModel.model_id, updated_model).value()); auto retrieved_model = model_list_.GetModelInfo(kTestModel.model_id); - EXPECT_EQ(retrieved_model.status, cortex::db::ModelStatus::RUNNING); + EXPECT_TRUE(retrieved_model); + EXPECT_EQ(retrieved_model.value().status, cortex::db::ModelStatus::RUNNING); updated_model.status = cortex::db::ModelStatus::READY; - EXPECT_TRUE(model_list_.UpdateModelEntry(kTestModel.model_id, updated_model)); + EXPECT_TRUE( + model_list_.UpdateModelEntry(kTestModel.model_id, updated_model).value()); // Clean up - EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id).value()); } TEST_F(ModelsTestSuite, TestDeleteModelEntry) { - EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel).value()); - EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); - EXPECT_THROW(model_list_.GetModelInfo(kTestModel.model_id), - std::runtime_error); + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id).value()); + EXPECT_TRUE(model_list_.GetModelInfo(kTestModel.model_id).has_error()); } TEST_F(ModelsTestSuite, TestGenerateShortenedAlias) { - EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel).value()); + auto models1 = model_list_.LoadModelList(); auto alias = model_list_.GenerateShortenedAlias( - "huggingface.co/bartowski/llama3.1-7b-gguf/Model_ID_Xxx.gguf"); + "huggingface.co/bartowski/llama3.1-7b-gguf/Model_ID_Xxx.gguf", + models1.value()); EXPECT_EQ(alias, "model_id_xxx"); - EXPECT_TRUE(model_list_.UpdateModelAlias(kTestModel.model_id, alias)); + EXPECT_TRUE(model_list_.UpdateModelAlias(kTestModel.model_id, alias).value()); // Test with existing entries to force longer alias + auto models2 = model_list_.LoadModelList(); alias = model_list_.GenerateShortenedAlias( - "huggingface.co/bartowski/llama3.1-7b-gguf/Model_ID_Xxx.gguf"); + "huggingface.co/bartowski/llama3.1-7b-gguf/Model_ID_Xxx.gguf", + models2.value()); EXPECT_EQ(alias, "llama3.1-7b-gguf:model_id_xxx"); // Clean up - EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id).value()); } TEST_F(ModelsTestSuite, TestPersistence) { - EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel).value()); // Create a new ModelListUtils instance to test if it loads from file cortex::db::Models new_model_list(db_); auto retrieved_model = new_model_list.GetModelInfo(kTestModel.model_id); - - EXPECT_EQ(retrieved_model.model_id, kTestModel.model_id); - EXPECT_EQ(retrieved_model.author_repo_id, kTestModel.author_repo_id); - EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); + EXPECT_TRUE(retrieved_model); + EXPECT_EQ(retrieved_model.value().model_id, kTestModel.model_id); + EXPECT_EQ(retrieved_model.value().author_repo_id, kTestModel.author_repo_id); + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id).value()); } TEST_F(ModelsTestSuite, TestUpdateModelAlias) { @@ -115,44 +123,48 @@ TEST_F(ModelsTestSuite, TestUpdateModelAlias) { constexpr const auto kFinalTestAlias = "final_test_alias"; constexpr const auto kAnotherModelId = "another_model_id"; // Add the test model - ASSERT_TRUE(model_list_.AddModelEntry(kTestModel)); + ASSERT_TRUE(model_list_.AddModelEntry(kTestModel).value()); // Test successful update - EXPECT_TRUE(model_list_.UpdateModelAlias(kTestModel.model_id, kNewTestAlias)); + EXPECT_TRUE( + model_list_.UpdateModelAlias(kTestModel.model_id, kNewTestAlias).value()); auto updated_model = model_list_.GetModelInfo(kNewTestAlias); - EXPECT_EQ(updated_model.model_alias, kNewTestAlias); - EXPECT_EQ(updated_model.model_id, kTestModel.model_id); + EXPECT_TRUE(updated_model); + EXPECT_EQ(updated_model.value().model_alias, kNewTestAlias); + EXPECT_EQ(updated_model.value().model_id, kTestModel.model_id); // Test update with non-existent model - EXPECT_FALSE(model_list_.UpdateModelAlias(kNonExistentModel, kAnotherAlias)); + EXPECT_FALSE( + model_list_.UpdateModelAlias(kNonExistentModel, kAnotherAlias).value()); // Test update with non-unique alias cortex::db::ModelEntry another_model = kTestModel; another_model.model_id = kAnotherModelId; another_model.model_alias = kAnotherAlias; - ASSERT_TRUE(model_list_.AddModelEntry(another_model)); + ASSERT_TRUE(model_list_.AddModelEntry(another_model).value()); EXPECT_FALSE( - model_list_.UpdateModelAlias(kTestModel.model_id, kAnotherAlias)); + model_list_.UpdateModelAlias(kTestModel.model_id, kAnotherAlias).value()); // Test update using model alias instead of model ID EXPECT_TRUE(model_list_.UpdateModelAlias(kNewTestAlias, kFinalTestAlias)); updated_model = model_list_.GetModelInfo(kFinalTestAlias); - EXPECT_EQ(updated_model.model_alias, kFinalTestAlias); - EXPECT_EQ(updated_model.model_id, kTestModel.model_id); + EXPECT_TRUE(updated_model); + EXPECT_EQ(updated_model.value().model_alias, kFinalTestAlias); + EXPECT_EQ(updated_model.value().model_id, kTestModel.model_id); // Clean up - EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); - EXPECT_TRUE(model_list_.DeleteModelEntry(kAnotherModelId)); + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id).value()); + EXPECT_TRUE(model_list_.DeleteModelEntry(kAnotherModelId).value()); } TEST_F(ModelsTestSuite, TestHasModel) { - EXPECT_TRUE(model_list_.AddModelEntry(kTestModel)); + EXPECT_TRUE(model_list_.AddModelEntry(kTestModel).value()); EXPECT_TRUE(model_list_.HasModel(kTestModel.model_id)); EXPECT_TRUE(model_list_.HasModel("test_alias")); EXPECT_FALSE(model_list_.HasModel("non_existent_model")); // Clean up - EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id)); + EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model_id).value()); } } // namespace cortex::db \ No newline at end of file From e32963443238c772809ae7d7cef5507096104196 Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 12:35:29 +0700 Subject: [PATCH 11/15] fix: transaction --- engine/database/models.cc | 38 +++++++++++++++++++++++++------------- engine/database/models.h | 2 ++ 2 files changed, 27 insertions(+), 13 deletions(-) diff --git a/engine/database/models.cc b/engine/database/models.cc index aac233c6a..739674bbd 100644 --- a/engine/database/models.cc +++ b/engine/database/models.cc @@ -38,6 +38,29 @@ Models::~Models() {} cpp::result, std::string> Models::LoadModelList() const { + try { + db_.exec("BEGIN TRANSACTION;"); + utils::ScopeExit se([this] { db_.exec("COMMIT;"); }); + return LoadModelListNoLock(); + } catch (const std::exception& e) { + CTL_WRN(e.what()); + return cpp::fail(e.what()); + } +} + +bool Models::IsUnique(const std::vector& entries, + const std::string& model_id, + const std::string& model_alias) const { + return std::none_of( + entries.begin(), entries.end(), [&](const ModelEntry& entry) { + return entry.model_id == model_id || entry.model_alias == model_id || + entry.model_id == model_alias || + entry.model_alias == model_alias; + }); +} + +cpp::result, std::string> Models::LoadModelListNoLock() + const { try { std::vector entries; SQLite::Statement query( @@ -64,17 +87,6 @@ cpp::result, std::string> Models::LoadModelList() } } -bool Models::IsUnique(const std::vector& entries, - const std::string& model_id, - const std::string& model_alias) const { - return std::none_of( - entries.begin(), entries.end(), [&](const ModelEntry& entry) { - return entry.model_id == model_id || entry.model_alias == model_id || - entry.model_id == model_alias || - entry.model_alias == model_alias; - }); -} - std::string Models::GenerateShortenedAlias( const std::string& model_id, const std::vector& entries) const { std::vector parts; @@ -180,7 +192,7 @@ cpp::result Models::AddModelEntry(ModelEntry new_entry, try { db_.exec("BEGIN TRANSACTION;"); utils::ScopeExit se([this] { db_.exec("COMMIT;"); }); - auto model_list = LoadModelList(); + auto model_list = LoadModelListNoLock(); if (model_list.has_error()) { CTL_WRN(model_list.error()); std::cout << "Test: " << model_list.error(); @@ -243,7 +255,7 @@ cpp::result Models::UpdateModelAlias( try { db_.exec("BEGIN TRANSACTION;"); utils::ScopeExit se([this] { db_.exec("COMMIT;"); }); - auto model_list = LoadModelList(); + auto model_list = LoadModelListNoLock(); if (model_list.has_error()) { CTL_WRN(model_list.error()); return cpp::fail(model_list.error()); diff --git a/engine/database/models.h b/engine/database/models.h index fca823368..23c91b9f0 100644 --- a/engine/database/models.h +++ b/engine/database/models.h @@ -28,6 +28,8 @@ class Models { const std::string& model_id, const std::string& model_alias) const; + cpp::result, std::string> LoadModelListNoLock() const; + public: static const std::string kModelListPath; cpp::result, std::string> LoadModelList() const; From f0a15de08d6e4c791af0d2a21f0182330dffb13a Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 12:39:54 +0700 Subject: [PATCH 12/15] fix: format --- engine/commands/chat_cmd.cc | 4 ++-- engine/commands/model_get_cmd.cc | 4 ++-- engine/commands/model_status_cmd.cc | 4 ++-- engine/commands/model_upd_cmd.cc | 2 +- engine/commands/model_upd_cmd.h | 2 +- engine/services/model_service.cc | 17 ++++++++--------- 6 files changed, 16 insertions(+), 17 deletions(-) diff --git a/engine/commands/chat_cmd.cc b/engine/commands/chat_cmd.cc index b1e56773f..bb44b476b 100644 --- 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 "database/models.h" namespace commands { namespace { @@ -43,7 +43,7 @@ void ChatCmd::Exec(const std::string& host, int port, config::YamlHandler yaml_handler; try { auto model_entry = modellist_handler.GetModelInfo(model_handle); - if(model_entry.has_error()) { + if (model_entry.has_error()) { CLI_LOG("Error: " + model_entry.error()); return; } diff --git a/engine/commands/model_get_cmd.cc b/engine/commands/model_get_cmd.cc index df11f982d..5f6658cba 100644 --- a/engine/commands/model_get_cmd.cc +++ b/engine/commands/model_get_cmd.cc @@ -5,9 +5,9 @@ #include #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 "database/models.h" namespace commands { @@ -16,7 +16,7 @@ void ModelGetCmd::Exec(const std::string& model_handle) { config::YamlHandler yaml_handler; try { auto model_entry = modellist_handler.GetModelInfo(model_handle); - if(model_entry.has_error()) { + if (model_entry.has_error()) { CLI_LOG("Error: " + model_entry.error()); return; } diff --git a/engine/commands/model_status_cmd.cc b/engine/commands/model_status_cmd.cc index 12b939f13..38ff17bdc 100644 --- a/engine/commands/model_status_cmd.cc +++ b/engine/commands/model_status_cmd.cc @@ -1,9 +1,9 @@ #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 "database/models.h" namespace commands { bool ModelStatusCmd::IsLoaded(const std::string& host, int port, @@ -12,7 +12,7 @@ bool ModelStatusCmd::IsLoaded(const std::string& host, int port, config::YamlHandler yaml_handler; try { auto model_entry = modellist_handler.GetModelInfo(model_handle); - if(model_entry.has_error()) { + if (model_entry.has_error()) { CLI_LOG("Error: " + model_entry.error()); return false; } diff --git a/engine/commands/model_upd_cmd.cc b/engine/commands/model_upd_cmd.cc index e2c52ab46..ea10d2d95 100644 --- a/engine/commands/model_upd_cmd.cc +++ b/engine/commands/model_upd_cmd.cc @@ -11,7 +11,7 @@ void ModelUpdCmd::Exec( const std::unordered_map& options) { try { auto model_entry = model_list_utils_.GetModelInfo(model_handle_); - if(model_entry.has_error()) { + if (model_entry.has_error()) { CLI_LOG("Error: " + model_entry.error()); return; } diff --git a/engine/commands/model_upd_cmd.h b/engine/commands/model_upd_cmd.h index c2bffc81a..49c104157 100644 --- a/engine/commands/model_upd_cmd.h +++ b/engine/commands/model_upd_cmd.h @@ -5,8 +5,8 @@ #include #include #include "config/model_config.h" -#include "database/models.h" #include "config/yaml_config.h" +#include "database/models.h" namespace commands { class ModelUpdCmd { public: diff --git a/engine/services/model_service.cc b/engine/services/model_service.cc index 4612c0948..22638edd8 100644 --- a/engine/services/model_service.cc +++ b/engine/services/model_service.cc @@ -4,11 +4,11 @@ #include #include "config/gguf_parser.h" #include "config/yaml_config.h" +#include "database/models.h" #include "utils/cli_selection_utils.h" #include "utils/file_manager_utils.h" #include "utils/huggingface_utils.h" #include "utils/logging_utils.h" -#include "database/models.h" #include "utils/result.hpp" #include "utils/string_utils.h" @@ -38,13 +38,12 @@ void ParseGguf(const DownloadItem& ggufDownloadItem, auto author_id = author.has_value() ? author.value() : "cortexso"; cortex::db::Models modellist_utils_obj; - cortex::db::ModelEntry model_entry{ - .model_id = ggufDownloadItem.id, - .author_repo_id = author_id, - .branch_name = branch, - .path_to_model_yaml = yaml_name.string(), - .model_alias = ggufDownloadItem.id, - .status = cortex::db::ModelStatus::READY}; + cortex::db::ModelEntry model_entry{.model_id = ggufDownloadItem.id, + .author_repo_id = author_id, + .branch_name = branch, + .path_to_model_yaml = yaml_name.string(), + .model_alias = ggufDownloadItem.id, + .status = cortex::db::ModelStatus::READY}; modellist_utils_obj.AddModelEntry(model_entry, true); } @@ -342,7 +341,7 @@ cpp::result ModelService::DeleteModel( try { auto model_entry = modellist_handler.GetModelInfo(model_handle); - if(model_entry.has_error()) { + if (model_entry.has_error()) { CLI_LOG("Error: " + model_entry.error()); return cpp::fail(model_entry.error()); } From b95666b6c05520be24512d5ac76e26b05db8daf0 Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 13:10:19 +0700 Subject: [PATCH 13/15] fix: make alias unique --- engine/database/models.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/engine/database/models.cc b/engine/database/models.cc index 739674bbd..4eaceb11b 100644 --- a/engine/database/models.cc +++ b/engine/database/models.cc @@ -30,7 +30,7 @@ Models::Models(SQLite::Database& db) : db_(db) { "author_repo_id TEXT," "branch_name TEXT," "path_to_model_yaml TEXT," - "model_alias TEXT," + "model_alias TEXT UNIQUE," "status TEXT);"); } From a1d52ba59b22b48bdef185921879a8c20af8a3c1 Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 18:47:46 +0700 Subject: [PATCH 14/15] fix: remove ModelStatus --- engine/commands/model_import_cmd.cc | 2 +- engine/controllers/models.cc | 2 +- engine/database/models.cc | 21 ++++----------------- engine/database/models.h | 4 ---- engine/services/model_service.cc | 6 ++---- engine/test/components/test_models_db.cc | 5 +---- 6 files changed, 9 insertions(+), 31 deletions(-) diff --git a/engine/commands/model_import_cmd.cc b/engine/commands/model_import_cmd.cc index 688753511..c12af5054 100644 --- a/engine/commands/model_import_cmd.cc +++ b/engine/commands/model_import_cmd.cc @@ -24,7 +24,7 @@ void ModelImportCmd::Exec() { .string(); cortex::db::ModelEntry model_entry{ model_handle_, "local", "imported", - model_yaml_path, model_handle_, cortex::db::ModelStatus::READY}; + model_yaml_path, model_handle_}; try { std::filesystem::create_directories( std::filesystem::path(model_yaml_path).parent_path()); diff --git a/engine/controllers/models.cc b/engine/controllers/models.cc index af7a9fd85..6ac0c1664 100644 --- a/engine/controllers/models.cc +++ b/engine/controllers/models.cc @@ -228,7 +228,7 @@ void Models::ImportModel( .string(); cortex::db::ModelEntry model_entry{ modelHandle, "local", "imported", - model_yaml_path, modelHandle, cortex::db::ModelStatus::READY}; + model_yaml_path, modelHandle}; try { std::filesystem::create_directories( std::filesystem::path(model_yaml_path).parent_path()); diff --git a/engine/database/models.cc b/engine/database/models.cc index 4eaceb11b..df78da115 100644 --- a/engine/database/models.cc +++ b/engine/database/models.cc @@ -19,8 +19,7 @@ Models::Models() : db_(cortex::db::Database::GetInstance().db()) { "author_repo_id TEXT," "branch_name TEXT," "path_to_model_yaml TEXT," - "model_alias TEXT," - "status TEXT);"); + "model_alias TEXT);"); } Models::Models(SQLite::Database& db) : db_(db) { @@ -30,8 +29,7 @@ Models::Models(SQLite::Database& db) : db_(db) { "author_repo_id TEXT," "branch_name TEXT," "path_to_model_yaml TEXT," - "model_alias TEXT UNIQUE," - "status TEXT);"); + "model_alias TEXT UNIQUE);"); } Models::~Models() {} @@ -76,8 +74,6 @@ cpp::result, std::string> Models::LoadModelListNoLock() entry.path_to_model_yaml = query.getColumn(3).getString(); entry.model_alias = query.getColumn(4).getString(); std::string status_str = query.getColumn(5).getString(); - entry.status = - (status_str == "RUNNING") ? ModelStatus::RUNNING : ModelStatus::READY; entries.push_back(entry); } return entries; @@ -166,8 +162,6 @@ cpp::result Models::GetModelInfo( entry.path_to_model_yaml = query.getColumn(3).getString(); entry.model_alias = query.getColumn(4).getString(); std::string status_str = query.getColumn(5).getString(); - entry.status = - (status_str == "RUNNING") ? ModelStatus::RUNNING : ModelStatus::READY; return entry; } else { return cpp::fail("Model not found: " + identifier); @@ -183,8 +177,6 @@ void Models::PrintModelInfo(const ModelEntry& entry) const { LOG_INFO << "Branch Name: " << entry.branch_name; LOG_INFO << "Path to model.yaml: " << entry.path_to_model_yaml; LOG_INFO << "Model Alias: " << entry.model_alias; - LOG_INFO << "Status: " - << (entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY"); } cpp::result Models::AddModelEntry(ModelEntry new_entry, @@ -203,8 +195,7 @@ cpp::result Models::AddModelEntry(ModelEntry new_entry, if (use_short_alias) { new_entry.model_alias = GenerateShortenedAlias(new_entry.model_id, model_list.value()); - } - new_entry.status = ModelStatus::READY; // Set default status to READY + } SQLite::Statement insert( db_, @@ -216,8 +207,6 @@ cpp::result Models::AddModelEntry(ModelEntry new_entry, insert.bind(3, new_entry.branch_name); insert.bind(4, new_entry.path_to_model_yaml); insert.bind(5, new_entry.model_alias); - insert.bind( - 6, (new_entry.status == ModelStatus::RUNNING ? "RUNNING" : "READY")); insert.exec(); return true; @@ -240,10 +229,8 @@ cpp::result Models::UpdateModelEntry( upd.bind(1, updated_entry.author_repo_id); upd.bind(2, updated_entry.branch_name); upd.bind(3, updated_entry.path_to_model_yaml); - upd.bind(4, (updated_entry.status == ModelStatus::RUNNING ? "RUNNING" - : "READY")); + upd.bind(4, identifier); upd.bind(5, identifier); - upd.bind(6, identifier); return upd.exec() == 1; } catch (const std::exception& e) { return cpp::fail(e.what()); diff --git a/engine/database/models.h b/engine/database/models.h index 23c91b9f0..184f1c6a6 100644 --- a/engine/database/models.h +++ b/engine/database/models.h @@ -7,16 +7,12 @@ #include "utils/result.hpp" namespace cortex::db { - -enum class ModelStatus { READY, RUNNING }; - struct ModelEntry { std::string model_id; std::string author_repo_id; std::string branch_name; std::string path_to_model_yaml; std::string model_alias; - ModelStatus status; }; class Models { diff --git a/engine/services/model_service.cc b/engine/services/model_service.cc index 22638edd8..0e25e71ab 100644 --- a/engine/services/model_service.cc +++ b/engine/services/model_service.cc @@ -42,8 +42,7 @@ void ParseGguf(const DownloadItem& ggufDownloadItem, .author_repo_id = author_id, .branch_name = branch, .path_to_model_yaml = yaml_name.string(), - .model_alias = ggufDownloadItem.id, - .status = cortex::db::ModelStatus::READY}; + .model_alias = ggufDownloadItem.id}; modellist_utils_obj.AddModelEntry(model_entry, true); } @@ -283,8 +282,7 @@ cpp::result ModelService::DownloadModelFromCortexso( .author_repo_id = "cortexso", .branch_name = branch, .path_to_model_yaml = model_yml_item->localPath.string(), - .model_alias = model_id, - .status = cortex::db::ModelStatus::READY}; + .model_alias = model_id}; modellist_utils_obj.AddModelEntry(model_entry); }; diff --git a/engine/test/components/test_models_db.cc b/engine/test/components/test_models_db.cc index 67f1bd0be..ee418d851 100644 --- a/engine/test/components/test_models_db.cc +++ b/engine/test/components/test_models_db.cc @@ -25,7 +25,7 @@ class ModelsTestSuite : public ::testing::Test { const cortex::db::ModelEntry kTestModel{ "test_model_id", "test_author", "main", - "/path/to/model.yaml", "test_alias", cortex::db::ModelStatus::READY}; + "/path/to/model.yaml", "test_alias"}; }; TEST_F(ModelsTestSuite, TestAddModelEntry) { @@ -61,15 +61,12 @@ TEST_F(ModelsTestSuite, TestUpdateModelEntry) { EXPECT_TRUE(model_list_.AddModelEntry(kTestModel).value()); cortex::db::ModelEntry updated_model = kTestModel; - updated_model.status = cortex::db::ModelStatus::RUNNING; EXPECT_TRUE( model_list_.UpdateModelEntry(kTestModel.model_id, updated_model).value()); auto retrieved_model = model_list_.GetModelInfo(kTestModel.model_id); EXPECT_TRUE(retrieved_model); - EXPECT_EQ(retrieved_model.value().status, cortex::db::ModelStatus::RUNNING); - updated_model.status = cortex::db::ModelStatus::READY; EXPECT_TRUE( model_list_.UpdateModelEntry(kTestModel.model_id, updated_model).value()); From b016a9e125e694688a79ac028d7422efefdd3299 Mon Sep 17 00:00:00 2001 From: vansangpfiev Date: Thu, 26 Sep 2024 18:58:57 +0700 Subject: [PATCH 15/15] fix: models --- engine/database/models.cc | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/engine/database/models.cc b/engine/database/models.cc index df78da115..cfaf275e7 100644 --- a/engine/database/models.cc +++ b/engine/database/models.cc @@ -61,10 +61,9 @@ cpp::result, std::string> Models::LoadModelListNoLock() const { try { std::vector entries; - SQLite::Statement query( - db_, - "SELECT model_id, author_repo_id, branch_name, " - "path_to_model_yaml, model_alias, status FROM models"); + SQLite::Statement query(db_, + "SELECT model_id, author_repo_id, branch_name, " + "path_to_model_yaml, model_alias FROM models"); while (query.executeStep()) { ModelEntry entry; @@ -73,7 +72,6 @@ cpp::result, std::string> Models::LoadModelListNoLock() entry.branch_name = query.getColumn(2).getString(); entry.path_to_model_yaml = query.getColumn(3).getString(); entry.model_alias = query.getColumn(4).getString(); - std::string status_str = query.getColumn(5).getString(); entries.push_back(entry); } return entries; @@ -146,11 +144,10 @@ std::string Models::GenerateShortenedAlias( cpp::result Models::GetModelInfo( const std::string& identifier) const { try { - SQLite::Statement query( - db_, - "SELECT model_id, author_repo_id, branch_name, " - "path_to_model_yaml, model_alias, status FROM models " - "WHERE model_id = ? OR model_alias = ?"); + SQLite::Statement query(db_, + "SELECT model_id, author_repo_id, branch_name, " + "path_to_model_yaml, model_alias FROM models " + "WHERE model_id = ? OR model_alias = ?"); query.bind(1, identifier); query.bind(2, identifier); @@ -161,7 +158,6 @@ cpp::result Models::GetModelInfo( entry.branch_name = query.getColumn(2).getString(); entry.path_to_model_yaml = query.getColumn(3).getString(); entry.model_alias = query.getColumn(4).getString(); - std::string status_str = query.getColumn(5).getString(); return entry; } else { return cpp::fail("Model not found: " + identifier); @@ -195,13 +191,13 @@ cpp::result Models::AddModelEntry(ModelEntry new_entry, if (use_short_alias) { new_entry.model_alias = GenerateShortenedAlias(new_entry.model_id, model_list.value()); - } + } SQLite::Statement insert( db_, "INSERT INTO models (model_id, author_repo_id, " - "branch_name, path_to_model_yaml, model_alias, status) VALUES (?, ?, " - "?, ?, ?, ?)"); + "branch_name, path_to_model_yaml, model_alias) VALUES (?, ?, " + "?, ?, ?)"); insert.bind(1, new_entry.model_id); insert.bind(2, new_entry.author_repo_id); insert.bind(3, new_entry.branch_name); @@ -224,7 +220,7 @@ cpp::result Models::UpdateModelEntry( SQLite::Statement upd(db_, "UPDATE models " "SET author_repo_id = ?, branch_name = ?, " - "path_to_model_yaml = ?, status = ? " + "path_to_model_yaml = ? " "WHERE model_id = ? OR model_alias = ?"); upd.bind(1, updated_entry.author_repo_id); upd.bind(2, updated_entry.branch_name);