diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py index bbfdc702b..b43ebdfb0 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py @@ -59,8 +59,6 @@ class BaseChatHandler: routing_type: HandlerRoutingType = ... - API_KEY_EXCEPTIONS = (OpenAIAuthenticationError,) - def __init__( self, log: Logger, @@ -117,25 +115,32 @@ async def handle_exc(self, e: Exception, message: HumanChatMessage): implementation is provided, however chat handlers (subclasses) should implement this method to provide a more helpful error response. """ - if self.is_api_key_error(e): - self.handle_api_key_error(e, message) + if self.is_api_key_exc(e): + self.handle_api_key_exc(e, message) else: await self._default_handle_exc(e, message) - def is_api_key_error(self, e: Exception): + def get_api_key_excs(self): + """ + Rerurns a list of API key exceptions based on models that are supported by default. + """ + return (OpenAIAuthenticationError,) + + def is_api_key_exc(self, e: Exception): """ - Checks if the exception is a known API key error. + Checks if the exception is an expected API key exception. """ - return isinstance(e, self.API_KEY_EXCEPTIONS) + api_excs = self.get_api_key_excs() + return isinstance(e, api_excs) - def handle_api_key_error(self, e: Exception, message: HumanChatMessage): + def handle_api_key_exc(self, e: Exception, message: HumanChatMessage): provider_name = "" if hasattr(self.config_manager, "lm_provider") and hasattr( self.config_manager.lm_provider, "name" ): name = getattr(self.config_manager.lm_provider, "name", "") provider_name = f" {name}" if name else "" - response = f"Oops! There's a problem with your {provider_name} API key. Please update your{provider_name} API key in the chat settings." + response = f"Oops! There's a problem with your{provider_name} API key. Please update your{provider_name} API key in the chat settings." self.reply(response, message) async def _default_handle_exc(self, e: Exception, message: HumanChatMessage):