Update main.py

main.py

@@ -42,7 +42,7 @@ class CodeExecutionResult:
 
 API_URL = "https://pvanand-code-execution-files-v5.hf.space"
 
-@tool
+@tool(response_format="content_and_artifact")
 def execute_python(code: str, config: RunnableConfig):
     """Execute Python code in an IPython interactiveshell and return the output.
        The returned artifacts (if present) are automatically rendered in the UI and visible to the user.
@@ -62,17 +62,41 @@ def execute_python(code: str, config: RunnableConfig):
         "session_token": session_token,
         "code": code
     }
-    response = requests.post(
-        f'{API_URL}/v0/execute',
-        headers=headers,
-        data=json.dumps(data)
-    )
-
-    if response.status_code != 200:
-        return f"Error: Request failed with status code {response.status_code}. Response: {response.text}"
-    else:
+    try:
+        response = requests.post(
+            f'{API_URL}/v1/execute',
+            headers=headers,
+            json=data
+        )
+        
+        if response.status_code != 200:
+            return (
+                f"Error: Request failed with status code {response.status_code}. Response: {response.text}",
+                None
+            )
+        
+        # Get the response JSON
         response_json = response.json()
-        return f"data: {json.dumps(response_json)} \ndata:"
+
+        # extract artifacts if they exist 
+        artifacts_data = response_json.get("artifacts_data", {})
+
+        # Create a clean response without artifacts
+        execution_response = {
+            "status": response_json.get("status"),
+            "text": response_json.get("text"),
+            "error_message": response_json.get("error_message"),
+            "artifacts": response_json.get("artifacts")
+        }
+        
+        return (
+            f"Execution completed successfully: {json.dumps(execution_response)}",
+            {"artifacts_data": artifacts_data} if artifacts_data else None
+        )
+        
+    except Exception as e:
+        return (f"Error executing code: {str(e)}", None)
+
 
 memory = MemorySaver()
 model = ChatOpenAI(model="gpt-4o-mini", streaming=True)
@@ -149,28 +173,9 @@ async def chat(input_data: ChatInput):
             
             elif kind == "on_tool_end":
                 tool_output = event['data'].get('output', '').content
-                print(type(tool_output))
-                #print(dir(tool_output))
-                #print the keys
-                pattern = r'data: (.*?)\ndata:'
-                match = re.search(pattern, tool_output)
-                print(tool_output)
-
-                if match:
-                    tool_output_json = match.group(1).strip()
-                    try:
-                        tool_output = json.loads(tool_output_json)
-                        if "artifacts" in tool_output:
-                            for artifact in tool_output["artifacts"]:
-                                artifact_content = requests.get(f"{API_URL}/artifact/{artifact['artifact_id']}").content
-                                #print(artifact_content)
-                                #tool_output["artifacts"][artifact["artifact_id"]] = artifact_content
-                    except Exception as e:
-                        print(e)
-                        print("Error parsing tool output as json: ", tool_output)
-                else:
-                    print("No match found in tool output")  
-                yield f"{json.dumps({'type': 'tool_end', 'tool': event['name'], 'output': tool_output})}\n"
+                artifact_output = event['data'].get('output', '').artifact.get('artifacts_data') if event['data'].get('output', '').artifact else None
+                yield f"{json.dumps({'type': 'tool_end', 'tool': event['name'], 'output': tool_output, 'artifacts_data': artifact_output})}\n"
+                print(f"{json.dumps({'type': 'tool_end', 'tool': event['name'], 'output': tool_output, 'artifacts_data': artifact_output})}\n")
     return EventSourceResponse(
         generate(),
         media_type="text/event-stream"