From bd49fd70aaff43b70863ba1a430630c5783a869d Mon Sep 17 00:00:00 2001 From: kpdev <156195510+kiannidev@users.noreply.github.com> Date: Thu, 4 Jun 2026 19:08:53 -0700 Subject: [PATCH] fix(api): set SDK document download Content-Type from filename (#15112) (#15113) ## Summary - Infer `Content-Type` from the stored document filename on SDK download routes. - Covers `GET /api/v1/datasets//documents/` and `GET /api/v1/documents/`. - Aligns with REST preview/download via `CONTENT_TYPE_MAP`. ## Test plan - [x] `pytest test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py::TestDocRoutesUnit::test_download_mimetype_from_filename` - [x] Manual: `curl -sSI` on SDK dataset document download for a PDF; expect `Content-Type: application/pdf` Fixes #15112. --- api/apps/restful_apis/document_api.py | 14 ++++++- test/testcases/restful_api/test_documents.py | 6 ++- .../test_doc_sdk_routes_unit.py | 38 +++++++++++++++++-- .../test_openai_stream_no_duplicate.py | 13 +++++-- 4 files changed, 60 insertions(+), 11 deletions(-) diff --git a/api/apps/restful_apis/document_api.py b/api/apps/restful_apis/document_api.py index ffb62f401..adf236f82 100644 --- a/api/apps/restful_apis/document_api.py +++ b/api/apps/restful_apis/document_api.py @@ -1956,6 +1956,16 @@ async def get(doc_id): except Exception as e: return server_error_response(e) + +def _mimetype_for_document(doc) -> str: + match = re.search(r"\.([^.]+)$", (doc.name or "").lower()) + if not match: + return "application/octet-stream" + ext = match.group(1) + fallback_prefix = "image" if doc.type == FileType.VISUAL.value else "application" + return CONTENT_TYPE_MAP.get(ext, f"{fallback_prefix}/{ext}") + + @manager.route("/datasets//documents/", methods=["GET"]) # noqa: F821 @login_required async def download(dataset_id, document_id): @@ -2010,7 +2020,7 @@ async def download(dataset_id, document_id): file, as_attachment=True, attachment_filename=doc[0].name, - mimetype="application/octet-stream", # Set a default MIME type + mimetype=_mimetype_for_document(doc[0]), ) @manager.route("/documents/", methods=["GET"]) # noqa: F821 @@ -2067,5 +2077,5 @@ async def download_document(document_id): file, as_attachment=True, attachment_filename=doc[0].name, - mimetype="application/octet-stream", # Set a default MIME type + mimetype=_mimetype_for_document(doc[0]), ) diff --git a/test/testcases/restful_api/test_documents.py b/test/testcases/restful_api/test_documents.py index e0084c13b..74dbffa7f 100644 --- a/test/testcases/restful_api/test_documents.py +++ b/test/testcases/restful_api/test_documents.py @@ -120,8 +120,10 @@ def _wait_document_runs(rest_client, dataset_id, document_ids, expected_run="DON def _download_document_to_file(rest_client, dataset_id, document_id, save_path): res = rest_client.get(f"/datasets/{dataset_id}/documents/{document_id}", timeout=60) - if res.status_code == 200 and res.headers.get("Content-Type", "").startswith("application/octet-stream"): - save_path.write_bytes(res.content) + if res.status_code == 200: + disposition = res.headers.get("Content-Disposition", "") + if disposition and "attachment" in disposition.lower(): + save_path.write_bytes(res.content) return res diff --git a/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py b/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py index 3d0a5b52e..66302fada 100644 --- a/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py +++ b/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py @@ -131,7 +131,7 @@ def _run(coro): return asyncio.run(coro) -def _load_doc_module(monkeypatch): +def _load_doc_module(monkeypatch, module_basename="chunk_api"): repo_root = Path(__file__).resolve().parents[4] common_pkg = ModuleType("common") common_pkg.__path__ = [str(repo_root / "common")] @@ -492,7 +492,25 @@ def _load_doc_module(monkeypatch): tenant_model_service_mod.get_tenant_default_model_by_type = _get_tenant_default_model_by_type monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod) - module_path = repo_root / "api" / "apps" / "restful_apis" / "chunk_api.py" + if module_basename == "document_api": + stub_apps_services = ModuleType("api.apps.services") + stub_apps_services.__path__ = [str(repo_root / "api" / "apps" / "services")] + monkeypatch.setitem(sys.modules, "api.apps.services", stub_apps_services) + + document_api_service_mod = ModuleType("api.apps.services.document_api_service") + document_api_service_mod.validate_document_update_fields = lambda *_args, **_kwargs: (None, None) + document_api_service_mod.map_doc_keys = lambda doc: doc.to_dict() if hasattr(doc, "to_dict") else doc + document_api_service_mod.map_doc_keys_with_run_status = lambda doc, run_status="0": { + **(doc if isinstance(doc, dict) else doc.to_dict()), + "run": run_status, + } + document_api_service_mod.update_document_name_only = lambda *_args, **_kwargs: None + document_api_service_mod.update_chunk_method = lambda *_args, **_kwargs: None + document_api_service_mod.update_document_status_only = lambda *_args, **_kwargs: None + document_api_service_mod.reset_document_for_reparse = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "api.apps.services.document_api_service", document_api_service_mod) + + module_path = repo_root / "api" / "apps" / "restful_apis" / f"{module_basename}.py" spec = importlib.util.spec_from_file_location("test_doc_sdk_routes_unit", module_path) module = importlib.util.module_from_spec(spec) module.manager = _DummyManager() @@ -515,7 +533,11 @@ def _route_core(func): def _patch_send_file(monkeypatch, module): async def _fake_send_file(file_obj, **kwargs): - return {"file": file_obj, "filename": kwargs.get("attachment_filename")} + return { + "file": file_obj, + "filename": kwargs.get("attachment_filename"), + "mimetype": kwargs.get("mimetype"), + } monkeypatch.setattr(module, "send_file", _fake_send_file) @@ -613,6 +635,16 @@ class TestDocRoutesUnit: res = _run(module.download_doc("doc-1")) assert res["filename"] == "doc.txt" + def test_download_mimetype_from_filename(self, monkeypatch): + module = _load_doc_module(monkeypatch, module_basename="document_api") + _patch_send_file(monkeypatch, module) + _patch_storage(monkeypatch, module, file_stream=b"pdf-bytes") + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc(name="report.pdf", doc_type=FileType.PDF)]) + monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n")) + res = _run(module.download.__wrapped__("ds-1", "doc-1")) + assert res["filename"] == "report.pdf" + assert res["mimetype"] == "application/pdf" + def test_parse_branches(self, monkeypatch): module = _load_doc_module(monkeypatch) diff --git a/test/unit_test/api/apps/restful_apis/test_openai_stream_no_duplicate.py b/test/unit_test/api/apps/restful_apis/test_openai_stream_no_duplicate.py index 410dd6104..3fad96476 100644 --- a/test/unit_test/api/apps/restful_apis/test_openai_stream_no_duplicate.py +++ b/test/unit_test/api/apps/restful_apis/test_openai_stream_no_duplicate.py @@ -51,15 +51,21 @@ def _stub(monkeypatch, name, **attrs): def _load_openai_api(monkeypatch): """Load api/apps/restful_apis/openai_api.py with the heavy deps stubbed.""" + repo_root = Path(__file__).resolve().parents[5] + apps_mod = ModuleType("api.apps") + apps_mod.__path__ = [str(repo_root / "api" / "apps")] + apps_mod.current_user = SimpleNamespace(id="tenant-1") + apps_mod.login_required = lambda func: func + monkeypatch.setitem(sys.modules, "api.apps", apps_mod) + _stub(monkeypatch, "quart", Response=object, jsonify=lambda *a, **k: None) - _stub(monkeypatch, "api.apps", current_user=SimpleNamespace(id="tenant-1"), login_required=lambda func: func) # Pre-register nested modules so importlib finds them directly in # sys.modules without trying to traverse the stubbed parent package. _stub( monkeypatch, "api.apps.restful_apis._generation_params", - extract_generation_config=lambda *a, **k: ({}, {}), - merge_generation_config=lambda *a, **k: None, + extract_generation_config=lambda req: {}, + merge_generation_config=lambda *_a, **_k: None, ) _stub(monkeypatch, "api.db.services.dialog_service", DialogService=SimpleNamespace(), async_chat=lambda *_a, **_k: None) _stub(monkeypatch, "api.db.services.doc_metadata_service", DocMetadataService=SimpleNamespace()) @@ -84,7 +90,6 @@ def _load_openai_api(monkeypatch): _stub(monkeypatch, "rag.prompts.generator", chunks_format=lambda reference: list(reference) if isinstance(reference, list) else []) _stub(monkeypatch, "api.utils.reference_metadata_utils", enrich_chunks_with_document_metadata=lambda *_a, **_k: None) - repo_root = Path(__file__).resolve().parents[5] module_path = repo_root / "api" / "apps" / "restful_apis" / "openai_api.py" spec = importlib.util.spec_from_file_location("test_openai_stream_openai_api", module_path) module = importlib.util.module_from_spec(spec)