129 lines
4.0 KiB
Python
129 lines
4.0 KiB
Python
<|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() |