diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py index 7788a8965..1f463dd9e 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py @@ -221,8 +221,14 @@ class GenerateChatHandler(BaseChatHandler): help = "Generate a Jupyter notebook from a text prompt" routing_type = SlashCommandRoutingType(slash_id="generate") - def __init__(self, *args, **kwargs): + def __init__(self, preferred_dir: str, log_dir: Optional[str], *args, **kwargs): super().__init__(*args, **kwargs) + self.log_dir = Path(log_dir) if log_dir else None + self.preferred_dir = ( + os.path.abspath(os.path.expanduser(preferred_dir)) + if preferred_dir != "" + else None + ) self.llm = None def create_llm_chain( @@ -251,7 +257,7 @@ async def _generate_notebook(self, prompt: str): # create and write the notebook to disk notebook = create_notebook(outline) - final_path = os.path.join(self.root_dir, outline["title"] + ".ipynb") + final_path = os.path.join(self._output_dir, outline["title"] + ".ipynb") nbformat.write(notebook, final_path) return final_path @@ -267,11 +273,22 @@ async def process_message(self, message: HumanChatMessage): self.reply(response, message) async def handle_exc(self, e: Exception, message: HumanChatMessage): - timestamp = time.strftime("%Y-%m-%d-%H:%M:%S") - log_path = Path(f"jupyter-ai-logs/generate-{timestamp}.log") - log_path.parent.mkdir(parents=True, exist_ok=True) + timestamp = time.strftime("%Y-%m-%d-%H.%M.%S") + default_log_dir = Path(self._output_dir) / "jupyter-ai-logs" + log_dir = self.log_dir or default_log_dir + log_dir.mkdir(parents=True, exist_ok=True) + log_path = log_dir / f"generate-{timestamp}.log" with log_path.open("w") as log: traceback.print_exc(file=log) response = f"An error occurred while generating the notebook. The error details have been saved to `./{log_path}`.\n\nTry running `/generate` again, as some language models require multiple attempts before a notebook is generated." self.reply(response, message) + + @property + def _output_dir(self): + # preferred dir is preferred, but if it is not specified, + # or if user removed it after startup, fallback to root. + if self.preferred_dir and os.path.exists(self.preferred_dir): + return self.preferred_dir + else: + return self.root_dir diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index 99ae63c99..a41c556c7 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -95,6 +95,15 @@ class AiExtension(ExtensionApp): config=True, ) + error_logs_dir = Unicode( + default_value=None, + help="""Path to a directory where the error logs should be + written to. Defaults to `jupyter-ai-logs/` in the preferred dir + (if defined) or in root dir otherwise.""", + allow_none=True, + config=True, + ) + def initialize_settings(self): start = time.time() @@ -172,10 +181,13 @@ def initialize_settings(self): "dask_client_future": dask_client_future, "model_parameters": self.settings["model_parameters"], } - default_chat_handler = DefaultChatHandler(**chat_handler_kwargs) clear_chat_handler = ClearChatHandler(**chat_handler_kwargs) - generate_chat_handler = GenerateChatHandler(**chat_handler_kwargs) + generate_chat_handler = GenerateChatHandler( + **chat_handler_kwargs, + preferred_dir=self.serverapp.contents_manager.preferred_dir, + log_dir=self.error_logs_dir, + ) learn_chat_handler = LearnChatHandler(**chat_handler_kwargs) retriever = Retriever(learn_chat_handler=learn_chat_handler) ask_chat_handler = AskChatHandler(**chat_handler_kwargs, retriever=retriever)