diff --git a/main.py b/main.py index 7e37827..7424d49 100644 --- a/main.py +++ b/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)) \ No newline at end of file + example_task = "Сравни Python и JavaScript" + result = asyncio.run(run_planning_agent(example_task)) + print(result) \ No newline at end of file