From c355235dba77dbddf0cca30489213c62af027fe3 Mon Sep 17 00:00:00 2001
From: yhyu13 <yohan680919@gmail.com>
Date: Wed, 24 May 2023 23:08:10 +0800
Subject: [PATCH] Add env handle

---
 models.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/models.py b/models.py
index 894ed89..05710d6 100644
--- a/models.py
+++ b/models.py
@@ -4,6 +4,17 @@ import backoff
 
 completion_tokens = prompt_tokens = 0
 
+api_key = os.getenv("OPENAI_API_KEY", "")
+if api_key != "":
+    openai.api_key = api_key
+else:
+    print("Warning: OPENAI_API_KEY is not set")
+    
+api_base = os.getenv("OPENAI_API_BASE", "")
+if api_base != "":
+    print("Warning: OPENAI_API_BASE is set to {}".format(api_base))
+    openai.api_base = api_base
+
 @backoff.on_exception(backoff.expo, openai.error.OpenAIError)
 def completions_with_backoff(**kwargs):
     return openai.ChatCompletion.create(**kwargs)