From ef4516251cf0a59a50ebf4c975fe65de5f90429c Mon Sep 17 00:00:00 2001 From: Sandipsinh Rathod Date: Mon, 12 Aug 2024 22:41:09 +0530 Subject: [PATCH 1/6] refactor(2671): Move retry logic from infer_type_name to wizard --- Cargo.lock | 12 +++++++ Cargo.toml | 1 + src/cli/llm/infer_type_name.rs | 66 +++++++++++++++------------------- src/cli/llm/wizard.rs | 42 ++++++++++++++++------ 4 files changed, 72 insertions(+), 49 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5f9809096c..c2d1a99a71 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..cc761a26b8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -169,6 +169,7 @@ indenter = "0.3.3" derive_more = { workspace = true } enum_dispatch = "0.3.13" strum = "0.26.2" +tokio-retry = "0.3.0" [dev-dependencies] tailcall-prettier = { path = "tailcall-prettier" } diff --git a/src/cli/llm/infer_type_name.rs b/src/cli/llm/infer_type_name.rs index ed0869dfd5..38548d48aa 100644 --- a/src/cli/llm/infer_type_name.rs +++ b/src/cli/llm/infer_type_name.rs @@ -102,47 +102,37 @@ 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; + let answer = wizard.ask(question).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; } - 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 + ); + + // 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]` + break; + } + Err(e) => { + tracing::warn!( + "Unable to retrieve a name for the type '{}', skipping with error {}", + type_name, + e + ); } } } diff --git a/src/cli/llm/wizard.rs b/src/cli/llm/wizard.rs index 862e205b84..42f82d5f2c 100644 --- a/src/cli/llm/wizard.rs +++ b/src/cli/llm/wizard.rs @@ -1,14 +1,17 @@ +use std::sync::Arc; + use derive_setters::Setters; use genai::adapter::AdapterKind; use genai::chat::{ChatOptions, ChatRequest, ChatResponse}; use genai::Client; +use tokio_retry::strategy::ExponentialBackoff; use super::Result; use crate::cli::llm::model::Model; #[derive(Setters, Clone)] pub struct Wizard { - client: Client, + client: Arc, model: Model, _q: std::marker::PhantomData, _a: std::marker::PhantomData, @@ -22,23 +25,23 @@ impl Wizard { } let adapter = AdapterKind::from_model(model.as_str()).unwrap_or(AdapterKind::Ollama); - + let client = Client::builder() + .with_chat_options( + ChatOptions::default() + .with_json_mode(true) + .with_temperature(0.0), + ) + .insert_adapter_config(adapter, config) + .build(); Self { - client: Client::builder() - .with_chat_options( - ChatOptions::default() - .with_json_mode(true) - .with_temperature(0.0), - ) - .insert_adapter_config(adapter, config) - .build(), + client: Arc::new(client), model, _q: Default::default(), _a: Default::default(), } } - pub async fn ask(&self, q: Q) -> Result + async fn ask_inner(&self, q: Q) -> Result where Q: TryInto, A: TryFrom, @@ -47,6 +50,23 @@ impl Wizard { .client .exec_chat(self.model.as_str(), q.try_into()?, None) .await?; + A::try_from(response) } + + pub async fn ask(&self, q: Q) -> Result + where + Q: TryInto + Clone, + A: TryFrom, + { + let retry_strategy = ExponentialBackoff::from_millis(3) + .map(tokio_retry::strategy::jitter) + .take(3); + + let result = + tokio_retry::Retry::spawn(retry_strategy, || async { self.ask_inner(q.clone()).await }) + .await; + + result + } } From 073305312dd87300aa3522ec4fc6c1a587686b38 Mon Sep 17 00:00:00 2001 From: Sandipsinh Rathod Date: Mon, 12 Aug 2024 22:42:25 +0530 Subject: [PATCH 2/6] lint fixes --- src/cli/llm/wizard.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/cli/llm/wizard.rs b/src/cli/llm/wizard.rs index 42f82d5f2c..21bb1b5155 100644 --- a/src/cli/llm/wizard.rs +++ b/src/cli/llm/wizard.rs @@ -63,10 +63,9 @@ impl Wizard { .map(tokio_retry::strategy::jitter) .take(3); - let result = - tokio_retry::Retry::spawn(retry_strategy, || async { self.ask_inner(q.clone()).await }) - .await; + - result + tokio_retry::Retry::spawn(retry_strategy, || async { self.ask_inner(q.clone()).await }) + .await } } From 07a7a22f645f896c7ea5000a1b62d59376d7ea3b Mon Sep 17 00:00:00 2001 From: Sandipsinh Rathod Date: Mon, 12 Aug 2024 22:47:13 +0530 Subject: [PATCH 3/6] merge main --- src/cli/llm/wizard.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/cli/llm/wizard.rs b/src/cli/llm/wizard.rs index 130e715abf..478f82f60c 100644 --- a/src/cli/llm/wizard.rs +++ b/src/cli/llm/wizard.rs @@ -1,5 +1,3 @@ -use std::sync::Arc; - use derive_setters::Setters; use genai::adapter::AdapterKind; use genai::chat::{ChatOptions, ChatRequest, ChatResponse}; @@ -64,9 +62,7 @@ impl Wizard { .map(tokio_retry::strategy::jitter) .take(3); - - tokio_retry::Retry::spawn(retry_strategy, || async { self.ask_inner(q.clone()).await }) - .await + .await } } From 8ddeb9dd4e18c2386b2cce4406a48e4d4841d8a8 Mon Sep 17 00:00:00 2001 From: Sandipsinh Rathod Date: Mon, 12 Aug 2024 22:49:28 +0530 Subject: [PATCH 4/6] simplify logic --- src/cli/llm/wizard.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/cli/llm/wizard.rs b/src/cli/llm/wizard.rs index 478f82f60c..8c21d495d2 100644 --- a/src/cli/llm/wizard.rs +++ b/src/cli/llm/wizard.rs @@ -62,7 +62,6 @@ impl Wizard { .map(tokio_retry::strategy::jitter) .take(3); - tokio_retry::Retry::spawn(retry_strategy, || async { self.ask_inner(q.clone()).await }) - .await + tokio_retry::Retry::spawn(retry_strategy, || self.ask_inner(q.clone())).await } } From fa6070cbc46bb91bc46f9ad222304ecbd0ca6f42 Mon Sep 17 00:00:00 2001 From: Sandipsinh Rathod Date: Mon, 12 Aug 2024 22:54:51 +0530 Subject: [PATCH 5/6] inline `ask_inner` --- src/cli/llm/wizard.rs | 23 +++++++++-------------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/src/cli/llm/wizard.rs b/src/cli/llm/wizard.rs index 8c21d495d2..c9fde06356 100644 --- a/src/cli/llm/wizard.rs +++ b/src/cli/llm/wizard.rs @@ -40,19 +40,6 @@ impl Wizard { } } - async fn ask_inner(&self, q: Q) -> Result - where - Q: TryInto, - A: TryFrom, - { - let response = self - .client - .exec_chat(self.model.as_str(), q.try_into()?, None) - .await?; - - A::try_from(response) - } - pub async fn ask(&self, q: Q) -> Result where Q: TryInto + Clone, @@ -62,6 +49,14 @@ impl Wizard { .map(tokio_retry::strategy::jitter) .take(3); - tokio_retry::Retry::spawn(retry_strategy, || self.ask_inner(q.clone())).await + tokio_retry::Retry::spawn(retry_strategy, || async { + let response = self + .client + .exec_chat(self.model.as_str(), q.clone().try_into()?, None) + .await?; + + A::try_from(response) + }) + .await } } From 64d7f325427235307757f089660268f4337c3308 Mon Sep 17 00:00:00 2001 From: Tushar Mathur Date: Tue, 13 Aug 2024 13:49:04 +0530 Subject: [PATCH 6/6] update retry-strategy --- src/cli/llm/wizard.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/cli/llm/wizard.rs b/src/cli/llm/wizard.rs index c9fde06356..25889db17c 100644 --- a/src/cli/llm/wizard.rs +++ b/src/cli/llm/wizard.rs @@ -1,3 +1,5 @@ +use std::time::Duration; + use derive_setters::Setters; use genai::adapter::AdapterKind; use genai::chat::{ChatOptions, ChatRequest, ChatResponse}; @@ -45,11 +47,12 @@ impl Wizard { Q: TryInto + Clone, A: TryFrom, { - let retry_strategy = ExponentialBackoff::from_millis(3) + let retry = ExponentialBackoff::from_millis(1000) + .max_delay(Duration::from_secs(60)) .map(tokio_retry::strategy::jitter) - .take(3); + .take(10); - tokio_retry::Retry::spawn(retry_strategy, || async { + tokio_retry::Retry::spawn(retry, || async { let response = self .client .exec_chat(self.model.as_str(), q.clone().try_into()?, None)