From 1288fa770ede26c0446f80593cca32889e4cceba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=9A=D1=83=D1=82?= =?UTF-8?q?=D0=BB=D0=B0=D1=85=D0=BC=D0=B5=D1=82=D0=BE=D0=B2?= Date: Thu, 4 Jun 2026 15:57:56 +0000 Subject: [PATCH] add main.py --- main.py | 133 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 133 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..455622b --- /dev/null +++ b/main.py @@ -0,0 +1,133 @@ +""" +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)) +"""