From 376d75c519674e9d8265099610bf59671ac7c131 Mon Sep 17 00:00:00 2001 From: David Anyatonwu Date: Mon, 12 Aug 2024 17:39:15 +0000 Subject: [PATCH 1/6] refactor(2671): implement tokio_retry for API rate limiting --- Cargo.lock | 12 +++++++ Cargo.toml | 2 ++ src/cli/llm/infer_type_name.rs | 60 ++++++++++++---------------------- src/cli/llm/wizard.rs | 26 ++++++++++----- 4 files changed, 52 insertions(+), 48 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0761018274..5965d11c65 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5530,6 +5530,7 @@ dependencies = [ "test-log", "thiserror", "tokio", + "tokio-retry", "tokio-test", "tonic 0.11.0", "tonic-types", @@ -5948,6 +5949,17 @@ dependencies = [ "syn 2.0.74", ] +[[package]] +name = "tokio-retry" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f57eb36ecbe0fc510036adff84824dd3c24bb781e21bfa67b69d556aa85214f" +dependencies = [ + "pin-project", + "rand", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.24.1" diff --git a/Cargo.toml b/Cargo.toml index 8c9ad84e2e..dc44b861cb 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/cli/llm/infer_type_name.rs b/src/cli/llm/infer_type_name.rs index ed0869dfd5..ffe0a634d1 100644 --- a/src/cli/llm/infer_type_name.rs +++ b/src/cli/llm/infer_type_name.rs @@ -6,6 +6,8 @@ use serde::{Deserialize, Serialize}; use super::model::groq; use super::{Error, Result, Wizard}; use crate::core::config::Config; +use tokio_retry::strategy::{jitter, ExponentialBackoff}; +use tokio_retry::Retry; #[derive(Default)] pub struct InferTypeName { @@ -79,9 +81,7 @@ impl InferTypeName { } pub async fn generate(&mut self, config: &Config) -> Result> { let secret = self.secret.as_ref().map(|s| s.to_owned()); - let wizard: Wizard = Wizard::new(groq::LLAMA38192, secret); - let mut new_name_mappings: HashMap = HashMap::new(); // removed root type from types. @@ -102,47 +102,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); } } } diff --git a/src/cli/llm/wizard.rs b/src/cli/llm/wizard.rs index 1604d7f15f..cfefd11562 100644 --- a/src/cli/llm/wizard.rs +++ b/src/cli/llm/wizard.rs @@ -3,8 +3,10 @@ use genai::adapter::AdapterKind; use genai::chat::{ChatOptions, ChatRequest, ChatResponse}; use genai::resolver::AuthResolver; use genai::Client; +use tokio_retry::strategy::{jitter, ExponentialBackoff}; +use tokio_retry::Retry; -use super::Result; +use super::error::{Error, Result}; use crate::cli::llm::model::Model; #[derive(Setters, Clone)] @@ -41,13 +43,21 @@ impl Wizard { pub async fn ask(&self, q: Q) -> Result where - Q: TryInto, - A: TryFrom, + Q: TryInto + Clone, + A: TryFrom, { - 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()?; + let response = self + .client + .exec_chat(self.model.as_str(), request, None) + .await + .map_err(Error::GenAI)?; + + A::try_from(response) + }) + .await } } From 463373c4dc07209be28f3f8094d4b8a3311bbb70 Mon Sep 17 00:00:00 2001 From: David Anyatonwu Date: Mon, 12 Aug 2024 19:39:01 +0000 Subject: [PATCH 2/6] lint fix --- src/cli/llm/infer_type_name.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/cli/llm/infer_type_name.rs b/src/cli/llm/infer_type_name.rs index ffe0a634d1..861300b5cf 100644 --- a/src/cli/llm/infer_type_name.rs +++ b/src/cli/llm/infer_type_name.rs @@ -6,8 +6,6 @@ use serde::{Deserialize, Serialize}; use super::model::groq; use super::{Error, Result, Wizard}; use crate::core::config::Config; -use tokio_retry::strategy::{jitter, ExponentialBackoff}; -use tokio_retry::Retry; #[derive(Default)] pub struct InferTypeName { From a86526f5e562c34833e94e84b0a2b697315ddf6b Mon Sep 17 00:00:00 2001 From: David Anyatonwu Date: Tue, 13 Aug 2024 11:41:39 +0000 Subject: [PATCH 3/6] Refine retry logic in Wizard's ask method --- src/cli/llm/wizard.rs | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/src/cli/llm/wizard.rs b/src/cli/llm/wizard.rs index cfefd11562..f2ef3ce388 100644 --- a/src/cli/llm/wizard.rs +++ b/src/cli/llm/wizard.rs @@ -50,13 +50,35 @@ impl Wizard { Retry::spawn(retry_strategy, || async { let request = q.clone().try_into()?; - let response = self + match self .client .exec_chat(self.model.as_str(), request, None) .await - .map_err(Error::GenAI)?; - - A::try_from(response) + { + Ok(response) => Ok(A::try_from(response)?), + Err(genai::Error::WebModelCall { webc_error, .. }) => { + if webc_error.to_string().contains("429") { + Err(Error::GenAI(genai::Error::WebModelCall { + model_info: genai::ModelInfo::new( + AdapterKind::from_model(self.model.as_str()) + .unwrap_or(AdapterKind::Ollama), + self.model.as_str(), + ), + webc_error, + })) + } else { + Ok(Err(Error::GenAI(genai::Error::WebModelCall { + model_info: genai::ModelInfo::new( + AdapterKind::from_model(self.model.as_str()) + .unwrap_or(AdapterKind::Ollama), + self.model.as_str(), + ), + webc_error, + }))?) + } + } + Err(e) => Ok(Err(Error::GenAI(e))?), + } }) .await } From 2a1e1d7b68707240cc4a4369810896b3514aaed6 Mon Sep 17 00:00:00 2001 From: David Anyatonwu Date: Tue, 20 Aug 2024 23:28:55 +0000 Subject: [PATCH 4/6] using match arm instead of string comparison Signed-off-by: David Anyatonwu --- src/cli/llm/error.rs | 45 ++++++++++++++++++++++++++++++++++++++----- src/cli/llm/wizard.rs | 34 +++++++++++++------------------- 2 files changed, 53 insertions(+), 26 deletions(-) diff --git a/src/cli/llm/error.rs b/src/cli/llm/error.rs index c2b44f9ca3..b0ceb2e7ed 100644 --- a/src/cli/llm/error.rs +++ b/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 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::() { + 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 = std::result::Result; +pub type Result = std::result::Result; diff --git a/src/cli/llm/wizard.rs b/src/cli/llm/wizard.rs index f2ef3ce388..12bfac079a 100644 --- a/src/cli/llm/wizard.rs +++ b/src/cli/llm/wizard.rs @@ -1,12 +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::error::{Error, Result}; +use super::error::{Error, Result, WebcError}; use crate::cli::llm::model::Model; #[derive(Setters, Clone)] @@ -56,28 +59,17 @@ impl Wizard { .await { Ok(response) => Ok(A::try_from(response)?), - Err(genai::Error::WebModelCall { webc_error, .. }) => { - if webc_error.to_string().contains("429") { - Err(Error::GenAI(genai::Error::WebModelCall { - model_info: genai::ModelInfo::new( - AdapterKind::from_model(self.model.as_str()) - .unwrap_or(AdapterKind::Ollama), - self.model.as_str(), - ), - webc_error, - })) - } else { - Ok(Err(Error::GenAI(genai::Error::WebModelCall { - model_info: genai::ModelInfo::new( - AdapterKind::from_model(self.model.as_str()) - .unwrap_or(AdapterKind::Ollama), - self.model.as_str(), - ), - webc_error, - }))?) + 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 } } - Err(e) => Ok(Err(Error::GenAI(e))?), } }) .await From 0be7933c6d1214bafb5bec765df3a2526db67fd5 Mon Sep 17 00:00:00 2001 From: David Anyatonwu <51977119+onyedikachi-david@users.noreply.github.com> Date: Tue, 27 Aug 2024 21:27:42 +0100 Subject: [PATCH 5/6] Update src/cli/llm/wizard.rs Co-authored-by: Tushar Mathur --- src/cli/llm/wizard.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/cli/llm/wizard.rs b/src/cli/llm/wizard.rs index 12bfac079a..b7ae0e6120 100644 --- a/src/cli/llm/wizard.rs +++ b/src/cli/llm/wizard.rs @@ -1,4 +1,3 @@ -// use std::borrow::Borrow; use derive_setters::Setters; use genai::adapter::AdapterKind; From 49863ddf66bf7bd33355ae11cc0140518be6a089 Mon Sep 17 00:00:00 2001 From: David Anyatonwu Date: Tue, 27 Aug 2024 21:03:43 +0000 Subject: [PATCH 6/6] fix(wizard): use RetryIf for more specific error handling Signed-off-by: David Anyatonwu --- src/cli/llm/error.rs | 4 ++-- src/cli/llm/wizard.rs | 40 ++++++++++++++-------------------------- 2 files changed, 16 insertions(+), 28 deletions(-) diff --git a/src/cli/llm/error.rs b/src/cli/llm/error.rs index b0ceb2e7ed..49f00cefd2 100644 --- a/src/cli/llm/error.rs +++ b/src/cli/llm/error.rs @@ -38,8 +38,8 @@ impl From for Error { } } } - } - Error::GenAI(err) + }; + err.into() } } diff --git a/src/cli/llm/wizard.rs b/src/cli/llm/wizard.rs index de5ba79003..2788926c47 100644 --- a/src/cli/llm/wizard.rs +++ b/src/cli/llm/wizard.rs @@ -1,4 +1,4 @@ - +use super::error::{Error, Result, WebcError}; use derive_setters::Setters; use genai::adapter::AdapterKind; use genai::chat::{ChatOptions, ChatRequest, ChatResponse}; @@ -6,10 +6,7 @@ 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 { @@ -50,27 +47,18 @@ impl Wizard { { 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 } }