add main.py

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 '<task>'")
+        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"]))