Solution ready for publish: update main.py

main.py

@@ -1,10 +1,10 @@
 """LangGraph code review agent.
 
-This implementation follows the assignment specification:
+Implementation follows assignment:
 - State: CodeReviewState with 4 criteria.
 - Nodes: draft_review, reflect, rewrite.
 - Graph: START -> draft_review -> reflect -> (ok -> END) or (needs_revision & round<max_rounds -> rewrite -> reflect).
-- Uses LangGraph and LangChain OpenAI (or Ollama) for LLM calls.
+- Uses LangGraph and LangChain OpenAI for LLM calls.
 - Structured output for reflect via Pydantic model.
 - Demo function sort_numbers.
 """
@@ -30,11 +30,8 @@ class CodeReviewState(TypedDict):
     max_rounds: int
 
 # ---------- LLM ----------
-# Use OpenAI if key present, else Ollama fallback
+# Use OpenAI only
-if os.getenv("OPENAI_API_KEY"):
 llm = ChatOpenAI(temperature=0)
-else:
-    llm = ChatOpenAI(model="ollama/llama3", temperature=0)
 
 # ---------- Nodes ----------
 
@@ -73,11 +70,7 @@ def reflect(state: CodeReviewState) -> CodeReviewState:
     Return a JSON object with keys: scores (dict), weakest_criterion (string), verdict ('ok' if all scores >=7 else 'needs_revision').
     """
     response = llm.invoke([HumanMessage(content=prompt)])
-    try:
     data = ReflectOutput.model_validate_json(response.content)
-    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
     state["weakest_criterion"] = data.weakest_criterion
     state["verdict"] = data.verdict