Repository: run-llama/llama_index
Pull request: `StorageContext.persist` supporting `Path`
# Description

Closes https://github.com/jerryjliu/llama_index/issues/7777.
- Looks like https://github.com/jerryjliu/llama_index/pull/3778 added support for a new control flow, that could have been just one, so I moved to just one control flow
- Improved `concat_dirs` to have better arg names and cleaner docstring

## Type of Change

Please delete options that are not relevant.

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] This change requires a documentation update

# How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

- [ ] Added new unit/integration tests
- [ ] Added new notebook (that tests end-to-end)
- [x] I stared at the code and made sure it makes sense


Problem statement:
[Bug]: typing errors in `OllamaEmbeddings` and `StorageContext.persist`
### Bug Description

Running a line like this through `mypy`:

```python
import pathlib

from langchain.embeddings import OllamaEmbeddings
from llama_index import StorageContext

STORAGE_FOLDER = pathlib.Path(__file__).parent / "storage"

_ = OllamaEmbeddings(model="llama2:13b")

StorageContext.from_defaults().persist(persist_dir=STORAGE_FOLDER)
```

### Version

0.8.30

### Steps to Reproduce

Running `mypy==1.5.1` on that file:

```none
a.py:8:5: error: Missing named argument "mirostat" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "mirostat_eta" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "mirostat_tau" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "num_ctx" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "num_gpu" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "num_thread" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "repeat_last_n" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "repeat_penalty" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "temperature" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "stop" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "tfs_z" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "top_k" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:8:5: error: Missing named argument "top_p" for "OllamaEmbeddings"  [call-arg]
    _ = OllamaEmbeddings(model="llama2:13b")
        ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
a.py:10:52: error: Argument "persist_dir" to "persist" of "StorageContext" has incompatible type "Path"; expected "str"  [arg-type]
    StorageContext.from_defaults().persist(persist_dir=STORAGE_FOLDER)
```

These imo are all type errors:
- In `OllamaEmbeddings`, `= None` needs to be present as a default for the `Optional` fields
- In `StorageContext.persist`, it should allow `os.PathLike`

Hints from issue discussion:
Actually, my apologies I realized `OllamaEmbeddings` stuff should be for LangChain (so I made https://github.com/langchain-ai/langchain/issues/10919).

However, type error for `StorageContext.persist` can be fixed here.  Signature can be `persist_dir: str | os.PathLike`
@jamesbraza we use `mypy==0.991` in our CI. So, these type errors are not an issue for now. (We also support python3.8+, so the `|` operator can't be used for typing, need to use a union)

Feel free to submit a PR if you feel this should be fixed, but going to close this for now.
Out of curiosity, why use a pre-stable release `mypy`?  It's much better these days.

Yeah I will open a PR shortly, thanks for `Union` heads up

Commit message: Simplified control flow in persist, and supporting Path

Patch to review:
diff --git a/llama_index/storage/storage_context.py b/llama_index/storage/storage_context.py
index c07b212e8b..78f99b4f12 100644
--- a/llama_index/storage/storage_context.py
+++ b/llama_index/storage/storage_context.py
@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Union
+import os
 
 import fsspec
 
@@ -89,7 +90,7 @@ def from_defaults(
 
     def persist(
         self,
-        persist_dir: str = DEFAULT_PERSIST_DIR,
+        persist_dir: Union[str, os.PathLike] = DEFAULT_PERSIST_DIR,
         docstore_fname: str = DOCSTORE_FNAME,
         index_store_fname: str = INDEX_STORE_FNAME,
         vector_store_fname: str = VECTOR_STORE_FNAME,
@@ -100,18 +101,12 @@ def persist(
 
         Args:
             persist_dir (str): directory to persist the storage context
-
         """
-        if fs is not None:
-            docstore_path = concat_dirs(persist_dir, docstore_fname)
-            index_store_path = concat_dirs(persist_dir, index_store_fname)
-            vector_store_path = concat_dirs(persist_dir, vector_store_fname)
-            graph_store_path = concat_dirs(persist_dir, graph_store_fname)
-        else:
-            docstore_path = str(Path(persist_dir) / docstore_fname)
-            index_store_path = str(Path(persist_dir) / index_store_fname)
-            vector_store_path = str(Path(persist_dir) / vector_store_fname)
-            graph_store_path = str(Path(persist_dir) / graph_store_fname)
+        persist_dir = str(persist_dir)
+        docstore_path = concat_dirs(persist_dir, docstore_fname)
+        index_store_path = concat_dirs(persist_dir, index_store_fname)
+        vector_store_path = concat_dirs(persist_dir, vector_store_fname)
+        graph_store_path = concat_dirs(persist_dir, graph_store_fname)
 
         self.docstore.persist(persist_path=docstore_path, fs=fs)
         self.index_store.persist(persist_path=index_store_path, fs=fs)
diff --git a/llama_index/utils.py b/llama_index/utils.py
index 52059d3813..6f77c4a5f3 100644
--- a/llama_index/utils.py
+++ b/llama_index/utils.py
@@ -205,14 +205,15 @@ def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
         yield b
 
 
-def concat_dirs(dir1: str, dir2: str) -> str:
+def concat_dirs(dirname: str, basename: str) -> str:
     """
-    Concat dir1 and dir2 while avoiding backslashes when running on windows.
-    os.path.join(dir1,dir2) will add a backslash before dir2 if dir1 does not
-    end with a slash, so we make sure it does.
+    Append basename to dirname, avoiding backslashes when running on windows.
+
+    os.path.join(dirname, basename) will add a backslash before dirname if
+    basename does not end with a slash, so we make sure it does.
     """
-    dir1 += "/" if dir1[-1] != "/" else ""
-    return os.path.join(dir1, dir2)
+    dirname += "/" if dirname[-1] != "/" else ""
+    return os.path.join(dirname, basename)
 
 
 def get_tqdm_iterable(items: Iterable, show_progress: bool, desc: str) -> Iterable:


Code review comments:
diff_hunk: @@ -100,18 +101,12 @@ def persist(
 
         Args:
             persist_dir (str): directory to persist the storage context
-
         """
-        if fs is not None:
-            docstore_path = concat_dirs(persist_dir, docstore_fname)
-            index_store_path = concat_dirs(persist_dir, index_store_fname)
-            vector_store_path = concat_dirs(persist_dir, vector_store_fname)
-            graph_store_path = concat_dirs(persist_dir, graph_store_fname)
-        else:
-            docstore_path = str(Path(persist_dir) / docstore_fname)
-            index_store_path = str(Path(persist_dir) / index_store_fname)
-            vector_store_path = str(Path(persist_dir) / vector_store_fname)
-            graph_store_path = str(Path(persist_dir) / graph_store_fname)
path: llama_index/storage/storage_context.py
text: @user1:
I think this breaks windows support?

The original change was trying to fix using `Path` + FSSpec on windows machines. Path was adding the wrong kind of separator.

This change feels like it will have the same issue?

@author:
You're right!  Excellent catch, I will fix this

@author:
Pushed the fix now