mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-09 12:54:47 +08:00
## 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:
@@ -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]),
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user