Fix/filename

openrag/components/files.py

@@ -1,11 +1,35 @@
+import asyncio
+import re
+import secrets
+import time
 from pathlib import Path
+from typing import Dict, Optional
+
 import aiofiles
-from fastapi import UploadFile
 import consts
-import time
-import secrets
-import asyncio
-from typing import Dict, Optional
+from fastapi import UploadFile
+
+
+def sanitize_filename(filename: str) -> str:
+    # Split filename into name and extension
+    path = Path(filename)
+    name = path.stem
+    ext = path.suffix
+
+    # Remove special characters (keep only word characters and hyphens temporarily)
+    name = re.sub(r"[^\w\-]", "_", name)
+
+    # Replace hyphens with underscores
+    name = name.replace("-", "_")
+
+    # Collapse multiple underscores
+    name = re.sub(r"_+", "_", name)
+
+    # Remove leading/trailing underscores
+    name = name.strip("_")
+
+    # Reconstruct filename
+    return name + ext
 
 
 def make_unique_filename(filename: str) -> Path:
@@ -64,6 +88,5 @@ async def serialize_file(task_id: str, path: str, metadata: Optional[Dict] = {})
         except Exception:
             raise
     else:
-
         ray.cancel(future, recursive=True)
         raise TimeoutError(f"Serialization task {task_id} timed out after seconds")

openrag/components/test_files.py

@@ -1,11 +1,9 @@
 import io
-import uuid
 from pathlib import Path
 
-from fastapi import UploadFile
 import pytest
-
-from components.files import save_file_to_disk
+from components.files import sanitize_filename, save_file_to_disk
+from fastapi import UploadFile
 
 
 @pytest.mark.asyncio
@@ -18,11 +16,7 @@ async def test_save_file_to_disk_writes_content(tmp_path: Path):
 
     dest_dir = tmp_path / "uploads"
 
-    saved_path = await save_file_to_disk(
-        file=upload,
-        dest_dir=dest_dir,
-        chunk_size=4
-    )
+    saved_path = await save_file_to_disk(file=upload, dest_dir=dest_dir, chunk_size=4)
 
     assert saved_path.exists()
     assert saved_path.parent == dest_dir
@@ -34,10 +28,8 @@ async def test_save_file_to_disk_writes_content(tmp_path: Path):
     assert saved_content == content
 
 
-
 @pytest.mark.asyncio
 async def test_save_file_to_disk_with_random_prefix(tmp_path, monkeypatch):
-
     def fake_make_unique_filename(filename: str) -> str:
         assert filename == "test.txt"
         return "PREFIX_1234_test.txt"
@@ -63,3 +55,29 @@ def fake_make_unique_filename(filename: str) -> str:
     assert saved_path.name == "PREFIX_1234_test.txt"
     assert saved_path.exists()
     assert saved_path.read_bytes() == file_content
+
+
+@pytest.mark.parametrize(
+    "input_name,expected",
+    [
+        # Basic cases
+        ("simple_file.txt", "simple_file.txt"),
+        ("file-name.pdf", "file_name.pdf"),
+        # Spaces and commas
+        ("my file.txt", "my_file.txt"),
+        ("file,name.txt", "file_name.txt"),
+        ("multiple   spaces.txt", "multiple_spaces.txt"),
+        # Special characters
+        ("file@name#2024.txt", "file_name_2024.txt"),
+        ("doc$with%special&chars.pdf", "doc_with_special_chars.pdf"),
+        # Multiple underscores
+        ("file___name.txt", "file_name.txt"),
+        ("file__name__here.txt", "file_name_here.txt"),
+        # Edge cases
+        ("", ""),
+        ("file(1).txt", "file_1.txt"),
+        ("file.with.dot.txt", "file_with_dot.txt"),
+    ],
+)
+def test_sanitize_filename(input_name, expected):
+    assert sanitize_filename(input_name) == expected

openrag/routers/utils.py

@@ -4,6 +4,7 @@
 from typing import Any, Optional
 
 import consts
+from components.files import sanitize_filename
 from config import load_config
 from fastapi import Depends, Form, HTTPException, Request, UploadFile, status
 from openai import AsyncOpenAI
@@ -230,6 +231,8 @@ async def validate_file_format(
             status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE,
             detail=details,
         )
+
+    file.filename = sanitize_filename(file.filename)
     return file