Add main.py

main.py

@@ -0,0 +1,115 @@
+import asyncio
+import os
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage
+from langchain.tools import tool
+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(
+    model="openai/gpt-4o-mini",
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv("OPENAI_API_KEY"),
+    temperature=0.0,
+)
+
+# Pydantic models for reflection
+class CriteriaScores(BaseModel):
+    pep8: int
+    type_hints: int
+    edge_cases: int
+    naming: int
+
+class ReviewResult(BaseModel):
+    scores: CriteriaScores
+    weakest: str
+    verdict: str
+
+# Tool to run shell commands
+@tool
+def run_command(command: str) -> str:
+    """Execute a shell command and return its output."""
+    try:
+        result = os.popen(command).read()
+        return result.strip() or "(no output)"
+    except Exception as e:
+        return f"Error: {e}"
+
+# Draft review node
+async def draft_review(state: dict):
+    code = state["code"]
+    prompt = f"Write a concise code review for the following Python function. Provide 3-6 bullet points highlighting strengths and areas for improvement.\n\n{code}"
+    response = llm.invoke([HumanMessage(content=prompt)])
+    state["draft_review"] = response.content
+    return state
+
+# Reflect node
+async def reflect(state: dict):
+    review = state["draft_review"]
+    prompt = f"Score the following code review on 4 criteria: PEP8, type hints, edge cases, naming. Return JSON with keys pep8, type_hints, edge_cases, naming (0-10). Also provide the weakest criterion and verdict ('ok' if all >=7 else 'needs_revision').\n\n{review}"
+    response = llm.invoke([HumanMessage(content=prompt)])
+    try:
+        data = ReviewResult.parse_raw(response.content)
+    except Exception:
+        # fallback simple parse
+        data = ReviewResult(scores=CriteriaScores(pep8=5,type_hints=5,edge_cases=5,naming=5),weakest="pep8",verdict="needs_revision")
+    state["criteria_scores"] = data.scores.dict()
+    state["weakest_criterion"] = data.weakest
+    state["verdict"] = data.verdict
+    return state
+
+# Rewrite node
+async def rewrite(state: dict):
+    weakest = state["weakest_criterion"]
+    review = state["draft_review"]
+    prompt = f"Improve the code review focusing on the {weakest} aspect. Keep the rest unchanged.\n\n{review}"
+    response = llm.invoke([HumanMessage(content=prompt)])
+    state["draft_review"] = response.content
+    state["round"] += 1
+    return state
+
+# Graph definition
+class CodeReviewState(TypedDict):
+    code: str
+    draft_review: str
+    criteria_scores: dict
+    weakest_criterion: str
+    verdict: str
+    round: int
+    max_rounds: int
+
+workflow = StateGraph(CodeReviewState)
+workflow.add_node("draft", draft_review)
+workflow.add_node("reflect", reflect)
+workflow.add_node("rewrite", rewrite)
+workflow.add_conditional_edges(START, lambda _: "draft")
+workflow.add_conditional_edges("draft", lambda _: "reflect")
+workflow.add_conditional_edges("reflect", lambda s: "rewrite" if s["verdict"]=="needs_revision" and s["round"]<s["max_rounds"] else "END")
+workflow.add_conditional_edges("rewrite", lambda _: "reflect")
+graph = workflow.compile()
+
+# Demo function
+async def run_demo():
+    code = """
+def sort_numbers(arr):
+    return sorted(arr)
+"""
+    state: CodeReviewState = {
+        "code": code,
+        "draft_review": "",
+        "criteria_scores": {},
+        "weakest_criterion": "",
+        "verdict": "",
+        "round": 0,
+        "max_rounds": 2,
+    }
+    final = await graph.ainvoke(state)
+    print("Final review:\n", final["draft_review"])
+    print("Scores:", final["criteria_scores"])
+
+if __name__ == "__main__":
+    asyncio.run(run_demo())