Update main.py

main.py

@@ -148,7 +148,28 @@ async def chat(input_data: ChatInput):
                 yield f"{json.dumps({'type': 'tool_start', 'tool': event['name'], 'input': tool_input})}\n"
             
             elif kind == "on_tool_end":
-                tool_output = event['data'].get('output', '').content 
+                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"
     return EventSourceResponse(
         generate(),