From ecf8aa322d8caec55a2d17935a0f20d15a856ab3 Mon Sep 17 00:00:00 2001
From: Dongri Jin <dongrium@gmail.com>
Date: Tue, 25 Jun 2024 11:44:38 +0900
Subject: [PATCH] Add timeout option

---
 src/v1/api.rs | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/src/v1/api.rs b/src/v1/api.rs
index b4c36ef..7b520bc 100644
--- a/src/v1/api.rs
+++ b/src/v1/api.rs
@@ -48,6 +48,7 @@ pub struct Client {
     pub api_key: String,
     pub organization: Option<String>,
     pub proxy: Option<String>,
+    pub timeout: Option<u64>,
 }
 
 impl Client {
@@ -62,6 +63,7 @@ impl Client {
             api_key,
             organization: None,
             proxy: None,
+            timeout: None,
         }
     }
 
@@ -72,6 +74,7 @@ impl Client {
             api_key,
             organization: organization.into(),
             proxy: None,
+            timeout: None,
         }
     }
 
@@ -82,6 +85,18 @@ impl Client {
             api_key,
             organization: None,
             proxy: Some(proxy),
+            timeout: None,
+        }
+    }
+
+    pub fn new_with_timeout(api_key: String, timeout: u64) -> Self {
+        let endpoint = std::env::var("OPENAI_API_BASE").unwrap_or_else(|_| API_URL_V1.to_owned());
+        Self {
+            api_endpoint: endpoint,
+            api_key,
+            organization: None,
+            proxy: None,
+            timeout: Some(timeout),
         }
     }
 
@@ -98,6 +113,9 @@ impl Client {
         if let Some(proxy) = &self.proxy {
             request = request.with_proxy(minreq::Proxy::new(proxy).unwrap());
         }
+        if let Some(timeout) = self.timeout {
+            request = request.with_timeout(timeout);
+        }
         request
     }