Update app.py

app.py

@@ -30,14 +30,14 @@ with open("ds-system-prompt.txt", "r") as f:
 
 
 def execute_jupyter_agent(
-    sytem_prompt, user_input, max_new_tokens, model, files, message_history
+    sytem_prompt, user_input, max_new_tokens, model, files, message_history, sbx
 ):
     client = InferenceClient(api_key=HF_TOKEN)
 
     tokenizer = AutoTokenizer.from_pretrained(model)
     # model = "meta-llama/Llama-3.1-8B-Instruct"
 
-    sbx = Sandbox(api_key=E2B_API_KEY)
+    
 
     filenames = []
     if files is not None:
@@ -67,9 +67,11 @@ def execute_jupyter_agent(
         yield notebook_html, message_history
 
 
-def clear(state):
-    state = []
-    return update_notebook_display(create_base_notebook([])[0]), state
+def clear(msg_state, sbx_state):
+    msg_state = []
+    sbx_state.kill()
+    sbx_state = Sandbox(api_key=E2B_API_KEY)
+    return update_notebook_display(create_base_notebook([])[0]), msg_state, sbx_state
 
 
 css = """
@@ -91,7 +93,8 @@ css = """
 
 # Create the interface
 with gr.Blocks() as demo:
-    state = gr.State(value=[])
+    msg_state = gr.State(value=[])
+    sbx_state = gr.State(value=Sandbox(api_key=E2B_API_KEY))
 
     html_output = gr.HTML(value=update_notebook_display(create_base_notebook([])[0]))
 
@@ -134,11 +137,11 @@ with gr.Blocks() as demo:
 
     generate_btn.click(
         fn=execute_jupyter_agent,
-        inputs=[system_input, user_input, max_tokens, model, files, state],
-        outputs=[html_output, state],
+        inputs=[system_input, user_input, max_tokens, model, files, msg_state, sbx_state],
+        outputs=[html_output, msg_state],
     )
 
-    clear_btn.click(fn=clear, inputs=[state], outputs=[html_output, state])
+    clear_btn.click(fn=clear, inputs=[msg_state, sbx_state], outputs=[html_output, msg_state, sbx_state])
 
     demo.load(
         fn=None,