diff --git a/docs/static/openapi/jan.json b/docs/static/openapi/jan.json index 352c8a5b2..8e05cf597 100644 --- a/docs/static/openapi/jan.json +++ b/docs/static/openapi/jan.json @@ -196,7 +196,66 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/PullModelResponse" + "type": "object", + "properties": { + "message": { + "type": "string" + }, + "task": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "items": { + "type": "array", + "items": { + "type": "object", + "properties": { + "bytes": { + "type": "integer" + }, + "checksum": { + "type": "string" + }, + "downloadUrl": { + "type": "string" + }, + "downloadedBytes": { + "type": "integer" + }, + "id": { + "type": "string" + }, + "localPath": { + "type": "string" + } + } + } + }, + "type": { + "type": "string" + } + } + } + } + }, + "example": { + "message": "Model start downloading!", + "task": { + "id": "TheBloke:Mistral-7B-Instruct-v0.1-GGUF:mistral-7b-instruct-v0.1.Q3_K_L.gguf", + "items": [ + { + "bytes": 3822024352, + "checksum": "N/A", + "downloadUrl": "https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.1-GGUF/resolve/main/mistral-7b-instruct-v0.1.Q3_K_L.gguf", + "downloadedBytes": 0, + "id": "TheBloke:Mistral-7B-Instruct-v0.1-GGUF:mistral-7b-instruct-v0.1.Q3_K_L.gguf", + "localPath": "/Users/jamesnguyen/cortexcpp/models/huggingface.co/TheBloke/Mistral-7B-Instruct-v0.1-GGUF/mistral-7b-instruct-v0.1.Q3_K_L.gguf" + } + ], + "type": "Model" + } } } } @@ -213,6 +272,99 @@ } }, "tags": ["Models"] + }, + "delete": { + "tags": ["Models"], + "summary": "Stop model download", + "description": "Stops the download of a model with the corresponding taskId provided in the request body", + "operationId": "ModelsController_stopModelDownload", + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "taskId": { + "type": "string", + "description": "The unique identifier of the download task to be stopped" + } + }, + "required": ["taskId"] + } + } + } + }, + "responses": { + "200": { + "description": "Download stopped successfully", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "message": { + "type": "string", + "example": "Download stopped successfully" + }, + "taskId": { + "type": "string", + "example": "task-123456" + } + } + } + } + } + }, + "400": { + "description": "Bad request", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "error": { + "type": "string", + "example": "Invalid taskId" + } + } + } + } + } + }, + "404": { + "description": "Task not found", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "error": { + "type": "string", + "example": "Download task not found" + } + } + } + } + } + }, + "500": { + "description": "Internal server error", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "error": { + "type": "string", + "example": "An unexpected error occurred" + } + } + } + } + } + } + } } }, "/v1/models": { diff --git a/engine/controllers/models.cc b/engine/controllers/models.cc index 0d5c19e66..7e66099dd 100644 --- a/engine/controllers/models.cc +++ b/engine/controllers/models.cc @@ -82,7 +82,8 @@ void Models::AbortPullModel( callback(resp); } else { Json::Value ret; - ret["message"] = "Task stopped!"; + ret["message"] = "Download stopped successfully"; + ret["taskId"] = result.value(); auto resp = cortex_utils::CreateCortexHttpJsonResponse(ret); resp->setStatusCode(k200OK); callback(resp); diff --git a/engine/controllers/models.h b/engine/controllers/models.h index 59bb6a51b..4482ebcbd 100644 --- a/engine/controllers/models.h +++ b/engine/controllers/models.h @@ -22,7 +22,7 @@ class Models : public drogon::HttpController { METHOD_ADD(Models::GetModelStatus, "/status/{1}", Get); ADD_METHOD_TO(Models::PullModel, "/v1/models/pull", Post); - ADD_METHOD_TO(Models::PullModel, "/v1/models/pull", Delete); + ADD_METHOD_TO(Models::AbortPullModel, "/v1/models/pull", Delete); ADD_METHOD_TO(Models::ListModel, "/v1/models", Get); ADD_METHOD_TO(Models::GetModel, "/v1/models/{1}", Get); ADD_METHOD_TO(Models::UpdateModel, "/v1/models/{1}", Patch); diff --git a/engine/services/download_service.cc b/engine/services/download_service.cc index 3966d8a98..5df2aaa96 100644 --- a/engine/services/download_service.cc +++ b/engine/services/download_service.cc @@ -306,6 +306,8 @@ void DownloadService::ProcessTask(DownloadTask& task) { downloading_data_.reset(); } + RemoveTaskFromStopList(task.id); + // if terminate by API calling and not from process stopping, we emit // DownloadStopped event if (is_terminated) { @@ -314,7 +316,6 @@ void DownloadService::ProcessTask(DownloadTask& task) { DownloadEvent{.type_ = DownloadEventType::DownloadStopped, .download_task_ = task}); } else { - RemoveTaskFromStopList(task.id); CTL_INF("Executing callback.."); ExecuteCallback(task); @@ -325,13 +326,13 @@ void DownloadService::ProcessTask(DownloadTask& task) { } } -cpp::result DownloadService::StopTask( +cpp::result DownloadService::StopTask( const std::string& task_id) { std::lock_guard lock(stop_mutex_); tasks_to_stop_.insert(task_id); CTL_INF("Added task to stop list: " << task_id); - return {}; + return task_id; } void DownloadService::ProcessCompletedTransfers() { diff --git a/engine/services/download_service.h b/engine/services/download_service.h index af46e2d47..9bf15efdd 100644 --- a/engine/services/download_service.h +++ b/engine/services/download_service.h @@ -75,7 +75,7 @@ class DownloadService { cpp::result GetFileSize( const std::string& url) const noexcept; - cpp::result StopTask(const std::string& task_id); + cpp::result StopTask(const std::string& task_id); private: struct DownloadingData { diff --git a/engine/services/model_service.cc b/engine/services/model_service.cc index bf6f5e9af..62cc92e5b 100644 --- a/engine/services/model_service.cc +++ b/engine/services/model_service.cc @@ -718,8 +718,7 @@ cpp::result ModelService::GetModelStatus( } } -cpp::result ModelService::AbortDownloadModel( +cpp::result ModelService::AbortDownloadModel( const std::string& task_id) { - auto result = download_service_->StopTask(task_id); - return result; + return download_service_->StopTask(task_id); } diff --git a/engine/services/model_service.h b/engine/services/model_service.h index c791c7fd7..43e597c40 100644 --- a/engine/services/model_service.h +++ b/engine/services/model_service.h @@ -17,7 +17,8 @@ class ModelService { */ cpp::result DownloadModel(const std::string& input); - cpp::result AbortDownloadModel(const std::string& task_id); + cpp::result AbortDownloadModel( + const std::string& task_id); cpp::result DownloadModelFromCortexso( const std::string& name, const std::string& branch = "main");