-
Notifications
You must be signed in to change notification settings - Fork 131
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: enhance API filetype detection (#445)
# Use the library for filetype detection
The mimetype detection has always been very naive in the API - we rely
on the file extension. If the user doesn't include a filename, we return
an error that `Filetype None is not supported`. The library has a
detect_filetype that actually inspects the file bytes, so let's reuse
this.
# Add a `content_type` param to override filetype detection
Add an optional `content_type` param that allows the user to override
the filetype detection. We'll use this value if it's set, or take the
`file.content_type` which is based on the multipart `Content-Type`
header. This provides an alternative when clients are unable to modify
the header.
# Testing
The important thing is that `test_happy_path_all_types` passes in the
docker smoke test - this contains all filetypes that we want the API to
support.
To test manually, you can try sending files to the server with and
without the filename/content_type defined.
Check out this branch and run `make run-web-app`.
Example sending with no extension in filename. This correctly processes
a pdf.
```
import requests
filename = "sample-docs/layout-parser-paper-fast.pdf"
url = "http://localhost:8000/general/v0/general"
with open(filename, 'rb') as f:
files = {'files': ("sample-doc", f)}
response = requests.post(url, files=files)
print(response.text)
```
For the new param, you can try modifying the content type for a text
based file.
Verify that you can change the `metadata.filetype` of the response using
the new param:
```
curl --location 'http://localhost:8000/general/v0/general' \
--form 'files=@"sample-docs/family-day.eml"' \
--form 'content_type="text/plain"'
[
{
"type": "UncategorizedText",
"element_id": "5cafe1ce2b0a96f8e3eba232e790db19",
"text": "MIME-Version: 1.0 Date: Wed, 21 Dec 2022 10:28:53 -0600 Message-ID: <CAPgNNXQKR=o6AsOTr74VMrsDNhUJW0Keou9n3vLa2UO_Nv+tZw@mail.gmail.com> Subject: Family Day From: Mallori Harrell <mallori@unstructured.io> To: Mallori Harrell <mallori@unstructured.io> Content-Type: multipart/alternative; boundary=\"0000000000005c115405f0590ce4\"",
"metadata": {
"filename": "family-day.eml",
"languages": [
"eng"
],
"filetype": "text/plain"
}
},
...
]
```- Loading branch information
Showing
9 changed files
with
152 additions
and
302 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,107 +1,55 @@ | ||
| import mimetypes | ||
| import os | ||
| from fastapi import UploadFile, HTTPException | ||
| from typing import Optional | ||
| from io import BytesIO | ||
|
|
||
| DEFAULT_MIMETYPES = ( | ||
| "application/pdf,application/msword,image/jpeg,image/png,text/markdown," | ||
| "text/x-markdown,text/html," | ||
| "application/vnd.openxmlformats-officedocument.wordprocessingml.document," | ||
| "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet," | ||
| "application/vnd.ms-excel,application/vnd.openxmlformats-officedocument." | ||
| "presentationml.presentation," | ||
| "application/json," | ||
| "application/vnd.ms-powerpoint," | ||
| "text/html,message/rfc822,text/plain,image/png," | ||
| "application/epub,application/epub+zip," | ||
| "application/rtf,text/rtf," | ||
| "application/vnd.oasis.opendocument.text," | ||
| "text/csv,text/x-csv,application/csv,application/x-csv," | ||
| "text/comma-separated-values,text/x-comma-separated-values," | ||
| "application/xml,text/xml,text/x-rst,text/prs.fallenstein.rst," | ||
| "text/tsv,text/tab-separated-values," | ||
| "application/x-ole-storage,application/vnd.ms-outlook," | ||
| "application/yaml," | ||
| "application/x-yaml," | ||
| "text/x-yaml," | ||
| "text/yaml," | ||
| "image/bmp," | ||
| "image/heic," | ||
| "image/tiff," | ||
| "text/org," | ||
| ) | ||
| from fastapi import HTTPException, UploadFile | ||
|
|
||
| if not os.environ.get("UNSTRUCTURED_ALLOWED_MIMETYPES", None): | ||
| os.environ["UNSTRUCTURED_ALLOWED_MIMETYPES"] = DEFAULT_MIMETYPES | ||
| from unstructured.file_utils.filetype import detect_filetype | ||
| from unstructured.file_utils.model import FileType | ||
|
|
||
|
|
||
| def _load_mimetypes() -> None: | ||
| """Call this on startup to ensure that all expected file extensions are present in the mimetypes | ||
| lib""" | ||
| expected_mimetypes = [ | ||
| (".bmp", "image/bmp"), | ||
| (".csv", "application/csv"), | ||
| (".doc", "application/msword"), | ||
| (".docx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document"), | ||
| (".eml", "message/rfc822"), | ||
| (".epub", "application/epub"), | ||
| (".gz", "application/gzip"), | ||
| (".heic", "image/heic"), | ||
| (".html", "text/html"), | ||
| (".jpeg", "image/jpeg"), | ||
| (".jpg", "image/jpeg"), | ||
| (".json", "application/json"), | ||
| (".md", "text/markdown"), | ||
| (".msg", "application/x-ole-storage"), | ||
| (".odt", "application/vnd.oasis.opendocument.text"), | ||
| (".org", "text/org"), | ||
| (".pdf", "application/pdf"), | ||
| (".png", "image/png"), | ||
| (".ppt", "application/vnd.ms-powerpoint"), | ||
| (".pptx", "application/vnd.openxmlformats-officedocument.presentationml.presentation"), | ||
| (".rst", "text/prs.fallenstein.rst"), | ||
| (".rtf", "application/rtf"), | ||
| (".tiff", "image/tiff"), | ||
| (".tsv", "text/tab-separated-values"), | ||
| (".txt", "text/plain"), | ||
| (".xlsx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"), | ||
| (".xml", "text/xml"), | ||
| ] | ||
| def _remove_optional_info_from_mime_type(content_type: str | None) -> str | None: | ||
| """removes charset information from mime types, e.g., | ||
| "application/json; charset=utf-8" -> "application/json" | ||
| """ | ||
| if not content_type: | ||
| return content_type | ||
| return content_type.split(";")[0] | ||
|
|
||
|
|
||
| for extension, mimetype in expected_mimetypes: | ||
| mimetypes.add_type(mimetype, extension) | ||
| def get_validated_mimetype(file: UploadFile, content_type_hint: str | None = None) -> Optional[str]: | ||
| """Given the incoming file, identify and return the correct mimetype. | ||
| Order of operations: | ||
| - If user passed content_type as a form param, take it as truth. | ||
| - Otherwise, use file.content_type (as set by the Content-Type header) | ||
| - If no content_type was passed and the header wasn't useful, call the library's detect_filetype | ||
| _load_mimetypes() | ||
| Once we have a filteype, check is_partitionable and return 400 if we don't support this file. | ||
| """ | ||
| content_type: str | None = None | ||
|
|
||
| if content_type_hint is not None: | ||
| content_type = content_type_hint | ||
| else: | ||
| content_type = _remove_optional_info_from_mime_type(file.content_type) | ||
|
|
||
| def get_validated_mimetype(file: UploadFile) -> Optional[str]: | ||
| """The MIME-type of `file`. | ||
| filetype = FileType.from_mime_type(content_type) | ||
|
|
||
| The mimetype is computed based on `file.content_type`, or the mimetypes lib if that's too | ||
| generic. If the user has set UNSTRUCTURED_ALLOWED_MIMETYPES, validate against this list and | ||
| return HTTP 400 for an invalid type. | ||
| """ | ||
| content_type = file.content_type | ||
| filename = str(file.filename) # -- "None" when file.filename is None -- | ||
| if not content_type or content_type == "application/octet-stream": | ||
| content_type = mimetypes.guess_type(filename)[0] | ||
| # If content_type was not specified, use the library to identify the file | ||
| # We inspect the bytes to do this, so we need to buffer the file | ||
| if not filetype or filetype == FileType.UNK: | ||
| file_buffer = BytesIO(file.file.read()) | ||
| file.file.seek(0) | ||
|
|
||
| # Some filetypes missing for this library, just hardcode them for now | ||
| if not content_type: | ||
| if filename.endswith(".md"): | ||
| content_type = "text/markdown" | ||
| elif filename.endswith(".msg"): | ||
| content_type = "message/rfc822" | ||
| file_buffer.name = file.filename | ||
|
|
||
| allowed_mimetypes_str = os.environ.get("UNSTRUCTURED_ALLOWED_MIMETYPES") | ||
| if allowed_mimetypes_str is not None: | ||
| allowed_mimetypes = allowed_mimetypes_str.split(",") | ||
| filetype = detect_filetype(file=file_buffer) | ||
|
|
||
| if content_type not in allowed_mimetypes: | ||
| raise HTTPException( | ||
| status_code=400, | ||
| detail=(f"File type {content_type} is not supported."), | ||
| ) | ||
| if not filetype.is_partitionable: | ||
| raise HTTPException( | ||
| status_code=400, | ||
| detail=(f"File type {filetype.mime_type} is not supported."), | ||
| ) | ||
|
|
||
| return content_type | ||
| return filetype.mime_type |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,2 +1,2 @@ | ||
| name: general | ||
| version: 0.0.75 | ||
| version: 0.0.76 |
Oops, something went wrong.