fix(needs_fixes): 1 исправлений, 0 отстояно — main.py

main.py

@@ -1,15 +1,12 @@
 import os
-import json
+from typing import TypedDict, Annotated
-import asyncio
-from typing import TypedDict
-
-from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage
-from deepagents import create_deep_agent
-from deepagents.backends import FilesystemBackend
 from langgraph.graph import StateGraph, START, END
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.output_parsers import PydanticOutputParser
+from pydantic import BaseModel, Field
 
-# ---------- LLM ----------
+# LLM configuration (OpenRouter)
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -17,101 +14,111 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# ---------- Backend & Agent ----------
-backend = FilesystemBackend()
-agent = create_deep_agent(
-    model=llm,
-    tools=[],
-    backend=backend,
-    system_prompt="You are a helpful assistant.",
-)
-
 # ---------- State ----------
 class ReflectState(TypedDict):
     question: str
     draft: str
     critique: str
-    verdict: str  # ok | needs_revision
+    verdict: str  # "ok" | "needs_revision"
     round: int
     max_rounds: int
 
+# ---------- Structured output models ----------
+class CritiqueModel(BaseModel):
+    verdict: str = Field(..., description="ok or needs_revision")
+    remarks: list[str] = Field(..., description="2-3 bullet points of critique")
+
+class RewriteModel(BaseModel):
+    draft: str = Field(..., description="Rewritten draft answer")
+
+critique_parser = PydanticOutputParser(pydantic_object=CritiqueModel)
+rewrite_parser = PydanticOutputParser(pydantic_object=RewriteModel)
+
 # ---------- Nodes ----------
 async def draft_answer(state: ReflectState) -> ReflectState:
-    prompt = f"Write a brief answer (5-10 sentences) to the following question: {state['question']}"
+    prompt = (
-    result = await agent.ainvoke(
+        f"Write a concise answer (5–10 sentences) to the following question:\n\n"
-        {"messages": [HumanMessage(content=prompt)], "configurable": {"thread_id": "draft"}},
+        f"Question: {state['question']}"
-        {},
     )
-    state["draft"] = result["messages"][-1].content
+    response = await llm.ainvoke([HumanMessage(content=prompt)])
+    state['draft'] = response.content.strip()
     return state
 
 async def reflect(state: ReflectState) -> ReflectState:
     prompt = (
-        f"Critique the following answer. Provide verdict ok or needs_revision and 2-3 points of critique in JSON format with keys verdict and critique.\nAnswer: {state['draft']}"
+        f"You are a critic evaluating the draft answer for completeness, concreteness, and absence of filler.\n"
+        f"Draft: {state['draft']}\n"
+        f"Provide a verdict ("ok" or "needs_revision") and 2–3 bullet points of critique.\n"
+        f"Return the result in JSON format: {{\"verdict\": "...", \"remarks\": ["...", ...]}}"
     )
-    result = await agent.ainvoke(
+    response = await llm.ainvoke([HumanMessage(content=prompt)])
-        {"messages": [HumanMessage(content=prompt)], "configurable": {"thread_id": "reflect"}},
+    parsed = critique_parser.parse(response.content)
-        {},
+    state['critique'] = "\n".join(parsed.remarks)
-    )
+    state['verdict'] = parsed.verdict
-    content = result["messages"][-1].content
-    try:
-        data = json.loads(content)
-        state["verdict"] = data.get("verdict", "").lower()
-        state["critique"] = data.get("critique", "")
-    except json.JSONDecodeError:
-        state["verdict"] = "needs_revision"
-        state["critique"] = content
     return state
 
 async def rewrite(state: ReflectState) -> ReflectState:
     prompt = (
-        f"Rewrite the answer to improve it based on the critique: {state['critique']}\nPrevious draft: {state['draft']}\nProvide the improved answer."
+        f"Rewrite the draft answer to address the following critique:\n"
+        f"Critique: {state['critique']}\n"
+        f"Original draft: {state['draft']}\n"
+        f"Provide the rewritten draft only.\n"
+        f"Return JSON: {{\"draft\": "..."}}"
     )
-    result = await agent.ainvoke(
+    response = await llm.ainvoke([HumanMessage(content=prompt)])
-        {"messages": [HumanMessage(content=prompt)], "configurable": {"thread_id": "rewrite"}},
+    parsed = rewrite_parser.parse(response.content)
-        {},
+    state['draft'] = parsed.draft
-    )
+    state['round'] += 1
-    state["draft"] = result["messages"][-1].content
-    state["round"] += 1
     return state
 
-# ---------- Conditional Edge ----------
-def reflect_cond(state: ReflectState):
-    if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"]:
-        return "rewrite"
-    return "end"
-
 # ---------- Graph ----------
+builder = StateGraph(ReflectState)
+builder.add_node("draft_answer", draft_answer)
+builder.add_node("reflect", reflect)
+builder.add_node("rewrite", rewrite)
 
-graph = StateGraph(ReflectState)
+builder.add_edge(START, "draft_answer")
-graph.add_node("draft_answer", draft_answer)
+builder.add_edge("draft_answer", "reflect")
-graph.add_node("reflect", reflect)
-graph.add_node("rewrite", rewrite)
 
-graph.add_edge(START, "draft_answer")
+# Conditional edges after reflect
-graph.add_edge("draft_answer", "reflect")
+builder.add_conditional_edges(
-graph.add_conditional_edges("reflect", reflect_cond, {"rewrite": "rewrite", "end": END})
+    "reflect",
-graph.add_edge("rewrite", "reflect")
+    lambda x: x['verdict'] == "ok",
+    {"ok": END, "needs_revision": "rewrite"},
+)
 
-graph.set_entry_point("draft_answer")
+# After rewrite go back to reflect
-graph.set_finish_point(END)
+builder.add_edge("rewrite", "reflect")
 
-executor = graph.compile()
+# If max rounds exceeded, end
+builder.add_conditional_edges(
+    "reflect",
+    lambda x: x['round'] >= x['max_rounds'] and x['verdict'] == "needs_revision",
+    {True: END, False: END},  # both lead to END, but loop already handled
+)
 
-# ---------- CLI ----------
+graph = builder.compile()
+
+# ---------- Demo ----------
 async def main():
-    question = input("Enter a question: ")
     initial_state: ReflectState = {
-        "question": question,
+        "question": "Объясни студенту разницу между tool и resource в MCP",
         "draft": "",
         "critique": "",
         "verdict": "",
-        "round": 1,
+        "round": 0,
         "max_rounds": 2,
     }
-    final_state = await executor(initial_state)
+    final_state = await graph.ainvoke(initial_state)
-    print("\nFinal answer:\n")
+    print("\n--- Final Draft ---")
-    print(final_state["draft"])
+    print(final_state['draft'])
+    print("\n--- Critique ---")
+    print(final_state['critique'])
+    print("\n--- Verdict ---")
+    print(final_state['verdict'])
+    print("\n--- Rounds Used ---")
+    print(final_state['round'])
 
 if __name__ == "__main__":
+    import asyncio
     asyncio.run(main())