Compare commits

...

9 Commits

Author SHA1 Message Date
Neeraj Pradhan 356f806688 compat python 3.9 2025-04-09 09:52:37 -07:00
Neeraj Pradhan d06c17ed37 backwards compatibility 2025-04-08 22:24:58 -07:00
Neeraj Pradhan e1ec231484 Fix bytes input for LlamaExtract 2025-04-08 21:33:41 -07:00
Jerry Liu 16f81bd7ee add due diligence notebook (#670) 2025-04-08 09:13:11 -07:00
Marplex 0ee049fd11 Add layout agent mode visual citation demo notebook (#672) 2025-04-07 09:54:06 -06:00
Neeraj Pradhan 7dba17e5bc Update extract.md (#671) 2025-04-06 22:18:03 -07:00
Jerry Liu eeb678b937 solar panel extraction workflow (#667)
* cr

* cr

* cr
2025-04-02 17:28:13 -07:00
Emanuel Ferreira fe4eb664fd chore: add base url documentation (#666)
* wip

* newline

* wip

* docs
2025-04-01 18:43:17 -03:00
Jerry Liu 257720e443 fix notebook (#665)
cr
2025-04-01 08:05:34 -07:00
16 changed files with 1551 additions and 28 deletions
+19
View File
@@ -38,6 +38,25 @@ See the quickstart guides for each service for more information:
- [LlamaReport (beta/invite-only)](./report.md)
- [LlamaExtract (beta/invite-only)](./extract.md)
## Switch to EU SaaS 🇪🇺
If you are interested in using LlamaCloud services in the EU, you can adjust your base URL to `https://api.cloud.eu.llamaindex.ai`.
You can also create your API key in the EU region [here](https://cloud.eu.llamaindex.ai).
```python
from llama_cloud_services import (
LlamaParse,
LlamaReport,
LlamaExtract,
EU_BASE_URL,
)
parser = LlamaParse(api_key="YOUR_API_KEY", base_url=EU_BASE_URL)
report = LlamaReport(api_key="YOUR_API_KEY", base_url=EU_BASE_URL)
extract = LlamaExtract(api_key="YOUR_API_KEY", base_url=EU_BASE_URL)
```
## Documentation
You can see complete SDK and API documentation for each service on [our official docs](https://docs.cloud.llamaindex.ai/).
File diff suppressed because one or more lines are too long
@@ -0,0 +1,10 @@
# 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%
Binary file not shown.

After

Width:  |  Height:  |  Size: 67 KiB

@@ -28,13 +28,12 @@
"\n",
"### Agent Workflow (Combination of Automation and Chat)\n",
"1. **Upload Datasheet:** The engineer uploads the LM317 datasheet PDF. \n",
"2. **Structured Extraction:** An automated agent processes the PDF and extracts key technical details into structured fields (e.g., output voltage as a range with separate min/max values). It also returns page citations for verification. \n",
"2. **Structured Extraction:** An automated agent processes the PDF and extracts key technical details into structured fields (e.g., output voltage as a range with separate min/max values).\n",
"3. **Interactive Verification:** The engineer can query the agent (via chat) for further details or clarification (e.g., \"Show me the detailed pin configuration extraction\") and review the cited pages.\n",
"\n",
"**Value Delivered:**\n",
"- Up to 70% reduction in manual data extraction time.\n",
"- Increased accuracy and standardization with structured fields.\n",
"- Full traceability with page citations for each extracted detail."
"- Increased accuracy and standardization with structured fields."
]
},
{
@@ -90,8 +89,7 @@
"\n",
"- The **output_voltage** and **input_voltage** fields are structured as ranges with separate minimum and maximum values and a unit.\n",
"- The **pin_configuration** field is structured to include a pin count and a descriptive layout.\n",
"- Additional technical fields (e.g., dropout voltage, max current) are captured as numbers.\n",
"- We include a **page_citations** mapping to track where each field was extracted from.\n"
"- Additional technical fields (e.g., dropout voltage, max current) are captured as numbers.\n"
]
},
{
@@ -231,7 +229,6 @@
"- **input_voltage** as a structured range\n",
"- **pin_configuration** with a `pin_count` and `layout`\n",
"- **features** (if available)\n",
"- **page_citations** mapping each field to a list of page numbers\n",
"\n",
"This structured approach makes it easier to standardize the information for downstream integration and verification. Engineers can click on the cited page numbers (in a UI that supports it) to validate the extraction."
]
@@ -0,0 +1,450 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "00f6713b-2a32-4f8f-80e5-9a7d9b6e3b90",
"metadata": {},
"source": [
"# Solar Panel Datasheet Comparison Workflow\n",
"\n",
"<a href=\"https://colab.research.google.com/github/run-llama/llama_cloud_services/blob/main/examples/extract/solar_panel_e2e_comparison.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
"\n",
"\n",
"This notebook demonstrates an endtoend agentic workflow using LlamaExtract and the LlamaIndex eventdriven workflow framework. In this workflow, we:\n",
"\n",
"1. **Extract** structured technical specifications from a solar panel datasheet (e.g. a PDF downloaded from a vendor).\n",
"2. **Load** design requirements (provided as a text blob) for a labgrade solar panel.\n",
"3. **Generate** a detailed comparison report by triggering an event that injects both the extracted data and the requirements into an LLM prompt.\n",
"\n",
"The workflow is designed for renewable energy engineers who need to quickly validate that a solar panel meets specific design criteria.\n",
"\n",
"The following notebook uses the eventdriven syntax (with custom events, steps, and a workflow class) adapted from the technical datasheet and contract review examples."
]
},
{
"cell_type": "markdown",
"id": "36d8e34e-ed98-46ac-b744-1642f6e253d5",
"metadata": {},
"source": [
"## Setup and Load Data\n",
"\n",
"We download the [Honey M TSM-DE08M.08(II) datasheet](https://static.trinasolar.com/sites/default/files/EU_Datasheet_HoneyM_DE08M.08%28II%29_2021_A.pdf) as a PDF.\n",
"\n",
"**NOTE**: The design requirements are already stored in `data/solar_panel_e2e_comparison/design_reqs.txt`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1de7b1b3-c285-492c-8b2e-b37974b4fc63",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"--2025-04-01 14:47:56-- https://static.trinasolar.com/sites/default/files/EU_Datasheet_HoneyM_DE08M.08%28II%29_2021_A.pdf\n",
"Resolving static.trinasolar.com (static.trinasolar.com)... 47.246.23.232, 47.246.23.234, 47.246.23.227, ...\n",
"Connecting to static.trinasolar.com (static.trinasolar.com)|47.246.23.232|:443... connected.\n",
"WARNING: cannot verify static.trinasolar.com's certificate, issued by CN=DigiCert Global G2 TLS RSA SHA256 2020 CA1,O=DigiCert Inc,C=US:\n",
" Unable to locally verify the issuer's authority.\n",
"HTTP request sent, awaiting response... 200 OK\n",
"Length: 1888183 (1.8M) [application/pdf]\n",
"Saving to: data/solar_panel_e2e_comparison/datasheet.pdf\n",
"\n",
"data/solar_panel_e2 100%[===================>] 1.80M 7.47MB/s in 0.2s \n",
"\n",
"2025-04-01 14:47:56 (7.47 MB/s) - data/solar_panel_e2e_comparison/datasheet.pdf saved [1888183/1888183]\n",
"\n"
]
}
],
"source": [
"!wget https://static.trinasolar.com/sites/default/files/EU_Datasheet_HoneyM_DE08M.08%28II%29_2021_A.pdf -O data/solar_panel_e2e_comparison/datasheet.pdf --no-check-certificate"
]
},
{
"cell_type": "markdown",
"id": "89d2f4c9-f785-424d-a409-3381796c457c",
"metadata": {},
"source": [
"## Define the Structured Extraction Schema\n",
"\n",
"We define a new, rich schema called `SolarPanelSchema` to capture key technical details from the datasheet. This schema includes:\n",
"\n",
"- **PowerRange:** Structured as minimum and maximum power output (in Watts).\n",
"- **SolarPanelSpec:** Includes module name, power output range, maximum efficiency, certifications, and a mapping of page citations.\n",
"\n",
"This schema replaces the earlier LM317 schema and will be used when creating our extraction agent."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "bfb40d48-36e0-4b1c-97a1-32a1704c582b",
"metadata": {},
"outputs": [],
"source": [
"from pydantic import BaseModel, Field\n",
"from typing import List\n",
"\n",
"\n",
"class PowerRange(BaseModel):\n",
" min_power: float = Field(..., description=\"Minimum power output in Watts\")\n",
" max_power: float = Field(..., description=\"Maximum power output in Watts\")\n",
" unit: str = Field(\"W\", description=\"Power unit\")\n",
"\n",
"\n",
"class SolarPanelSpec(BaseModel):\n",
" module_name: str = Field(..., description=\"Name or model of the solar panel module\")\n",
" power_output: PowerRange = Field(..., description=\"Power output range\")\n",
" maximum_efficiency: float = Field(\n",
" ..., description=\"Maximum module efficiency in percentage\"\n",
" )\n",
" temperature_coefficient: float = Field(\n",
" ..., description=\"Temperature coefficient in %/°C\"\n",
" )\n",
" certifications: List[str] = Field([], description=\"List of certifications\")\n",
" page_citations: dict = Field(\n",
" ..., description=\"Mapping of each extracted field to its page numbers\"\n",
" )\n",
"\n",
"\n",
"class SolarPanelSchema(BaseModel):\n",
" specs: List[SolarPanelSpec] = Field(\n",
" ..., description=\"List of extracted solar panel specifications\"\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "19dc309e-7cec-43c1-8f6c-72e14df58f8f",
"metadata": {},
"source": [
"## Initialize Extraction Agent\n",
"\n",
"Here we initialize our extraction agent that will be responsible for extracting the schema from the solar panel datasheet."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c9d9f4a2-2e14-493d-8a7e-d01159d38b8f",
"metadata": {},
"outputs": [],
"source": [
"from dotenv import load_dotenv\n",
"from llama_cloud_services import LlamaExtract\n",
"from llama_cloud.core.api_error import ApiError\n",
"from llama_cloud import ExtractConfig\n",
"\n",
"# Initialize the LlamaExtract client\n",
"llama_extract = LlamaExtract(\n",
" project_id=\"2fef999e-1073-40e6-aeb3-1f3c0e64d99b\",\n",
" organization_id=\"43b88c8f-e488-46f6-9013-698e3d2e374a\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ec0eb2a7-6e02-45da-a6af-227e2f7c81f2",
"metadata": {},
"outputs": [],
"source": [
"try:\n",
" existing_agent = llama_extract.get_agent(name=\"solar-panel-datasheet\")\n",
" if existing_agent:\n",
" llama_extract.delete_agent(existing_agent.id)\n",
"except ApiError as e:\n",
" if e.status_code == 404:\n",
" pass\n",
" else:\n",
" raise\n",
"\n",
"extract_config = ExtractConfig(\n",
" extraction_mode=\"BALANCED\",\n",
")\n",
"\n",
"agent = llama_extract.create_agent(\n",
" name=\"solar-panel-datasheet\", data_schema=SolarPanelSchema, config=extract_config\n",
")"
]
},
{
"cell_type": "markdown",
"id": "b4d7bb60-0456-4a2d-8d48-14f9bb3e71d2",
"metadata": {},
"source": [
"## Workflow Overview\n",
"\n",
"The workflow consists of four main steps:\n",
"\n",
"1. **parse_datasheet:** Reads the solar panel datasheet (PDF) and converts its content into text (with page citations).\n",
"2. **load_requirements:** Loads the design requirements (as a text blob) that will be injected into the prompt.\n",
"3. **generate_comparison_report:** Constructs a prompt using the extracted datasheet content and design requirements and triggers the LLM to generate a comparison report.\n",
"4. **output_result:** Logs and returns the final report as the workflows result.\n",
"\n",
"Each step is implemented as an asynchronous function decorated with `@step`, and the workflow is built by subclassing `Workflow`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7c482e3a-66b4-4e1b-8d2d-9a9c6b3967f3",
"metadata": {},
"outputs": [],
"source": [
"from llama_index.core.workflow import (\n",
" Event,\n",
" StartEvent,\n",
" StopEvent,\n",
" Context,\n",
" Workflow,\n",
" step,\n",
")\n",
"from llama_index.llms.openai import OpenAI\n",
"from llama_index.core.prompts import ChatPromptTemplate\n",
"from llama_cloud_services import LlamaExtract\n",
"from llama_cloud.core.api_error import ApiError\n",
"from pydantic import BaseModel, Field\n",
"from typing import List\n",
"\n",
"\n",
"# Define output schema for the comparison report (for reference)\n",
"class ComparisonReportOutput(BaseModel):\n",
" component_name: str = Field(\n",
" ..., description=\"The name of the component being evaluated.\"\n",
" )\n",
" meets_requirements: bool = Field(\n",
" ...,\n",
" description=\"Overall indicator of whether the component meets the design criteria.\",\n",
" )\n",
" summary: str = Field(..., description=\"A brief summary of the evaluation results.\")\n",
" details: dict = Field(\n",
" ..., description=\"Detailed comparisons for each key parameter.\"\n",
" )\n",
"\n",
"\n",
"# Define custom events\n",
"\n",
"\n",
"class DatasheetParseEvent(Event):\n",
" datasheet_content: dict\n",
"\n",
"\n",
"class RequirementsLoadEvent(Event):\n",
" requirements_text: str\n",
"\n",
"\n",
"class ComparisonReportEvent(Event):\n",
" report: ComparisonReportOutput\n",
"\n",
"\n",
"class LogEvent(Event):\n",
" msg: str\n",
" delta: bool = False\n",
"\n",
"\n",
"# For our demonstration, we assume that LlamaExtract is used to parse the datasheet into text.\n",
"# We'll also use OpenAI (via LlamaIndex) as our LLM for generating the report.\n",
"\n",
"llm = OpenAI(model=\"gpt-4o\") # or your preferred model"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "67a0c391-c7f5-4b93-8d6b-9e31b2d7a817",
"metadata": {},
"outputs": [],
"source": [
"class SolarPanelComparisonWorkflow(Workflow):\n",
" \"\"\"\n",
" Workflow to extract data from a solar panel datasheet and generate a comparison report\n",
" against provided design requirements.\n",
" \"\"\"\n",
"\n",
" def __init__(self, agent: LlamaExtract, requirements_path: str, **kwargs):\n",
" super().__init__(**kwargs)\n",
" self.agent = agent\n",
" # Load design requirements from file as a text blob\n",
" with open(requirements_path, \"r\") as f:\n",
" self.requirements_text = f.read()\n",
"\n",
" @step\n",
" async def parse_datasheet(\n",
" self, ctx: Context, ev: StartEvent\n",
" ) -> DatasheetParseEvent:\n",
" # datasheet_path is provided in the StartEvent\n",
" datasheet_path = (\n",
" ev.datasheet_path\n",
" ) # e.g., \"./data/solar_panel_comparison/datasheet.pdf\"\n",
" extraction_result = await self.agent.aextract(datasheet_path)\n",
" datasheet_dict = (\n",
" extraction_result.data\n",
" ) # assumed to be a string with page citations\n",
" await ctx.set(\"datasheet_content\", datasheet_dict)\n",
" ctx.write_event_to_stream(LogEvent(msg=\"Datasheet parsed successfully.\"))\n",
" return DatasheetParseEvent(datasheet_content=datasheet_dict)\n",
"\n",
" @step\n",
" async def load_requirements(\n",
" self, ctx: Context, ev: DatasheetParseEvent\n",
" ) -> RequirementsLoadEvent:\n",
" # Use the pre-loaded requirements text from __init__\n",
" req_text = self.requirements_text\n",
" ctx.write_event_to_stream(LogEvent(msg=\"Design requirements loaded.\"))\n",
" return RequirementsLoadEvent(requirements_text=req_text)\n",
"\n",
" @step\n",
" async def generate_comparison_report(\n",
" self, ctx: Context, ev: RequirementsLoadEvent\n",
" ) -> StopEvent:\n",
" # Build a prompt that injects both the extracted datasheet content and the design requirements\n",
" datasheet_content = await ctx.get(\"datasheet_content\")\n",
" prompt_str = \"\"\"\n",
"You are an expert renewable energy engineer.\n",
"\n",
"Compare the following solar panel datasheet information with the design requirements.\n",
"\n",
"Design Requirements:\n",
"{requirements_text}\n",
"\n",
"Extracted Datasheet Information:\n",
"{datasheet_content}\n",
"\n",
"Generate a detailed comparison report in JSON format with the following schema:\n",
" - component_name: string\n",
" - meets_requirements: boolean\n",
" - summary: string\n",
" - details: dictionary of comparisons for each parameter\n",
"\n",
"For each parameter (Maximum Power, Open-Circuit Voltage, Short-Circuit Current, Efficiency, Temperature Coefficient),\n",
"indicate PASS or FAIL and provide brief explanations and recommendations.\n",
"\"\"\"\n",
"\n",
" # extract from contract\n",
" prompt = ChatPromptTemplate.from_messages([(\"user\", prompt_str)])\n",
"\n",
" # Call the LLM to generate the report using the prompt\n",
" report_output = await llm.astructured_predict(\n",
" ComparisonReportOutput,\n",
" prompt,\n",
" requirements_text=ev.requirements_text,\n",
" datasheet_content=str(datasheet_content),\n",
" )\n",
" ctx.write_event_to_stream(LogEvent(msg=\"Comparison report generated.\"))\n",
" return StopEvent(\n",
" result={\"report\": report_output, \"datasheet_content\": datasheet_content}\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "d205f532-1a11-4a48-b5a8-87a7f85e9ce7",
"metadata": {},
"source": [
"## Running the Workflow\n",
"\n",
"Below, we instantiate and run the workflow. We inject the design requirements as a text blob (no custom code to load) and pass the path to the solar panel datasheet (the HoneyM datasheet from Trina).\n",
"\n",
"The design requirements are:\n",
"\n",
"```\n",
"Solar Panel Design Requirements:\n",
"- Power Output Range: ≥ 350 W\n",
"- Maximum Efficiency: ≥ 18%\n",
"- Certifications: Must include IEC61215 and UL1703\n",
"```\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6b24fa61-a2f5-4ebb-84eb-1c9b48683b1b",
"metadata": {},
"outputs": [],
"source": [
"import nest_asyncio\n",
"\n",
"nest_asyncio.apply()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "be3ebad5-1f70-4671-a2ec-17bf9e4d788f",
"metadata": {},
"outputs": [],
"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",
"\n",
"# Instantiate the workflow\n",
"workflow = SolarPanelComparisonWorkflow(\n",
" agent=agent, requirements_path=requirements_path, verbose=True, timeout=120\n",
")\n",
"\n",
"# 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\"])"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "llama_parse",
"language": "python",
"name": "llama_parse"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
File diff suppressed because one or more lines are too long
Binary file not shown.

After

Width:  |  Height:  |  Size: 96 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 828 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 626 KiB

-3
View File
@@ -1,8 +1,5 @@
# LlamaExtract
> **⚠️ EXPERIMENTAL**
> This library is under active development with frequent breaking changes. APIs and functionality may change significantly between versions. If you're interested in being an early adopter, please contact us at [support@llamaindex.ai](mailto:support@llamaindex.ai) or join our [Discord](https://discord.com/invite/eN6D2HQ4aX).
LlamaExtract provides a simple API for extracting structured data from unstructured documents like PDFs, text files and images (upcoming).
## Quick Start
+2
View File
@@ -1,6 +1,7 @@
from llama_cloud_services.parse import LlamaParse
from llama_cloud_services.report import ReportClient, LlamaReport
from llama_cloud_services.extract import LlamaExtract, ExtractionAgent
from llama_cloud_services.constants import EU_BASE_URL
__all__ = [
"LlamaParse",
@@ -8,4 +9,5 @@ __all__ = [
"LlamaReport",
"LlamaExtract",
"ExtractionAgent",
"EU_BASE_URL",
]
+1
View File
@@ -0,0 +1 @@
EU_BASE_URL = "https://api.cloud.eu.llamaindex.ai"
+6 -2
View File
@@ -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"]
+73 -16
View File
@@ -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."""
+11 -1
View File
@@ -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