169 lines
6.4 KiB
Python
169 lines
6.4 KiB
Python
# solution.py
|
|
# Планирующий агент на основе LangGraph, LangChain и questionary
|
|
# Запуск: python solution.py
|
|
|
|
import os
|
|
from typing import TypedDict, List, Optional, Dict, Any
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.graph import StateGraph, END
|
|
from langgraph.checkpoint.sqlite import SqliteSaver
|
|
from questionary import select, text
|
|
|
|
# ----------------------------------------------------------------------
|
|
# 1. Определяем состояние графа
|
|
# ----------------------------------------------------------------------
|
|
class AgentState(TypedDict):
|
|
# Ввод задачи от пользователя
|
|
task: str
|
|
# Список шагов, которые LLM распланировал
|
|
plan: List[str]
|
|
# Текущий индекс шага
|
|
step_index: int
|
|
# Список результатов выполнения шагов
|
|
step_results: List[str]
|
|
# Итоговый ответ
|
|
final_answer: Optional[str]
|
|
# Выбранный вариант (если есть)
|
|
choice: Optional[str]
|
|
|
|
# ----------------------------------------------------------------------
|
|
# 2. Инициализируем LLM
|
|
# ----------------------------------------------------------------------
|
|
# Ожидается, что переменная окружения OPENAI_API_KEY установлена
|
|
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)
|
|
|
|
# ----------------------------------------------------------------------
|
|
# 3. Определяем функции-узлы
|
|
# ----------------------------------------------------------------------
|
|
def plan(state: AgentState) -> Dict[str, Any]:
|
|
"""
|
|
Первый узел: LLM разбивает задачу на шаги.
|
|
"""
|
|
prompt = f"""
|
|
Задача: {state['task']}
|
|
|
|
Разбей её на последовательные шаги, которые можно выполнить один за другим.
|
|
Выведите список шагов, каждый шаг в отдельной строке, без префиксов.
|
|
"""
|
|
response = llm.invoke(prompt)
|
|
steps = [s.strip() for s in response.content.splitlines() if s.strip()]
|
|
return {"plan": steps, "step_index": 0, "step_results": []}
|
|
|
|
def execute_step(state: AgentState) -> Dict[str, Any]:
|
|
"""
|
|
Выполняем текущий шаг из плана.
|
|
"""
|
|
idx = state["step_index"]
|
|
if idx >= len(state["plan"]):
|
|
# План завершён
|
|
return {"final_answer": "План выполнен."}
|
|
current_step = state["plan"][idx]
|
|
prompt = f"""
|
|
Выполните следующий шаг задачи:
|
|
|
|
Шаг {idx + 1}: {current_step}
|
|
|
|
Предоставьте краткий результат выполнения этого шага.
|
|
"""
|
|
response = llm.invoke(prompt)
|
|
result = response.content.strip()
|
|
new_results = state["step_results"] + [result]
|
|
return {
|
|
"step_results": new_results,
|
|
"step_index": idx + 1
|
|
}
|
|
|
|
def ask_choice(state: AgentState) -> Dict[str, Any]:
|
|
"""
|
|
Если в результате выполнения шага есть варианты выбора,
|
|
предлагаем пользователю выбрать один.
|
|
"""
|
|
last_result = state["step_results"][-1]
|
|
# Простейший парсинг: ищем строки вида "1) ... 2) ..."
|
|
import re
|
|
matches = re.findall(r'\d+\)\s*([^\d]+)', last_result)
|
|
if not matches:
|
|
# Нет вариантов, просто продолжаем
|
|
return {}
|
|
# Показываем варианты пользователю
|
|
choice = select(
|
|
"Выберите вариант:",
|
|
choices=[f"{i+1}) {m.strip()}" for i, m in enumerate(matches)]
|
|
).ask()
|
|
# Сохраняем выбранный вариант
|
|
return {"choice": choice}
|
|
|
|
def finalize(state: AgentState) -> Dict[str, Any]:
|
|
"""
|
|
После завершения всех шагов формируем итоговый ответ.
|
|
"""
|
|
prompt = f"""
|
|
Задача: {state['task']}
|
|
|
|
Выполнены шаги:
|
|
{chr(10).join(f"{i+1}. {r}" for i, r in enumerate(state['step_results']))}
|
|
|
|
На основании результата сформируйте окончательный ответ на задачу.
|
|
"""
|
|
response = llm.invoke(prompt)
|
|
return {"final_answer": response.content.strip()}
|
|
|
|
# ----------------------------------------------------------------------
|
|
# 4. Строим граф
|
|
# ----------------------------------------------------------------------
|
|
builder = StateGraph(AgentState)
|
|
|
|
# Добавляем узлы
|
|
builder.add_node("plan", plan)
|
|
builder.add_node("execute_step", execute_step)
|
|
builder.add_node("ask_choice", ask_choice)
|
|
builder.add_node("finalize", finalize)
|
|
|
|
# Определяем переходы
|
|
builder.set_entry_point("plan")
|
|
builder.add_conditional_edges(
|
|
"plan",
|
|
lambda x: "execute_step",
|
|
{"execute_step": "execute_step"}
|
|
)
|
|
builder.add_conditional_edges(
|
|
"execute_step",
|
|
lambda x: "finalize" if x["step_index"] >= len(x["plan"]) else "ask_choice",
|
|
{"finalize": "finalize", "ask_choice": "ask_choice"}
|
|
)
|
|
builder.add_conditional_edges(
|
|
"ask_choice",
|
|
lambda x: "execute_step",
|
|
{"execute_step": "execute_step"}
|
|
)
|
|
builder.add_edge("finalize", END)
|
|
|
|
# Сохраняем состояние в SQLite (чекпоинтер)
|
|
checkpoint = SqliteSaver.from_conn_str("sqlite:///langgraph_checkpoint.db")
|
|
graph = builder.compile(checkpointer=checkpoint)
|
|
|
|
# ----------------------------------------------------------------------
|
|
# 5. Запуск агента
|
|
# ----------------------------------------------------------------------
|
|
def main() -> None:
|
|
# Получаем задачу от пользователя
|
|
task = text("Введите задачу для агента:").ask()
|
|
# Инициализируем состояние
|
|
init_state: AgentState = {
|
|
"task": task,
|
|
"plan": [],
|
|
"step_index": 0,
|
|
"step_results": [],
|
|
"final_answer": None,
|
|
"choice": None,
|
|
}
|
|
# Запускаем граф
|
|
for event in graph.stream(init_state):
|
|
# event - это обновлённое состояние
|
|
if event["final_answer"]:
|
|
print("\n=== Итоговый ответ ===")
|
|
print(event["final_answer"])
|
|
break
|
|
|
|
if __name__ == "__main__":
|
|
main() |