Udpate the azure type into azure_llm in secret manager (#88)

dsg-polymtl · Dec 8, 2024 · a9c9d67 · a9c9d67
1 parent a5bcafa
commit a9c9d67
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 12 deletions.
diff --git a/src/include/flockmtl/secret_manager/secret_manager.hpp b/src/include/flockmtl/secret_manager/secret_manager.hpp
@@ -22,6 +22,7 @@ extern const SecretDetails ollama_secret_details;
 class SecretManager {
 public:
     enum SupportedProviders { OPENAI, AZURE, OLLAMA };
+    static std::unordered_map<std::string, SupportedProviders> providerNames;
 
     static void Register(duckdb::DatabaseInstance& instance);
     static std::unordered_map<std::string, std::string> GetSecret(std::string provider);

diff --git a/src/model_manager/model.cpp b/src/model_manager/model.cpp
@@ -22,6 +22,8 @@ void Model::LoadModelDetails(const nlohmann::json& model_json) {
     model_details_.provider_name =
         model_json.contains("provider") ? model_json.at("provider").get<std::string>() : std::get<1>(query_result);
     auto secret_name = "__default_" + model_details_.provider_name;
+    if (model_details_.provider_name == AZURE)
+        secret_name += "_llm";
     if (model_json.contains("secret_name")) {
         secret_name = model_json["secret_name"].get<std::string>();
     }

diff --git a/src/secret_manager/secret_manager.cpp b/src/secret_manager/secret_manager.cpp
@@ -8,25 +8,20 @@ namespace flockmtl {
 
 const SecretDetails openai_secret_details = {"openai", "flockmtl", "openai://", {"api_key"}, {"api_key"}, {"api_key"}};
 
-const SecretDetails azure_secret_details = {"azure",     "flockmtl",
-                                            "azure://",  {"api_key", "resource_name", "api_version"},
-                                            {"api_key"}, {"api_key", "resource_name", "api_version"}};
+const SecretDetails azure_secret_details = {"azure_llm",    "flockmtl",
+                                            "azure_llm://", {"api_key", "resource_name", "api_version"},
+                                            {"api_key"},    {"api_key", "resource_name", "api_version"}};
 
 const SecretDetails ollama_secret_details = {"ollama", "flockmtl", "ollama://", {"api_url"}, {"api_url"}, {"api_url"}};
 
 const std::vector<const SecretDetails*> secret_details_list = {&openai_secret_details, &azure_secret_details,
                                                                &ollama_secret_details};
 
+std::unordered_map<std::string, SecretManager::SupportedProviders> SecretManager::providerNames = {
+    {"openai", OPENAI}, {"azure", AZURE}, {"ollama", OLLAMA}};
+
 SecretManager::SupportedProviders SecretManager::GetProviderType(std::string provider) {
-    if (provider == "openai") {
-        return OPENAI;
-    } else if (provider == "azure") {
-        return AZURE;
-    } else if (provider == "ollama") {
-        return OLLAMA;
-    } else {
-        throw duckdb::InvalidInputException("Unsupported secret type: %s", provider.c_str());
-    }
+    return providerNames[provider];
 }
 
 void SecretManager::Register(duckdb::DatabaseInstance& instance) {