fix: receive csv when output_format=text/csv (#373)

Fixed bug where `output_format=text/csv` would return JSON formatted output.
Made it so that `accept` headers are no longer used to determine output format except for multipart/mixed case.

---------

Co-authored-by: Filip Knefel <[email protected]>

Author: ds-filipknefel

README.md

@@ -322,14 +322,15 @@ The output format can also be set to `text/csv` to get the data in csv format ra
 
 The response will be a list of the extracted elements in csv format:
 ```
-"type,text,element_id,filename,page_number,url,sent_from,sent_to,subject,sender\n
-UncategorizedText,\"Hi,\",bc50944723f014607ad612b6983944a7,alert.eml,1,,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],ALERT: Stolen Lunch,Mallori Harrell <[email protected]>\n
-NarrativeText,\"It has come to our attention that as of 9:00am this morning, Harold's lunch is missing. If this was done in error please return the lunch immediately to the fridge on the 2nd floor by noon.\",51944d1f63f9472edb165fb3c9e5c525,alert.eml,1,,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],ALERT: Stolen Lunch,Mallori Harrell <[email protected]>\n
-NarrativeText,\"If the lunch has not been returned by noon, we will be reviewing camera footage to determine who stole Harold's lunch.\",8e8f9e2e50e39e072fda08d277aa77b9,alert.eml,1,,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],ALERT: Stolen Lunch,Mallori Harrell <[email protected]>\n
-NarrativeText,The perpetrators will be PUNISHED to the full extent of our employee code of conduct handbook.,736a826679b971f594103fd9751e5c8f,alert.eml,1,,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],ALERT: Stolen Lunch,Mallori Harrell <[email protected]>\n
-UncategorizedText,\"Thank you for your time,\",3eeae5f64dab54c52dd5fff779808071,alert.eml,1,,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],ALERT: Stolen Lunch,Mallori Harrell <[email protected]>\n
-Title,Unstructured Technologies,d5b612de8cd918addd9569b0255b65b2,alert.eml,1,,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],ALERT: Stolen Lunch,Mallori Harrell <[email protected]>\n
-Title,Data Scientist,46b174f1ec7c25d23e5e50ffff0cc55b,alert.eml,1,,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],ALERT: Stolen Lunch,Mallori Harrell <[email protected]>\n"
+type,element_id,text,filename,sent_from,sent_to,subject,languages,filetype
+UncategorizedText,db1ca22813f01feda8759ff04a844e56,"Hi All,",family-day.eml,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],Family Day,['eng'],message/rfc822
+NarrativeText,a663c393a5e143c01ef2bb5c98efa2c1,Get excited for our first annual family day! ,family-day.eml,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],Family Day,['eng'],message/rfc822
+NarrativeText,ce65ca3bef59957d3f1c2bab5725c82f,"There will be face painting, a petting zoo, funnel cake and more.",family-day.eml,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],Family Day,['eng'],message/rfc822
+NarrativeText,d7bcf988af9f06042d83e25c531e5744,Make sure to RSVP!,family-day.eml,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],Family Day,['eng'],message/rfc822
+Title,5550577db69c2c8aabcd90979698120a,Best.,family-day.eml,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],Family Day,['eng'],message/rfc822
+Title,ca1c571d993b6c1ed8ef56a06c16ba22,Mallori Harrell,family-day.eml,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],Family Day,['eng'],message/rfc822
+Title,d5b612de8cd918addd9569b0255b65b2,Unstructured Technologies,family-day.eml,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],Family Day,['eng'],message/rfc822
+Title,2e0b9e8ee04b9594a9c26d8535b818ff,Data Scientist,family-day.eml,['Mallori Harrell <[email protected]>'],['Mallori Harrell <[email protected]>'],Family Day,['eng'],message/rfc822
 ```
 
 #### Parallel Mode for PDFs

prepline_general/api/general.py

@@ -751,12 +751,6 @@ def general_partition(
                 files[file_index], form_params.gz_uncompressed_content_type
             )
 
-    default_response_type = form_params.output_format or "application/json"
-    if not content_type or content_type == "*/*" or content_type == "multipart/mixed":
-        media_type = default_response_type
-    else:
-        media_type = content_type
-
     def response_generator(is_multipart: bool):
         for file in files:
             file_content_type = get_validated_mimetype(file)
@@ -775,7 +769,7 @@ def response_generator(is_multipart: bool):
                 skip_infer_table_types=form_params.skip_infer_table_types,
                 strategy=form_params.strategy,
                 xml_keep_tags=form_params.xml_keep_tags,
-                response_type=media_type,
+                response_type=form_params.output_format,
                 filename=str(file.filename),
                 file_content_type=file_content_type,
                 languages=form_params.languages,
@@ -790,27 +784,12 @@ def response_generator(is_multipart: bool):
                 overlap_all=form_params.overlap_all,
             )
 
-            if not is_compatible_response_type(media_type, type(response)):
-                raise HTTPException(
-                    detail=(
-                        f"Conflict in media type {media_type}"
-                        f" with response type {type(response)}.\n"
-                    ),
-                    status_code=status.HTTP_406_NOT_ACCEPTABLE,
-                )
-
-            if media_type not in ["application/json", "text/csv", "*/*", "multipart/mixed"]:
-                raise HTTPException(
-                    detail=f"Unsupported media type {media_type}.\n",
-                    status_code=status.HTTP_406_NOT_ACCEPTABLE,
-                )
-
             yield (
                 json.dumps(response)
                 if is_multipart and type(response) not in [str, bytes]
                 else (
                     PlainTextResponse(response)
-                    if not is_multipart and media_type == "text/csv"
+                    if not is_multipart and form_params.output_format == "text/csv"
                     else response
                 )
             )
@@ -819,7 +798,7 @@ def join_responses(
         responses: Sequence[str | List[Dict[str, Any]] | PlainTextResponse]
     ) -> List[str | List[Dict[str, Any]]] | PlainTextResponse:
         """Consolidate partitionings from multiple documents into single response payload."""
-        if media_type != "text/csv":
+        if form_params.output_format != "text/csv":
             return cast(List[Union[str, List[Dict[str, Any]]]], responses)
         responses = cast(List[PlainTextResponse], responses)
         data = pd.read_csv(  # pyright: ignore[reportUnknownMemberType]
@@ -836,7 +815,9 @@ def join_responses(
         return PlainTextResponse(data.to_csv())
 
     return (
-        MultipartMixedResponse(response_generator(is_multipart=True), content_type=media_type)
+        MultipartMixedResponse(
+            response_generator(is_multipart=True), content_type=form_params.output_format
+        )
         if content_type == "multipart/mixed"
         else (
             list(response_generator(is_multipart=False))[0]

prepline_general/api/models/form_params.py

@@ -85,7 +85,7 @@ def as_form(
             ),
         ] = None,
         output_format: Annotated[
-            str,
+            Literal["application/json", "text/csv"],
             Form(
                 title="Output Format",
                 description="The format of the response. Supported formats are application/json and text/csv. Default: application/json.",

test_general/api/test_app.py

@@ -975,3 +975,32 @@ def test_get_request():
     response = client.get("/general/v0/general")
     assert response.status_code == 405
     assert response.json() == {"detail": "Only POST requests are supported."}
+
+
+def test_output_format_csv():
+    client = TestClient(app)
+    test_file = Path("sample-docs") / "family-day.eml"
+    response = client.post(
+        MAIN_API_ROUTE,
+        files=[("files", (str(test_file), open(test_file, "rb")))],
+        data={"output_format": "text/csv"},
+    )
+    assert response.status_code == 200
+    df = pd.read_csv(io.StringIO(response.text))
+    assert len(df) == 8
+    assert df["text"][3] == "Make sure to RSVP!"
+
+
+def test_output_format_csv_ignore_specified_accept_header():
+    client = TestClient(app)
+    test_file = Path("sample-docs") / "family-day.eml"
+    response = client.post(
+        MAIN_API_ROUTE,
+        files=[("files", (str(test_file), open(test_file, "rb")))],
+        data={"output_format": "text/csv"},
+        headers={"accept": "application/json"},
+    )
+    assert response.status_code == 200
+    df = pd.read_csv(io.StringIO(response.text))
+    assert len(df) == 8
+    assert df["text"][3] == "Make sure to RSVP!"