From 906e758e2c7e4f863afa74daf9c5b1a87ed7526b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A0=D0=B8=D0=BD=D0=B0=D1=80=20=D0=9C=D0=B8=D1=80=D0=B7?= =?UTF-8?q?=D0=B0=D0=B3=D0=B8=D1=82=D0=BE=D0=B2?= Date: Thu, 11 Jun 2026 16:11:40 +0000 Subject: [PATCH] Updated LangGraph reflection agent with max_rounds logic.: update main.py --- main.py | 132 ++++++++++++++++++++++++++++++++------------------------ 1 file changed, 76 insertions(+), 56 deletions(-) diff --git a/main.py b/main.py index 672df5f..3720b6e 100644 --- a/main.py +++ b/main.py @@ -1,86 +1,105 @@ -"""LangGraph reflection agent example. - -The agent writes a short answer, critiques it, and rewrites if needed. """ +LangGraph Reflection Agent -from typing import TypedDict, Dict +Usage: + python main.py "Your question" +""" import sys +from typing import TypedDict, Dict -from langgraph.graph import StateGraph, END -from langgraph.prebuilt import create_react_agent -from langchain_openai import ChatOpenAI - -# Define state +# State definition class ReflectState(TypedDict): question: str draft: str critique: str - verdict: str # "ok" | "needs_revision" + verdict: str # 'ok' or '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 +# Node functions +async def draft_answer(state: Dict) -> Dict: + from langchain_openai import ChatOpenAI + llm = ChatOpenAI(model="gpt-3.5-turbo") + prompt = f"Write a concise answer (5–10 sentences) to the following question:\n\n{state['question']}" + response = await llm.invoke(prompt) + state["draft"] = response.content return state -# Critique node -async def reflect(state: ReflectState) -> Dict: +async def reflect(state: Dict) -> Dict: + from langchain_openai import ChatOpenAI + llm = ChatOpenAI(model="gpt-3.5-turbo") prompt = ( - f"You are a critic. Evaluate the following draft answer for completeness, specificity, and lack of filler.\n" + f"You are a critic evaluating the 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." + "Provide verdict (ok or needs_revision) and 2–3 bullet 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 + response = await llm.invoke(prompt) + # Simple parsing + text = response.content.strip() + if "needs_revision" in text.lower(): + state["verdict"] = "needs_revision" + else: + state["verdict"] = "ok" + state["critique"] = text return state -# Rewrite node -async def rewrite(state: ReflectState) -> Dict: +async def rewrite(state: Dict) -> Dict: + from langchain_openai import ChatOpenAI + llm = ChatOpenAI(model="gpt-3.5-turbo") prompt = ( - f"Rewrite the draft answer taking into account the following critique: {state['critique']}\n" - f"Original draft: {state['draft']}" + f"Rewrite the draft answer incorporating the following critique:\n" + f"Critique: {state['critique']}\n" + "Provide a revised concise answer (5–10 sentences)." ) - response = await llm.ainvoke(prompt) - state['draft'] = response.content - state['round'] += 1 + response = await llm.invoke(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) +from langgraph.graph import StateGraph -builder.set_entry_point("draft_answer") -builder.add_edge("draft_answer", "reflect") -builder.add_conditional_edges( +graph_builder = StateGraph(ReflectState) + +graph_builder.add_node("draft_answer", draft_answer) + +graph_builder.add_node("reflect", reflect) + +graph_builder.add_node("rewrite", rewrite) + +# Connections +start_edge = "draft_answer" +end_edge = None # will be set in condition + +def should_end(state: Dict) -> str: + if state["verdict"] == "ok": + return "END" + if state["round"] >= state.get("max_rounds", 2): + return "END" + return "rewrite" + +# Add edges with condition +from langgraph.graph import END + +graph_builder.set_entry_point(start_edge) + +graph_builder.add_conditional_edges( "reflect", - lambda state: ( - "END" - if state["verdict"] == "ok" - or state.get("round", 0) >= state.get("max_rounds", 2) - else "rewrite" - ), + should_end, + { + "rewrite": "rewrite", + "END": END, + }, ) -builder.add_edge("rewrite", "reflect") +# rewrite -> reflect -graph = builder.compile() +graph_builder.add_edge("rewrite", "reflect") + +graph = graph_builder.compile() if __name__ == "__main__": if len(sys.argv) < 2: - print("Usage: python main.py ") + print("Usage: python main.py \"Your question\"") sys.exit(1) question = sys.argv[1] initial_state: ReflectState = { @@ -92,6 +111,7 @@ if __name__ == "__main__": "max_rounds": 2, } result = graph.invoke(initial_state) - print("\nFinal answer:\n", result["draft"]) - print("\nCritique:\n", result["critique"]) - print("\nVerdict:\n", result["verdict"]) + print("\n--- Final Answer ---") + print(result["draft"]) + print("\n--- Critique ---") + print(result["critique"])