diff --git a/engine/commands/model_alias_cmd.cc b/engine/commands/model_alias_cmd.cc new file mode 100644 index 000000000..2123d06cf --- /dev/null +++ b/engine/commands/model_alias_cmd.cc @@ -0,0 +1,23 @@ +#include "model_alias_cmd.h" +#include "utils/modellist_utils.h" + +namespace commands { + +void ModelAliasCmd::Exec(const std::string& model_handle, + const std::string& model_alias) { + modellist_utils::ModelListUtils modellist_handler; + try { + if (modellist_handler.UpdateModelAlias(model_handle, model_alias)) { + CLI_LOG("Successfully set model alias '" + model_alias + + "' for modeID '" + model_handle + "'."); + } else { + CLI_LOG("Unable to set model alias for modelID '" + model_handle + + "': model alias '" + model_alias + "' is not unique!"); + } + } catch (const std::exception& e) { + CLI_LOG("Error when setting model alias ('" + model_alias + + "') for modelID '" + model_handle + "':" + e.what()); + } +} + +} // namespace commands \ No newline at end of file diff --git a/engine/commands/model_alias_cmd.h b/engine/commands/model_alias_cmd.h new file mode 100644 index 000000000..dcec44947 --- /dev/null +++ b/engine/commands/model_alias_cmd.h @@ -0,0 +1,10 @@ +#pragma once +#include +#include "utils/logging_utils.h" +namespace commands { + +class ModelAliasCmd { + public: + void Exec(const std::string& model_handle, const std::string& model_alias); +}; +} // namespace commands \ No newline at end of file diff --git a/engine/controllers/command_line_parser.cc b/engine/controllers/command_line_parser.cc index 0d74a749b..6357a11ff 100644 --- a/engine/controllers/command_line_parser.cc +++ b/engine/controllers/command_line_parser.cc @@ -6,6 +6,7 @@ #include "commands/engine_install_cmd.h" #include "commands/engine_list_cmd.h" #include "commands/engine_uninstall_cmd.h" +#include "commands/model_alias_cmd.h" #include "commands/model_del_cmd.h" #include "commands/model_get_cmd.h" #include "commands/model_list_cmd.h" @@ -150,6 +151,18 @@ bool CommandLineParser::SetupCommand(int argc, char** argv) { mdc.Exec(model_id); }); + std::string model_alias; + auto model_alias_cmd = + models_cmd->add_subcommand("alias", "Add alias name for a modelID"); + model_alias_cmd->add_option("--model_id", model_id, "Can be modelID or model alias to identify model"); + model_alias_cmd->require_option(); + model_alias_cmd->add_option("--alias", model_alias, "new alias to be set"); + model_alias_cmd->require_option(); + model_alias_cmd->callback([&model_id, &model_alias]() { + commands::ModelAliasCmd mdc; + mdc.Exec(model_id, model_alias); + }); + auto model_update_cmd = models_cmd->add_subcommand("update", "Update configuration of a model"); diff --git a/engine/controllers/models.cc b/engine/controllers/models.cc index 1d3157fcb..45a5bf60e 100644 --- a/engine/controllers/models.cc +++ b/engine/controllers/models.cc @@ -5,11 +5,11 @@ #include "utils/cortex_utils.h" #include "utils/file_manager_utils.h" #include "utils/model_callback_utils.h" +#include "utils/modellist_utils.h" - void - Models::PullModel( - const HttpRequestPtr& req, - std::function&& callback) const { +void Models::PullModel( + const HttpRequestPtr& req, + std::function&& callback) const { if (!http_util::HasFieldInReq(req, callback, "modelId")) { return; } @@ -192,4 +192,56 @@ void Models::DeleteModel(const HttpRequestPtr& req, resp->setStatusCode(k404NotFound); callback(resp); } +} + +void Models::SetModelAlias( + const HttpRequestPtr& req, + std::function&& callback) const { + if (!http_util::HasFieldInReq(req, callback, "modelId") || + !http_util::HasFieldInReq(req, callback, "modelAlias")) { + return; + } + auto model_handle = (*(req->getJsonObject())).get("modelId", "").asString(); + auto model_alias = (*(req->getJsonObject())).get("modelAlias", "").asString(); + LOG_DEBUG << "GetModel, Model handle: " << model_handle + << ", Model alias: " << model_alias; + + modellist_utils::ModelListUtils modellist_handler; + try { + if (modellist_handler.UpdateModelAlias(model_handle, model_alias)) { + std::string message = "Successfully set model alias '" + model_alias + + "' for modeID '" + model_handle + "'."; + LOG_INFO << message; + Json::Value ret; + ret["result"] = "OK"; + ret["modelHandle"] = model_handle; + ret["message"] = message; + auto resp = cortex_utils::CreateCortexHttpJsonResponse(ret); + resp->setStatusCode(k200OK); + callback(resp); + } else { + std::string message = "Unable to set model alias for modelID '" + + model_handle + "': model alias '" + model_alias + + "' is not unique!"; + LOG_ERROR << message; + Json::Value ret; + ret["result"] = "Set alias failed!"; + ret["modelHandle"] = model_handle; + ret["message"] = message; + auto resp = cortex_utils::CreateCortexHttpJsonResponse(ret); + resp->setStatusCode(k400BadRequest); + callback(resp); + } + } catch (const std::exception& e) { + std::string message = "Error when setting model alias ('" + model_alias + + "') for modelID '" + model_handle + "':" + e.what(); + LOG_ERROR << message; + Json::Value ret; + ret["result"] = "Set alias failed!"; + ret["modelHandle"] = model_handle; + ret["message"] = message; + auto resp = cortex_utils::CreateCortexHttpJsonResponse(ret); + resp->setStatusCode(k400BadRequest); + callback(resp); + } } \ No newline at end of file diff --git a/engine/controllers/models.h b/engine/controllers/models.h index 907ab3917..65a3641f3 100644 --- a/engine/controllers/models.h +++ b/engine/controllers/models.h @@ -16,6 +16,7 @@ class Models : public drogon::HttpController { METHOD_ADD(Models::ListModel, "/list", Get); METHOD_ADD(Models::GetModel, "/get", Post); METHOD_ADD(Models::DeleteModel, "/{1}", Delete); + METHOD_ADD(Models::SetModelAlias, "/alias", Post); METHOD_LIST_END void PullModel(const HttpRequestPtr& req, @@ -27,4 +28,7 @@ class Models : public drogon::HttpController { void DeleteModel(const HttpRequestPtr& req, std::function&& callback, const std::string& model_id) const; + void SetModelAlias( + const HttpRequestPtr& req, + std::function&& callback) const; }; \ No newline at end of file diff --git a/engine/e2e-test/test_cli_engine_install.py b/engine/e2e-test/test_cli_engine_install.py index 89d49401d..dfb4e9599 100644 --- a/engine/e2e-test/test_cli_engine_install.py +++ b/engine/e2e-test/test_cli_engine_install.py @@ -8,7 +8,7 @@ class TestCliEngineInstall: def test_engines_install_llamacpp_should_be_successfully(self): exit_code, output, error = run( - "Install Engine", ["engines", "install", "cortex.llamacpp"], timeout=60 + "Install Engine", ["engines", "install", "cortex.llamacpp"], timeout=600 ) assert "Start downloading" in output, "Should display downloading message" assert exit_code == 0, f"Install engine failed with error: {error}" @@ -31,7 +31,7 @@ def test_engines_install_onnx_on_tensorrt_should_be_failed(self): def test_engines_install_pre_release_llamacpp(self): exit_code, output, error = run( - "Install Engine", ["engines", "install", "cortex.llamacpp", "-v", "v0.1.29"], timeout=None + "Install Engine", ["engines", "install", "cortex.llamacpp", "-v", "v0.1.29"], timeout=600 ) assert "Start downloading" in output, "Should display downloading message" assert exit_code == 0, f"Install engine failed with error: {error}" diff --git a/engine/e2e-test/test_cli_engine_uninstall.py b/engine/e2e-test/test_cli_engine_uninstall.py index 525c0ad63..685e5387f 100644 --- a/engine/e2e-test/test_cli_engine_uninstall.py +++ b/engine/e2e-test/test_cli_engine_uninstall.py @@ -8,7 +8,7 @@ class TestCliEngineUninstall: def setup_and_teardown(self): # Setup # Preinstall llamacpp engine - run("Install Engine", ["engines", "install", "cortex.llamacpp"],timeout=None) + run("Install Engine", ["engines", "install", "cortex.llamacpp"],timeout = None) yield diff --git a/engine/e2e-test/test_cli_model_delete.py b/engine/e2e-test/test_cli_model_delete.py index eda21d1ce..8bbe95b35 100644 --- a/engine/e2e-test/test_cli_model_delete.py +++ b/engine/e2e-test/test_cli_model_delete.py @@ -8,6 +8,7 @@ class TestCliModelDelete: def setup_and_teardown(self): # Setup # Pull model + stdout, stderr, return_code = popen(["pull", "tinyllama"], "1\n") yield diff --git a/engine/test/components/test_modellist_utils.cc b/engine/test/components/test_modellist_utils.cc index cb46b2b70..2a7abc05a 100644 --- a/engine/test/components/test_modellist_utils.cc +++ b/engine/test/components/test_modellist_utils.cc @@ -88,4 +88,36 @@ TEST_F(ModelListUtilsTestSuite, TestPersistence) { 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"); +} + +TEST_F(ModelListUtilsTestSuite, TestUpdateModelAlias) { + // 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"); + + // Test update with non-existent model + EXPECT_FALSE(model_list_.UpdateModelAlias("non_existent_model", "another_alias")); + + // 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"; + ASSERT_TRUE(model_list_.AddModelEntry(another_model)); + + EXPECT_FALSE(model_list_.UpdateModelAlias("test_model_id", "another_alias")); + + // 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"); + + // Clean up + model_list_.DeleteModelEntry("test_model_id"); + model_list_.DeleteModelEntry("another_model_id"); } \ No newline at end of file diff --git a/engine/utils/modellist_utils.cc b/engine/utils/modellist_utils.cc index a9527cbc1..261bf58d5 100644 --- a/engine/utils/modellist_utils.cc +++ b/engine/utils/modellist_utils.cc @@ -198,6 +198,28 @@ bool ModelListUtils::UpdateModelEntry(const std::string& identifier, return false; // Entry not found } +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 +} + bool ModelListUtils::DeleteModelEntry(const std::string& identifier) { std::lock_guard lock(mutex_); auto entries = LoadModelList(); diff --git a/engine/utils/modellist_utils.h b/engine/utils/modellist_utils.h index e8efab0d2..7625c264b 100644 --- a/engine/utils/modellist_utils.h +++ b/engine/utils/modellist_utils.h @@ -40,5 +40,6 @@ class ModelListUtils { bool UpdateModelEntry(const std::string& identifier, const ModelEntry& updated_entry); bool DeleteModelEntry(const std::string& identifier); + bool UpdateModelAlias(const std::string& model_id, const std::string& model_alias); }; } // namespace modellist_utils \ No newline at end of file