Why is the on_retry method of my callback not called?

[PepijnBoers]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.runnables.retry import RunnableRetry


class MyCustomHandler(BaseCallbackHandler):
    def on_llm_new_token(self, token: str, **kwargs) -> None:
        print(f"My custom handler, token: {token}")

    def on_retry(self, *args, **kwargs):
        print("on_retry")

    def on_llm_start(self, *args, **kwargs):
        print("on_llm_start")


class SimulatedRetryableError(Exception):
    pass


class FailingModel(LLM):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def _call(self, *args, **kwargs):
        raise SimulatedRetryableError("Simulated failure")
        return super().invoke(*args, **kwargs)

    def _llm_type(self, *args, **kwargs):
        return "custom"

prompt = ChatPromptTemplate.from_messages(["Tell me a joke about {animal}"])

model = FailingModel(n=1, callbacks=[MyCustomHandler()])

chain = prompt | model

runnable_with_retries = RunnableRetry(
    bound=chain,
    config={"callbacks": [MyCustomHandler()]},
    max_attempt_number=5,
    wait_exponential_jitter=True
)

response = runnable_with_retries.invoke({"animal": "bears"})

Description

I'm expecting to see an on_retry statement, but I do not manage to get it. When is the on_retry method triggered?

System Info

langchain==0.2.10
langchain-core==0.2.22
langchain-openai==0.1.17
langchain-text-splitters==0.2.2
MacOS 14.2.1
Python 3.9.6

[dosubot[bot]]

@PepijnBoers Hello! I'm here to assist you with any bugs, questions, or contributions you have. Let's take a look at the issue you've encountered.

The on_retry method in your callback might not be called due to a few potential reasons:

1. Exception Type: Ensure that the exception being raised (SimulatedRetryableError) is included in the retry_exception_types for RunnableRetry. By default, RunnableRetry retries on all exceptions, but it's good to verify this.

2. Retry Configuration: Verify that the retry configuration is correctly set up. The max_attempt_number should be greater than 1, which you have set to 5, so this should be fine.

3. Callback Registration: Ensure that the MyCustomHandler is correctly registered in the RunnableRetry configuration.

Here is a revised version of your code with some adjustments to ensure the on_retry method is called:

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.runnables.retry import RunnableRetry
from langchain_core.llms import LLM
from tenacity import retry_if_exception_type, stop_after_attempt


class MyCustomHandler(BaseCallbackHandler):
    def on_llm_new_token(self, token: str, **kwargs) -> None:
        print(f"My custom handler, token: {token}")

    def on_retry(self, *args, **kwargs):
        print("on_retry")

    def on_llm_start(self, *args, **kwargs):
        print("on_llm_start")


class SimulatedRetryableError(Exception):
    pass


class FailingModel(LLM):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def _call(self, *args, **kwargs):
        raise SimulatedRetryableError("Simulated failure")
        return super().invoke(*args, **kwargs)

    def _llm_type(self, *args, **kwargs):
        return "custom"

prompt = ChatPromptTemplate.from_messages(["Tell me a joke about {animal}"])

model = FailingModel(n=1, callbacks=[MyCustomHandler()])

chain = prompt | model

runnable_with_retries = RunnableRetry(
    bound=chain,
    config={"callbacks": [MyCustomHandler()]},
    max_attempt_number=5,
    retry_exception_types=(SimulatedRetryableError,),  # Ensure the correct exception type
    wait_exponential_jitter=True
)

response = runnable_with_retries.invoke({"animal": "bears"})

Make sure that the retry_exception_types includes SimulatedRetryableError to ensure that the retry mechanism is triggered for this specific exception. This should help in calling the on_retry method as expected ^[1].

To continue talking to Dosu, mention @dosu.

[PepijnBoers]

Only on_llm_start is called.

[lei-lqc]

the same with my code, only on_llm_start and on_chain_start are called, but not on_retry()

[eyurtsev]

Could you check what happens in this case:

response = runnable_with_retries.invoke({"animal": "bears"}, {'callbacks': [MyCustomHandler()]})

[PepijnBoers]

The behaviour is the same, i.e. only the on_llm_start is printed.

[eyurtsev]

I'll create an issue to investigate -- I don't see anything obvious

[PepijnBoers]

@eyurtsev Great, thanks!

[lei-lqc]

I'll create an issue to investigate -- I don't see anything obvious

Hello @eyurtsev, may I have the issue id?

[PepijnBoers]

@eyurtsev Any updates on this? If not, I would be helped tremendously by seeing a working example of the on_retry method. Do you know of a source that has this functionality implemented? Thank you.