Skip to content

Commit

Permalink
Respect user preferred dir and allow to configure logs dir
Browse files Browse the repository at this point in the history
  • Loading branch information
dlqqq committed Dec 19, 2023
1 parent 5e4b660 commit 86bc91e
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 7 deletions.
27 changes: 22 additions & 5 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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

Expand All @@ -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
16 changes: 14 additions & 2 deletions packages/jupyter-ai/jupyter_ai/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 86bc91e

Please sign in to comment.