Add Optional files/todos Propagation and Fix Sub-Agent State Issues

pyproject.toml

@@ -22,4 +22,4 @@ packages = ["deepagents"]
 "deepagents" = "src/deepagents"
 
 [tool.setuptools.package-data]
-"*" = ["py.typed"]
+"*" = ["py.typed"]

src/deepagents/graph.py

@@ -94,4 +94,4 @@ def create_deep_agent(
         post_model_hook=selected_post_model_hook,
         config_schema=config_schema,
         checkpointer=checkpointer,
-    )
+    )

src/deepagents/sub_agent.py

@@ -17,20 +17,31 @@ class SubAgent(TypedDict):
     description: str
     prompt: str
     tools: NotRequired[list[str]]
+    include_files: NotRequired[bool]
+    include_todos: NotRequired[bool]
     # Optional per-subagent model configuration
     model_settings: NotRequired[dict[str, Any]]
 
 
 def _create_task_tool(tools, instructions, subagents: list[SubAgent], model, state_schema):
     agents = {
-        "general-purpose": create_react_agent(model, prompt=instructions, tools=tools, checkpointer=False)
+        "general-purpose": create_react_agent(
+            model, prompt=instructions, tools=tools, state_schema=state_schema
+        )
+    }
+    subagent_configs = {
+        "general-purpose": {"include_files": True, "include_todos": True}
     }
     tools_by_name = {}
     for tool_ in tools:
         if not isinstance(tool_, BaseTool):
             tool_ = tool(tool_)
         tools_by_name[tool_.name] = tool_
     for _agent in subagents:
+        subagent_configs[_agent["name"]] = {
+            "include_files": _agent.get("include_files", False),
+            "include_todos": _agent.get("include_todos", False)
+        }
         if "tools" in _agent:
             _tools = [tools_by_name[t] for t in _agent["tools"]]
         else:
@@ -63,17 +74,22 @@ async def task(
         if subagent_type not in agents:
             return f"Error: invoked agent of type {subagent_type}, the only allowed types are {[f'`{k}`' for k in agents]}"
         sub_agent = agents[subagent_type]
-        state["messages"] = [{"role": "user", "content": description}]
-        result = await sub_agent.ainvoke(state)
-        return Command(
-            update={
-                "files": result.get("files", {}),
-                "messages": [
-                    ToolMessage(
-                        result["messages"][-1].content, tool_call_id=tool_call_id
-                    )
-                ],
-            }
-        )
+        sub_state = {"messages": [{"role": "user", "content": description}]}
+        config = subagent_configs.get(subagent_type, {"include_files": False, "include_todos": False})
+        if config["include_files"]:
+            sub_state["files"] = state.get("files", {})
+        if config["include_todos"]:
+            sub_state["todos"] = state.get("todos", [])
+        result = await sub_agent.ainvoke(sub_state)
+        update = {
+            "messages": [
+                ToolMessage(result["messages"][-1].content, tool_call_id=tool_call_id)
+            ]
+        }
+        if config["include_files"]:
+            update["files"] = result.get("files", {})
+        if config["include_todos"]:
+            update["todos"] = result.get("todos", [])
+        return Command(update=update)
 
-    return task
+    return task