From 78b248cc61e608f5ae9cb45e79f5bb255925d63a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Krassowski?=
 <5832902+krassowski@users.noreply.github.com>
Date: Fri, 15 Dec 2023 19:05:45 +0000
Subject: [PATCH] Respect user preferred dir and allow to configure logs dir
 (#490)

---
 .../jupyter_ai/chat_handlers/generate.py      | 25 ++++++++++++++++---
 packages/jupyter-ai/jupyter_ai/extension.py   | 16 ++++++++++--
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py
index b1ccb3b6e..a17f51390 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 = "Generates a Jupyter notebook, including name, outline, and section contents"
     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
 
@@ -268,10 +274,21 @@ async def process_message(self, message: HumanChatMessage):
 
     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)
+        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 ec08d9962..c6ce4f5d9 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)
         help_chat_handler = HelpChatHandler(**chat_handler_kwargs)
         retriever = Retriever(learn_chat_handler=learn_chat_handler)