fix: main.py — Повторный экзамен: Граф с рефлексией и доработкой
This commit is contained in:
@@ -1,15 +1,6 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
# DESIGN DECISION: deepagents is required by the course assignment to build the agent.
|
||||
# NECESSITY: The assignment explicitly requires using create_deep_agent from deepagents; without it the agent cannot be instantiated.
|
||||
# OPTIMALITY: Using deepagents ensures consistent agent behavior and simplifies tool integration; alternative frameworks would violate the course constraints.
|
||||
# ALTERNATIVES CONSIDERED: Using plain langgraph without deepagents would miss the required framework; manually handling tool calls would increase boilerplate.
|
||||
|
||||
import os
|
||||
import asyncio
|
||||
import argparse
|
||||
from typing import TypedDict, Annotated
|
||||
from typing import TypedDict
|
||||
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_core.messages import HumanMessage
|
||||
@@ -17,8 +8,6 @@ from langchain.tools import tool
|
||||
from deepagents import create_deep_agent
|
||||
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
|
||||
from langgraph.graph import StateGraph, START, END
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
# LLM configuration - OpenRouter
|
||||
llm = ChatOpenAI(
|
||||
@@ -36,8 +25,7 @@ backend = CompositeBackend(
|
||||
]
|
||||
)
|
||||
|
||||
# ---------- LangGraph components ----------
|
||||
|
||||
# ---------- LangGraph definition ----------
|
||||
class ReflectState(TypedDict):
|
||||
question: str
|
||||
draft: str
|
||||
@@ -46,61 +34,70 @@ class ReflectState(TypedDict):
|
||||
round: int
|
||||
max_rounds: int
|
||||
|
||||
class CritiqueOutput(BaseModel):
|
||||
verdict: str = Field(description="ok or needs_revision")
|
||||
critique: str = Field(description="2-3 bullet points of critique")
|
||||
|
||||
critique_parser = PydanticOutputParser(pydantic_object=CritiqueOutput)
|
||||
|
||||
def draft_answer(state: dict) -> dict:
|
||||
prompt = f"Write a brief answer (5-10 sentences) to the following question:\n\n{state['question']}"
|
||||
result = llm.invoke([HumanMessage(content=prompt)])
|
||||
state["draft"] = result.content.strip()
|
||||
def draft_answer(state: ReflectState) -> ReflectState:
|
||||
prompt = f"Write a short answer (5-10 sentences) to the following question: {state['question']}"
|
||||
msg = HumanMessage(content=prompt)
|
||||
response = llm.invoke([msg])
|
||||
state["draft"] = response.content
|
||||
state["round"] = 0
|
||||
return state
|
||||
|
||||
def reflect(state: dict) -> dict:
|
||||
def reflect(state: ReflectState) -> ReflectState:
|
||||
prompt = (
|
||||
f"You are a critic. Evaluate the following draft answer for completeness, specificity, and lack of fluff.\n\nDraft:\n{state['draft']}\n\n"
|
||||
"Respond with JSON containing 'verdict' ('ok' or 'needs_revision') and 'critique' (2-3 bullet points)."
|
||||
f"You are a critic. Evaluate the following draft answer:\n{state['draft']}\n\n"
|
||||
"Provide verdict 'ok' or 'needs_revision' and 2-3 points of critique."
|
||||
)
|
||||
result = llm.invoke([HumanMessage(content=prompt)])
|
||||
critique = critique_parser.parse(result.content)
|
||||
state["verdict"] = critique.verdict
|
||||
state["critique"] = critique.critique
|
||||
msg = HumanMessage(content=prompt)
|
||||
response = llm.invoke([msg])
|
||||
text = response.content.strip()
|
||||
verdict = "ok"
|
||||
critique = ""
|
||||
if "needs_revision" in text.lower():
|
||||
verdict = "needs_revision"
|
||||
# Extract critique after the word 'Critique:' if present
|
||||
lower_text = text.lower()
|
||||
if "critique:" in lower_text:
|
||||
idx = lower_text.find("critique:")
|
||||
critique = text[idx + len("critique:") :].strip()
|
||||
else:
|
||||
parts = text.split("\n")
|
||||
if len(parts) > 1:
|
||||
critique = "\n".join(parts[1:]).strip()
|
||||
state["verdict"] = verdict
|
||||
state["critique"] = critique
|
||||
return state
|
||||
|
||||
def rewrite(state: dict) -> dict:
|
||||
def rewrite(state: ReflectState) -> ReflectState:
|
||||
prompt = (
|
||||
f"Rewrite the draft answer to address the following critique:\n\n{state['critique']}\n\n"
|
||||
"Keep the answer brief (5-10 sentences)."
|
||||
f"Rewrite the draft answer to address the following critique:\n{state['critique']}\n\n"
|
||||
"Keep the answer short (5-10 sentences)."
|
||||
)
|
||||
result = llm.invoke([HumanMessage(content=prompt)])
|
||||
state["draft"] = result.content.strip()
|
||||
msg = HumanMessage(content=prompt)
|
||||
response = llm.invoke([msg])
|
||||
state["draft"] = response.content
|
||||
state["round"] += 1
|
||||
return state
|
||||
|
||||
def build_graph() -> StateGraph:
|
||||
graph = StateGraph(ReflectState)
|
||||
graph.add_node("draft_answer", draft_answer)
|
||||
graph.add_node("reflect", reflect)
|
||||
graph.add_node("rewrite", rewrite)
|
||||
graph = StateGraph(ReflectState)
|
||||
graph.add_node("draft_answer", draft_answer)
|
||||
graph.add_node("reflect", reflect)
|
||||
graph.add_node("rewrite", rewrite)
|
||||
graph.set_entry_point("draft_answer")
|
||||
graph.add_edge("draft_answer", "reflect")
|
||||
graph.add_conditional_edges(
|
||||
"reflect",
|
||||
lambda s: (
|
||||
"ok"
|
||||
if s["verdict"] == "ok"
|
||||
else ("rewrite" if s["round"] < s["max_rounds"] else "stop")
|
||||
),
|
||||
{"ok": END, "rewrite": "rewrite", "stop": END},
|
||||
)
|
||||
graph.add_edge("rewrite", "reflect")
|
||||
compiled_graph = graph.compile()
|
||||
|
||||
graph.add_edge(START, "draft_answer")
|
||||
graph.add_edge("draft_answer", "reflect")
|
||||
|
||||
def decide_next(state: dict) -> str:
|
||||
if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"]:
|
||||
return "rewrite"
|
||||
return END
|
||||
|
||||
graph.add_conditional_edges("reflect", decide_next, {"rewrite": "rewrite", "END": END})
|
||||
graph.add_edge("rewrite", "reflect")
|
||||
|
||||
return graph.compile()
|
||||
|
||||
def run_graph(question: str) -> str:
|
||||
graph = build_graph()
|
||||
initial_state: ReflectState = {
|
||||
def run_reflection(question: str) -> str:
|
||||
state: ReflectState = {
|
||||
"question": question,
|
||||
"draft": "",
|
||||
"critique": "",
|
||||
@@ -108,44 +105,42 @@ def run_graph(question: str) -> str:
|
||||
"round": 0,
|
||||
"max_rounds": 2,
|
||||
}
|
||||
final_state = graph.run(initial_state)
|
||||
return final_state["draft"]
|
||||
final_state = compiled_graph.invoke(state)
|
||||
output = (
|
||||
f"Draft:\n{final_state['draft']}\n\n"
|
||||
f"Critique:\n{final_state['critique']}\n\n"
|
||||
f"Verdict: {final_state['verdict']}\n\n"
|
||||
f"Final answer:\n{final_state['draft']}\n"
|
||||
)
|
||||
return output
|
||||
|
||||
# ---------- DeepAgents tool ----------
|
||||
|
||||
@tool
|
||||
def run_graph_tool(question: str) -> str:
|
||||
"""Run the LangGraph to produce a refined answer."""
|
||||
return run_graph(question)
|
||||
|
||||
# ---------- Agent ----------
|
||||
def answer_with_reflection(question: str) -> str:
|
||||
"""Generate a short answer with self-reflection and rewrite if needed."""
|
||||
return run_reflection(question)
|
||||
|
||||
# ---------- DeepAgents agent ----------
|
||||
agent = create_deep_agent(
|
||||
model=llm,
|
||||
tools=[run_graph_tool],
|
||||
tools=[answer_with_reflection],
|
||||
backend=backend,
|
||||
system_prompt="You are a helpful agent that answers questions by running the run_graph tool.",
|
||||
system_prompt=(
|
||||
"You are a helpful agent that writes short answers and self-reflects. "
|
||||
"Use the tool 'answer_with_reflection' to answer questions."
|
||||
),
|
||||
)
|
||||
|
||||
# ---------- CLI ----------
|
||||
|
||||
# ---------- Demo ----------
|
||||
async def main():
|
||||
parser = argparse.ArgumentParser(description="Answer a question with self-reflection.")
|
||||
parser.add_argument("question", nargs="*", help="The question to answer.")
|
||||
args = parser.parse_args()
|
||||
if not args.question:
|
||||
question = input("Enter your question: ").strip()
|
||||
else:
|
||||
question = " ".join(args.question).strip()
|
||||
|
||||
question = (
|
||||
"Explain to a student the difference between tool and resource in MCP."
|
||||
)
|
||||
result = await agent.ainvoke(
|
||||
{"messages": [HumanMessage(content=question)]},
|
||||
{"configurable": {"thread_id": "session-1"}},
|
||||
)
|
||||
# The agent will return the final answer in the last message
|
||||
final_message = result["messages"][-1].content
|
||||
print("\nAnswer:\n")
|
||||
print(final_message)
|
||||
print(result["messages"][-1].content)
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user