Skip to content

Commit 49863dd

Browse files
fix(wizard): use RetryIf for more specific error handling
Signed-off-by: David Anyatonwu <davidanyatonwu@gmail.com>
1 parent 5ee7b1a commit 49863dd

File tree

2 files changed

+16
-28
lines changed

2 files changed

+16
-28
lines changed

src/cli/llm/error.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,8 @@ impl From<genai::Error> for Error {
3838
}
3939
}
4040
}
41-
}
42-
Error::GenAI(err)
41+
};
42+
err.into()
4343
}
4444
}
4545

src/cli/llm/wizard.rs

Lines changed: 14 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,12 @@
1-
1+
use super::error::{Error, Result, WebcError};
22
use derive_setters::Setters;
33
use genai::adapter::AdapterKind;
44
use genai::chat::{ChatOptions, ChatRequest, ChatResponse};
55
use genai::resolver::AuthResolver;
66
use genai::Client;
77
use reqwest::StatusCode;
88
use tokio_retry::strategy::{jitter, ExponentialBackoff};
9-
use tokio_retry::Retry;
10-
use super::error::{Error, Result, WebcError};
11-
use crate::cli::llm::model::Model;
12-
9+
use tokio_retry::RetryIf;
1310

1411
#[derive(Setters, Clone)]
1512
pub struct Wizard<Q, A> {
@@ -50,27 +47,18 @@ impl<Q, A> Wizard<Q, A> {
5047
{
5148
let retry_strategy = ExponentialBackoff::from_millis(1000).map(jitter).take(5);
5249

53-
Retry::spawn(retry_strategy, || async {
54-
let request = q.clone().try_into()?;
55-
match self
56-
.client
57-
.exec_chat(self.model.as_str(), request, None)
58-
.await
59-
{
60-
Ok(response) => Ok(A::try_from(response)?),
61-
Err(err) => {
62-
let error = Error::from(err);
63-
match &error {
64-
Error::Webc(WebcError::ResponseFailedStatus { status, .. })
65-
if *status == StatusCode::TOO_MANY_REQUESTS =>
66-
{
67-
Err(error) // Propagate the error to trigger a retry
68-
}
69-
_ => Ok(Err(error)?), // Other errors are returned without retrying
70-
}
71-
}
72-
}
73-
})
50+
RetryIf::spawn(
51+
retry_strategy,
52+
|| async {
53+
let request = q.clone().try_into()?;
54+
self.client
55+
.exec_chat(self.model.as_str(), request, None)
56+
.await
57+
.map_err(Error::from)
58+
.and_then(A::try_from)
59+
},
60+
|err: &Error| matches!(err, Error::Webc(WebcError::ResponseFailedStatus { status, .. }) if *status == StatusCode::TOO_MANY_REQUESTS)
61+
)
7462
.await
7563
}
7664
}

0 commit comments

Comments
 (0)