Kimmy7/chat_over_docs

1from langchain_core.callbacks.base import BaseCallbackHandler2from langchain.agents.agent import AgentAction3from typing import Any, Union4 5 6class RetryCallbackHandler(BaseCallbackHandler):7    def on_tool_error(8        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any9    ) -> Any:10        # Log the error or perform any necessary actions11 12        # Retry logic13        max_retries = 114        current_retry = kwargs.get("retry_count", 0)15        if current_retry < max_retries:16            print(f"Retrying tool execution (Attempt {current_retry + 1})")17            # Increment retry count18            kwargs["retry_count"] = current_retry + 119            # Re-run the tool20            return AgentAction.RERUN21        else:22            print("Maximum retries reached. Switching to another tool.")23            return AgentAction.CHANGE_TOOL24