test_tool_retry_custom_exception_filter() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_tool_retry.py lines 480–543

def test_tool_retry_custom_exception_filter() -> None:
    """Test ToolRetryMiddlewarewith custom exception filter function."""

    class CustomError(Exception):
        """Custom exception with retry_me attribute."""

        def __init__(self, message: str, *, retry_me: bool):
            """Initialize custom error.

            Args:
                message: Error message.
                retry_me: Whether this error should be retried.
            """
            super().__init__(message)
            self.retry_me = retry_me

    attempt_count = {"value": 0}

    @tool
    def custom_error_tool(val: str) -> str:
        """Tool that raises CustomError."""
        attempt_count["value"] += 1
        if attempt_count["value"] == 1:
            msg = "Retryable error"
            raise CustomError(msg, retry_me=True)
        msg = "Non-retryable error"
        raise CustomError(msg, retry_me=False)

    def should_retry(exc: Exception) -> bool:
        return isinstance(exc, CustomError) and exc.retry_me

    model = FakeToolCallingModel(
        tool_calls=[
            [ToolCall(name="custom_error_tool", args={"val": "test"}, id="1")],
            [],
        ]
    )

    retry = ToolRetryMiddleware(
        max_retries=3,
        retry_on=should_retry,
        initial_delay=0.01,
        jitter=False,
        on_failure="continue",
    )

    agent = create_agent(
        model=model,
        tools=[custom_error_tool],
        middleware=[retry],
        checkpointer=InMemorySaver(),
    )

    result = agent.invoke(
        {"messages": [HumanMessage("Use custom error tool")]},
        {"configurable": {"thread_id": "test"}},
    )

    tool_messages = [m for m in result["messages"] if isinstance(m, ToolMessage)]
    assert len(tool_messages) == 1

    # Should retry once (attempt 1 with retry_me=True), then fail on attempt 2 (retry_me=False)
    assert attempt_count["value"] == 2
    assert "2 attempts" in tool_messages[0].content