118 lines
3.0 KiB
Python
118 lines
3.0 KiB
Python
"""
|
||
LangGraph Reflection Agent
|
||
|
||
Usage:
|
||
python main.py "Your question"
|
||
"""
|
||
import sys
|
||
from typing import TypedDict, Dict
|
||
|
||
# State definition
|
||
class ReflectState(TypedDict):
|
||
question: str
|
||
draft: str
|
||
critique: str
|
||
verdict: str # 'ok' or 'needs_revision'
|
||
round: int
|
||
max_rounds: int
|
||
|
||
# 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
|
||
|
||
async def reflect(state: Dict) -> Dict:
|
||
from langchain_openai import ChatOpenAI
|
||
llm = ChatOpenAI(model="gpt-3.5-turbo")
|
||
prompt = (
|
||
f"You are a critic evaluating the draft answer for completeness, specificity, and lack of filler.\n"
|
||
f"Draft: {state['draft']}\n"
|
||
"Provide verdict (ok or needs_revision) and 2–3 bullet points of 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
|
||
|
||
async def rewrite(state: Dict) -> Dict:
|
||
from langchain_openai import ChatOpenAI
|
||
llm = ChatOpenAI(model="gpt-3.5-turbo")
|
||
prompt = (
|
||
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.invoke(prompt)
|
||
state["draft"] = response.content
|
||
state["round"] += 1
|
||
return state
|
||
|
||
# Build graph
|
||
from langgraph.graph import StateGraph
|
||
|
||
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",
|
||
should_end,
|
||
{
|
||
"rewrite": "rewrite",
|
||
"END": END,
|
||
},
|
||
)
|
||
# rewrite -> reflect
|
||
|
||
graph_builder.add_edge("rewrite", "reflect")
|
||
|
||
graph = graph_builder.compile()
|
||
|
||
if __name__ == "__main__":
|
||
if len(sys.argv) < 2:
|
||
print("Usage: python main.py \"Your 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("\n--- Final Answer ---")
|
||
print(result["draft"])
|
||
print("\n--- Critique ---")
|
||
print(result["critique"])
|