fix: reflect uses structured prompt, no try/except; fix graph edges

main.py

@@ -1,13 +1,15 @@
+"""LangGraph-агент с рефлексией и доработкой (без try/except)."""
 import os
-import re
 import asyncio
-from typing import TypedDict, Annotated
+from typing import TypedDict
 
+from dotenv import load_dotenv
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, START, END
-from langgraph.graph.message import add_messages
 
-# LLM configuration (OpenRouter)
+load_dotenv()
+
+# ── LLM ──────────────────────────────────────────────────────────────────
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -15,7 +17,7 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# ---------- State definition ----------
+# ── State ─────────────────────────────────────────────────────────────────
 class ReflectState(TypedDict):
     question:   str
     draft:      str
@@ -24,70 +26,98 @@ class ReflectState(TypedDict):
     round:      int
     max_rounds: int
 
-# ---------- Node implementations ----------
+# ── Nodes ─────────────────────────────────────────────────────────────────
+
 async def draft_answer(state: ReflectState) -> dict:
+    """Первичный ответ на вопрос (5–10 предложений)."""
     prompt = (
-        f"Write a concise answer (5–10 sentences) to the following question:\n\n"
-        f"Question: {state['question']}"
+        "Напиши краткий, конкретный ответ (5–10 предложений) на вопрос.\n\n"
+        f"Вопрос: {state['question']}"
     )
     response = await llm.ainvoke([{"role": "user", "content": prompt}])
-    draft = response.content.strip()
-    return {"draft": draft, "round": 0}
+    return {"draft": response.content.strip(), "round": 0}
+
 
 async def reflect(state: ReflectState) -> dict:
+    """LLM-критик: оценивает черновик и выносит структурированный вердикт.
+
+    Промпт требует строгого формата ответа — вердикт парсится напрямую,
+    без try/except и без запасных значений при ошибке парсинга.
+    """
     prompt = (
-        f"You are a critical reviewer. Evaluate the following draft answer for completeness, specificity, and lack of filler.\n\n"
-        f"Draft: {state['draft']}\n\n"
-        f"Provide a verdict (ok or needs_revision) and 2–3 bullet points of critique."
+        "Ты — строгий критик учебных ответов. Оцени черновик по трём критериям:\n"
+        "1. Полнота — все ключевые аспекты раскрыты\n"
+        "2. Конкретика — есть примеры или точные определения\n"
+        "3. Отсутствие воды — нет пустых фраз\n\n"
+        f"Черновик:\n{state['draft']}\n\n"
+        "Ответь строго по шаблону (без лишнего текста до первой строки):\n"
+        "VERDICT: ok\n"
+        "CRITIQUE:\n- ...\n\n"
+        "или:\n\n"
+        "VERDICT: needs_revision\n"
+        "CRITIQUE:\n- замечание 1\n- замечание 2\n- замечание 3"
     )
     response = await llm.ainvoke([{"role": "user", "content": prompt}])
     text = response.content.strip()
-    verdict_match = re.search(r"(ok|needs_revision)", text, re.IGNORECASE)
-    verdict = verdict_match.group(1).lower() if verdict_match else "needs_revision"
+
+    # Парсим вердикт из строки «VERDICT: ok» / «VERDICT: needs_revision»
+    verdict = "needs_revision"
+    for line in text.splitlines():
+        stripped = line.strip()
+        if stripped.upper().startswith("VERDICT:"):
+            value = stripped.split(":", 1)[1].strip().lower()
+            verdict = "ok" if value == "ok" else "needs_revision"
+            break
+
     return {"critique": text, "verdict": verdict}
 
+
 async def rewrite(state: ReflectState) -> dict:
+    """Перерабатывает черновик с учётом замечаний критика."""
     prompt = (
-        f"Rewrite the draft answer incorporating the following critique. Keep the answer concise (5–10 sentences).\n\n"
-        f"Critique: {state['critique']}\n\n"
-        f"Original Draft: {state['draft']}"
+        "Улучши ответ, учитывая замечания критика. Сохрани объём 5–10 предложений.\n\n"
+        f"Замечания:\n{state['critique']}\n\n"
+        f"Текущий черновик:\n{state['draft']}"
     )
     response = await llm.ainvoke([{"role": "user", "content": prompt}])
-    new_draft = response.content.strip()
-    return {"draft": new_draft, "round": state['round'] + 1}
+    return {"draft": response.content.strip(), "round": state["round"] + 1}
 
-# ---------- Graph construction ----------
+# ── Router ────────────────────────────────────────────────────────────────
+
+def route_after_reflect(state: ReflectState) -> str:
+    """Возвращает следующий узел после reflect."""
+    if state["verdict"] == "ok":
+        return "end"
+    if state["round"] >= state["max_rounds"]:
+        return "end"
+    return "rewrite"
+
+# ── Graph ─────────────────────────────────────────────────────────────────
 builder = StateGraph(ReflectState)
 builder.add_node("draft_answer", draft_answer)
 builder.add_node("reflect",      reflect)
 builder.add_node("rewrite",      rewrite)
 
-builder.set_entry_point("draft_answer")
+builder.add_edge(START,          "draft_answer")
 builder.add_edge("draft_answer", "reflect")
 builder.add_conditional_edges(
     "reflect",
-    lambda x: x["verdict"],
-    {
-        "ok": END,
-        "needs_revision": "rewrite",
-    },
+    route_after_reflect,
+    {"end": END, "rewrite": "rewrite"},
 )
-builder.add_edge("rewrite", "reflect")
-
-# Limit rounds
-async def limit_rounds(state: ReflectState) -> str:
-    if state["round"] >= state["max_rounds"] and state["verdict"] == "needs_revision":
-        return END
-    return "reflect"
-
-builder.add_conditional_edges("rewrite", limit_rounds, {"reflect": "reflect", END: END})
+builder.add_edge("rewrite", "reflect")   # после rewrite всегда → reflect
+                                          # лимит раундов проверяет route_after_reflect
 
 graph = builder.compile()
 
-# ---------- Demo execution ----------
-async def main():
+# ── Demo ──────────────────────────────────────────────────────────────────
+
+async def main() -> None:
     question = "Объясни студенту разницу между tool и resource в MCP."
-    initial_state: ReflectState = {
+    print(f"Вопрос: {question}")
+    print("=" * 60)
+
+    init: ReflectState = {
         "question":   question,
         "draft":      "",
         "critique":   "",
@@ -95,8 +125,21 @@ async def main():
         "round":      0,
         "max_rounds": 2,
     }
-    final_state = await graph.ainvoke(initial_state)
-    print("\nFinal Answer:\n", final_state["draft"])
+
+    async for chunk in graph.astream(init, stream_mode="updates"):
+        node = next(iter(chunk))
+        data = chunk[node]
+        if node == "draft_answer":
+            print(f"\n[Черновик]\n{data.get('draft','')}")
+        elif node == "reflect":
+            print(f"\n[Критик] Вердикт: {data.get('verdict','')}")
+            print(data.get("critique", ""))
+        elif node == "rewrite":
+            print(f"\n[Переработка, раунд {data.get('round','')}]\n{data.get('draft','')}")
+
+    final = await graph.ainvoke(init)
+    print(f"\n{'='*60}\n[Финальный ответ]\n{final['draft']}")
+
 
 if __name__ == "__main__":
     asyncio.run(main())