compat python 3.9

llama_cloud_services/extract/__init__.py

@@ -1,3 +1,7 @@
-from llama_cloud_services.extract.extract import LlamaExtract, ExtractionAgent
+from llama_cloud_services.extract.extract import (
+    LlamaExtract,
+    ExtractionAgent,
+    SourceText,
+)
 
-__all__ = ["LlamaExtract", "ExtractionAgent"]
+__all__ = ["LlamaExtract", "ExtractionAgent", "SourceText"]

llama_cloud_services/extract/extract.py

@@ -34,7 +34,7 @@ from concurrent.futures import ThreadPoolExecutor
 
 T = TypeVar("T")
 
-FileInput = Union[str, Path, bytes, BufferedIOBase]
+
 SchemaInput = Union[JSONObjectType, Type[BaseModel]]
 
 DEFAULT_EXTRACT_CONFIG = ExtractConfig(
@@ -43,6 +43,30 @@ DEFAULT_EXTRACT_CONFIG = ExtractConfig(
 )
 
 
+class SourceText:
+    def __init__(
+        self,
+        file: Union[bytes, BufferedIOBase, str, Path],
+        filename: Optional[str] = None,
+    ):
+        self.file = file
+        self.filename = filename
+        self._validate()
+
+    def _validate(self) -> None:
+        """Ensure filename is provided when needed."""
+        if isinstance(self.file, (bytes, BufferedIOBase)):
+            if not self.filename and hasattr(self.file, "name"):
+                self.filename = os.path.basename(str(self.file.name))
+            elif not hasattr(self.file, "name") and self.filename is None:
+                raise ValueError(
+                    "filename must be provided when file is bytes or a file-like object without a name"
+                )
+
+
+FileInput = Union[str, Path, BufferedIOBase, SourceText]
+
+
 class ExtractionAgent:
     """Class representing a single extraction agent with methods for extraction operations."""
 
@@ -141,26 +165,59 @@ class ExtractionAgent:
     def config(self, config: ExtractConfig) -> None:
         self._config = config
 
-    async def _upload_file(self, file_input: FileInput) -> File:
-        """Upload a file for extraction."""
-        if isinstance(file_input, BufferedIOBase):
-            upload_file = file_input
-        elif isinstance(file_input, bytes):
-            upload_file = BytesIO(file_input)
-        elif isinstance(file_input, (str, Path)):
-            upload_file = open(file_input, "rb")
-        else:
-            raise ValueError(
-                "file_input must be either a file path string, file bytes, or buffer object"
-            )
+    async def upload_file(self, file_input: SourceText) -> File:
+        """Upload a file for extraction.
 
+        Args:
+            file_input: The file to upload (path, bytes, or file-like object)
+
+        Raises:
+            ValueError: If filename is not provided for bytes input or for file-like objects
+                       without a name attribute.
+        """
         try:
+            file_contents: Union[BufferedIOBase, BytesIO]
+            if isinstance(file_input.file, (str, Path)):
+                file_contents = open(file_input.file, "rb")
+            elif isinstance(file_input.file, bytes):
+                file_contents = BytesIO(file_input.file)
+            else:
+                file_contents = file_input.file
+            # Add name attribute to file object if needed
+            if not hasattr(file_contents, "name"):
+                file_contents.name = file_input.filename  # type: ignore
+
             return await self._client.files.upload_file(
-                project_id=self._project_id, upload_file=upload_file
+                project_id=self._project_id, upload_file=file_contents
             )
         finally:
-            if isinstance(upload_file, BufferedReader):
-                upload_file.close()
+            if isinstance(file_contents, BufferedReader):
+                file_contents.close()
+
+    async def _upload_file(self, file_input: FileInput) -> File:
+        source_text = None
+        if isinstance(file_input, SourceText):
+            source_text = file_input
+        elif isinstance(file_input, (str, Path)):
+            path = Path(file_input)
+            source_text = SourceText(file=path, filename=path.name)
+        else:
+            # Try to get filename from the file object if not provided
+            filename = None
+            if hasattr(file_input, "name"):
+                filename = os.path.basename(str(file_input.name))
+            if filename is None:
+                raise ValueError(
+                    "Use SourceText to provide filename when uploading bytes or file-like objects."
+                )
+
+            warnings.warn(
+                "Use SourceText instead of bytes or file-like objects",
+                DeprecationWarning,
+            )
+            source_text = SourceText(file=file_input, filename=filename)
+
+        return await self.upload_file(source_text)
 
     async def _wait_for_job_result(self, job_id: str) -> Optional[ExtractRun]:
         """Wait for and return the results of an extraction job."""

tests/extract/test_extract_api.py

@@ -3,7 +3,7 @@ import pytest
 from pathlib import Path
 from pydantic import BaseModel
 
-from llama_cloud_services.extract import LlamaExtract, ExtractionAgent
+from llama_cloud_services.extract import LlamaExtract, ExtractionAgent, SourceText
 from tests.extract.util import load_test_dotenv
 
 load_test_dotenv()
@@ -150,6 +150,16 @@ class TestExtractionAgent:
         assert "title" in result.data
         assert "summary" in result.data
 
+    def test_extract_file_from_bytes(self, test_agent):
+        with open(TEST_PDF, "rb") as f:
+            file_bytes = f.read()
+        result = test_agent.extract(SourceText(file=file_bytes, filename=TEST_PDF.name))
+        assert result.status == "SUCCESS"
+        assert result.data is not None
+        assert isinstance(result.data, dict)
+        assert "title" in result.data
+        assert "summary" in result.data
+
     @pytest.mark.asyncio
     async def test_extract_multiple_files(self, test_agent):
         files = [TEST_PDF, TEST_PDF]  # Using same file twice for testing