diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py index 2d09aa9a5..8d6fb09aa 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py @@ -19,6 +19,7 @@ IndexMetadata, ) from jupyter_core.paths import jupyter_data_dir +from jupyter_core.utils import ensure_dir_exists from langchain.schema import BaseRetriever, Document from langchain.text_splitter import ( LatexTextSplitter, @@ -97,11 +98,12 @@ def __init__(self, *args, **kwargs): self.metadata = IndexMetadata(dirs=[]) self.prev_em_id = None - if not os.path.exists(INDEX_SAVE_DIR): - os.makedirs(INDEX_SAVE_DIR) - + self._ensure_dirs() self._load() + def _ensure_dirs(self): + ensure_dir_exists(INDEX_SAVE_DIR, mode=0o700) + def _load(self): """Loads the vector store.""" if self.index is not None: diff --git a/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py b/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py index 0194dec73..496cde371 100644 --- a/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py +++ b/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py @@ -1,3 +1,24 @@ +import os +import stat +from unittest import mock + +from jupyter_ai.chat_handlers import learn + + +class MockLearnHandler(learn.LearnChatHandler): + def __init__(self): + pass + + +def test_learn_index_permissions(tmp_path): + test_dir = tmp_path / "test" + with mock.patch.object(learn, "INDEX_SAVE_DIR", new=test_dir): + handler = MockLearnHandler() + handler._ensure_dirs() + mode = os.stat(test_dir).st_mode + assert stat.filemode(mode) == "drwx------" + + # TODO # import json