🐛 fix(chat.py): handle exceptions while building the flow and add a final response to the event stream

The code now handles exceptions that may occur while building the flow. A final response is added to the event stream to indicate the end of the stream.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-13 20:41:58 -03:00
commit 9f5591639e

View file

@ -69,6 +69,8 @@ async def stream_build(flow_id: str):
"""Stream the build process based on stored flow data.""" """Stream the build process based on stored flow data."""
async def event_stream(flow_id): async def event_stream(flow_id):
final_response = json.dumps({"end_of_stream": True})
try:
if flow_id not in flow_data_store: if flow_id not in flow_data_store:
error_message = "Invalid session ID" error_message = "Invalid session ID"
yield f"data: {json.dumps({'error': error_message})}\n\n" yield f"data: {json.dumps({'error': error_message})}\n\n"
@ -102,11 +104,13 @@ async def stream_build(flow_id: str):
"id": node.id, "id": node.id,
} }
) )
yield f"data: {response}\n\n" # SSE format yield f"data: {response}\n\n"
chat_manager.set_cache(flow_id, graph.build()) chat_manager.set_cache(flow_id, graph.build())
final_response = json.dumps({"end_of_stream": True}) except Exception:
yield f"data: {final_response}\n\n" # SSE format logger.error("Error while building the flow")
finally:
yield f"data: {final_response}\n\n"
try: try:
return StreamingResponse(event_stream(flow_id), media_type="text/event-stream") return StreamingResponse(event_stream(flow_id), media_type="text/event-stream")