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

solution.py

@@ -1,81 +1,90 @@
-import os
+import json
-from typing import List, Dict, Any
+from typing import Dict, Any
 
-# Библиотеки, которые должны быть установлены в окружении:
+from langchain_openai import ChatOpenAI
-# - langchain
+from langgraph.graph import StateGraph, END
-# - fastapi (если понадобится)
+from langgraph.prebuilt import create_react_agent
-# - pydantic
-# - rich (для красивого вывода)
 
-try:
-    from langchain.llms import OpenAI
-    from langchain.prompts import PromptTemplate
-    from langchain.chains import LLMChain
-except ImportError as exc:
-    raise RuntimeError(
-        "Не удалось импортировать необходимые библиотеки. "
-        "Убедитесь, что установлены langchain и его зависимости."
-    ) from exc
 
-# Конфигурация модели (можно заменить на другую)
+# ------------------------------------------------------------
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+# 1. Определяем состояние графа
-if not OPENAI_API_KEY:
+# ------------------------------------------------------------
-    raise RuntimeError(
+class PlannerState:
-        "Переменная окружения OPENAI_API_KEY не задана. "
+    """Состояние агента: цель и список шагов."""
-        "Пожалуйста, установите ключ API для OpenAI."
+    goal: str = ""
-    )
+    steps: list[str] | None = None
 
-# Основной класс планирующего агента
+
-class PlanningAgent:
+# ------------------------------------------------------------
+# 2. Функция планирования (использует LLM)
+# ------------------------------------------------------------
+def plan_steps(state: Dict[str, Any]) -> Dict[str, Any]:
     """
-    Простая реализация планировщика задач на основе LLM.
+    Принимает цель и генерирует список шагов в формате JSON.
-    Принимает список целей и генерирует последовательность действий,
+    Возвращает обновлённое состояние с полем `steps`.
-    которые необходимо выполнить для достижения каждой цели.
     """
+    goal = state["goal"]
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)
 
-    def __init__(self, llm_model: str = "gpt-3.5-turbo"):
+    prompt = f"""
-        self.llm = OpenAI(model=llm_model, temperature=0.7)
+    Ты планировщик задач. Тебе дана цель:
-        # Шаблон запроса к LLM
+    {goal}
-        self.prompt_template = PromptTemplate(
-            input_variables=["goal"],
-            template=(
-                "Ты планировщик задач. Для цели: '{goal}' "
-                "предложи пошаговый план действий в виде списка.\n"
-                "Ответ должен быть коротким и конкретным, каждый шаг на одной строке."
-            ),
-        )
-        self.chain = LLMChain(llm=self.llm, prompt=self.prompt_template)
 
-    def plan_for_goal(self, goal: str) -> List[str]:
+    Сгенерируй список конкретных шагов, которые нужно выполнить,
-        """
+    чтобы достичь этой цели. Каждый шаг должен быть короткой фразой.
-        Генерирует план действий для одной цели.
+    Выведите результат в виде JSON-массива строк, например:
-        Возвращает список строк – шагов плана.
+    ["Шаг 1", "Шаг 2", ...]
-        """
+    """
-        raw_output = self.chain.run({"goal": goal})
+    response = llm.invoke(prompt)
-        # Разделяем вывод по строкам и убираем пустые
+    try:
-        steps = [step.strip() for step in raw_output.splitlines() if step.strip()]
+        steps = json.loads(response.content.strip())
-        return steps
+        if not isinstance(steps, list):
+            raise ValueError
+    except Exception:
+        # Если LLM не вернул корректный JSON – создаём простую fallback‑планировку
+        steps = [f"Шаг {i+1} для выполнения цели: {goal}" for i in range(3)]
 
-    def plan_for_goals(self, goals: List[str]) -> Dict[str, List[str]]:
+    return {"steps": steps}
-        """
-        Генерирует планы для списка целей.
-        Возвращает словарь: цель → список шагов.
-        """
-        plans = {}
-        for goal in goals:
-            plans[goal] = self.plan_for_goal(goal)
-        return plans
 
-# Тестовый пример использования
+
+# ------------------------------------------------------------
+# 3. Создаём граф планировщика
+# ------------------------------------------------------------
+builder = StateGraph(PlannerState)
+
+# Входной узел – задаёт цель и сразу переходит к планированию
+def start_node(state: PlannerState) -> Dict[str, Any]:
+    return {"goal": state.goal}
+
+builder.add_node("start", start_node)
+builder.add_node("plan_steps", plan_steps)
+
+# Переходы
+builder.set_entry_point("start")
+builder.add_edge("start", "plan_steps")
+builder.add_edge("plan_steps", END)
+
+planner_graph = builder.compile()
+
+
+# ------------------------------------------------------------
+# 4. Функция запуска планировщика
+# ------------------------------------------------------------
+def run_planner(goal: str) -> list[str]:
+    """
+    Запускает граф планировщика для заданной цели.
+    Возвращает список шагов.
+    """
+    result = planner_graph.invoke({"goal": goal})
+    return result.get("steps", [])
+
+
+# ------------------------------------------------------------
+# 5. Тестовый запуск (можно удалить при интеграции)
+# ------------------------------------------------------------
 if __name__ == "__main__":
-    agent = PlanningAgent()
+    sample_goal = "Составить план подготовки к экзамену по Python"
-    sample_goals = [
+    steps = run_planner(sample_goal)
-        "Запустить новый веб-сайт",
+    print("\nПлан действий:")
-        "Оптимизировать производительность базы данных",
+    for i, step in enumerate(steps, 1):
-        "Подготовить презентацию для инвесторов"
+        print(f"{i}. {step}")
-    ]
-    plans = agent.plan_for_goals(sample_goals)
-    for goal, steps in plans.items():
-        print(f"\nЦель: {goal}")
-        for i, step in enumerate(steps, 1):
-            print(f"  Шаг {i}: {step}")