fix memory bug and clean up code

app.py

@@ -1,9 +1,4 @@
 import os
-# os.environ["OPENAI_API_TYPE"]=os.getenv("OPENAI_API_TYPE")
-# os.environ["OPENAI_API_VERSION"]=os.getenv("OPENAI_API_VERSION")
-# os.environ["OPENAI_API_BASE"]=os.getenv("OPENAI_API_BASE")
-# os.environ["OPENAI_API_KEY"]=os.getenv("OPENAI_API_KEY")
-
 import gradio as gr
 from langchain.chat_models import ChatOpenAI
 from langchain import LLMMathChain
@@ -11,19 +6,11 @@ from langchain.utilities import GoogleSerperAPIWrapper, WikipediaAPIWrapper
 from langchain.agents import initialize_agent, Tool, AgentType
 
 
-def set_key(api_key):
-    if not api_key:
-        return "Key can't be empty!"
-    
-    os.environ["OPENAI_API_KEY"] = api_key
-    return "Key received"
-
-
 def initialize():
     # Define models
     llm = ChatOpenAI(temperature=0) 
 
-    # search = GoogleSerperAPIWrapper(serper_api_key='845bee8ad6861cd07d480cfae78df20352481282')
+    search = GoogleSerperAPIWrapper()
     wiki = WikipediaAPIWrapper(top_k_results = 1)
     llm_math_chain = LLMMathChain(llm=llm)
     tools = [
@@ -37,11 +24,11 @@ def initialize():
             func=wiki.run,
             description="useful for when you need to answer questions about historical entity. the input to this should be a single search term."
         ),
-        # Tool(
-        #     name = "Current Search",
-        #     func=search.run,
-        #     description="useful for when you need to answer questions about current events or the current state of the world, also useful if there is no wikipedia result. the input to this should be a single search term."
-        # ),
+        Tool(
+            name = "Current Search",
+            func=search.run,
+            description="useful for when you need to answer questions about current events or the current state of the world, also useful if there is no wikipedia result. the input to this should be a single search term."
+        )
     ]
 
     from langchain.memory import ConversationBufferMemory
@@ -57,11 +44,17 @@ def initialize():
     return chatbot_engine
 
 
-def chat(chat_history, chatbot_engine, message=""):
-    if not chatbot_engine:
-        print("no engine")
-        chatbot_engine = initialize()
+def set_key(api_key):
+    if not api_key:
+        return "Key can't be empty!", None
+    
+    os.environ["OPENAI_API_KEY"] = api_key
+
+    chatbot_engine = initialize()
+    return "Key received", chatbot_engine
 
+
+def chat(chat_history, chatbot_engine, message=""):
     # Empty msg
     if not message.strip():
         return chat_history, chat_history, ""
@@ -94,15 +87,15 @@ with gr.Blocks() as demo:
         api_key_set.click(
             fn=set_key,
             inputs=[openai_api_key_textbox],
-            outputs=[api_key_set],
+            outputs=[api_key_set, chatbot_engine],
         )
 
-
-    gr.Markdown("""<h1><center>Chat with your personal assistant!</center></h1>""")
+    gr.Markdown("""<h1><center>Chat with your online-connected bot!</center></h1>""")
     chatbot = gr.Chatbot()
     message = gr.Textbox()
     submit = gr.Button("SEND")
     submit.click(chat, inputs=[chat_history, chatbot_engine, message], outputs=[chatbot, chat_history, message])
 
+
 if __name__ == "__main__":
     demo.launch(debug = True)