nits

* add alternative builder method

* fix test

examples/README.md

@@ -5,5 +5,6 @@ In this folder you will find several python notebooks that contain examples rega
 - [LlamaParse](./parse/)
 - [LlamaExtract](./extract/)
 - [LlamaReport](./report/)
+- [LlamaCloudIndex](./index/)
 
 Follow the instructions in each notebook to get started!

py/llama_cloud_services/beta/classifier/client.py

@@ -17,6 +17,9 @@ from llama_cloud_services.files.client import FileClient
 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):
@@ -52,6 +55,24 @@ class ClassifyClient:
         self.file_client = FileClient(client, project_id, organization_id)
         self.polling_timeout = polling_timeout
 
+    @classmethod
+    def from_api_key(
+        cls,
+        api_key: str,
+        project_id: Optional[str] = None,
+        organization_id: Optional[str] = None,
+        base_url: Optional[str] = None,
+    ) -> "ClassifyClient":
+        """
+        Create a classify client from an API key.
+        """
+        client = AsyncLlamaCloud(token=api_key, base_url=base_url)
+        return cls(
+            client,
+            project_id,
+            organization_id,
+        )
+
     async def acreate_classify_job(
         self,
         rules: list[ClassifierRule],
@@ -152,11 +173,12 @@ class ClassifyClient:
         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,
@@ -164,7 +186,7 @@ class ClassifyClient:
         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(
@@ -180,7 +202,7 @@ class ClassifyClient:
         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,
@@ -188,9 +210,10 @@ class ClassifyClient:
             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,
@@ -198,7 +221,7 @@ class ClassifyClient:
         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 os
 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
@@ -130,6 +131,44 @@ async def test_classify_file_ids(
         assert item.result.type == expected_type
 
 
+@pytest.mark.asyncio
+async def test_classify_file_ids_from_api_key(
+    e2e_test_settings: EndToEndTestSettings,
+    file_client: FileClient,
+    simple_pdf_file_path: str,
+    research_paper_path: str,
+    classification_rules: list[ClassifierRule],
+):
+    """Test classifying files by their IDs"""
+    # Upload test files first to get their IDs
+    pdf_file = await file_client.upload_file(simple_pdf_file_path)
+    research_paper_file = await file_client.upload_file(research_paper_path)
+
+    classify_client = ClassifyClient.from_api_key(
+        api_key=e2e_test_settings.LLAMA_CLOUD_API_KEY.get_secret_value(),
+        base_url=e2e_test_settings.LLAMA_CLOUD_BASE_URL,
+        project_id=pdf_file.project_id,
+        organization_id=e2e_test_settings.LLAMA_CLOUD_ORGANIZATION_ID,
+    )
+
+    # Classify the uploaded files
+    results = await classify_client.aclassify_file_ids(
+        rules=classification_rules, file_ids=[pdf_file.id, research_paper_file.id]
+    )
+
+    assert isinstance(results, ClassifyJobResults)
+    assert len(results.items) == 2
+
+    file_id_to_expected_type = {
+        pdf_file.id: "number",
+        research_paper_file.id: "research_paper",
+    }
+    # Verify each file got classified
+    for item in results.items:
+        expected_type = file_id_to_expected_type[item.file_id]
+        assert item.result.type == expected_type
+
+
 @parameterize_sync_and_async
 @pytest.mark.asyncio
 async def test_classify_file_path(
@@ -149,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
@@ -180,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 = {
@@ -189,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