Add main.py

main.py

@@ -1,96 +1,146 @@
 import os
-import asyncio
+import json
+from typing import TypedDict, Dict
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import create_agent
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.messages import HumanMessage, AIMessage
-from langgraph.graph import StateGraph, START, END
-from typing import TypedDict, Annotated
-from langgraph.graph.message import add_messages
-from pydantic import BaseModel, Field
-from langchain_core.output_parsers import PydanticOutputParser
 
-# LLM setup
+# ---------------------
-llm = ChatOpenAI(
+# 1. State definition
-    model="gpt-4o-mini",
+# ---------------------
-    base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENAI_API_KEY"),
-    temperature=0.0,
-)
-
-# State definition
 class CodeReviewState(TypedDict):
     code: str
     draft_review: str
-    criteria_scores: dict[str, int]
+    criteria_scores: Dict[str, int]  # {"pep8": 0-10, "type_hints": 0-10, "edge_cases": 0-10, "naming": 0-10}
     weakest_criterion: str
-    verdict: str
+    verdict: str  # "ok" | "needs_revision"
     round: int
     max_rounds: int
 
-# Node: draft_review
+# ---------------------
-async def draft_review(state: CodeReviewState):
+# 2. LLM setup
-    prompt = f"Write a concise code review (3-6 points) for the following Python function:\n\n{state['code']}"
+# ---------------------
-    msg = await llm.ainvoke([HumanMessage(content=prompt)])
+# Use OpenAI or Ollama based on env variable
-    state['draft_review'] = msg.content
+if os.getenv("USE_OLLAMA", "false").lower() == "true":
-    return state
+    from langchain_ollama import ChatOllama
+    llm = ChatOllama(model="llama3", temperature=0.2)
+else:
+    llm = ChatOpenAI(temperature=0.2, model_name="gpt-4o-mini")
 
-# Node: reflect
+# ---------------------
-class ReflectOutput(BaseModel):
+# 3. Node definitions
-    scores: dict[str, int]
+# ---------------------
-    weakest: str
-    verdict: str
 
-parser = PydanticOutputParser(pydantic_object=ReflectOutput)
+def draft_review_fn(state: CodeReviewState) -> Dict:
+    code = state["code"]
+    prompt = f"""
+You are a senior Python developer. Provide a concise code review (3-6 bullet points) for the following function. Focus on style, correctness, and potential improvements.
 
-async def reflect(state: CodeReviewState):
+Function:
-    prompt = f"Evaluate the draft review and assign scores 0-10 for PEP8, type_hints, edge_cases, naming. Return JSON with keys scores, weakest, verdict (ok or needs_revision).\n\nDraft review:\n{state['draft_review']}"
+{code}
-    msg = await llm.ainvoke([HumanMessage(content=prompt)])
-    out = parser.parse(msg.content)
-    state['criteria_scores'] = out.scores
-    state['weakest_criterion'] = out.weakest
-    state['verdict'] = out.verdict
-    return state
 
-# Node: rewrite
+Review:
-async def rewrite(state: CodeReviewState):
+"""
-    prompt = f"Rewrite the part of the draft review that addresses the weakest criterion '{state['weakest_criterion']}'. Keep other points unchanged.\n\nOriginal draft:\n{state['draft_review']}"
+    response = llm.invoke([HumanMessage(content=prompt)])
-    msg = await llm.ainvoke([HumanMessage(content=prompt)])
+    review = response.content.strip()
-    state['draft_review'] = msg.content
+    return {"draft_review": review}
-    state['round'] += 1
-    return state
 
-# Graph
-graph = StateGraph(CodeReviewState)
-graph.add_node("draft_review", draft_review)
-graph.add_node("reflect", reflect)
-graph.add_node("rewrite", rewrite)
 
-graph.set_entry_point("draft_review")
+def reflect_fn(state: CodeReviewState) -> Dict:
-graph.add_edge("draft_review", "reflect")
+    review = state["draft_review"]
-graph.add_conditional_edges(
+    prompt = f"""
+You are an automated code review critic. Evaluate the following code review on four criteria: PEP8 compliance, type hints usage, edge case handling, and naming conventions. Assign each a score from 0 to 10. Also determine the weakest criterion and a verdict: "ok" if all scores are 7 or higher, otherwise "needs_revision".
+
+Review:
+{review}
+
+Respond in JSON with keys: "pep8", "type_hints", "edge_cases", "naming", "weakest_criterion", "verdict".
+"""
+    response = llm.invoke([HumanMessage(content=prompt)])
+    try:
+        data = json.loads(response.content)
+    except Exception:
+        # Fallback: simple parsing
+        data = {
+            "pep8": 5,
+            "type_hints": 5,
+            "edge_cases": 5,
+            "naming": 5,
+            "weakest_criterion": "pep8",
+            "verdict": "needs_revision"
+        }
+    return {
+        "criteria_scores": {
+            "pep8": int(data.get("pep8", 0)),
+            "type_hints": int(data.get("type_hints", 0)),
+            "edge_cases": int(data.get("edge_cases", 0)),
+            "naming": int(data.get("naming", 0))
+        },
+        "weakest_criterion": data.get("weakest_criterion", "pep8"),
+        "verdict": data.get("verdict", "needs_revision")
+    }
+
+
+def rewrite_fn(state: CodeReviewState) -> Dict:
+    weakest = state["weakest_criterion"]
+    review = state["draft_review"]
+    prompt = f"""
+You are a senior Python developer. The following code review has been identified as weak in the "{weakest}" criterion. Rewrite only the part of the review that addresses this criterion, improving it significantly. Keep the rest of the review unchanged.
+
+Original Review:
+{review}
+
+Rewritten Review:
+"""
+    response = llm.invoke([HumanMessage(content=prompt)])
+    new_review = response.content.strip()
+    return {"draft_review": new_review, "round": state["round"] + 1}
+
+# ---------------------
+# 4. Graph construction
+# ---------------------
+builder = StateGraph(CodeReviewState)
+
+builder.add_node("draft_review", draft_review_fn)
+builder.add_node("reflect", reflect_fn)
+builder.add_node("rewrite", rewrite_fn)
+
+# Entry point
+builder.set_entry_point("draft_review")
+
+# Transitions
+builder.add_edge("draft_review", "reflect")
+builder.add_conditional_edges(
     "reflect",
-    lambda s: "rewrite" if s['verdict']=='needs_revision' and s['round']<s['max_rounds'] else END,
+    lambda x: x["verdict"],
+    {
+        "ok": END,
+        "needs_revision": "rewrite"
+    }
 )
-graph.add_edge("rewrite", "reflect")
+builder.add_edge("rewrite", "reflect")
 
-app = graph.compile()
+# Final graph
+graph = builder.compile()
 
-# Demo function
+# ---------------------
-async def demo():
+# 5. Demo CLI
-    code = """def sort_numbers(arr):
+# ---------------------
-    return sorted(arr)"""
+if __name__ == "__main__":
-    state: CodeReviewState = {
+    sample_code = """
-        "code": code,
+    def sort_numbers(arr):
+        return sorted(arr)
+    """
+    initial_state: CodeReviewState = {
+        "code": sample_code.strip(),
         "draft_review": "",
         "criteria_scores": {},
         "weakest_criterion": "",
         "verdict": "",
         "round": 0,
-        "max_rounds": 2,
+        "max_rounds": 2
     }
-    final = await app.ainvoke(state)
+    result = graph.invoke(initial_state)
-    print("Final draft review:\n", final['draft_review'])
+    print("\n--- Final State ---")
-    print("Scores:", final['criteria_scores'])
+    print(json.dumps(result, indent=2))
-    print("Verdict:", final['verdict'])
-
-if __name__ == "__main__":
-    asyncio.run(demo())