diff --git a/main.py b/main.py index 1113bed..133a813 100644 --- a/main.py +++ b/main.py @@ -76,7 +76,7 @@ def reflect(state: CodeReviewState) -> CodeReviewState: response = llm.invoke([HumanMessage(content=prompt)]) try: data = ReflectOutput.model_validate_json(response.content) - except Exception as e: + except Exception: # Fallback simple parsing data = ReflectOutput.model_validate_json("{\"scores\":{\"pep8\":5,\"type_hints\":5,\"edge_cases\":5,\"naming\":5},\"weakest_criterion\":\"pep8\",\"verdict\":\"needs_revision\"}") state["criteria_scores"] = data.scores @@ -111,13 +111,12 @@ builder.add_conditional_edges( lambda state: state["verdict"] == "ok", {"ok": END, "needs_revision": "rewrite"}, ) +# After rewrite, decide to reflect again or end if max rounds reached builder.add_conditional_edges( "rewrite", - lambda state: state["round"] < state["max_rounds"], - {"rewrite": "reflect", "maxed": END}, + lambda state: "reflect" if state["round"] < state["max_rounds"] else END, + {"reflect": "reflect", "END": END} ) -# If maxed, go to END -builder.add_edge("rewrite", "maxed") graph = builder.compile()