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

main.py

@@ -3,12 +3,13 @@ import asyncio
 import json
 from typing import TypedDict, Annotated, List
 
-from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_openai import ChatOpenAI
 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
 
@@ -28,19 +29,18 @@ backend = CompositeBackend(
     ]
 )
 
-# ---------- Tools (optional, can be used by LLM) ----------
+# ---------- Tools (optional example) ----------
 @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)."
+def echo_tool(text: str) -> str:
+    """Return the given text unchanged."""
+    return text
 
 # ---------- Deep Agent ----------
 deep_agent = create_deep_agent(
     model=llm,
-    tools=[search_web],
+    tools=[echo_tool],
     backend=backend,
-    system_prompt="You are a helpful planning assistant.",
+    system_prompt="You are a helpful planning agent.",
 )
 
 # ---------- State ----------
@@ -54,50 +54,42 @@ class PlanningState(TypedDict):
 # ---------- 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".
+Return the plan as a JSON array of strings, e.g. ["step 1", "step 2", ...].
 Task: {state['task']}"""
     response = deep_agent.invoke(
         {"messages": [HumanMessage(content=prompt)]},
-        {"configurable": {"thread_id": f"planner-{state['task'][:10]}"}},
+        {"configurable": {"thread_id": "planner"}},
     )
     content = response["messages"][-1].content
-    parser = JsonOutputParser()
     try:
-        plan = parser.parse(content)
+        plan = json.loads(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": [],
-    }
+    state["plan"] = plan
+    state["current_step"] = 0
+    state["results"] = []
+    return state
 
-# ---------- Executor Node ----------
-def executor_node(state: PlanningState) -> PlanningState:
+# ---------- Execution Node ----------
+def execution_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}
+Task: {state['task']}
+Step: {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]}"}},
+        {"configurable": {"thread_id": f"exec-{step_idx}"}},
     )
     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,
-    }
+    state["results"].append(f"[Step {step_idx + 1}] {result}")
+    state["current_step"] += 1
+    return state
 
 # ---------- Conditional Edge ----------
 def should_continue(state: PlanningState) -> str:
@@ -105,11 +97,11 @@ def should_continue(state: PlanningState) -> str:
         return "finish"
     return "execute"
 
-# ---------- Graph ----------
+# ---------- Build Graph ----------
 graph = StateGraph(PlanningState)
 
 graph.add_node("planning", planner_node)
-graph.add_node("execution", executor_node)
+graph.add_node("execution", execution_node)
 
 graph.add_edge(START, "planning")
 graph.add_edge("planning", "execution")
@@ -120,43 +112,50 @@ graph.add_conditional_edges(
 )
 
 graph.set_entry_point("planning")
-app = graph.compile()
+graph = graph.compile()
 
-# ---------- Demo ----------
-async def run_demo(task: str):
-    # Initialize empty state
-    init_state: PlanningState = {
+# ---------- Runner ----------
+async def run_planning_agent(task: str) -> str:
+    # Initialise state
+    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
+    # Run graph
+    async for event in graph.astream(state):
+        # We only need 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'])}"""
+    # Build final answer
+    plan_text = "\n".join(f"{i+1}. {step}" for i, step in enumerate(final_state["plan"]))
+    steps_text = "\n".join(final_state["results"])
+    summary_prompt = f"""You have completed all steps of the following task.
+Task: {task}
+Plan:
+{plan_text}
+Steps results:
+{steps_text}
+Provide a concise final summary."""
     summary_resp = deep_agent.invoke(
         {"messages": [HumanMessage(content=summary_prompt)]},
         {"configurable": {"thread_id": "summary"}},
     )
-    print(summary_resp["messages"][-1].content)
+    summary = summary_resp["messages"][-1].content
+    output = f"""Задача: {task}
 
+План:
+{plan_text}
+
+{steps_text}
+Итог: {summary}
+"""
+    return output
+
+# ---------- Main ----------
 if __name__ == "__main__":
-    demo_task = "Сравни Python и JavaScript"
-    asyncio.run(run_demo(demo_task))
+    example_task = "Сравни Python и JavaScript"
+    result = asyncio.run(run_planning_agent(example_task))
+    print(result)