Update main.py

main.py

@@ -147,53 +147,56 @@ async def chat(input_data: ChatInput):
     }
     
     input_message = HumanMessage(content=input_data.message)
+    try:
+        async def generate():
+            async for event in agent.astream_events(
+                {"messages": [input_message]}, 
+                config,
+                version="v2"
+            ):
+                kind = event["event"]
+                
+                if kind == "on_chat_model_stream":
+                    content = event["data"]["chunk"].content
+                    if content:
+                        yield f"{json.dumps({'type': 'token', 'content': content})}\n"
     
-    async def generate():
-        async for event in agent.astream_events(
-            {"messages": [input_message]}, 
-            config,
-            version="v2"
-        ):
-            kind = event["event"]
-            
-            if kind == "on_chat_model_stream":
-                content = event["data"]["chunk"].content
-                if content:
-                    yield f"{json.dumps({'type': 'token', 'content': content})}\n"
-
-            if kind == "on_tool_start":
-                print(f"Debug - Tool being called: {event['name']}")  # Add this debug line
-                tool_input = event['data'].get('input', '')
-                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
-                #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(),
-        media_type="text/event-stream"
-    )
+                if kind == "on_tool_start":
+                    print(f"Debug - Tool being called: {event['name']}")  # Add this debug line
+                    tool_input = event['data'].get('input', '')
+                    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
+                    #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(),
+            media_type="text/event-stream"
+        )
+    except Exception as e:
+        print(f"Error during event streaming: {str(e)}")
+        yield f"{json.dumps({'type': 'error', 'content': str(e)})}\n"
 
 @app.get("/health")
 async def health_check():