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/<dataset_id>/documents/<document_id>` and
`GET /api/v1/documents/<document_id>`.
- 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.
This commit is contained in:
kpdev
2026-06-04 19:08:53 -07:00
committed by GitHub
parent 794c1f4b25
commit bd49fd70aa
4 changed files with 60 additions and 11 deletions

View File

@@ -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/<dataset_id>/documents/<document_id>", 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/<document_id>", 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]),
)

View File

@@ -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

View File

@@ -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)

View File

@@ -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)