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

// 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.

#include "core/column/column_array_view.h"
#include "exprs/function/ai/ai_classify.h"
#include "exprs/function/ai/ai_extract.h"
#include "exprs/function/ai/ai_filter.h"
#include "exprs/function/ai/ai_fix_grammar.h"
#include "exprs/function/ai/ai_generate.h"
#include "exprs/function/ai/ai_mask.h"
#include "exprs/function/ai/ai_sentiment.h"
#include "exprs/function/ai/ai_similarity.h"
#include "exprs/function/ai/ai_summarize.h"
#include "exprs/function/ai/ai_translate.h"
#include "exprs/function/ai/embed.h"
#include "exprs/function/simple_function_factory.h"

namespace doris {
static Status format_labels(const ColumnPtr& labels_column, size_t row_num,
                            std::string_view function_name, std::string& labels_str) {
    auto readable_column = check_column_const_set_readability(*labels_column, row_num);
    if (!is_column<ColumnArray>(*readable_column.first)) {
        return Status::InternalError(
                "labels argument for {} must be Array(String) or Array(Varchar)", function_name);
    }

    auto labels_view = ColumnArrayView<TYPE_STRING>::create(labels_column);
    auto labels = labels_view[row_num];
    labels_str = "[";
    bool is_first_label = true;
    for (size_t i = 0; i < labels.size(); ++i) {
        if (labels.is_null_at(i)) {
            continue;
        }
        if (!is_first_label) {
            labels_str += ", ";
        }
        StringRef label = labels.value_at(i);
        labels_str += "\"";
        labels_str.append(label.data, label.size);
        labels_str += "\"";
        is_first_label = false;
    }
    labels_str += "]";
    return Status::OK();
}

Status FunctionAIClassify::build_prompt(const Columns& prompt_columns, size_t row_num,
                                        std::string& prompt) const {
    // Get the text column
    StringRef text = prompt_columns[0]->get_data_at(row_num);
    std::string text_str = std::string(text.data, text.size);

    std::string labels_str;
    RETURN_IF_ERROR(format_labels(prompt_columns[1], row_num, name, labels_str));

    prompt = "Labels: " + labels_str + "\nText: " + text_str;

    return Status::OK();
}

Status FunctionAIExtract::build_prompt(const Columns& prompt_columns, size_t row_num,
                                       std::string& prompt) const {
    // Get the text column
    StringRef text = prompt_columns[0]->get_data_at(row_num);
    std::string text_str = std::string(text.data, text.size);

    std::string labels_str;
    RETURN_IF_ERROR(format_labels(prompt_columns[1], row_num, name, labels_str));

    prompt = "Labels: " + labels_str + "\nText: " + text_str;

    return Status::OK();
}

Status FunctionAIGenerate::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 FunctionAIMask::build_prompt(const Columns& prompt_columns, size_t row_num,
                                    std::string& prompt) const {
    // Get the text column
    StringRef text = prompt_columns[0]->get_data_at(row_num);
    std::string text_str = std::string(text.data, text.size);

    std::string labels_str;
    RETURN_IF_ERROR(format_labels(prompt_columns[1], row_num, name, labels_str));

    prompt = "Labels: " + labels_str + "\nText: " + text_str;

    return Status::OK();
}

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

    // text2
    StringRef text_2 = prompt_columns[1]->get_data_at(row_num);
    std::string text_str_2 = std::string(text_2.data, text_2.size);

    prompt = "Text 1: " + text_str_1 + "\nText 2: " + text_str_2;

    return Status::OK();
}

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

    // target language
    StringRef lang = prompt_columns[1]->get_data_at(row_num);
    std::string target_lang = std::string(lang.data, lang.size);

    prompt = "Translate the following text to " + target_lang + ".\nText: " + text_str;

    return Status::OK();
}

void register_function_ai(SimpleFunctionFactory& factory) {
    factory.register_function<FunctionEmbed>();
    factory.register_function<FunctionAIClassify>();
    factory.register_function<FunctionAIExtract>();
    factory.register_function<FunctionAIFilter>();
    factory.register_function<FunctionAIFixGrammar>();
    factory.register_function<FunctionAIGenerate>();
    factory.register_function<FunctionAIMask>();
    factory.register_function<FunctionAISentiment>();
    factory.register_function<FunctionAISimilarity>();
    factory.register_function<FunctionAISummarize>();
    factory.register_function<FunctionAITranslate>();
}

} // namespace doris