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

main.py

@@ -0,0 +1,162 @@
+import os
+import asyncio
+import json
+from typing import TypedDict, Annotated, List
+
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.output_parsers import JsonOutputParser
+from langchain.tools import tool
+from deepagents import create_deep_agent
+from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
+from langgraph.graph import StateGraph, START, END
+from langgraph.graph.message import add_messages
+
+# ---------- 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,
+)
+
+# ---------- Backend ----------
+backend = CompositeBackend(
+    [
+        LocalShellBackend(workspace_dir="./workspace"),
+        FilesystemBackend(),
+    ]
+)
+
+# ---------- Tools (optional, can be used by LLM) ----------
+@tool
+def search_web(query: str) -> str:
+    """Search the web for the given query and return a short summary."""
+    # Placeholder implementation - in real use you could call an API.
+    return f"Search result for '{query}' (mock)."
+
+# ---------- Deep Agent ----------
+deep_agent = create_deep_agent(
+    model=llm,
+    tools=[search_web],
+    backend=backend,
+    system_prompt="You are a helpful planning assistant.",
+)
+
+# ---------- State ----------
+class PlanningState(TypedDict):
+    messages: Annotated[List[AIMessage | HumanMessage], add_messages]
+    task: str
+    plan: List[str] | None
+    current_step: int
+    results: List[str]
+
+# ---------- Planner Node ----------
+def planner_node(state: PlanningState) -> PlanningState:
+    prompt = f"""You are given a task. Break it into 3-6 concrete steps.
+Return the plan as a JSON array of strings under the key "plan".
+Task: {state['task']}"""
+    response = deep_agent.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        {"configurable": {"thread_id": f"planner-{state['task'][:10]}"}},
+    )
+    content = response["messages"][-1].content
+    parser = JsonOutputParser()
+    try:
+        plan = parser.parse(content)
+        if not isinstance(plan, list):
+            raise ValueError
+    except Exception:
+        # Fallback: try to extract lines starting with numbers
+        lines = [line.strip() for line in content.splitlines() if line.strip()]
+        plan = [line.split(".", 1)[-1].strip() for line in lines if line[0].isdigit()]
+    return {
+        "messages": state["messages"],
+        "task": state["task"],
+        "plan": plan,
+        "current_step": 0,
+        "results": [],
+    }
+
+# ---------- Executor Node ----------
+def executor_node(state: PlanningState) -> PlanningState:
+    step_idx = state["current_step"]
+    step_instruction = state["plan"][step_idx]
+    prompt = f"""You are executing step {step_idx + 1} of a plan.
+Step description: {step_instruction}
+Provide a concise answer for this step."""
+    response = deep_agent.invoke(
+        {"messages": [HumanMessage(content=prompt)]},
+        {"configurable": {"thread_id": f"executor-{state['task'][:10]}"}},
+    )
+    result = response["messages"][-1].content
+    new_results = state["results"] + [f"[Step {step_idx + 1}] {result}"]
+    return {
+        "messages": state["messages"],
+        "task": state["task"],
+        "plan": state["plan"],
+        "current_step": step_idx + 1,
+        "results": new_results,
+    }
+
+# ---------- Conditional Edge ----------
+def should_continue(state: PlanningState) -> str:
+    if state["current_step"] >= len(state["plan"]):
+        return "finish"
+    return "execute"
+
+# ---------- Graph ----------
+graph = StateGraph(PlanningState)
+
+graph.add_node("planning", planner_node)
+graph.add_node("execution", executor_node)
+
+graph.add_edge(START, "planning")
+graph.add_edge("planning", "execution")
+graph.add_conditional_edges(
+    "execution",
+    should_continue,
+    {"execute": "execution", "finish": END},
+)
+
+graph.set_entry_point("planning")
+app = graph.compile()
+
+# ---------- Demo ----------
+async def run_demo(task: str):
+    # Initialize empty state
+    init_state: PlanningState = {
+        "messages": [],
+        "task": task,
+        "plan": None,
+        "current_step": 0,
+        "results": [],
+    }
+    async for event in app.astream(
+        init_state,
+        {"configurable": {"thread_id": "demo-session"}},
+    ):
+        # We only care about final state
+        pass
+    final_state = event
+    print(f"Задача: {task}\n")
+    print("План:")
+    for i, step in enumerate(final_state["plan"], 1):
+        print(f"{i}. {step}")
+    print()
+    for res in final_state["results"]:
+        print(res)
+    print("\nИтог:")
+    summary_prompt = f"""Based on the following step results, provide a concise final summary.
+
+Results:
+{chr(10).join(final_state['results'])}"""
+    summary_resp = deep_agent.invoke(
+        {"messages": [HumanMessage(content=summary_prompt)]},
+        {"configurable": {"thread_id": "summary"}},
+    )
+    print(summary_resp["messages"][-1].content)
+
+if __name__ == "__main__":
+    demo_task = "Сравни Python и JavaScript"
+    asyncio.run(run_demo(demo_task))