|
1 |
| - |
| 1 | +use super::error::{Error, Result, WebcError}; |
2 | 2 | use derive_setters::Setters;
|
3 | 3 | use genai::adapter::AdapterKind;
|
4 | 4 | use genai::chat::{ChatOptions, ChatRequest, ChatResponse};
|
5 | 5 | use genai::resolver::AuthResolver;
|
6 | 6 | use genai::Client;
|
7 | 7 | use reqwest::StatusCode;
|
8 | 8 | 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; |
13 | 10 |
|
14 | 11 | #[derive(Setters, Clone)]
|
15 | 12 | pub struct Wizard<Q, A> {
|
@@ -50,27 +47,18 @@ impl<Q, A> Wizard<Q, A> {
|
50 | 47 | {
|
51 | 48 | let retry_strategy = ExponentialBackoff::from_millis(1000).map(jitter).take(5);
|
52 | 49 |
|
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 | + ) |
74 | 62 | .await
|
75 | 63 | }
|
76 | 64 | }
|
0 commit comments