fix: main.py — Повторный экзамен: Граф с рефлексией и доработкой
This commit is contained in:
@@ -1,17 +1,26 @@
|
||||
#!/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 langchain_openai import ChatOpenAI
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.graph import StateGraph, START, END
|
||||
from langgraph.graph.message import add_messages
|
||||
from langchain.tools import tool
|
||||
from deepagents import create_deep_agent
|
||||
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
|
||||
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, required by the course)
|
||||
# ----------------------------------------------------------------------
|
||||
# LLM configuration - OpenRouter
|
||||
llm = ChatOpenAI(
|
||||
model="openai/gpt-oss-20b:free",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
@@ -19,9 +28,7 @@ llm = ChatOpenAI(
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Backend for deepagents (required by the framework)
|
||||
# ----------------------------------------------------------------------
|
||||
# Backend for deepagents
|
||||
backend = CompositeBackend(
|
||||
[
|
||||
LocalShellBackend(workspace_dir="./workspace"),
|
||||
@@ -29,134 +36,70 @@ backend = CompositeBackend(
|
||||
]
|
||||
)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# State definition for the reflection loop
|
||||
# ----------------------------------------------------------------------
|
||||
# ---------- LangGraph components ----------
|
||||
|
||||
class ReflectState(TypedDict):
|
||||
question: str
|
||||
draft: str
|
||||
critique: str
|
||||
verdict: str # "ok" or "needs_revision"
|
||||
verdict: str # ok | needs_revision
|
||||
round: int
|
||||
max_rounds: int
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Helper function to build a simple LLM chain for a given prompt
|
||||
# ----------------------------------------------------------------------
|
||||
def llm_call(prompt: str, state: ReflectState) -> str:
|
||||
"""Invoke the LLM with a system prompt and the current state."""
|
||||
messages = [
|
||||
HumanMessage(content=prompt.format(**state))
|
||||
]
|
||||
response = llm.invoke(messages)
|
||||
return response.content.strip()
|
||||
class CritiqueOutput(BaseModel):
|
||||
verdict: str = Field(description="ok or needs_revision")
|
||||
critique: str = Field(description="2-3 bullet points of critique")
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Node: draft_answer - produce the first answer
|
||||
# ----------------------------------------------------------------------
|
||||
def draft_answer(state: ReflectState) -> ReflectState:
|
||||
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()
|
||||
return state
|
||||
|
||||
def reflect(state: dict) -> dict:
|
||||
prompt = (
|
||||
"You are an expert educator. Answer the following question in 5-10 sentences, "
|
||||
"clear and concise, without unnecessary filler. Question: {question}"
|
||||
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)."
|
||||
)
|
||||
draft = llm_call(prompt, state)
|
||||
state["draft"] = draft
|
||||
state["round"] = 0
|
||||
result = llm.invoke([HumanMessage(content=prompt)])
|
||||
critique = critique_parser.parse(result.content)
|
||||
state["verdict"] = critique.verdict
|
||||
state["critique"] = critique.critique
|
||||
return state
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Node: reflect - LLM critic evaluates the draft
|
||||
# ----------------------------------------------------------------------
|
||||
def reflect(state: ReflectState) -> ReflectState:
|
||||
critique_prompt = (
|
||||
"You are a reviewer. Evaluate the draft answer provided below. "
|
||||
"Assess completeness, concreteness and absence of filler. "
|
||||
"Return a verdict ('ok' or 'needs_revision') and list 2-3 short remarks. "
|
||||
"Format exactly as:\n"
|
||||
"Verdict: <verdict>\n"
|
||||
"Critique:\n"
|
||||
"- <remark 1>\n"
|
||||
"- <remark 2>\n"
|
||||
"Draft:\n{draft}"
|
||||
def rewrite(state: dict) -> dict:
|
||||
prompt = (
|
||||
f"Rewrite the draft answer to address the following critique:\n\n{state['critique']}\n\n"
|
||||
"Keep the answer brief (5-10 sentences)."
|
||||
)
|
||||
critique_raw = llm_call(critique_prompt, state)
|
||||
# Parse the structured response
|
||||
lines = critique_raw.splitlines()
|
||||
verdict_line = next((l for l in lines if l.lower().startswith("verdict:")), "")
|
||||
verdict = verdict_line.split(":", 1)[1].strip().lower()
|
||||
critique_start = lines.index("Critique:") + 1 if "Critique:" in lines else 0
|
||||
critique_items = [l.lstrip("- ").strip() for l in lines[critique_start:] if l.startswith("-")]
|
||||
state["verdict"] = verdict
|
||||
state["critique"] = "\n".join(critique_items)
|
||||
return state
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Node: rewrite - improve the draft based on critique
|
||||
# ----------------------------------------------------------------------
|
||||
def rewrite(state: ReflectState) -> ReflectState:
|
||||
rewrite_prompt = (
|
||||
"You are a writer. Improve the previous draft according to the following critique points. "
|
||||
"Make the answer clearer, more concrete and remove any filler. Keep the length 5-10 sentences.\n"
|
||||
"Critique:\n{critique}\n\nCurrent draft:\n{draft}"
|
||||
)
|
||||
new_draft = llm_call(rewrite_prompt, state)
|
||||
state["draft"] = new_draft
|
||||
result = llm.invoke([HumanMessage(content=prompt)])
|
||||
state["draft"] = result.content.strip()
|
||||
state["round"] += 1
|
||||
return state
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# DESIGN DECISION: Use a pure LangGraph state machine for reflection.
|
||||
# NECESSITY: The assignment explicitly requires a separate reflect node and
|
||||
# iteration via rewrite → reflect, not a try/except retry loop.
|
||||
# OPTIMALITY: Graph representation makes the flow declarative, guarantees
|
||||
# max_rounds enforcement, and isolates each responsibility.
|
||||
# ALTERNATIVES CONSIDERED: A manual while-loop with try/except was removed
|
||||
# because it mixes error handling with logical revision, violating
|
||||
# the task specification.
|
||||
# ----------------------------------------------------------------------
|
||||
def build_graph() -> StateGraph:
|
||||
graph = StateGraph(ReflectState)
|
||||
|
||||
graph.add_node("draft_answer", draft_answer)
|
||||
graph.add_node("reflect", reflect)
|
||||
graph.add_node("rewrite", rewrite)
|
||||
|
||||
# START → draft_answer
|
||||
graph.add_edge(START, "draft_answer")
|
||||
# draft_answer → reflect
|
||||
graph.add_edge("draft_answer", "reflect")
|
||||
|
||||
# reflect → END if ok
|
||||
graph.add_conditional_edges(
|
||||
"reflect",
|
||||
lambda s: END if s["verdict"] == "ok" else "rewrite",
|
||||
)
|
||||
# rewrite → reflect (if rounds left)
|
||||
def rewrite_condition(s: ReflectState):
|
||||
if s["round"] < s["max_rounds"]:
|
||||
return "reflect"
|
||||
def decide_next(state: dict) -> str:
|
||||
if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"]:
|
||||
return "rewrite"
|
||||
return END
|
||||
|
||||
graph.add_edge("rewrite", rewrite_condition)
|
||||
graph.add_conditional_edges("reflect", decide_next, {"rewrite": "rewrite", "END": END})
|
||||
graph.add_edge("rewrite", "reflect")
|
||||
|
||||
graph.set_entry_point("draft_answer")
|
||||
return graph
|
||||
return graph.compile()
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# DeepAgent wrapper - required by the course
|
||||
# ----------------------------------------------------------------------
|
||||
agent = create_deep_agent(
|
||||
model=llm,
|
||||
tools=[], # No external tools needed for this assignment
|
||||
backend=backend,
|
||||
system_prompt="You are a reflective assistant that writes concise answers and improves them based on critique.",
|
||||
)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Demo execution
|
||||
# ----------------------------------------------------------------------
|
||||
async def main():
|
||||
question = "Объясни студенту разницу между tool и resource в MCP"
|
||||
def run_graph(question: str) -> str:
|
||||
graph = build_graph()
|
||||
initial_state: ReflectState = {
|
||||
"question": question,
|
||||
"draft": "",
|
||||
@@ -165,18 +108,44 @@ async def main():
|
||||
"round": 0,
|
||||
"max_rounds": 2,
|
||||
}
|
||||
final_state = graph.run(initial_state)
|
||||
return final_state["draft"]
|
||||
|
||||
graph = build_graph()
|
||||
# Run the graph synchronously (LangGraph supports async, but our nodes are sync)
|
||||
final_state = await graph.ainvoke(initial_state, config={"configurable": {"thread_id": "demo-1"}})
|
||||
# ---------- DeepAgents tool ----------
|
||||
|
||||
print("=== Final Answer ===")
|
||||
print(final_state["draft"])
|
||||
print("\n=== Verdict ===")
|
||||
print(final_state["verdict"])
|
||||
if final_state["critique"]:
|
||||
print("\n=== Critique ===")
|
||||
print(final_state["critique"])
|
||||
@tool
|
||||
def run_graph_tool(question: str) -> str:
|
||||
"""Run the LangGraph to produce a refined answer."""
|
||||
return run_graph(question)
|
||||
|
||||
# ---------- Agent ----------
|
||||
|
||||
agent = create_deep_agent(
|
||||
model=llm,
|
||||
tools=[run_graph_tool],
|
||||
backend=backend,
|
||||
system_prompt="You are a helpful agent that answers questions by running the run_graph tool.",
|
||||
)
|
||||
|
||||
# ---------- CLI ----------
|
||||
|
||||
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()
|
||||
|
||||
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)
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user