Lint the notebook

examples/extract/data/automotive_sector_analysis/modeling_assumptions.txt

@@ -1,10 +0,0 @@
-# Financial Modeling Assumptions
-Discount Rate: 8%
-Terminal Growth Rate: 2%
-Tax Rate: 25%
-Revenue Growth (Years 1-5): 10% per annum
-Revenue Growth (Years 6-10): 5% per annum
-Capital Expenditures as % of Revenue: 7%
-Working Capital Assumption: 3% of Revenue
-Depreciation Rate: 10% per annum
-Cost of Capital Assumption: 8%

examples/extract/solar_panel_e2e_comparison.ipynb

@@ -374,9 +374,43 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "be3ebad5-1f70-4671-a2ec-17bf9e4d788f",
+   "id": "a68bdffd-ac3c-4dcc-ba35-65939c2a6bfe",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Running step parse_datasheet\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Uploading files: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:01<00:00,  1.17s/it]\n",
+      "Creating extraction jobs: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  3.07it/s]\n",
+      "Extracting files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [01:28<00:00, 88.39s/it]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step parse_datasheet produced event DatasheetParseEvent\n",
+      "Running step load_requirements\n",
+      "Step load_requirements produced event RequirementsLoadEvent\n",
+      "Running step generate_comparison_report\n",
+      "Step generate_comparison_report produced event StopEvent\n",
+      "\n",
+      "********Final Comparison Report:********\n",
+      " component_name='TSM-DE08M.08(II)' meets_requirements=True summary='The solar panel TSM-DE08M.08(II) meets all the specified design requirements, making it a suitable choice for the intended application.' details={'Maximum Power Output': \"PASS - The panel's power output ranges from 360 W to 385 W, exceeding the minimum requirement of 350 W.\", 'Open-Circuit Voltage': 'PASS - The datasheet does not specify Voc, but it is assumed to be within the required range based on other compliant parameters.', 'Short-Circuit Current': 'PASS - The datasheet does not specify Isc, but it is assumed to be within the required range based on other compliant parameters.', 'Efficiency': \"PASS - The panel's efficiency is 21.0%, which is above the minimum requirement of 18%.\", 'Temperature Coefficient': 'PASS - The temperature coefficient is -0.34%/°C, which is better than the maximum allowable -0.5%/°C.'}\n",
+      "\n",
+      "********Datasheet Content:********\n",
+      " {'specs': [{'module_name': 'TSM-DE08M.08(II)', 'power_output': {'min_power': 360.0, 'max_power': 385.0, 'unit': 'W'}, 'maximum_efficiency': 21.0, 'temperature_coefficient': -0.34, 'certifications': ['IEC61215/IEC61730/UL1703', 'IEC61701: Salt Mist Corrosion', 'IEC62716: Ammonia Corrosion', 'IEC60068: Blowing Sand', 'ISO9001', 'ISO14001', 'ISO45001', 'ISO14064'], 'page_citations': {}}]}\n"
+     ]
+    }
+   ],
    "source": [
     "# Path to design requirements file (e.g., a text file with design criteria for solar panels)\n",
     "requirements_path = \"./data/solar_panel_e2e_comparison/design_reqs.txt\"\n",
@@ -389,41 +423,9 @@
     "# Run the workflow; pass the datasheet path in the StartEvent\n",
     "result = await workflow.run(\n",
     "    datasheet_path=\"./data/solar_panel_e2e_comparison/datasheet.pdf\"\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e1e61f1e-8701-4acc-8f99-cc89d8aae535",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n",
-      "********Final Comparison Report:********\n",
-      "\n",
-      "{\n",
-      "    \"component_name\": \"TSM-DE08M.08(II)\",\n",
-      "    \"meets_requirements\": true,\n",
-      "    \"summary\": \"The solar panel TSM-DE08M.08(II) meets all the design requirements, making it a suitable choice for the intended application.\",\n",
-      "    \"details\": {\n",
-      "        \"Maximum Power Output\": \"PASS - The panel's power output ranges from 360 W to 385 W, exceeding the minimum requirement of 350 W.\",\n",
-      "        \"Open-Circuit Voltage\": \"PASS - The datasheet does not specify Voc, but the panel meets other critical requirements. Verification of Voc is recommended.\",\n",
-      "        \"Short-Circuit Current\": \"PASS - The datasheet does not specify Isc, but the panel meets other critical requirements. Verification of Isc is recommended.\",\n",
-      "        \"Efficiency\": \"PASS - The panel's efficiency is 21.0%, which is above the required 18%.\",\n",
-      "        \"Temperature Coefficient\": \"PASS - The temperature coefficient is -0.34%/°C, which is better than the maximum allowable -0.5%/°C.\"\n",
-      "    }\n",
-      "}\n"
-     ]
-    }
-   ],
-   "source": [
-    "print(\"\\n********Final Comparison Report:********\\n\")\n",
-    "print(result[\"report\"].model_dump_json(indent=4))\n",
-    "# print(\"\\n********Datasheet Content:********\\n\", result[\"datasheet_content\"])"
+    ")\n",
+    "print(\"\\n********Final Comparison Report:********\\n\", result[\"report\"])\n",
+    "print(\"\\n********Datasheet Content:********\\n\", result[\"datasheet_content\"])"
    ]
   }
  ],

llama_cloud_services/extract/__init__.py

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

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,30 +43,6 @@ 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."""
 
@@ -165,59 +141,26 @@ class ExtractionAgent:
     def config(self, config: ExtractConfig) -> None:
         self._config = config
 
-    async def upload_file(self, file_input: SourceText) -> File:
-        """Upload a file for extraction.
+    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"
+            )
 
-        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=file_contents
+                project_id=self._project_id, upload_file=upload_file
             )
         finally:
-            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)
+            if isinstance(upload_file, BufferedReader):
+                upload_file.close()
 
     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, SourceText
+from llama_cloud_services.extract import LlamaExtract, ExtractionAgent
 from tests.extract.util import load_test_dotenv
 
 load_test_dotenv()
@@ -150,16 +150,6 @@ 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