chore(langchain): speed up todo list middleware init#36311
chore(langchain): speed up todo list middleware init#36311Eugene Yurtsev (eyurtsev) merged 3 commits intomasterfrom
Conversation
Eugene Yurtsev (eyurtsev)
changed the title
github-actions
bot
added
infra
Eugene Yurtsev (eyurtsev)
changed the title
There was a problem hiding this comment.
Pull request overview
This PR aims to speed up TodoListMiddleware initialization by avoiding per-instance nested tool function definition and switching to a StructuredTool built from a module-level implementation.
Changes:
- Replace the nested
@toolfunction defined inTodoListMiddleware.__init__with a module-level_write_todosimplementation. - Build the
write_todostool viaStructuredTool.from_function(...)with an explicit Pydantic args schema (WriteTodosInput). - Adjust imports to support the new tool construction and runtime access.
Comments suppressed due to low confidence (1)
libs/langchain_v1/langchain/agents/middleware/todo.py:128
- This module now contains both a
write_todosfunction and a_write_todosfunction, while the tool itself is namedwrite_todos. That overlap makes it harder to tell which implementation is the actual tool entrypoint (and increases the chance of accidental misuse during future refactors). Consider renaming_write_todosto a clearer tool-specific name (e.g.,write_todos_tool_impl) or reusing the existingwrite_todosfunction as the tool implementation (if compatible).
def write_todos(
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def _write_todos( | ||
| runtime: ToolRuntime[ContextT, PlanningState[ResponseT]], todos: list[Todo] | ||
| ) -> Command[Any]: | ||
| """Create and manage a structured task list for your current work session.""" | ||
| return Command( | ||
| update={ | ||
| "todos": todos, | ||
| "messages": [ | ||
| ToolMessage(f"Updated todo list to {todos}", tool_call_id=runtime.tool_call_id) | ||
| ], | ||
| } | ||
| ) |
There was a problem hiding this comment.
StructuredTool.from_function(..., args_schema=WriteTodosInput, infer_schema=False) will validate/provide only the fields from WriteTodosInput (currently just todos). However _write_todos requires a runtime positional argument as well, which is not part of the schema; this is likely to raise a TypeError at invocation time unless StructuredTool has an explicit injection mechanism for ToolRuntime in this codebase. To make this robust, align the tool function signature with the args schema (e.g., accept only todos) and use the established injection mechanism for the tool call id (such as an injected tool_call_id argument), or otherwise ensure runtime is provided by the tool runner in a documented/supported way.
| args_schema=WriteTodosInput, | ||
| infer_schema=False, |
There was a problem hiding this comment.
StructuredTool.from_function(..., args_schema=WriteTodosInput, infer_schema=False) will validate/provide only the fields from WriteTodosInput (currently just todos). However _write_todos requires a runtime positional argument as well, which is not part of the schema; this is likely to raise a TypeError at invocation time unless StructuredTool has an explicit injection mechanism for ToolRuntime in this codebase. To make this robust, align the tool function signature with the args schema (e.g., accept only todos) and use the established injection mechanism for the tool call id (such as an injected tool_call_id argument), or otherwise ensure runtime is provided by the tool runner in a documented/supported way.
| args_schema=WriteTodosInput, | |
| infer_schema=False, |
|
|
||
|
|
||
| # Dynamically create the write_todos tool with the custom description | ||
| def _write_todos( |
There was a problem hiding this comment.
This module now contains both a write_todos function and a _write_todos function, while the tool itself is named write_todos. That overlap makes it harder to tell which implementation is the actual tool entrypoint (and increases the chance of accidental misuse during future refactors). Consider renaming _write_todos to a clearer tool-specific name (e.g., write_todos_tool_impl) or reusing the existing write_todos function as the tool implementation (if compatible).
| from collections.abc import Awaitable, Callable | ||
| from typing import Annotated, Any, Literal, cast | ||
|
|
||
| from langchain_core.messages import AIMessage, SystemMessage, ToolMessage | ||
| from langchain_core.tools import tool | ||
| from langchain_core.tools import InjectedToolCallId, StructuredTool, tool | ||
| from langgraph.runtime import Runtime | ||
| from langgraph.types import Command | ||
| from pydantic import BaseModel, Field |
There was a problem hiding this comment.
This change removes from __future__ import annotations and the TYPE_CHECKING-guarded imports, which means langgraph.runtime.Runtime (and other types) are now imported/evaluated at module import time. That can increase cold-start/import overhead and may counteract the stated goal of speeding up initialization. If runtime evaluation isn’t needed, consider restoring postponed annotations (or moving heavier type-only imports behind TYPE_CHECKING) so these imports don’t execute on every import.
Speed up todo list middleware initialization