not focusing more on api state

app.py

@@ -1,20 +1,19 @@
 import gradio as gr
 from src.main import ChatWrapper
 
-agent = ChatWrapper('openai')  # default agent_state
+agent = ChatWrapper('openai', '')  # default agnet_state
 
 def update_agent(api_key: str, selection: str):
     global agent
-    agent = ChatWrapper(chain_type=selection)
+    agent = ChatWrapper(chain_type=selection, api_key=api_key)
     return agent # This is agent state
 
-def chat(message, api_key):
+def chat(message):
     global agent
-    agent(message, api_key)  # Get a response to the current message
+    agent(message)  # Get a response to the current message
     history = agent.history  # Access the entire chat history
     return history, history # Return the history twice to update both the chatbot and the state
 
-
 block = gr.Blocks(css=".gradio-container {background-color: lightgray}")
 
 with block:
@@ -53,8 +52,8 @@ with block:
     state = gr.State()
     agent_state = gr.State()
 
-    submit.click(chat, inputs=[message, api_key_textbox], outputs=[chatbot, state])
-    message.submit(chat, inputs=[message, api_key_textbox], outputs=[chatbot, state])
+    submit.click(chat, inputs=[message], outputs=[chatbot, state])
+    message.submit(chat, inputs=[message], outputs=[chatbot, state])
 
     api_key_textbox.change(update_agent, inputs=[api_key_textbox, selection], outputs=[agent_state])
     selection.change(update_agent, inputs=[api_key_textbox, selection], outputs=[agent_state])

src/main.py

@@ -21,23 +21,29 @@ def load_chain_falcon(api_key: str):
     return chain
 
 class ChatWrapper:
-    def __init__(self, chain_type: str):
+    def __init__(self, chain_type: str, api_key: str = ''):
+        self.api_key = api_key
         self.chain_type = chain_type
         self.history = []
         self.lock = Lock()
-        self.chain = None
 
-    def __call__(self, inp: str, api_key: str = ''):
+        if self.api_key:
+            if chain_type == 'openai':
+                self.chain = load_chain_openai(self.api_key)
+            elif chain_type == 'falcon':
+                self.chain = load_chain_falcon(self.api_key)
+            else:
+                raise ValueError(f'Invalid chain_type: {chain_type}')
+        else:
+            self.chain = None
+    
+    def clear_api_key(self):
+        if hasattr(self, 'api_key'):
+            del self.api_key
+
+    def __call__(self, inp: str):
         self.lock.acquire()
         try:
-            if api_key:
-                if self.chain_type == 'openai':
-                    self.chain = load_chain_openai(api_key)
-                elif self.chain_type == 'falcon':
-                    self.chain = load_chain_falcon(api_key)
-                else:
-                    raise ValueError(f'Invalid chain_type: {self.chain_type}')
-
             if self.chain is None:
                 self.history.append((inp, "Please add your API key to proceed."))
                 return self.history
@@ -47,6 +53,7 @@ class ChatWrapper:
         except Exception as e:
             self.history.append((inp, f"An error occurred: {e}"))
         finally:
+            self.clear_api_key() # API key is cleared after running each chain in the class
             self.lock.release()
 
         return self.history