diff --git a/main.py b/main.py new file mode 100644 index 0000000..40d2f2e --- /dev/null +++ b/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 ") + 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"])