Update app.py

app.py

@@ -2,28 +2,32 @@ import gradio as gr
 import matplotlib.pyplot as plt
 import io
 import numpy as np
-import base64
 from PIL import Image
 import requests
 import json
+import re
 
 # 将图像转换为 base64，以便在 gradio 中显示
-def get_image_data(plt):
+def get_image_data(fig):
     buf = io.BytesIO()
-    plt.savefig(buf, format='PNG')
+    fig.savefig(buf, format='PNG')
     buf.seek(0)
     img = Image.open(buf)
     return img
 
 # 执行 Python 代码并生成图像
 def execute_code(code):
-    exec(code)
-    return get_image_data(plt)
+    namespace = {}
+    exec(code, namespace)
+    fig = namespace.get('fig')  # Assume the code generates a matplotlib figure named 'fig'
+    if fig:
+        return get_image_data(fig)
+    else:
+        raise ValueError("The code did not generate a matplotlib figure named 'fig'")
 
 def gpt_inference(base_url, model, openai_key, prompt):
     
     newprompt = f'Write Python code that does the following: \n\n{prompt}\n\nNote, the code is going to be executed in a Jupyter Python kernel.\n\nLast instruction, and this is the most important, just return code. No other outputs, as your full response will directly be executed in the kernel.'
-    
 
     data = {
       "model": model,
@@ -41,9 +45,7 @@ def gpt_inference(base_url, model, openai_key, prompt):
         "Authorization": f"Bearer {openai_key}",
     }
 
-    
     response = requests.post(f"{base_url}/v1/chat/completions", headers=headers, data=json.dumps(data))
-    print(f"text: {response.text}")
 
     def extract_code(text):
         # Match triple backtick blocks first