diff --git a/main.py b/main.py new file mode 100644 index 0000000..3e74a87 --- /dev/null +++ b/main.py @@ -0,0 +1,114 @@ +""" +LangGraph planning agent demo. + +Run with: + python main.py "Сравни Python и JavaScript" + +Requires OPENAI_API_KEY environment variable. +""" +import os +from typing import TypedDict, List, Optional + +from langchain_openai import ChatOpenAI +from langgraph.graph import StateGraph, END +from langgraph.prebuilt import create_chat_agent + +# ---------- State definition ---------- +class PlanningState(TypedDict): + task: str + plan: Optional[List[str]] + current_step: int + results: List[str] + +# ---------- LLM ---------- +llm = ChatOpenAI(model="gpt-4o-mini", temperature=0) + +# ---------- Planning node ---------- +def planning(state: PlanningState) -> PlanningState: + prompt = ( + "Разбей задачу на 3–6 конкретных шагов.\n" + "Ответ в формате JSON массива строк, например:\n" + "[""\n" + " \"Шаг 1: ...\",\n" + " \"Шаг 2: ...\"\n" + "]\n" + f"Задача: {state['task']}" + ) + response = llm.invoke(prompt) + import json + try: + plan = json.loads(response.content) + if not isinstance(plan, list): + raise ValueError("Not a list") + except Exception as e: + # fallback: split by lines starting with digits + plan = [] + for line in response.content.splitlines(): + line = line.strip() + if line and (line[0].isdigit() or line.startswith("Шаг")): + plan.append(line) + return { + **state, + "plan": plan, + "current_step": 0, + "results": [], + } + +# ---------- Execution node ---------- +def execution(state: PlanningState) -> PlanningState: + step = state["plan"][state["current_step"]] + prompt = f"Выполни шаг:\n{step}\nОтвет в одном абзаце." + response = llm.invoke(prompt) + result = response.content.strip() + new_results = state["results"].copy() + new_results.append(result) + return { + **state, + "current_step": state["current_step"] + 1, + "results": new_results, + } + +# ---------- Should continue node ---------- +def should_continue(state: PlanningState) -> str: + if state["current_step"] >= len(state["plan"]): + return "finish" + return "execute" + +# ---------- Graph construction ---------- +builder = StateGraph(PlanningState) +builder.add_node("planning", planning) +builder.add_node("execution", execution) +builder.add_conditional_edges( + "execution", + should_continue, + { + "finish": END, + "execute": "execution", + }, +) +builder.set_entry_point("planning") +graph = builder.compile() + +# ---------- Run ---------- +if __name__ == "__main__": + import sys + if len(sys.argv) < 2: + print("Usage: python main.py ''") + sys.exit(1) + task = sys.argv[1] + initial_state: PlanningState = { + "task": task, + "plan": None, + "current_step": 0, + "results": [], + } + result = graph.invoke(initial_state) + plan = result["plan"] or [] + print(f"\nЗадача: {task}\n") + print("План:") + for i, step in enumerate(plan, 1): + print(f"{i}. {step}") + print("\n[Шаги]") + for i, res in enumerate(result["results"], 1): + print(f"[{i}] {res}\n") + print("Итог: ", "\n".join(result["results"]))