test_selective_error_handling() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 412–443

    def test_selective_error_handling(self) -> None:
        """Test middleware that only handles specific errors."""

        class SpecificErrorModel(GenericFakeChatModel):
            @override
            def _generate(
                self,
                messages: list[BaseMessage],
                stop: list[str] | None = None,
                run_manager: CallbackManagerForLLMRun | None = None,
                **kwargs: Any,
            ) -> ChatResult:
                msg = "Network error"
                raise ConnectionError(msg)

        class SelectiveErrorMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                try:
                    return handler(request)
                except ConnectionError:
                    return AIMessage(content="Network issue, try again later")

        model = SpecificErrorModel(messages=iter([]))
        agent = create_agent(model=model, middleware=[SelectiveErrorMiddleware()])

        result = agent.invoke({"messages": [HumanMessage("Test")]})

        assert result["messages"][1].content == "Network issue, try again later"