Pass filename when uploading local files with HumeBatchClient (#76)

Author: gregorybchris

hume/_batch/hume_batch_client.py

@@ -1,9 +1,9 @@
 """Batch API client."""
 import json
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
-import requests
+from requests import Session
 
 from hume._batch.batch_job import BatchJob
 from hume._batch.batch_job_details import BatchJobDetails
@@ -55,6 +55,7 @@ def __init__(
             timeout (int): Time in seconds before canceling long-running Hume API requests.
         """
         self._timeout = timeout
+        self._session = Session()
         super().__init__(api_key, *args, **kwargs)
 
     @classmethod
@@ -118,7 +119,7 @@ def get_job_details(self, job_id: str) -> BatchJobDetails:
             BatchJobDetails: Batch job details.
         """
         endpoint = self._construct_endpoint(f"jobs/{job_id}")
-        response = requests.get(
+        response = self._session.get(
             endpoint,
             timeout=self._timeout,
             headers=self._get_client_headers(),
@@ -148,7 +149,7 @@ def get_job_predictions(self, job_id: str) -> Any:
             Any: Batch job predictions.
         """
         endpoint = self._construct_endpoint(f"jobs/{job_id}/predictions")
-        response = requests.get(
+        response = self._session.get(
             endpoint,
             timeout=self._timeout,
             headers=self._get_client_headers(),
@@ -179,7 +180,7 @@ def download_job_artifacts(self, job_id: str, filepath: Union[str, Path]) -> Non
             Any: Batch job artifacts.
         """
         endpoint = self._construct_endpoint(f"jobs/{job_id}/artifacts")
-        response = requests.get(
+        response = self._session.get(
             endpoint,
             timeout=self._timeout,
             headers=self._get_client_headers(),
@@ -212,7 +213,7 @@ def _construct_request(
     def _submit_job(
         self,
         request_body: Any,
-        files: Optional[List[Union[str, Path]]],
+        filepaths: Optional[List[Union[str, Path]]],
     ) -> BatchJob:
         """Start a job for batch processing by passing a JSON request body.
 
@@ -221,7 +222,7 @@ def _submit_job(
 
         Args:
             request_body (Any): JSON request body to be passed to the batch API.
-            files (Optional[List[Union[str, Path]]]): List of paths to files on the local disk to be processed.
+            filepaths (Optional[List[Union[str, Path]]]): List of paths to files on the local disk to be processed.
 
         Raises:
             HumeClientException: If the batch job fails to start.
@@ -231,21 +232,20 @@ def _submit_job(
         """
         endpoint = self._construct_endpoint("jobs")
 
-        if files is None:
-            response = requests.post(
+        if filepaths is None:
+            response = self._session.post(
                 endpoint,
                 json=request_body,
                 timeout=self._timeout,
                 headers=self._get_client_headers(),
             )
         else:
-            post_files = [("file", Path(path).read_bytes()) for path in files]
-            post_files.append(("json", json.dumps(request_body).encode("utf-8")))
-            response = requests.post(
+            form_data = self._get_multipart_form_data(request_body, filepaths)
+            response = self._session.post(
                 endpoint,
                 timeout=self._timeout,
                 headers=self._get_client_headers(),
-                files=post_files,
+                files=form_data,
             )
 
         try:
@@ -268,3 +268,30 @@ def _submit_job(
             raise HumeClientException(f"Unexpected error when starting batch job: {body}")
 
         return BatchJob(self, body["job_id"])
+
+    def _get_multipart_form_data(
+        self,
+        request_body: Any,
+        filepaths: List[Union[str, Path]],
+    ) -> List[Tuple[str, Union[bytes, Tuple[str, bytes]]]]:
+        """Convert a list of filepaths into a list of multipart form data.
+
+        Multipart form data allows the client to attach files to the POST request,
+        including both the raw file bytes and the filename.
+
+        Args:
+            request_body (Any): JSON request body to be passed to the batch API.
+            filepaths (List[Union[str, Path]]): List of paths to files on the local disk to be processed.
+
+        Returns:
+            List[Tuple[str, Union[bytes, Tuple[str, bytes]]]]: A list of tuples representing
+                the multipart form data for the POST request.
+        """
+        form_data: List[Tuple[str, Union[bytes, Tuple[str, bytes]]]] = []
+        for filepath in filepaths:
+            path = Path(filepath)
+            post_file = ("file", (path.name, path.read_bytes()))
+            form_data.append(post_file)
+
+        form_data.append(("json", json.dumps(request_body).encode("utf-8")))
+        return form_data

pyproject.toml

@@ -25,7 +25,7 @@ license = "Proprietary"
 name = "hume"
 readme = "README.md"
 repository = "https://github.com/HumeAI/hume-python-sdk"
-version = "0.3.4"
+version = "0.3.5"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4"

tests/batch/test_hume_batch_client.py

@@ -1,7 +1,7 @@
 from unittest.mock import MagicMock
 
 import pytest
-from pytest import MonkeyPatch
+from pytest import MonkeyPatch, TempPathFactory
 
 from hume import BatchJob, HumeBatchClient
 from hume.models.config import BurstConfig, FaceConfig, LanguageConfig, ProsodyConfig
@@ -96,3 +96,36 @@ def test_language(self, batch_client: HumeBatchClient):
     def test_get_job(self, batch_client: HumeBatchClient):
         job = batch_client.get_job("mock-job-id")
         assert job.id == "mock-job-id"
+
+    def test_files(self, batch_client: HumeBatchClient):
+        mock_filepath = "my-audio.mp3"
+        config = ProsodyConfig(identify_speakers=True)
+        job = batch_client.submit_job(None, [config], files=[mock_filepath])
+        assert isinstance(job, BatchJob)
+        assert job.id == "mock-job-id"
+        batch_client._submit_job.assert_called_once_with(
+            {
+                "urls": None,
+                "models": {
+                    "prosody": {
+                        "identify_speakers": True,
+                    },
+                }
+            },
+            ['my-audio.mp3'],
+        )
+
+    def test_get_multipart_form_data(self, batch_client: HumeBatchClient, tmp_path_factory: TempPathFactory):
+        dirpath = tmp_path_factory.mktemp("multipart")
+        filepath = dirpath / "my-audio.mp3"
+        with filepath.open("w") as f:
+            f.write("I can't believe this test passed!")
+
+        request_body = {"mock": "body"}
+        filepaths = [filepath]
+        result = batch_client._get_multipart_form_data(request_body, filepaths)
+
+        assert result == [
+            ('file', ('my-audio.mp3', b"I can't believe this test passed!")),
+            ('json', b'{"mock": "body"}'),
+        ]

tests/batch/test_service_hume_batch_client.py

@@ -121,8 +121,8 @@ def test_job_invalid_api_key(self, eval_data: EvalData, batch_client: HumeBatchC
             rehydrated_job = BatchJob(invalid_client, job.id)
             rehydrated_job.await_complete(10)
 
-    def test_local_file_upload(self, eval_data: EvalData, batch_client: HumeBatchClient,
-                               tmp_path_factory: TempPathFactory):
+    def test_local_file_upload_simple(self, eval_data: EvalData, batch_client: HumeBatchClient,
+                                      tmp_path_factory: TempPathFactory):
         data_url = eval_data["image-obama-face"]
         data_filepath = tmp_path_factory.mktemp("data-dir") / "obama.png"
         urlretrieve(data_url, data_filepath)
@@ -131,6 +131,9 @@ def test_local_file_upload(self, eval_data: EvalData, batch_client: HumeBatchCli
         job = batch_client.submit_job([], [config], files=[data_filepath])
         self.check_job(job, config, FaceConfig, job_files_dirpath, complete_config=False)
 
+        predictions = job.get_predictions()
+        assert predictions[0]["source"]["filename"] == "obama.png"
+
     def test_local_file_upload_configure(self, eval_data: EvalData, batch_client: HumeBatchClient,
                                          tmp_path_factory: TempPathFactory):
         data_url = eval_data["text-happy-place"]