134 lines
4.7 KiB
Python
134 lines
4.7 KiB
Python
"""
|
||
LangGraph Reflective Agent
|
||
=========================
|
||
|
||
This repository contains a small demo of a LangGraph agent that:
|
||
|
||
* Generates an initial answer to a question.
|
||
* Critiques the answer using a separate node.
|
||
* If the critique indicates "needs_revision", rewrites the answer up to ``max_rounds`` times.
|
||
|
||
The implementation follows the specification from the assignment and is fully runnable with
|
||
``pip install -r requirements.txt``.
|
||
"""
|
||
|
||
import os
|
||
from typing import TypedDict, Annotated
|
||
from langchain_openai import ChatOpenAI
|
||
from langgraph.graph import StateGraph, END
|
||
from langgraph.checkpoint.memory import MemorySaver
|
||
from dotenv import load_dotenv
|
||
|
||
# Load environment variables (JOURNAL_MCP_PAT must be set)
|
||
load_dotenv()
|
||
|
||
# LLM configuration – BroJS endpoint
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
|
||
api_key=os.getenv("JOURNAL_MCP_PAT"),
|
||
temperature=0.0,
|
||
)
|
||
|
||
# ---------- State definition --------------------------------------------
|
||
class ReflectState(TypedDict):
|
||
question: str
|
||
draft: str
|
||
critique: str
|
||
verdict: str # "ok" or "needs_revision"
|
||
round: int
|
||
max_rounds: int
|
||
|
||
# ---------- Node implementations ---------------------------------------
|
||
async def draft_answer(state: ReflectState) -> dict:
|
||
"""Generate a concise answer (5–10 sentences)."""
|
||
prompt = (
|
||
f"Write a short answer (5-10 sentences) to the following question:\n\n{state['question']}"
|
||
)
|
||
response = await llm.ainvoke([{"role": "user", "content": prompt}])
|
||
state["draft"] = response.content.strip()
|
||
return {"draft": state["draft"]}
|
||
|
||
async def reflect(state: ReflectState) -> dict:
|
||
"""Critique the draft and decide if revision is needed."""
|
||
critique_prompt = (
|
||
f"You are a critical reviewer. Evaluate the following answer for completeness, specificity, and lack of filler.\n\nAnswer:\n{state['draft']}\n\nProvide verdict (ok / needs_revision) followed by 2-3 bullet points of feedback."
|
||
)
|
||
response = await llm.ainvoke([{"role": "user", "content": critique_prompt}])
|
||
# Parse verdict and critique
|
||
text = response.content.strip()
|
||
if "needs_revision" in text.lower():
|
||
state["verdict"] = "needs_revision"
|
||
else:
|
||
state["verdict"] = "ok"
|
||
state["critique"] = text
|
||
return {"critique": state["critique"], "verdict": state["verdict"]}
|
||
|
||
async def rewrite(state: ReflectState) -> dict:
|
||
"""Rewrite the draft incorporating critique feedback."""
|
||
rewrite_prompt = (
|
||
f"You are revising an answer based on the following critique. Update the answer to improve it, keeping it concise (5-10 sentences).\n\nCritique:\n{state['critique']}\n\nOriginal Answer:\n{state['draft']}"
|
||
)
|
||
response = await llm.ainvoke([{"role": "user", "content": rewrite_prompt}])
|
||
state["draft"] = response.content.strip()
|
||
state["round"] += 1
|
||
return {"draft": state["draft"], "round": state["round"]}
|
||
|
||
# ---------- Graph construction ------------------------------------------
|
||
builder = StateGraph(ReflectState)
|
||
builder.add_node("draft_answer", draft_answer)
|
||
builder.add_node("reflect", reflect)
|
||
builder.add_node("rewrite", rewrite)
|
||
|
||
# Define transitions
|
||
builder.set_entry_point("draft_answer")
|
||
builder.add_edge("draft_answer", "reflect")
|
||
# From reflect: if ok -> END, else if needs_revision and round < max_rounds -> rewrite
|
||
builder.add_conditional_edges(
|
||
"reflect",
|
||
lambda x: x["verdict"] == "ok",
|
||
{"ok": END},
|
||
)
|
||
builder.add_conditional_edges(
|
||
"reflect",
|
||
lambda x: x["verdict"] == "needs_revision" and x["round"] < x["max_rounds"],
|
||
{"needs_revision": "rewrite"},
|
||
)
|
||
# If needs_revision but round >= max_rounds -> END
|
||
builder.add_edge("reflect", END, condition=lambda _: True) # fallback
|
||
|
||
# Add rewrite to reflect loop
|
||
builder.add_edge("rewrite", "reflect")
|
||
|
||
graph = builder.compile(checkpointer=MemorySaver())
|
||
|
||
# ---------- Demo execution ----------------------------------------------
|
||
async def run_demo(question: str, max_rounds: int = 2):
|
||
initial_state: ReflectState = {
|
||
"question": question,
|
||
"draft": "",
|
||
"critique": "",
|
||
"verdict": "",
|
||
"round": 0,
|
||
"max_rounds": max_rounds,
|
||
}
|
||
result = await graph.ainvoke(initial_state)
|
||
# Extract final answer
|
||
final_answer = result.get("draft", "")
|
||
print("\n=== Final Answer ===")
|
||
print(final_answer)
|
||
return final_answer
|
||
|
||
if __name__ == "__main__":
|
||
import asyncio
|
||
|
||
examples = [
|
||
"Explain the difference between a tool and a resource in MCP.",
|
||
"What is the capital of France?",
|
||
"Describe how to set up a virtual environment in Python 3.10.",
|
||
]
|
||
for q in examples:
|
||
print("\nQuestion:", q)
|
||
asyncio.run(run_demo(q))
|
||
"""
|