add file names in return values (#888)

Author: sourabhdesai

py/llama_cloud_services/beta/classifier/client.py

@@ -17,6 +17,9 @@
 from llama_cloud_services.constants import POLLING_TIMEOUT_SECONDS
 from llama_cloud_services.utils import is_terminal_status, augment_async_errors
 from llama_index.core.async_utils import DEFAULT_NUM_WORKERS, run_jobs
+from llama_cloud_services.beta.classifier.types import (
+    ClassifyJobResultsWithFiles,
+)
 
 
 class ClassificationOutput(BaseModel):
@@ -170,19 +173,20 @@ async def aclassify_file_path(
         file_input_path: str,
         parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
         raise_on_error: bool = True,
-    ) -> ClassifyJobResults:
+    ) -> ClassifyJobResultsWithFiles:
         file = await self.file_client.upload_file(file_input_path)
-        return await self.aclassify_file_ids(
+        results = await self.aclassify_file_ids(
             rules, [file.id], parsing_configuration, raise_on_error
         )
+        return ClassifyJobResultsWithFiles.from_classify_job_results(results, [file])
 
     def classify_file_path(
         self,
         rules: list[ClassifierRule],
         file_input_path: str,
         parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
         raise_on_error: bool = True,
-    ) -> ClassifyJobResults:
+    ) -> ClassifyJobResultsWithFiles:
         with augment_async_errors():
             return asyncio.run(
                 self.aclassify_file_path(
@@ -198,25 +202,26 @@ async def aclassify_file_paths(
         raise_on_error: bool = True,
         workers: int = DEFAULT_NUM_WORKERS,
         show_progress: bool = False,
-    ) -> ClassifyJobResults:
+    ) -> ClassifyJobResultsWithFiles:
         coroutines = [self.file_client.upload_file(path) for path in file_input_paths]
         files: list[File] = await run_jobs(
             coroutines,
             show_progress=show_progress,
             workers=workers,
             desc="Uploading files for classification",
         )
-        return await self.aclassify_file_ids(
+        results = await self.aclassify_file_ids(
             rules, [file.id for file in files], parsing_configuration, raise_on_error
         )
+        return ClassifyJobResultsWithFiles.from_classify_job_results(results, files)
 
     def classify_file_paths(
         self,
         rules: list[ClassifierRule],
         file_input_paths: list[str],
         parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
         raise_on_error: bool = True,
-    ) -> ClassifyJobResults:
+    ) -> ClassifyJobResultsWithFiles:
         with augment_async_errors():
             return asyncio.run(
                 self.aclassify_file_paths(

py/llama_cloud_services/beta/classifier/types.py

@@ -0,0 +1,59 @@
+from llama_cloud.types.classify_job_results import ClassifyJobResults
+from llama_cloud.types.file_classification import FileClassification
+from llama_cloud.types.file import File
+
+
+class FileClassificationWithFile(FileClassification):
+    """
+    File classification with file object.
+    """
+
+    file: File
+
+    @classmethod
+    def from_file_classification(
+        cls, file_classification: FileClassification, file: File
+    ) -> "FileClassificationWithFile":
+        if file_classification.file_id != file.id:
+            raise ValueError(
+                f"File classification ID {file_classification.id} does not match file ID {file.id}"
+            )
+        ctor_args = {
+            **file_classification.dict(),
+            "file": file,
+        }
+        return cls(**ctor_args)
+
+
+class ClassifyJobResultsWithFiles(ClassifyJobResults):
+    """
+    Classify job results with file objects.
+    """
+
+    items: list[FileClassificationWithFile]
+
+    @classmethod
+    def from_classify_job_results(
+        cls, classify_job_results: ClassifyJobResults, files: list[File]
+    ) -> "ClassifyJobResultsWithFiles":
+        if len(classify_job_results.items) != len(files):
+            raise ValueError(
+                f"Number of classify job results {len(classify_job_results.items)} does not match number of files {len(files)}"
+            )
+        # create mapping of file classification result to file object
+        file_id_to_file: dict[str, File] = {file.id: file for file in files}
+        file_classification_to_file: list[tuple[FileClassification, File]] = []
+        for item in classify_job_results.items:
+            if item.file_id not in file_id_to_file:
+                raise ValueError(
+                    f"File classification result {item.id} has file ID {item.file_id} that does not match any provided file ID"
+                )
+            file_classification_to_file.append((item, file_id_to_file[item.file_id]))
+
+        # create a list of file classification with file objects
+        ctor_args = classify_job_results.dict()
+        ctor_args["items"] = [
+            FileClassificationWithFile.from_file_classification(item, file)
+            for item, file in file_classification_to_file
+        ]
+        return cls(**ctor_args)

py/tests/classifier/test_client.py

@@ -2,6 +2,7 @@
 import pytest
 from llama_cloud.client import AsyncLlamaCloud
 from llama_cloud.types import Project, ClassifierRule, ClassifyJobResults
+from llama_cloud_services.beta.classifier.types import ClassifyJobResultsWithFiles
 from llama_cloud_services.beta.classifier.client import ClassifyClient
 from llama_cloud_services.files.client import FileClient
 from llama_cloud.errors.unprocessable_entity_error import UnprocessableEntityError
@@ -187,7 +188,7 @@ async def test_classify_file_path(
             rules=classification_rules, file_input_path=simple_pdf_file_path
         )
 
-    assert isinstance(results, ClassifyJobResults)
+    assert isinstance(results, ClassifyJobResultsWithFiles)
     assert len(results.items) == 1
 
     # Verify the file got classified
@@ -218,7 +219,7 @@ async def test_classify_file_paths(
             file_input_paths=[simple_pdf_file_path, research_paper_path],
         )
 
-    assert isinstance(results, ClassifyJobResults)
+    assert isinstance(results, ClassifyJobResultsWithFiles)
     assert len(results.items) == 2
 
     file_name_to_expected_type = {
@@ -227,8 +228,7 @@ async def test_classify_file_paths(
     }
     # Verify each file got classified
     for item in results.items:
-        file = await file_client.get_file(item.file_id)
-        expected_type = file_name_to_expected_type[file.name]
+        expected_type = file_name_to_expected_type[item.file.name]
         assert item.result.type == expected_type