From 01835de61f48e6c704ff68158b88efad6737df4b Mon Sep 17 00:00:00 2001 From: lmangani Date: Sun, 29 Dec 2024 13:41:21 +0000 Subject: [PATCH 1/8] secret manager --- CMakeLists.txt | 2 +- docs/README.md | 11 ++++++ src/include/open_prompt_secret.hpp | 13 +++++++ src/open_prompt_extension.cpp | 48 +++++++++++++++++++++--- src/open_prompt_secret.cpp | 60 ++++++++++++++++++++++++++++++ 5 files changed, 127 insertions(+), 7 deletions(-) create mode 100644 src/include/open_prompt_secret.hpp create mode 100644 src/open_prompt_secret.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 8551c08..3014c0d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -14,7 +14,7 @@ set(LOADABLE_EXTENSION_NAME ${TARGET_NAME}_loadable_extension) project(${TARGET_NAME}) include_directories(src/include duckdb/third_party/httplib) -set(EXTENSION_SOURCES src/open_prompt_extension.cpp) +set(EXTENSION_SOURCES src/open_prompt_extension.cpp src/open_prompt_secret.cpp) if(MINGW) set(OPENSSL_USE_STATIC_LIBS TRUE) diff --git a/docs/README.md b/docs/README.md index 07dea1f..0fa4564 100644 --- a/docs/README.md +++ b/docs/README.md @@ -28,7 +28,18 @@ Setup the completions API configuration w/ optional auth token and model name SET VARIABLE openprompt_api_url = 'http://localhost:11434/v1/chat/completions'; SET VARIABLE openprompt_api_token = 'your_api_key_here'; SET VARIABLE openprompt_model_name = 'qwen2.5:0.5b'; +``` +For persistent usage, configure parameters using DuckDB SECRETS +```sql +CREATE SECRET IF NOT EXISTS open_prompt ( + TYPE open_prompt, + PROVIDER config, + api_token 'your-api-token', + api_url 'http://localhost:11434/v1/chat/completions', + model_name 'qwen2.5:0.5b', + api_timeout '30' + ); ``` ### Usage diff --git a/src/include/open_prompt_secret.hpp b/src/include/open_prompt_secret.hpp new file mode 100644 index 0000000..d1f902f --- /dev/null +++ b/src/include/open_prompt_secret.hpp @@ -0,0 +1,13 @@ +#pragma once + +#include "duckdb/main/secret/secret.hpp" +#include "duckdb/main/extension_util.hpp" + +namespace duckdb { + +struct CreateOpenPromptSecretFunctions { +public: + static void Register(DatabaseInstance &instance); +}; + +} // namespace duckdb diff --git a/src/open_prompt_extension.cpp b/src/open_prompt_extension.cpp index d4c07e1..99d8de5 100644 --- a/src/open_prompt_extension.cpp +++ b/src/open_prompt_extension.cpp @@ -7,6 +7,12 @@ #include "duckdb/common/exception/http_exception.hpp" #include +#include "duckdb/main/secret/secret_manager.hpp" +#include "duckdb/main/secret/secret.hpp" +#include "duckdb/main/secret/secret_storage.hpp" + +#include "open_prompt_secret.hpp" + #ifdef USE_ZLIB #define CPPHTTPLIB_ZLIB_SUPPORT #endif @@ -142,14 +148,41 @@ namespace duckdb { // Settings management static std::string GetConfigValue(ClientContext &context, const string &var_name, const string &default_value) { - Value value; - auto &config = ClientConfig::GetConfig(context); - if (!config.GetUserVariable(var_name, value) || value.IsNull()) { - return default_value; - } - return value.ToString(); + // First try to get from secrets + auto &secret_manager = SecretManager::Get(context); + try { + auto transaction = CatalogTransaction::GetSystemCatalogTransaction(context); + auto secret_match = secret_manager.LookupSecret(transaction, "open_prompt", "open_prompt"); + if (secret_match.HasMatch()) { + auto &secret = secret_match.GetSecret(); + if (secret.GetType() != "open_prompt") { + throw InvalidInputException("Invalid secret type. Expected 'open_prompt', got '%s'", secret.GetType()); + } + + const auto *kv_secret = dynamic_cast(&secret); + if (!kv_secret) { + throw InvalidInputException("Invalid secret format for 'open_prompt' secret"); + } + + Value secret_value; + if (kv_secret->TryGetValue(var_name, secret_value)) { + return secret_value.ToString(); + } + } + } catch (...) { + // If secret lookup fails, fall back to user variables + } + + // Fall back to user variables if secret not found + Value value; + auto &config = ClientConfig::GetConfig(context); + if (!config.GetUserVariable(var_name, value) || value.IsNull()) { + return default_value; + } + return value.ToString(); } + static void SetConfigValue(DataChunk &args, ExpressionState &state, Vector &result, const string &var_name, const string &value_type) { UnaryExecutor::Execute(args.data[0], result, args.size(), @@ -356,6 +389,9 @@ namespace duckdb { LogicalType::VARCHAR, OpenPromptRequestFunction, OpenPromptBind)); + // Register Secret functions + CreateOpenPromptSecretFunctions::Register(instance); + ExtensionUtil::RegisterFunction(instance, open_prompt); ExtensionUtil::RegisterFunction(instance, ScalarFunction( diff --git a/src/open_prompt_secret.cpp b/src/open_prompt_secret.cpp new file mode 100644 index 0000000..4beeda1 --- /dev/null +++ b/src/open_prompt_secret.cpp @@ -0,0 +1,60 @@ +#include "open_prompt_secret.hpp" +#include "duckdb/common/exception.hpp" +#include "duckdb/main/secret/secret.hpp" +#include "duckdb/main/extension_util.hpp" + +namespace duckdb { + +static void CopySecret(const std::string &key, const CreateSecretInput &input, KeyValueSecret &result) { + auto val = input.options.find(key); + if (val != input.options.end()) { + result.secret_map[key] = val->second; + } +} + +static void RegisterCommonSecretParameters(CreateSecretFunction &function) { + // Register open_prompt common parameters + function.named_parameters["api_token"] = LogicalType::VARCHAR; + function.named_parameters["api_url"] = LogicalType::VARCHAR; + function.named_parameters["model_name"] = LogicalType::VARCHAR; + function.named_parameters["api_timeout"] = LogicalType::VARCHAR; +} + +static void RedactCommonKeys(KeyValueSecret &result) { + // Redact sensitive information + result.redact_keys.insert("api_token"); +} + +static unique_ptr CreateOpenPromptSecretFromConfig(ClientContext &context, CreateSecretInput &input) { + auto scope = input.scope; + auto result = make_uniq(scope, input.type, input.provider, input.name); + + // Copy all relevant secrets + CopySecret("api_token", input, *result); + CopySecret("api_url", input, *result); + CopySecret("model_name", input, *result); + CopySecret("api_timeout", input, *result); + + // Redact sensitive keys + RedactCommonKeys(*result); + + return std::move(result); +} + +void CreateOpenPromptSecretFunctions::Register(DatabaseInstance &instance) { + string type = "open_prompt"; + + // Register the new type + SecretType secret_type; + secret_type.name = type; + secret_type.deserializer = KeyValueSecret::Deserialize; + secret_type.default_provider = "config"; + ExtensionUtil::RegisterSecretType(instance, secret_type); + + // Register the config secret provider + CreateSecretFunction config_function = {type, "config", CreateOpenPromptSecretFromConfig}; + RegisterCommonSecretParameters(config_function); + ExtensionUtil::RegisterFunction(instance, config_function); +} + +} // namespace duckdb From b86b137982601e9667679c3ec031c60d9d500ba0 Mon Sep 17 00:00:00 2001 From: lmangani Date: Sun, 29 Dec 2024 14:02:47 +0000 Subject: [PATCH 2/8] ENV support --- docs/README.md | 9 +++++++++ src/open_prompt_extension.cpp | 14 +++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/docs/README.md b/docs/README.md index 0fa4564..30803ea 100644 --- a/docs/README.md +++ b/docs/README.md @@ -42,6 +42,15 @@ CREATE SECRET IF NOT EXISTS open_prompt ( ); ``` +Alternatively the following ENV variables can be used at runtime +``` + OPEN_PROMPT_API_URL + OPEN_PROMPT_API_TOKEN + OPEN_PROMPT_MODEL_NAME + OPEN_PROMPT_API_TIMEOUT +``` + + ### Usage ```sql D SELECT open_prompt('Write a one-line poem about ducks') AS response; diff --git a/src/open_prompt_extension.cpp b/src/open_prompt_extension.cpp index 99d8de5..bb09eb7 100644 --- a/src/open_prompt_extension.cpp +++ b/src/open_prompt_extension.cpp @@ -20,6 +20,9 @@ #define CPPHTTPLIB_OPENSSL_SUPPORT #include "httplib.hpp" +#include +#include +#include #include #include #include @@ -148,7 +151,16 @@ namespace duckdb { // Settings management static std::string GetConfigValue(ClientContext &context, const string &var_name, const string &default_value) { - // First try to get from secrets + + // Try environment variables + std::string env_var_name = "OPEN_PROMPT_" + var_name; + std::transform(env_var_name.begin(), env_var_name.end(), env_var_name.begin(), ::toupper); + const char* env_value = std::getenv(env_var_name.c_str()); + if (env_value != nullptr && strlen(env_value) > 0) { + return std::string(env_value); + } + + // Try to get from secrets auto &secret_manager = SecretManager::Get(context); try { auto transaction = CatalogTransaction::GetSystemCatalogTransaction(context); From 5e053207f2f23b739441258922496b6503b79df9 Mon Sep 17 00:00:00 2001 From: Lorenzo Mangani Date: Tue, 7 Jan 2025 12:15:10 +0100 Subject: [PATCH 3/8] cast unique_ptr to the correct type --- src/open_prompt_extension.cpp | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/open_prompt_extension.cpp b/src/open_prompt_extension.cpp index bb09eb7..61da7b8 100644 --- a/src/open_prompt_extension.cpp +++ b/src/open_prompt_extension.cpp @@ -38,13 +38,13 @@ namespace duckdb { idx_t model_idx; idx_t json_schema_idx; idx_t json_system_prompt_idx; - unique_ptr Copy() const { - auto res = make_uniq(); - res->model_idx = model_idx; - res->json_schema_idx = json_schema_idx; - res->json_system_prompt_idx = json_system_prompt_idx; - return res; - }; + unique_ptr Copy() const override { + auto res = make_uniq(); + res->model_idx = model_idx; + res->json_schema_idx = json_schema_idx; + res->json_system_prompt_idx = json_system_prompt_idx; + return std::unique_ptr(std::move(res)); + }; bool Equals(const FunctionData &other) const { return model_idx == other.Cast().model_idx && json_schema_idx == other.Cast().json_schema_idx && From e443363d24c545720e594e3b101ecb11afa7880a Mon Sep 17 00:00:00 2001 From: Lorenzo Mangani Date: Tue, 7 Jan 2025 12:29:33 +0100 Subject: [PATCH 4/8] cast unique_ptr to the correct type --- src/open_prompt_extension.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/open_prompt_extension.cpp b/src/open_prompt_extension.cpp index 61da7b8..4676dc8 100644 --- a/src/open_prompt_extension.cpp +++ b/src/open_prompt_extension.cpp @@ -43,7 +43,7 @@ namespace duckdb { res->model_idx = model_idx; res->json_schema_idx = json_schema_idx; res->json_system_prompt_idx = json_system_prompt_idx; - return std::unique_ptr(std::move(res)); + return unique_ptr(std::move(res)); }; bool Equals(const FunctionData &other) const { return model_idx == other.Cast().model_idx && From ba088287a3aa526f6431e4b33fb2d88ff09ecbda Mon Sep 17 00:00:00 2001 From: lmangani Date: Tue, 7 Jan 2025 14:05:13 +0000 Subject: [PATCH 5/8] resync --- duckdb | 2 +- extension-ci-tools | 2 +- src/open_prompt_extension.cpp | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/duckdb b/duckdb index af39bd0..b9e368e 160000 --- a/duckdb +++ b/duckdb @@ -1 +1 @@ -Subproject commit af39bd0dcf66876e09ac2a7c3baa28fe1b301151 +Subproject commit b9e368e888fed036f598acf4994bd30fbfe97472 diff --git a/extension-ci-tools b/extension-ci-tools index 00831df..f473553 160000 --- a/extension-ci-tools +++ b/extension-ci-tools @@ -1 +1 @@ -Subproject commit 00831df06713072df217d3fb2f6b5e0fae78742f +Subproject commit f473553168fd1db490aaa9f440b8f812af0568da diff --git a/src/open_prompt_extension.cpp b/src/open_prompt_extension.cpp index 4676dc8..3177d39 100644 --- a/src/open_prompt_extension.cpp +++ b/src/open_prompt_extension.cpp @@ -38,7 +38,7 @@ namespace duckdb { idx_t model_idx; idx_t json_schema_idx; idx_t json_system_prompt_idx; - unique_ptr Copy() const override { + unique_ptr Copy() const override { auto res = make_uniq(); res->model_idx = model_idx; res->json_schema_idx = json_schema_idx; From c981e4ac231df98a0254afedca37fd05321611eb Mon Sep 17 00:00:00 2001 From: lmangani Date: Tue, 7 Jan 2025 14:14:14 +0000 Subject: [PATCH 6/8] add tests --- test/sql/open_prompt.test | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 test/sql/open_prompt.test diff --git a/test/sql/open_prompt.test b/test/sql/open_prompt.test new file mode 100644 index 0000000..276d7e5 --- /dev/null +++ b/test/sql/open_prompt.test @@ -0,0 +1,31 @@ +# name: test/sql/rusty_quack.test +# description: test rusty_quack extension +# group: [quack] + +# Before we load the extension, this will fail +statement error +SELECT open_prompt('error'); +---- +Catalog Error: Scalar Function with name open_prompt does not exist! + +# Require statement will ensure the extension is loaded from now on +require open_prompt + +# Confirm the extension works by setting a secret +query I +CREATE SECRET IF NOT EXISTS open_prompt ( + TYPE open_prompt, + PROVIDER config, + api_token 'xxxxx', + api_url 'https://api.groq.com/openai/v1/chat/completions', + model_name 'llama-3.3-70b-versatile', + api_timeout '30' + ); +---- +true + +# Confirm the secret exists +query I +SELECT name FROM duckdb_secrets() WHERE name = 'open_prompt' ; +---- +open_prompt From 0172f3b2a3f363cf7a69765ec2832c821859f325 Mon Sep 17 00:00:00 2001 From: lmangani Date: Tue, 7 Jan 2025 18:44:11 +0000 Subject: [PATCH 7/8] fix env, secrets handling --- src/open_prompt_extension.cpp | 105 +++++++++++++++++++++------------- 1 file changed, 65 insertions(+), 40 deletions(-) diff --git a/src/open_prompt_extension.cpp b/src/open_prompt_extension.cpp index 3177d39..4daf812 100644 --- a/src/open_prompt_extension.cpp +++ b/src/open_prompt_extension.cpp @@ -151,47 +151,72 @@ namespace duckdb { // Settings management static std::string GetConfigValue(ClientContext &context, const string &var_name, const string &default_value) { + // Try environment variables + { + // Create uppercase ENV version: OPEN_PROMPT_SETTING + std::string stripped_name = var_name; + const std::string prefix = "openprompt_"; + if (stripped_name.substr(0, prefix.length()) == prefix) { + stripped_name = stripped_name.substr(prefix.length()); + } + std::string env_var_name = "OPEN_PROMPT_" + stripped_name; + std::transform(env_var_name.begin(), env_var_name.end(), env_var_name.begin(), ::toupper); + // std::cout << "SEARCH ENV FOR " << env_var_name << "\n"; + + const char* env_value = std::getenv(env_var_name.c_str()); + if (env_value != nullptr && strlen(env_value) > 0) { + // std::cout << "USING ENV FOR " << var_name << "\n"; + std::string result(env_value); + return result; + } + } + + // Try to get from secrets + { + // Create lowercase secret version: open_prompt_setting + std::string secret_key = var_name; + const std::string prefix = "openprompt_"; + if (secret_key.substr(0, prefix.length()) == prefix) { + secret_key = secret_key.substr(prefix.length()); + } + // secret_key = "open_prompt_" + secret_key; + std::transform(secret_key.begin(), secret_key.end(), secret_key.begin(), ::tolower); + + auto &secret_manager = SecretManager::Get(context); + try { + // std::cout << "SEARCH SECRET FOR " << secret_key << "\n"; + auto transaction = CatalogTransaction::GetSystemCatalogTransaction(context); + auto secret_match = secret_manager.LookupSecret(transaction, "open_prompt", "open_prompt"); + if (secret_match.HasMatch()) { + auto &secret = secret_match.GetSecret(); + if (secret.GetType() != "open_prompt") { + throw InvalidInputException("Invalid secret type. Expected 'open_prompt', got '%s'", secret.GetType()); + } + const auto *kv_secret = dynamic_cast(&secret); + if (!kv_secret) { + throw InvalidInputException("Invalid secret format for 'open_prompt' secret"); + } + Value secret_value; + if (kv_secret->TryGetValue(secret_key, secret_value)) { + // std::cout << "USING SECRET FOR " << var_name << "\n"; + return secret_value.ToString(); + } + } + } catch (...) { + // If secret lookup fails, fall back to user variables + } + } + + // Fall back to user variables if secret not found (using original var_name) + Value value; + auto &config = ClientConfig::GetConfig(context); + if (!config.GetUserVariable(var_name, value) || value.IsNull()) { + // std::cout << "USING SET FOR " << var_name << "\n"; + return default_value; + } - // Try environment variables - std::string env_var_name = "OPEN_PROMPT_" + var_name; - std::transform(env_var_name.begin(), env_var_name.end(), env_var_name.begin(), ::toupper); - const char* env_value = std::getenv(env_var_name.c_str()); - if (env_value != nullptr && strlen(env_value) > 0) { - return std::string(env_value); - } - - // Try to get from secrets - auto &secret_manager = SecretManager::Get(context); - try { - auto transaction = CatalogTransaction::GetSystemCatalogTransaction(context); - auto secret_match = secret_manager.LookupSecret(transaction, "open_prompt", "open_prompt"); - if (secret_match.HasMatch()) { - auto &secret = secret_match.GetSecret(); - if (secret.GetType() != "open_prompt") { - throw InvalidInputException("Invalid secret type. Expected 'open_prompt', got '%s'", secret.GetType()); - } - - const auto *kv_secret = dynamic_cast(&secret); - if (!kv_secret) { - throw InvalidInputException("Invalid secret format for 'open_prompt' secret"); - } - - Value secret_value; - if (kv_secret->TryGetValue(var_name, secret_value)) { - return secret_value.ToString(); - } - } - } catch (...) { - // If secret lookup fails, fall back to user variables - } - - // Fall back to user variables if secret not found - Value value; - auto &config = ClientConfig::GetConfig(context); - if (!config.GetUserVariable(var_name, value) || value.IsNull()) { - return default_value; - } - return value.ToString(); + // std::cout << "USING DEFAULT FOR " << var_name << "\n"; + return value.ToString(); } From 9928d4886f1250b21fb02cb21c2bd87af6ecbb97 Mon Sep 17 00:00:00 2001 From: Lorenzo Mangani Date: Tue, 7 Jan 2025 19:57:50 +0100 Subject: [PATCH 8/8] Update README.md --- docs/README.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/README.md b/docs/README.md index 30803ea..5d1da08 100644 --- a/docs/README.md +++ b/docs/README.md @@ -30,6 +30,14 @@ SET VARIABLE openprompt_api_token = 'your_api_key_here'; SET VARIABLE openprompt_model_name = 'qwen2.5:0.5b'; ``` +Alternatively the following ENV variables can be used at runtime +``` + OPEN_PROMPT_API_URL='http://localhost:11434/v1/chat/completions' + OPEN_PROMPT_API_TOKEN='your_api_key_here' + OPEN_PROMPT_MODEL_NAME='qwen2.5:0.5b' + OPEN_PROMPT_API_TIMEOUT='30' +``` + For persistent usage, configure parameters using DuckDB SECRETS ```sql CREATE SECRET IF NOT EXISTS open_prompt ( @@ -42,14 +50,6 @@ CREATE SECRET IF NOT EXISTS open_prompt ( ); ``` -Alternatively the following ENV variables can be used at runtime -``` - OPEN_PROMPT_API_URL - OPEN_PROMPT_API_TOKEN - OPEN_PROMPT_MODEL_NAME - OPEN_PROMPT_API_TIMEOUT -``` - ### Usage ```sql