diff --git a/main.py b/main.py
index 84bc222..0284284 100644
--- a/main.py
+++ b/main.py
@@ -75,14 +75,14 @@ Provide the output in the following JSON-like format:
 
 async def rewrite(state: CodeReviewState) -> CodeReviewState:
     # Simple rewrite: add a sentence addressing the weakest criterion
-    additional = f"Additionally, the review should pay more attention to {state['weakest_criterion']}."
+    additional = f"Additionally, the review should pay more attention to {state['weakest_criterion']}.",
     state['draft_review'] = state['draft_review'] + "\n" + additional
     state['round'] += 1
     return state
 
 # ---------- Graph ----------
 def build_graph() -> StateGraph[CodeReviewState]:
-    graph = StateGraph(CodeGraphState)
+    graph = StateGraph(CodeReviewState)
     graph.add_node("draft_review", draft_review)
     graph.add_node("reflect", reflect)
     graph.add_node("rewrite", rewrite)