diff --git a/solution.py b/solution.py new file mode 100644 index 0000000..5155f3a --- /dev/null +++ b/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() \ No newline at end of file