mirror of
https://github.com/run-llama/llama_cloud_services.git
synced 2026-07-19 16:43:32 -04:00
Compare commits
18 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 7551e9fc0f | |||
| dfd83c6039 | |||
| 31f54bca55 | |||
| b1ae7bb736 | |||
| 31fe12e0da | |||
| 90b0c5e295 | |||
| 79fe1930cf | |||
| ab225c3eab | |||
| 6f1de75909 | |||
| 230ed64e41 | |||
| ef126c3a93 | |||
| 51a7534733 | |||
| 4f5d2bde13 | |||
| 3d05fe5d77 | |||
| c16ca673af | |||
| 6619034bce | |||
| c56fb5d8f7 | |||
| b407a5edb5 |
@@ -31,7 +31,7 @@ jobs:
|
||||
|
||||
- name: Run Tests
|
||||
working-directory: py
|
||||
run: uv run pytest unit_tests/ tests/ -v
|
||||
run: make e2e
|
||||
|
||||
- name: Remove virtual environment
|
||||
working-directory: py
|
||||
|
||||
+93
-21
@@ -2,6 +2,29 @@
|
||||
|
||||
LlamaExtract provides a simple API for extracting structured data from unstructured documents like PDFs, text files and images.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Quick Start](#quick-start)
|
||||
- [Supported File Types](#supported-file-types)
|
||||
- [Different Input Types](#different-input-types)
|
||||
- [Async Extraction](#async-extraction)
|
||||
- [Core Concepts](#core-concepts)
|
||||
- [Defining Schemas](#defining-schemas)
|
||||
- [Using Pydantic (Recommended)](#using-pydantic-recommended)
|
||||
- [Using JSON Schema](#using-json-schema)
|
||||
- [Important restrictions on JSON/Pydantic Schema](#important-restrictions-on-jsonpydantic-schema)
|
||||
- [Extraction Configuration](#extraction-configuration)
|
||||
- [Configuration Options](#configuration-options)
|
||||
- [Extraction Agents (Advanced)](#extraction-agents-advanced)
|
||||
- [Creating Agents](#creating-agents)
|
||||
- [Agent Batch Processing](#agent-batch-processing)
|
||||
- [Updating Agent Schemas](#updating-agent-schemas)
|
||||
- [Managing Agents](#managing-agents)
|
||||
- [When to Use Agents vs Direct Extraction](#when-to-use-agents-vs-direct-extraction)
|
||||
- [Installation](#installation)
|
||||
- [Tips & Best Practices](#tips--best-practices)
|
||||
- [Additional Resources](#additional-resources)
|
||||
|
||||
## Quick Start
|
||||
|
||||
The simplest way to get started is to use the stateless API with the extraction configuration and the file/text to extract from:
|
||||
@@ -12,7 +35,7 @@ from llama_cloud import ExtractConfig, ExtractMode
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
# Initialize client
|
||||
extractor = LlamaExtract()
|
||||
extractor = LlamaExtract(api_key="YOUR_API_KEY")
|
||||
|
||||
|
||||
# Define schema using Pydantic
|
||||
@@ -64,7 +87,9 @@ result = extractor.extract(Resume, config, SourceText(text_content=text))
|
||||
|
||||
### Async Extraction
|
||||
|
||||
For better performance with multiple files or when integrating with async applications:
|
||||
For better performance with multiple files or when integrating with async applications.
|
||||
Here `queue_extraction` will enqueue the extraction jobs and exit. Alternatively, you
|
||||
can use `aextract` to poll for the job and return the extraction results.
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
@@ -80,10 +105,18 @@ async def extract_resumes():
|
||||
Resume, config, ["resume1.pdf", "resume2.pdf"]
|
||||
)
|
||||
print(f"Queued {len(jobs)} extraction jobs")
|
||||
return jobs
|
||||
|
||||
|
||||
# Run async function
|
||||
asyncio.run(extract_resumes())
|
||||
jobs = asyncio.run(extract_resumes())
|
||||
# Check job status
|
||||
for job in jobs:
|
||||
status = agent.get_extraction_job(job.id).status
|
||||
print(f"Job {job.id}: {status}")
|
||||
|
||||
# Get results when complete
|
||||
results = [agent.get_extraction_run_for_job(job.id) for job in jobs]
|
||||
```
|
||||
|
||||
## Core Concepts
|
||||
@@ -159,24 +192,6 @@ config = ExtractConfig(extraction_mode=ExtractMode.FAST)
|
||||
result = extractor.extract(schema, config, "resume.pdf")
|
||||
```
|
||||
|
||||
## Extraction Configuration
|
||||
|
||||
Configure how extraction is performed using `ExtractConfig`:
|
||||
|
||||
```python
|
||||
from llama_cloud import ExtractConfig, ExtractMode
|
||||
|
||||
# Fast extraction (lower accuracy, faster processing)
|
||||
fast_config = ExtractConfig(extraction_mode=ExtractMode.FAST)
|
||||
|
||||
# Balanced extraction (good balance of speed and accuracy)
|
||||
balanced_config = ExtractConfig(extraction_mode=ExtractMode.BALANCED)
|
||||
|
||||
# Use different configs for different needs
|
||||
result = extractor.extract(schema, fast_config, "simple_document.pdf")
|
||||
result = extractor.extract(schema, balanced_config, "complex_document.pdf")
|
||||
```
|
||||
|
||||
### Important restrictions on JSON/Pydantic Schema
|
||||
|
||||
_LlamaExtract only supports a subset of the JSON Schema specification._ While limited, it should
|
||||
@@ -194,6 +209,62 @@ be sufficient for a wide variety of use-cases.
|
||||
your extraction workflow to fit within these constraints, e.g. by extracting subset of fields
|
||||
and later merging them together.
|
||||
|
||||
## Extraction Configuration
|
||||
|
||||
Configure how extraction is performed using `ExtractConfig`. The schema is the most important part, but several configuration options can significantly impact the extraction process.
|
||||
|
||||
```python
|
||||
from llama_cloud import ExtractConfig, ExtractMode, ChunkMode, ExtractTarget
|
||||
|
||||
# Basic configuration
|
||||
config = ExtractConfig(
|
||||
extraction_mode=ExtractMode.BALANCED, # FAST, BALANCED, MULTIMODAL, PREMIUM
|
||||
extraction_target=ExtractTarget.PER_DOC, # PER_DOC, PER_PAGE
|
||||
system_prompt="Focus on the most recent data",
|
||||
page_range="1-5,10-15", # Extract from specific pages
|
||||
)
|
||||
|
||||
# Advanced configuration
|
||||
advanced_config = ExtractConfig(
|
||||
extraction_mode=ExtractMode.MULTIMODAL,
|
||||
chunk_mode=ChunkMode.PAGE, # PAGE, SECTION
|
||||
high_resolution_mode=True, # Better OCR accuracy
|
||||
invalidate_cache=False, # Bypass cached results
|
||||
cite_sources=True, # Enable source citations
|
||||
use_reasoning=True, # Enable reasoning (not in FAST mode)
|
||||
confidence_scores=True, # MULTIMODAL/PREMIUM only
|
||||
)
|
||||
```
|
||||
|
||||
### Key Configuration Options
|
||||
|
||||
**Extraction Mode**: Controls processing quality and speed
|
||||
|
||||
- `FAST`: Fastest processing, suitable for simple documents with no OCR
|
||||
- `BALANCED`: Good speed/accuracy tradeoff for text-rich documents
|
||||
- `MULTIMODAL`: For visually rich documents with text, tables, and images (recommended)
|
||||
- `PREMIUM`: Highest accuracy with OCR, complex table/header detection
|
||||
|
||||
**Extraction Target**: Defines extraction scope
|
||||
|
||||
- `PER_DOC`: Apply schema to entire document (default)
|
||||
- `PER_PAGE`: Apply schema to each page, returns array of results
|
||||
|
||||
**Advanced Options**:
|
||||
|
||||
- `system_prompt`: Additional system-level instructions
|
||||
- `page_range`: Specific pages to extract (e.g., "1,3,5-7,9")
|
||||
- `chunk_mode`: Document splitting strategy (`PAGE` or `SECTION`)
|
||||
- `high_resolution_mode`: Better OCR for small text (slower processing)
|
||||
|
||||
**Extensions** (return additional metadata):
|
||||
|
||||
- `cite_sources`: Source tracing for extracted fields
|
||||
- `use_reasoning`: Explanations for extraction decisions
|
||||
- `confidence_scores`: Quantitative confidence measures (MULTIMODAL/PREMIUM only)
|
||||
|
||||
For complete configuration options, advanced settings, and detailed examples, see the [LlamaExtract Configuration Documentation](https://docs.cloud.llamaindex.ai/llamaextract/features/options).
|
||||
|
||||
## Extraction Agents (Advanced)
|
||||
|
||||
For reusable extraction workflows, you can create extraction agents that encapsulate both schema and configuration:
|
||||
@@ -326,6 +397,7 @@ Another option (orthogonal to the above) is to break the document into smaller s
|
||||
|
||||
## Additional Resources
|
||||
|
||||
- [Extract Documentation](https://docs.cloud.llamaindex.ai/llamaextract/getting_started) - Details on Extract features, API and examples.
|
||||
- [Example Notebook](docs/examples-py/extract/resume_screening.ipynb) - Detailed walkthrough of resume parsing
|
||||
- [Example Application with TypeScript](./examples-ts/extract/) - End-to-end examples using LlamaExtract TypeScript client.
|
||||
- [Discord Community](https://discord.com/invite/eN6D2HQ4aX) - Get help and share feedback
|
||||
|
||||
+9
-5
@@ -4,11 +4,15 @@ help: ## Show all Makefile targets.
|
||||
@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
|
||||
|
||||
format: ## Run code autoformatters (black).
|
||||
pre-commit install
|
||||
git ls-files | xargs pre-commit run black --files
|
||||
uv run pre-commit install
|
||||
git ls-files | xargs uv run pre-commit run black --files
|
||||
|
||||
lint: ## Run linters: pre-commit (black, ruff, codespell) and mypy
|
||||
pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
|
||||
uv run pre-commit install && git ls-files | xargs uv run pre-commit run --show-diff-on-failure --files
|
||||
|
||||
test: ## Run tests via pytest
|
||||
pytest tests
|
||||
test: ## Run unit tests via pytest
|
||||
uv run pytest -v unit_tests/
|
||||
|
||||
.PHONY: e2e
|
||||
e2e: ## Run all tests. Run with high parallelism using xdist since tests are bottlenecked bound by the slow backend parsing
|
||||
uv run pytest -v -n 32 tests/
|
||||
|
||||
@@ -37,11 +37,10 @@ Example Usage:
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
import numbers
|
||||
from llama_cloud import ExtractRun
|
||||
from llama_cloud.types.agent_data import AgentData
|
||||
from llama_cloud.types.aggregate_group import AggregateGroup
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
from pydantic import BaseModel, Field, ValidationError, model_validator, ConfigDict
|
||||
from typing import (
|
||||
Generic,
|
||||
List,
|
||||
@@ -176,6 +175,16 @@ class TypedAgentDataItems(BaseModel, Generic[AgentDataT]):
|
||||
)
|
||||
|
||||
|
||||
class FieldCitation(BaseModel):
|
||||
page: Optional[int] = Field(
|
||||
None, description="The page number that the field occurred on"
|
||||
)
|
||||
matching_text: Optional[str] = Field(
|
||||
None,
|
||||
description="The original text this field's value was derived from",
|
||||
)
|
||||
|
||||
|
||||
class ExtractedFieldMetadata(BaseModel):
|
||||
"""
|
||||
Metadata for an extracted data field, such as confidence, and citation information.
|
||||
@@ -193,14 +202,14 @@ class ExtractedFieldMetadata(BaseModel):
|
||||
None,
|
||||
description="The confidence score for the field based on the extracted text only",
|
||||
)
|
||||
page_number: Optional[int] = Field(
|
||||
None, description="The page number that the field occurred on"
|
||||
)
|
||||
matching_text: Optional[str] = Field(
|
||||
citation: Optional[List[FieldCitation]] = Field(
|
||||
None,
|
||||
description="The original text this field's value was derived from",
|
||||
description="The citation for the field, including page number and matching text",
|
||||
)
|
||||
|
||||
# Forbid unknown keys to avoid swallowing nested dicts
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
ExtractedFieldMetaDataDict = Dict[
|
||||
str, Union[ExtractedFieldMetadata, Dict[str, Any], list[Any]]
|
||||
@@ -214,10 +223,12 @@ def parse_extracted_field_metadata(
|
||||
k: _parse_extracted_field_metadata_recursive(v)
|
||||
for k, v in field_metadata.items()
|
||||
if k not in _METADATA_FIELDS_SIBLING_TO_LEAF
|
||||
and k not in _ADDITIONAL_ROOT_METADATA_FIELDS
|
||||
}
|
||||
|
||||
|
||||
_METADATA_FIELDS_SIBLING_TO_LEAF = {"reasoning"}
|
||||
_ADDITIONAL_ROOT_METADATA_FIELDS = {"error"}
|
||||
|
||||
|
||||
def _parse_extracted_field_metadata_recursive(
|
||||
@@ -238,19 +249,10 @@ def _parse_extracted_field_metadata_recursive(
|
||||
if len(indicator_fields.intersection(field_value.keys())) > 0:
|
||||
try:
|
||||
merged = {**field_value, **additional_fields}
|
||||
allowed_fields = ExtractedFieldMetadata.model_fields.keys()
|
||||
merged = {k: v for k, v in merged.items() if k in allowed_fields}
|
||||
validated = ExtractedFieldMetadata.model_validate(merged)
|
||||
|
||||
# grab the citation from the array. This is just an array for backwards compatibility.
|
||||
if "citation" in field_value and len(field_value["citation"]) > 0:
|
||||
first_citation = field_value["citation"][0]
|
||||
if "page" in first_citation and isinstance(
|
||||
first_citation["page"], numbers.Number
|
||||
):
|
||||
validated.page_number = int(first_citation["page"]) # type: ignore
|
||||
if "matching_text" in first_citation and isinstance(
|
||||
first_citation["matching_text"], str
|
||||
):
|
||||
validated.matching_text = first_citation["matching_text"]
|
||||
return validated
|
||||
except ValidationError:
|
||||
pass
|
||||
@@ -340,6 +342,28 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
|
||||
description="Additional metadata about the extracted data, such as errors, tokens, etc.",
|
||||
)
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def _normalize_field_metadata_on_input(cls, value: Any) -> Any:
|
||||
# Ensure any inbound representation (including JSON round-trips)
|
||||
# gets normalized so nested dicts become ExtractedFieldMetadata where appropriate.
|
||||
if (
|
||||
isinstance(value, dict)
|
||||
and "field_metadata" in value
|
||||
and isinstance(value["field_metadata"], dict)
|
||||
):
|
||||
try:
|
||||
value = {
|
||||
**value,
|
||||
"field_metadata": parse_extracted_field_metadata(
|
||||
value["field_metadata"]
|
||||
),
|
||||
}
|
||||
except Exception:
|
||||
# Let pydantic surface detailed errors later rather than swallowing completely
|
||||
pass
|
||||
return value
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
@@ -395,11 +419,14 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
|
||||
"""
|
||||
file_id = file_id or result.file.id
|
||||
file_name = file_name or result.file.name
|
||||
job_id = result.job_id
|
||||
job_field_metadata = result.extraction_metadata.get("field_metadata", {})
|
||||
errors = job_field_metadata.get("error", None)
|
||||
if not isinstance(errors, str):
|
||||
errors = None
|
||||
|
||||
try:
|
||||
field_metadata = parse_extracted_field_metadata(
|
||||
result.extraction_metadata.get("field_metadata", {})
|
||||
)
|
||||
field_metadata = parse_extracted_field_metadata(job_field_metadata)
|
||||
except ValidationError:
|
||||
field_metadata = {}
|
||||
|
||||
@@ -408,11 +435,15 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
|
||||
return cls.create(
|
||||
data=data,
|
||||
status=status,
|
||||
field_metadata=field_metadata,
|
||||
field_metadata=job_field_metadata,
|
||||
file_id=file_id,
|
||||
file_name=file_name,
|
||||
file_hash=file_hash,
|
||||
metadata=metadata or {},
|
||||
metadata={
|
||||
**({"field_errors": errors} if errors else {}),
|
||||
"job_id": job_id,
|
||||
**(metadata or {}),
|
||||
},
|
||||
)
|
||||
except ValidationError as e:
|
||||
invalid_item = ExtractedData[Dict[str, Any]].create(
|
||||
|
||||
@@ -0,0 +1,187 @@
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Optional
|
||||
from pydantic import BaseModel
|
||||
from llama_cloud.client import AsyncLlamaCloud
|
||||
from llama_cloud.types import (
|
||||
ClassifierRule,
|
||||
ClassifyJobResults,
|
||||
ClassifyParsingConfiguration,
|
||||
StatusEnum,
|
||||
ClassifyJobWithStatus,
|
||||
File,
|
||||
)
|
||||
from llama_cloud.resources.classifier.client import OMIT
|
||||
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
|
||||
|
||||
|
||||
class ClassificationOutput(BaseModel):
|
||||
file_id: str
|
||||
classification: str
|
||||
|
||||
|
||||
class ClassifyClient:
|
||||
"""
|
||||
Experimental - Client for interacting with the LlamaCloud Classifier API.
|
||||
The Classification API is currently in beta and may change in the future without notice.
|
||||
|
||||
Args:
|
||||
client: The LlamaCloud client to use.
|
||||
project_id: The project ID to use.
|
||||
organization_id: The organization ID to use.
|
||||
polling_interval: The interval to poll for job completion in seconds.
|
||||
polling_timeout: The timeout for the job to complete in seconds.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: AsyncLlamaCloud,
|
||||
project_id: Optional[str] = None,
|
||||
organization_id: Optional[str] = None,
|
||||
polling_interval: float = 1.0,
|
||||
polling_timeout: float = POLLING_TIMEOUT_SECONDS,
|
||||
):
|
||||
self.client = client
|
||||
self.project_id = project_id
|
||||
self.organization_id = organization_id
|
||||
self.polling_interval = polling_interval
|
||||
self.file_client = FileClient(client, project_id, organization_id)
|
||||
self.polling_timeout = polling_timeout
|
||||
|
||||
async def aclassify_file_ids(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_ids: list[str],
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
) -> ClassifyJobResults:
|
||||
"""
|
||||
Classify a list of files by their IDs.
|
||||
Note that even if a job fails, some of the files may have been classified successfully.
|
||||
In this case, you may want to set raise_on_error to False and check the results for successful classifications.
|
||||
|
||||
Args:
|
||||
rules: The rules to use for classification.
|
||||
file_ids: The IDs of the files to classify.
|
||||
parsing_configuration: The parsing configuration to use for classification.
|
||||
raise_on_error: Whether to raise an error if the classification job fails.
|
||||
|
||||
Returns:
|
||||
The results of the classification job.
|
||||
"""
|
||||
classify_job = await self.client.classifier.create_classify_job(
|
||||
rules=rules,
|
||||
file_ids=file_ids,
|
||||
parsing_configuration=parsing_configuration or OMIT,
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
)
|
||||
|
||||
classify_job_with_status = await self._wait_for_job_completion(classify_job.id)
|
||||
|
||||
if raise_on_error and classify_job_with_status.status == StatusEnum.ERROR:
|
||||
raise ValueError(
|
||||
f"Error classifying files under job ID {classify_job_with_status.id}"
|
||||
)
|
||||
|
||||
results = await self.client.classifier.get_classification_job_results(
|
||||
classify_job_with_status.id,
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
def classify_file_ids(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_ids: list[str],
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
) -> ClassifyJobResults:
|
||||
with augment_async_errors():
|
||||
return asyncio.run(
|
||||
self.aclassify_file_ids(
|
||||
rules, file_ids, parsing_configuration, raise_on_error
|
||||
)
|
||||
)
|
||||
|
||||
async def aclassify_file_path(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_input_path: str,
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
) -> ClassifyJobResults:
|
||||
file = await self.file_client.upload_file(file_input_path)
|
||||
return await self.aclassify_file_ids(
|
||||
rules, [file.id], parsing_configuration, raise_on_error
|
||||
)
|
||||
|
||||
def classify_file_path(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_input_path: str,
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
) -> ClassifyJobResults:
|
||||
with augment_async_errors():
|
||||
return asyncio.run(
|
||||
self.aclassify_file_path(
|
||||
rules, file_input_path, parsing_configuration, raise_on_error
|
||||
)
|
||||
)
|
||||
|
||||
async def aclassify_file_paths(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_input_paths: list[str],
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
workers: int = DEFAULT_NUM_WORKERS,
|
||||
show_progress: bool = False,
|
||||
) -> ClassifyJobResults:
|
||||
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(
|
||||
rules, [file.id for file in files], parsing_configuration, raise_on_error
|
||||
)
|
||||
|
||||
def classify_file_paths(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_input_paths: list[str],
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
) -> ClassifyJobResults:
|
||||
with augment_async_errors():
|
||||
return asyncio.run(
|
||||
self.aclassify_file_paths(
|
||||
rules, file_input_paths, parsing_configuration, raise_on_error
|
||||
)
|
||||
)
|
||||
|
||||
async def _wait_for_job_completion(self, job_id: str) -> ClassifyJobWithStatus:
|
||||
job = await self.client.classifier.get_classify_job(
|
||||
job_id, project_id=self.project_id, organization_id=self.organization_id
|
||||
)
|
||||
start_time = time.time()
|
||||
while not is_terminal_status(job.status):
|
||||
polling_duration = time.time() - start_time
|
||||
if polling_duration > self.polling_timeout:
|
||||
raise TimeoutError(
|
||||
f"Job {job_id} timed out after {polling_duration} seconds"
|
||||
)
|
||||
await asyncio.sleep(self.polling_interval)
|
||||
job = await self.client.classifier.get_classify_job(
|
||||
job_id, project_id=self.project_id, organization_id=self.organization_id
|
||||
)
|
||||
return job
|
||||
@@ -1 +1,2 @@
|
||||
EU_BASE_URL = "https://api.cloud.eu.llamaindex.ai"
|
||||
POLLING_TIMEOUT_SECONDS = 300.0
|
||||
|
||||
@@ -33,9 +33,9 @@ from llama_cloud.client import AsyncLlamaCloud
|
||||
from llama_cloud.core.api_error import ApiError
|
||||
from llama_cloud_services.extract.utils import (
|
||||
JSONObjectType,
|
||||
augment_async_errors,
|
||||
ExperimentalWarning,
|
||||
)
|
||||
from llama_cloud_services.utils import augment_async_errors
|
||||
from llama_index.core.schema import BaseComponent
|
||||
from llama_index.core.async_utils import run_jobs
|
||||
from llama_index.core.bridge.pydantic import Field, PrivateAttr
|
||||
@@ -227,7 +227,7 @@ class SourceText:
|
||||
raise ValueError(f"Unsupported file type: {type(self.file)}")
|
||||
|
||||
|
||||
FileInput = Union[str, Path, BufferedIOBase, SourceText]
|
||||
FileInput = Union[str, Path, BufferedIOBase, SourceText, File]
|
||||
|
||||
|
||||
def run_in_thread(
|
||||
@@ -406,6 +406,8 @@ class ExtractionAgent:
|
||||
|
||||
async def _upload_file(self, file_input: FileInput) -> File:
|
||||
source_text = None
|
||||
if isinstance(file_input, File):
|
||||
return file_input
|
||||
if isinstance(file_input, SourceText):
|
||||
source_text = file_input
|
||||
elif isinstance(file_input, (str, Path)):
|
||||
@@ -533,7 +535,7 @@ class ExtractionAgent:
|
||||
|
||||
upload_tasks = [self._upload_file(file) for file in files]
|
||||
with augment_async_errors():
|
||||
uploaded_files = await run_jobs(
|
||||
uploaded_files: List[File] = await run_jobs(
|
||||
upload_tasks,
|
||||
workers=self.num_workers,
|
||||
desc="Uploading files",
|
||||
@@ -987,8 +989,13 @@ class LlamaExtract(BaseComponent):
|
||||
f"Could not determine file type. Please provide a filename with one of these supported extensions: {supported_list}"
|
||||
)
|
||||
|
||||
def _convert_file_to_file_data(self, file_input: FileInput) -> Union[FileData, str]:
|
||||
def _convert_file_to_file_data(
|
||||
self, file_input: FileInput
|
||||
) -> Union[FileData, str, File]:
|
||||
"""Convert FileInput to FileData or text string for stateless extraction."""
|
||||
if isinstance(file_input, File):
|
||||
return file_input
|
||||
|
||||
if isinstance(file_input, SourceText):
|
||||
if file_input.text_content is not None:
|
||||
return file_input.text_content
|
||||
@@ -1084,24 +1091,23 @@ class LlamaExtract(BaseComponent):
|
||||
for file_input in files:
|
||||
file_data_or_text = self._convert_file_to_file_data(file_input)
|
||||
|
||||
if isinstance(file_data_or_text, str):
|
||||
if isinstance(file_data_or_text, File):
|
||||
file_args = {"file_id": file_data_or_text.id}
|
||||
|
||||
elif isinstance(file_data_or_text, str):
|
||||
# It's text content
|
||||
job = await self._async_client.llama_extract.extract_stateless(
|
||||
project_id=self._project_id,
|
||||
organization_id=self._organization_id,
|
||||
data_schema=processed_schema,
|
||||
config=config,
|
||||
text=file_data_or_text,
|
||||
)
|
||||
file_args = {"text": file_data_or_text}
|
||||
else:
|
||||
# It's FileData
|
||||
job = await self._async_client.llama_extract.extract_stateless(
|
||||
project_id=self._project_id,
|
||||
organization_id=self._organization_id,
|
||||
data_schema=processed_schema,
|
||||
config=config,
|
||||
file=file_data_or_text,
|
||||
)
|
||||
file_args = {"file": file_data_or_text}
|
||||
|
||||
job = await self._async_client.llama_extract.extract_stateless(
|
||||
project_id=self._project_id,
|
||||
organization_id=self._organization_id,
|
||||
data_schema=processed_schema,
|
||||
config=config,
|
||||
**file_args,
|
||||
)
|
||||
jobs.append(job)
|
||||
|
||||
return jobs[0] if len(jobs) == 1 else jobs
|
||||
|
||||
@@ -1,12 +1,4 @@
|
||||
from typing import Any, Dict, List, Union, Generator
|
||||
from contextlib import contextmanager
|
||||
|
||||
# Asyncio error messages
|
||||
nest_asyncio_err = "cannot be called from a running event loop"
|
||||
nest_asyncio_msg = (
|
||||
"The event loop is already running. "
|
||||
"Add `import nest_asyncio; nest_asyncio.apply()` to your code to fix this issue."
|
||||
)
|
||||
from typing import Any, Dict, List, Union
|
||||
|
||||
|
||||
def is_jupyter() -> bool:
|
||||
@@ -19,17 +11,6 @@ def is_jupyter() -> bool:
|
||||
return False
|
||||
|
||||
|
||||
@contextmanager
|
||||
def augment_async_errors() -> Generator[None, None, None]:
|
||||
"""Context manager to add helpful information for errors due to nested event loops."""
|
||||
try:
|
||||
yield
|
||||
except RuntimeError as e:
|
||||
if nest_asyncio_err in str(e):
|
||||
raise RuntimeError(nest_asyncio_msg)
|
||||
raise
|
||||
|
||||
|
||||
JSONType = Union[Dict[str, Any], List[Any], str, int, float, bool, None]
|
||||
JSONObjectType = Dict[str, JSONType]
|
||||
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
from .client import FileClient
|
||||
|
||||
__all__ = ["FileClient"]
|
||||
@@ -0,0 +1,97 @@
|
||||
from io import BytesIO
|
||||
from typing import BinaryIO
|
||||
import os
|
||||
from llama_cloud.client import AsyncLlamaCloud
|
||||
from llama_cloud.types import File, FileCreate
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class FileClient:
|
||||
"""
|
||||
Higher-level client for interacting with the LlamaCloud Files API.
|
||||
Uses presigned URLs for uploads by default.
|
||||
|
||||
Args:
|
||||
client: The LlamaCloud client to use.
|
||||
project_id: The project ID to use.
|
||||
organization_id: The organization ID to use.
|
||||
use_presigned_url: Whether to use presigned URLs for uploads (set to False when uploading to BYOC deployments).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: AsyncLlamaCloud,
|
||||
project_id: Optional[str] = None,
|
||||
organization_id: Optional[str] = None,
|
||||
use_presigned_url: bool = True,
|
||||
):
|
||||
self.client = client
|
||||
self.project_id = project_id
|
||||
self.organization_id = organization_id
|
||||
self.use_presigned_url = use_presigned_url
|
||||
|
||||
async def get_file(self, file_id: str) -> File:
|
||||
return await self.client.files.get_file(
|
||||
file_id, project_id=self.project_id, organization_id=self.organization_id
|
||||
)
|
||||
|
||||
async def read_file_content(self, file_id: str) -> bytes:
|
||||
presigned_url = await self.client.files.read_file_content(
|
||||
file_id,
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
)
|
||||
httpx_client = self.client._client_wrapper.httpx_client
|
||||
response = await httpx_client.get(presigned_url.url)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
|
||||
async def upload_file(
|
||||
self, file_path: str, external_file_id: Optional[str] = None
|
||||
) -> File:
|
||||
external_file_id = external_file_id or file_path
|
||||
file_size = os.path.getsize(file_path)
|
||||
with open(file_path, "rb") as file:
|
||||
return await self.upload_buffer(file, external_file_id, file_size)
|
||||
|
||||
async def upload_bytes(self, bytes: bytes, external_file_id: str) -> File:
|
||||
return await self.upload_buffer(BytesIO(bytes), external_file_id, len(bytes))
|
||||
|
||||
async def upload_buffer(
|
||||
self,
|
||||
buffer: BinaryIO,
|
||||
external_file_id: str,
|
||||
file_size: int,
|
||||
) -> File:
|
||||
if self.use_presigned_url:
|
||||
if getattr(buffer, "name", None):
|
||||
name = os.path.basename(str(getattr(buffer, "name", external_file_id)))
|
||||
else:
|
||||
name = external_file_id
|
||||
presigned_url = await self.client.files.generate_presigned_url(
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
request=FileCreate(
|
||||
name=name,
|
||||
external_file_id=external_file_id,
|
||||
file_size=file_size,
|
||||
),
|
||||
)
|
||||
httpx_client = self.client._client_wrapper.httpx_client
|
||||
upload_response = await httpx_client.put(
|
||||
presigned_url.url,
|
||||
data=buffer.read(),
|
||||
)
|
||||
upload_response.raise_for_status()
|
||||
return await self.client.files.get_file(
|
||||
presigned_url.file_id,
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
)
|
||||
else:
|
||||
return await self.client.files.upload_file(
|
||||
upload_file=buffer,
|
||||
external_file_id=external_file_id,
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
)
|
||||
@@ -222,11 +222,18 @@ class LlamaCloudCompositeRetriever(BaseRetriever):
|
||||
rerank_config if rerank_config is not None else self._rerank_config
|
||||
)
|
||||
|
||||
# Inject rerank_top_n into rerank_config if specified
|
||||
if rerank_top_n is not None and rerank_top_n != OMIT:
|
||||
if rerank_config is None or rerank_config == OMIT:
|
||||
rerank_config = ReRankConfig(top_n=rerank_top_n)
|
||||
else:
|
||||
# Update existing rerank_config with top_n
|
||||
rerank_config = rerank_config.copy(update={"top_n": rerank_top_n})
|
||||
|
||||
if self._persisted:
|
||||
result = self._client.retrievers.retrieve(
|
||||
self.retriever.id, # type: ignore [union-attr]
|
||||
mode=mode,
|
||||
rerank_top_n=rerank_top_n,
|
||||
rerank_config=rerank_config,
|
||||
query=query_bundle.query_str,
|
||||
)
|
||||
@@ -234,7 +241,6 @@ class LlamaCloudCompositeRetriever(BaseRetriever):
|
||||
result = self._client.retrievers.direct_retrieve(
|
||||
project_id=self.project.id,
|
||||
mode=mode,
|
||||
rerank_top_n=rerank_top_n,
|
||||
rerank_config=rerank_config,
|
||||
query=query_bundle.query_str,
|
||||
pipelines=self.retriever.pipelines, # type: ignore [union-attr]
|
||||
@@ -263,19 +269,25 @@ class LlamaCloudCompositeRetriever(BaseRetriever):
|
||||
rerank_config if rerank_config is not None else self._rerank_config
|
||||
)
|
||||
|
||||
# Inject rerank_top_n into rerank_config if specified
|
||||
if rerank_top_n is not None and rerank_top_n != OMIT:
|
||||
if rerank_config is None or rerank_config == OMIT:
|
||||
rerank_config = ReRankConfig(top_n=rerank_top_n)
|
||||
else:
|
||||
# Update existing rerank_config with top_n
|
||||
rerank_config = rerank_config.copy(update={"top_n": rerank_top_n})
|
||||
|
||||
if self._persisted:
|
||||
result = await self._aclient.retrievers.retrieve(
|
||||
self.retriever.id, # type: ignore [union-attr]
|
||||
mode=mode,
|
||||
rerank_config=rerank_config,
|
||||
rerank_top_n=rerank_top_n,
|
||||
query=query_bundle.query_str,
|
||||
)
|
||||
else:
|
||||
result = await self._aclient.retrievers.direct_retrieve(
|
||||
project_id=self.project.id,
|
||||
mode=mode,
|
||||
rerank_top_n=rerank_top_n,
|
||||
rerank_config=rerank_config,
|
||||
query=query_bundle.query_str,
|
||||
pipelines=self.retriever.pipelines, # type: ignore [union-attr]
|
||||
|
||||
@@ -52,6 +52,10 @@ class PageItem(BaseModel):
|
||||
bBox: Optional[BBox] = Field(
|
||||
default=None, description="The bounding box of the item."
|
||||
)
|
||||
html: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The HTML-formatted content of the item. Only applicable for table items when output_tables_as_HTML=True.",
|
||||
)
|
||||
|
||||
|
||||
class ImageItem(BaseModel):
|
||||
@@ -108,7 +112,7 @@ class ChartItem(BaseModel):
|
||||
class Page(BaseModel):
|
||||
"""A page of the document."""
|
||||
|
||||
page: int = Field(description="The page number.")
|
||||
page: int = Field(default=0, description="The page number.")
|
||||
text: Optional[str] = Field(default=None, description="The text of the page.")
|
||||
md: Optional[str] = Field(default=None, description="The markdown of the page.")
|
||||
images: List[ImageItem] = Field(
|
||||
@@ -149,6 +153,12 @@ class Page(BaseModel):
|
||||
noTextContent: bool = Field(
|
||||
default=False, description="Whether the page has no text content."
|
||||
)
|
||||
isAudioTranscript: bool = Field(
|
||||
default=False, description="Whether the page is an audio transcript."
|
||||
)
|
||||
durationInSeconds: Optional[float] = Field(
|
||||
default=None, description="The duration of the audio transcript in seconds."
|
||||
)
|
||||
|
||||
|
||||
class JobResult(BaseModel):
|
||||
|
||||
@@ -1,12 +1,21 @@
|
||||
import os
|
||||
import importlib.metadata
|
||||
|
||||
from contextlib import contextmanager
|
||||
from typing import Generator
|
||||
import difflib
|
||||
from llama_cloud.types import StatusEnum
|
||||
import httpx
|
||||
import packaging.version
|
||||
from pydantic import BaseModel
|
||||
from typing import Any, Dict, List, Tuple, Type
|
||||
|
||||
# Asyncio error messages
|
||||
nest_asyncio_err = "cannot be called from a running event loop"
|
||||
nest_asyncio_msg = (
|
||||
"The event loop is already running. "
|
||||
"Add `import nest_asyncio; nest_asyncio.apply()` to your code to fix this issue."
|
||||
)
|
||||
|
||||
|
||||
def check_extra_params(
|
||||
model_cls: Type[BaseModel], data: Dict[str, Any]
|
||||
@@ -34,6 +43,25 @@ def check_extra_params(
|
||||
return extra_params, suggestions
|
||||
|
||||
|
||||
def is_terminal_status(status: StatusEnum) -> bool:
|
||||
"""
|
||||
Check if a status is terminal, i.e. the job is done and no more updates are expected.
|
||||
Note: this must be updated if the status enum is updated.
|
||||
|
||||
Args:
|
||||
status: The status to check
|
||||
|
||||
Returns:
|
||||
True if the status is terminal, False otherwise
|
||||
"""
|
||||
return status in {
|
||||
StatusEnum.SUCCESS,
|
||||
StatusEnum.ERROR,
|
||||
StatusEnum.CANCELLED,
|
||||
StatusEnum.PARTIAL_SUCCESS,
|
||||
}
|
||||
|
||||
|
||||
async def check_for_updates(client: httpx.AsyncClient, quiet: bool = True) -> bool:
|
||||
"""Check if an SDK update is available.
|
||||
|
||||
@@ -65,3 +93,14 @@ async def check_for_updates(client: httpx.AsyncClient, quiet: bool = True) -> bo
|
||||
elif not quiet:
|
||||
print(f"{package_name} is up to date")
|
||||
return False
|
||||
|
||||
|
||||
@contextmanager
|
||||
def augment_async_errors() -> Generator[None, None, None]:
|
||||
"""Context manager to add helpful information for errors due to nested event loops."""
|
||||
try:
|
||||
yield
|
||||
except RuntimeError as e:
|
||||
if nest_asyncio_err in str(e):
|
||||
raise RuntimeError(nest_asyncio_msg)
|
||||
raise
|
||||
|
||||
@@ -11,13 +11,13 @@ dev = [
|
||||
|
||||
[project]
|
||||
name = "llama-parse"
|
||||
version = "0.6.55"
|
||||
version = "0.6.60"
|
||||
description = "Parse files into RAG-Optimized formats."
|
||||
authors = [{name = "Logan Markewich", email = "logan@llamaindex.ai"}]
|
||||
requires-python = ">=3.9,<4.0"
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
dependencies = ["llama-cloud-services>=0.6.55"]
|
||||
dependencies = ["llama-cloud-services>=0.6.60"]
|
||||
|
||||
[project.scripts]
|
||||
llama-parse = "llama_parse.cli.main:parse"
|
||||
|
||||
+4
-2
@@ -5,6 +5,7 @@ build-backend = "hatchling.build"
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"pytest>=8.0.0,<9",
|
||||
"pytest-xdist>=3.6.1,<4",
|
||||
"pytest-asyncio",
|
||||
"ipykernel>=6.29.0,<7",
|
||||
"pre-commit==3.2.0",
|
||||
@@ -12,12 +13,13 @@ dev = [
|
||||
"deepdiff>=8.1.1,<9",
|
||||
"ipython>=8.12.3,<9",
|
||||
"jupyter>=1.1.1,<2",
|
||||
"mypy>=1.14.1,<2"
|
||||
"mypy>=1.14.1,<2",
|
||||
"pydantic-settings>=2.10.1"
|
||||
]
|
||||
|
||||
[project]
|
||||
name = "llama-cloud-services"
|
||||
version = "0.6.55"
|
||||
version = "0.6.61"
|
||||
description = "Tailored SDK clients for LlamaCloud services."
|
||||
authors = [{name = "Logan Markewich", email = "logan@runllama.ai"}]
|
||||
requires-python = ">=3.9,<4.0"
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
[pytest]
|
||||
asyncio_mode = auto
|
||||
@@ -0,0 +1,212 @@
|
||||
import os
|
||||
import pytest
|
||||
from llama_cloud.client import AsyncLlamaCloud
|
||||
from llama_cloud.types import Project, ClassifierRule, ClassifyJobResults
|
||||
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
|
||||
from tests.conftest import EndToEndTestSettings
|
||||
import nest_asyncio
|
||||
|
||||
# Skip all tests if API key is not set
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not os.getenv("LLAMA_CLOUD_API_KEY"), reason="LLAMA_CLOUD_API_KEY not set"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_llama_cloud_client(
|
||||
e2e_test_settings: EndToEndTestSettings,
|
||||
) -> AsyncLlamaCloud:
|
||||
return AsyncLlamaCloud(
|
||||
token=e2e_test_settings.LLAMA_CLOUD_API_KEY.get_secret_value(),
|
||||
base_url=e2e_test_settings.LLAMA_CLOUD_BASE_URL,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def project(
|
||||
async_llama_cloud_client: AsyncLlamaCloud, e2e_test_settings: EndToEndTestSettings
|
||||
) -> Project:
|
||||
projects = await async_llama_cloud_client.projects.list_projects(
|
||||
project_name=e2e_test_settings.LLAMA_CLOUD_PROJECT_NAME,
|
||||
organization_id=e2e_test_settings.LLAMA_CLOUD_ORGANIZATION_ID,
|
||||
)
|
||||
assert len(projects) == 1
|
||||
return projects[0]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def classify_client(
|
||||
async_llama_cloud_client: AsyncLlamaCloud, project: Project
|
||||
) -> ClassifyClient:
|
||||
return ClassifyClient(
|
||||
async_llama_cloud_client,
|
||||
project_id=project.id,
|
||||
organization_id=project.organization_id,
|
||||
polling_interval=1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def file_client(
|
||||
async_llama_cloud_client: AsyncLlamaCloud, project: Project
|
||||
) -> FileClient:
|
||||
return FileClient(
|
||||
async_llama_cloud_client,
|
||||
project_id=project.id,
|
||||
organization_id=project.organization_id,
|
||||
use_presigned_url=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def apply_nest_asyncio():
|
||||
nest_asyncio.apply()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def simple_pdf_file_path() -> str:
|
||||
return "tests/test_files/index/Simple PDF Slides.pdf"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def research_paper_path() -> str:
|
||||
return "tests/test_files/attention_is_all_you_need_chart.pdf"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def classification_rules() -> list[ClassifierRule]:
|
||||
return [
|
||||
ClassifierRule(
|
||||
type="number",
|
||||
description="Documents with numbers",
|
||||
),
|
||||
ClassifierRule(
|
||||
type="research_paper",
|
||||
description="Research papers",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
parameterize_sync_and_async = pytest.mark.parametrize("sync", [True, False])
|
||||
|
||||
|
||||
@parameterize_sync_and_async
|
||||
@pytest.mark.asyncio
|
||||
async def test_classify_file_ids(
|
||||
classify_client: ClassifyClient,
|
||||
file_client: FileClient,
|
||||
simple_pdf_file_path: str,
|
||||
research_paper_path: str,
|
||||
classification_rules: list[ClassifierRule],
|
||||
sync: bool,
|
||||
):
|
||||
"""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 the uploaded files
|
||||
if sync:
|
||||
results = classify_client.classify_file_ids(
|
||||
rules=classification_rules, file_ids=[pdf_file.id, research_paper_file.id]
|
||||
)
|
||||
else:
|
||||
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(
|
||||
classify_client: ClassifyClient,
|
||||
simple_pdf_file_path: str,
|
||||
classification_rules: list[ClassifierRule],
|
||||
sync: bool,
|
||||
):
|
||||
"""Test classifying a single file by path"""
|
||||
# Classify the file
|
||||
if sync:
|
||||
results = classify_client.classify_file_path(
|
||||
rules=classification_rules, file_input_path=simple_pdf_file_path
|
||||
)
|
||||
else:
|
||||
results = await classify_client.aclassify_file_path(
|
||||
rules=classification_rules, file_input_path=simple_pdf_file_path
|
||||
)
|
||||
|
||||
assert isinstance(results, ClassifyJobResults)
|
||||
assert len(results.items) == 1
|
||||
|
||||
# Verify the file got classified
|
||||
item = results.items[0]
|
||||
assert item.result.type == "number"
|
||||
|
||||
|
||||
@parameterize_sync_and_async
|
||||
@pytest.mark.asyncio
|
||||
async def test_classify_file_paths(
|
||||
classify_client: ClassifyClient,
|
||||
file_client: FileClient,
|
||||
simple_pdf_file_path: str,
|
||||
research_paper_path: str,
|
||||
classification_rules: list[ClassifierRule],
|
||||
sync: bool,
|
||||
):
|
||||
"""Test classifying multiple files by paths"""
|
||||
# Classify all test files
|
||||
if sync:
|
||||
results = classify_client.classify_file_paths(
|
||||
rules=classification_rules,
|
||||
file_input_paths=[simple_pdf_file_path, research_paper_path],
|
||||
)
|
||||
else:
|
||||
results = await classify_client.aclassify_file_paths(
|
||||
rules=classification_rules,
|
||||
file_input_paths=[simple_pdf_file_path, research_paper_path],
|
||||
)
|
||||
|
||||
assert isinstance(results, ClassifyJobResults)
|
||||
assert len(results.items) == 2
|
||||
|
||||
file_name_to_expected_type = {
|
||||
os.path.basename(simple_pdf_file_path): "number",
|
||||
os.path.basename(research_paper_path): "research_paper",
|
||||
}
|
||||
# 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]
|
||||
assert item.result.type == expected_type
|
||||
|
||||
|
||||
@parameterize_sync_and_async
|
||||
@pytest.mark.asyncio
|
||||
async def test_classify_empty_file_list(
|
||||
classify_client: ClassifyClient,
|
||||
classification_rules: list[ClassifierRule],
|
||||
sync: bool,
|
||||
):
|
||||
"""Test classifying an empty list of files"""
|
||||
# This should throw an error
|
||||
with pytest.raises(UnprocessableEntityError):
|
||||
if sync:
|
||||
classify_client.classify_file_ids(rules=classification_rules, file_ids=[])
|
||||
else:
|
||||
await classify_client.aclassify_file_ids(
|
||||
rules=classification_rules, file_ids=[]
|
||||
)
|
||||
@@ -0,0 +1,26 @@
|
||||
import pytest
|
||||
from llama_index.core.constants import DEFAULT_BASE_URL
|
||||
from pydantic import Field, SecretStr
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class EndToEndTestSettings(BaseSettings):
|
||||
LLAMA_CLOUD_BASE_URL: str = Field(
|
||||
description="The base URL of the LlamaCloud API", default=DEFAULT_BASE_URL
|
||||
)
|
||||
LLAMA_CLOUD_API_KEY: SecretStr = Field(
|
||||
description="The API key for the LlamaCloud API"
|
||||
)
|
||||
LLAMA_CLOUD_ORGANIZATION_ID: str | None = Field(
|
||||
description="The organization ID for the LlamaCloud API",
|
||||
default=None,
|
||||
)
|
||||
LLAMA_CLOUD_PROJECT_NAME: str = Field(
|
||||
description="The project name for the LlamaCloud API",
|
||||
default="framework_integration_test",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def e2e_test_settings() -> EndToEndTestSettings:
|
||||
return EndToEndTestSettings()
|
||||
@@ -27,8 +27,6 @@ api_key = os.environ.get("LLAMA_CLOUD_API_KEY", None)
|
||||
organization_id = os.environ.get("LLAMA_CLOUD_ORGANIZATION_ID", None)
|
||||
project_name = os.environ.get("LLAMA_CLOUD_PROJECT_NAME", "framework_integration_test")
|
||||
|
||||
print("api-key", api_key, "base-url", base_url)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def remote_file() -> Tuple[str, str]:
|
||||
@@ -414,12 +412,54 @@ async def test_composite_retriever(index_name: str):
|
||||
# Assertions to verify the retrieval
|
||||
assert len(nodes) >= 2
|
||||
|
||||
# Test additional rerank_top_n configurations to cover the injection logic
|
||||
|
||||
# Test retriever with only rerank_top_n=1 (no existing rerank_config)
|
||||
retriever_with_rerank_top_n = LlamaCloudCompositeRetriever(
|
||||
name="composite_retriever_test_2",
|
||||
project_name=project_name,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
create_if_not_exists=True,
|
||||
mode=CompositeRetrievalMode.FULL,
|
||||
rerank_top_n=1,
|
||||
)
|
||||
retriever_with_rerank_top_n.add_index(index1)
|
||||
retriever_with_rerank_top_n.add_index(index2)
|
||||
nodes = retriever_with_rerank_top_n.retrieve("Hello world.")
|
||||
assert len(nodes) <= 1 # Should be limited to 1 result by rerank_top_n
|
||||
|
||||
# Test retriever with both rerank_top_n and custom rerank_config
|
||||
custom_config = ReRankConfig(top_n=10, model="test-model")
|
||||
retriever_with_both = LlamaCloudCompositeRetriever(
|
||||
name="composite_retriever_test_3",
|
||||
project_name=project_name,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
create_if_not_exists=True,
|
||||
mode=CompositeRetrievalMode.FULL,
|
||||
rerank_top_n=2,
|
||||
rerank_config=custom_config,
|
||||
)
|
||||
retriever_with_both.add_index(index1)
|
||||
retriever_with_both.add_index(index2)
|
||||
nodes = retriever_with_both.retrieve("Hello world.")
|
||||
assert len(nodes) >= 2 # Should have results from both indices
|
||||
|
||||
# Retrieve nodes using the composite retriever
|
||||
nodes = await retriever.aretrieve("Hello world.")
|
||||
|
||||
# Assertions to verify the retrieval
|
||||
assert len(nodes) >= 2
|
||||
|
||||
# Test async retrieve with the rerank_top_n only retriever
|
||||
nodes = await retriever_with_rerank_top_n.aretrieve("Hello world.")
|
||||
assert len(nodes) >= 1
|
||||
|
||||
# Test async retrieve with the both rerank_top_n and rerank_config retriever
|
||||
nodes = await retriever_with_both.aretrieve("Hello world.")
|
||||
assert len(nodes) >= 2
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not base_url or not api_key, reason="No platform base url or api key set"
|
||||
|
||||
@@ -202,3 +202,12 @@ async def test_get_result(markdown_parser: LlamaParse) -> None:
|
||||
result = await markdown_parser.aget_result(expected.job_id)
|
||||
assert result.job_id == expected.job_id
|
||||
assert len(result.pages) == len(expected.pages)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_audio() -> None:
|
||||
parser = LlamaParse()
|
||||
filepath = "tests/test_files/hello_world.m4a"
|
||||
|
||||
result = await parser.aparse(filepath)
|
||||
assert result.job_id is not None
|
||||
|
||||
Binary file not shown.
@@ -0,0 +1,147 @@
|
||||
import os
|
||||
import pytest
|
||||
from io import BytesIO
|
||||
from llama_cloud.client import AsyncLlamaCloud
|
||||
from llama_cloud.types import Project, File
|
||||
from llama_cloud_services.files.client import FileClient
|
||||
from tests.conftest import EndToEndTestSettings
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def llama_cloud_client(
|
||||
e2e_test_settings: EndToEndTestSettings,
|
||||
) -> AsyncLlamaCloud:
|
||||
return AsyncLlamaCloud(
|
||||
token=e2e_test_settings.LLAMA_CLOUD_API_KEY.get_secret_value(),
|
||||
base_url=e2e_test_settings.LLAMA_CLOUD_BASE_URL,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def project(
|
||||
llama_cloud_client: AsyncLlamaCloud, e2e_test_settings: EndToEndTestSettings
|
||||
) -> Project:
|
||||
projects = await llama_cloud_client.projects.list_projects(
|
||||
project_name=e2e_test_settings.LLAMA_CLOUD_PROJECT_NAME,
|
||||
organization_id=e2e_test_settings.LLAMA_CLOUD_ORGANIZATION_ID,
|
||||
)
|
||||
assert len(projects) == 1
|
||||
return projects[0]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def file_client(
|
||||
llama_cloud_client: AsyncLlamaCloud, project: Project, use_presigned_url: bool
|
||||
) -> FileClient:
|
||||
return FileClient(
|
||||
llama_cloud_client,
|
||||
project_id=project.id,
|
||||
organization_id=project.organization_id,
|
||||
use_presigned_url=use_presigned_url,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_file() -> str:
|
||||
return "tests/test_files/index/Simple PDF Slides.pdf"
|
||||
|
||||
|
||||
parametrize_use_presigned_url = pytest.mark.parametrize(
|
||||
"use_presigned_url", [True, False]
|
||||
)
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_from_path(file_client: FileClient, test_file: str):
|
||||
"""Test uploading a file from file path"""
|
||||
external_file_id = f"test_upload_path_{os.getpid()}"
|
||||
uploaded_file = await file_client.upload_file(test_file, external_file_id)
|
||||
|
||||
assert isinstance(uploaded_file, File)
|
||||
expected_name = os.path.basename(test_file)
|
||||
assert uploaded_file.name == expected_name
|
||||
assert uploaded_file.external_file_id == external_file_id
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_bytes(
|
||||
file_client: FileClient, test_file: str, use_presigned_url: bool
|
||||
):
|
||||
"""Test uploading a file from bytes"""
|
||||
# Read file as bytes
|
||||
with open(test_file, "rb") as f:
|
||||
file_bytes = f.read()
|
||||
|
||||
external_file_id = f"test_upload_bytes_{os.getpid()}"
|
||||
uploaded_file = await file_client.upload_bytes(file_bytes, external_file_id)
|
||||
|
||||
assert isinstance(uploaded_file, File)
|
||||
expected_name = external_file_id if use_presigned_url else "upload"
|
||||
assert uploaded_file.name == expected_name
|
||||
assert uploaded_file.external_file_id == external_file_id
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_buffer(
|
||||
file_client: FileClient, test_file: str, use_presigned_url: bool
|
||||
):
|
||||
"""Test uploading a file from buffer"""
|
||||
# Read file as bytes and create buffer
|
||||
with open(test_file, "rb") as f:
|
||||
file_bytes = f.read()
|
||||
|
||||
buffer = BytesIO(file_bytes)
|
||||
file_size = len(file_bytes)
|
||||
external_file_id = f"test_upload_buffer_{os.getpid()}"
|
||||
|
||||
uploaded_file = await file_client.upload_buffer(buffer, external_file_id, file_size)
|
||||
|
||||
assert isinstance(uploaded_file, File)
|
||||
expected_name = external_file_id if use_presigned_url else "upload"
|
||||
assert uploaded_file.name == expected_name
|
||||
assert uploaded_file.external_file_id == external_file_id
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_file(file_client: FileClient, test_file: str):
|
||||
"""Test retrieving a file by ID"""
|
||||
# Upload a file first
|
||||
external_file_id = f"test_get_file_{os.getpid()}"
|
||||
uploaded_file = await file_client.upload_file(test_file, external_file_id)
|
||||
|
||||
# Retrieve the file by ID
|
||||
retrieved_file = await file_client.get_file(uploaded_file.id)
|
||||
|
||||
assert isinstance(retrieved_file, File)
|
||||
assert retrieved_file == uploaded_file
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_with_default_external_id(file_client: FileClient, test_file: str):
|
||||
"""Test uploading file with default external_file_id"""
|
||||
# Upload file without specifying external_file_id
|
||||
uploaded_file = await file_client.upload_file(test_file)
|
||||
|
||||
assert isinstance(uploaded_file, File)
|
||||
assert uploaded_file.name == os.path.basename(test_file)
|
||||
assert uploaded_file.external_file_id == test_file
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_file_content(file_client: FileClient, test_file: str):
|
||||
"""Test reading a file content"""
|
||||
# Upload a file first
|
||||
external_file_id = f"test_read_file_content_{os.getpid()}"
|
||||
uploaded_file = await file_client.upload_file(test_file, external_file_id)
|
||||
|
||||
# Read the file content
|
||||
file_content = await file_client.read_file_content(uploaded_file.id)
|
||||
with open(test_file, "rb") as f:
|
||||
expected_file_content = f.read()
|
||||
assert file_content == expected_file_content
|
||||
@@ -1,15 +1,18 @@
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import pytest
|
||||
from llama_cloud import ExtractRun, File
|
||||
from llama_cloud.types.agent_data import AgentData
|
||||
from llama_cloud.types.aggregate_group import AggregateGroup
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
|
||||
from llama_cloud_services.beta.agent_data.schema import (
|
||||
ExtractedData,
|
||||
ExtractedFieldMetadata,
|
||||
FieldCitation,
|
||||
InvalidExtractionData,
|
||||
TypedAgentData,
|
||||
TypedAggregateGroup,
|
||||
@@ -21,8 +24,8 @@ from llama_cloud_services.beta.agent_data.schema import (
|
||||
# Test data models
|
||||
class Person(BaseModel):
|
||||
name: str
|
||||
age: int
|
||||
email: str
|
||||
age: Optional[int] = None
|
||||
email: Optional[str] = None
|
||||
|
||||
|
||||
class Company(BaseModel):
|
||||
@@ -81,8 +84,12 @@ def test_extracted_data_create_method():
|
||||
|
||||
# Test with custom values using ExtractedFieldMetadata
|
||||
field_metadata = {
|
||||
"name": ExtractedFieldMetadata(confidence=0.99, page_number=1),
|
||||
"age": ExtractedFieldMetadata(confidence=0.85, page_number=1),
|
||||
"name": ExtractedFieldMetadata(
|
||||
confidence=0.99, citation=[FieldCitation(page=1)]
|
||||
),
|
||||
"age": ExtractedFieldMetadata(
|
||||
confidence=0.85, citation=[FieldCitation(page=1)]
|
||||
),
|
||||
}
|
||||
extracted_custom = ExtractedData.create(
|
||||
person, status="accepted", field_metadata=field_metadata
|
||||
@@ -254,14 +261,16 @@ def test_parse_extracted_field_metadata():
|
||||
# name should have parsed citation data
|
||||
assert isinstance(result["name"], ExtractedFieldMetadata)
|
||||
assert result["name"].confidence == 0.95
|
||||
assert result["name"].page_number == 1
|
||||
assert result["name"].matching_text == "John Smith"
|
||||
assert result["name"].citation == [
|
||||
FieldCitation(page=1, matching_text="John Smith")
|
||||
]
|
||||
|
||||
# age should handle float page number
|
||||
assert isinstance(result["age"], ExtractedFieldMetadata)
|
||||
assert result["age"].confidence == 0.87
|
||||
assert result["age"].page_number == 2 # Should be converted to int
|
||||
assert result["age"].matching_text == "25 years old"
|
||||
assert result["age"].citation == [
|
||||
FieldCitation(page=2, matching_text="25 years old")
|
||||
]
|
||||
|
||||
# email should handle empty citations
|
||||
assert isinstance(result["email"], ExtractedFieldMetadata)
|
||||
@@ -327,30 +336,38 @@ def test_parse_extracted_field_metadata_complex():
|
||||
reasoning="Combined key parametrics and construction from the datasheet for a structured title.",
|
||||
confidence=0.9470628580889779,
|
||||
extraction_confidence=0.9470628580889779,
|
||||
page_number=1,
|
||||
matching_text="PHE844/F844, Film, Metallized Polypropylene, Safety, 0.47 uF",
|
||||
citation=[
|
||||
FieldCitation(
|
||||
page=1,
|
||||
matching_text="PHE844/F844, Film, Metallized Polypropylene, Safety, 0.47 uF",
|
||||
)
|
||||
],
|
||||
),
|
||||
"manufacturer": ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.9997446550976602,
|
||||
extraction_confidence=0.9997446550976602,
|
||||
page_number=1,
|
||||
matching_text="YAGEO KEMET",
|
||||
citation=[FieldCitation(page=1, matching_text="YAGEO KEMET")],
|
||||
),
|
||||
"features": [
|
||||
ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.9999308195540074,
|
||||
extraction_confidence=0.9999308195540074,
|
||||
page_number=1,
|
||||
matching_text="Features</td><td>EMI Safety",
|
||||
citation=[
|
||||
FieldCitation(
|
||||
page=1,
|
||||
matching_text="Features</td><td>EMI Safety",
|
||||
)
|
||||
],
|
||||
),
|
||||
ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.8642493886452225,
|
||||
extraction_confidence=0.8642493886452225,
|
||||
page_number=1,
|
||||
matching_text="THB Performance</td><td>Yes",
|
||||
citation=[
|
||||
FieldCitation(page=1, matching_text="THB Performance</td><td>Yes")
|
||||
],
|
||||
),
|
||||
],
|
||||
"dimensions": {
|
||||
@@ -358,15 +375,13 @@ def test_parse_extracted_field_metadata_complex():
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.8986941382802304,
|
||||
extraction_confidence=0.8986941382802304,
|
||||
page_number=1,
|
||||
matching_text="L</td><td>41mm MAX",
|
||||
citation=[FieldCitation(page=1, matching_text="L</td><td>41mm MAX")],
|
||||
),
|
||||
"width": ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.9999377974447091,
|
||||
extraction_confidence=0.9999377974447091,
|
||||
page_number=1,
|
||||
matching_text="T</td><td>13mm MAX",
|
||||
citation=[FieldCitation(page=1, matching_text="T</td><td>13mm MAX")],
|
||||
),
|
||||
},
|
||||
}
|
||||
@@ -390,6 +405,7 @@ def create_file(
|
||||
|
||||
def create_extract_run(
|
||||
id: str = "extract-123",
|
||||
job_id: str = "job-123",
|
||||
data: Dict[str, Any] = {"name": "John Doe", "age": 30, "email": "john@example.com"},
|
||||
extraction_metadata: Dict[str, Any] = {
|
||||
"name": {
|
||||
@@ -408,6 +424,7 @@ def create_extract_run(
|
||||
return ExtractRun.parse_obj(
|
||||
{
|
||||
"id": id,
|
||||
"job_id": job_id,
|
||||
"data": data,
|
||||
"extraction_metadata": {
|
||||
"field_metadata": extraction_metadata,
|
||||
@@ -450,8 +467,9 @@ def test_extracted_data_from_extraction_result_success():
|
||||
# Verify field metadata was parsed
|
||||
assert isinstance(extracted.field_metadata["name"], ExtractedFieldMetadata)
|
||||
assert extracted.field_metadata["name"].confidence == 0.95
|
||||
assert extracted.field_metadata["name"].page_number == 1
|
||||
assert extracted.field_metadata["name"].matching_text == "John Doe"
|
||||
assert extracted.field_metadata["name"].citation == [
|
||||
FieldCitation(page=1, matching_text="John Doe")
|
||||
]
|
||||
|
||||
# Verify overall confidence was calculated
|
||||
expected_confidence = (0.95 + 0.87 + 0.92) / 3
|
||||
@@ -485,10 +503,9 @@ def test_extracted_data_from_extraction_result_invalid_data():
|
||||
# Create ExtractRun with data that doesn't match Person schema
|
||||
extract_run = create_extract_run(
|
||||
data={
|
||||
"name": "Valid Name",
|
||||
"missing_name": "Valid Name",
|
||||
"age": "not_a_number",
|
||||
"missing_email": True,
|
||||
}, # Invalid age, missing email
|
||||
}, # Invalid age, missing name
|
||||
extraction_metadata={
|
||||
"name": {"confidence": 0.9},
|
||||
},
|
||||
@@ -507,9 +524,8 @@ def test_extracted_data_from_extraction_result_invalid_data():
|
||||
assert isinstance(invalid_data, ExtractedData)
|
||||
assert invalid_data.status == "error"
|
||||
assert invalid_data.data == {
|
||||
"name": "Valid Name",
|
||||
"missing_name": "Valid Name",
|
||||
"age": "not_a_number",
|
||||
"missing_email": True,
|
||||
}
|
||||
assert invalid_data.file_id == "error-file"
|
||||
assert invalid_data.file_name == "bad_data.pdf"
|
||||
@@ -523,3 +539,77 @@ def test_extracted_data_from_extraction_result_invalid_data():
|
||||
assert isinstance(invalid_data.field_metadata["name"], ExtractedFieldMetadata)
|
||||
assert invalid_data.field_metadata["name"].confidence == 0.9
|
||||
assert invalid_data.overall_confidence == 0.9
|
||||
|
||||
|
||||
class Dimensions(BaseModel):
|
||||
length: Optional[str] = Field(
|
||||
None, description="Length in mm (Size, Longest Side, L)"
|
||||
)
|
||||
width: Optional[str] = Field(
|
||||
None, description="Width in mm (Breadth, Side Width, W)"
|
||||
)
|
||||
height: Optional[str] = Field(
|
||||
None, description="Height in mm (Thickness, Vertical Size, H)"
|
||||
)
|
||||
diameter: Optional[str] = Field(
|
||||
None,
|
||||
description="Diameter in mm (for radial or cylindrical types) (Outer Diameter, dt, OD, D, d<sub>t</sub>)",
|
||||
)
|
||||
lead_spacing: Optional[str] = Field(
|
||||
None, description="Lead spacing in mm (Pin Pitch, Terminal Gap, LS)"
|
||||
)
|
||||
|
||||
|
||||
class Capacitor(BaseModel):
|
||||
dimensions: Optional[Dimensions] = None
|
||||
|
||||
|
||||
def test_full_parse_nested_dimensions():
|
||||
with open(Path(__file__).parent.parent.parent / "data" / "capacitor.json") as f:
|
||||
data = json.load(f)
|
||||
result = ExtractedData.from_extraction_result(ExtractRun.parse_obj(data), Capacitor)
|
||||
expected = {
|
||||
"dimensions": {
|
||||
"diameter": ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=1.0,
|
||||
extraction_confidence=1.0,
|
||||
),
|
||||
"lead_spacing": ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.9999999031936799,
|
||||
extraction_confidence=0.9999999031936799,
|
||||
),
|
||||
"length": ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.9999968039036192,
|
||||
extraction_confidence=0.9999968039036192,
|
||||
),
|
||||
}
|
||||
}
|
||||
assert result.field_metadata == expected
|
||||
parsed = ExtractedData.model_validate_json(result.model_dump_json())
|
||||
assert parsed.field_metadata == expected
|
||||
|
||||
|
||||
def test_parses_field_metadata_with_error_field():
|
||||
extract_run = create_extract_run(
|
||||
extraction_metadata={
|
||||
"name": {
|
||||
"confidence": 0.95,
|
||||
"citation": [{"page": 1, "matching_text": "John Smith"}],
|
||||
},
|
||||
"error": "This is an error",
|
||||
},
|
||||
)
|
||||
|
||||
parsed = ExtractedData.from_extraction_result(extract_run, Person)
|
||||
|
||||
assert parsed.field_metadata == {
|
||||
"name": ExtractedFieldMetadata(
|
||||
confidence=0.95,
|
||||
citation=[FieldCitation(page=1, matching_text="John Smith")],
|
||||
),
|
||||
}
|
||||
assert parsed.metadata.get("field_errors") == "This is an error"
|
||||
assert parsed.metadata.get("job_id") == "job-123"
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
{
|
||||
"id": "de058dda-6ca7-4eea-a426-da802f84f971",
|
||||
"created_at": "2025-08-13T15:45:39.286921Z",
|
||||
"updated_at": "2025-08-13T15:47:04.878069Z",
|
||||
"extraction_agent_id": "e834e99f-1f35-4748-b82f-03de4bd07ca6",
|
||||
"data_schema": {
|
||||
"additionalProperties": false,
|
||||
"properties": {
|
||||
"dimensions": {
|
||||
"anyOf": [
|
||||
{
|
||||
"additionalProperties": false,
|
||||
"properties": {
|
||||
"length": {
|
||||
"anyOf": [{ "type": "string" }, { "type": "null" }],
|
||||
"description": "Length in mm (Size, Longest Side, L)"
|
||||
},
|
||||
"width": {
|
||||
"anyOf": [{ "type": "string" }, { "type": "null" }],
|
||||
"description": "Width in mm (Breadth, Side Width, W)"
|
||||
},
|
||||
"height": {
|
||||
"anyOf": [{ "type": "string" }, { "type": "null" }],
|
||||
"description": "Height in mm (Thickness, Vertical Size, H)"
|
||||
},
|
||||
"diameter": {
|
||||
"anyOf": [{ "type": "string" }, { "type": "null" }],
|
||||
"description": "Diameter in mm (for radial or cylindrical types) (Outer Diameter, OD, D)"
|
||||
},
|
||||
"lead_spacing": {
|
||||
"anyOf": [{ "type": "string" }, { "type": "null" }],
|
||||
"description": "Lead spacing in mm (Pin Pitch, Terminal Gap, LS)"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"length",
|
||||
"width",
|
||||
"height",
|
||||
"diameter",
|
||||
"lead_spacing"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
{ "type": "null" }
|
||||
]
|
||||
}
|
||||
},
|
||||
"required": ["dimensions"],
|
||||
"type": "object"
|
||||
},
|
||||
"config": {
|
||||
"priority": null,
|
||||
"extraction_target": "PER_DOC",
|
||||
"extraction_mode": "PREMIUM",
|
||||
"multimodal_fast_mode": false,
|
||||
"system_prompt": "",
|
||||
"use_reasoning": true,
|
||||
"cite_sources": false,
|
||||
"confidence_scores": true,
|
||||
"chunk_mode": "PAGE",
|
||||
"high_resolution_mode": false,
|
||||
"invalidate_cache": false,
|
||||
"page_range": null
|
||||
},
|
||||
"file": {
|
||||
"id": "9233cc2b-00e3-4ddd-b426-b8b59357d4cb",
|
||||
"created_at": "2025-08-12T18:31:54.269440Z",
|
||||
"updated_at": "2025-08-13T15:45:39.064906Z",
|
||||
"name": "document (2).pdf.txt",
|
||||
"external_file_id": "document (2).pdf.txt",
|
||||
"file_size": 2562,
|
||||
"file_type": "txt",
|
||||
"project_id": "77bdc79f-fb69-49ae-a783-fcc573eec7ce",
|
||||
"last_modified_at": "2025-08-13T15:45:39Z",
|
||||
"resource_info": {
|
||||
"file_size": 2562,
|
||||
"last_modified_at": "2025-08-13T15:45:39"
|
||||
},
|
||||
"permission_info": null,
|
||||
"data_source_id": null
|
||||
},
|
||||
"status": "SUCCESS",
|
||||
"error": null,
|
||||
"job_id": "5bb8a583-366c-416c-ba55-4f5724fef9a9",
|
||||
"data": {
|
||||
"dimensions": {
|
||||
"length": "82 mm",
|
||||
"width": null,
|
||||
"height": null,
|
||||
"diameter": "35 mm",
|
||||
"lead_spacing": "6.0 mm"
|
||||
}
|
||||
},
|
||||
"extraction_metadata": {
|
||||
"field_metadata": {
|
||||
"dimensions": {
|
||||
"length": {
|
||||
"extraction_confidence": 0.9999968039036192,
|
||||
"confidence": 0.9999968039036192
|
||||
},
|
||||
"diameter": { "extraction_confidence": 1.0, "confidence": 1.0 },
|
||||
"lead_spacing": {
|
||||
"extraction_confidence": 0.9999999031936799,
|
||||
"confidence": 0.9999999031936799
|
||||
},
|
||||
"reasoning": "VERBATIM EXTRACTION"
|
||||
}
|
||||
},
|
||||
"usage": {
|
||||
"num_pages_extracted": 2,
|
||||
"num_document_tokens": 1034,
|
||||
"num_output_tokens": 3440
|
||||
}
|
||||
},
|
||||
"from_ui": false
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from llama_cloud.types import File as CloudFile
|
||||
from llama_cloud_services.extract import LlamaExtract, ExtractionAgent
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _set_dummy_env(monkeypatch):
|
||||
monkeypatch.setenv("LLAMA_CLOUD_API_KEY", "test-api-key")
|
||||
monkeypatch.setenv("LLAMA_CLOUD_BASE_URL", "https://example.test")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def llama_file() -> CloudFile:
|
||||
return CloudFile(
|
||||
id="file_123",
|
||||
name="sample.pdf",
|
||||
external_file_id="ext_123",
|
||||
project_id="proj_123",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def extractor() -> LlamaExtract:
|
||||
return LlamaExtract(
|
||||
api_key=os.environ["LLAMA_CLOUD_API_KEY"],
|
||||
base_url=os.environ["LLAMA_CLOUD_BASE_URL"],
|
||||
verify=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_external_validation(monkeypatch):
|
||||
import llama_cloud_services.extract.extract as extract_mod
|
||||
|
||||
async def _noop_validate_schema(client, data_schema):
|
||||
return data_schema
|
||||
|
||||
# Disable config warnings and external schema validation
|
||||
monkeypatch.setattr(
|
||||
extract_mod, "_extraction_config_warning", lambda *_args, **_kwargs: None
|
||||
)
|
||||
monkeypatch.setattr(extract_mod, "_validate_schema", _noop_validate_schema)
|
||||
|
||||
|
||||
def test_convert_fileinput_accepts_llama_file_directly(
|
||||
extractor: LlamaExtract, llama_file: CloudFile
|
||||
):
|
||||
result = extractor._convert_file_to_file_data(llama_file)
|
||||
assert result is llama_file
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queue_extraction_with_llama_file_uses_file_id(
|
||||
extractor: LlamaExtract, llama_file: CloudFile, no_external_validation, monkeypatch
|
||||
):
|
||||
calls = []
|
||||
|
||||
async def fake_extract_stateless(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return SimpleNamespace(id="job_1")
|
||||
|
||||
# Patch the client's method that would normally hit the network
|
||||
monkeypatch.setattr(
|
||||
extractor._async_client.llama_extract,
|
||||
"extract_stateless",
|
||||
fake_extract_stateless,
|
||||
)
|
||||
|
||||
# Minimal schema and dummy config (warnings disabled by fixture)
|
||||
schema = {"type": "object", "properties": {}}
|
||||
dummy_config = SimpleNamespace()
|
||||
|
||||
job = await extractor.queue_extraction(schema, dummy_config, llama_file)
|
||||
|
||||
assert getattr(job, "id") == "job_1"
|
||||
assert len(calls) == 1
|
||||
kwargs = calls[0]
|
||||
assert "file_id" in kwargs and kwargs["file_id"] == llama_file.id
|
||||
assert "file" not in kwargs
|
||||
assert "text" not in kwargs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extraction_agent_upload_file_accepts_llama_file_directly(
|
||||
llama_file: CloudFile,
|
||||
):
|
||||
# Build a minimal agent without hitting external services
|
||||
dummy_async_client = SimpleNamespace()
|
||||
dummy_agent = SimpleNamespace(id="agent_1", name="dummy", data_schema={}, config={})
|
||||
|
||||
agent = ExtractionAgent(
|
||||
client=dummy_async_client,
|
||||
agent=dummy_agent,
|
||||
project_id=None,
|
||||
organization_id=None,
|
||||
check_interval=0,
|
||||
max_timeout=0,
|
||||
num_workers=1,
|
||||
show_progress=False,
|
||||
verbose=False,
|
||||
verify=False,
|
||||
httpx_timeout=1,
|
||||
)
|
||||
|
||||
result = await agent._upload_file(llama_file)
|
||||
assert result is llama_file
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extraction_agent_aextract_accepts_llama_file(
|
||||
monkeypatch, llama_file: CloudFile
|
||||
):
|
||||
# Build a minimal agent without network
|
||||
dummy_llama_extract_iface = SimpleNamespace()
|
||||
|
||||
async def fake_run_job(**kwargs):
|
||||
# Ensure we are receiving a request with the right file_id
|
||||
request = kwargs.get("request")
|
||||
assert hasattr(request, "file_id")
|
||||
assert request.file_id == llama_file.id
|
||||
return SimpleNamespace(id="job_42")
|
||||
|
||||
dummy_llama_extract_iface.run_job = fake_run_job
|
||||
dummy_async_client = SimpleNamespace(llama_extract=dummy_llama_extract_iface)
|
||||
dummy_agent = SimpleNamespace(id="agent_1", name="dummy", data_schema={}, config={})
|
||||
|
||||
agent = ExtractionAgent(
|
||||
client=dummy_async_client,
|
||||
agent=dummy_agent,
|
||||
project_id=None,
|
||||
organization_id=None,
|
||||
check_interval=0,
|
||||
max_timeout=0,
|
||||
num_workers=1,
|
||||
show_progress=False,
|
||||
verbose=False,
|
||||
verify=False,
|
||||
httpx_timeout=1,
|
||||
)
|
||||
|
||||
# Ensure _upload_file returns the File directly and is called with our File
|
||||
calls = {}
|
||||
|
||||
async def fake_upload_file(file_input):
|
||||
calls["upload_called_with"] = file_input
|
||||
assert file_input is llama_file
|
||||
return file_input
|
||||
|
||||
monkeypatch.setattr(agent, "_upload_file", fake_upload_file)
|
||||
|
||||
# Avoid polling logic by short-circuiting result wait
|
||||
async def fake_wait(job_id: str):
|
||||
assert job_id == "job_42"
|
||||
return SimpleNamespace(id="run_42", status="SUCCESS", data={})
|
||||
|
||||
monkeypatch.setattr(agent, "_wait_for_job_result", fake_wait)
|
||||
|
||||
result = await agent.aextract(llama_file)
|
||||
|
||||
assert calls.get("upload_called_with") is llama_file
|
||||
assert getattr(result, "status") == "SUCCESS"
|
||||
assert getattr(result, "id") == "run_42"
|
||||
Generated
+41
-1
@@ -741,6 +741,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/36/f4/c6e662dade71f56cd2f3735141b265c3c79293c109549c1e6933b0651ffc/exceptiongroup-1.3.0-py3-none-any.whl", hash = "sha256:4d111e6e0c13d0644cad6ddaa7ed0261a0b36971f6d23e7ec9b4b9097da78a10", size = 16674, upload-time = "2025-05-10T17:42:49.33Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "execnet"
|
||||
version = "2.1.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/bb/ff/b4c0dc78fbe20c3e59c0c7334de0c27eb4001a2b2017999af398bf730817/execnet-2.1.1.tar.gz", hash = "sha256:5189b52c6121c24feae288166ab41b32549c7e2348652736540b9e6e7d4e72e3", size = 166524, upload-time = "2024-04-08T09:04:19.245Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/43/09/2aea36ff60d16dd8879bdb2f5b3ee0ba8d08cbbdcdfe870e695ce3784385/execnet-2.1.1-py3-none-any.whl", hash = "sha256:26dee51f1b80cebd6d0ca8e74dd8745419761d3bef34163928cbebbdc4749fdc", size = 40612, upload-time = "2024-04-08T09:04:17.414Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "executing"
|
||||
version = "2.2.0"
|
||||
@@ -1587,7 +1596,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "llama-cloud-services"
|
||||
version = "0.6.54"
|
||||
version = "0.6.61"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click", version = "8.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
|
||||
@@ -1612,8 +1621,10 @@ dev = [
|
||||
{ name = "jupyter" },
|
||||
{ name = "mypy" },
|
||||
{ name = "pre-commit" },
|
||||
{ name = "pydantic-settings" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "pytest-xdist" },
|
||||
]
|
||||
|
||||
[package.metadata]
|
||||
@@ -1638,8 +1649,10 @@ dev = [
|
||||
{ name = "jupyter", specifier = ">=1.1.1,<2" },
|
||||
{ name = "mypy", specifier = ">=1.14.1,<2" },
|
||||
{ name = "pre-commit", specifier = "==3.2.0" },
|
||||
{ name = "pydantic-settings", specifier = ">=2.10.1" },
|
||||
{ name = "pytest", specifier = ">=8.0.0,<9" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "pytest-xdist", specifier = ">=3.6.1,<4" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2855,6 +2868,20 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d4/29/3cade8a924a61f60ccfa10842f75eb12787e1440e2b8660ceffeb26685e7/pydantic_core-2.33.2-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:2807668ba86cb38c6817ad9bc66215ab8584d1d304030ce4f0887336f28a5e27", size = 2066661, upload-time = "2025-04-23T18:33:49.995Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pydantic-settings"
|
||||
version = "2.10.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pydantic" },
|
||||
{ name = "python-dotenv" },
|
||||
{ name = "typing-inspection" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/68/85/1ea668bbab3c50071ca613c6ab30047fb36ab0da1b92fa8f17bbc38fd36c/pydantic_settings-2.10.1.tar.gz", hash = "sha256:06f0062169818d0f5524420a360d632d5857b83cffd4d42fe29597807a1614ee", size = 172583, upload-time = "2025-06-24T13:26:46.841Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/58/f0/427018098906416f580e3cf1366d3b1abfb408a0652e9f31600c24a1903c/pydantic_settings-2.10.1-py3-none-any.whl", hash = "sha256:a60952460b99cf661dc25c29c0ef171721f98bfcb52ef8d9ea4c943d7c8cc796", size = 45235, upload-time = "2025-06-24T13:26:45.485Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pygments"
|
||||
version = "2.19.2"
|
||||
@@ -2896,6 +2923,19 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c7/9d/bf86eddabf8c6c9cb1ea9a869d6873b46f105a5d292d3a6f7071f5b07935/pytest_asyncio-1.1.0-py3-none-any.whl", hash = "sha256:5fe2d69607b0bd75c656d1211f969cadba035030156745ee09e7d71740e58ecf", size = 15157, upload-time = "2025-07-16T04:29:24.929Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest-xdist"
|
||||
version = "3.8.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "execnet" },
|
||||
{ name = "pytest" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/78/b4/439b179d1ff526791eb921115fca8e44e596a13efeda518b9d845a619450/pytest_xdist-3.8.0.tar.gz", hash = "sha256:7e578125ec9bc6050861aa93f2d59f1d8d085595d6551c2c90b6f4fad8d3a9f1", size = 88069, upload-time = "2025-07-01T13:30:59.346Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl", hash = "sha256:202ca578cfeb7370784a8c33d6d05bc6e13b4f25b5053c30a152269fd10f0b88", size = 46396, upload-time = "2025-07-01T13:30:56.632Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "python-dateutil"
|
||||
version = "2.9.0.post0"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "llama-cloud-services",
|
||||
"version": "0.3.0",
|
||||
"version": "0.3.3",
|
||||
"type": "module",
|
||||
"license": "MIT",
|
||||
"scripts": {
|
||||
|
||||
@@ -4,6 +4,8 @@ export type {
|
||||
AggregateAgentDataOptions,
|
||||
ComparisonOperator,
|
||||
ExtractedData,
|
||||
ExtractedFieldMetadata,
|
||||
ExtractedFieldMetadataDict,
|
||||
FilterOperation,
|
||||
SearchAgentDataOptions,
|
||||
StatusType,
|
||||
|
||||
@@ -38,8 +38,12 @@ export interface ExtractedFieldMetadata {
|
||||
confidence?: number;
|
||||
/** The confidence score for the field based on the extracted text only */
|
||||
extraction_confidence?: number;
|
||||
citation: FieldCitation[];
|
||||
}
|
||||
|
||||
export interface FieldCitation {
|
||||
/** The page number that the field occurred on */
|
||||
page_number?: number;
|
||||
page?: number;
|
||||
/** The original text this field's value was derived from */
|
||||
matching_text?: string;
|
||||
}
|
||||
|
||||
@@ -14,7 +14,8 @@ import {
|
||||
getJobTextResultApiV1ParsingJobJobIdResultTextGet,
|
||||
uploadFileApiV1ParsingUploadPost,
|
||||
} from "./api";
|
||||
import { sleep } from "./utils";
|
||||
import { sleep, getSavePath } from "./utils";
|
||||
import type { ParseResult } from "./type";
|
||||
|
||||
export type Language = ParserLanguages;
|
||||
export type ResultType = "text" | "markdown" | "json";
|
||||
@@ -594,15 +595,15 @@ export class LlamaParseReader extends FileReader {
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads data from a file and returns an array of JSON objects.
|
||||
* Loads data from a file and returns an array of ParseResult objects.
|
||||
* To be used with resultType "json".
|
||||
*
|
||||
* @param filePathOrContent - The file path or the file content as a Uint8Array.
|
||||
* @returns A Promise that resolves to an array of JSON objects.
|
||||
* @returns A Promise that resolves to an array of ParseResult objects, encapsulating the JSON array resulting from the parsing within the pages attribute.
|
||||
*/
|
||||
async loadJson(
|
||||
filePathOrContent: string | Uint8Array,
|
||||
): Promise<Record<string, any>[]> {
|
||||
): Promise<ParseResult[]> {
|
||||
let jobId;
|
||||
const isFilePath =
|
||||
typeof filePathOrContent === "string" &&
|
||||
@@ -628,7 +629,7 @@ export class LlamaParseReader extends FileReader {
|
||||
const resultJson = await this.getJobResult(jobId, "json");
|
||||
resultJson.job_id = jobId;
|
||||
resultJson.file_path = isFilePath ? filePathOrContent : undefined;
|
||||
return [resultJson];
|
||||
return [resultJson] as ParseResult[];
|
||||
} catch (e) {
|
||||
console.error(`Error while parsing the file under job id ${jobId}`, e);
|
||||
if (this.ignoreErrors) {
|
||||
@@ -639,6 +640,80 @@ export class LlamaParseReader extends FileReader {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads data from a file or file bytes (or an array of those) and returns an array of ParseResult objects.
|
||||
*
|
||||
* @param filePathOrContent - The file path or the file content as a Uint8Array, or an array of one of those two types.
|
||||
* @returns A Promise that resolves to an array of ParseResult objects, encapsulating the JSON array resulting from the parsing within the pages attribute.
|
||||
*/
|
||||
async parse(
|
||||
filePathOrContent: string | Uint8Array | string[] | Uint8Array[],
|
||||
): Promise<ParseResult[]> {
|
||||
const jsonResults: Record<string, any>[][] = [];
|
||||
if (!Array.isArray(filePathOrContent)) {
|
||||
const jsonResult = await this.loadJson(filePathOrContent);
|
||||
jsonResults.push(jsonResult);
|
||||
} else {
|
||||
for (let i = 0; i < filePathOrContent.length; i++) {
|
||||
console.log(
|
||||
`Processing file ${i + 1} of ${filePathOrContent.length}...`,
|
||||
);
|
||||
const jsonResult = await this.loadJson(
|
||||
filePathOrContent[i] as string | Uint8Array,
|
||||
);
|
||||
jsonResults.push(jsonResult);
|
||||
}
|
||||
}
|
||||
const parseResults: ParseResult[] = [];
|
||||
for (const jsonResult of jsonResults) {
|
||||
for (const result of jsonResult) {
|
||||
const parseResult = {
|
||||
pages: result.pages,
|
||||
job_metadata: result.job_metadata,
|
||||
job_id: result.job_id,
|
||||
file_path: result?.file_path ?? "",
|
||||
is_completed: true,
|
||||
} as ParseResult;
|
||||
parseResults.push(parseResult);
|
||||
}
|
||||
}
|
||||
return parseResults;
|
||||
}
|
||||
|
||||
/**
|
||||
* Downloads and saves tables from a given array of ParseResult to a specified download path.
|
||||
*
|
||||
* @param jsonResults - The array of ParseResult containing table information.
|
||||
* @param downloadPath - The path where the downloaded tables will be saved as CSV files.
|
||||
* @returns A Promise that resolves to an array of strings representing the paths to the tables.
|
||||
*/
|
||||
async getTables(
|
||||
jsonResults: ParseResult[],
|
||||
downloadPath: string,
|
||||
): Promise<string[]> {
|
||||
const tables: string[] = [];
|
||||
for (const result of jsonResults) {
|
||||
for (const page of result.pages) {
|
||||
if ("items" in page && Array.isArray(page.items)) {
|
||||
for (let i = 0; i < page.items.length; i++) {
|
||||
if (
|
||||
"type" in page.items[i] &&
|
||||
page.items[i].type === "table" &&
|
||||
"csv" in page.items[i] &&
|
||||
typeof page.items[i].csv === "string" &&
|
||||
page.items[i].csv != ""
|
||||
) {
|
||||
const savePath = getSavePath(downloadPath, i);
|
||||
await fs.writeFile(savePath, page.items[i].csv);
|
||||
tables.push(savePath);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return tables;
|
||||
}
|
||||
|
||||
/**
|
||||
* Downloads and saves images from a given JSON result to a specified download path.
|
||||
* Currently only supports resultType "json".
|
||||
@@ -648,7 +723,7 @@ export class LlamaParseReader extends FileReader {
|
||||
* @returns A Promise that resolves to an array of image objects.
|
||||
*/
|
||||
async getImages(
|
||||
jsonResult: Record<string, any>[],
|
||||
jsonResult: ParseResult[],
|
||||
downloadPath: string,
|
||||
): Promise<Record<string, any>[]> {
|
||||
try {
|
||||
|
||||
@@ -49,3 +49,13 @@ export type ExtractResult = {
|
||||
| null;
|
||||
};
|
||||
};
|
||||
|
||||
export type ParseResult = {
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
pages: Record<string, any>[];
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
job_metadata: Record<string, any>;
|
||||
job_id: string;
|
||||
is_completed: boolean;
|
||||
file_path: string;
|
||||
};
|
||||
|
||||
@@ -6,6 +6,8 @@ import {
|
||||
import { DEFAULT_BASE_URL } from "@llamaindex/core/global";
|
||||
import { getEnv } from "@llamaindex/env";
|
||||
import type { ClientParams } from "./type.js";
|
||||
import fs from "fs";
|
||||
import path from "path";
|
||||
|
||||
function getBaseUrl(baseUrl?: string): string {
|
||||
return baseUrl ?? getEnv("LLAMA_CLOUD_BASE_URL") ?? DEFAULT_BASE_URL;
|
||||
@@ -99,3 +101,19 @@ export async function getPipelineId(
|
||||
export async function sleep(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
export function getSavePath(downloadPath: string, i: number): string {
|
||||
const now = new Date();
|
||||
const formatted =
|
||||
now.toISOString().replace(/[-:T]/g, "_").replace(/\..+/, "") +
|
||||
"_" +
|
||||
now.getMilliseconds().toString().padStart(3, "0");
|
||||
|
||||
let savePath = path.join(downloadPath, `table_${formatted}.csv`);
|
||||
|
||||
if (fs.existsSync(savePath)) {
|
||||
savePath = savePath.replace(".csv", "_") + i.toString() + ".csv";
|
||||
}
|
||||
|
||||
return savePath;
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ import { LlamaExtract, LlamaExtractAgent } from "../src/LlamaExtract.js";
|
||||
import { Document } from "@llamaindex/core/schema";
|
||||
import { fs } from "@llamaindex/env";
|
||||
import { ExtractConfig } from "../src/api.js";
|
||||
import { ParseResult } from "../src/type.js";
|
||||
|
||||
// Integration tests that require actual API keys and files
|
||||
describe("Integration Tests", () => {
|
||||
@@ -284,6 +285,78 @@ describe("Integration Tests", () => {
|
||||
},
|
||||
60000,
|
||||
);
|
||||
|
||||
it.skipIf(skipIfNoApiKey)(
|
||||
"should parse a file and return a ParseResult array",
|
||||
async () => {
|
||||
const parseReader = new LlamaParseReader({
|
||||
apiKey: process.env.LLAMA_CLOUD_API_KEY!,
|
||||
verbose: false,
|
||||
});
|
||||
|
||||
const testContent = "Test document for JSON parsing";
|
||||
const testFilePath = "test-json.txt";
|
||||
|
||||
await fs.writeFile(testFilePath, new TextEncoder().encode(testContent));
|
||||
|
||||
try {
|
||||
const result = await parseReader.parse(testFilePath);
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(Array.isArray(result)).toBe(true);
|
||||
expect(result.length).toBeGreaterThan(0);
|
||||
expect(result[0]).toHaveProperty("job_id");
|
||||
expect(result[0]).toHaveProperty("job_metadata");
|
||||
expect(result[0]).toHaveProperty("file_path");
|
||||
expect(result[0]).toHaveProperty("pages");
|
||||
|
||||
await fs.unlink(testFilePath);
|
||||
} catch (error) {
|
||||
try {
|
||||
await fs.unlink(testFilePath);
|
||||
} catch {}
|
||||
throw error;
|
||||
}
|
||||
},
|
||||
60000,
|
||||
);
|
||||
|
||||
it.skipIf(skipIfNoApiKey)(
|
||||
"should extract tables correctly from a JSON result",
|
||||
async () => {
|
||||
const parseReader = new LlamaParseReader({
|
||||
apiKey: process.env.LLAMA_CLOUD_API_KEY!,
|
||||
verbose: false,
|
||||
});
|
||||
const pseudoJsonResult = [
|
||||
{
|
||||
pages: [
|
||||
{
|
||||
items: [
|
||||
{
|
||||
type: "table",
|
||||
csv: "Name,Age,Height (cm)\nAnna,12,140\nBob,22,175\nClaire,33,173\nDenis,44,185\n",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
job_id: "jobId",
|
||||
job_metadata: { job_id: "jobId" },
|
||||
file_path: "table.csv",
|
||||
is_completed: true,
|
||||
},
|
||||
] as ParseResult[];
|
||||
|
||||
const tmpdir = await fs.mkdtemp("tables");
|
||||
const result = await parseReader.getTables(pseudoJsonResult, tmpdir);
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(Array.isArray(result)).toBe(true);
|
||||
expect(result.length).toBeGreaterThan(0);
|
||||
expect(typeof result[0] === "string").toBe(true);
|
||||
},
|
||||
60000,
|
||||
);
|
||||
});
|
||||
|
||||
describe("LlamaCloudIndex Integration", () => {
|
||||
|
||||
Reference in New Issue
Block a user