diff --git a/.changeset/llamacloud-local-file-name-extension.md b/.changeset/llamacloud-local-file-name-extension.md new file mode 100644 index 000000000..09cc3db2b --- /dev/null +++ b/.changeset/llamacloud-local-file-name-extension.md @@ -0,0 +1,10 @@ +--- +"@create-llama/llama-index-server": patch +--- + +fix: only strip the trailing extension when building the local LlamaCloud file name + +`get_local_file_name` used `str.replace(file_ext, "")`, which removes every +occurrence of the extension string. For files whose stem also ends with the +extension (e.g. `report.pdf.pdf`) this corrupted the stem (`report` instead of +`report.pdf`). It now uses the root returned by `os.path.splitext`. diff --git a/python/llama-index-server/llama_index/server/utils/llamacloud.py b/python/llama-index-server/llama_index/server/utils/llamacloud.py index 42b37a870..842eaf407 100644 --- a/python/llama-index-server/llama_index/server/utils/llamacloud.py +++ b/python/llama-index-server/llama_index/server/utils/llamacloud.py @@ -22,8 +22,10 @@ def get_local_file_name( if llamacloud_file_name is None or pipeline_id is None: raise ValueError("Couldn't find llamacloud_file_name and pipeline_id") # Construct the local file name - file_ext = os.path.splitext(llamacloud_file_name)[1] - file_name = llamacloud_file_name.replace(file_ext, "") + file_root, file_ext = os.path.splitext(llamacloud_file_name) + # Only strip the trailing extension; using str.replace would remove every + # occurrence of the extension (e.g. "report.pdf.pdf" -> "report"). + file_name = file_root sanitized_file_name = re.sub(r"[^A-Za-z0-9_\-]", "_", file_name) return f"{sanitized_file_name}_{pipeline_id}{file_ext}" diff --git a/python/llama-index-server/tests/utils/test_llamacloud.py b/python/llama-index-server/tests/utils/test_llamacloud.py new file mode 100644 index 000000000..7cf2ccb69 --- /dev/null +++ b/python/llama-index-server/tests/utils/test_llamacloud.py @@ -0,0 +1,37 @@ +from llama_index.server.utils.llamacloud import ( + get_local_file_name, + is_llamacloud_file, +) + + +def test_get_local_file_name_single_extension() -> None: + assert ( + get_local_file_name(llamacloud_file_name="report.pdf", pipeline_id="123") + == "report_123.pdf" + ) + + +def test_get_local_file_name_keeps_repeated_extension_in_stem() -> None: + # The stem also ends with the extension string. Only the trailing + # extension must be stripped, otherwise the stem is corrupted. + assert ( + get_local_file_name(llamacloud_file_name="report.pdf.pdf", pipeline_id="123") + == "report_pdf_123.pdf" + ) + + +def test_get_local_file_name_extension_substring_in_stem() -> None: + assert ( + get_local_file_name(llamacloud_file_name="backup.tar.tar", pipeline_id="abc") + == "backup_tar_abc.tar" + ) + + +def test_get_local_file_name_from_metadata() -> None: + metadata = {"file_name": "my.report.pdf.pdf", "pipeline_id": "p1"} + assert get_local_file_name(node_metadata=metadata) == "my_report_pdf_p1.pdf" + + +def test_is_llamacloud_file() -> None: + assert is_llamacloud_file({"pipeline_id": "p1"}) is True + assert is_llamacloud_file({"file_name": "a.pdf"}) is False