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

#include <rapidjson/rapidjson.h>

#include <cctype>

#include <cstdint>

#include <memory>

#include <optional>

#include <string>

#include <string_view>

#include <unordered_map>

#include <vector>

#include "common/status.h"

#include "core/string_buffer.hpp"

#include "rapidjson/document.h"

#include "rapidjson/stringbuffer.h"

#include "rapidjson/writer.h"

#include "service/http/http_client.h"

#include "service/http/http_headers.h"

#include "util/security.h"

namespace doris {

// Missing counts remain unknown; totals are never inferred from other fields.

struct AIUsage {

    std::optional<int64_t> prompt_tokens;

    std::optional<int64_t> completion_tokens;

    std::optional<int64_t> total_tokens;

};

struct AIResource {

    AIResource() = default;

    AIResource(const TAIResource& tai)

            : 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;

    std::string model_name;

    std::string api_key;

    double temperature;

    int64_t max_tokens;

    int32_t max_retries;

    int32_t retry_delay_second;

    std::string anthropic_version;

    int32_t dimensions;

    std::string effort;

    // Optional provider output constraints for scalar batches; not serialized in AI_AGG state.

    bool enable_structured_output = true;

    void serialize(BufferWritable& buf) const {

        buf.write_binary(endpoint);

        buf.write_binary(provider_type);

        buf.write_binary(model_name);

        buf.write_binary(api_key);

        buf.write_binary(temperature);

        buf.write_binary(max_tokens);

        buf.write_binary(max_retries);

        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) {

        buf.read_binary(endpoint);

        buf.read_binary(provider_type);

        buf.read_binary(model_name);

        buf.read_binary(api_key);

        buf.read_binary(temperature);

        buf.read_binary(max_tokens);

        buf.read_binary(max_retries);

        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),

              enable_structured_output(!tai.__isset.enable_structured_output ||

                                       tai.enable_structured_output) {}

};

enum class MultimodalType { IMAGE, VIDEO, AUDIO };

inline const char* multimodal_type_to_string(MultimodalType type) {

    switch (type) {

    case MultimodalType::IMAGE:

        return "image";

    case MultimodalType::VIDEO:

        return "video";

    case MultimodalType::AUDIO:

        return "audio";

    }

    return "unknown";

}

class AIAdapter {

public:

    explicit AIAdapter(bool structured_batch = false) : _structured_batch(structured_batch) {}

    virtual ~AIAdapter() = default;

    // Usage is independent of result parsing and may also be present in error responses.

    // OpenAI-compatible providers use either chat or Responses API field names.

    virtual AIUsage parse_usage(const std::string& response_body) const {

        return parse_usage_fields(response_body, "usage", "prompt_tokens", "completion_tokens",

                                  "total_tokens", "input_tokens", "output_tokens");

    }

    // Set authentication headers for the HTTP client

    virtual Status set_authentication(HttpClient* client) const = 0;

    virtual void init(const TAIResource& config) { _config = config; }

    virtual void init(const AIResource& config) {

        _config.endpoint = config.endpoint;

        _config.provider_type = config.provider_type;

        _config.model_name = config.model_name;

        _config.api_key = config.api_key;

        _config.temperature = config.temperature;

        _config.max_tokens = config.max_tokens;

        _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. The format fields in provider

    // examples apply only to structured batches; AI_AGG uses the same transport without them.

    virtual Status build_request_payload(const std::vector<std::string>& inputs,

                                         const char* const system_prompt, std::string& request_body,

                                         bool structured_batch = false) const {

        return Status::NotSupported("{} don't support text generation", _config.provider_type);

    }

    bool supports_structured_batch() const { return _structured_batch; }

    Status parse_batch_response(const std::string& response_body,

                                std::vector<std::string>& results) const {

        if (!supports_structured_batch()) {

            return parse_response(response_body, results);

        }

        std::vector<std::string> texts;

        RETURN_IF_ERROR(parse_response(response_body, texts, false));

        // Multiple choices are alternatives, not additional rows of the same batch.

        if (texts.size() != 1) {

            return Status::InternalError("Expected one structured batch response, got {}",

                                         texts.size());

        }

        rapidjson::Document doc;

        doc.Parse(texts[0].data(), texts[0].size());

        // A bare array, fenced JSON, or a missing/wrong results field violates our schema,

        // even for a batch of one row; never return that envelope as a successful SQL value.

        if (doc.HasParseError() || !doc.IsObject() || doc.MemberCount() != 1 ||

            !doc.HasMember("results") || !doc["results"].IsArray()) {

            return Status::InternalError("Invalid structured batch result: expected results array");

        }

        const auto& rows = doc["results"];

        for (rapidjson::SizeType i = 0; i < rows.Size(); ++i) {

            if (!rows[i].IsString()) {

                return Status::InternalError(

                        "Invalid structured batch result: element {} is not a string", i);

            }

        }

        for (const auto& row : rows.GetArray()) {

            results.emplace_back(row.GetString(), row.GetStringLength());

        }

        return Status::OK();

    }

    // Serialize the request while appending the provider's immutable result-format fragment.

    // RawValue avoids constructing/parsing a schema DOM; nested fields retain siblings like effort.

    static void append_result_format(const rapidjson::Document& doc, std::string& request_body,

                                     const char* key, std::string_view fragment,

                                     const char* parent = nullptr) {

        rapidjson::StringBuffer buffer;

        rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);

        auto write_format = [&]() {

            writer.Key(key);

            writer.RawValue(fragment.data(), fragment.size(), rapidjson::kObjectType);

        };

        writer.StartObject();

        for (const auto& member : doc.GetObject()) {

            writer.Key(member.name.GetString(), member.name.GetStringLength());

            if (parent && member.name == parent) {

                writer.StartObject();

                for (const auto& nested : member.value.GetObject()) {

                    writer.Key(nested.name.GetString(), nested.name.GetStringLength());

                    nested.value.Accept(writer);

                }

                write_format();

                writer.EndObject();

            } else {

                member.value.Accept(writer);

            }

        }

        if (!parent) {

            write_format();

        } else if (parent && !doc.HasMember(parent)) {

            writer.Key(parent);

            writer.StartObject();

            write_format();

            writer.EndObject();

        }

        writer.EndObject();

        request_body.assign(buffer.GetString(), buffer.GetSize());

    }

    // Parse response from AI service and extract generated text results

    virtual Status parse_response(const std::string& response_body,

                                  std::vector<std::string>& results,

                                  bool /* expand_batch */ = true) const {

        return Status::NotSupported("{} don't support text generation", _config.provider_type);

    }

    virtual Status build_embedding_request(const std::vector<std::string>& inputs,

                                           std::string& request_body) const {

        return embed_not_supported_status();

    }

    virtual Status build_multimodal_embedding_request(

            const std::vector<MultimodalType>& /*media_types*/,

            const std::vector<std::string>& /*media_urls*/,

            const std::vector<std::string>& /*media_content_types*/,

            std::string& /*request_body*/) const {

        return Status::NotSupported("{} does not support multimodal Embed feature.",

                                    _config.provider_type);

    }

    virtual Status parse_embedding_response(const std::string& response_body,

                                            std::vector<std::vector<float>>& results) const {

        return embed_not_supported_status();

    }

protected:

    static std::optional<int64_t> usage_count(const rapidjson::Value& object, const char* field,

                                              const char* fallback = nullptr) {

        if (field == nullptr) {

            return std::nullopt;

        }

        auto member = object.FindMember(field);

        if (member == object.MemberEnd() && fallback != nullptr) {

            member = object.FindMember(fallback);

        }

        if (member == object.MemberEnd() || !member->value.IsInt64() ||

            member->value.GetInt64() < 0) {

            return std::nullopt;

        }

        return member->value.GetInt64();

    }

    static AIUsage parse_usage_fields(const std::string& response_body, const char* envelope,

                                      const char* prompt, const char* completion, const char* total,

                                      const char* prompt_fallback = nullptr,

                                      const char* completion_fallback = nullptr) {

        rapidjson::Document doc;

        doc.Parse(response_body.data(), response_body.size());

        if (doc.HasParseError() || !doc.IsObject()) {

            return {};

        }

        const auto member = doc.FindMember(envelope);

        if (member == doc.MemberEnd() || !member->value.IsObject()) {

            return {};

        }

        const auto& usage = member->value;

        return {.prompt_tokens = usage_count(usage, prompt, prompt_fallback),

                .completion_tokens = usage_count(usage, completion, completion_fallback),

                .total_tokens = usage_count(usage, total)};

    }

    TAIResource _config;

    Status embed_not_supported_status() const {

        return Status::NotSupported(

                "{} does not support the Embed feature. Currently supported providers are "

                "OpenAI, Gemini, Voyage, Jina, Qwen, and Minimax.",

                _config.provider_type);

    }

    // Appends one provider-parsed text result to `results`.

    // The adapter has already parsed the provider's outer response envelope before calling here.

    // Example:

    // provider response -> choices[0].message.content = "[\"1\",\"0\",\"1\"]"

    // this helper       -> appends "1", "0", "1" into `results`

    // Set expand_batch to false when AI_AGG needs the complete generated text as one result.

    static Status append_parsed_text_result(std::string_view text,

                                            std::vector<std::string>& results,

                                            bool expand_batch = true) {

        if (!expand_batch) {

            results.emplace_back(text.data(), text.size());

            return Status::OK();

        }

        size_t begin = 0;

        size_t end = text.size();

        while (begin < end && std::isspace(static_cast<unsigned char>(text[begin]))) {

            ++begin;

        }

        while (begin < end && std::isspace(static_cast<unsigned char>(text[end - 1]))) {

            --end;

        }

        if (begin < end && text[begin] == '[' && text[end - 1] == ']') {

            rapidjson::Document doc;

            doc.Parse(text.data() + begin, end - begin);

            if (!doc.HasParseError() && doc.IsArray()) {

                for (rapidjson::SizeType i = 0; i < doc.Size(); ++i) {

                    if (!doc[i].IsString()) {

                        return Status::InternalError(

                                "Invalid batch result format, array element {} is not a string", i);

                    }

                    results.emplace_back(doc[i].GetString(), doc[i].GetStringLength());

                }

                return Status::OK();

            }

        }

        results.emplace_back(text.data(), text.size());

        return Status::OK();

    }

    Status append_parsed_embedding_result(const rapidjson::Value& embedding,

                                          std::vector<std::vector<float>>& results,

                                          const std::string& response_body) const {

        if (!embedding.IsArray()) {

            return Status::InternalError("Invalid {} response format: {}", _config.provider_type,

                                         response_body);

        }

        std::vector<float> parsed_embedding;

        parsed_embedding.reserve(embedding.Size());

        for (const auto& value : embedding.GetArray()) {

            if (!value.IsNumber()) {

                return Status::InternalError("Invalid {} response format: {}",

                                             _config.provider_type, response_body);

            }

            parsed_embedding.emplace_back(value.GetFloat());

        }

        results.emplace_back(std::move(parsed_embedding));

        return Status::OK();

    }

    // return true if the model support dimension parameter

    virtual bool supports_dimension_param(const std::string& model_name) const { return false; }

    // Different providers may have different dimension parameter names.

    virtual std::string get_dimension_param_name() const { return "dimensions"; }

    virtual void add_dimension_params(rapidjson::Value& doc,

                                      rapidjson::Document::AllocatorType& allocator) const {

        if (_config.dimensions != -1 && supports_dimension_param(_config.model_name)) {

            std::string param_name = get_dimension_param_name();

            rapidjson::Value name(param_name.c_str(), allocator);

            doc.AddMember(name, _config.dimensions, allocator);

        }

    }

    // Validates common multimodal embedding request invariants shared by providers.

    Status validate_multimodal_embedding_inputs(

            std::string_view provider_name, const std::vector<MultimodalType>& media_types,

            const std::vector<std::string>& media_urls,

            std::initializer_list<MultimodalType> supported_types) const {

        if (media_urls.empty()) {

            return Status::InvalidArgument("{} multimodal embed inputs can not be empty",

                                           provider_name);

        }

        if (media_types.size() != media_urls.size()) {

            return Status::InvalidArgument(

                    "{} multimodal embed input size mismatch, media_types={}, media_urls={}",

                    provider_name, media_types.size(), media_urls.size());

        }

        for (MultimodalType media_type : media_types) {

            bool supported = false;

            for (MultimodalType supported_type : supported_types) {

                if (media_type == supported_type) {

                    supported = true;

                    break;

                }

            }

            if (!supported) [[unlikely]] {

                return Status::InvalidArgument(

                        "{} only supports {} multimodal embed, got {}", provider_name,

                        supported_multimodal_types_to_string(supported_types),

                        multimodal_type_to_string(media_type));

            }

        }

        return Status::OK();

    }

    static std::string supported_multimodal_types_to_string(

            std::initializer_list<MultimodalType> supported_types) {

        std::string result;

        for (MultimodalType type : supported_types) {

            if (!result.empty()) {

                result += "/";

            }

            result += multimodal_type_to_string(type);

        }

        return result;

    }

private:

    // The batch protocol is fixed by the concrete adapter, including inherited transports.

    const bool _structured_batch;

};

// Most LLM-providers' Embedding formats are based on VoyageAI.

// The following adapters inherit from VoyageAIAdapter to directly reuse its embedding logic.

class VoyageAIAdapter : public AIAdapter {

public:

    using AIAdapter::AIAdapter;

    AIUsage parse_usage(const std::string& response_body) const override {

        return parse_usage_fields(response_body, "usage", nullptr, nullptr, "total_tokens");

    }

    Status set_authentication(HttpClient* client) const override {

        client->set_header(HttpHeaders::AUTHORIZATION, "Bearer " + _config.api_key);

        client->set_content_type("application/json");

        return Status::OK();

    }

    Status build_embedding_request(const std::vector<std::string>& inputs,

                                   std::string& request_body) const override {

        rapidjson::Document doc;

        doc.SetObject();

        auto& allocator = doc.GetAllocator();

        /*{

            "model": "xxx",

            "input": [

              "xxx",

              "xxx",

              ...

            ],

            "output_dimensions": 512

        }*/

        doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator), allocator);

        add_dimension_params(doc, allocator);

        rapidjson::Value input(rapidjson::kArrayType);

        for (const auto& msg : inputs) {

            input.PushBack(rapidjson::Value(msg.c_str(), allocator), allocator);

        }

        doc.AddMember("input", input, allocator);

        rapidjson::StringBuffer buffer;

        rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);

        doc.Accept(writer);

        request_body = buffer.GetString();

        return Status::OK();

    }

    Status build_multimodal_embedding_request(

            const std::vector<MultimodalType>& media_types,

            const std::vector<std::string>& media_urls,

            const std::vector<std::string>& /*media_content_types*/,

            std::string& request_body) const override {

        RETURN_IF_ERROR(validate_multimodal_embedding_inputs(

                "VoyageAI", media_types, media_urls,

                {MultimodalType::IMAGE, MultimodalType::VIDEO}));

        if (_config.dimensions != -1) {

            LOG(WARNING) << "VoyageAI multimodal embedding currently ignores dimensions parameter, "

                         << "model=" << _config.model_name << ", dimensions=" << _config.dimensions;

        }

        rapidjson::Document doc;

        doc.SetObject();

        auto& allocator = doc.GetAllocator();

        /*{

            "inputs": [

              {

                "content": [

                  {"type": "image_url", "image_url": "<url>"}

                ]

              },

              {

                "content": [

                  {"type": "video_url", "video_url": "<url>"}

                ]

              }

            ],

            "model": "voyage-multimodal-3.5"

        }*/

        doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator), allocator);

        rapidjson::Value request_inputs(rapidjson::kArrayType);

        for (size_t i = 0; i < media_urls.size(); ++i) {

            rapidjson::Value input(rapidjson::kObjectType);

            rapidjson::Value content(rapidjson::kArrayType);

            rapidjson::Value media_item(rapidjson::kObjectType);

            if (media_types[i] == MultimodalType::IMAGE) {

                media_item.AddMember("type", "image_url", allocator);

                media_item.AddMember("image_url",

                                     rapidjson::Value(media_urls[i].c_str(), allocator), allocator);

            } else {

                media_item.AddMember("type", "video_url", allocator);

                media_item.AddMember("video_url",

                                     rapidjson::Value(media_urls[i].c_str(), allocator), allocator);

            }

            content.PushBack(media_item, allocator);

            input.AddMember("content", content, allocator);

            request_inputs.PushBack(input, allocator);

        }

        doc.AddMember("inputs", request_inputs, allocator);

        rapidjson::StringBuffer buffer;

        rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);

        doc.Accept(writer);

        request_body = buffer.GetString();

        return Status::OK();

    }

    Status parse_embedding_response(const std::string& response_body,

                                    std::vector<std::vector<float>>& results) const override {

        rapidjson::Document doc;

        doc.Parse(response_body.c_str());

        if (doc.HasParseError() || !doc.IsObject()) {

            return Status::InternalError("Failed to parse {} response: {}", _config.provider_type,

                                         response_body);

        }

        if (!doc.HasMember("data") || !doc["data"].IsArray()) {

            return Status::InternalError("Invalid {} response format: {}", _config.provider_type,

                                         response_body);

        }

        /*{

            "data":[

              {

                "object": "embedding",

                "embedding": [...], <- only need this

                "index": 0

              },

              {

                "object": "embedding",

                "embedding": [...],

                "index": 1

              }, ...

            ],

            "model"....

        }*/

        const auto& data = doc["data"];

        results.reserve(data.Size());

        for (rapidjson::SizeType i = 0; i < data.Size(); i++) {

            if (!data[i].IsObject() || !data[i].HasMember("embedding")) {

                return Status::InternalError("Invalid {} response format: {}",

                                             _config.provider_type, response_body);

            }

            RETURN_IF_ERROR(

                    append_parsed_embedding_result(data[i]["embedding"], results, response_body));

        }

        return Status::OK();

    }

protected:

    bool supports_dimension_param(const std::string& model_name) const override {

        static const std::unordered_set<std::string> no_dimension_models = {

                "voyage-law-2", "voyage-2", "voyage-code-2", "voyage-finance-2",

                "voyage-multimodal-3"};

        return !no_dimension_models.contains(model_name);

    }

    std::string get_dimension_param_name() const override { return "output_dimension"; }

};

// Local AI adapter for locally hosted models (Ollama, LLaMA, etc.)

class LocalAdapter : public AIAdapter {

public:

    AIUsage parse_usage(const std::string& response_body) const override {

        rapidjson::Document doc;

        doc.Parse(response_body.data(), response_body.size());

        if (doc.HasParseError() || !doc.IsObject()) {

            return {};

        }

        if (doc.HasMember("prompt_eval_count") || doc.HasMember("eval_count")) {

            return {.prompt_tokens = usage_count(doc, "prompt_eval_count"),

                    .completion_tokens = usage_count(doc, "eval_count"),

                    .total_tokens = std::nullopt};

        }

        const auto member = doc.FindMember("usage");

        if (member == doc.MemberEnd() || !member->value.IsObject()) {

            return {};

        }

        const auto& usage = member->value;

        return {.prompt_tokens = usage_count(usage, "prompt_tokens", "input_tokens"),

                .completion_tokens = usage_count(usage, "completion_tokens", "output_tokens"),

                .total_tokens = usage_count(usage, "total_tokens")};

    }

    // Local deployments typically don't need authentication

    Status set_authentication(HttpClient* client) const override {

        client->set_content_type("application/json");

        return Status::OK();

    }

    Status build_request_payload(const std::vector<std::string>& inputs,

                                 const char* const system_prompt, std::string& request_body,

                                 bool structured_batch = false) const override {

        rapidjson::Document doc;

        doc.SetObject();

        auto& allocator = doc.GetAllocator();

        std::string end_point = _config.endpoint;

        if (end_point.ends_with("chat") || end_point.ends_with("generate")) {

            RETURN_IF_ERROR(

                    build_ollama_request(doc, allocator, inputs, system_prompt, request_body));

        } else {

            RETURN_IF_ERROR(

                    build_default_request(doc, allocator, inputs, system_prompt, request_body));

        }

        rapidjson::StringBuffer buffer;

        rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);

        doc.Accept(writer);

        request_body = buffer.GetString();

        return Status::OK();

    }

    Status parse_response(const std::string& response_body, std::vector<std::string>& results,

                          bool expand_batch = true) const override {

        rapidjson::Document doc;

        doc.Parse(response_body.c_str());

        if (doc.HasParseError() || !doc.IsObject()) {

            return Status::InternalError("Failed to parse {} response: {}", _config.provider_type,

                                         response_body);

        }

        // Handle various response formats from local LLMs

        // Format 1: OpenAI-compatible format with choices/message/content

        if (doc.HasMember("choices") && doc["choices"].IsArray()) {

            const auto& choices = doc["choices"];

            results.reserve(choices.Size());

            for (rapidjson::SizeType i = 0; i < choices.Size(); i++) {

                if (!choices[i].IsObject()) {

                    return Status::InternalError("Invalid {} response format: {}",

                                                 _config.provider_type, response_body);

                }

                if (choices[i].HasMember("message") && !choices[i]["message"].IsObject()) {

                    return Status::InternalError("Invalid {} response format: {}",

                                                 _config.provider_type, response_body);

                }

                if (choices[i].HasMember("message") && choices[i]["message"].HasMember("content") &&

                    choices[i]["message"]["content"].IsString()) {

                    RETURN_IF_ERROR(append_parsed_text_result(

                            std::string_view(choices[i]["message"]["content"].GetString(),

                                             choices[i]["message"]["content"].GetStringLength()),

                            results, expand_batch));

                } else if (choices[i].HasMember("text") && choices[i]["text"].IsString()) {

                    // Some local LLMs use a simpler format

                    RETURN_IF_ERROR(append_parsed_text_result(

                            std::string_view(choices[i]["text"].GetString(),

                                             choices[i]["text"].GetStringLength()),

                            results, expand_batch));

                }

            }

        } else if (doc.HasMember("text") && doc["text"].IsString()) {

            // Format 2: Simple response with just "text" or "content" field

            RETURN_IF_ERROR(append_parsed_text_result(

                    std::string_view(doc["text"].GetString(), doc["text"].GetStringLength()),

                    results, expand_batch));

        } else if (doc.HasMember("content") && doc["content"].IsString()) {

            RETURN_IF_ERROR(append_parsed_text_result(

                    std::string_view(doc["content"].GetString(), doc["content"].GetStringLength()),

                    results, expand_batch));

        } else if (doc.HasMember("response") && doc["response"].IsString()) {

            // Format 3: Response field (Ollama `generate` format)

            RETURN_IF_ERROR(

                    append_parsed_text_result(std::string_view(doc["response"].GetString(),

                                                               doc["response"].GetStringLength()),

                                              results, expand_batch));

        } else if (doc.HasMember("message") && doc["message"].IsObject() &&

                   doc["message"].HasMember("content") && doc["message"]["content"].IsString()) {

            // Format 4: message/content field (Ollama `chat` format)

            RETURN_IF_ERROR(append_parsed_text_result(

                    std::string_view(doc["message"]["content"].GetString(),

                                     doc["message"]["content"].GetStringLength()),

                    results, expand_batch));

        } else {

            return Status::NotSupported("Unsupported response format from local AI.");

        }

        return Status::OK();

    }

    Status build_embedding_request(const std::vector<std::string>& inputs,

                                   std::string& request_body) const override {

        rapidjson::Document doc;

        doc.SetObject();

        auto& allocator = doc.GetAllocator();

        if (!_config.model_name.empty()) {

            doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator),

                          allocator);

        }

        add_dimension_params(doc, allocator);

        rapidjson::Value input(rapidjson::kArrayType);

        for (const auto& msg : inputs) {

            input.PushBack(rapidjson::Value(msg.c_str(), allocator), allocator);

        }

        doc.AddMember("input", input, allocator);

        rapidjson::StringBuffer buffer;

        rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);

        doc.Accept(writer);

        request_body = buffer.GetString();

        return Status::OK();

    }

    Status build_multimodal_embedding_request(

            const std::vector<MultimodalType>& /*media_types*/,

            const std::vector<std::string>& /*media_urls*/,

            const std::vector<std::string>& /*media_content_types*/,

            std::string& /*request_body*/) const override {

        return Status::NotSupported("{} does not support multimodal Embed feature.",

                                    _config.provider_type);

    }

    Status parse_embedding_response(const std::string& response_body,

                                    std::vector<std::vector<float>>& results) const override {

        rapidjson::Document doc;

        doc.Parse(response_body.c_str());

        if (doc.HasParseError() || !doc.IsObject()) {

            return Status::InternalError("Failed to parse {} response: {}", _config.provider_type,

                                         response_body);

        }

        // parse different response format

        if (doc.HasMember("data") && doc["data"].IsArray()) {

            // "data":["object":"embedding", "embedding":[0.1, 0.2...], "index":0]

            const auto& data = doc["data"];

            results.reserve(data.Size());

            for (rapidjson::SizeType i = 0; i < data.Size(); i++) {

                if (!data[i].IsObject() || !data[i].HasMember("embedding")) {

                    return Status::InternalError("Invalid {} response format",

                                                 _config.provider_type);

                }

                RETURN_IF_ERROR(append_parsed_embedding_result(data[i]["embedding"], results,

                                                               response_body));

            }

        } else if (doc.HasMember("embeddings") && doc["embeddings"].IsArray()) {

            // "embeddings":[[0.1, 0.2, ...]]

            const auto& embeddings = doc["embeddings"];

            results.reserve(embeddings.Size());

            for (rapidjson::SizeType i = 0; i < embeddings.Size(); i++) {

                RETURN_IF_ERROR(

                        append_parsed_embedding_result(embeddings[i], results, response_body));

            }

        } else if (doc.HasMember("embedding") && doc["embedding"].IsArray()) {

            // "embedding":[0.1, 0.2, ...]

            results.reserve(1);

            RETURN_IF_ERROR(

                    append_parsed_embedding_result(doc["embedding"], results, response_body));

        } else {

            return Status::InternalError("Invalid {} response format: {}", _config.provider_type,

                                         response_body);

        }

        return Status::OK();

    }

private:

    Status build_ollama_request(rapidjson::Document& doc,

                                rapidjson::Document::AllocatorType& allocator,

                                const std::vector<std::string>& inputs,

                                const char* const system_prompt, std::string& request_body) const {

        /*

        for endpoints end_with `/chat` like 'http://localhost:11434/api/chat':

        {

            "model": <model_name>,

            "stream": false,

            "think": false,

            "options": {

                "temperature": <temperature>,

                "max_token": <max_token>

            },

            "messages": [

                {"role": "system", "content": <system_prompt>},

                {"role": "user", "content": <user_prompt>}

            ]

        }

        for endpoints end_with `/generate` like 'http://localhost:11434/api/generate':

        {

            "model": <model_name>,

            "stream": false,

            "think": false

            "options": {

                "temperature": <temperature>,

                "max_token": <max_token>

            },

            "system": <system_prompt>,

            "prompt": <user_prompt>

        }

        */

        // For Ollama, only the prompt section ("system" + "prompt" or "role" + "content") is affected by the endpoint;

        // The rest remains identical.

        doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator), allocator);

        doc.AddMember("stream", false, allocator);

        doc.AddMember("think", false, allocator);

        // option section

        rapidjson::Value options(rapidjson::kObjectType);

        if (_config.temperature != -1) {

            options.AddMember("temperature", _config.temperature, allocator);

        }

        if (_config.max_tokens != -1) {

            options.AddMember("max_token", _config.max_tokens, allocator);

        }

        doc.AddMember("options", options, allocator);

        // prompt section

        if (_config.endpoint.ends_with("chat")) {

            rapidjson::Value messages(rapidjson::kArrayType);

            if (system_prompt && *system_prompt) {

                rapidjson::Value sys_msg(rapidjson::kObjectType);

                sys_msg.AddMember("role", "system", allocator);

                sys_msg.AddMember("content", rapidjson::Value(system_prompt, allocator), allocator);

                messages.PushBack(sys_msg, allocator);

            }

            for (const auto& input : inputs) {

                rapidjson::Value message(rapidjson::kObjectType);

                message.AddMember("role", "user", allocator);

                message.AddMember("content", rapidjson::Value(input.c_str(), allocator), allocator);

                messages.PushBack(message, allocator);

            }

            doc.AddMember("messages", messages, allocator);

        } else {

            if (system_prompt && *system_prompt) {

                doc.AddMember("system", rapidjson::Value(system_prompt, allocator), allocator);

            }

            doc.AddMember("prompt", rapidjson::Value(inputs[0].c_str(), allocator), allocator);

        }

        return Status::OK();

    }

    Status build_default_request(rapidjson::Document& doc,

                                 rapidjson::Document::AllocatorType& allocator,

                                 const std::vector<std::string>& inputs,

                                 const char* const system_prompt, std::string& request_body) const {

        /*

        Default format(OpenAI-compatible):

        {

            "model": <model_name>,

            "temperature": <temperature>,

            "max_tokens": <max_tokens>,

            "messages": [

                {"role": "system", "content": <system_prompt>},

                {"role": "user", "content": <user_prompt>}

            ]

        }

        */

        doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator), allocator);

        // If 'temperature' and 'max_tokens' are set, add them to the request body.

        if (_config.temperature != -1) {

            doc.AddMember("temperature", _config.temperature, allocator);

        }

        if (_config.max_tokens != -1) {

            doc.AddMember("max_tokens", _config.max_tokens, allocator);

        }

        rapidjson::Value messages(rapidjson::kArrayType);

        if (system_prompt && *system_prompt) {

            rapidjson::Value sys_msg(rapidjson::kObjectType);

            sys_msg.AddMember("role", "system", allocator);

            sys_msg.AddMember("content", rapidjson::Value(system_prompt, allocator), allocator);

            messages.PushBack(sys_msg, allocator);

        }

        for (const auto& input : inputs) {

            rapidjson::Value message(rapidjson::kObjectType);

            message.AddMember("role", "user", allocator);

            message.AddMember("content", rapidjson::Value(input.c_str(), allocator), allocator);

            messages.PushBack(message, allocator);

        }

        doc.AddMember("messages", messages, allocator);

        return Status::OK();

    }

};

// The OpenAI API format can be reused with some compatible AIs.

class OpenAIAdapter : public VoyageAIAdapter {

public:

    explicit OpenAIAdapter(bool chat_json_schema = true)

            : VoyageAIAdapter(true), _chat_json_schema(chat_json_schema) {}

    AIUsage parse_usage(const std::string& response_body) const override {

        return AIAdapter::parse_usage(response_body);

    }

    Status set_authentication(HttpClient* client) const override {

        client->set_header(HttpHeaders::AUTHORIZATION, "Bearer " + _config.api_key);

        client->set_content_type("application/json");

        return Status::OK();

    }

    Status build_request_payload(const std::vector<std::string>& inputs,

                                 const char* const system_prompt, std::string& request_body,

                                 bool structured_batch = false) const override {

        rapidjson::Document doc;

        doc.SetObject();

        auto& allocator = doc.GetAllocator();

        if (_config.endpoint.ends_with("responses")) {

            /*{

              "model": "gpt-4.1-mini",

              "input": [

                {"role": "system", "content": "system_prompt here"},

                {"role": "user", "content": "xxx"}

              ],

              "temperature": 0.7,

              "max_output_tokens": 150,

              "reasoning": {"effort": "max"},

              "text": {"format": <responses_format>}

            }*/

            doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator),

                          allocator);

            // If 'temperature' and 'max_tokens' are set, add them to the request body.

            if (_config.temperature != -1) {

                doc.AddMember("temperature", _config.temperature, allocator);

            }

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

            if (system_prompt && *system_prompt) {

                rapidjson::Value sys_msg(rapidjson::kObjectType);

                sys_msg.AddMember("role", "system", allocator);

                sys_msg.AddMember("content", rapidjson::Value(system_prompt, allocator), allocator);

                input.PushBack(sys_msg, allocator);

            }

            for (const auto& msg : inputs) {

                rapidjson::Value message(rapidjson::kObjectType);

                message.AddMember("role", "user", allocator);

                message.AddMember("content", rapidjson::Value(msg.c_str(), allocator), allocator);

                input.PushBack(message, allocator);

            }

            doc.AddMember("input", input, allocator);

        } else {

            /*{

              "model": "gpt-4",

              "messages": [

                {"role": "system", "content": "system_prompt here"},

                {"role": "user", "content": "xxx"}

              ],

              "temperature": x,

              "max_tokens": x,

              "reasoning_effort": "low",

              "response_format": <chat_format>

            }*/

            doc.AddMember("model", rapidjson::Value(_config.model_name.c_str(), allocator),

                          allocator);

            // If 'temperature' and 'max_tokens' are set, add them to the request body.

            if (_config.temperature != -1) {

                doc.AddMember("temperature", _config.temperature, allocator);

            }

            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) {

                rapidjson::Value sys_msg(rapidjson::kObjectType);

                sys_msg.AddMember("role", "system", allocator);

                sys_msg.AddMember("content", rapidjson::Value(system_prompt, allocator), allocator);