Update graph_builder.py

graph_builder.py

@@ -712,6 +712,7 @@ def build_document_extraction_graph():
     workflow.add_edge("store_in_knowledge_base", END)
     
     # Handle errors - any node can output an error
+    # FIX: Properly define the conditional edges to handle errors
     for node in workflow.nodes:
         # Check if status is error, if yes, go to END
         workflow.add_conditional_edges(
@@ -719,7 +720,7 @@ def build_document_extraction_graph():
             lambda state: "status" in state and state["status"] == "error",
             {
                 True: END,
-                False: None
+                False: workflow.get_next_node(node) if node != "store_in_knowledge_base" else None
             }
         )