commit 22367665ca9b346eb1a301f0d48913ede65450c0 Author: Глеб Никишин Date: Thu May 28 17:25:32 2026 +0000 add main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..3f03911 --- /dev/null +++ b/main.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python +# main.py +# LangGraph planning agent example + +import os +from typing import TypedDict, List, Dict, Any + +from langgraph.graph import StateGraph, END +from langgraph.prebuilt import create_chat_agent +from langchain_openai import ChatOpenAI +from langchain_core.messages import HumanMessage, AIMessage + +# ---------- LLM ---------- +llm = ChatOpenAI( + model="openai/gpt-oss-20b:free", + base_url="https://openrouter.ai/api/v1", + api_key=os.getenv("OPENAI_API_KEY"), + temperature=0.0, +) + +# ---------- State ---------- +class PlanningState(TypedDict): + task: str + plan: List[str] | None + current_step: int + results: List[str] + +# ---------- Planning node ---------- +planning_prompt = ( + "You are a helpful assistant. Given the following task, break it into 3-6 concrete steps. " + "Return the steps as a numbered list or JSON array. The steps should be short, actionable, and in Russian." +) + +async def planning(state: PlanningState) -> PlanningState: + task = state["task"] + messages = [HumanMessage(content=f"{planning_prompt}\nTask: {task}")] + response = await llm.ainvoke(messages) + text = response.content.strip() + + # Try to parse JSON first + plan: List[str] | None = None + try: + import json + data = json.loads(text) + if isinstance(data, list): + plan = [str(item).strip() for item in data] + except Exception: + pass + + # If JSON parsing failed, try to extract numbered list + if plan is None: + import re + lines = re.findall(r"\d+\.\s*(.+)", text) + if lines: + plan = [line.strip() for line in lines] + + if plan is None: + raise ValueError("Could not parse plan from LLM response") + + return { + "task": task, + "plan": plan, + "current_step": 0, + "results": [], + } + +# ---------- Execution node ---------- +async def execution(state: PlanningState) -> PlanningState: + plan = state["plan"] + idx = state["current_step"] + if plan is None or idx >= len(plan): + return state + + step = plan[idx] + # Execute the step: ask LLM to produce result for this step + messages = [HumanMessage(content=f"Task: {state['task']}\nStep {idx+1}: {step}\nProvide the result for this step.")] + response = await llm.ainvoke(messages) + result = response.content.strip() + + new_results = state["results"].copy() + new_results.append(result) + + return { + "task": state["task"], + "plan": plan, + "current_step": idx + 1, + "results": new_results, + } + +# ---------- Should continue ---------- +async def should_continue(state: PlanningState) -> str: + if state["current_step"] >= len(state["plan"]): + return "finish" + return "execute" + +# ---------- Build graph ---------- +builder = StateGraph(PlanningState) +builder.add_node("planning", planning) +builder.add_node("execution", execution) +builder.add_conditional_edges( + "planning", + lambda _: "execute", +) +builder.add_conditional_edges( + "execution", + should_continue, + { + "execute": "execution", + "finish": END, + }, +) +builder.set_entry_point("planning") +graph = builder.compile() + +# ---------- Run example ---------- +if __name__ == "__main__": + task = "Сравни Python и JavaScript" + initial_state: PlanningState = { + "task": task, + "plan": None, + "current_step": 0, + "results": [], + } + result = graph.invoke(initial_state) + + print(f"\nЗадача: {task}\n") + print("План:") + for i, step in enumerate(result["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"]))