|
7 | 7 | if TYPE_CHECKING: |
8 | 8 | from collections.abc import Awaitable, Callable |
9 | 9 |
|
| 10 | + from langchain_core.tools import BaseTool |
| 11 | + |
10 | 12 | from langchain_core.messages import SystemMessage, ToolMessage |
11 | 13 | from langchain_core.tools import tool |
12 | 14 | from langgraph.types import Command |
@@ -115,15 +117,37 @@ class PlanningState(AgentState): |
115 | 117 | - Don't be afraid to revise the To-Do list as you go. New information may reveal new tasks that need to be done, or old tasks that are irrelevant.""" # noqa: E501 |
116 | 118 |
|
117 | 119 |
|
118 | | -@tool(description=WRITE_TODOS_TOOL_DESCRIPTION) |
119 | | -def write_todos(todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]) -> Command: |
120 | | - """Create and manage a structured task list for your current work session.""" |
121 | | - return Command( |
122 | | - update={ |
123 | | - "todos": todos, |
124 | | - "messages": [ToolMessage(f"Updated todo list to {todos}", tool_call_id=tool_call_id)], |
125 | | - } |
126 | | - ) |
| 120 | +def _create_write_todos_tool( |
| 121 | + description: str = WRITE_TODOS_TOOL_DESCRIPTION, |
| 122 | +) -> BaseTool: |
| 123 | + """Factory function to create a write_todos tool with custom description. |
| 124 | +
|
| 125 | + Args: |
| 126 | + description: The description for the tool. |
| 127 | +
|
| 128 | + Returns: |
| 129 | + A configured write_todos tool instance. |
| 130 | + """ |
| 131 | + |
| 132 | + @tool(description=description) |
| 133 | + def write_todos( |
| 134 | + todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId] |
| 135 | + ) -> Command: |
| 136 | + """Create and manage a structured task list for your current work session.""" |
| 137 | + return Command( |
| 138 | + update={ |
| 139 | + "todos": todos, |
| 140 | + "messages": [ |
| 141 | + ToolMessage(f"Updated todo list to {todos}", tool_call_id=tool_call_id) |
| 142 | + ], |
| 143 | + } |
| 144 | + ) |
| 145 | + |
| 146 | + return write_todos |
| 147 | + |
| 148 | + |
| 149 | +# Module-level tool with default description |
| 150 | +write_todos = _create_write_todos_tool() |
127 | 151 |
|
128 | 152 |
|
129 | 153 | class TodoListMiddleware(AgentMiddleware): |
@@ -169,23 +193,7 @@ def __init__( |
169 | 193 | super().__init__() |
170 | 194 | self.system_prompt = system_prompt |
171 | 195 | self.tool_description = tool_description |
172 | | - |
173 | | - # Dynamically create the write_todos tool with the custom description |
174 | | - @tool(description=self.tool_description) |
175 | | - def write_todos( |
176 | | - todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId] |
177 | | - ) -> Command: |
178 | | - """Create and manage a structured task list for your current work session.""" |
179 | | - return Command( |
180 | | - update={ |
181 | | - "todos": todos, |
182 | | - "messages": [ |
183 | | - ToolMessage(f"Updated todo list to {todos}", tool_call_id=tool_call_id) |
184 | | - ], |
185 | | - } |
186 | | - ) |
187 | | - |
188 | | - self.tools = [write_todos] |
| 196 | + self.tools = [_create_write_todos_tool(self.tool_description)] |
189 | 197 |
|
190 | 198 | def wrap_model_call( |
191 | 199 | self, |
|
0 commit comments