From eead409a06dba8082b186a09735de26b306abf3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=B8=D0=BB=20=D0=92=D0=B8=D0=BA?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BE=D0=B2?= Date: Thu, 2 Jul 2026 08:41:03 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20main.py=20=E2=80=94=20=D0=9F=D0=BE=D0=B2?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BD=D1=8B=D0=B9=20=D1=8D=D0=BA=D0=B7=D0=B0?= =?UTF-8?q?=D0=BC=D0=B5=D0=BD:=20=D0=93=D1=80=D0=B0=D1=84=20=D1=81=20?= =?UTF-8?q?=D1=80=D0=B5=D1=84=D0=BB=D0=B5=D0=BA=D1=81=D0=B8=D0=B5=D0=B9=20?= =?UTF-8?q?=D0=B8=20=D0=B4=D0=BE=D1=80=D0=B0=D0=B1=D0=BE=D1=82=D0=BA=D0=BE?= =?UTF-8?q?=D0=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 157 +++++++++++++++++++++++++++----------------------------- 1 file changed, 76 insertions(+), 81 deletions(-) diff --git a/main.py b/main.py index a35f1bd..25843b9 100644 --- a/main.py +++ b/main.py @@ -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()) \ No newline at end of file