after_model() — langchain Function Reference

Architecture documentation for the after_model() function in pii.py from the langchain codebase.

Function python LangChainCore LanguageModelBase calls 1 called by 1

Entity Profile

LangChainCore→ LanguageModelBase→ after_model() — langchain Function Reference

Dependency Diagram

graph TD
  27d4728e_2258_70e9_e8c8_f68e263ca8fb["after_model()"]
  dc527b75_da40_4d1c_e472_82d1252ba70b["PIIMiddleware"]
  27d4728e_2258_70e9_e8c8_f68e263ca8fb -->|defined in| dc527b75_da40_4d1c_e472_82d1252ba70b
  074d1c14_b6ba_2cd0_af55_0985e89d7adf["aafter_model()"]
  074d1c14_b6ba_2cd0_af55_0985e89d7adf -->|calls| 27d4728e_2258_70e9_e8c8_f68e263ca8fb
  ffe55823_7c55_8b73_2be6_a2b281bb64cf["_process_content()"]
  27d4728e_2258_70e9_e8c8_f68e263ca8fb -->|calls| ffe55823_7c55_8b73_2be6_a2b281bb64cf
  style 27d4728e_2258_70e9_e8c8_f68e263ca8fb fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

libs/langchain_v1/langchain/agents/middleware/pii.py lines 287–344

    def after_model(
        self,
        state: AgentState[Any],
        runtime: Runtime[ContextT],
    ) -> dict[str, Any] | None:
        """Check AI messages for PII after model invocation.

        Args:
            state: The current agent state.
            runtime: The langgraph runtime.

        Returns:
            Updated state with PII handled according to strategy, or None if no PII
                detected.

        Raises:
            PIIDetectionError: If PII is detected and strategy is `'block'`.
        """
        if not self.apply_to_output:
            return None

        messages = state["messages"]
        if not messages:
            return None

        # Get last AI message
        last_ai_msg = None
        last_ai_idx = None
        for i in range(len(messages) - 1, -1, -1):
            msg = messages[i]
            if isinstance(msg, AIMessage):
                last_ai_msg = msg
                last_ai_idx = i
                break

        if last_ai_idx is None or not last_ai_msg or not last_ai_msg.content:
            return None

        # Detect PII in message content
        content = str(last_ai_msg.content)
        new_content, matches = self._process_content(content)

        if not matches:
            return None

        # Create updated message
        updated_message = AIMessage(
            content=new_content,
            id=last_ai_msg.id,
            name=last_ai_msg.name,
            tool_calls=last_ai_msg.tool_calls,
        )

        # Return updated messages
        new_messages = list(messages)
        new_messages[last_ai_idx] = updated_message

        return {"messages": new_messages}