parallelization_workflow.h

#pragma once
 
#include <agents-cpp/workflow.h>
#include <functional>
#include <vector>
 
namespace agents {
namespace workflows {

class ParallelizationWorkflow : public Workflow {
public:
    enum class Strategy {
        SECTIONING,
        VOTING
    };

    struct Task {
        String name;
        String prompt_template;
        JsonObject context;
        std::function<String(const String&)> prompt_fn;
        std::function<JsonObject(const String&)> result_parser;

        Task(
            const String& name,
            const String& prompt_template,
            const JsonObject& context = JsonObject(),
            std::function<String(const String&)> prompt_fn = nullptr,
            std::function<JsonObject(const String&)> result_parser = nullptr
        ) : name(name), prompt_template(prompt_template), context(context), prompt_fn(prompt_fn), result_parser(result_parser) {}
    };

    ParallelizationWorkflow(
        std::shared_ptr<AgentContext> context,
        Strategy strategy = Strategy::SECTIONING
    );

    void addTask(const Task& task);

    void addTask(
        const String& name,
        const String& prompt_template,
        const JsonObject& context = JsonObject()
    );

    void addTask(
        const String& name,
        const String& prompt_template,
        std::function<String(const String&)> prompt_fn,
        std::function<JsonObject(const String&)> result_parser,
        const JsonObject& context = JsonObject()
    );

    void setAggregator(std::function<JsonObject(const std::vector<JsonObject>&)> aggregator);

    void setStrategy(Strategy strategy);

    void init();

    void setVotingThreshold(double threshold);

    JsonObject run(const String& input) override;

    JsonObject run();
 
private:
    std::vector<Task> tasks_;

    Strategy strategy_;

    std::function<JsonObject(const std::vector<JsonObject>&)> aggregator_;

    double voting_threshold_ = 0.5;

    std::shared_ptr<LLMInterface> llm_;

    std::vector<JsonObject> runTasksInParallel(const String& input);

    static JsonObject defaultSectioningAggregator(const std::vector<JsonObject>& results);

    static JsonObject defaultVotingAggregator(const std::vector<JsonObject>& results);
};
 
} // namespace workflows
} // namespace agents