fix(wizard): use RetryIf for more specific error handling

Signed-off-by: David Anyatonwu <[email protected]>

src/cli/llm/error.rs

@@ -38,8 +38,8 @@ impl From<genai::Error> for Error {
                     }
                 }
             }
-        }
-        Error::GenAI(err)
+        };
+        err.into()
     }
 }
 

src/cli/llm/wizard.rs

@@ -1,15 +1,12 @@
-
+use super::error::{Error, Result, WebcError};
 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::error::{Error, Result, WebcError};
-use crate::cli::llm::model::Model;
-
+use tokio_retry::RetryIf;
 
 #[derive(Setters, Clone)]
 pub struct Wizard<Q, A> {
@@ -50,27 +47,18 @@ impl<Q, A> Wizard<Q, A> {
     {
         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
-                    }
-                }
-            }
-        })
+        RetryIf::spawn(
+            retry_strategy,
+            || async {
+                let request = q.clone().try_into()?;
+                self.client
+                    .exec_chat(self.model.as_str(), request, None)
+                    .await
+                    .map_err(Error::from)
+                    .and_then(A::try_from)
+            },
+            |err: &Error| matches!(err, Error::Webc(WebcError::ResponseFailedStatus { status, .. }) if *status == StatusCode::TOO_MANY_REQUESTS)
+        )
         .await
     }
 }