This is an automated email from the ASF dual-hosted git repository.
yiguolei pushed a commit to branch branch-4.1
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/branch-4.1 by this push:
new 2eb1ec36b09 branch-4.1: [Enhance](ai_func) Support dedicate embed
properties in AI RESOURCE (#67673) (#68384)
2eb1ec36b09 is described below
commit 2eb1ec36b098cd12b3a60b26cb8994d829feabb0
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 9e6d34c7c39..6753dbe2f4c 100644
--- a/gensrc/thrift/PaloInternalService.thrift
+++ b/gensrc/thrift/PaloInternalService.thrift
@@ -673,9 +673,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
@@ -683,6 +683,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]
