Экзамен: Планирующий агент: solution.py

solution.py

@@ -0,0 +1,169 @@
+# 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()