This is an automated email from the ASF dual-hosted git repository. yiguolei pushed a commit to branch branch-4.2 in repository https://gitbox.apache.org/repos/asf/doris.git
commit 97976581cc58fca14582102f14aa55d84825917d Author: linrrarity <[email protected]> AuthorDate: Wed Sep 23 08:52:32 2026 +0800 branch-4.1: [Enhance](ai_func) Support dedicate embed properties in AI RESOURCE (#67673) (#68384) pick: https://github.com/apache/doris/pull/67673 --- be/src/core/string_buffer.hpp | 7 +- be/src/exprs/function/ai/ai_adapter.h | 77 ++- be/src/exprs/function/ai/ai_functions.h | 25 +- be/src/exprs/function/ai/embed.h | 34 +- be/test/ai/ai_adapter_test.cpp | 15 + be/test/ai/ai_function_test.cpp | 12 +- be/test/ai/embed_test.cpp | 148 ++++++ be/test/core/string_buffer_test.cpp | 9 + .../java/org/apache/doris/catalog/AIResource.java | 123 +++-- .../org/apache/doris/common/util/PrintableMap.java | 3 + .../property/constants/AIProperties.java | 88 +++- .../trees/expressions/functions/agg/AIAgg.java | 4 + .../trees/expressions/functions/ai/AIFunction.java | 4 + .../trees/expressions/functions/ai/Embed.java | 17 +- .../org/apache/doris/catalog/AIResourceTest.java | 515 +++++++++++++++++++++ .../apache/doris/common/util/PrintableMapTest.java | 7 + gensrc/thrift/PaloInternalService.thrift | 15 +- 17 files changed, 1008 insertions(+), 95 deletions(-) diff --git a/be/src/core/string_buffer.hpp b/be/src/core/string_buffer.hpp index 81223f99b0d..4f93776ab2a 100644 --- a/be/src/core/string_buffer.hpp +++ b/be/src/core/string_buffer.hpp @@ -227,8 +227,8 @@ using BufferWriter = BufferWritable; // There is consumption of the buffer in the read method. class BufferReadable { public: - explicit BufferReadable(StringRef& ref) : _data(ref.data) {} - explicit BufferReadable(StringRef&& ref) : _data(ref.data) {} + explicit BufferReadable(StringRef& ref) : _data(ref.data), _end(ref.data + ref.size) {} + explicit BufferReadable(StringRef&& ref) : _data(ref.data), _end(ref.data + ref.size) {} ~BufferReadable() = default; StringRef read(size_t len) { @@ -244,6 +244,8 @@ public: const char* data() { return _data; } + bool has_remaining() const { return _data < _end; } + void add_offset(size_t len) { _data += len; } void read_var_uint(UInt64& x) { @@ -316,6 +318,7 @@ public: private: const char* _data; + const char* _end; }; using VectorBufferReader = BufferReadable; diff --git a/be/src/exprs/function/ai/ai_adapter.h b/be/src/exprs/function/ai/ai_adapter.h index eff48ef83cd..4a527e96761 100644 --- a/be/src/exprs/function/ai/ai_adapter.h +++ b/be/src/exprs/function/ai/ai_adapter.h @@ -42,16 +42,17 @@ namespace doris { struct AIResource { AIResource() = default; AIResource(const TAIResource& tai) - : endpoint(tai.endpoint), - provider_type(tai.provider_type), - model_name(tai.model_name), - api_key(tai.api_key), - temperature(tai.temperature), - max_tokens(tai.max_tokens), - max_retries(tai.max_retries), - retry_delay_second(tai.retry_delay_second), - anthropic_version(tai.anthropic_version), - dimensions(tai.dimensions) {} + : AIResource(tai, tai.endpoint, tai.provider_type, tai.model_name, tai.api_key) {} + + static AIResource from_embed(const TAIResource& tai) { + return AIResource(tai, tai.embed_endpoint, tai.embed_provider_type, tai.embed_model_name, + tai.embed_api_key); + } + + static AIResource from_multimodal_embed(const TAIResource& tai) { + return AIResource(tai, tai.embed_mm_endpoint, tai.embed_mm_provider_type, + tai.embed_mm_model_name, tai.embed_mm_api_key); + } std::string endpoint; std::string provider_type; @@ -63,6 +64,7 @@ struct AIResource { int32_t retry_delay_second; std::string anthropic_version; int32_t dimensions; + std::string effort; void serialize(BufferWritable& buf) const { buf.write_binary(endpoint); @@ -75,6 +77,9 @@ struct AIResource { buf.write_binary(retry_delay_second); buf.write_binary(anthropic_version); buf.write_binary(dimensions); + if (!effort.empty()) { + buf.write_binary(effort); + } } void deserialize(BufferReadable& buf) { @@ -88,7 +93,26 @@ struct AIResource { buf.read_binary(retry_delay_second); buf.read_binary(anthropic_version); buf.read_binary(dimensions); + if (buf.has_remaining()) { + buf.read_binary(effort); + } } + +private: + AIResource(const TAIResource& tai, const std::string& selected_endpoint, + const std::string& selected_provider_type, const std::string& selected_model_name, + const std::string& selected_api_key) + : endpoint(selected_endpoint), + provider_type(selected_provider_type), + model_name(selected_model_name), + api_key(selected_api_key), + temperature(tai.temperature), + max_tokens(tai.max_tokens), + max_retries(tai.max_retries), + retry_delay_second(tai.retry_delay_second), + anthropic_version(tai.anthropic_version), + dimensions(tai.dimensions), + effort(tai.effort) {} }; enum class MultimodalType { IMAGE, VIDEO, AUDIO }; @@ -123,6 +147,8 @@ public: _config.max_retries = config.max_retries; _config.retry_delay_second = config.retry_delay_second; _config.anthropic_version = config.anthropic_version; + _config.dimensions = config.dimensions; + _config.effort = config.effort; } // Build request payload based on input text strings @@ -767,7 +793,8 @@ public: {"role": "user", "content": "xxx"} ], "temperature": 0.7, - "max_output_tokens": 150 + "max_output_tokens": 150, + "reasoning": {"effort": "max"} }*/ doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator), allocator); @@ -779,6 +806,12 @@ public: if (_config.max_tokens != -1) { doc.AddMember("max_output_tokens", _config.max_tokens, allocator); } + if (!_config.effort.empty()) { + rapidjson::Value reasoning(rapidjson::kObjectType); + reasoning.AddMember("effort", rapidjson::Value(_config.effort.c_str(), allocator), + allocator); + doc.AddMember("reasoning", reasoning, allocator); + } // input rapidjson::Value input(rapidjson::kArrayType); @@ -804,6 +837,7 @@ public: ], "temperature": x, "max_tokens": x, + "reasoning_effort": "low" }*/ doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator), allocator); @@ -815,6 +849,10 @@ public: if (_config.max_tokens != -1) { doc.AddMember("max_tokens", _config.max_tokens, allocator); } + if (!_config.effort.empty()) { + doc.AddMember("reasoning_effort", + rapidjson::Value(_config.effort.c_str(), allocator), allocator); + } rapidjson::Value messages(rapidjson::kArrayType); if (system_prompt && *system_prompt) { @@ -1264,7 +1302,8 @@ public: ], "generationConfig": { "temperature": 0.7, - "maxOutputTokens": 1024 + "maxOutputTokens": 1024, + "thinkingConfig": {"thinkingLevel": "high"} } }*/ @@ -1302,6 +1341,13 @@ public: if (_config.max_tokens != -1) { generationConfig.AddMember("maxOutputTokens", _config.max_tokens, allocator); } + if (!_config.effort.empty()) { + rapidjson::Value thinking_config(rapidjson::kObjectType); + thinking_config.AddMember("thinkingLevel", + rapidjson::Value(_config.effort.c_str(), allocator), + allocator); + generationConfig.AddMember("thinkingConfig", thinking_config, allocator); + } doc.AddMember("generationConfig", generationConfig, allocator); rapidjson::StringBuffer buffer; @@ -1580,6 +1626,7 @@ public: /* "model": "claude-opus-4-1-20250805", "max_tokens": 1024, + "output_config": {"effort": "medium"}, "system": "system_prompt here", "messages": [ {"role": "user", "content": "xxx"} @@ -1598,6 +1645,12 @@ public: // Keep the default value, Anthropic requires this parameter doc.AddMember("max_tokens", 2048, allocator); } + if (!_config.effort.empty()) { + rapidjson::Value output_config(rapidjson::kObjectType); + output_config.AddMember("effort", rapidjson::Value(_config.effort.c_str(), allocator), + allocator); + doc.AddMember("output_config", output_config, allocator); + } if (system_prompt && *system_prompt) { doc.AddMember("system", rapidjson::Value(system_prompt, allocator), allocator); } diff --git a/be/src/exprs/function/ai/ai_functions.h b/be/src/exprs/function/ai/ai_functions.h index 6fdd0414e47..fd2aabdfe34 100644 --- a/be/src/exprs/function/ai/ai_functions.h +++ b/be/src/exprs/function/ai/ai_functions.h @@ -97,7 +97,7 @@ public: return Status::OK(); } - TAIResource config; + AIResource config; std::shared_ptr<AIAdapter> adapter; if (Status status = this->_init_from_resource(context, block, arguments, config, adapter); !status.ok()) { @@ -131,7 +131,7 @@ protected: return Status::OK(); } - static void normalize_endpoint(TAIResource& config) { + static void normalize_endpoint(AIResource& config) { // 1. If users configure only the version root like `.../v1` or `.../v1beta`, append // `models/<model>:batchEmbedContents` for `embed`, and `models/<model>:generateContent` // for other AI scalar functions. @@ -172,7 +172,7 @@ protected: // Executes one HTTP POST request and validates transport-level success. Status do_send_request(HttpClient* client, const std::string& request_body, - std::string& response, const TAIResource& config, + std::string& response, const AIResource& config, std::shared_ptr<AIAdapter>& adapter, FunctionContext* context) const { RETURN_IF_ERROR(client->init(config.endpoint, false)); @@ -208,7 +208,7 @@ protected: // Sends the request with retry mechanism for handling transient failures Status send_request_to_llm(const std::string& request_body, std::string& response, - const TAIResource& config, std::shared_ptr<AIAdapter>& adapter, + const AIResource& config, std::shared_ptr<AIAdapter>& adapter, FunctionContext* context) const { return HttpClient::execute_with_retry(config.max_retries, config.retry_delay_second, [this, &request_body, &response, &config, &adapter, @@ -229,7 +229,7 @@ protected: // Provider-reusable helper for string-returning functions. // Executes one batch request and parses the provider result into one string per input row. Status execute_batch_request(const std::vector<std::string>& batch_prompts, - std::vector<std::string>& results, const TAIResource& config, + std::vector<std::string>& results, const AIResource& config, std::shared_ptr<AIAdapter>& adapter, FunctionContext* context) const { #ifdef BE_TEST @@ -293,7 +293,7 @@ protected: // Runs the common batch execution flow; derived classes only need to define how one batch of // string results is inserted into the final output column. Status execute(FunctionContext* context, Block& block, const ColumnNumbers& arguments, - uint32_t result, size_t input_rows_count, const TAIResource& config, + uint32_t result, size_t input_rows_count, const AIResource& config, std::shared_ptr<AIAdapter>& adapter) const { Columns prompt_columns; prompt_columns.reserve(arguments.size() - 1); @@ -413,7 +413,7 @@ protected: private: // The ai resource must be literal Status _init_from_resource(FunctionContext* context, const Block& block, - const ColumnNumbers& arguments, TAIResource& config, + const ColumnNumbers& arguments, AIResource& config, std::shared_ptr<AIAdapter>& adapter) const { const ColumnWithTypeAndName& resource_column = block.get_by_position(arguments[0]); StringRef resource_name_ref = resource_column.column->get_data_at(0); @@ -424,7 +424,12 @@ private: DORIS_CHECK(ai_resources); auto it = ai_resources->find(resource_name); DORIS_CHECK(it != ai_resources->end()); - config = it->second; + PrimitiveType input_type = INVALID_TYPE; + if (arguments.size() > 1) { + input_type = + remove_nullable(block.get_by_position(arguments[1]).type)->get_primitive_type(); + } + config = assert_cast<const Derived&>(*this).select_ai_resource(it->second, input_type); normalize_endpoint(config); @@ -435,6 +440,10 @@ private: return Status::OK(); } + AIResource select_ai_resource(const TAIResource& resource, PrimitiveType /*input_type*/) const { + return AIResource(resource); + } + // Serializes one text batch into the shared JSON-array prompt format consumed by LLM // providers for batch string functions. Status build_batch_prompt(const std::vector<std::string>& batch_prompts, diff --git a/be/src/exprs/function/ai/embed.h b/be/src/exprs/function/ai/embed.h index f193a1c171c..7ed1665b2d5 100644 --- a/be/src/exprs/function/ai/embed.h +++ b/be/src/exprs/function/ai/embed.h @@ -44,8 +44,22 @@ public: using PreparedFunctionImpl::execute; + AIResource select_ai_resource(const TAIResource& resource, PrimitiveType input_type) const { + bool has_complete_multimodal_embed_properties = _has_complete_resource_properties( + resource.embed_mm_endpoint, resource.embed_mm_provider_type, + resource.embed_mm_model_name, resource.embed_mm_api_key); + if (input_type == PrimitiveType::TYPE_JSONB && has_complete_multimodal_embed_properties) { + return AIResource::from_multimodal_embed(resource); + } + bool has_complete_embed_properties = _has_complete_resource_properties( + resource.embed_endpoint, resource.embed_provider_type, resource.embed_model_name, + resource.embed_api_key); + return has_complete_embed_properties ? AIResource::from_embed(resource) + : AIResource(resource); + } + Status execute(FunctionContext* context, Block& block, const ColumnNumbers& arguments, - uint32_t result, size_t input_rows_count, const TAIResource& config, + uint32_t result, size_t input_rows_count, const AIResource& config, std::shared_ptr<AIAdapter>& adapter) const { if (arguments.size() != 2) { return Status::InvalidArgument("Function EMBED expects 2 arguments, but got {}", @@ -90,6 +104,14 @@ public: static FunctionPtr create() { return std::make_shared<FunctionEmbed>(); } private: + static bool _has_complete_resource_properties(std::string_view endpoint, + std::string_view provider_type, + std::string_view model_name, + std::string_view api_key) { + return !endpoint.empty() && !provider_type.empty() && !model_name.empty() && + (provider_type == "LOCAL" || !api_key.empty()); + } + static int32_t _get_embed_max_batch_size(FunctionContext* context) { QueryContext* query_ctx = context->state()->get_query_ctx(); DORIS_CHECK(query_ctx != nullptr); @@ -98,7 +120,7 @@ private: } Status _execute_text_embed(FunctionContext* context, Block& block, uint32_t result, - size_t input_rows_count, const TAIResource& config, + size_t input_rows_count, const AIResource& config, std::shared_ptr<AIAdapter>& adapter, const ColumnPtr& input_column, ColumnUInt8::MutablePtr result_null_map) const { auto col_result = ColumnArray::create( @@ -155,7 +177,7 @@ private: } Status _execute_multimodal_embed(FunctionContext* context, Block& block, uint32_t result, - size_t input_rows_count, const TAIResource& config, + size_t input_rows_count, const AIResource& config, std::shared_ptr<AIAdapter>& adapter, const ColumnPtr& input_column, ColumnUInt8::MutablePtr result_null_map) const { @@ -218,7 +240,7 @@ private: // Sends one embedding request with a prebuilt request body and validates returned row count. Status _execute_prebuilt_embedding_request(const std::string& request_body, std::vector<std::vector<float>>& results, - size_t expected_size, const TAIResource& config, + size_t expected_size, const AIResource& config, std::shared_ptr<AIAdapter>& adapter, FunctionContext* context) const { std::string response; @@ -251,7 +273,7 @@ private: // EMBED-private helper. // Flushes one accumulated text embedding batch into the output array column. Status _flush_text_embedding_batch(std::vector<std::string>& batch_prompts, - ColumnArray& col_result, const TAIResource& config, + ColumnArray& col_result, const AIResource& config, std::shared_ptr<AIAdapter>& adapter, FunctionContext* context) const { if (batch_prompts.empty()) { @@ -275,7 +297,7 @@ private: Status _flush_multimodal_embedding_batch(std::vector<MultimodalType>& batch_media_types, std::vector<std::string>& batch_media_content_types, std::vector<std::string>& batch_media_urls, - ColumnArray& col_result, const TAIResource& config, + ColumnArray& col_result, const AIResource& config, std::shared_ptr<AIAdapter>& adapter, FunctionContext* context) const { if (batch_media_urls.empty()) { diff --git a/be/test/ai/ai_adapter_test.cpp b/be/test/ai/ai_adapter_test.cpp index 1eac053d205..747f15b6597 100644 --- a/be/test/ai/ai_adapter_test.cpp +++ b/be/test/ai/ai_adapter_test.cpp @@ -248,6 +248,7 @@ TEST(AI_ADAPTER_TEST, openai_adapter_completions_request) { config.temperature = 0.5; config.max_tokens = 64; config.api_key = "test_openai_key"; + config.effort = "low"; adapter.init(config); // header @@ -284,6 +285,8 @@ TEST(AI_ADAPTER_TEST, openai_adapter_completions_request) { ASSERT_TRUE(doc.HasMember("max_tokens")) << "Missing max_tokens field"; ASSERT_TRUE(doc["max_tokens"].IsInt()) << "Max_tokens field is not an integer"; ASSERT_EQ(doc["max_tokens"].GetInt(), 64); + ASSERT_TRUE(doc.HasMember("reasoning_effort")); + ASSERT_STREQ(doc["reasoning_effort"].GetString(), "low"); // msg ASSERT_TRUE(doc.HasMember("messages")) << "Missing messages field"; ASSERT_TRUE(doc["messages"].IsArray()) << "Messages is not an array"; @@ -321,6 +324,7 @@ TEST(AI_ADAPTER_TEST, openai_adatper_responses_request) { config.max_tokens = 64; config.api_key = "test_openai_key"; config.endpoint = "https://api.openai.com/v1/responses";; + config.effort = "max"; adapter.init(config); // header @@ -357,6 +361,9 @@ TEST(AI_ADAPTER_TEST, openai_adatper_responses_request) { ASSERT_TRUE(doc.HasMember("max_output_tokens")) << "Missing max_output_tokens field"; ASSERT_TRUE(doc["max_output_tokens"].IsInt()) << "max_output_tokens field is not an integer"; ASSERT_EQ(doc["max_output_tokens"].GetInt(), 64); + ASSERT_TRUE(doc.HasMember("reasoning")); + ASSERT_TRUE(doc["reasoning"].IsObject()); + ASSERT_STREQ(doc["reasoning"]["effort"].GetString(), "max"); // input ASSERT_TRUE(doc.HasMember("input")) << "Missing input field"; @@ -671,6 +678,7 @@ TEST(AI_ADAPTER_TEST, gemini_adapter_request) { config.temperature = 0.2; config.max_tokens = 32; config.api_key = "test_gemini_key"; + config.effort = "high"; adapter.init(config); // header test @@ -705,6 +713,9 @@ TEST(AI_ADAPTER_TEST, gemini_adapter_request) { ASSERT_TRUE(gen_cfg.HasMember("maxOutputTokens")) << "Missing maxOutputTokens field"; ASSERT_TRUE(gen_cfg["maxOutputTokens"].IsInt()); ASSERT_EQ(gen_cfg["maxOutputTokens"].GetInt(), 32); + ASSERT_TRUE(gen_cfg.HasMember("thinkingConfig")); + ASSERT_TRUE(gen_cfg["thinkingConfig"].IsObject()); + ASSERT_STREQ(gen_cfg["thinkingConfig"]["thinkingLevel"].GetString(), "high"); // system_prompt ASSERT_TRUE(doc.HasMember("systemInstruction")) << "Missing system field"; @@ -753,6 +764,7 @@ TEST(AI_ADAPTER_TEST, anthropic_adapter_request) { config.max_tokens = 256; config.api_key = "test_anthropic_key"; config.anthropic_version = "2023-06-01"; + config.effort = "medium"; adapter.init(config); // header @@ -790,6 +802,9 @@ TEST(AI_ADAPTER_TEST, anthropic_adapter_request) { ASSERT_TRUE(doc.HasMember("max_tokens")) << "Missing max_tokens field"; ASSERT_TRUE(doc["max_tokens"].IsInt()) << "Max_tokens field is not an integer"; ASSERT_EQ(doc["max_tokens"].GetInt(), 256); + ASSERT_TRUE(doc.HasMember("output_config")); + ASSERT_TRUE(doc["output_config"].IsObject()); + ASSERT_STREQ(doc["output_config"]["effort"].GetString(), "medium"); // system_prompt ASSERT_TRUE(doc.HasMember("system")) << "Missing system field"; diff --git a/be/test/ai/ai_function_test.cpp b/be/test/ai/ai_function_test.cpp index 976eaddd5bc..40db74c1344 100644 --- a/be/test/ai/ai_function_test.cpp +++ b/be/test/ai/ai_function_test.cpp @@ -1558,14 +1558,14 @@ public: }; TEST(AIFunctionTest, NormalizeLegacyCompletionsEndpoint) { - TAIResource resource; + AIResource resource; resource.endpoint = "https://api.openai.com/v1/completions";; FunctionAISentimentTestHelper::normalize_endpoint(resource); ASSERT_EQ(resource.endpoint, "https://api.openai.com/v1/chat/completions";); } TEST(AIFunctionTest, NormalizeEndpointNoopForOtherPaths) { - TAIResource resource; + AIResource resource; resource.endpoint = "https://api.openai.com/v1/chat/completions";; FunctionAISentimentTestHelper::normalize_endpoint(resource); ASSERT_EQ(resource.endpoint, "https://api.openai.com/v1/chat/completions";); @@ -1576,7 +1576,7 @@ TEST(AIFunctionTest, NormalizeEndpointNoopForOtherPaths) { } TEST(AIFunctionTest, NormalizeGeminiGenerateEndpointFromBaseVersion) { - TAIResource resource; + AIResource resource; resource.provider_type = "gemini"; resource.model_name = "gemini-pro"; resource.endpoint = "https://generativelanguage.googleapis.com/v1beta";; @@ -1587,7 +1587,7 @@ TEST(AIFunctionTest, NormalizeGeminiGenerateEndpointFromBaseVersion) { } TEST(AIFunctionTest, NormalizeGeminiEmbedEndpointFromBaseVersion) { - TAIResource resource; + AIResource resource; resource.provider_type = "GEMINI"; resource.model_name = "gemini-embedding-2-preview"; resource.endpoint = "https://generativelanguage.googleapis.com/v1beta";; @@ -1599,7 +1599,7 @@ TEST(AIFunctionTest, NormalizeGeminiEmbedEndpointFromBaseVersion) { } TEST(AIFunctionTest, NormalizeGeminiEndpointNoopForNonBasePath) { - TAIResource resource; + AIResource resource; resource.provider_type = "gemini"; resource.model_name = "gemini-pro"; resource.endpoint = @@ -1611,7 +1611,7 @@ TEST(AIFunctionTest, NormalizeGeminiEndpointNoopForNonBasePath) { } TEST(AIFunctionTest, NormalizeGeminiEmbedLegacySingleEndpointToBatchEndpoint) { - TAIResource resource; + AIResource resource; resource.provider_type = "gemini"; resource.model_name = "gemini-embedding-2-preview"; resource.endpoint = diff --git a/be/test/ai/embed_test.cpp b/be/test/ai/embed_test.cpp index 3ed13921d19..e45082193af 100644 --- a/be/test/ai/embed_test.cpp +++ b/be/test/ai/embed_test.cpp @@ -281,6 +281,154 @@ TEST(EMBED_TEST, embed_function_test) { } } +TEST(EMBED_TEST, prefer_embed_resource_properties) { + TQueryOptions query_options = create_fake_query_options(); + auto query_ctx = MockQueryContext::create(TUniqueId(), ExecEnv::GetInstance(), query_options); + + TAIResource ai_resource; + ai_resource.__set_endpoint("invalid://general-endpoint"); + ai_resource.__set_provider_type("OPENAI"); + ai_resource.__set_model_name("general-model"); + ai_resource.__set_api_key("general-api-key"); + ai_resource.__set_embed_endpoint("http://localhost";); + ai_resource.__set_embed_provider_type("MOCK"); + ai_resource.__set_embed_model_name("embed-model"); + ai_resource.__set_embed_api_key("embed-api-key"); + ai_resource.__set_temperature(0.5); + ai_resource.__set_max_tokens(16); + ai_resource.__set_max_retries(1); + ai_resource.__set_retry_delay_second(1); + ai_resource.__set_dimensions(514); + query_ctx->set_ai_resources( + std::map<std::string, TAIResource> {{"embed_resource", ai_resource}}); + + TQueryGlobals query_globals; + RuntimeState runtime_state(TUniqueId(), 0, query_options, query_globals, nullptr, + query_ctx.get()); + auto ctx = FunctionContext::create_context(&runtime_state, {}, {}); + + auto col_resource = ColumnHelper::create_column<DataTypeString>( + std::vector<std::string> {"embed_resource"}); + auto col_text = + ColumnHelper::create_column<DataTypeString>(std::vector<std::string> {"test input"}); + + Block block; + block.insert({std::move(col_resource), std::make_shared<DataTypeString>(), "resource"}); + block.insert({std::move(col_text), std::make_shared<DataTypeString>(), "text"}); + block.insert( + {nullptr, + std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeFloat32>())), + "result"}); + + auto embed_func = FunctionEmbed::create(); + Status exec_status = embed_func->execute_impl(ctx.get(), block, {0, 1}, 2, 1); + + ASSERT_TRUE(exec_status.ok()) << exec_status.to_string(); +} + +TEST(EMBED_TEST, select_resource_properties_by_input_type) { + TAIResource ai_resource; + ai_resource.__set_endpoint("general-endpoint"); + ai_resource.__set_provider_type("OPENAI"); + ai_resource.__set_model_name("general-model"); + ai_resource.__set_api_key("general-api-key"); + ai_resource.__set_embed_endpoint("embed-endpoint"); + ai_resource.__set_embed_provider_type("QWEN"); + ai_resource.__set_embed_model_name("embed-model"); + ai_resource.__set_embed_api_key("embed-api-key"); + ai_resource.__set_embed_mm_endpoint("multimodal-endpoint"); + ai_resource.__set_embed_mm_provider_type("GEMINI"); + ai_resource.__set_embed_mm_model_name("multimodal-model"); + ai_resource.__set_embed_mm_api_key("multimodal-api-key"); + + FunctionEmbed embed_function; + AIResource text_resource = + embed_function.select_ai_resource(ai_resource, PrimitiveType::TYPE_STRING); + EXPECT_EQ(text_resource.endpoint, "embed-endpoint"); + EXPECT_EQ(text_resource.provider_type, "QWEN"); + + AIResource multimodal_resource = + embed_function.select_ai_resource(ai_resource, PrimitiveType::TYPE_JSONB); + EXPECT_EQ(multimodal_resource.endpoint, "multimodal-endpoint"); + EXPECT_EQ(multimodal_resource.provider_type, "GEMINI"); +} + +TEST(EMBED_TEST, multimodal_resource_fallback) { + TAIResource ai_resource; + ai_resource.__set_endpoint("general-endpoint"); + ai_resource.__set_provider_type("OPENAI"); + ai_resource.__set_model_name("general-model"); + ai_resource.__set_api_key("general-api-key"); + ai_resource.__set_embed_endpoint("embed-endpoint"); + ai_resource.__set_embed_provider_type("QWEN"); + ai_resource.__set_embed_model_name("embed-model"); + ai_resource.__set_embed_api_key("embed-api-key"); + + FunctionEmbed embed_function; + AIResource embed_resource = + embed_function.select_ai_resource(ai_resource, PrimitiveType::TYPE_JSONB); + EXPECT_EQ(embed_resource.endpoint, "embed-endpoint"); + + TAIResource general_resource; + general_resource.__set_endpoint("general-endpoint"); + general_resource.__set_provider_type("OPENAI"); + general_resource.__set_model_name("general-model"); + general_resource.__set_api_key("general-api-key"); + AIResource fallback_resource = + embed_function.select_ai_resource(general_resource, PrimitiveType::TYPE_JSONB); + EXPECT_EQ(fallback_resource.endpoint, "general-endpoint"); + + general_resource.__set_embed_mm_endpoint("multimodal-endpoint"); + general_resource.__set_embed_mm_provider_type("GEMINI"); + general_resource.__set_embed_mm_model_name("multimodal-model"); + general_resource.__set_embed_mm_api_key("multimodal-api-key"); + AIResource text_resource = + embed_function.select_ai_resource(general_resource, PrimitiveType::TYPE_STRING); + EXPECT_EQ(text_resource.endpoint, "general-endpoint"); +} + +TEST(EMBED_TEST, incomplete_dedicated_resource_fallback) { + TAIResource ai_resource; + ai_resource.__set_endpoint("general-endpoint"); + ai_resource.__set_provider_type("OPENAI"); + ai_resource.__set_model_name("general-model"); + ai_resource.__set_api_key("general-api-key"); + ai_resource.__set_embed_endpoint("incomplete-embed-endpoint"); + + FunctionEmbed embed_function; + AIResource text_resource = + embed_function.select_ai_resource(ai_resource, PrimitiveType::TYPE_STRING); + EXPECT_EQ(text_resource.endpoint, "general-endpoint"); + EXPECT_EQ(text_resource.provider_type, "OPENAI"); + + ai_resource.__set_embed_provider_type("QWEN"); + ai_resource.__set_embed_model_name("embed-model"); + ai_resource.__set_embed_api_key("embed-api-key"); + ai_resource.__set_embed_mm_endpoint("incomplete-multimodal-endpoint"); + AIResource multimodal_resource = + embed_function.select_ai_resource(ai_resource, PrimitiveType::TYPE_JSONB); + EXPECT_EQ(multimodal_resource.endpoint, "incomplete-embed-endpoint"); + EXPECT_EQ(multimodal_resource.provider_type, "QWEN"); +} + +TEST(EMBED_TEST, empty_dedicated_resource_properties_fallback) { + TAIResource ai_resource; + ai_resource.__set_endpoint("general-endpoint"); + ai_resource.__set_provider_type("OPENAI"); + ai_resource.__set_model_name("general-model"); + ai_resource.__set_api_key("general-api-key"); + ai_resource.__set_embed_endpoint(""); + ai_resource.__set_embed_provider_type(""); + ai_resource.__set_embed_model_name(""); + ai_resource.__set_embed_api_key(""); + + FunctionEmbed embed_function; + AIResource text_resource = + embed_function.select_ai_resource(ai_resource, PrimitiveType::TYPE_STRING); + EXPECT_EQ(text_resource.endpoint, "general-endpoint"); + EXPECT_EQ(text_resource.provider_type, "OPENAI"); +} + TEST(EMBED_TEST, embed_function_text_multi_rows) { auto runtime_state = std::make_unique<MockRuntimeState>(); auto ctx = FunctionContext::create_context(runtime_state.get(), {}, {}); diff --git a/be/test/core/string_buffer_test.cpp b/be/test/core/string_buffer_test.cpp index 926e61ca721..7f230ab6c86 100644 --- a/be/test/core/string_buffer_test.cpp +++ b/be/test/core/string_buffer_test.cpp @@ -118,6 +118,15 @@ TEST(StringBufferTest, TestWriteReadBinary) { ASSERT_EQ(str_ref_val.to_string(), read_str_ref_val.to_string()); } +TEST(StringBufferTest, HasRemaining) { + StringRef ref("doris", 5); + BufferReadable reader(ref); + + EXPECT_TRUE(reader.has_remaining()); + reader.read(5); + EXPECT_FALSE(reader.has_remaining()); +} + // This test may fail due to a bug in read_var_uint, where it can read out of bounds. // The loop condition `i < 9` should probably be `i < len`. //TEST(StringBufferTest, TestVarUInt) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java index 346e58285fd..91553c0fc4d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/AIResource.java @@ -23,14 +23,14 @@ import org.apache.doris.datasource.property.constants.AIProperties; import org.apache.doris.thrift.TAIResource; import com.google.common.base.Preconditions; +import com.google.common.base.Strings; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.gson.annotations.SerializedName; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; import java.util.HashMap; +import java.util.List; import java.util.Map; /** @@ -54,7 +54,6 @@ import java.util.Map; */ public class AIResource extends Resource { - private static final Logger LOG = LogManager.getLogger(AIResource.class); @SerializedName(value = "properties") private Map<String, String> properties; @SerializedName(value = "createdByRoot") @@ -83,12 +82,6 @@ public class AIResource extends Resource { this.properties = Maps.newHashMap(newProperties); AIProperties.requiredAIProperties(properties); - - boolean needCheck = isNeedCheck(properties); - if (LOG.isDebugEnabled()) { - LOG.debug("AI resource need check validity: {}", needCheck); - } - AIProperties.optionalAIProperties(this.properties); } @@ -96,43 +89,54 @@ public class AIResource extends Resource { return properties.get(propertyKey); } - private boolean isNeedCheck(Map<String, String> newProperties) { - boolean needCheck = !this.properties.containsKey(AIProperties.VALIDITY_CHECK) - || Boolean.parseBoolean(this.properties.get(AIProperties.VALIDITY_CHECK)); + public boolean hasCompleteGeneralProperties() { + return AIProperties.REQUIRED_FIELDS.stream() + .allMatch(field -> !Strings.isNullOrEmpty(properties.get(field))) + && ("LOCAL".equalsIgnoreCase(properties.get(AIProperties.PROVIDER_TYPE)) + || !Strings.isNullOrEmpty(properties.get(AIProperties.API_KEY))); + } - if (newProperties != null && newProperties.containsKey(AIProperties.VALIDITY_CHECK)) { - needCheck = Boolean.parseBoolean(newProperties.get(AIProperties.VALIDITY_CHECK)); - } + public boolean hasCompleteEmbedProperties() { + return hasCompleteProperties(AIProperties.EMBED_REQUIRED_FIELDS, + AIProperties.EMBED_PROVIDER_TYPE, AIProperties.EMBED_API_KEY); + } - if ("LOCAL".equalsIgnoreCase(this.properties.getOrDefault(AIProperties.PROVIDER_TYPE, ""))) { - needCheck = false; - } - return needCheck; + public boolean hasCompleteMultimodalEmbedProperties() { + return hasCompleteProperties(AIProperties.MULTIMODAL_EMBED_REQUIRED_FIELDS, + AIProperties.MULTIMODAL_EMBED_PROVIDER_TYPE, AIProperties.MULTIMODAL_EMBED_API_KEY); + } + + private boolean hasCompleteProperties(List<String> requiredFields, + String providerTypeField, String apiKeyField) { + return requiredFields.stream() + .allMatch(field -> !Strings.isNullOrEmpty(properties.get(field))) + && ("LOCAL".equalsIgnoreCase(properties.get(providerTypeField)) + || !Strings.isNullOrEmpty(properties.get(apiKeyField))); } @Override public void modifyProperties(Map<String, String> properties) throws DdlException { - boolean needCheck = isNeedCheck(properties); - if (LOG.isDebugEnabled()) { - LOG.debug("AI resource need check validity: {}", needCheck); - } - - if (needCheck) { + writeLock(); + try { Map<String, String> changedProperties = new HashMap<>(this.properties); - changedProperties.putAll(properties); + for (Map.Entry<String, String> kv : properties.entrySet()) { + replaceIfEffectiveValue(changedProperties, kv.getKey(), kv.getValue()); + if (kv.getKey().equals(AIProperties.API_KEY) + || kv.getKey().equals(AIProperties.EMBED_API_KEY) + || kv.getKey().equals(AIProperties.MULTIMODAL_EMBED_API_KEY)) { + changedProperties.put(kv.getKey(), kv.getValue()); + } else if (kv.getKey().equals(AIProperties.EFFORT) + && Strings.isNullOrEmpty(kv.getValue())) { + changedProperties.remove(kv.getKey()); + } + } AIProperties.requiredAIProperties(changedProperties); - } - // modify properties - writeLock(); - for (Map.Entry<String, String> kv : properties.entrySet()) { - replaceIfEffectiveValue(this.properties, kv.getKey(), kv.getValue()); - if (kv.getKey().equals(AIProperties.API_KEY)) { - this.properties.put(kv.getKey(), kv.getValue()); - } + this.properties = changedProperties; + ++version; + } finally { + writeUnlock(); } - ++version; - writeUnlock(); super.modifyProperties(properties); } @@ -148,7 +152,9 @@ public class AIResource extends Resource { readLock(); result.addRow(Lists.newArrayList(name, lowerCaseType, "version", String.valueOf(version))); for (Map.Entry<String, String> entry : properties.entrySet()) { - if (entry.getKey().equals(AIProperties.API_KEY)) { + if (entry.getKey().equals(AIProperties.API_KEY) + || entry.getKey().equals(AIProperties.EMBED_API_KEY) + || entry.getKey().equals(AIProperties.MULTIMODAL_EMBED_API_KEY)) { result.addRow(Lists.newArrayList(name, lowerCaseType, entry.getKey(), "******")); } else { result.addRow(Lists.newArrayList(name, lowerCaseType, entry.getKey(), entry.getValue())); @@ -159,10 +165,45 @@ public class AIResource extends Resource { public TAIResource toThrift() throws NumberFormatException { TAIResource tAIResource = new TAIResource(); - tAIResource.setProviderType(properties.get(AIProperties.PROVIDER_TYPE)); - tAIResource.setEndpoint(properties.get(AIProperties.ENDPOINT)); - tAIResource.setApiKey(properties.get(AIProperties.API_KEY)); - tAIResource.setModelName(properties.get(AIProperties.MODEL_NAME)); + if (properties.containsKey(AIProperties.PROVIDER_TYPE)) { + tAIResource.setProviderType(properties.get(AIProperties.PROVIDER_TYPE)); + } + if (properties.containsKey(AIProperties.ENDPOINT)) { + tAIResource.setEndpoint(properties.get(AIProperties.ENDPOINT)); + } + if (properties.containsKey(AIProperties.API_KEY)) { + tAIResource.setApiKey(properties.get(AIProperties.API_KEY)); + } + if (properties.containsKey(AIProperties.MODEL_NAME)) { + tAIResource.setModelName(properties.get(AIProperties.MODEL_NAME)); + } + if (properties.containsKey(AIProperties.EMBED_PROVIDER_TYPE)) { + tAIResource.setEmbedProviderType(properties.get(AIProperties.EMBED_PROVIDER_TYPE)); + } + if (properties.containsKey(AIProperties.EMBED_ENDPOINT)) { + tAIResource.setEmbedEndpoint(properties.get(AIProperties.EMBED_ENDPOINT)); + } + if (properties.containsKey(AIProperties.EMBED_API_KEY)) { + tAIResource.setEmbedApiKey(properties.get(AIProperties.EMBED_API_KEY)); + } + if (properties.containsKey(AIProperties.EMBED_MODEL_NAME)) { + tAIResource.setEmbedModelName(properties.get(AIProperties.EMBED_MODEL_NAME)); + } + if (properties.containsKey(AIProperties.MULTIMODAL_EMBED_PROVIDER_TYPE)) { + tAIResource.setEmbedMmProviderType(properties.get(AIProperties.MULTIMODAL_EMBED_PROVIDER_TYPE)); + } + if (properties.containsKey(AIProperties.MULTIMODAL_EMBED_ENDPOINT)) { + tAIResource.setEmbedMmEndpoint(properties.get(AIProperties.MULTIMODAL_EMBED_ENDPOINT)); + } + if (properties.containsKey(AIProperties.MULTIMODAL_EMBED_API_KEY)) { + tAIResource.setEmbedMmApiKey(properties.get(AIProperties.MULTIMODAL_EMBED_API_KEY)); + } + if (properties.containsKey(AIProperties.MULTIMODAL_EMBED_MODEL_NAME)) { + tAIResource.setEmbedMmModelName(properties.get(AIProperties.MULTIMODAL_EMBED_MODEL_NAME)); + } + if (!Strings.isNullOrEmpty(properties.get(AIProperties.EFFORT))) { + tAIResource.setEffort(properties.get(AIProperties.EFFORT)); + } tAIResource.setAnthropicVersion(properties.get(AIProperties.ANTHROPIC_VERSION)); try { diff --git a/fe/fe-core/src/main/java/org/apache/doris/common/util/PrintableMap.java b/fe/fe-core/src/main/java/org/apache/doris/common/util/PrintableMap.java index 6a2a3170cef..813fb2b9839 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/common/util/PrintableMap.java +++ b/fe/fe-core/src/main/java/org/apache/doris/common/util/PrintableMap.java @@ -61,6 +61,9 @@ public class PrintableMap<K, V> { SENSITIVE_KEY.add("bos_secret_accesskey"); SENSITIVE_KEY.add("jdbc.password"); SENSITIVE_KEY.add("elasticsearch.password"); + SENSITIVE_KEY.add("ai.api_key"); + SENSITIVE_KEY.add("ai.embed.api_key"); + SENSITIVE_KEY.add("ai.embed.mm.api_key"); SENSITIVE_KEY.add("lance.rest.bearer-token"); SENSITIVE_KEY.add("lance.rest.api-key"); SENSITIVE_KEY.addAll(Arrays.asList( diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java index 6019aaef303..43890d5110a 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/property/constants/AIProperties.java @@ -32,15 +32,24 @@ public class AIProperties extends BaseProperties { public static final String ENDPOINT = "ai.endpoint"; public static final String PROVIDER_TYPE = "ai.provider_type"; public static final String MODEL_NAME = "ai.model_name"; + public static final String EMBED_ENDPOINT = "ai.embed.endpoint"; + public static final String EMBED_PROVIDER_TYPE = "ai.embed.provider_type"; + public static final String EMBED_MODEL_NAME = "ai.embed.model_name"; + public static final String MULTIMODAL_EMBED_ENDPOINT = "ai.embed.mm.endpoint"; + public static final String MULTIMODAL_EMBED_PROVIDER_TYPE = "ai.embed.mm.provider_type"; + public static final String MULTIMODAL_EMBED_MODEL_NAME = "ai.embed.mm.model_name"; // optional public static final String API_KEY = "ai.api_key"; + public static final String EMBED_API_KEY = "ai.embed.api_key"; + public static final String MULTIMODAL_EMBED_API_KEY = "ai.embed.mm.api_key"; public static final String TEMPERATURE = "ai.temperature"; public static final String MAX_TOKEN = "ai.max_token"; public static final String MAX_RETRIES = "ai.max_retries"; public static final String RETRY_DELAY_SECOND = "ai.retry_delay_second"; public static final String ANTHROPIC_VERSION = "ai.anthropic_version"; public static final String DIMENSIONS = "ai.dimensions"; + public static final String EFFORT = "ai.effort"; // default_val public static final String DEFAULT_TEMPERATURE = "-1"; @@ -53,29 +62,54 @@ public class AIProperties extends BaseProperties { public static final String VALIDITY_CHECK = "ai.validity_check"; public static final List<String> REQUIRED_FIELDS = Arrays.asList(ENDPOINT, PROVIDER_TYPE, MODEL_NAME); + public static final List<String> EMBED_REQUIRED_FIELDS = + Arrays.asList(EMBED_ENDPOINT, EMBED_PROVIDER_TYPE, EMBED_MODEL_NAME); + public static final List<String> MULTIMODAL_EMBED_REQUIRED_FIELDS = + Arrays.asList(MULTIMODAL_EMBED_ENDPOINT, MULTIMODAL_EMBED_PROVIDER_TYPE, + MULTIMODAL_EMBED_MODEL_NAME); public static final List<String> PROVIDERS = Arrays.asList("OPENAI", "LOCAL", "GEMINI", "DEEPSEEK", "ANTHROPIC", "MOONSHOT", "QWEN", "MINIMAX", "ZHIPU", "BAICHUAN", "VOYAGEAI", "JINA"); + private static final List<String> ALL_EFFORT_LEVELS = + Arrays.asList("none", "minimal", "low", "medium", "high", "xhigh", "max"); + private static final List<String> ANTHROPIC_EFFORT_LEVELS = + Arrays.asList("low", "medium", "high", "xhigh", "max"); + private static final List<String> GEMINI_EFFORT_LEVELS = + Arrays.asList("minimal", "low", "medium", "high"); public static void requiredAIProperties(Map<String, String> properties) throws DdlException { - // Check required field - for (String field : REQUIRED_FIELDS) { - if (Strings.isNullOrEmpty(properties.get(field))) { - throw new DdlException("Missing [" + field + "] in properties."); - } + boolean hasGeneralProperties = hasAnyProperty(properties, REQUIRED_FIELDS, API_KEY); + boolean hasEmbedProperties = hasAnyProperty(properties, EMBED_REQUIRED_FIELDS, EMBED_API_KEY); + boolean hasMultimodalEmbedProperties = hasAnyProperty(properties, + MULTIMODAL_EMBED_REQUIRED_FIELDS, MULTIMODAL_EMBED_API_KEY); + if (!hasGeneralProperties && !hasEmbedProperties && !hasMultimodalEmbedProperties) { + throw new DdlException("At least one complete AI property group must be configured."); } - // Check the provider is valid - properties.put(PROVIDER_TYPE, properties.get(PROVIDER_TYPE).toUpperCase()); - if (PROVIDERS.stream().noneMatch(s -> s.equals(properties.get(PROVIDER_TYPE).toUpperCase()))) { - throw new DdlException("Provider must be one of " + PROVIDERS); + if (hasGeneralProperties) { + validatePropertyGroup(properties, REQUIRED_FIELDS, PROVIDER_TYPE, API_KEY); + } + if (hasEmbedProperties) { + validatePropertyGroup(properties, EMBED_REQUIRED_FIELDS, EMBED_PROVIDER_TYPE, EMBED_API_KEY); + } + if (hasMultimodalEmbedProperties) { + validatePropertyGroup(properties, MULTIMODAL_EMBED_REQUIRED_FIELDS, + MULTIMODAL_EMBED_PROVIDER_TYPE, MULTIMODAL_EMBED_API_KEY); } - // Only the 'local' provider can ignore the 'api-key' - if (!"LOCAL".equals(properties.get(AIProperties.PROVIDER_TYPE)) - && Strings.isNullOrEmpty(properties.get(AIProperties.API_KEY))) { - throw new DdlException("Missing [" + API_KEY + "] in properties for provider: " - + properties.get(AIProperties.PROVIDER_TYPE)); + String effort = properties.get(EFFORT); + if (!Strings.isNullOrEmpty(effort)) { + String providerType = properties.get(PROVIDER_TYPE); + List<String> effortLevels = ALL_EFFORT_LEVELS; + if ("ANTHROPIC".equals(providerType)) { + effortLevels = ANTHROPIC_EFFORT_LEVELS; + } else if ("GEMINI".equals(providerType)) { + effortLevels = GEMINI_EFFORT_LEVELS; + } + if (!effortLevels.contains(effort)) { + throw new DdlException("[" + EFFORT + "] must be one of " + effortLevels + + " for provider: " + providerType); + } } // Check weather the 'temperature' is valid @@ -89,7 +123,7 @@ public class AIProperties extends BaseProperties { // Check 'dimensions' temp = properties.get(AIProperties.DIMENSIONS); - if (!Strings.isNullOrEmpty(temp) && temp.equals("-1")) { + if (!Strings.isNullOrEmpty(temp) && !temp.equals("-1")) { int tempVal = Integer.parseInt(temp); if (tempVal <= 0) { throw new DdlException("Dimensions must be a positive integer"); @@ -97,6 +131,30 @@ public class AIProperties extends BaseProperties { } } + private static boolean hasAnyProperty(Map<String, String> properties, List<String> requiredFields, + String apiKeyField) { + return properties.containsKey(apiKeyField) || requiredFields.stream().anyMatch(properties::containsKey); + } + + private static void validatePropertyGroup(Map<String, String> properties, List<String> requiredFields, + String providerTypeField, String apiKeyField) throws DdlException { + for (String field : requiredFields) { + if (Strings.isNullOrEmpty(properties.get(field))) { + throw new DdlException("Missing [" + field + "] in properties."); + } + } + + String providerType = properties.get(providerTypeField).toUpperCase(); + properties.put(providerTypeField, providerType); + if (!PROVIDERS.contains(providerType)) { + throw new DdlException("Provider must be one of " + PROVIDERS); + } + + if (!"LOCAL".equals(providerType) && Strings.isNullOrEmpty(properties.get(apiKeyField))) { + throw new DdlException("Missing [" + apiKeyField + "] in properties for provider: " + providerType); + } + } + public static void optionalAIProperties(Map<String, String> properties) { properties.putIfAbsent(AIProperties.TEMPERATURE, AIProperties.DEFAULT_TEMPERATURE); properties.putIfAbsent(AIProperties.MAX_TOKEN, AIProperties.DEFAULT_MAX_TOKEN); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java index 833f1b7db4d..40238a936ad 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/AIAgg.java @@ -90,6 +90,10 @@ public class AIAgg extends NullableAggregateFunction if (!(resource instanceof AIResource)) { throw new AnalysisException("AI resource '" + resourceName + "' does not exist"); } + if (!((AIResource) resource).hasCompleteGeneralProperties()) { + throw new AnalysisException("AI resource '" + resourceName + + "' does not contain complete general AI properties"); + } Resource.registerUsedAIResourceName(resourceName); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java index 7399f2a348f..22b0bfe825a 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/AIFunction.java @@ -61,6 +61,10 @@ public abstract class AIFunction extends ScalarFunction if (!(resource instanceof AIResource)) { throw new AnalysisException("AI resource '" + resourceName + "' does not exist"); } + if (!((AIResource) resource).hasCompleteGeneralProperties()) { + throw new AnalysisException("AI resource '" + resourceName + + "' does not contain complete general AI properties"); + } Resource.registerUsedAIResourceName(resourceName); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java index e3e8725856f..e6bb28a82bd 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/ai/Embed.java @@ -102,12 +102,17 @@ public class Embed extends AIFunction { if (arity() == 2) { String aiResourceName = requireStringLiteral(child(0), "resource name", "AI Function must accept literal for the resource name."); - validateAIResource(aiResourceName); + validateAIResource(aiResourceName, getArgument(1).getDataType().isJsonType()); return; } throw new AnalysisException("Function EMBED only accepts 1 or 2 arguments"); } + @Override + public void checkLegalityAfterRewrite() { + checkLegalityBeforeTypeCoercion(); + } + private static String requireStringLiteral(Expression arg, String argName, String errorMsg) { if (!(arg instanceof StringLikeLiteral)) { throw new AnalysisException(errorMsg); @@ -119,11 +124,19 @@ public class Embed extends AIFunction { return value; } - private static void validateAIResource(String resourceName) { + private static void validateAIResource(String resourceName, boolean isMultimodal) { Resource resource = Env.getCurrentEnv().getResourceMgr().getResource(resourceName); if (!(resource instanceof AIResource)) { throw new AnalysisException("AI resource '" + resourceName + "' does not exist"); } + AIResource aiResource = (AIResource) resource; + boolean hasSupportedProperties = aiResource.hasCompleteEmbedProperties() + || aiResource.hasCompleteGeneralProperties() + || (isMultimodal && aiResource.hasCompleteMultimodalEmbedProperties()); + if (!hasSupportedProperties) { + throw new AnalysisException("AI resource '" + resourceName + + "' does not have properties required by EMBED"); + } Resource.registerUsedAIResourceName(resourceName); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java index b191c324990..680c94ba584 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/catalog/AIResourceTest.java @@ -23,10 +23,17 @@ import org.apache.doris.common.FeConstants; import org.apache.doris.common.FeMetaVersion; import org.apache.doris.common.UserException; import org.apache.doris.common.io.Text; +import org.apache.doris.common.proc.BaseProcResult; import org.apache.doris.datasource.property.constants.AIProperties; import org.apache.doris.meta.MetaContext; import org.apache.doris.mysql.privilege.AccessControllerManager; import org.apache.doris.mysql.privilege.PrivPredicate; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.functions.agg.AIAgg; +import org.apache.doris.nereids.trees.expressions.functions.ai.AISentiment; +import org.apache.doris.nereids.trees.expressions.functions.ai.Embed; +import org.apache.doris.nereids.trees.expressions.literal.JsonLiteral; +import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; import org.apache.doris.nereids.trees.plans.commands.CreateResourceCommand; import org.apache.doris.nereids.trees.plans.commands.info.CreateResourceInfo; import org.apache.doris.persist.gson.GsonUtils; @@ -43,14 +50,24 @@ import org.apache.logging.log4j.Logger; import org.junit.Assert; import org.junit.Before; import org.junit.Test; +import org.junit.jupiter.api.Assertions; +import org.mockito.MockedStatic; +import org.mockito.Mockito; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.io.DataInputStream; import java.io.DataOutputStream; import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; +import java.util.Arrays; import java.util.HashMap; +import java.util.List; import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; public class AIResourceTest { private static final Logger LOG = LogManager.getLogger(AIResourceTest.class); @@ -150,6 +167,371 @@ public class AIResourceTest { Assert.assertEquals(retryDelaySecond, aiResource.getProperty(AIProperties.RETRY_DELAY_SECOND)); } + @Test + public void testEmbedOnlyResource() throws DdlException { + Map<String, String> properties = new HashMap<>(); + properties.put(AIProperties.EMBED_ENDPOINT, "https://api.example.com/v1/embeddings";); + properties.put(AIProperties.EMBED_PROVIDER_TYPE, "openai"); + properties.put(AIProperties.EMBED_MODEL_NAME, "text-embedding-model"); + properties.put(AIProperties.EMBED_API_KEY, "embed-api-key"); + + AIResource aiResource = new AIResource("embed-only-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + + Assertions.assertEquals("OPENAI", aiResource.getProperty(AIProperties.EMBED_PROVIDER_TYPE)); + Assertions.assertEquals("https://api.example.com/v1/embeddings";, + aiResource.toThrift().getEmbedEndpoint()); + Assertions.assertEquals("OPENAI", aiResource.toThrift().getEmbedProviderType()); + Assertions.assertEquals("text-embedding-model", aiResource.toThrift().getEmbedModelName()); + Assertions.assertEquals("embed-api-key", aiResource.toThrift().getEmbedApiKey()); + Assertions.assertFalse(aiResource.toThrift().isSetEndpoint()); + } + + @Test + public void testMultimodalEmbedOnlyResource() throws DdlException { + Map<String, String> properties = new HashMap<>(); + properties.put(AIProperties.MULTIMODAL_EMBED_ENDPOINT, + "https://api.example.com/v1/multimodal-embeddings";); + properties.put(AIProperties.MULTIMODAL_EMBED_PROVIDER_TYPE, "qwen"); + properties.put(AIProperties.MULTIMODAL_EMBED_MODEL_NAME, "multimodal-embedding-model"); + properties.put(AIProperties.MULTIMODAL_EMBED_API_KEY, "multimodal-embed-api-key"); + + AIResource aiResource = new AIResource("multimodal-embed-only-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + + Assertions.assertEquals("QWEN", aiResource.getProperty(AIProperties.MULTIMODAL_EMBED_PROVIDER_TYPE)); + Assertions.assertEquals("https://api.example.com/v1/multimodal-embeddings";, + aiResource.getProperty(AIProperties.MULTIMODAL_EMBED_ENDPOINT)); + Assertions.assertEquals("multimodal-embedding-model", + aiResource.getProperty(AIProperties.MULTIMODAL_EMBED_MODEL_NAME)); + Assertions.assertEquals("multimodal-embed-api-key", + aiResource.getProperty(AIProperties.MULTIMODAL_EMBED_API_KEY)); + Assertions.assertEquals("https://api.example.com/v1/multimodal-embeddings";, + aiResource.toThrift().getEmbedMmEndpoint()); + Assertions.assertEquals("QWEN", aiResource.toThrift().getEmbedMmProviderType()); + Assertions.assertEquals("multimodal-embedding-model", aiResource.toThrift().getEmbedMmModelName()); + Assertions.assertEquals("multimodal-embed-api-key", aiResource.toThrift().getEmbedMmApiKey()); + Assertions.assertFalse(aiResource.toThrift().isSetEndpoint()); + Assertions.assertFalse(aiResource.toThrift().isSetEmbedEndpoint()); + + BaseProcResult result = new BaseProcResult(); + aiResource.getProcNodeData(result); + Assertions.assertTrue(result.getRows().stream().anyMatch(row -> + AIProperties.MULTIMODAL_EMBED_API_KEY.equals(row.get(2)) && "******".equals(row.get(3)))); + Assertions.assertFalse(result.getRows().stream().anyMatch(row -> row.contains("multimodal-embed-api-key"))); + } + + @Test + public void testRejectPartialMultimodalEmbedProperties() { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.MULTIMODAL_EMBED_ENDPOINT, + "https://api.example.com/v1/multimodal-embeddings";); + + AIResource aiResource = new AIResource("partial-multimodal-embed-resource"); + DdlException exception = Assertions.assertThrows(DdlException.class, + () -> aiResource.setProperties(ImmutableMap.copyOf(properties))); + Assertions.assertTrue(exception.getMessage().contains(AIProperties.MULTIMODAL_EMBED_PROVIDER_TYPE)); + } + + @Test + public void testRejectEmbedOnlyResourceForNonEmbedScalarFunction() throws DdlException { + AIResource aiResource = createEmbedOnlyResource(); + try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) { + Env env = Mockito.mock(Env.class); + ResourceMgr resourceMgr = Mockito.mock(ResourceMgr.class); + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + Mockito.when(env.getResourceMgr()).thenReturn(resourceMgr); + Mockito.when(resourceMgr.getResource("embed-only-resource")).thenReturn(aiResource); + + AISentiment function = new AISentiment(new StringLiteral("embed-only-resource"), + new StringLiteral("text")); + Assertions.assertThrows(AnalysisException.class, function::checkLegalityAfterRewrite); + } + } + + @Test + public void testRejectEmbedOnlyResourceForAiAgg() throws DdlException { + AIResource aiResource = createEmbedOnlyResource(); + try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) { + Env env = Mockito.mock(Env.class); + ResourceMgr resourceMgr = Mockito.mock(ResourceMgr.class); + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + Mockito.when(env.getResourceMgr()).thenReturn(resourceMgr); + Mockito.when(resourceMgr.getResource("embed-only-resource")).thenReturn(aiResource); + + AIAgg function = new AIAgg(new StringLiteral("embed-only-resource"), + new StringLiteral("text"), new StringLiteral("task")); + Assertions.assertThrows(AnalysisException.class, function::checkLegalityAfterRewrite); + } + } + + @Test + public void testAcceptEmbedOnlyResourceForEmbedFunction() throws DdlException { + AIResource aiResource = createEmbedOnlyResource(); + try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) { + Env env = Mockito.mock(Env.class); + ResourceMgr resourceMgr = Mockito.mock(ResourceMgr.class); + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + Mockito.when(env.getResourceMgr()).thenReturn(resourceMgr); + Mockito.when(resourceMgr.getResource("embed-only-resource")).thenReturn(aiResource); + + Embed function = new Embed(new StringLiteral("embed-only-resource"), + new StringLiteral("text")); + Assertions.assertDoesNotThrow(function::checkLegalityBeforeTypeCoercion); + Assertions.assertDoesNotThrow(function::checkLegalityAfterRewrite); + } + } + + @Test + public void testAcceptMultimodalEmbedOnlyResourceForJsonEmbed() throws DdlException { + AIResource aiResource = createMultimodalEmbedOnlyResource(); + try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) { + Env env = Mockito.mock(Env.class); + ResourceMgr resourceMgr = Mockito.mock(ResourceMgr.class); + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + Mockito.when(env.getResourceMgr()).thenReturn(resourceMgr); + Mockito.when(resourceMgr.getResource("multimodal-embed-only-resource")).thenReturn(aiResource); + + Embed function = new Embed(new StringLiteral("multimodal-embed-only-resource"), + new JsonLiteral("{\"text\":\"hello\"}")); + Assertions.assertDoesNotThrow(function::checkLegalityBeforeTypeCoercion); + Assertions.assertDoesNotThrow(function::checkLegalityAfterRewrite); + } + } + + @Test + public void testRejectMultimodalEmbedOnlyResourceForTextEmbed() throws DdlException { + AIResource aiResource = createMultimodalEmbedOnlyResource(); + try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) { + Env env = Mockito.mock(Env.class); + ResourceMgr resourceMgr = Mockito.mock(ResourceMgr.class); + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + Mockito.when(env.getResourceMgr()).thenReturn(resourceMgr); + Mockito.when(resourceMgr.getResource("multimodal-embed-only-resource")).thenReturn(aiResource); + + Embed function = new Embed(new StringLiteral("multimodal-embed-only-resource"), + new StringLiteral("hello")); + Assertions.assertThrows(AnalysisException.class, function::checkLegalityBeforeTypeCoercion); + } + } + + private AIResource createEmbedOnlyResource() throws DdlException { + Map<String, String> properties = new HashMap<>(); + properties.put(AIProperties.EMBED_ENDPOINT, "https://api.example.com/v1/embeddings";); + properties.put(AIProperties.EMBED_PROVIDER_TYPE, "openai"); + properties.put(AIProperties.EMBED_MODEL_NAME, "text-embedding-model"); + properties.put(AIProperties.EMBED_API_KEY, "embed-api-key"); + + AIResource aiResource = new AIResource("embed-only-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + return aiResource; + } + + private AIResource createMultimodalEmbedOnlyResource() throws DdlException { + Map<String, String> properties = new HashMap<>(); + properties.put(AIProperties.MULTIMODAL_EMBED_ENDPOINT, + "https://api.example.com/v1/multimodal-embeddings";); + properties.put(AIProperties.MULTIMODAL_EMBED_PROVIDER_TYPE, "qwen"); + properties.put(AIProperties.MULTIMODAL_EMBED_MODEL_NAME, "multimodal-embedding-model"); + properties.put(AIProperties.MULTIMODAL_EMBED_API_KEY, "multimodal-embed-api-key"); + + AIResource aiResource = new AIResource("multimodal-embed-only-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + return aiResource; + } + + @Test + public void testLocalEmbedResourceWithoutApiKey() throws DdlException { + Map<String, String> properties = new HashMap<>(); + properties.put(AIProperties.EMBED_ENDPOINT, "http://localhost:8000/v1/embeddings";); + properties.put(AIProperties.EMBED_PROVIDER_TYPE, "local"); + properties.put(AIProperties.EMBED_MODEL_NAME, "local-embedding-model"); + + AIResource aiResource = new AIResource("local-embed-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + + Assertions.assertEquals("LOCAL", aiResource.getProperty(AIProperties.EMBED_PROVIDER_TYPE)); + } + + @Test + public void testRejectEmbedResourceWithoutApiKey() { + Map<String, String> properties = new HashMap<>(); + properties.put(AIProperties.EMBED_ENDPOINT, "https://api.example.com/v1/embeddings";); + properties.put(AIProperties.EMBED_PROVIDER_TYPE, "openai"); + properties.put(AIProperties.EMBED_MODEL_NAME, "text-embedding-model"); + + AIResource aiResource = new AIResource("embed-resource-without-api-key"); + DdlException exception = Assertions.assertThrows(DdlException.class, + () -> aiResource.setProperties(ImmutableMap.copyOf(properties))); + Assertions.assertTrue(exception.getMessage().contains("ai.embed.api_key")); + } + + @Test + public void testMaskEmbedApiKey() throws DdlException { + Map<String, String> properties = new HashMap<>(); + properties.put(AIProperties.EMBED_ENDPOINT, "https://api.example.com/v1/embeddings";); + properties.put(AIProperties.EMBED_PROVIDER_TYPE, "openai"); + properties.put(AIProperties.EMBED_MODEL_NAME, "text-embedding-model"); + properties.put(AIProperties.EMBED_API_KEY, "embed-api-key"); + + AIResource aiResource = new AIResource("embed-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + BaseProcResult result = new BaseProcResult(); + aiResource.getProcNodeData(result); + + Assertions.assertTrue(result.getRows().stream().anyMatch(row -> + AIProperties.EMBED_API_KEY.equals(row.get(2)) && "******".equals(row.get(3)))); + Assertions.assertFalse(result.getRows().stream().anyMatch(row -> row.contains("embed-api-key"))); + } + + @Test + public void testRejectPartialEmbedProperties() throws DdlException { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.EMBED_ENDPOINT, "https://api.example.com/v1/embeddings";); + + AIResource aiResource = new AIResource("partial-embed-resource"); + DdlException exception = Assertions.assertThrows(DdlException.class, + () -> aiResource.setProperties(ImmutableMap.copyOf(properties))); + Assertions.assertTrue(exception.getMessage().contains("ai.embed.provider_type")); + } + + @Test + public void testRejectResourceWithoutCompletePropertyGroup() { + Map<String, String> properties = new HashMap<>(); + properties.put("ai.validity_check", "false"); + + AIResource aiResource = new AIResource("empty-ai-resource"); + Assertions.assertThrows(DdlException.class, + () -> aiResource.setProperties(ImmutableMap.copyOf(properties))); + } + + @Test + public void testRejectInvalidEffort() { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.EFFORT, "invalid"); + + AIResource aiResource = new AIResource("invalid-effort-resource"); + DdlException exception = Assertions.assertThrows(DdlException.class, + () -> aiResource.setProperties(ImmutableMap.copyOf(properties))); + Assertions.assertTrue(exception.getMessage().contains(AIProperties.EFFORT)); + } + + @Test + public void testOpenAiDeepSeekAndQwenAcceptAllEffortLevels() throws DdlException { + List<String> providers = Arrays.asList("OPENAI", "DEEPSEEK", "QWEN"); + List<String> effortLevels = Arrays.asList("none", "minimal", "low", "medium", "high", "xhigh", "max"); + for (String provider : providers) { + for (String effort : effortLevels) { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.PROVIDER_TYPE, provider); + properties.put(AIProperties.EFFORT, effort); + + AIResource aiResource = new AIResource("effort-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + + Assertions.assertEquals(effort, aiResource.toThrift().getEffort()); + } + } + } + + @Test + public void testAnthropicEffortLevels() throws DdlException { + for (String effort : Arrays.asList("low", "medium", "high", "xhigh", "max")) { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.PROVIDER_TYPE, "ANTHROPIC"); + properties.put(AIProperties.EFFORT, effort); + + AIResource aiResource = new AIResource("effort-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + + Assertions.assertEquals(effort, aiResource.toThrift().getEffort()); + } + + for (String effort : Arrays.asList("none", "minimal")) { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.PROVIDER_TYPE, "ANTHROPIC"); + properties.put(AIProperties.EFFORT, effort); + + AIResource aiResource = new AIResource("invalid-anthropic-effort-resource"); + Assertions.assertThrows(DdlException.class, + () -> aiResource.setProperties(ImmutableMap.copyOf(properties))); + } + } + + @Test + public void testGeminiEffortLevels() throws DdlException { + for (String effort : Arrays.asList("minimal", "low", "medium", "high")) { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.PROVIDER_TYPE, "GEMINI"); + properties.put(AIProperties.EFFORT, effort); + + AIResource aiResource = new AIResource("effort-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + + Assertions.assertEquals(effort, aiResource.toThrift().getEffort()); + } + + for (String effort : Arrays.asList("none", "xhigh", "max")) { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.PROVIDER_TYPE, "GEMINI"); + properties.put(AIProperties.EFFORT, effort); + + AIResource aiResource = new AIResource("invalid-gemini-effort-resource"); + Assertions.assertThrows(DdlException.class, + () -> aiResource.setProperties(ImmutableMap.copyOf(properties))); + } + } + + @Test + public void testOtherProvidersAcceptAllEffortLevels() throws DdlException { + List<String> providers = Arrays.asList("LOCAL", "MOONSHOT", "MINIMAX", "ZHIPU", "BAICHUAN", "VOYAGEAI", "JINA"); + List<String> effortLevels = Arrays.asList("none", "minimal", "low", "medium", "high", "xhigh", "max"); + for (String provider : providers) { + for (String effort : effortLevels) { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.PROVIDER_TYPE, provider); + properties.put(AIProperties.EFFORT, effort); + + AIResource aiResource = new AIResource("effort-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + + Assertions.assertEquals(effort, aiResource.toThrift().getEffort()); + } + } + } + + @Test + public void testEmptyEffortIsNotForwarded() throws DdlException { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.EFFORT, ""); + + AIResource aiResource = new AIResource("empty-effort-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + + Assertions.assertFalse(aiResource.toThrift().isSetEffort()); + } + + @Test + public void testClearEffortOnModifyAndPersistence() throws Exception { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.EFFORT, "high"); + + AIResource aiResource = new AIResource("clear-effort-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + aiResource.modifyProperties(ImmutableMap.of(AIProperties.EFFORT, "")); + + Assertions.assertNull(aiResource.getProperty(AIProperties.EFFORT)); + Assertions.assertFalse(aiResource.toThrift().isSetEffort()); + + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + aiResource.write(new DataOutputStream(bytes)); + AIResource restoredResource = (AIResource) Resource.read( + new DataInputStream(new ByteArrayInputStream(bytes.toByteArray()))); + + Assertions.assertNull(restoredResource.getProperty(AIProperties.EFFORT)); + Assertions.assertFalse(restoredResource.toThrift().isSetEffort()); + } + @Test public void testAnthropic(@Mocked Env env, @Injectable AccessControllerManager accessManager) throws UserException { @@ -299,6 +681,139 @@ public class AIResourceTest { Assert.assertEquals("0.9", aiResource.getProperty(AIProperties.TEMPERATURE)); } + @Test + public void testModifyPropertiesPersistsNormalizedProviders() throws Exception { + Map<String, String> properties = new HashMap<>(aiProperties); + properties.put(AIProperties.VALIDITY_CHECK, "true"); + properties.put(AIProperties.EMBED_ENDPOINT, "https://api.openai.com/v1/embeddings";); + properties.put(AIProperties.EMBED_PROVIDER_TYPE, "openai"); + properties.put(AIProperties.EMBED_MODEL_NAME, "text-embedding-3-small"); + properties.put(AIProperties.EMBED_API_KEY, "embed-api-key"); + properties.put(AIProperties.DIMENSIONS, "8"); + + AIResource aiResource = new AIResource("normalized-provider-resource"); + aiResource.setProperties(ImmutableMap.copyOf(properties)); + aiResource.modifyProperties(ImmutableMap.of( + AIProperties.PROVIDER_TYPE, "openai", + AIProperties.EMBED_PROVIDER_TYPE, "qwen")); + + Assertions.assertEquals("OPENAI", aiResource.getProperty(AIProperties.PROVIDER_TYPE)); + Assertions.assertEquals("QWEN", aiResource.getProperty(AIProperties.EMBED_PROVIDER_TYPE)); + Assertions.assertEquals("OPENAI", aiResource.toThrift().getProviderType()); + Assertions.assertEquals("QWEN", aiResource.toThrift().getEmbedProviderType()); + + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + aiResource.write(new DataOutputStream(bytes)); + AIResource restoredResource = (AIResource) Resource.read( + new DataInputStream(new ByteArrayInputStream(bytes.toByteArray()))); + + Assertions.assertEquals("OPENAI", restoredResource.getProperty(AIProperties.PROVIDER_TYPE)); + Assertions.assertEquals("QWEN", restoredResource.getProperty(AIProperties.EMBED_PROVIDER_TYPE)); + Assertions.assertEquals("OPENAI", restoredResource.toThrift().getProviderType()); + Assertions.assertEquals("QWEN", restoredResource.toThrift().getEmbedProviderType()); + } + + @Test + public void testModifyPropertiesNormalizesEmbedProviderForLocalResource() throws Exception { + AIResource aiResource = new AIResource("local-resource"); + aiResource.setProperties(ImmutableMap.of( + AIProperties.ENDPOINT, "http://127.0.0.1:8000/v1/chat/completions";, + AIProperties.PROVIDER_TYPE, "local", + AIProperties.MODEL_NAME, "local-model", + AIProperties.DIMENSIONS, "8")); + + aiResource.modifyProperties(ImmutableMap.of( + AIProperties.EMBED_ENDPOINT, "http://127.0.0.1:8000/v1/embeddings";, + AIProperties.EMBED_PROVIDER_TYPE, "openai", + AIProperties.EMBED_MODEL_NAME, "text-embedding-3-small", + AIProperties.EMBED_API_KEY, "embed-api-key")); + + Assertions.assertEquals("OPENAI", aiResource.getProperty(AIProperties.EMBED_PROVIDER_TYPE)); + Assertions.assertEquals("OPENAI", aiResource.toThrift().getEmbedProviderType()); + } + + @Test + public void testModifyPropertiesNormalizesMultimodalProviderWhenValidityCheckDisabled() + throws Exception { + AIResource aiResource = new AIResource("validity-check-disabled-resource"); + aiProperties.put(AIProperties.DIMENSIONS, "8"); + aiResource.setProperties(ImmutableMap.copyOf(aiProperties)); + + aiResource.modifyProperties(ImmutableMap.of( + AIProperties.MULTIMODAL_EMBED_ENDPOINT, "https://example.com/multimodal-embeddings";, + AIProperties.MULTIMODAL_EMBED_PROVIDER_TYPE, "qwen", + AIProperties.MULTIMODAL_EMBED_MODEL_NAME, "qwen3-vl-embedding", + AIProperties.MULTIMODAL_EMBED_API_KEY, "multimodal-api-key")); + + Assertions.assertEquals("QWEN", + aiResource.getProperty(AIProperties.MULTIMODAL_EMBED_PROVIDER_TYPE)); + Assertions.assertEquals("QWEN", aiResource.toThrift().getEmbedMmProviderType()); + } + + @Test + public void testConcurrentModifyPropertiesPreservesIndependentUpdates() throws Exception { + AIResource aiResource = new AIResource("concurrent-resource"); + aiProperties.put(AIProperties.DIMENSIONS, "8"); + aiResource.setProperties(ImmutableMap.copyOf(aiProperties)); + + CountDownLatch ready = new CountDownLatch(2); + CountDownLatch start = new CountDownLatch(1); + AtomicReference<Throwable> failure = new AtomicReference<>(); + Thread updateTemperature = new Thread(() -> modifyPropertiesAfterStart( + aiResource, ImmutableMap.of(AIProperties.TEMPERATURE, "0.8"), ready, start, failure)); + Thread updateMaxToken = new Thread(() -> modifyPropertiesAfterStart( + aiResource, ImmutableMap.of(AIProperties.MAX_TOKEN, "4096"), ready, start, failure)); + + boolean bothWaiting; + aiResource.writeLock(); + try { + updateTemperature.start(); + updateMaxToken.start(); + Assertions.assertTrue(ready.await(5, TimeUnit.SECONDS)); + start.countDown(); + bothWaiting = waitUntilBlocked(updateTemperature, updateMaxToken); + } finally { + aiResource.writeUnlock(); + } + + updateTemperature.join(5000); + updateMaxToken.join(5000); + Assertions.assertTrue(bothWaiting); + Assertions.assertFalse(updateTemperature.isAlive()); + Assertions.assertFalse(updateMaxToken.isAlive()); + Assertions.assertNull(failure.get(), () -> "Concurrent ALTER failed: " + failure.get()); + Assertions.assertEquals("0.8", aiResource.getProperty(AIProperties.TEMPERATURE)); + Assertions.assertEquals("4096", aiResource.getProperty(AIProperties.MAX_TOKEN)); + } + + private static void modifyPropertiesAfterStart(AIResource aiResource, + Map<String, String> properties, CountDownLatch ready, CountDownLatch start, + AtomicReference<Throwable> failure) { + try { + ready.countDown(); + start.await(); + aiResource.modifyProperties(properties); + } catch (Throwable t) { + failure.compareAndSet(null, t); + } + } + + private static boolean waitUntilBlocked(Thread... threads) throws InterruptedException { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (System.nanoTime() < deadline) { + if (Arrays.stream(threads).allMatch(AIResourceTest::isBlocked)) { + return true; + } + Thread.sleep(10); + } + return false; + } + + private static boolean isBlocked(Thread thread) { + return thread.getState() == Thread.State.WAITING + || thread.getState() == Thread.State.BLOCKED; + } + @Test public void testDifferentProviders() throws DdlException { // 1. OpenAI diff --git a/fe/fe-core/src/test/java/org/apache/doris/common/util/PrintableMapTest.java b/fe/fe-core/src/test/java/org/apache/doris/common/util/PrintableMapTest.java index b260c90ebfb..d59869a20e7 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/common/util/PrintableMapTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/common/util/PrintableMapTest.java @@ -55,6 +55,13 @@ public class PrintableMapTest { Assertions.assertTrue(PrintableMap.SENSITIVE_KEY.contains("oss.secret_key")); } + @Test + public void testSensitiveKeysContainAiApiKeys() { + Assertions.assertTrue(PrintableMap.SENSITIVE_KEY.contains("ai.api_key")); + Assertions.assertTrue(PrintableMap.SENSITIVE_KEY.contains("ai.embed.api_key")); + Assertions.assertTrue(PrintableMap.SENSITIVE_KEY.contains("ai.embed.mm.api_key")); + } + @Test public void testBasicConstructor() { Map<String, String> testMap = new HashMap<>(); diff --git a/gensrc/thrift/PaloInternalService.thrift b/gensrc/thrift/PaloInternalService.thrift index e2ba8937684..24292edb7f9 100644 --- a/gensrc/thrift/PaloInternalService.thrift +++ b/gensrc/thrift/PaloInternalService.thrift @@ -679,9 +679,9 @@ enum TCompoundType { } struct TAIResource { - 1: required string endpoint - 2: required string provider_type - 3: required string model_name + 1: optional string endpoint + 2: optional string provider_type + 3: optional string model_name 4: optional string api_key 5: optional double temperature 6: optional i64 max_tokens @@ -689,6 +689,15 @@ struct TAIResource { 8: optional i32 retry_delay_second 9: optional string anthropic_version 10: optional i32 dimensions + 11: optional string embed_endpoint + 12: optional string embed_provider_type + 13: optional string embed_model_name + 14: optional string embed_api_key + 15: optional string effort + 16: optional string embed_mm_endpoint + 17: optional string embed_mm_provider_type + 18: optional string embed_mm_model_name + 19: optional string embed_mm_api_key } struct TCondition { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
