Повторный экзамен: Граф с рефлексией и доработкой: solution.py

solution.py

@@ -0,0 +1,129 @@
+<|channel|>final code<|message|>import os
+from typing import TypedDict, Literal
+
+from langgraph.graph import StateGraph, END
+from langchain_openai import ChatOpenAI
+
+
+# ---------- 1. Состояние ----------
+class ReflectState(TypedDict):
+    question: str
+    draft: str
+    critique: str
+    verdict: Literal["ok", "needs_revision"]
+    round: int
+    max_rounds: int
+
+
+# ---------- 2. Узлы ----------
+def draft_answer(state: ReflectState) -> dict:
+    """Первичный ответ на вопрос."""
+    llm = ChatOpenAI(model="gpt-4o-mini")
+    prompt = f"Напиши краткий ответ (5–10 предложений) на следующий вопрос:\n\n{state['question']}"
+    response = llm.invoke(prompt)
+    state["draft"] = response.content.strip()
+    return {"draft": state["draft"]}
+
+
+def reflect(state: ReflectState) -> dict:
+    """LLM‑критик оценивает ответ."""
+    llm = ChatOpenAI(model="gpt-4o-mini")
+    prompt = (
+        f"Оцени следующий ответ на вопрос \"{state['question']}\":\n\n"
+        f"{state['draft']}\n\n"
+        "Дай verdict: ok или needs_revision. "
+        "Если needs_revision, перечисли 2–3 конкретных замечания."
+    )
+    response = llm.invoke(prompt)
+    text = response.content.strip()
+    # Разделяем verdict и critique
+    if "needs_revision" in text:
+        state["verdict"] = "needs_revision"
+        parts = text.split("needs_revision", 1)[-1].strip()
+        state["critique"] = parts
+    else:
+        state["verdict"] = "ok"
+        state["critique"] = ""
+    return {"verdict": state["verdict"], "critique": state["critique"]}
+
+
+def rewrite(state: ReflectState) -> dict:
+    """Учитываем critique и переписываем ответ."""
+    llm = ChatOpenAI(model="gpt-4o-mini")
+    prompt = (
+        f"Перепиши ответ, учитывая следующие замечания:\n\n"
+        f"{state['critique']}\n\n"
+        f"Оригинальный ответ:\n{state['draft']}"
+    )
+    response = llm.invoke(prompt)
+    state["draft"] = response.content.strip()
+    state["round"] += 1
+    return {"draft": state["draft"], "round": state["round"]}
+
+
+# ---------- 3. Граф ----------
+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("draft_answer", "reflect")
+builder.add_conditional_edges(
+    "reflect",
+    lambda x: x["verdict"],
+    {
+        "ok": END,
+        "needs_revision": "rewrite",
+    },
+)
+builder.add_edge("rewrite", "reflect")
+
+graph = builder.compile()
+
+
+# ---------- 4. CLI ----------
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser(description="LangGraph рефлексия")
+    parser.add_argument(
+        "--question",
+        type=str,
+        required=True,
+        help="Вопрос, на который нужно ответить.",
+    )
+    args = parser.parse_args()
+
+    # Инициализируем состояние
+    state: ReflectState = {
+        "question": args.question,
+        "draft": "",
+        "critique": "",
+        "verdict": "",
+        "round": 0,
+        "max_rounds": 2,
+    }
+
+    # Запускаем граф
+    result = graph.invoke(state)
+
+    # Выводим финальный ответ и логи
+    print("\n=== Финальный ответ ===")
+    print(result["draft"])
+    print("\n=== Лог рефлексии ===")
+    print(f"Round: {result['round']}")
+    print(f"Verdict: {result['verdict']}")
+    if result["critique"]:
+        print("Critique:")
+        print(result["critique"])
+
+
+if __name__ == "__main__":
+    # Проверяем наличие ключа OpenAI
+    if "OPENAI_API_KEY" not in os.environ:
+        raise RuntimeError(
+            "Пожалуйста, установите переменную окружения OPENAI_API_KEY."
+        )
+    main()