fix: get downloaded models

Author: sangjanai

engine/database/models.cc

@@ -270,23 +270,6 @@ bool Models::HasModel(const std::string& identifier) const {
   }
 }
 
-cpp::result<std::vector<std::string>, std::string> Models::GetModelSources()
-    const {
-  try {
-    std::vector<std::string> sources;
-    SQLite::Statement query(db_,
-                            "SELECT DISTINCT model_source FROM models WHERE "
-                            "status = \"downloadable\"");
-
-    while (query.executeStep()) {
-      sources.push_back(query.getColumn(0).getString());
-    }
-    return sources;
-  } catch (const std::exception& e) {
-    return cpp::fail(e.what());
-  }
-}
-
 cpp::result<std::vector<ModelEntry>, std::string> Models::GetModels(
     const std::string& model_src) const {
   try {
@@ -318,16 +301,17 @@ cpp::result<std::vector<ModelEntry>, std::string> Models::GetModels(
   }
 }
 
-cpp::result<std::vector<ModelEntry>, std::string>
-Models::GetDownloadableModels() const {
+cpp::result<std::vector<ModelEntry>, std::string> Models::GetModelSources()
+    const {
   try {
     std::vector<ModelEntry> res;
     SQLite::Statement query(
         db_,
         "SELECT model_id, author_repo_id, branch_name, "
         "path_to_model_yaml, model_alias, model_format, "
         "model_source, status, engine, metadata FROM models "
-        "WHERE status = \"downloadable\"");
+        "WHERE model_source != \"\" AND (status = \"downloaded\" OR status = "
+        "\"downloadable\")");
     while (query.executeStep()) {
       ModelEntry entry;
       entry.model = query.getColumn(0).getString();

engine/database/models.h

@@ -56,10 +56,9 @@ class Models {
   cpp::result<std::vector<std::string>, std::string> FindRelatedModel(
       const std::string& identifier) const;
   bool HasModel(const std::string& identifier) const;
-  cpp::result<std::vector<std::string>, std::string> GetModelSources() const;
   cpp::result<std::vector<ModelEntry>, std::string> GetModels(
       const std::string& model_src) const;
-  cpp::result<std::vector<ModelEntry>, std::string> GetDownloadableModels() const;
+  cpp::result<std::vector<ModelEntry>, std::string> GetModelSources() const;
 };
 
 }  // namespace cortex::db

engine/services/database_service.cc

@@ -118,18 +118,13 @@ bool DatabaseService::HasModel(const std::string& identifier) const {
   return cortex::db::Models().HasModel(identifier);
 }
 
-cpp::result<std::vector<std::string>, std::string>
-DatabaseService::GetModelSources() const {
-  return cortex::db::Models().GetModelSources();
-}
-
 cpp::result<std::vector<ModelEntry>, std::string> DatabaseService::GetModels(
     const std::string& model_src) const {
   return cortex::db::Models().GetModels(model_src);
 }
 
 cpp::result<std::vector<ModelEntry>, std::string>
-DatabaseService::GetDownloadableModels() const {
-  return cortex::db::Models().GetDownloadableModels();
+DatabaseService::GetModelSources() const {
+  return cortex::db::Models().GetModelSources();
 }
 // end models

engine/services/database_service.h

@@ -60,10 +60,9 @@ class DatabaseService {
   cpp::result<std::vector<std::string>, std::string> FindRelatedModel(
       const std::string& identifier) const;
   bool HasModel(const std::string& identifier) const;
-  cpp::result<std::vector<std::string>, std::string> GetModelSources() const;
   cpp::result<std::vector<ModelEntry>, std::string> GetModels(
       const std::string& model_src) const;
-  cpp::result<std::vector<ModelEntry>, std::string> GetDownloadableModels()
+  cpp::result<std::vector<ModelEntry>, std::string> GetModelSources()
       const;
 
  private:

engine/services/model_source_service.cc

@@ -113,7 +113,14 @@ cpp::result<bool, std::string> ModelSourceService::RemoveModelSource(
     return cpp::fail(srcs.error());
   } else {
     auto& v = srcs.value();
-    if (std::find(v.begin(), v.end(), model_source) == v.end()) {
+    auto exists = [&v, &model_source]() {
+      for (auto const& m : v) {
+        if (m.model_source == model_source)
+          return true;
+      }
+      return false;
+    }();
+    if (!exists) {
       return cpp::fail("Model source does not exist: " + model_source);
     }
   }
@@ -146,7 +153,7 @@ cpp::result<bool, std::string> ModelSourceService::RemoveModelSource(
 
 cpp::result<std::unordered_map<std::string, ModelSource>, std::string>
 ModelSourceService::GetModelSources() {
-  auto res = db_service_->GetDownloadableModels();
+  auto res = db_service_->GetModelSources();
   if (res.has_error()) {
     return cpp::fail(res.error());
   }
@@ -476,18 +483,18 @@ void ModelSourceService::SyncModelSource() {
         CTL_INF(res.error());
       } else {
         for (auto const& src : res.value()) {
-          CTL_DBG(src);
+          CTL_DBG(src.model_source);
         }
 
         std::unordered_set<std::string> orgs;
         std::vector<std::string> repos;
         for (auto const& src : res.value()) {
-          auto url_res = url_parser::FromUrlString(src);
+          auto url_res = url_parser::FromUrlString(src.model_source);
           if (url_res.has_value()) {
             if (url_res->pathParams.size() == 1) {
-              orgs.insert(src);
+              orgs.insert(src.model_source);
             } else if (url_res->pathParams.size() == 2) {
-              repos.push_back(src);
+              repos.push_back(src.model_source);
             }
           }
         }