fix: add filter by tag for repository list (#2021)

Co-authored-by: sangjanai <sang@jan.ai>

Author: vansangpfiev

engine/controllers/models.cc

@@ -859,10 +859,11 @@ void Models::GetModelSource(
 void Models::GetRepositoryList(
     const HttpRequestPtr& req,
     std::function<void(const HttpResponsePtr&)>&& callback,
-    std::optional<std::string> author) {
+    std::optional<std::string> author, std::optional<std::string> tag) {
   if (!author.has_value())
     author = "cortexso";
-  auto res = model_src_svc_->GetRepositoryList(author.value());
+  auto res =
+      model_src_svc_->GetRepositoryList(author.value(), tag.value_or(""));
   if (res.has_error()) {
     Json::Value ret;
     ret["message"] = res.error();

engine/controllers/models.h

@@ -44,8 +44,8 @@ class Models : public drogon::HttpController<Models, false> {
   ADD_METHOD_TO(Models::DeleteModelSource, "/v1/models/sources", Delete);
   ADD_METHOD_TO(Models::GetModelSources, "/v1/models/sources", Get);
   ADD_METHOD_TO(Models::GetModelSource, "/v1/models/sources/{src}", Get);
-  ADD_METHOD_TO(Models::GetRepositoryList, "/v1/models/hub?author={author}",
-                Get);
+  ADD_METHOD_TO(Models::GetRepositoryList,
+                "/v1/models/hub?author={author}&tag={tag}", Get);
   METHOD_LIST_END
 
   explicit Models(std::shared_ptr<DatabaseService> db_service,
@@ -115,7 +115,8 @@ class Models : public drogon::HttpController<Models, false> {
 
   void GetRepositoryList(const HttpRequestPtr& req,
                          std::function<void(const HttpResponsePtr&)>&& callback,
-                         std::optional<std::string> author);
+                         std::optional<std::string> author,
+                         std::optional<std::string> tag);
 
  private:
   std::shared_ptr<DatabaseService> db_service_;

engine/services/model_source_service.cc

@@ -14,17 +14,6 @@
 namespace hu = huggingface_utils;
 
 namespace {
-struct ModelInfo {
-  std::string id;
-  int likes;
-  int trending_score;
-  bool is_private;
-  int downloads;
-  std::vector<std::string> tags;
-  std::string created_at;
-  std::string model_id;
-};
-
 std::vector<ModelInfo> ParseJsonString(const std::string& json_str) {
   std::vector<ModelInfo> models;
 
@@ -201,20 +190,37 @@ cpp::result<ModelSource, std::string> ModelSourceService::GetModelSource(
 }
 
 cpp::result<std::vector<std::string>, std::string>
-ModelSourceService::GetRepositoryList(std::string_view author) {
+ModelSourceService::GetRepositoryList(std::string_view author,
+                                      std::string_view tag_filter) {
   std::string as(author);
+  auto get_repo_list = [this, &as, &tag_filter] {
+    std::vector<std::string> repo_list;
+    auto const& mis = cortexso_repos_.at(as);
+    for (auto const& mi : mis) {
+      if (!tag_filter.empty()) {
+        if (std::count(mi.tags.begin(), mi.tags.end(), tag_filter)) {
+          repo_list.push_back(mi.id);
+        }
+      } else {
+        repo_list.push_back(mi.id);
+      }
+    }
+    return repo_list;
+  };
   if (cortexso_repos_.find(as) != cortexso_repos_.end() &&
-      !cortexso_repos_.at(as).empty())
-    return cortexso_repos_.at(as);
+      !cortexso_repos_.at(as).empty()) {
+    return get_repo_list();
+  }
+
   const auto begin = std::chrono::high_resolution_clock::now();
   auto res =
       curl_utils::SimpleGet("https://huggingface.co/api/models?author=" + as);
   if (res.has_value()) {
     auto repos = ParseJsonString(res.value());
     for (auto& r : repos) {
-      cortexso_repos_[as].push_back(r.id);
+      cortexso_repos_[as].push_back(r);
     }
-    return cortexso_repos_.at(as);
+    return get_repo_list();
   } else {
     return cpp::fail(res.error());
   }

engine/services/model_source_service.h

@@ -37,6 +37,17 @@ struct ModelSource {
   };
 };
 
+struct ModelInfo {
+  std::string id;
+  int likes;
+  int trending_score;
+  bool is_private;
+  int downloads;
+  std::vector<std::string> tags;
+  std::string created_at;
+  std::string model_id;
+};
+
 class ModelSourceService {
  public:
   explicit ModelSourceService(std::shared_ptr<DatabaseService> db_service);
@@ -54,7 +65,7 @@ class ModelSourceService {
   cpp::result<ModelSource, std::string> GetModelSource(const std::string& src);
 
   cpp::result<std::vector<std::string>, std::string> GetRepositoryList(
-      std::string_view author);
+      std::string_view author, std::string_view tag_filter);
 
  private:
   cpp::result<bool, std::string> AddHfOrg(const std::string& model_source,
@@ -75,12 +86,10 @@ class ModelSourceService {
       const std::string& model_source, const std::string& author,
       const std::string& model_name);
 
-  cpp::result<std::string, std::string>
-  AddCortexsoRepoBranch(const std::string& model_source,
-                        const std::string& author,
-                        const std::string& model_name,
-                        const std::string& branch, const std::string& metadata,
-                        const std::string& desc);
+  cpp::result<std::string, std::string> AddCortexsoRepoBranch(
+      const std::string& model_source, const std::string& author,
+      const std::string& model_name, const std::string& branch,
+      const std::string& metadata, const std::string& desc);
 
   void SyncModelSource();
 
@@ -89,5 +98,5 @@ class ModelSourceService {
   std::thread sync_db_thread_;
   std::atomic<bool> running_;
 
-  std::unordered_map<std::string, std::vector<std::string>> cortexso_repos_;
+  std::unordered_map<std::string, std::vector<ModelInfo>> cortexso_repos_;
 };