feat: add model status and information dump

janhq · Nov 13, 2023 · c5e97fe · c5e97fe
1 parent f49bd78
commit c5e97fe
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/controllers/llamaCPP.cc b/controllers/llamaCPP.cc
@@ -6,6 +6,7 @@
 #include <drogon/HttpResponse.h>
 #include <drogon/HttpTypes.h>
 #include <regex>
+#include <string>
 #include <thread>
 #include <trantor/utils/Logger.h>
 
@@ -214,6 +215,17 @@ void llamaCPP::unloadModel(
   callback(resp);
   return;
 }
+void llamaCPP::modelStatus(
+    const HttpRequestPtr &req,
+    std::function<void(const HttpResponsePtr &)> &&callback) {
+  Json::Value jsonResp;
+  jsonResp["model_loaded"] = this->model_loaded.load();
+  jsonResp["model_data"] = llama.get_model_props().dump();
+
+  auto resp = nitro_utils::nitroHttpJsonResponse(jsonResp);
+  callback(resp);
+  return;
+}
 
 void llamaCPP::loadModel(
     const HttpRequestPtr &req,

diff --git a/controllers/llamaCPP.h b/controllers/llamaCPP.h
@@ -2125,6 +2125,7 @@ class llamaCPP : public drogon::HttpController<llamaCPP> {
   METHOD_ADD(llamaCPP::embedding, "embedding", Post);
   METHOD_ADD(llamaCPP::loadModel, "loadmodel", Post);
   METHOD_ADD(llamaCPP::unloadModel, "unloadmodel", Get);
+  METHOD_ADD(llamaCPP::modelStatus, "modelstatus", Get);
 
   // PATH_ADD("/llama/chat_completion", Post);
   METHOD_LIST_END
@@ -2136,6 +2137,10 @@ class llamaCPP : public drogon::HttpController<llamaCPP> {
                  std::function<void(const HttpResponsePtr &)> &&callback);
   void unloadModel(const HttpRequestPtr &req,
                    std::function<void(const HttpResponsePtr &)> &&callback);
+
+  void modelStatus(const HttpRequestPtr &req,
+                   std::function<void(const HttpResponsePtr &)> &&callback);
+
   void warmupModel();
 
   void backgroundTask();