refactor(2671): move retry logic and implement tokio_retry

Cargo.toml

@@ -14,6 +14,7 @@ futures-util = { version = "0.3.30" }
 indexmap = "2.2.6"
 insta = { version = "1.38.0", features = ["json"] }
 tokio = { version = "1.37.0", features = ["rt", "time"] }
+tokio-retry = "0.3"
 reqwest = { version = "0.11", features = [
     "json",
     "rustls-tls",
@@ -61,6 +62,7 @@ rustls-pemfile = { version = "1.0.4" }
 schemars = { version = "0.8.17", features = ["derive"] }
 hyper = { version = "0.14.28", features = ["server"], default-features = false }
 tokio = { workspace = true }
+tokio-retry = { workspace = true }
 anyhow = { workspace = true }
 reqwest = { workspace = true }
 derive_setters = "0.1.6"

src/cli/llm/error.rs

@@ -1,11 +1,46 @@
-use derive_more::From;
-use strum_macros::Display;
+use reqwest::StatusCode;
+use thiserror::Error;
 
-#[derive(Debug, From, Display, thiserror::Error)]
+#[derive(Debug, Error)]
+pub enum WebcError {
+    #[error("Response failed with status {status}: {body}")]
+    ResponseFailedStatus { status: StatusCode, body: String },
+    #[error("Reqwest error: {0}")]
+    Reqwest(#[from] reqwest::Error),
+}
+
+#[derive(Debug, Error)]
 pub enum Error {
+    #[error("GenAI error: {0}")]
     GenAI(genai::Error),
+    #[error("Webc error: {0}")]
+    Webc(WebcError),
+    #[error("Empty response")]
     EmptyResponse,
-    Serde(serde_json::Error),
+    #[error("Serde error: {0}")]
+    Serde(#[from] serde_json::Error),
+}
+
+impl From<genai::Error> for Error {
+    fn from(err: genai::Error) -> Self {
+        if let genai::Error::WebModelCall { webc_error, .. } = &err {
+            let error_str = webc_error.to_string();
+            if error_str.contains("ResponseFailedStatus") {
+                // Extract status and body from the error message
+                let parts: Vec<&str> = error_str.splitn(3, ": ").collect();
+                if parts.len() >= 3 {
+                    if let Ok(status) = parts[1].parse::<u16>() {
+                        return Error::Webc(WebcError::ResponseFailedStatus {
+                            status: StatusCode::from_u16(status)
+                                .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR),
+                            body: parts[2].to_string(),
+                        });
+                    }
+                }
+            }
+        }
+        Error::GenAI(err)
+    }
 }
 
-pub type Result<A> = std::result::Result<A, Error>;
+pub type Result<T> = std::result::Result<T, Error>;

src/cli/llm/infer_type_name.rs

@@ -79,9 +79,7 @@ impl InferTypeName {
     }
     pub async fn generate(&mut self, config: &Config) -> Result<HashMap<String, String>> {
         let secret = self.secret.as_ref().map(|s| s.to_owned());
-
         let wizard: Wizard<Question, Answer> = Wizard::new(groq::LLAMA38192, secret);
-
         let mut new_name_mappings: HashMap<String, String> = HashMap::new();
 
         // removed root type from types.
@@ -102,47 +100,27 @@ impl InferTypeName {
                     .collect(),
             };
 
-            let mut delay = 3;
-            loop {
-                let answer = wizard.ask(question.clone()).await;
-                match answer {
-                    Ok(answer) => {
-                        let name = &answer.suggestions.join(", ");
-                        for name in answer.suggestions {
-                            if config.types.contains_key(&name)
-                                || new_name_mappings.contains_key(&name)
-                            {
-                                continue;
-                            }
-                            new_name_mappings.insert(name, type_name.to_owned());
-                            break;
+            match wizard.ask(question).await {
+                Ok(answer) => {
+                    let name = &answer.suggestions.join(", ");
+                    for name in answer.suggestions {
+                        if config.types.contains_key(&name) || new_name_mappings.contains_key(&name)
+                        {
+                            continue;
                         }
-                        tracing::info!(
-                            "Suggestions for {}: [{}] - {}/{}",
-                            type_name,
-                            name,
-                            i + 1,
-                            total
-                        );
-
-                        // TODO: case where suggested names are already used, then extend the base
-                        // question with `suggest different names, we have already used following
-                        // names: [names list]`
+                        new_name_mappings.insert(name, type_name.to_owned());
                         break;
                     }
-                    Err(e) => {
-                        // TODO: log errors after certain number of retries.
-                        if let Error::GenAI(_) = e {
-                            // TODO: retry only when it's required.
-                            tracing::warn!(
-                                "Unable to retrieve a name for the type '{}'. Retrying in {}s",
-                                type_name,
-                                delay
-                            );
-                            tokio::time::sleep(tokio::time::Duration::from_secs(delay)).await;
-                            delay *= std::cmp::min(delay * 2, 60);
-                        }
-                    }
+                    tracing::info!(
+                        "Suggestions for {}: [{}] - {}/{}",
+                        type_name,
+                        name,
+                        i + 1,
+                        total
+                    );
+                }
+                Err(e) => {
+                    tracing::error!("Failed to generate name for {}: {:?}", type_name, e);
                 }
             }
         }

src/cli/llm/wizard.rs

@@ -1,10 +1,15 @@
+// use std::borrow::Borrow;
+
 use derive_setters::Setters;
 use genai::adapter::AdapterKind;
 use genai::chat::{ChatOptions, ChatRequest, ChatResponse};
 use genai::resolver::AuthResolver;
 use genai::Client;
+use reqwest::StatusCode;
+use tokio_retry::strategy::{jitter, ExponentialBackoff};
+use tokio_retry::Retry;
 
-use super::Result;
+use super::error::{Error, Result, WebcError};
 use crate::cli::llm::model::Model;
 
 #[derive(Setters, Clone)]
@@ -41,13 +46,32 @@ impl<Q, A> Wizard<Q, A> {
 
     pub async fn ask(&self, q: Q) -> Result<A>
     where
-        Q: TryInto<ChatRequest, Error = super::Error>,
-        A: TryFrom<ChatResponse, Error = super::Error>,
+        Q: TryInto<ChatRequest, Error = Error> + Clone,
+        A: TryFrom<ChatResponse, Error = Error>,
     {
-        let response = self
-            .client
-            .exec_chat(self.model.as_str(), q.try_into()?, None)
-            .await?;
-        A::try_from(response)
+        let retry_strategy = ExponentialBackoff::from_millis(1000).map(jitter).take(5);
+
+        Retry::spawn(retry_strategy, || async {
+            let request = q.clone().try_into()?;
+            match self
+                .client
+                .exec_chat(self.model.as_str(), request, None)
+                .await
+            {
+                Ok(response) => Ok(A::try_from(response)?),
+                Err(err) => {
+                    let error = Error::from(err);
+                    match &error {
+                        Error::Webc(WebcError::ResponseFailedStatus { status, .. })
+                            if *status == StatusCode::TOO_MANY_REQUESTS =>
+                        {
+                            Err(error) // Propagate the error to trigger a retry
+                        }
+                        _ => Ok(Err(error)?), // Other errors are returned without retrying
+                    }
+                }
+            }
+        })
+        .await
     }
 }