rate_limited_provider.h

Go to the documentation of this file.

 
#pragma once
 
#include <neograph/api.h>
#include <neograph/provider.h>
 
#include <memory>
#include <string>
 
namespace neograph::llm {
 
class NEOGRAPH_API RateLimitedProvider : public Provider {
public:
    struct Config {
        int max_retries          = 3;   
        int default_wait_seconds = 30;  
        int max_wait_seconds     = 120; 
        int max_total_wait_seconds = 0;
    };
 
    static std::unique_ptr<RateLimitedProvider> create(
        std::shared_ptr<Provider> inner, Config cfg);
 
    static std::unique_ptr<RateLimitedProvider> create(
        std::shared_ptr<Provider> inner) {
        return create(std::move(inner), Config{});
    }
 
    asio::awaitable<ChatCompletion>
    complete_async(const CompletionParams& params) override;
 
    ChatCompletion complete_stream(const CompletionParams& params,
                                   const StreamCallback& on_chunk) override;
 
    asio::awaitable<ChatCompletion>
    invoke(const CompletionParams& params, StreamCallback on_chunk) override;
 
    std::string get_name() const override;
 
private:
    RateLimitedProvider(std::shared_ptr<Provider> inner, Config cfg);
 
    std::shared_ptr<Provider> inner_;
    Config                    cfg_;
};
 
} // namespace neograph::llm