Retry logic

[jonnyforsterMSG]

I've got a simple chain like this that takes an input from langserve, adds it to the prompt template and returns the response

chain = ( RunnableParallel({"transcript": RunnablePassthrough()}) | prompt | model_with_langfuse | JsonOutputParser() )

answer = chain.with_types(input_type=Transcript)

The model has the default retries but I've changed max_retries to a different value and can see the effect working

model = AzureChatOpenAI( azure_deployment=os.getenv("AZURE_OPENAI_DEPLOYMENT"), temperature=0.7, max_tokens=800, top_p=0.95, frequency_penalty=0, presence_penalty=0, stop=None )

I want to add some retry logic for specific error codes so tried adding this

from openai import RateLimitError, InternalServerError

chain = ( RunnableParallel({"transcript": RunnablePassthrough()}) | prompt | model_with_langfuse | JsonOutputParser() ).with_retry(retry_if_exception_type=(RateLimitError, InternalServerError), stop_after_attempt=10)

But it doesn't seem to have any effect. When I point at a model with no quota so only returns 429s the only thing that alters the retry logic is max_retries on the model itself. I've tried adding with_retry on the chain and model but neither seem to do anything.

Any help much appreciated