langchain 1.2.3__py3-none-any.whl → 1.2.4__py3-none-any.whl

langchain/agents/middleware/human_in_the_loop.py

@@ -102,7 +102,9 @@ class HITLResponse(TypedDict):
 class _DescriptionFactory(Protocol):
     """Callable that generates a description for a tool call."""
 
-    def __call__(self, tool_call: ToolCall, state: AgentState, runtime: Runtime[ContextT]) -> str:
+    def __call__(
+        self, tool_call: ToolCall, state: AgentState[Any], runtime: Runtime[ContextT]
+    ) -> str:
         """Generate a description for a tool call."""
         ...
 
@@ -203,7 +205,7 @@ class HumanInTheLoopMiddleware(AgentMiddleware[StateT, ContextT]):
         self,
         tool_call: ToolCall,
         config: InterruptOnConfig,
-        state: AgentState,
+        state: AgentState[Any],
         runtime: Runtime[ContextT],
     ) -> tuple[ActionRequest, ReviewConfig]:
         """Create an ActionRequest and ReviewConfig for a tool call."""
@@ -235,8 +237,8 @@ class HumanInTheLoopMiddleware(AgentMiddleware[StateT, ContextT]):
 
         return action_request, review_config
 
+    @staticmethod
     def _process_decision(
-        self,
         decision: Decision,
         tool_call: ToolCall,
         config: InterruptOnConfig,
@@ -277,8 +279,22 @@ class HumanInTheLoopMiddleware(AgentMiddleware[StateT, ContextT]):
         )
         raise ValueError(msg)
 
-    def after_model(self, state: AgentState, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
-        """Trigger interrupt flows for relevant tool calls after an `AIMessage`."""
+    def after_model(
+        self, state: AgentState[Any], runtime: Runtime[ContextT]
+    ) -> dict[str, Any] | None:
+        """Trigger interrupt flows for relevant tool calls after an `AIMessage`.
+
+        Args:
+            state: The current agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Updated message with the revised tool calls.
+
+        Raises:
+            ValueError: If the number of human decisions does not match the number of
+                interrupted tool calls.
+        """
         messages = state["messages"]
         if not messages:
             return None
@@ -351,7 +367,15 @@ class HumanInTheLoopMiddleware(AgentMiddleware[StateT, ContextT]):
         return {"messages": [last_ai_msg, *artificial_tool_messages]}
 
     async def aafter_model(
-        self, state: AgentState, runtime: Runtime[ContextT]
+        self, state: AgentState[Any], runtime: Runtime[ContextT]
     ) -> dict[str, Any] | None:
-        """Async trigger interrupt flows for relevant tool calls after an `AIMessage`."""
+        """Async trigger interrupt flows for relevant tool calls after an `AIMessage`.
+
+        Args:
+            state: The current agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Updated message with the revised tool calls.
+        """
         return self.after_model(state, runtime)

langchain/agents/middleware/model_call_limit.py

@@ -19,7 +19,7 @@ if TYPE_CHECKING:
     from langgraph.runtime import Runtime
 
 
-class ModelCallLimitState(AgentState):
+class ModelCallLimitState(AgentState[Any]):
     """State schema for `ModelCallLimitMiddleware`.
 
     Extends `AgentState` with model call tracking fields.

langchain/agents/middleware/model_retry.py

@@ -163,7 +163,8 @@ class ModelRetryMiddleware(AgentMiddleware):
         self.max_delay = max_delay
         self.jitter = jitter
 
-    def _format_failure_message(self, exc: Exception, attempts_made: int) -> AIMessage:
+    @staticmethod
+    def _format_failure_message(exc: Exception, attempts_made: int) -> AIMessage:
         """Format the failure message when retries are exhausted.
 
         Args:
@@ -218,6 +219,9 @@ class ModelRetryMiddleware(AgentMiddleware):
 
         Returns:
             `ModelResponse` or `AIMessage` (the final result).
+
+        Raises:
+            RuntimeError: If the retry loop completes without returning. (This should not happen.)
         """
         # Initial attempt + retries
         for attempt in range(self.max_retries + 1):
@@ -265,6 +269,9 @@ class ModelRetryMiddleware(AgentMiddleware):
 
         Returns:
             `ModelResponse` or `AIMessage` (the final result).
+
+        Raises:
+            RuntimeError: If the retry loop completes without returning. (This should not happen.)
         """
         # Initial attempt + retries
         for attempt in range(self.max_retries + 1):

langchain/agents/middleware/pii.py

@@ -164,7 +164,7 @@ class PIIMiddleware(AgentMiddleware):
     @override
     def before_model(
         self,
-        state: AgentState,
+        state: AgentState[Any],
         runtime: Runtime,
     ) -> dict[str, Any] | None:
         """Check user messages and tool results for PII before model invocation.
@@ -259,7 +259,7 @@ class PIIMiddleware(AgentMiddleware):
     @hook_config(can_jump_to=["end"])
     async def abefore_model(
         self,
-        state: AgentState,
+        state: AgentState[Any],
         runtime: Runtime,
     ) -> dict[str, Any] | None:
         """Async check user messages and tool results for PII before model invocation.
@@ -280,7 +280,7 @@ class PIIMiddleware(AgentMiddleware):
     @override
     def after_model(
         self,
-        state: AgentState,
+        state: AgentState[Any],
         runtime: Runtime,
     ) -> dict[str, Any] | None:
         """Check AI messages for PII after model invocation.
@@ -339,7 +339,7 @@ class PIIMiddleware(AgentMiddleware):
 
     async def aafter_model(
         self,
-        state: AgentState,
+        state: AgentState[Any],
         runtime: Runtime,
     ) -> dict[str, Any] | None:
         """Async check AI messages for PII after model invocation.

langchain/agents/middleware/shell_tool.py

@@ -78,7 +78,7 @@ class _SessionResources:
     session: ShellSession
     tempdir: tempfile.TemporaryDirectory[str] | None
     policy: BaseExecutionPolicy
-    finalizer: weakref.finalize = field(init=False, repr=False)
+    finalizer: weakref.finalize = field(init=False, repr=False)  # type: ignore[type-arg]
 
     def __post_init__(self) -> None:
         self.finalizer = weakref.finalize(
@@ -90,7 +90,7 @@ class _SessionResources:
         )
 
 
-class ShellToolState(AgentState):
+class ShellToolState(AgentState[Any]):
     """Agent state extension for tracking shell session resources."""
 
     shell_session_resources: NotRequired[
@@ -134,7 +134,11 @@ class ShellSession:
         self._terminated = False
 
     def start(self) -> None:
-        """Start the shell subprocess and reader threads."""
+        """Start the shell subprocess and reader threads.
+
+        Raises:
+            RuntimeError: If the shell session pipes cannot be initialized.
+        """
         if self._process and self._process.poll() is None:
             return
 
@@ -604,19 +608,35 @@ class ShellToolMiddleware(AgentMiddleware[ShellToolState, Any]):
         normalized: dict[str, str] = {}
         for key, value in env.items():
             if not isinstance(key, str):
-                msg = "Environment variable names must be strings."
+                msg = "Environment variable names must be strings."  # type: ignore[unreachable]
                 raise TypeError(msg)
             normalized[key] = str(value)
         return normalized
 
     @override
     def before_agent(self, state: ShellToolState, runtime: Runtime) -> dict[str, Any] | None:
-        """Start the shell session and run startup commands."""
+        """Start the shell session and run startup commands.
+
+        Args:
+            state: The current agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Shell session resources to be stored in the agent state.
+        """
         resources = self._get_or_create_resources(state)
         return {"shell_session_resources": resources}
 
     async def abefore_agent(self, state: ShellToolState, runtime: Runtime) -> dict[str, Any] | None:
-        """Async start the shell session and run startup commands."""
+        """Async start the shell session and run startup commands.
+
+        Args:
+            state: The current agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Shell session resources to be stored in the agent state.
+        """
         return self.before_agent(state, runtime)
 
     @override

langchain/agents/middleware/summarization.py

@@ -269,8 +269,16 @@ class SummarizationMiddleware(AgentMiddleware):
             raise ValueError(msg)
 
     @override
-    def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
-        """Process messages before model invocation, potentially triggering summarization."""
+    def before_model(self, state: AgentState[Any], runtime: Runtime) -> dict[str, Any] | None:
+        """Process messages before model invocation, potentially triggering summarization.
+
+        Args:
+            state: The agent state.
+            runtime: The runtime environment.
+
+        Returns:
+            An updated state with summarized messages if summarization was performed.
+        """
         messages = state["messages"]
         self._ensure_message_ids(messages)
 
@@ -297,8 +305,18 @@ class SummarizationMiddleware(AgentMiddleware):
         }
 
     @override
-    async def abefore_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
-        """Process messages before model invocation, potentially triggering summarization."""
+    async def abefore_model(
+        self, state: AgentState[Any], runtime: Runtime
+    ) -> dict[str, Any] | None:
+        """Process messages before model invocation, potentially triggering summarization.
+
+        Args:
+            state: The agent state.
+            runtime: The runtime environment.
+
+        Returns:
+            An updated state with summarized messages if summarization was performed.
+        """
         messages = state["messages"]
         self._ensure_message_ids(messages)
 
@@ -449,7 +467,8 @@ class SummarizationMiddleware(AgentMiddleware):
 
         return max_input_tokens
 
-    def _validate_context_size(self, context: ContextSize, parameter_name: str) -> ContextSize:
+    @staticmethod
+    def _validate_context_size(context: ContextSize, parameter_name: str) -> ContextSize:
         """Validate context configuration tuples."""
         kind, value = context
         if kind == "fraction":
@@ -465,19 +484,24 @@ class SummarizationMiddleware(AgentMiddleware):
             raise ValueError(msg)
         return context
 
-    def _build_new_messages(self, summary: str) -> list[HumanMessage]:
+    @staticmethod
+    def _build_new_messages(summary: str) -> list[HumanMessage]:
         return [
-            HumanMessage(content=f"Here is a summary of the conversation to date:\n\n{summary}")
+            HumanMessage(
+                content=f"Here is a summary of the conversation to date:\n\n{summary}",
+                additional_kwargs={"lc_source": "summarization"},
+            )
         ]
 
-    def _ensure_message_ids(self, messages: list[AnyMessage]) -> None:
+    @staticmethod
+    def _ensure_message_ids(messages: list[AnyMessage]) -> None:
         """Ensure all messages have unique IDs for the add_messages reducer."""
         for msg in messages:
             if msg.id is None:
                 msg.id = str(uuid.uuid4())
 
+    @staticmethod
     def _partition_messages(
-        self,
         conversation_messages: list[AnyMessage],
         cutoff_index: int,
     ) -> tuple[list[AnyMessage], list[AnyMessage]]:
@@ -502,7 +526,8 @@ class SummarizationMiddleware(AgentMiddleware):
         target_cutoff = len(messages) - messages_to_keep
         return self._find_safe_cutoff_point(messages, target_cutoff)
 
-    def _find_safe_cutoff_point(self, messages: list[AnyMessage], cutoff_index: int) -> int:
+    @staticmethod
+    def _find_safe_cutoff_point(messages: list[AnyMessage], cutoff_index: int) -> int:
         """Find a safe cutoff point that doesn't split AI/Tool message pairs.
 
         If the message at `cutoff_index` is a `ToolMessage`, search backward for the

langchain/agents/middleware/todo.py

@@ -12,7 +12,7 @@ if TYPE_CHECKING:
 from langchain_core.messages import AIMessage, SystemMessage, ToolMessage
 from langchain_core.tools import tool
 from langgraph.types import Command
-from typing_extensions import NotRequired, TypedDict
+from typing_extensions import NotRequired, TypedDict, override
 
 from langchain.agents.middleware.types import (
     AgentMiddleware,
@@ -35,7 +35,7 @@ class Todo(TypedDict):
     """The current status of the todo item."""
 
 
-class PlanningState(AgentState):
+class PlanningState(AgentState[Any]):
     """State schema for the todo middleware."""
 
     todos: Annotated[NotRequired[list[Todo]], OmitFromInput]
@@ -118,7 +118,9 @@ Writing todos takes time and tokens, use it when it is helpful for managing comp
 
 
 @tool(description=WRITE_TODOS_TOOL_DESCRIPTION)
-def write_todos(todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]) -> Command:
+def write_todos(
+    todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]
+) -> Command[Any]:
     """Create and manage a structured task list for your current work session."""
     return Command(
         update={
@@ -178,7 +180,7 @@ class TodoListMiddleware(AgentMiddleware):
         @tool(description=self.tool_description)
         def write_todos(
             todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]
-        ) -> Command:
+        ) -> Command[Any]:
             """Create and manage a structured task list for your current work session."""
             return Command(
                 update={
@@ -196,7 +198,16 @@ class TodoListMiddleware(AgentMiddleware):
         request: ModelRequest,
         handler: Callable[[ModelRequest], ModelResponse],
     ) -> ModelCallResult:
-        """Update the system message to include the todo system prompt."""
+        """Update the system message to include the todo system prompt.
+
+        Args:
+            request: Model request to execute (includes state and runtime).
+            handler: Async callback that executes the model request and returns
+                `ModelResponse`.
+
+        Returns:
+            The model call result.
+        """
         if request.system_message is not None:
             new_system_content = [
                 *request.system_message.content_blocks,
@@ -214,7 +225,16 @@ class TodoListMiddleware(AgentMiddleware):
         request: ModelRequest,
         handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
     ) -> ModelCallResult:
-        """Update the system message to include the todo system prompt (async version)."""
+        """Update the system message to include the todo system prompt.
+
+        Args:
+            request: Model request to execute (includes state and runtime).
+            handler: Async callback that executes the model request and returns
+                `ModelResponse`.
+
+        Returns:
+            The model call result.
+        """
         if request.system_message is not None:
             new_system_content = [
                 *request.system_message.content_blocks,
@@ -227,11 +247,8 @@ class TodoListMiddleware(AgentMiddleware):
         )
         return await handler(request.override(system_message=new_system_message))
 
-    def after_model(
-        self,
-        state: AgentState,
-        runtime: Runtime,  # noqa: ARG002
-    ) -> dict[str, Any] | None:
+    @override
+    def after_model(self, state: AgentState[Any], runtime: Runtime) -> dict[str, Any] | None:
         """Check for parallel write_todos tool calls and return errors if detected.
 
         The todo list is designed to be updated at most once per model turn. Since
@@ -280,11 +297,8 @@ class TodoListMiddleware(AgentMiddleware):
 
         return None
 
-    async def aafter_model(
-        self,
-        state: AgentState,
-        runtime: Runtime,
-    ) -> dict[str, Any] | None:
+    @override
+    async def aafter_model(self, state: AgentState[Any], runtime: Runtime) -> dict[str, Any] | None:
         """Check for parallel write_todos tool calls and return errors if detected.
 
         Async version of `after_model`. The todo list is designed to be updated at

langchain/agents/middleware/tool_emulator.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import HumanMessage, ToolMessage
@@ -109,8 +109,8 @@ class LLMToolEmulator(AgentMiddleware):
     def wrap_tool_call(
         self,
         request: ToolCallRequest,
-        handler: Callable[[ToolCallRequest], ToolMessage | Command],
-    ) -> ToolMessage | Command:
+        handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
+    ) -> ToolMessage | Command[Any]:
         """Emulate tool execution using LLM if tool should be emulated.
 
         Args:
@@ -159,8 +159,8 @@ class LLMToolEmulator(AgentMiddleware):
     async def awrap_tool_call(
         self,
         request: ToolCallRequest,
-        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
-    ) -> ToolMessage | Command:
+        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]],
+    ) -> ToolMessage | Command[Any]:
         """Async version of `wrap_tool_call`.
 
         Emulate tool execution using LLM if tool should be emulated.

langchain/agents/middleware/tool_retry.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import asyncio
 import time
 import warnings
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from langchain_core.messages import ToolMessage
 
@@ -189,14 +189,14 @@ class ToolRetryMiddleware(AgentMiddleware):
 
         # Handle backwards compatibility for deprecated on_failure values
         if on_failure == "raise":  # type: ignore[comparison-overlap]
-            msg = (
+            msg = (  # type: ignore[unreachable]
                 "on_failure='raise' is deprecated and will be removed in a future version. "
                 "Use on_failure='error' instead."
             )
             warnings.warn(msg, DeprecationWarning, stacklevel=2)
             on_failure = "error"
         elif on_failure == "return_message":  # type: ignore[comparison-overlap]
-            msg = (
+            msg = (  # type: ignore[unreachable]
                 "on_failure='return_message' is deprecated and will be removed "
                 "in a future version. Use on_failure='continue' instead."
             )
@@ -233,7 +233,8 @@ class ToolRetryMiddleware(AgentMiddleware):
             return True
         return tool_name in self._tool_filter
 
-    def _format_failure_message(self, tool_name: str, exc: Exception, attempts_made: int) -> str:
+    @staticmethod
+    def _format_failure_message(tool_name: str, exc: Exception, attempts_made: int) -> str:
         """Format the failure message when retries are exhausted.
 
         Args:
@@ -287,8 +288,8 @@ class ToolRetryMiddleware(AgentMiddleware):
     def wrap_tool_call(
         self,
         request: ToolCallRequest,
-        handler: Callable[[ToolCallRequest], ToolMessage | Command],
-    ) -> ToolMessage | Command:
+        handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
+    ) -> ToolMessage | Command[Any]:
         """Intercept tool execution and retry on failure.
 
         Args:
@@ -297,6 +298,9 @@ class ToolRetryMiddleware(AgentMiddleware):
 
         Returns:
             `ToolMessage` or `Command` (the final result).
+
+        Raises:
+            RuntimeError: If the retry loop completes without returning. This should not happen.
         """
         tool_name = request.tool.name if request.tool else request.tool_call["name"]
 
@@ -342,8 +346,8 @@ class ToolRetryMiddleware(AgentMiddleware):
     async def awrap_tool_call(
         self,
         request: ToolCallRequest,
-        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
-    ) -> ToolMessage | Command:
+        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]],
+    ) -> ToolMessage | Command[Any]:
         """Intercept and control async tool execution with retry logic.
 
         Args:
@@ -353,6 +357,9 @@ class ToolRetryMiddleware(AgentMiddleware):
 
         Returns:
             `ToolMessage` or `Command` (the final result).
+
+        Raises:
+            RuntimeError: If the retry loop completes without returning. This should not happen.
         """
         tool_name = request.tool.name if request.tool else request.tool_call["name"]
 

langchain/agents/middleware/tool_selection.py

@@ -4,12 +4,7 @@ from __future__ import annotations
 
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Annotated, Literal, Union
-
-if TYPE_CHECKING:
-    from collections.abc import Awaitable, Callable
-
-    from langchain.tools import BaseTool
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Union
 
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import HumanMessage
@@ -24,6 +19,11 @@ from langchain.agents.middleware.types import (
 )
 from langchain.chat_models.base import init_chat_model
 
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from langchain.tools import BaseTool
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_SYSTEM_PROMPT = (
@@ -42,7 +42,7 @@ class _SelectionRequest:
     valid_tool_names: list[str]
 
 
-def _create_tool_selection_response(tools: list[BaseTool]) -> TypeAdapter:
+def _create_tool_selection_response(tools: list[BaseTool]) -> TypeAdapter[Any]:
     """Create a structured output schema for tool selection.
 
     Args:
@@ -51,6 +51,9 @@ def _create_tool_selection_response(tools: list[BaseTool]) -> TypeAdapter:
     Returns:
         `TypeAdapter` for a schema where each tool name is a `Literal` with its
             description.
+
+    Raises:
+        AssertionError: If `tools` is empty.
     """
     if not tools:
         msg = "Invalid usage: tools must be non-empty"
@@ -153,9 +156,16 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
     def _prepare_selection_request(self, request: ModelRequest) -> _SelectionRequest | None:
         """Prepare inputs for tool selection.
 
+        Args:
+            request: the model request.
+
         Returns:
             `SelectionRequest` with prepared inputs, or `None` if no selection is
-                needed.
+            needed.
+
+        Raises:
+            ValueError: If tools in `always_include` are not found in the request.
+            AssertionError: If no user message is found in the request messages.
         """
         # If no tools available, return None
         if not request.tools or len(request.tools) == 0:
@@ -217,7 +227,7 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
 
     def _process_selection_response(
         self,
-        response: dict,
+        response: dict[str, Any],
         available_tools: list[BaseTool],
         valid_tool_names: list[str],
         request: ModelRequest,
@@ -262,7 +272,19 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
         request: ModelRequest,
         handler: Callable[[ModelRequest], ModelResponse],
     ) -> ModelCallResult:
-        """Filter tools based on LLM selection before invoking the model via handler."""
+        """Filter tools based on LLM selection before invoking the model via handler.
+
+        Args:
+            request: Model request to execute (includes state and runtime).
+            handler: Async callback that executes the model request and returns
+                `ModelResponse`.
+
+        Returns:
+            The model call result.
+
+        Raises:
+            AssertionError: If the selection model response is not a dict.
+        """
         selection_request = self._prepare_selection_request(request)
         if selection_request is None:
             return handler(request)
@@ -293,7 +315,19 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
         request: ModelRequest,
         handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
     ) -> ModelCallResult:
-        """Filter tools based on LLM selection before invoking the model via handler."""
+        """Filter tools based on LLM selection before invoking the model via handler.
+
+        Args:
+            request: Model request to execute (includes state and runtime).
+            handler: Async callback that executes the model request and returns
+                `ModelResponse`.
+
+        Returns:
+            The model call result.
+
+        Raises:
+            AssertionError: If the selection model response is not a dict.
+        """
         selection_request = self._prepare_selection_request(request)
         if selection_request is None:
             return await handler(request)