Update app.py

app.py

@@ -24,29 +24,44 @@ 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.'
     
-    headers = {
-        'Content-Type': 'application/json',
-        'Authorization': f'Bearer {openai_key}'
-    }
 
     data = {
       "model": model,
       "messages": [
-        {
-          "role": "system",
-          "content": "You are a helpful assistant."
-        },
         {
           "role": "user",
           "content": newprompt
         }
-      ]
+      ],
+      "temperature": 0.7,
+    }
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {openai_key}",
     }
 
+    
     response = requests.post(f"{base_url}/v1/chat/completions", headers=headers, data=json.dumps(data))
-    print(f"xxxxxx: {response.text}")
-    response_json = response.json()
-    code = response_json['choices'][0]['message']['content'].strip()
+    print(f"text: {response.text}")
+
+    def extract_code(text):
+        # Match triple backtick blocks first
+        triple_match = re.search(r'```(?:\w+\n)?(.+?)```', text, re.DOTALL)
+        if triple_match:
+            return triple_match.group(1).strip()
+        else:
+            # If no triple backtick blocks, match single backtick blocks
+            single_match = re.search(r'`(.+?)`', text, re.DOTALL)
+            if single_match:
+                return single_match.group(1).strip()
+        # If no code blocks found, return original text
+        return text
+
+    if response.status_code != 200:
+        return "Error: " + response.text, 500
+    
+    code = extract_code(response.json()["choices"][0]["message"]["content"])
     img = execute_code(code)
     return img