diff --git a/solution.py b/solution.py new file mode 100644 index 0000000..6154801 --- /dev/null +++ b/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() \ No newline at end of file