engine.h

Go to the documentation of this file.

 
#pragma once
 
#include <neograph/api.h>
#include <neograph/graph/cancel.h>
#include <neograph/graph/checkpoint.h>
#include <neograph/graph/compiler.h>
#include <neograph/graph/coordinator.h>
#include <neograph/graph/executor.h>
#include <neograph/graph/node.h>
#include <neograph/graph/node_cache.h>
#include <neograph/graph/scheduler.h>
#include <neograph/graph/state.h>
#include <neograph/graph/store.h>
#include <neograph/graph/types.h>
 
#include <asio/awaitable.hpp>
#include <asio/thread_pool.hpp>
 
#include <atomic>
#include <chrono>
#include <cstddef>
#include <memory>
#include <optional>
#include <set>
#include <string>
 
namespace neograph::graph {
 
struct RunConfig {
    std::string thread_id;                          
    json        input;                              
    int         max_steps  = 50;                    
    StreamMode  stream_mode = StreamMode::ALL;      
 
    std::shared_ptr<CancelToken> cancel_token;
 
    bool resume_if_exists = false;
};
 
struct RunContext {
    std::shared_ptr<CancelToken> cancel_token;
 
    std::optional<std::chrono::steady_clock::time_point> deadline;
 
    std::string trace_id;
 
    std::string thread_id;
 
    int step = 0;
 
    StreamMode stream_mode = StreamMode::ALL;
 
    std::shared_ptr<Store> store;
};
 
struct RunResult {
    json        output;                             
    bool        interrupted       = false;          
    std::string interrupt_node;                     
    json        interrupt_value;                    
    std::string checkpoint_id;                      
    std::vector<std::string> execution_trace;       
 
    template <typename T>
    T channel(const std::string& name) const {
        return channel_raw(name).get<T>();
    }
 
    json channel_raw(const std::string& name) const {
        if (output.contains("channels") && output["channels"].contains(name)) {
            const auto wrapped = output["channels"][name];
            if (wrapped.contains("value")) return wrapped["value"];
        }
        if (output.contains(name)) return output[name];
        throw json::out_of_range(
            "RunResult::channel: no such channel '" + name + "' in output "
            "(checked output[\"channels\"][\"" + name + "\"][\"value\"] and output[\"" + name + "\"])");
    }
 
    bool has_channel(const std::string& name) const noexcept {
        if (output.contains("channels") && output["channels"].contains(name)
            && output["channels"][name].contains("value")) {
            return true;
        }
        return output.contains(name);
    }
};
 
class NEOGRAPH_API GraphEngine {
public:
    static std::unique_ptr<GraphEngine> compile(
        const json& definition,
        const NodeContext& default_context,
        std::shared_ptr<CheckpointStore> store = nullptr);
 
    RunResult run(const RunConfig& config);
 
    asio::awaitable<RunResult> run_async(RunConfig config);
 
    RunResult run_stream(const RunConfig& config,
                         const GraphStreamCallback& cb);
 
    asio::awaitable<RunResult> run_stream_async(
        RunConfig config, GraphStreamCallback cb);
 
    RunResult resume(const std::string& thread_id,
                     const json& resume_value = json(),
                     const GraphStreamCallback& cb = nullptr);
 
    asio::awaitable<RunResult> resume_async(
        const std::string& thread_id,
        const json& resume_value = json(),
        const GraphStreamCallback& cb = nullptr);
 
    // ── State inspection & manipulation (LangGraph Checkpointer API) ──
 
    std::optional<json> get_state(const std::string& thread_id) const;
 
    std::vector<Checkpoint> get_state_history(const std::string& thread_id,
                                              int limit = 100) const;
 
    void update_state(const std::string& thread_id,
                      const json& channel_writes,
                      const std::string& as_node = "");
 
    std::string fork(const std::string& source_thread_id,
                     const std::string& new_thread_id,
                     const std::string& checkpoint_id = "");
 
    // ── Configuration ──
 
    void own_tools(std::vector<std::unique_ptr<Tool>> tools);
 
    void set_checkpoint_store(std::shared_ptr<CheckpointStore> store);
 
    void set_store(std::shared_ptr<Store> store);
 
    std::shared_ptr<Store> get_store() const { return store_; }
 
    void set_retry_policy(const RetryPolicy& policy);
 
    void set_node_retry_policy(const std::string& node_name, const RetryPolicy& policy);
 
    void set_worker_count(std::size_t n);
 
    void set_worker_count_auto();
 
    void set_node_cache_enabled(const std::string& node_name, bool enabled);
 
    void clear_node_cache();
 
    const NodeCache& node_cache() const { return node_cache_; }
 
    const std::string& get_graph_name() const { return name_; }
 
private:
    GraphEngine() = default;
 
    void init_state(GraphState& state) const;
    void apply_input(GraphState& state, const json& input) const;
 
    asio::awaitable<RunResult> execute_graph_async(
        const RunConfig& config,
        const GraphStreamCallback& cb,
        const std::vector<std::string>& resume_from = {},
        const json& resume_value = json());
 
    RetryPolicy get_retry_policy(const std::string& node_name) const;
 
    // --- Graph definition ---
    std::string name_;
 
    std::vector<ChannelDef> channel_defs_;
 
    std::map<std::string, std::unique_ptr<GraphNode>> nodes_;
    std::vector<Edge>            edges_;
    std::vector<ConditionalEdge> conditional_edges_;
 
    std::unique_ptr<Scheduler> scheduler_;
 
    std::unique_ptr<NodeExecutor> executor_;
 
    std::set<std::string> interrupt_before_;
    std::set<std::string> interrupt_after_;
 
    std::shared_ptr<CheckpointStore> checkpoint_store_;
    std::shared_ptr<Store>           store_;
    std::vector<std::unique_ptr<Tool>> owned_tools_;
 
    // Retry policies
    RetryPolicy default_retry_policy_;
    std::map<std::string, RetryPolicy> node_retry_policies_;
 
    std::unique_ptr<asio::thread_pool> pool_;
 
    NodeCache node_cache_;
 
    std::atomic<int> active_runs_{0};
};
 
} // namespace neograph::graph