langchain 1.0.0a4__py3-none-any.whl → 1.0.0a6__py3-none-any.whl

langchain/agents/middleware/summarization.py

@@ -48,7 +48,7 @@ Respond ONLY with the extracted context. Do not include any additional informati
 <messages>
 Messages to summarize:
 {messages}
-</messages>"""
+</messages>"""  # noqa: E501
 
 SUMMARY_PREFIX = "## Previous conversation summary:"
 
@@ -98,7 +98,7 @@ class SummarizationMiddleware(AgentMiddleware):
         self.summary_prompt = summary_prompt
         self.summary_prefix = summary_prefix
 
-    def before_model(self, state: AgentState) -> dict[str, Any] | None:
+    def before_model(self, state: AgentState) -> dict[str, Any] | None:  # type: ignore[override]
         """Process messages before model invocation, potentially triggering summarization."""
         messages = state["messages"]
         self._ensure_message_ids(messages)

langchain/agents/middleware/types.py

@@ -8,15 +8,27 @@ from typing import TYPE_CHECKING, Annotated, Any, Generic, Literal, cast
 # needed as top level import for pydantic schema generation on AgentState
 from langchain_core.messages import AnyMessage  # noqa: TC002
 from langgraph.channels.ephemeral_value import EphemeralValue
-from langgraph.graph.message import Messages, add_messages
+from langgraph.graph.message import add_messages
+from langgraph.runtime import Runtime
+from langgraph.typing import ContextT
 from typing_extensions import NotRequired, Required, TypedDict, TypeVar
 
 if TYPE_CHECKING:
     from langchain_core.language_models.chat_models import BaseChatModel
     from langchain_core.tools import BaseTool
+    from langgraph.runtime import Runtime
 
     from langchain.agents.structured_output import ResponseFormat
 
+__all__ = [
+    "AgentMiddleware",
+    "AgentState",
+    "ContextT",
+    "ModelRequest",
+    "OmitFromSchema",
+    "PublicAgentState",
+]
+
 JumpTo = Literal["tools", "model", "__end__"]
 """Destination to jump to when a middleware node returns."""
 
@@ -36,29 +48,53 @@ class ModelRequest:
     model_settings: dict[str, Any] = field(default_factory=dict)
 
 
+@dataclass
+class OmitFromSchema:
+    """Annotation used to mark state attributes as omitted from input or output schemas."""
+
+    input: bool = True
+    """Whether to omit the attribute from the input schema."""
+
+    output: bool = True
+    """Whether to omit the attribute from the output schema."""
+
+
+OmitFromInput = OmitFromSchema(input=True, output=False)
+"""Annotation used to mark state attributes as omitted from input schema."""
+
+OmitFromOutput = OmitFromSchema(input=False, output=True)
+"""Annotation used to mark state attributes as omitted from output schema."""
+
+PrivateStateAttr = OmitFromSchema(input=True, output=True)
+"""Annotation used to mark state attributes as purely internal for a given middleware."""
+
+
 class AgentState(TypedDict, Generic[ResponseT]):
     """State schema for the agent."""
 
     messages: Required[Annotated[list[AnyMessage], add_messages]]
-    model_request: NotRequired[Annotated[ModelRequest | None, EphemeralValue]]
-    jump_to: NotRequired[Annotated[JumpTo | None, EphemeralValue]]
+    jump_to: NotRequired[Annotated[JumpTo | None, EphemeralValue, PrivateStateAttr]]
     response: NotRequired[ResponseT]
 
 
 class PublicAgentState(TypedDict, Generic[ResponseT]):
-    """Input / output schema for the agent."""
+    """Public state schema for the agent.
 
-    messages: Required[Messages]
+    Just used for typing purposes.
+    """
+
+    messages: Required[Annotated[list[AnyMessage], add_messages]]
     response: NotRequired[ResponseT]
 
 
-StateT = TypeVar("StateT", bound=AgentState)
+StateT = TypeVar("StateT", bound=AgentState, default=AgentState)
 
 
-class AgentMiddleware(Generic[StateT]):
+class AgentMiddleware(Generic[StateT, ContextT]):
     """Base middleware class for an agent.
 
-    Subclass this and implement any of the defined methods to customize agent behavior between steps in the main agent loop.
+    Subclass this and implement any of the defined methods to customize agent behavior
+    between steps in the main agent loop.
     """
 
     state_schema: type[StateT] = cast("type[StateT]", AgentState)
@@ -67,12 +103,17 @@ class AgentMiddleware(Generic[StateT]):
     tools: list[BaseTool]
     """Additional tools registered by the middleware."""
 
-    def before_model(self, state: StateT) -> dict[str, Any] | None:
+    def before_model(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
         """Logic to run before the model is called."""
 
-    def modify_model_request(self, request: ModelRequest, state: StateT) -> ModelRequest:  # noqa: ARG002
+    def modify_model_request(
+        self,
+        request: ModelRequest,
+        state: StateT,  # noqa: ARG002
+        runtime: Runtime[ContextT],  # noqa: ARG002
+    ) -> ModelRequest:
         """Logic to modify request kwargs before the model is called."""
         return request
 
-    def after_model(self, state: StateT) -> dict[str, Any] | None:
+    def after_model(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
         """Logic to run after the model is called."""

langchain/agents/middleware_agent.py

@@ -2,7 +2,8 @@
 
 import itertools
 from collections.abc import Callable, Sequence
-from typing import Any, Union
+from inspect import signature
+from typing import Annotated, Any, cast, get_args, get_origin, get_type_hints
 
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import AIMessage, AnyMessage, SystemMessage, ToolMessage
@@ -10,18 +11,19 @@ from langchain_core.runnables import Runnable
 from langchain_core.tools import BaseTool
 from langgraph.constants import END, START
 from langgraph.graph.state import StateGraph
+from langgraph.runtime import Runtime
+from langgraph.types import Send
 from langgraph.typing import ContextT
-from typing_extensions import TypedDict, TypeVar
+from typing_extensions import NotRequired, Required, TypedDict, TypeVar
 
 from langchain.agents.middleware.types import (
     AgentMiddleware,
     AgentState,
     JumpTo,
     ModelRequest,
+    OmitFromSchema,
     PublicAgentState,
 )
-
-# Import structured output classes from the old implementation
 from langchain.agents.structured_output import (
     MultipleStructuredOutputsError,
     OutputToolBinding,
@@ -37,29 +39,52 @@ from langchain.chat_models import init_chat_model
 STRUCTURED_OUTPUT_ERROR_TEMPLATE = "Error: {error}\n Please fix your mistakes."
 
 
-def _merge_state_schemas(schemas: list[type]) -> type:
-    """Merge multiple TypedDict schemas into a single schema with all fields."""
-    if not schemas:
-        return AgentState
+def _resolve_schema(schemas: set[type], schema_name: str, omit_flag: str | None = None) -> type:
+    """Resolve schema by merging schemas and optionally respecting OmitFromSchema annotations.
 
+    Args:
+        schemas: List of schema types to merge
+        schema_name: Name for the generated TypedDict
+        omit_flag: If specified, omit fields with this flag set ('input' or 'output')
+    """
     all_annotations = {}
 
     for schema in schemas:
-        all_annotations.update(schema.__annotations__)
+        hints = get_type_hints(schema, include_extras=True)
+
+        for field_name, field_type in hints.items():
+            should_omit = False
 
-    return TypedDict("MergedState", all_annotations)  # type: ignore[operator]
+            if omit_flag:
+                # Check for omission in the annotation metadata
+                metadata = _extract_metadata(field_type)
+                for meta in metadata:
+                    if isinstance(meta, OmitFromSchema) and getattr(meta, omit_flag) is True:
+                        should_omit = True
+                        break
 
+            if not should_omit:
+                all_annotations[field_name] = field_type
 
-def _filter_state_for_schema(state: dict[str, Any], schema: type) -> dict[str, Any]:
-    """Filter state to only include fields defined in the given schema."""
-    if not hasattr(schema, "__annotations__"):
-        return state
+    return TypedDict(schema_name, all_annotations)  # type: ignore[operator]
 
-    schema_fields = set(schema.__annotations__.keys())
-    return {k: v for k, v in state.items() if k in schema_fields}
 
+def _extract_metadata(type_: type) -> list:
+    """Extract metadata from a field type, handling Required/NotRequired and Annotated wrappers."""
+    # Handle Required[Annotated[...]] or NotRequired[Annotated[...]]
+    if get_origin(type_) in (Required, NotRequired):
+        inner_type = get_args(type_)[0]
+        if get_origin(inner_type) is Annotated:
+            return list(get_args(inner_type)[1:])
 
-def _supports_native_structured_output(model: Union[str, BaseChatModel]) -> bool:
+    # Handle direct Annotated[...]
+    elif get_origin(type_) is Annotated:
+        return list(get_args(type_)[1:])
+
+    return []
+
+
+def _supports_native_structured_output(model: str | BaseChatModel) -> bool:
     """Check if a model supports native structured output."""
     model_name: str | None = None
     if isinstance(model, str):
@@ -113,7 +138,7 @@ def create_agent(  # noqa: PLR0915
     model: str | BaseChatModel,
     tools: Sequence[BaseTool | Callable | dict[str, Any]] | ToolNode | None = None,
     system_prompt: str | None = None,
-    middleware: Sequence[AgentMiddleware] = (),
+    middleware: Sequence[AgentMiddleware[AgentState[ResponseT], ContextT]] = (),
     response_format: ResponseFormat[ResponseT] | type[ResponseT] | None = None,
     context_schema: type[ContextT] | None = None,
 ) -> StateGraph[
@@ -198,46 +223,30 @@ def create_agent(  # noqa: PLR0915
         m for m in middleware if m.__class__.after_model is not AgentMiddleware.after_model
     ]
 
-    # Collect all middleware state schemas and create merged schema
-    merged_state_schema: type[AgentState] = _merge_state_schemas(
-        [m.state_schema for m in middleware]
-    )
+    state_schemas = {m.state_schema for m in middleware}
+    state_schemas.add(AgentState)
+
+    state_schema = _resolve_schema(state_schemas, "StateSchema", None)
+    input_schema = _resolve_schema(state_schemas, "InputSchema", "input")
+    output_schema = _resolve_schema(state_schemas, "OutputSchema", "output")
 
     # create graph, add nodes
-    graph = StateGraph(
-        merged_state_schema,
-        input_schema=PublicAgentState,
-        output_schema=PublicAgentState,
+    graph: StateGraph[
+        AgentState[ResponseT], ContextT, PublicAgentState[ResponseT], PublicAgentState[ResponseT]
+    ] = StateGraph(
+        state_schema=state_schema,
+        input_schema=input_schema,
+        output_schema=output_schema,
         context_schema=context_schema,
     )
 
-    def _prepare_model_request(state: dict[str, Any]) -> tuple[ModelRequest, list[AnyMessage]]:
-        """Prepare model request and messages."""
-        request = state.get("model_request") or ModelRequest(
-            model=model,
-            tools=default_tools,
-            system_prompt=system_prompt,
-            response_format=response_format,
-            messages=state["messages"],
-            tool_choice=None,
-        )
-
-        # prepare messages
-        messages = request.messages
-        if request.system_prompt:
-            messages = [SystemMessage(request.system_prompt), *messages]
-
-        return request, messages
-
-    def _handle_model_output(state: dict[str, Any], output: AIMessage) -> dict[str, Any]:
+    def _handle_model_output(output: AIMessage) -> dict[str, Any]:
         """Handle model output including structured responses."""
         # Handle structured output with native strategy
         if isinstance(response_format, ProviderStrategy):
             if not output.tool_calls and native_output_binding:
                 structured_response = native_output_binding.parse(output)
                 return {"messages": [output], "response": structured_response}
-            if state.get("response") is not None:
-                return {"messages": [output], "response": None}
             return {"messages": [output]}
 
         # Handle structured output with tools strategy
@@ -315,9 +324,6 @@ def create_agent(  # noqa: PLR0915
                         ],
                     }
 
-        # Standard response handling
-        if state.get("response") is not None:
-            return {"messages": [output], "response": None}
         return {"messages": [output]}
 
     def _get_bound_model(request: ModelRequest) -> Runnable:
@@ -340,37 +346,67 @@ def create_agent(  # noqa: PLR0915
             )
         return request.model.bind(**request.model_settings)
 
-    def model_request(state: dict[str, Any]) -> dict[str, Any]:
+    model_request_signatures: list[
+        tuple[bool, AgentMiddleware[AgentState[ResponseT], ContextT]]
+    ] = [
+        ("runtime" in signature(m.modify_model_request).parameters, m)
+        for m in middleware_w_modify_model_request
+    ]
+
+    def model_request(state: AgentState, runtime: Runtime[ContextT]) -> dict[str, Any]:
         """Sync model request handler with sequential middleware processing."""
-        # Start with the base model request
-        request, messages = _prepare_model_request(state)
+        request = ModelRequest(
+            model=model,
+            tools=default_tools,
+            system_prompt=system_prompt,
+            response_format=response_format,
+            messages=state["messages"],
+            tool_choice=None,
+        )
 
         # Apply modify_model_request middleware in sequence
-        for m in middleware_w_modify_model_request:
-            # Filter state to only include fields defined in this middleware's schema
-            filtered_state = _filter_state_for_schema(state, m.state_schema)
-            request = m.modify_model_request(request, filtered_state)
+        for use_runtime, m in model_request_signatures:
+            if use_runtime:
+                m.modify_model_request(request, state, runtime)
+            else:
+                m.modify_model_request(request, state)  # type: ignore[call-arg]
 
-        # Get the bound model with the final request
+        # Get the final model and messages
         model_ = _get_bound_model(request)
+        messages = request.messages
+        if request.system_prompt:
+            messages = [SystemMessage(request.system_prompt), *messages]
+
         output = model_.invoke(messages)
-        return _handle_model_output(state, output)
+        return _handle_model_output(output)
 
-    async def amodel_request(state: dict[str, Any]) -> dict[str, Any]:
+    async def amodel_request(state: AgentState, runtime: Runtime[ContextT]) -> dict[str, Any]:
         """Async model request handler with sequential middleware processing."""
         # Start with the base model request
-        request, messages = _prepare_model_request(state)
+        request = ModelRequest(
+            model=model,
+            tools=default_tools,
+            system_prompt=system_prompt,
+            response_format=response_format,
+            messages=state["messages"],
+            tool_choice=None,
+        )
 
         # Apply modify_model_request middleware in sequence
-        for m in middleware_w_modify_model_request:
-            # Filter state to only include fields defined in this middleware's schema
-            filtered_state = _filter_state_for_schema(state, m.state_schema)
-            request = m.modify_model_request(request, filtered_state)
+        for use_runtime, m in model_request_signatures:
+            if use_runtime:
+                m.modify_model_request(request, state, runtime)
+            else:
+                m.modify_model_request(request, state)  # type: ignore[call-arg]
 
-        # Get the bound model with the final request
+        # Get the final model and messages
         model_ = _get_bound_model(request)
+        messages = request.messages
+        if request.system_prompt:
+            messages = [SystemMessage(request.system_prompt), *messages]
+
         output = await model_.ainvoke(messages)
-        return _handle_model_output(state, output)
+        return _handle_model_output(output)
 
     # Use sync or async based on model capabilities
     from langgraph._internal._runnable import RunnableCallable
@@ -385,16 +421,12 @@ def create_agent(  # noqa: PLR0915
     for m in middleware:
         if m.__class__.before_model is not AgentMiddleware.before_model:
             graph.add_node(
-                f"{m.__class__.__name__}.before_model",
-                m.before_model,
-                input_schema=m.state_schema,
+                f"{m.__class__.__name__}.before_model", m.before_model, input_schema=state_schema
             )
 
         if m.__class__.after_model is not AgentMiddleware.after_model:
             graph.add_node(
-                f"{m.__class__.__name__}.after_model",
-                m.after_model,
-                input_schema=m.state_schema,
+                f"{m.__class__.__name__}.after_model", m.after_model, input_schema=state_schema
             )
 
     # add start edge
@@ -414,12 +446,12 @@ def create_agent(  # noqa: PLR0915
     if tool_node is not None:
         graph.add_conditional_edges(
             "tools",
-            _make_tools_to_model_edge(tool_node, first_node),
+            _make_tools_to_model_edge(tool_node, first_node, structured_output_tools),
             [first_node, END],
         )
         graph.add_conditional_edges(
             last_node,
-            _make_model_to_tools_edge(first_node, structured_output_tools),
+            _make_model_to_tools_edge(first_node, structured_output_tools, tool_node),
             [first_node, "tools", END],
         )
     elif last_node == "model_request":
@@ -478,27 +510,48 @@ def _resolve_jump(jump_to: JumpTo | None, first_node: str) -> str | None:
     return None
 
 
+def _fetch_last_ai_and_tool_messages(
+    messages: list[AnyMessage],
+) -> tuple[AIMessage, list[ToolMessage]]:
+    last_ai_index: int
+    last_ai_message: AIMessage
+
+    for i in range(len(messages) - 1, -1, -1):
+        if isinstance(messages[i], AIMessage):
+            last_ai_index = i
+            last_ai_message = cast("AIMessage", messages[i])
+            break
+
+    tool_messages = [m for m in messages[last_ai_index + 1 :] if isinstance(m, ToolMessage)]
+    return last_ai_message, tool_messages
+
+
 def _make_model_to_tools_edge(
-    first_node: str, structured_output_tools: dict[str, OutputToolBinding]
-) -> Callable[[AgentState], str | None]:
-    def model_to_tools(state: AgentState) -> str | None:
+    first_node: str, structured_output_tools: dict[str, OutputToolBinding], tool_node: ToolNode
+) -> Callable[[AgentState], str | list[Send] | None]:
+    def model_to_tools(state: AgentState) -> str | list[Send] | None:
         if jump_to := state.get("jump_to"):
             return _resolve_jump(jump_to, first_node)
 
-        message = state["messages"][-1]
-
-        # Check if this is a ToolMessage from structured output - if so, end
-        if isinstance(message, ToolMessage) and message.name in structured_output_tools:
-            return END
-
-        # Check for tool calls
-        if isinstance(message, AIMessage) and message.tool_calls:
-            # If all tool calls are for structured output, don't go to tools
-            non_structured_calls = [
-                tc for tc in message.tool_calls if tc["name"] not in structured_output_tools
+        last_ai_message, tool_messages = _fetch_last_ai_and_tool_messages(state["messages"])
+        tool_message_ids = [m.tool_call_id for m in tool_messages]
+
+        pending_tool_calls = [
+            c
+            for c in last_ai_message.tool_calls
+            if c["id"] not in tool_message_ids and c["name"] not in structured_output_tools
+        ]
+
+        if pending_tool_calls:
+            # imo we should not be injecting state, store here,
+            # this should be done by the tool node itself ideally but this is a consequence
+            # of using Send w/ tool calls directly which allows more intuitive interrupt behavior
+            # largely internal so can be fixed later
+            pending_tool_calls = [
+                tool_node.inject_tool_args(call, state, None)  # type: ignore[arg-type]
+                for call in pending_tool_calls
             ]
-            if non_structured_calls:
-                return "tools"
+            return [Send("tools", [tool_call]) for tool_call in pending_tool_calls]
 
         return END
 
@@ -506,17 +559,21 @@ def _make_model_to_tools_edge(
 
 
 def _make_tools_to_model_edge(
-    tool_node: ToolNode, next_node: str
+    tool_node: ToolNode, next_node: str, structured_output_tools: dict[str, OutputToolBinding]
 ) -> Callable[[AgentState], str | None]:
     def tools_to_model(state: AgentState) -> str | None:
-        ai_message = [m for m in state["messages"] if isinstance(m, AIMessage)][-1]
+        last_ai_message, tool_messages = _fetch_last_ai_and_tool_messages(state["messages"])
+
         if all(
             tool_node.tools_by_name[c["name"]].return_direct
-            for c in ai_message.tool_calls
+            for c in last_ai_message.tool_calls
             if c["name"] in tool_node.tools_by_name
         ):
             return END
 
+        if any(t.name in structured_output_tools for t in tool_messages):
+            return END
+
         return next_node
 
     return tools_to_model