run_context.h

Go to the documentation of this file.

 
#pragma once
 
#include <neograph/graph/cancel.h>
#include <neograph/graph/store.h>
#include <neograph/graph/types.h>
#include <neograph/graph/provider_call_broker.h>
#include <neograph/tool_dispatch.h>
#include <neograph/types.h>
 
#include <asio/any_io_executor.hpp>
 
#include <atomic>
#include <chrono>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
namespace sp { class NativeArchive; }
 
namespace neograph::graph {
 
struct ProviderOutcomes {
    mutable std::mutex mutex;
    std::vector<sp::runtime::Result> values;
    void add(sp::runtime::Result value) {
        std::lock_guard lock(mutex);
        values.push_back(std::move(value));
    }
    std::vector<sp::runtime::Result> snapshot() const {
        std::lock_guard lock(mutex);
        return values;
    }
};
 
struct ProviderLoopHistory {
    struct Entry {
        std::vector<sp::Message> messages;
        std::uint64_t turns = 0;
        bool client_calls_ready = false;
    };
    mutable std::mutex mutex;
    std::map<std::string, Entry> entries;
    Entry get(const std::string& task) const {
        std::lock_guard lock(mutex);
        const auto found = entries.find(task);
        return found == entries.end() ? Entry{} : found->second;
    }
    void set(const std::string& task, Entry entry) {
        std::lock_guard lock(mutex);
        entries.insert_or_assign(task, std::move(entry));
    }
    std::map<std::string, Entry> snapshot() const {
        std::lock_guard lock(mutex);
        return entries;
    }
};
 
struct RunContext {
    std::shared_ptr<CancelToken> cancel_token;
 
    std::shared_ptr<UsageAccumulator> usage;
    std::shared_ptr<ProviderOutcomes> provider_outcomes;
    std::shared_ptr<ProviderLoopHistory> provider_loop_history;
    std::function<void(const sp::Event&)> on_provider_event;
    std::shared_ptr<sp::NativeArchive> native_history_archive;
 
    std::uint64_t                     model_token_budget = 0;
    std::shared_ptr<std::atomic_bool> budget_exhausted;
    std::string run_id;
    std::shared_ptr<CancelToken> budget_cancel_token;
    std::shared_ptr<OwnedManagedBudgetLease> managed_budget_lease;
    std::shared_ptr<CheckpointStore> managed_budget_store;
 
    std::optional<std::chrono::steady_clock::time_point> deadline;
 
    bool shared_scheduler = false;
 
    asio::any_io_executor dispatch_executor;
 
    std::string trace_id;
 
    std::string thread_id;
 
    std::uint64_t cache_execution_id = 0;
 
    int step = 0;
 
    StreamMode stream_mode = StreamMode::ALL;
 
    std::optional<json> resume_value;
 
    std::string resume_interrupt_id;
 
    std::shared_ptr<Store> store;
 
    ToolGate tool_gate;
 
    std::shared_ptr<ToolExecutionController> tool_execution_controller;
 
    ToolExecutionIdentity tool_execution_identity;
 
};
 
NEOGRAPH_API ToolExecutionContext make_tool_execution_context(const RunContext& ctx);
NEOGRAPH_API std::shared_ptr<ProviderCallBroker> provider_call_broker(const RunContext& ctx);
inline ProviderCallIdentity make_provider_call_identity(const RunContext& ctx,
    const std::string& node, std::uint64_t ordinal = 0) {
    ProviderCallIdentity identity;
    identity.owner_scope = ctx.tool_execution_identity.owner_scope;
    identity.run_id = ctx.run_id;
    identity.thread_id = ctx.thread_id;
    identity.task_id = make_tool_execution_context(ctx).effect_task_id;
    identity.node_name = node;
    identity.call_ordinal = ordinal;
    identity.usage = ctx.usage;
    identity.model_token_budget = ctx.model_token_budget;
    identity.budget_exhausted = ctx.budget_exhausted;
    identity.budget_cancel_token = ctx.budget_cancel_token;
    identity.managed_budget_lease = ctx.managed_budget_lease;
    identity.managed_budget_store = ctx.managed_budget_store;
    return identity;
}
 
inline void record_usage(const RunContext& ctx, const sp::runtime::Result& result) {
    if (!result) throw std::runtime_error("Provider returned no outcome");
    if (ctx.usage && ctx.model_token_budget == 0 && !provider_call_broker(ctx))
        ctx.usage->add(outcome_usage(*result));
    if (ctx.provider_outcomes) ctx.provider_outcomes->add(result);
}
 
inline asio::awaitable<sp::runtime::Result> observe_provider_result(
    const RunContext& ctx, asio::awaitable<sp::runtime::Result> operation) {
    try {
        co_return co_await std::move(operation);
    } catch (const ProviderOutcomeError& error) {
        // An uncertain dispatch carries no outcome, so there is nothing to
        // record; it must still leave as the typed uncertain failure.
        if (error.outcome()) record_usage(ctx, error.outcome());
        throw;
    } catch (const ProviderFailure& error) {
        record_usage(ctx, error.outcome());
        throw;
    }
}
 
}  // namespace neograph::graph