should work?

Author: fabnemEPFL

src/mmore/run_index_api.py

@@ -34,6 +34,19 @@
 logger = logging.getLogger(__name__)
 
 
+def _apply_uploaded_file_metadata(documents, file_id: str, filename: str) -> None:
+    """Bind processed chunks to the API file ID and persist the original filename."""
+    for doc in documents:
+        default_doc_id = doc.document_id
+        doc.document_id = file_id
+        if default_doc_id and doc.id.startswith(default_doc_id):
+            doc.id = f"{file_id}{doc.id[len(default_doc_id) :]}"
+        else:
+            doc.id = file_id
+
+        doc.metadata.extra["filename"] = filename
+
+
 def make_router(config_path: str) -> APIRouter:
     router = APIRouter()
 
@@ -97,10 +110,7 @@ async def upload_file(
                     temp_dir, COLLECTION_NAME, [file_extension]
                 )
 
-                for doc in documents:
-                    defDocId = doc.document_id
-                    doc.document_id = fileId
-                    doc.id = doc.id.replace(defDocId, fileId)
+                _apply_uploaded_file_metadata(documents, fileId, file.filename)
 
                 # Get indexer and index the document
                 try:
@@ -257,9 +267,8 @@ async def update_file(
                     temp_dir, COLLECTION_NAME, [file_extension]
                 )
 
-                # Set the custom ID
-                for doc in documents:
-                    doc.id = fileId
+                # Set the custom ID and preserve the original upload filename
+                _apply_uploaded_file_metadata(documents, fileId, file.filename)
 
                 # Get indexer and reindex the document
                 try:

tests/test_live_retriever_api.py

@@ -18,9 +18,14 @@
 
 from mmore.index.indexer import Indexer
 from mmore.rag.model import DenseModelConfig, SparseModelConfig
-from mmore.run_index_api import make_router as make_index_router
+from mmore.run_index_api import (
+    _apply_uploaded_file_metadata,
+)
+from mmore.run_index_api import (
+    make_router as make_index_router,
+)
 from mmore.run_retriever import make_router, save_results
-from mmore.type import MultimodalSample
+from mmore.type import DocumentMetadata, MultimodalSample
 
 _COLLECTION = "my_docs"
 
@@ -216,13 +221,16 @@ def test_save_results_writes_valid_json(tmp_path):
     docs = [
         Document(
             page_content="Paris is the capital.",
-            metadata={
-                "rank": 1,
-                "similarity": 0.9,
-                "id": "1",
-                "page_numbers": [],
-                "paragraph_numbers": [],
-            },
+            metadata=DocumentMetadata(
+                file_path="paris.txt",
+                extra={
+                    "rank": 1,
+                    "similarity": 0.9,
+                    "id": "1",
+                    "page_numbers": [],
+                    "paragraph_numbers": [],
+                },
+            ).to_dict(),
         )
     ]
     results = [docs]
@@ -246,25 +254,31 @@ def test_save_results_multiple_queries(tmp_path):
         [
             Document(
                 page_content="doc A",
-                metadata={
-                    "rank": 1,
-                    "similarity": 0.8,
-                    "id": "a",
-                    "page_numbers": [],
-                    "paragraph_numbers": [],
-                },
+                metadata=DocumentMetadata(
+                    file_path="doc-a.txt",
+                    extra={
+                        "rank": 1,
+                        "similarity": 0.8,
+                        "id": "a",
+                        "page_numbers": [],
+                        "paragraph_numbers": [],
+                    },
+                ).to_dict(),
             )
         ],
         [
             Document(
                 page_content="doc B",
-                metadata={
-                    "rank": 1,
-                    "similarity": 0.7,
-                    "id": "b",
-                    "page_numbers": [],
-                    "paragraph_numbers": [],
-                },
+                metadata=DocumentMetadata(
+                    file_path="doc-b.txt",
+                    extra={
+                        "rank": 1,
+                        "similarity": 0.7,
+                        "id": "b",
+                        "page_numbers": [],
+                        "paragraph_numbers": [],
+                    },
+                ).to_dict(),
             )
         ],
     ]
@@ -294,10 +308,20 @@ def _fake_doc(file_path: str, document_id: str = "doc") -> MultimodalSample:
         document_id=document_id,
         text="Test document content.",
         modalities=[],
-        metadata={"file_path": file_path},
+        metadata=DocumentMetadata(file_path=file_path),
     )
 
 
+def test_apply_uploaded_file_metadata_preserves_chunk_suffix():
+    doc = _fake_doc("/tmp/original-name.txt", document_id="default-doc")
+
+    _apply_uploaded_file_metadata([doc], "client-doc", "original-name.txt")
+
+    assert doc.document_id == "client-doc"
+    assert doc.id == "client-doc+0"
+    assert doc.metadata.extra["filename"] == "original-name.txt"
+
+
 @pytest.fixture(scope="module")
 def indexer_client(tmp_path_factory):
     """Builds the indexer FastAPI app."""
@@ -392,6 +416,81 @@ def test_upload_file_success(indexer_client):
     assert Path(upload_dir, "new-doc").exists()
 
 
+def test_uploaded_file_has_filename_in_list_files(tmp_path):
+    upload_dir = tmp_path / "uploads"
+    upload_dir.mkdir()
+    db_path = str(tmp_path / "uploaded_list_files.db")
+    config_file = tmp_path / "config.yaml"
+    cfg = {
+        "db": {"uri": db_path, "name": "my_db"},
+        "hybrid_search_weight": 0.5,
+        "k": 2,
+        "collection_name": _COLLECTION,
+        "use_web": False,
+        "reranker_model_name": None,
+    }
+    with open(config_file, "w") as f:
+        yaml.dump(cfg, f)
+
+    with ExitStack() as stack:
+        stack.enter_context(
+            patch(
+                "mmore.index.indexer.SparseModel.from_config",
+                return_value=FakeSparseEmbedding(),
+            )
+        )
+        milvus_client = MilvusClient(db_path, enable_sparse=True)
+        the_indexer = Indexer(
+            dense_model_config=DenseModelConfig(model_name="debug"),
+            sparse_model_config=SparseModelConfig(
+                model_name="naver/splade-cocondenser-selfdistil"
+            ),
+            client=milvus_client,
+        )
+        stack.enter_context(patch("mmore.run_index_api.UPLOAD_DIR", str(upload_dir)))
+        stack.enter_context(patch("mmore.run_index_api.register_all_processors"))
+        stack.enter_context(
+            patch("mmore.run_index_api.get_indexer", return_value=the_indexer)
+        )
+
+        index_app = FastAPI()
+        index_app.include_router(make_index_router(str(config_file)))
+        index_client = TestClient(index_app, raise_server_exceptions=False)
+
+        uploaded_path = str(upload_dir / "listed-doc.txt")
+        stack.enter_context(
+            patch(
+                "mmore.run_index_api.process_files_default",
+                return_value=[_fake_doc(uploaded_path)],
+            )
+        )
+        response = index_client.post(
+            "/v1/files",
+            data={"fileId": "listed-doc"},
+            files={"file": ("listed-doc.txt", b"Hello list files", "text/plain")},
+        )
+        assert response.status_code == 201
+
+        stack.enter_context(
+            patch(
+                "mmore.rag.retriever.SparseModel.from_config",
+                return_value=FakeSparseEmbedding(),
+            )
+        )
+        retriever_app = FastAPI()
+        retriever_app.include_router(make_router(str(config_file)))
+        retriever_client = TestClient(retriever_app)
+
+        response = retriever_client.get(
+            "/list_files", params={"collection_name": _COLLECTION}
+        )
+
+    assert response.status_code == 200
+    files_by_id = {file["id"]: file["filename"] for file in response.json()}
+    assert files_by_id["listed-doc"] == "listed-doc.txt"
+    assert files_by_id["listed-doc"] != "Unknown"
+
+
 def test_upload_duplicate_file_returns_400(indexer_client):
     tc, upload_dir, _ = indexer_client
     duplicate_id = "duplicate-doc"