146 lines
6.5 KiB
Python
146 lines
6.5 KiB
Python
"""LangGraph-агент с рефлексией и доработкой (без try/except)."""
|
|
import os
|
|
import asyncio
|
|
from typing import TypedDict
|
|
|
|
from dotenv import load_dotenv
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.graph import StateGraph, START, END
|
|
|
|
load_dotenv()
|
|
|
|
# ── LLM ──────────────────────────────────────────────────────────────────
|
|
llm = ChatOpenAI(
|
|
model="openai/gpt-oss-20b:free",
|
|
base_url="https://openrouter.ai/api/v1",
|
|
api_key=os.getenv("OPENAI_API_KEY"),
|
|
temperature=0.0,
|
|
)
|
|
|
|
# ── State ─────────────────────────────────────────────────────────────────
|
|
class ReflectState(TypedDict):
|
|
question: str
|
|
draft: str
|
|
critique: str
|
|
verdict: str # ok | needs_revision
|
|
round: int
|
|
max_rounds: int
|
|
|
|
# ── Nodes ─────────────────────────────────────────────────────────────────
|
|
|
|
async def draft_answer(state: ReflectState) -> dict:
|
|
"""Первичный ответ на вопрос (5–10 предложений)."""
|
|
prompt = (
|
|
"Напиши краткий, конкретный ответ (5–10 предложений) на вопрос.\n\n"
|
|
f"Вопрос: {state['question']}"
|
|
)
|
|
response = await llm.ainvoke([{"role": "user", "content": prompt}])
|
|
return {"draft": response.content.strip(), "round": 0}
|
|
|
|
|
|
async def reflect(state: ReflectState) -> dict:
|
|
"""LLM-критик: оценивает черновик и выносит структурированный вердикт.
|
|
|
|
Промпт требует строгого формата ответа — вердикт парсится напрямую,
|
|
без try/except и без запасных значений при ошибке парсинга.
|
|
"""
|
|
prompt = (
|
|
"Ты — строгий критик учебных ответов. Оцени черновик по трём критериям:\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: 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 = (
|
|
"Улучши ответ, учитывая замечания критика. Сохрани объём 5–10 предложений.\n\n"
|
|
f"Замечания:\n{state['critique']}\n\n"
|
|
f"Текущий черновик:\n{state['draft']}"
|
|
)
|
|
response = await llm.ainvoke([{"role": "user", "content": prompt}])
|
|
return {"draft": response.content.strip(), "round": state["round"] + 1}
|
|
|
|
# ── 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.add_edge(START, "draft_answer")
|
|
builder.add_edge("draft_answer", "reflect")
|
|
builder.add_conditional_edges(
|
|
"reflect",
|
|
route_after_reflect,
|
|
{"end": END, "rewrite": "rewrite"},
|
|
)
|
|
builder.add_edge("rewrite", "reflect") # после rewrite всегда → reflect
|
|
# лимит раундов проверяет route_after_reflect
|
|
|
|
graph = builder.compile()
|
|
|
|
# ── Demo ──────────────────────────────────────────────────────────────────
|
|
|
|
async def main() -> None:
|
|
question = "Объясни студенту разницу между tool и resource в MCP."
|
|
print(f"Вопрос: {question}")
|
|
print("=" * 60)
|
|
|
|
init: ReflectState = {
|
|
"question": question,
|
|
"draft": "",
|
|
"critique": "",
|
|
"verdict": "",
|
|
"round": 0,
|
|
"max_rounds": 2,
|
|
}
|
|
|
|
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())
|