Make preferred-dir the default read/write directory for slash commands

packages/jupyter-ai/jupyter_ai/chat_handlers/base.py

@@ -33,6 +33,15 @@
     from jupyter_ai.handlers import RootChatHandler
 
 
+def get_preferred_dir(root_dir: str, preferred_dir: str) -> str | None:
+    if preferred_dir != "":
+        preferred_dir = os.path.expanduser(preferred_dir)
+        if not preferred_dir.startswith(root_dir):
+            preferred_dir = os.path.join(root_dir, preferred_dir)
+        return os.path.abspath(preferred_dir)
+    return None
+
+
 # Chat handler type, with specific attributes for each
 class HandlerRoutingType(BaseModel):
     routing_method: ClassVar[Union[Literal["slash_command"]]] = ...
@@ -82,6 +91,7 @@ def __init__(
         model_parameters: Dict[str, Dict],
         chat_history: List[ChatMessage],
         root_dir: str,
+        preferred_dir: Optional[str],
         dask_client_future: Awaitable[DaskClient],
     ):
         self.log = log
@@ -91,6 +101,7 @@ def __init__(
         self._chat_history = chat_history
         self.parser = argparse.ArgumentParser()
         self.root_dir = os.path.abspath(os.path.expanduser(root_dir))
+        self.preferred_dir = get_preferred_dir(self.root_dir, preferred_dir)
         self.dask_client_future = dask_client_future
         self.llm = None
         self.llm_params = None
@@ -302,3 +313,12 @@ def parse_args(self, message):
             self.reply(response, message)
             return None
         return args
+
+    @property
+    def output_dir(self) -> str:
+        # 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

packages/jupyter-ai/jupyter_ai/chat_handlers/export.py

@@ -42,7 +42,7 @@ async def process_message(self, message: HumanChatMessage):
             else f"chat_history-{datetime.now():%Y-%m-%d-%H-%M-%S}.md"
         )  # Handles both empty args and double tap <Enter> key
         chat_file = os.path.join(
-            self.root_dir, chat_filename
+            self.output_dir, chat_filename
         )  # Do not use timestamp if filename is entered as argument
         with open(chat_file, "w") as chat_history:
             chat_history.write(markdown_content)

packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py

@@ -223,14 +223,9 @@ class GenerateChatHandler(BaseChatHandler):
 
     uses_llm = True
 
-    def __init__(self, preferred_dir: str, log_dir: Optional[str], *args, **kwargs):
+    def __init__(self, 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(
@@ -262,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._output_dir, outline["title"] + ".ipynb")
+        final_path = os.path.join(self.output_dir, outline["title"] + ".ipynb")
         nbformat.write(notebook, final_path)
         return final_path
 
@@ -279,7 +274,7 @@ 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")
-        default_log_dir = Path(self._output_dir) / "jupyter-ai-logs"
+        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"
@@ -288,12 +283,3 @@ async def handle_exc(self, e: Exception, message: HumanChatMessage):
 
         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

packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py

@@ -118,7 +118,7 @@ async def process_message(self, message: HumanChatMessage):
             if remote_type == "arxiv":
                 try:
                     id = args.path[0]
-                    args.path = [arxiv_to_text(id, self.root_dir)]
+                    args.path = [arxiv_to_text(id, self.output_dir)]
                     self.reply(
                         f"Learning arxiv file with id **{id}**, saved in **{args.path[0]}**.",
                         message,
@@ -142,7 +142,7 @@ async def process_message(self, message: HumanChatMessage):
             self.reply(f"{self.parser.format_usage()}", message)
             return
         short_path = args.path[0]
-        load_path = os.path.join(self.root_dir, short_path)
+        load_path = os.path.join(self.output_dir, short_path)
         if not os.path.exists(load_path):
             response = f"Sorry, that path doesn't exist: {load_path}"
             self.reply(response, message)

packages/jupyter-ai/jupyter_ai/extension.py

@@ -255,12 +255,12 @@ def initialize_settings(self):
             "root_dir": self.serverapp.root_dir,
             "dask_client_future": self.settings["dask_client_future"],
             "model_parameters": self.settings["model_parameters"],
+            "preferred_dir": self.serverapp.contents_manager.preferred_dir,
         }
         default_chat_handler = DefaultChatHandler(**chat_handler_kwargs)
         clear_chat_handler = ClearChatHandler(**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)