routing_workflow.h

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

class RoutingWorkflow : public Workflow {
public:
    struct RouteHandler {
        String name;
        String description;
        String prompt_template;
        std::shared_ptr<LLMInterface> llm;  // Optional separate LLM for this route
        std::shared_ptr<Workflow> workflow;  // Optional workflow to delegate to
        std::function<JsonObject(const String&, const JsonObject&)> handler; // Optional function handler

        RouteHandler(
            const String& name,
            const String& description,
            const String& prompt_template = "",
            std::shared_ptr<LLMInterface> llm = nullptr,
            std::shared_ptr<Workflow> workflow = nullptr,
            std::function<JsonObject(const String&, const JsonObject&)> handler = nullptr
        ) : name(name), description(description),
            prompt_template(prompt_template), llm(llm), workflow(workflow), handler(handler) {}
    };

    RoutingWorkflow(
        std::shared_ptr<AgentContext> context,
        const String& router_prompt_template = ""
    );

    void addRouteHandler(const RouteHandler& handler);

    void addRouteHandler(
        const String& name,
        const String& description,
        const String& prompt_template = "",
        std::shared_ptr<LLMInterface> handler_llm = nullptr,
        std::shared_ptr<Workflow> workflow = nullptr
    );

    void init();

    JsonObject run(const String& input) override;

    void setRouterPromptTemplate(const String& prompt_template);

    void setRouterPrompt(const String& prompt_template) { setRouterPromptTemplate(prompt_template); }

    void setDefaultHandler(const RouteHandler& handler);

    void setDefaultRoute(std::function<JsonObject(const String&, const JsonObject&)> handler);

    JsonObject getRoutesSchema() const;

    void addRoute(
        const String& name,
        const String& description,
        std::function<JsonObject(const String&, const JsonObject&)> handler
    );
 
private:
    String router_prompt_template_;

    std::map<String, RouteHandler> route_handlers_;

    std::optional<RouteHandler> default_handler_;

    String createRouterSystemPrompt() const;
};
 
} // namespace workflows
} // namespace agents