Put OpenAI API key only temporarily in environment variable

app.py

@@ -8,22 +8,26 @@ import datetime
 gpt_only_prompt = "Calculate the following, giving only the final answer:\n"
 prompt = ""
 
-os.environ["OPENAI_API_KEY"] = ""
+# os.environ["OPENAI_API_KEY"] = ""
 
 
-def set_openai_api_key(api_key, pal_chain_sta):
-    if api_key != "":
-        os.environ["OPENAI_API_KEY"] = api_key
+def set_openai_api_key(api_key, openai_api_key, pal_chain):
+    if api_key:
+        openai_api_key = api_key
 
+        os.environ["OPENAI_API_KEY"] = api_key
         llm = OpenAI(model_name='code-davinci-002', temperature=0, max_tokens=512)
+        os.environ["OPENAI_API_KEY"] = ""
+
         pal_chain = PALChain.from_math_prompt(llm, verbose=True)
-        return pal_chain
+        return openai_api_key, pal_chain
 
 
-def openai_create(prompt):
+def openai_create(prompt, openai_api_key):
     print("prompt: " + prompt)
 
     # We use temperature of 0.0 because it gives the most predictable, factual answer (i.e. avoids hallucination).
+    os.environ["OPENAI_API_KEY"] = openai_api_key
     response = openai.Completion.create(
         model="text-davinci-003",
         prompt=prompt,
@@ -33,12 +37,13 @@ def openai_create(prompt):
         frequency_penalty=0,
         presence_penalty=0
     )
+    os.environ["OPENAI_API_KEY"] = ""
 
     return response.choices[0].text
 
 
-def calc_gpt_only(math_problem):
-    answer = openai_create(gpt_only_prompt + math_problem + "\n")
+def calc_gpt_only(math_problem, openai_api_key):
+    answer = openai_create(gpt_only_prompt + math_problem + "\n", openai_api_key)
 
     print("\n==== date/time: " + str(datetime.datetime.now()) + " ====")
     print("calc_gpt_only math problem: " + math_problem)
@@ -98,13 +103,14 @@ with block:
 
     gr.HTML("<center>Powered by <a href='https://github.com/hwchase17/langchain'>LangChain 🦜️🔗</a></center>")
 
+    openai_api_key_state = gr.State()
     pal_chain_state = gr.State()
 
-    gpt_only.click(calc_gpt_only, inputs=[request], outputs=[answer_html])
+    gpt_only.click(calc_gpt_only, inputs=[request, openai_api_key_state], outputs=[answer_html])
     gpt_pal.click(calc_gpt_pal, inputs=[request, pal_chain_state], outputs=[answer_html])
 
     openai_api_key_textbox.change(set_openai_api_key,
-                                  inputs=[openai_api_key_textbox, pal_chain_state],
-                                  outputs=[pal_chain_state])
+                                  inputs=[openai_api_key_textbox, openai_api_key_state, pal_chain_state],
+                                  outputs=[openai_api_key_state, pal_chain_state])
 
 block.launch()