Initial implementation of LangGraph reflection agent: add main.py

main.py

@@ -0,0 +1,92 @@
+"""LangGraph reflection agent example.
+
+The agent writes a short answer, critiques it, and rewrites if needed.
+"""
+
+from typing import TypedDict, Dict
+import sys
+
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import create_react_agent
+from langchain_openai import ChatOpenAI
+
+# Define state
+class ReflectState(TypedDict):
+    question: str
+    draft: str
+    critique: str
+    verdict: str  # "ok" | "needs_revision"
+    round: int
+    max_rounds: int
+
+# LLM
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+# Draft node
+async def draft_answer(state: ReflectState) -> Dict:
+    prompt = f"Write a short answer (5–10 sentences) to the following question: {state['question']}"
+    response = await llm.ainvoke(prompt)
+    state['draft'] = response.content
+    return state
+
+# Critique node
+async def reflect(state: ReflectState) -> Dict:
+    prompt = (
+        f"You are a critic. Evaluate the following draft answer for completeness, specificity, and lack of filler.\n"
+        f"Draft: {state['draft']}\n"
+        f"Give a verdict: 'ok' or 'needs_revision'.\n"
+        f"If needs_revision, provide 2–3 points of critique."
+    )
+    response = await llm.ainvoke(prompt)
+    # Simple parsing: first line verdict, rest critique
+    lines = response.content.strip().splitlines()
+    verdict = lines[0].strip().lower()
+    critique = "\n".join(lines[1:]).strip()
+    state['verdict'] = verdict
+    state['critique'] = critique
+    return state
+
+# Rewrite node
+async def rewrite(state: ReflectState) -> Dict:
+    prompt = (
+        f"Rewrite the draft answer taking into account the following critique: {state['critique']}\n"
+        f"Original draft: {state['draft']}"
+    )
+    response = await llm.ainvoke(prompt)
+    state['draft'] = response.content
+    state['round'] += 1
+    return state
+
+# Build graph
+builder = StateGraph(ReflectState)
+builder.add_node("draft_answer", draft_answer)
+builder.add_node("reflect", reflect)
+builder.add_node("rewrite", rewrite)
+
+builder.set_entry_point("draft_answer")
+builder.add_edge("draft_answer", "reflect")
+builder.add_conditional_edges(
+    "reflect",
+    lambda state: "END" if state["verdict"] == "ok" else "rewrite",
+)
+builder.add_edge("rewrite", "reflect")
+
+graph = builder.compile()
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python main.py <question>")
+        sys.exit(1)
+    question = sys.argv[1]
+    initial_state: ReflectState = {
+        "question": question,
+        "draft": "",
+        "critique": "",
+        "verdict": "",
+        "round": 0,
+        "max_rounds": 2,
+    }
+    result = graph.invoke(initial_state)
+    print("\nFinal answer:\n", result["draft"])
+    print("\nCritique:\n", result["critique"])
+    print("\nVerdict:\n", result["verdict"])