be/src/exprs/function/ai/ai_functions.h

Source
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#pragma once

#include <gen_cpp/FrontendService.h>
#include <gen_cpp/PaloInternalService_types.h>
#include <glog/logging.h>

#include <algorithm>
#include <cstdlib>
#include <memory>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

#include "common/config.h"
#include "common/status.h"
#include "core/column/column_array.h"
#include "core/column/column_const.h"
#include "core/column/column_nullable.h"
#include "core/cow.h"
#include "core/data_type/data_type_array.h"
#include "core/data_type/data_type_nullable.h"
#include "core/data_type/data_type_number.h"
#include "core/data_type/define_primitive_type.h"
#include "core/data_type/primitive_type.h"
#include "exec/common/util.hpp"
#include "exprs/function/ai/ai_adapter.h"
#include "exprs/function/ai/ai_request_executor.h"
#include "exprs/function/function.h"
#include "runtime/query_context.h"
#include "runtime/runtime_state.h"
#include "util/defer_op.h"
#include "util/string_util.h"
#include "util/threadpool.h"
#include "util/time.h"

namespace doris {

// Base class for AI-based functions
template <typename Derived>
class AIFunction : public IFunction {
public:
    std::string get_name() const override { return assert_cast<const Derived&>(*this).name; }

    // If the user doesn't provide the first arg, `resource_name`
    // FE will add the `resource_name` to the arguments list using the Session Variable.
    // So the value here should be the maximum number that the function can accept.
    size_t get_number_of_arguments() const override {
        return assert_cast<const Derived&>(*this).number_of_arguments;
    }

    bool is_blockable() const override { return true; }

    Status close(FunctionContext* context, FunctionContext::FunctionStateScope scope) override {
        if (scope == FunctionContext::FunctionStateScope::THREAD_LOCAL) {
            // Release this execution context's cached clients.
            context->set_function_state(scope, nullptr);
        }
        return Status::OK();
    }

    bool use_default_implementation_for_nulls() const final { return false; }

    DataTypePtr get_return_type_impl(const DataTypes& arguments) const final {
        bool has_nullable_argument = std::ranges::any_of(
                arguments, [](const auto& argument) { return argument->is_nullable(); });
        DataTypePtr return_type =
                assert_cast<const Derived&>(*this).get_nested_return_type_impl(arguments);
        return has_nullable_argument ? make_nullable(return_type) : return_type;
    }

    using PreparedFunctionImpl::execute;

    virtual Status build_prompt(const Columns& prompt_columns, size_t row_num,
                                std::string& prompt) const {
        StringRef text_ref = prompt_columns[0]->get_data_at(row_num);
        prompt = std::string(text_ref.data, text_ref.size);

        return Status::OK();
    }

    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                        uint32_t result, size_t input_rows_count) const override {
        auto* executor = get_executor(context);
        if (block.get_by_position(arguments[0]).column->only_null()) {
            if (executor->statistics()) {
                executor->statistics()->record_input(input_rows_count, input_rows_count);
            }
            block.get_by_position(result).column =
                    block.get_by_position(result).type->create_column_const(input_rows_count,
                                                                            Field());
            return Status::OK();
        }

        AIResource config;
        std::shared_ptr<AIAdapter> adapter;
        if (Status status = this->_init_from_resource(context, block, arguments, config, adapter);
            !status.ok()) {
            return status;
        }

        return assert_cast<const Derived&>(*this).execute(context, block, arguments, result,
                                                          input_rows_count, config, adapter);
    }

protected:
    // Reads the shared AI context window size from query options. String AI batch functions and
    // ai_agg both use the same byte-based session variable so batching behavior stays consistent.
    static int64_t get_ai_context_window_size(FunctionContext* context) {
        DORIS_CHECK(context != nullptr);
        QueryContext* query_ctx = context->state()->get_query_ctx();
        DORIS_CHECK(query_ctx != nullptr);

        return query_ctx->query_options().ai_context_window_size;
    }

    MutableColumnPtr create_result_column() const { return ColumnString::create(); }

    // Provider-reusable hook for AI functions(string) -> string.
    Status append_batch_results(const std::vector<std::string>& batch_results,
                                IColumn& col_result) const {
        auto& string_col = assert_cast<ColumnString&>(col_result);
        for (const auto& batch_result : batch_results) {
            string_col.insert_data(batch_result.data(), batch_result.size());
        }
        return Status::OK();
    }

    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.
        // 2. `:embedContent` -> `:batchEmbedContents`
        if (iequal(config.provider_type, "GEMINI")) {
            if (iequal(Derived::name, "embed") && config.endpoint.ends_with(":embedContent")) {
                static constexpr std::string_view legacy_suffix = ":embedContent";
                config.endpoint.replace(config.endpoint.size() - legacy_suffix.size(),
                                        legacy_suffix.size(), ":batchEmbedContents");
                return;
            }

            if (!config.endpoint.ends_with("v1") && !config.endpoint.ends_with("v1beta")) {
                return;
            }

            std::string model_name = config.model_name;
            if (!model_name.starts_with("models/")) {
                model_name = "models/" + model_name;
            }

            config.endpoint += "/";
            config.endpoint += model_name;
            config.endpoint +=
                    iequal(Derived::name, "embed") ? ":batchEmbedContents" : ":generateContent";
            return;
        }

        // The endpoint `v1/completions` does not support `system_prompt`.
        // To ensure a clear structure and stable AI results.
        // Convert from `v1/completions` to `v1/chat/completions`
        if (config.endpoint.ends_with("v1/completions")) {
            static constexpr std::string_view legacy_suffix = "v1/completions";
            config.endpoint.replace(config.endpoint.size() - legacy_suffix.size(),
                                    legacy_suffix.size(), "v1/chat/completions");
        }
    }

    static AIRequestExecutor* get_executor(FunctionContext* context) {
        // Retain the cache across batches/Blocks, including Blocks with only NULL inputs.
        auto* executor = static_cast<AIRequestExecutor*>(
                context->get_function_state(FunctionContext::FunctionStateScope::THREAD_LOCAL));
        if (executor == nullptr) {
            auto* runtime_state = context->state();
            // Disabled profiles need the connection cache, but no statistics owner.
            auto statistics = runtime_state->enable_profile()
                                      ? runtime_state->get_query_ctx()->get_ai_statistics(
                                                runtime_state->fragment_id())
                                      : nullptr;
            auto state = std::make_shared<AIRequestExecutor>(std::move(statistics));
            executor = state.get();
            context->set_function_state(FunctionContext::FunctionStateScope::THREAD_LOCAL,
                                        std::move(state));
        }
        return executor;
    }

    static void record_input(FunctionContext* context, size_t rows,
                             const ColumnUInt8::MutablePtr& null_map) {
        if (const auto& statistics = get_executor(context)->statistics()) {
            const int64_t skipped = null_map ? std::count_if(null_map->get_data().begin(),
                                                             null_map->get_data().end(),
                                                             [](UInt8 value) { return value != 0; })
                                             : 0;
            statistics->record_input(rows, skipped);
        }
    }

    Status send_request_to_llm(const std::string& request_body, std::string& response,
                               const AIResource& config, std::shared_ptr<AIAdapter>& adapter,
                               FunctionContext* context) const {
        return get_executor(context)->execute(request_body, response, config, *adapter,
                                              context->state()->get_query_ctx());
    }

    // Provider-reusable helper for string-returning functions.
    // Estimates one batch entry size using the raw prompt length plus the fixed JSON wrapper cost.
    size_t estimate_batch_entry_size(size_t idx, const std::string& prompt) const {
        static constexpr size_t json_wrapper_size = 20;
        return prompt.size() + std::to_string(idx).size() + json_wrapper_size;
    }

    // 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 AIResource& config,
                                 std::shared_ptr<AIAdapter>& adapter,
                                 FunctionContext* context) const {
#ifdef BE_TEST
        const char* test_result = std::getenv("AI_TEST_RESULT");
        if (test_result != nullptr) {
            const auto& statistics = get_executor(context)->statistics();
            const auto parse_start = statistics ? MonotonicNanos() : 0;
            bool parse_failed = true;
            Defer record_parse([&] {
                if (statistics) {
                    statistics->record_parse(MonotonicNanos() - parse_start, parse_failed);
                }
            });
            std::vector<std::string> parsed_test_response;
            RETURN_IF_ERROR(
                    adapter->parse_response(std::string(test_result), parsed_test_response));
            if (parsed_test_response.empty()) {
                return Status::InternalError("AI returned empty result");
            }
            if (parsed_test_response.size() != batch_prompts.size()) {
                return Status::RuntimeError(
                        "Failed to parse {} batch result, expected {} items but got {}", get_name(),
                        batch_prompts.size(), parsed_test_response.size());
            }
            results = std::move(parsed_test_response);
            parse_failed = false;
            return Status::OK();
        }
        if (config.provider_type == "MOCK") {
            results.clear();
            results.reserve(batch_prompts.size());
            for (const auto& prompt : batch_prompts) {
                results.emplace_back("this is a mock response. " + prompt);
            }
            return Status::OK();
        }
#endif

        std::string batch_prompt;
        RETURN_IF_ERROR(build_batch_prompt(batch_prompts, batch_prompt));

        std::vector<std::string> inputs = {batch_prompt};
        std::vector<std::string> parsed_response;

        std::string request_body;
        RETURN_IF_ERROR(adapter->build_request_payload(
                inputs, assert_cast<const Derived&>(*this).system_prompt, request_body));

        std::string response;
        RETURN_IF_ERROR(send_request_to_llm(request_body, response, config, adapter, context));
        const auto& statistics = get_executor(context)->statistics();
        const auto parse_start = statistics ? MonotonicNanos() : 0;
        bool parse_failed = true;
        Defer record_parse([&] {
            if (statistics) {
                statistics->record_parse(MonotonicNanos() - parse_start, parse_failed);
            }
        });
        RETURN_IF_ERROR(adapter->parse_response(response, parsed_response));
        if (parsed_response.empty()) {
            return Status::InternalError("AI returned empty result");
        }
        if (parsed_response.size() != batch_prompts.size()) {
            LOG(WARNING) << "AI batch result size mismatch, function=" << get_name()
                         << ", provider=" << config.provider_type << ", model=" << config.model_name
                         << ", expected_rows=" << batch_prompts.size()
                         << ", actual_rows=" << parsed_response.size()
                         << ", response_body=" << response;
            return Status::RuntimeError(
                    "Failed to parse {} batch result, expected {} items but got {}", get_name(),
                    batch_prompts.size(), parsed_response.size());
        }
        results = std::move(parsed_response);
        parse_failed = false;
        return Status::OK();
    }

    Status append_results_with_statistics(FunctionContext* context,
                                          const std::vector<std::string>& results,
                                          IColumn& column) const {
        const auto& statistics = get_executor(context)->statistics();
        const auto parse_start = statistics ? MonotonicNanos() : 0;
        Status status = assert_cast<const Derived&>(*this).append_batch_results(results, column);
        if (statistics) {
            statistics->record_parse(MonotonicNanos() - parse_start, !status.ok());
        }
        return status;
    }

    // Provider-reusable helper for string-returning functions.
    // 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 AIResource& config,
                   std::shared_ptr<AIAdapter>& adapter) const {
        Columns prompt_columns;
        prompt_columns.reserve(arguments.size() - 1);
        ColumnUInt8::MutablePtr result_null_map;
        for (size_t i = 1; i < arguments.size(); ++i) {
            const auto& argument = block.get_by_position(arguments[i]);
            if (argument.type->is_nullable()) {
                const auto& [column, is_const] = unpack_if_const(argument.column);
                const auto& nullable =
                        assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*column);
                if (!result_null_map) {
                    result_null_map = ColumnUInt8::create(input_rows_count, 0);
                }
                VectorizedUtils::update_null_map(result_null_map->get_data(),
                                                 nullable.get_null_map_data(), is_const);
            }
            prompt_columns.emplace_back(
                    argument.unnest_nullable(argument.type->is_nullable()
                                                     ? argument.get_nullable_column_info()
                                                     : NullableColumnInfo {},
                                             false)
                            .column);
        }

        record_input(context, input_rows_count, result_null_map);

        if (result_null_map &&
            !simd::contain_zero(result_null_map->get_data().data(), input_rows_count)) {
            block.get_by_position(result).column =
                    block.get_by_position(result).type->create_column_const(input_rows_count,
                                                                            Field());
            return Status::OK();
        }

        auto col_result = assert_cast<const Derived&>(*this).create_result_column();
        std::vector<std::string> batch_prompts;
        size_t current_batch_size = 2; // []
        const size_t max_batch_prompt_size =
                static_cast<size_t>(get_ai_context_window_size(context));
        const NullMap* null_map = result_null_map ? &result_null_map->get_data() : nullptr;

        for (size_t i = 0; i < input_rows_count; ++i) {
            if (null_map && (*null_map)[i]) {
                continue;
            }

            std::string prompt;
            RETURN_IF_ERROR(
                    assert_cast<const Derived&>(*this).build_prompt(prompt_columns, i, prompt));

            size_t entry_size = estimate_batch_entry_size(batch_prompts.size(), prompt);
            if (entry_size > max_batch_prompt_size) {
                if (!batch_prompts.empty()) {
                    std::vector<std::string> batch_results;
                    RETURN_IF_ERROR(this->execute_batch_request(batch_prompts, batch_results,
                                                                config, adapter, context));
                    RETURN_IF_ERROR(
                            append_results_with_statistics(context, batch_results, *col_result));
                    batch_prompts.clear();
                    current_batch_size = 2;
                }

                std::vector<std::string> single_prompts;
                single_prompts.emplace_back(std::move(prompt));
                std::vector<std::string> single_results;
                RETURN_IF_ERROR(this->execute_batch_request(single_prompts, single_results, config,
                                                            adapter, context));
                RETURN_IF_ERROR(
                        append_results_with_statistics(context, single_results, *col_result));
                continue;
            }

            size_t additional_size = entry_size + (batch_prompts.empty() ? 0 : 1);
            if (!batch_prompts.empty() &&
                current_batch_size + additional_size > max_batch_prompt_size) {
                std::vector<std::string> batch_results;
                RETURN_IF_ERROR(this->execute_batch_request(batch_prompts, batch_results, config,
                                                            adapter, context));
                RETURN_IF_ERROR(
                        append_results_with_statistics(context, batch_results, *col_result));
                batch_prompts.clear();
                current_batch_size = 2;
                additional_size = entry_size;
            }

            batch_prompts.emplace_back(std::move(prompt));
            current_batch_size += additional_size;
        }

        if (!batch_prompts.empty()) {
            std::vector<std::string> batch_results;
            RETURN_IF_ERROR(this->execute_batch_request(batch_prompts, batch_results, config,
                                                        adapter, context));
            RETURN_IF_ERROR(append_results_with_statistics(context, batch_results, *col_result));
        }

        if (!result_null_map) {
            block.replace_by_position(result, std::move(col_result));
            return Status::OK();
        }

        if (!simd::contain_one(result_null_map->get_data().data(), input_rows_count)) {
            block.replace_by_position(result, ColumnNullable::create(std::move(col_result),
                                                                     std::move(result_null_map)));
            return Status::OK();
        }

        auto nested_result = col_result->clone_empty();
        size_t result_row = 0;
        for (UInt8 is_null : result_null_map->get_data()) {
            if (is_null) {
                nested_result->insert_default();
            } else {
                nested_result->insert_from(*col_result, result_row++);
            }
        }

        block.replace_by_position(result, ColumnNullable::create(std::move(nested_result),
                                                                 std::move(result_null_map)));
        return Status::OK();
    }

private:
    // The ai resource must be literal
    Status _init_from_resource(FunctionContext* context, const Block& block,
                               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);
        std::string resource_name = std::string(resource_name_ref.data, resource_name_ref.size);

        const std::shared_ptr<std::map<std::string, TAIResource>>& ai_resources =
                context->state()->get_query_ctx()->get_ai_resources();
        DORIS_CHECK(ai_resources);
        auto it = ai_resources->find(resource_name);
        DORIS_CHECK(it != ai_resources->end());
        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);

        adapter = AIAdapterFactory::create_adapter(config.provider_type);
        DORIS_CHECK(adapter);

        adapter->init(config);
        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,
                              std::string& prompt) const {
        rapidjson::StringBuffer buffer;
        rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);

        writer.StartArray();
        for (size_t i = 0; i < batch_prompts.size(); ++i) {
            writer.StartObject();
            writer.Key("idx");
            writer.Uint64(i);
            writer.Key("input");
            writer.String(batch_prompts[i].data(),
                          static_cast<rapidjson::SizeType>(batch_prompts[i].size()));
            writer.EndObject();
        }
        writer.EndArray();

        prompt = buffer.GetString();
        return Status::OK();
    }
};

} // namespace doris

Coverage Report

Created: 2026-10-20 19:14