From 66a73c7b84a9cffc67f1dc0b6b836b4d9e5d0c42 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Thu, 16 Mar 2023 17:11:25 +0800
Subject: [PATCH] feat: add `config.connect_timeout` (#76)

---
 src/client.rs     | 4 ++--
 src/config/mod.rs | 9 +++++++++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index b438440..14d65d0 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -10,7 +10,6 @@ use std::time::Duration;
 use tokio::runtime::Runtime;
 use tokio::time::sleep;
 
-const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
 const API_URL: &str = "https://api.openai.com/v1/chat/completions";
 
 #[derive(Debug)]
@@ -128,8 +127,9 @@ impl ChatGptClient {
             builder = builder
                 .proxy(Proxy::all(proxy).with_context(|| format!("Invalid proxy `{proxy}`"))?);
         }
+        let timeout = self.config.read().get_connect_timeout();
         let client = builder
-            .connect_timeout(CONNECT_TIMEOUT)
+            .connect_timeout(timeout)
             .build()
             .with_context(|| "Failed to build http client")?;
         Ok(client)
diff --git a/src/config/mod.rs b/src/config/mod.rs
index f543ee9..1607711 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -12,6 +12,7 @@ use anyhow::{anyhow, bail, Context, Result};
 use inquire::{Confirm, Text};
 use parking_lot::RwLock;
 use serde::Deserialize;
+use std::time::Duration;
 use std::{
     env,
     fs::{create_dir_all, read_to_string, File, OpenOptions},
@@ -65,6 +66,8 @@ pub struct Config {
     pub conversation_first: bool,
     /// Is ligth theme
     pub light_theme: bool,
+    /// Set a timeout in seconds for connect to gpt
+    pub connect_timeout: usize,
     /// Predefined roles
     #[serde(skip)]
     pub roles: Vec<Role>,
@@ -90,6 +93,7 @@ impl Default for Config {
             dry_run: false,
             conversation_first: false,
             light_theme: false,
+            connect_timeout: 10,
             roles: vec![],
             role: None,
             conversation: None,
@@ -267,6 +271,10 @@ impl Config {
         }
     }
 
+    pub fn get_connect_timeout(&self) -> Duration {
+        Duration::from_secs(self.connect_timeout as u64)
+    }
+
     pub fn get_model(&self) -> (String, usize) {
         self.model.clone()
     }
@@ -328,6 +336,7 @@ impl Config {
             ("proxy", proxy),
             ("conversation_first", self.conversation_first.to_string()),
             ("light_theme", self.light_theme.to_string()),
+            ("connect_timeout", self.connect_timeout.to_string()),
             ("dry_run", self.dry_run.to_string()),
         ];
         let mut output = String::new();