From 510d7849a6b6804e64c5c5fbe7bc93ffd7b02688 Mon Sep 17 00:00:00 2001
From: Roger Zurawicki <zurawiki@users.noreply.github.com>
Date: Thu, 2 Mar 2023 23:41:22 -0500
Subject: [PATCH] Add OpenAI API retry functionality and backoff crate.

- Add retry logic to the OpenAI API configuration
- Implement a new setting for OpenAI API retries
---
 Cargo.lock         | 1 +
 Cargo.toml         | 1 +
 src/llms/openai.rs | 9 ++++++++-
 src/settings.rs    | 3 +++
 src/toml.rs        | 2 ++
 5 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 07f517e..690bdb0 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -876,6 +876,7 @@ dependencies = [
  "async-openai",
  "async-std",
  "async-trait",
+ "backoff",
  "clap",
  "colored",
  "config",
diff --git a/Cargo.toml b/Cargo.toml
index ec49174..fe96083 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,6 +21,7 @@ anyhow = "1.0.69"
 async-openai = "0.8.0"
 async-std = "1.12.0"
 async-trait = "0.1.64"
+backoff = "0.4.0"
 clap = { version = "4.1.8", features = ["derive"] }
 colored = "2.0.0"
 config = { version = "0.13.3", features = ["toml"] }
diff --git a/src/llms/openai.rs b/src/llms/openai.rs
index 5d21b50..912a85b 100644
--- a/src/llms/openai.rs
+++ b/src/llms/openai.rs
@@ -32,7 +32,14 @@ impl OpenAIClient {
             bail!("No OpenAI model configured.")
         }
 
-        let client = Client::new().with_api_key(&api_key);
+        let mut client = Client::new().with_api_key(&api_key);
+
+        if settings.retries.unwrap_or_default() > 0 {
+            let backoff = backoff::ExponentialBackoffBuilder::new()
+                .with_max_elapsed_time(Some(std::time::Duration::from_secs(60)))
+                .build();
+            client = client.with_backoff(backoff);
+        }
         Ok(Self { model, client })
     }
 
diff --git a/src/settings.rs b/src/settings.rs
index a526839..f139b7e 100644
--- a/src/settings.rs
+++ b/src/settings.rs
@@ -73,6 +73,7 @@ impl<'de> serde::Deserialize<'de> for ModelProvider {
 pub(crate) struct OpenAISettings {
     pub api_key: Option<String>,
     pub model: Option<String>,
+    pub retries: Option<u16>,
 }
 
 // implement the trait `From<OpenAISettings>` for `ValueKind`
@@ -81,6 +82,7 @@ impl From<OpenAISettings> for config::ValueKind {
         let mut properties = HashMap::new();
         properties.insert("api_key".to_string(), config::Value::from(settings.api_key));
         properties.insert("model".to_string(), config::Value::from(settings.model));
+        properties.insert("retries".to_string(), config::Value::from(settings.retries));
         Self::Table(properties)
     }
 }
@@ -198,6 +200,7 @@ impl Settings {
                 Some(OpenAISettings {
                     api_key: None,
                     model: Some(DEFAULT_OPENAI_MODEL.to_string()),
+                    retries: Some(2),
                 }),
             )?
             .set_default(
diff --git a/src/toml.rs b/src/toml.rs
index 59df7f0..eb75471 100644
--- a/src/toml.rs
+++ b/src/toml.rs
@@ -87,6 +87,7 @@ the-force = { value = "surrounds-you" }
                 "model_provider",
                 "openai.api_key",
                 "openai.model",
+                "openai.retries",
                 "output.lang",
                 "prompt.commit_summary",
                 "prompt.commit_title",
@@ -108,6 +109,7 @@ the-force = { value = "surrounds-you" }
                 "model_provider",
                 "openai.api_key",
                 "openai.model",
+                "openai.retries",
                 "output.lang",
                 "prompt.commit_summary",
                 "prompt.commit_title",