Skip to content

branch-4.1: [Enhance](ai_func) Support dedicate embed properties in AI RESOURCE (#67673) #68384

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yiguolei merged 1 commit into from
Sep 23, 2026
+1,008 −95
Merged

branch-4.1: [Enhance](ai_func) Support dedicate embed properties in AI RESOURCE (#67673) #68384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions be/src/core/string_buffer.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -244,6 +244,8 @@ class BufferReadable {

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) {
Expand Down Expand Up @@ -316,6 +318,7 @@ class BufferReadable {

private:
const char* _data;
const char* _end;
};

using VectorBufferReader = BufferReadable;
Expand Down
77 changes: 65 additions & 12 deletions be/src/exprs/function/ai/ai_adapter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,16 +43,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;
Expand All @@ -64,6 +65,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);
Expand All @@ -76,6 +78,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) {
Expand All @@ -89,7 +94,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 };
Expand Down Expand Up @@ -124,6 +148,8 @@ class AIAdapter {
_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
Expand Down Expand Up @@ -747,7 +773,8 @@ class OpenAIAdapter : public VoyageAIAdapter {
{"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);
Expand All @@ -759,6 +786,12 @@ class OpenAIAdapter : public VoyageAIAdapter {
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);
Expand All @@ -784,6 +817,7 @@ class OpenAIAdapter : public VoyageAIAdapter {
],
"temperature": x,
"max_tokens": x,
"reasoning_effort": "low"
}*/
doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator),
allocator);
Expand All @@ -795,6 +829,10 @@ class OpenAIAdapter : public VoyageAIAdapter {
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) {
Expand Down Expand Up @@ -1245,7 +1283,8 @@ class GeminiAdapter : public AIAdapter {
],
"generationConfig": {
"temperature": 0.7,
"maxOutputTokens": 1024
"maxOutputTokens": 1024,
"thinkingConfig": {"thinkingLevel": "high"}
}

}*/
Expand Down Expand Up @@ -1283,6 +1322,13 @@ class GeminiAdapter : public AIAdapter {
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;
Expand Down Expand Up @@ -1561,6 +1607,7 @@ class AnthropicAdapter : public VoyageAIAdapter {
/*
"model": "claude-opus-4-1-20250805",
"max_tokens": 1024,
"output_config": {"effort": "medium"},
"system": "system_prompt here",
"messages": [
{"role": "user", "content": "xxx"}
Expand All @@ -1579,6 +1626,12 @@ class AnthropicAdapter : public VoyageAIAdapter {
// 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);
}
Expand Down
25 changes: 17 additions & 8 deletions be/src/exprs/function/ai/ai_functions.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ class AIFunction : public IFunction {
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()) {
Expand Down Expand Up @@ -131,7 +131,7 @@ class AIFunction : public IFunction {
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.
Expand Down Expand Up @@ -172,7 +172,7 @@ class AIFunction : public IFunction {

// 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));

Expand Down Expand Up @@ -208,7 +208,7 @@ class AIFunction : public IFunction {

// 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,
Expand All @@ -229,7 +229,7 @@ class AIFunction : public IFunction {
// 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
Expand Down Expand Up @@ -293,7 +293,7 @@ class AIFunction : public IFunction {
// 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);
Expand Down Expand Up @@ -413,7 +413,7 @@ class AIFunction : public IFunction {
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);
Expand All @@ -424,7 +424,12 @@ class AIFunction : public IFunction {
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);

Expand All @@ -435,6 +440,10 @@ class AIFunction : public IFunction {
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,
Expand Down
34 changes: 28 additions & 6 deletions be/src/exprs/function/ai/embed.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,22 @@ class FunctionEmbed : public AIFunction<FunctionEmbed> {

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 {}",
Expand Down Expand Up @@ -90,6 +104,14 @@ class FunctionEmbed : public AIFunction<FunctionEmbed> {
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);
Expand All @@ -98,7 +120,7 @@ class FunctionEmbed : public AIFunction<FunctionEmbed> {
}

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(
Expand Down Expand Up @@ -155,7 +177,7 @@ class FunctionEmbed : public AIFunction<FunctionEmbed> {
}

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 {
Expand Down Expand Up @@ -218,7 +240,7 @@ class FunctionEmbed : public AIFunction<FunctionEmbed> {
// 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;
Expand Down Expand Up @@ -251,7 +273,7 @@ class FunctionEmbed : public AIFunction<FunctionEmbed> {
// 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()) {
Expand All @@ -275,7 +297,7 @@ class FunctionEmbed : public AIFunction<FunctionEmbed> {
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()) {
Expand Down
Loading
Loading