add main.py

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))
+"""