Compare commits

...

28 Commits

Author SHA1 Message Date
Adrian Lyjak 9ce2044995 bump to 0.6.44 2025-07-08 17:40:39 -04:00
Adrian Lyjak 90d1608a71 Add nicer hand-written agent data interface 2025-07-08 21:12:42 +00:00
Logan 2448a42b90 relax pydantic job object (#784) 2025-07-08 12:12:56 -06:00
Neeraj Pradhan c75a900174 Bump up version to 0.6.42 (#783) 2025-07-08 09:16:46 -07:00
Peter Rowlands (변기호) 2fb7adfe0e parse: loosen PageItem.rows type hint (v0.6.41) (#776)
* parse: loosen PageItem.rows type hint

* bump version to 0.6.41
2025-06-30 21:47:40 +09:00
Pierre-Loic Doulcet dc82270724 header footer control in llamaparse (#775) 2025-06-30 16:02:59 +08:00
Neeraj Pradhan d880a48dd0 Bump to version 0.6.39 (#772)
* Bump to version 0.6.39

* lock file update
2025-06-27 16:04:40 -07:00
Logan 7567e8b45e except one more error type (#771) 2025-06-27 10:17:57 -06:00
Neeraj Pradhan 0d59a90151 Relax tenacity version; bump up version to 0.6.37 (#769) 2025-06-25 15:32:20 -07:00
Neeraj Pradhan 98ad550b1a Manage extract agent lifecycle in pytest (#766) 2025-06-24 08:59:38 -07:00
Neeraj Pradhan b58f43ce9f Bump up version to 0.6.36 (#763) 2025-06-23 14:26:05 -07:00
Neeraj Pradhan acf6adcd91 Make job fetching more robust to connection errors (#764) 2025-06-23 13:17:28 -07:00
Neeraj Pradhan daf6576c3c Bump version to 0.6.35 (#762) 2025-06-20 09:33:21 -07:00
Logan 8caa4defa6 fix partition (#758) 2025-06-16 17:37:52 -06:00
Pierre-Loic Doulcet 26918b8de4 add high_res_ocr to the package (#757) 2025-06-16 16:28:23 +08:00
Pierre-Loic Doulcet 6fb5ebe2f9 6.32 warning on unused parameters (#755) 2025-06-12 22:35:48 -06:00
dependabot[bot] c0aa67995b Bump requests from 2.32.3 to 2.32.4 in /llama_parse (#754) 2025-06-10 18:14:44 -06:00
dependabot[bot] 9f841f8328 Bump tornado from 6.4.2 to 6.5.1 in /llama_parse (#753) 2025-06-10 18:14:35 -06:00
dependabot[bot] 99c75eece9 Bump h11 from 0.14.0 to 0.16.0 in /llama_parse (#752) 2025-06-10 18:14:27 -06:00
Logan 57d2586ee3 v0.6.31 (#751) 2025-06-10 17:58:36 -06:00
Jerry Liu 4280a43ec8 add multi-fund analysis notebook (#739) 2025-06-07 11:25:25 -07:00
Neeraj Pradhan 7f1082bbb2 Bump to version 0.6.30 (#748) 2025-06-05 14:34:20 -07:00
Simon Suo 57cfc45804 Directly pass None project_id (#743) 2025-06-05 14:16:54 -07:00
Soumil.Binhani 30e8913875 0.6.29: Standerdize the parsing input format for both .aget_json() and .aload_data() (#745) 2025-06-05 10:58:07 -06:00
Logan 0ce6d4d7a4 more optional types marked (#747) 2025-06-05 10:50:29 -06:00
Peter Rowlands (변기호) 584ba8d48e 0.6.28: fix job result format after partitioning changes (#741)
* parse: fix job result format

* bump to 0.6.28
2025-06-02 15:25:30 -07:00
Peter Rowlands (변기호) 925805ee11 parse: support partitioning files before parsing (#709)
* parse: add utils for handling target_pages

* parse: support partitioning docs into multiple parse jobs

* tests: add tests for partitioned parse

* drop unneeded get_job_result call

* add parse JobFailedException and expected error handling

* bump to 0.6.27
2025-06-02 12:27:58 -07:00
Logan 76fb73c971 v0.6.26 (#740) 2025-06-02 09:59:45 -06:00
25 changed files with 3492 additions and 195 deletions
+1 -1
View File
@@ -41,7 +41,7 @@ jobs:
- name: Wait for PyPI to update
run: |
sleep 60
sleep 120
- name: Update llama-parse lock file
run: |
File diff suppressed because it is too large Load Diff
Binary file not shown.

After

Width:  |  Height:  |  Size: 3.3 MiB

@@ -0,0 +1,19 @@
from .schema import (
TypedAgentData,
ExtractedData,
TypedAgentDataItems,
StatusType,
ExtractedT,
AgentDataT,
)
from .client import AsyncAgentDataClient
__all__ = [
"TypedAgentData",
"AsyncAgentDataClient",
"ExtractedData",
"TypedAgentDataItems",
"StatusType",
"ExtractedT",
"AgentDataT",
]
@@ -0,0 +1,267 @@
import os
from typing import Dict, Generic, List, Optional, Type
from llama_cloud import FilterOperation
from llama_cloud.client import AsyncLlamaCloud
from tenacity import (
WrappedFn,
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
)
import httpx
from .schema import (
AgentDataT,
TypedAgentData,
TypedAgentDataItems,
TypedAggregateGroup,
TypedAggregateGroupItems,
)
def agent_data_retry(func: WrappedFn) -> WrappedFn:
"""
Decorator that adds automatic retry logic to agent data API calls.
Applies exponential backoff retry strategy for common network-related exceptions:
- Up to 3 retry attempts
- Exponential wait time between 0.5s and 10s
- Retries on timeout, connection, and HTTP status errors
This ensures resilient API communication in distributed environments where
temporary network issues or service unavailability may occur.
"""
return retry(
stop=stop_after_attempt(3),
wait=wait_exponential(min=0.5, max=10),
retry=retry_if_exception_type(
(httpx.TimeoutException, httpx.ConnectError, httpx.HTTPStatusError)
),
)(func)
def get_default_agent_id() -> Optional[str]:
"""
Retrieve the default agent ID from environment variables.
Returns:
The value of LLAMA_DEPLOY_DEPLOYMENT_NAME environment variable,
or None if not set
Note:
This provides a convenient way to configure agent ID globally
via environment variables instead of passing it explicitly
to each client instance.
"""
return os.getenv("LLAMA_DEPLOY_DEPLOYMENT_NAME")
class AsyncAgentDataClient(Generic[AgentDataT]):
"""
Async client for managing agent-generated structured data with type safety.
This client provides a high-level interface for CRUD operations, searching, and
aggregation of structured data created by agents. It enforces type safety by
validating all data against a specified Pydantic model type.
The client is generic over AgentDataT, which must be a Pydantic BaseModel that
defines the structure of your agent's data output.
Example:
```python
from pydantic import BaseModel
from llama_cloud.client import AsyncLlamaCloud
from llama_cloud_services.beta.agent_data import AsyncAgentDataClient
class ExtractedPerson(BaseModel):
name: str
age: int
email: str
# Initialize client
llama_client = AsyncLlamaCloud(token="your-api-key")
agent_client = AsyncAgentDataClient(
client=llama_client,
type=ExtractedPerson,
collection_name="extracted_people",
agent_url_id="person-extraction-agent"
)
# Create data
person = ExtractedPerson(name="John Doe", age=30, email="john@example.com")
result = await agent_client.create_agent_data(person)
# Search data
results = await agent_client.search_agent_data(
filter={"age": FilterOperation(gt=25)},
order_by="data.name",
page_size=20
)
```
Type Parameters:
AgentDataT: Pydantic BaseModel type that defines the structure of agent data
"""
def __init__(
self,
client: AsyncLlamaCloud,
type: Type[AgentDataT],
collection_name: str = "default",
agent_url_id: Optional[str] = None,
):
"""
Initialize the AsyncAgentDataClient.
Args:
client: AsyncLlamaCloud client instance for API communication
type: Pydantic BaseModel class that defines the data structure.
All agent data will be validated against this type.
collection_name: Named collection within the agent for organizing data.
Defaults to "default". Collections allow logical separation of
different data types or workflows within the same agent.
agent_url_id: Unique identifier for the agent. This normally appears in the
url of an agent within the llama cloud platform. If not provided,
will attempt to use the LLAMA_DEPLOY_DEPLOYMENT_NAME environment
variable. Data can only be added to an already existing agent in the
platform.
Raises:
ValueError: If agent_url_id is not provided and the
LLAMA_DEPLOY_DEPLOYMENT_NAME environment variable is not set
Note:
The client automatically applies retry logic to all API calls with
exponential backoff for timeout, connection, and HTTP status errors.
"""
self.agent_url_id = agent_url_id or get_default_agent_id()
if not self.agent_url_id:
raise ValueError(
"Agent ID is required, or set the LLAMA_DEPLOY_DEPLOYMENT_NAME environment variable"
)
self.collection_name = collection_name
self.client = client
self.type = type
@agent_data_retry
async def get_agent_data(self, item_id: str) -> TypedAgentData[AgentDataT]:
raw_data = await self.client.beta.get_agent_data(
item_id=item_id,
)
return TypedAgentData.from_raw(raw_data, validator=self.type)
@agent_data_retry
async def create_agent_data(self, data: AgentDataT) -> TypedAgentData[AgentDataT]:
raw_data = await self.client.beta.create_agent_data(
agent_slug=self.agent_url_id,
collection=self.collection_name,
data=data.model_dump(),
)
return TypedAgentData.from_raw(raw_data, validator=self.type)
@agent_data_retry
async def update_agent_data(
self, item_id: str, data: AgentDataT
) -> TypedAgentData[AgentDataT]:
raw_data = await self.client.beta.update_agent_data(
item_id=item_id,
data=data.model_dump(),
)
return TypedAgentData.from_raw(raw_data, validator=self.type)
@agent_data_retry
async def delete_agent_data(self, item_id: str) -> None:
await self.client.beta.delete_agent_data(item_id=item_id)
@agent_data_retry
async def search_agent_data(
self,
filter: Optional[Dict[str, Optional[FilterOperation]]] = None,
order_by: Optional[str] = None,
offset: Optional[int] = None,
page_size: Optional[int] = None,
include_total: bool = False,
) -> TypedAgentDataItems[AgentDataT]:
"""
Search agent data with filtering, sorting, and pagination.
Args:
filter: Filter conditions to apply to the search. Dict mapping field names to FilterOperation objects. Filters only by data fields
Examples:
- {"age": FilterOperation(gt=18)} - age greater than 18
- {"status": FilterOperation(eq="active")} - status equals "active"
- {"tags": FilterOperation(includes=["python", "ml"])} - tags include "python" or "ml"
- {"created_at": FilterOperation(gte="2024-01-01")} - created after date
- {"score": FilterOperation(lt=100, gte=50)} - score between 50 and 100
order_by: Comma delimited list of fields to sort results by. Can order by standard agent fields like created_at, or by data fields. Data fields must be prefixed with "data.". If ordering desceding, use a " desc" suffix.
Examples:
- "data.name desc, created_at" - sort by name in descending order, and then by creation date
page_size: Maximum number of items to return per page. Defaults to 10.
offset: Number of items to skip from the beginning. Defaults to 0.
include_total: Whether to include the total count in the response. Defaults to False to improve performance. It's recommended to only request on the first page.
"""
raw = await self.client.beta.search_agent_data_api_v_1_beta_agent_data_search_post(
agent_slug=self.agent_url_id,
collection=self.collection_name,
filter=filter,
order_by=order_by,
offset=offset,
page_size=page_size,
include_total=include_total,
)
return TypedAgentDataItems(
items=[
TypedAgentData.from_raw(item, validator=self.type) for item in raw.items
],
has_more=raw.next_page_token is not None,
total=raw.total_size,
)
@agent_data_retry
async def aggregate_agent_data(
self,
filter: Optional[Dict[str, Optional[FilterOperation]]] = None,
group_by: Optional[List[str]] = None,
count: Optional[bool] = None,
first: Optional[bool] = None,
order_by: Optional[str] = None,
offset: Optional[int] = None,
page_size: Optional[int] = None,
) -> TypedAggregateGroupItems[AgentDataT]:
"""
Aggregate agent data into groups according to the group_by fields.
Args:
filter: Filter conditions to apply to the search. Dict mapping field names to FilterOperation objects. Filters only by data fields
See search_agent_data for more details on filtering.
group_by: List of fields to group by. Groups strictly by equality. Can only group by data fields.
Examples:
- ["name"] - group by name
- ["name", "age"] - group by name and age
count: Whether to include the count of items in each group.
first: Whether to include the first item in each group.
order_by: Comma delimited list of fields to sort results by. See search_agent_data for more details on ordering.
offset: Number of groups to skip from the beginning. Defaults to 0.
page_size: Maximum number of groups to return per page.
"""
raw = await self.client.beta.aggregate_agent_data_api_v_1_beta_agent_data_aggregate_post(
agent_slug=self.agent_url_id,
collection=self.collection_name,
page_size=page_size,
filter=filter,
order_by=order_by,
group_by=group_by,
count=count,
first=first,
offset=offset,
)
return TypedAggregateGroupItems(
items=[
TypedAggregateGroup.from_raw(item, validator=self.type)
for item in raw.items
],
has_more=raw.next_page_token is not None,
total=raw.total_size,
)
@@ -0,0 +1,357 @@
"""
Agent Data API Schema Definitions
This module provides typed wrappers around the raw LlamaCloud agent data API,
enabling type-safe interactions with agent-generated structured data.
The agent data API serves as a persistent storage system for structured data
produced by LlamaCloud agents (particularly extraction agents). It provides
CRUD operations, search capabilities, filtering, and aggregation functionality
for managing agent-generated data at scale.
Key Concepts:
- Agent Slug: Unique identifier for an agent instance
- Collection: Named grouping of data within an agent (defaults to "default"). Data within a collection should be of the same type.
- Agent Data: Individual structured data records with metadata and timestamps
Example Usage:
```python
from pydantic import BaseModel
class Person(BaseModel):
name: str
age: int
client = AsyncAgentDataClient(
client=async_llama_cloud,
type=Person,
collection="people",
agent_url_id="my-extraction-agent-xyz"
)
# Create typed data
person = Person(name="John", age=30)
result = await client.create_agent_data(person)
print(result.data.name) # Type-safe access
```
"""
from datetime import datetime
from llama_cloud.types.agent_data import AgentData
from llama_cloud.types.aggregate_group import AggregateGroup
from pydantic import BaseModel, Field
from typing import (
Generic,
List,
Literal,
Optional,
Dict,
Type,
TypeVar,
Union,
Any,
)
# Type variable for user-defined data models
AgentDataT = TypeVar("AgentDataT", bound=BaseModel)
# Type variable for extracted data (can be dict or Pydantic model)
ExtractedT = TypeVar("ExtractedT", bound=Union[BaseModel, dict])
# Status types for extracted data workflow
StatusType = Union[Literal["error", "accepted", "rejected", "in_review"], str]
class TypedAgentData(BaseModel, Generic[AgentDataT]):
"""
Type-safe wrapper for agent data records.
This class represents a single data record stored in the agent data API,
combining the structured data payload with metadata about when and where
it was created.
Attributes:
id: Unique identifier for this data record
agent_url_id: Identifier of the agent that created this data
collection: Named collection within the agent (used for organization)
data: The actual structured data payload (typed as AgentDataT)
created_at: Timestamp when the record was first created
updated_at: Timestamp when the record was last modified
Example:
```python
# Access typed data
person_data: TypedAgentData[Person] = await client.get_agent_data(id)
print(person_data.data.name) # Type-safe access to Person fields
print(person_data.created_at) # Access metadata
```
"""
id: Optional[str] = Field(description="Unique identifier for this data record")
agent_url_id: str = Field(
description="Identifier of the agent that created this data"
)
collection: Optional[str] = Field(
description="Named collection within the agent for data organization"
)
data: AgentDataT = Field(description="The structured data payload")
created_at: Optional[datetime] = Field(description="When this record was created")
updated_at: Optional[datetime] = Field(
description="When this record was last modified"
)
@classmethod
def from_raw(
cls, raw_data: AgentData, validator: Type[AgentDataT]
) -> "TypedAgentData[AgentDataT]":
"""
Convert raw API response to typed agent data.
Args:
raw_data: Raw agent data from the API
validator: Pydantic model class to validate the data field
Returns:
TypedAgentData instance with validated data
"""
data: AgentDataT = validator.model_validate(raw_data.data)
return cls(
id=raw_data.id,
agent_url_id=raw_data.agent_slug,
collection=raw_data.collection,
data=data,
created_at=raw_data.created_at,
updated_at=raw_data.updated_at,
)
class TypedAgentDataItems(BaseModel, Generic[AgentDataT]):
"""
Paginated collection of agent data records.
This class represents a page of search results from the agent data API,
providing both the data records and pagination metadata.
Attributes:
items: List of agent data records in this page
total: Total number of records matching the query (only present if requested)
has_more: Whether there are more records available beyond this page
Example:
```python
# Search with pagination
results = await client.search_agent_data(
page_size=10,
include_total=True
)
for item in results.items:
print(item.data.name)
if results.has_more:
# Load next page
next_page = await client.search_agent_data(
page_size=10,
offset=10
)
```
"""
items: List[TypedAgentData[AgentDataT]] = Field(
description="List of agent data records in this page"
)
total: Optional[int] = Field(
description="Total number of records matching the query (only present if requested)"
)
has_more: bool = Field(
description="Whether there are more records available beyond this page"
)
class ExtractedData(BaseModel, Generic[ExtractedT]):
"""
Wrapper for extracted data with workflow status tracking.
This class is designed for extraction workflows where data goes through
review and approval stages. It maintains both the original extracted data
and the current state after any modifications.
Attributes:
original_data: The data as originally extracted from the source
data: The current state of the data (may differ from original after edits)
status: Current workflow status (in_review, accepted, rejected, error)
confidence: Confidence scores for individual fields (if available)
Status Workflow:
- "in_review": Initial state, awaiting human review
- "accepted": Data approved and ready for use
- "rejected": Data rejected, needs re-extraction or manual fix
- "error": Processing error occurred
Example:
```python
# Create extracted data for review
extracted = ExtractedData.create(
extracted_data=person_data,
status="in_review",
confidence={"name": 0.95, "age": 0.87}
)
# Later, after review
if extracted.status == "accepted":
# Use the data
process_person(extracted.data)
```
"""
original_data: ExtractedT = Field(
description="The original data that was extracted from the document"
)
data: ExtractedT = Field(
description="The latest state of the data. Will differ if data has been updated"
)
status: Union[Literal["error", "accepted", "rejected", "in_review"], str] = Field(
description="The status of the extracted data"
)
confidence: Dict[str, Union[float, Dict]] = Field(
default_factory=dict,
description="Confidence scores, if any, for each primitive field in the original_data data",
)
@classmethod
def create(
cls,
extracted_data: ExtractedT,
status: StatusType = "in_review",
confidence: Optional[Dict[str, Union[float, Dict]]] = None,
) -> "ExtractedData[ExtractedT]":
"""
Create a new ExtractedData instance with sensible defaults.
Args:
extracted_data: The extracted data payload
status: Initial workflow status
confidence: Optional confidence scores for fields
Returns:
New ExtractedData instance ready for storage
"""
return cls(
original_data=extracted_data,
data=extracted_data,
status=status,
confidence=confidence or {},
)
class TypedAggregateGroup(BaseModel, Generic[AgentDataT]):
"""
Represents a group of agent data records aggregated by common field values.
This class is used for grouping and analyzing agent data based on shared
characteristics. It's particularly useful for generating summaries and
statistics across large datasets.
Attributes:
group_key: The field values that define this group
count: Number of records in this group (if count aggregation was requested)
first_item: Representative data record from this group (if requested)
Example:
```python
# Group by age range
groups = await client.aggregate_agent_data(
group_by=["age_range"],
count=True,
first=True
)
for group in groups.items:
print(f"Age range {group.group_key['age_range']}: {group.count} people")
if group.first_item:
print(f"Example: {group.first_item.name}")
```
"""
group_key: Dict[str, Any] = Field(
description="The field values that define this group"
)
count: Optional[int] = Field(
description="Number of records in this group (if count aggregation was requested)"
)
first_item: Optional[AgentDataT] = Field(
description="Representative data record from this group (if requested)"
)
@classmethod
def from_raw(
cls, raw_data: AggregateGroup, validator: Type[AgentDataT]
) -> "TypedAggregateGroup[AgentDataT]":
"""
Convert raw API response to typed aggregate group.
Args:
raw_data: Raw aggregate group from the API
validator: Pydantic model class to validate the first_item field
Returns:
TypedAggregateGroup instance with validated first_item
"""
first_item: Optional[AgentDataT] = raw_data.first_item
if first_item is not None:
first_item = validator.model_validate(first_item)
return cls(
group_key=raw_data.group_key,
count=raw_data.count,
first_item=first_item,
)
class TypedAggregateGroupItems(BaseModel, Generic[AgentDataT]):
"""
Paginated collection of aggregate groups.
This class represents a page of aggregation results from the agent data API,
providing both the grouped data and pagination metadata.
Attributes:
items: List of aggregate groups in this page
total: Total number of groups matching the query (only present if requested)
has_more: Whether there are more groups available beyond this page
Example:
```python
# Get first page of groups
results = await client.aggregate_agent_data(
group_by=["department"],
count=True,
page_size=20
)
for group in results.items:
dept = group.group_key["department"]
print(f"{dept}: {group.count} employees")
# Load more if needed
if results.has_more:
next_page = await client.aggregate_agent_data(
group_by=["department"],
count=True,
page_size=20,
offset=20
)
```
"""
items: List[TypedAggregateGroup[AgentDataT]] = Field(
description="List of aggregate groups in this page"
)
total: Optional[int] = Field(
description="Total number of groups matching the query (only present if requested)"
)
has_more: bool = Field(
description="Whether there are more groups available beyond this page"
)
+66 -41
View File
@@ -8,6 +8,12 @@ import secrets
import warnings
import httpx
from pydantic import BaseModel
from tenacity import (
retry_if_exception,
stop_after_attempt,
wait_exponential_jitter,
AsyncRetrying,
)
from llama_cloud import (
ExtractAgent as CloudExtractAgent,
ExtractConfig,
@@ -17,12 +23,12 @@ from llama_cloud import (
File,
ExtractMode,
StatusEnum,
Project,
ExtractTarget,
LlamaExtractSettings,
PaginatedExtractRunsResponse,
)
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,
@@ -45,6 +51,17 @@ DEFAULT_EXTRACT_CONFIG = ExtractConfig(
)
def _is_retryable_error(exception: BaseException) -> bool:
"""Check if an exception is retryable."""
if isinstance(exception, ApiError):
return exception.status_code in (502, 503, 504, 425, 408)
elif isinstance(
exception, (httpx.HTTPStatusError, httpx.RequestError, httpx.TimeoutException)
):
return True
return False
class SourceText:
def __init__(
self,
@@ -120,9 +137,6 @@ def run_in_thread(
def _extraction_config_warning(config: ExtractConfig) -> None:
if config.extraction_mode == ExtractMode.ACCURATE:
warnings.warn("ACCURATE extraction mode is deprecated. Using BALANCED instead.")
config.extraction_mode = ExtractMode.BALANCED
if config.use_reasoning:
warnings.warn(
"`use_reasoning` is an experimental feature. Results will be available in "
@@ -232,9 +246,8 @@ class ExtractionAgent:
ValueError: If filename is not provided for bytes input or for file-like objects
without a name attribute.
"""
file_contents: Optional[Union[BufferedIOBase, BytesIO]] = None
try:
file_contents: Union[BufferedIOBase, BytesIO]
if file_input.text_content is not None:
# Handle direct text content
file_contents = BytesIO(file_input.text_content.encode("utf-8"))
@@ -261,7 +274,7 @@ class ExtractionAgent:
project_id=self._project_id, upload_file=file_contents
)
finally:
if isinstance(file_contents, BufferedReader):
if file_contents is not None and isinstance(file_contents, BufferedReader):
file_contents.close()
async def _upload_file(self, file_input: FileInput) -> File:
@@ -289,35 +302,60 @@ class ExtractionAgent:
return await self.upload_file(source_text)
async def _get_job_with_retry(self, job_id: str) -> ExtractJob:
"""Get job with retry logic for transient errors."""
async for attempt in AsyncRetrying(
retry=retry_if_exception(_is_retryable_error),
stop=stop_after_attempt(5),
wait=wait_exponential_jitter(initial=1, max=60, jitter=5),
reraise=True,
):
with attempt:
return await self._client.llama_extract.get_job(job_id=job_id)
async def _get_run_with_retry(self, job_id: str) -> ExtractRun:
"""Get extraction run with retry logic for transient errors."""
async for attempt in AsyncRetrying(
retry=retry_if_exception(_is_retryable_error),
stop=stop_after_attempt(3),
wait=wait_exponential_jitter(initial=1, max=20, jitter=3),
reraise=True,
):
with attempt:
return await self._client.llama_extract.get_run_by_job_id(job_id=job_id)
async def _wait_for_job_result(self, job_id: str) -> Optional[ExtractRun]:
"""Wait for and return the results of an extraction job."""
start = time.perf_counter()
tries = 0
while True:
await asyncio.sleep(self.check_interval)
tries += 1
job = await self._client.llama_extract.get_job(
job_id=job_id,
)
if job.status == StatusEnum.SUCCESS:
return await self._client.llama_extract.get_run_by_job_id(
job_id=job_id,
)
elif job.status == StatusEnum.PENDING:
end = time.perf_counter()
if end - start > self.max_timeout:
raise Exception(f"Timeout while extracting the file: {job_id}")
if self._verbose and tries % 10 == 0:
print(".", end="", flush=True)
continue
else:
warnings.warn(
f"Failure in job: {job_id}, status: {job.status}, error: {job.error}"
)
return await self._client.llama_extract.get_run_by_job_id(
job_id=job_id,
)
try:
job = await self._get_job_with_retry(job_id)
if job.status == StatusEnum.SUCCESS:
return await self._get_run_with_retry(job_id)
elif job.status == StatusEnum.PENDING:
end = time.perf_counter()
if end - start > self.max_timeout:
raise Exception(f"Timeout while extracting the file: {job_id}")
if self._verbose and tries % 10 == 0:
print(".", end="", flush=True)
continue
else:
warnings.warn(
f"Failure in job: {job_id}, status: {job.status}, error: {job.error}"
)
return await self._get_run_with_retry(job_id)
except Exception as e:
# If we get a non-retryable error or all retries are exhausted, re-raise
if self._verbose:
print(f"\nError in job polling for {job_id}: {e}")
raise e
def save(self) -> None:
"""Persist the extraction agent's schema and config to the database.
@@ -633,21 +671,8 @@ class LlamaExtract(BaseComponent):
self._thread_pool = ThreadPoolExecutor(
max_workers=min(10, (os.cpu_count() or 1) + 4)
)
# Fetch default project id if not provided
if not project_id:
project_id = os.getenv("LLAMA_CLOUD_PROJECT_ID", None)
if not project_id:
print("No project_id provided, fetching default project.")
projects: List[Project] = self._run_in_thread(
self._async_client.projects.list_projects()
)
default_project = [p for p in projects if p.is_default]
if not default_project:
raise ValueError(
"No default project found. Please provide a project_id."
)
project_id = default_project[0].id
self._project_id = project_id
self._organization_id = organization_id
+295 -61
View File
@@ -2,6 +2,7 @@ import asyncio
import mimetypes
import os
import time
import warnings
from contextlib import asynccontextmanager
from copy import deepcopy
from enum import Enum
@@ -13,21 +14,29 @@ from urllib.parse import urlparse
import httpx
from fsspec import AbstractFileSystem
from llama_index.core.async_utils import asyncio_run, run_jobs
from llama_index.core.bridge.pydantic import Field, PrivateAttr, field_validator
from llama_index.core.bridge.pydantic import (
Field,
PrivateAttr,
field_validator,
model_validator,
)
from llama_index.core.constants import DEFAULT_BASE_URL
from llama_index.core.readers.base import BasePydanticReader
from llama_index.core.readers.file.base import get_default_fs
from llama_index.core.schema import Document
from llama_cloud_services.utils import check_extra_params
from llama_cloud_services.parse.types import JobResult
from llama_cloud_services.parse.utils import (
SUPPORTED_FILE_TYPES,
ResultType,
ParsingMode,
FailedPageMode,
expand_target_pages,
nest_asyncio_err,
nest_asyncio_msg,
make_api_request,
partition_pages,
)
# can put in a path to the file or the file bytes itself
@@ -57,6 +66,36 @@ def build_url(
return base_url
class JobFailedException(Exception):
"""Parse job failed exception."""
def __init__(
self,
job_id: str,
status: str,
error_code: Optional[str] = None,
error_message: Optional[str] = None,
):
exception_str = (
f"Job ID: {job_id} failed with status: {status}, "
f'Error code: {error_code or "No error code found"}, '
f'Error message: {error_message or "No error message found"}'
)
super().__init__(exception_str)
self.job_id = job_id
self.status = status
self.error_code = error_code
self.error_message = error_message
@classmethod
def from_result(cls, result_json: Dict[str, Any]) -> "JobFailedException":
job_id = result_json["id"]
status = result_json["status"]
error_code = result_json.get("error_code")
error_message = result_json.get("error_message")
return cls(job_id, status, error_code=error_code, error_message=error_message)
class BackoffPattern(str, Enum):
"""Backoff pattern for polling."""
@@ -234,6 +273,10 @@ class LlamaParse(BasePydanticReader):
default=False,
description="Whether to guess the sheet names of the xlsx file.",
)
high_res_ocr: Optional[bool] = Field(
default=False,
description="If set to true, the parser will use high resolution OCR to extract text from images. This will increase the accuracy of the parsing job, but reduce the speed.",
)
html_make_all_elements_visible: Optional[bool] = Field(
default=False,
description="If set to true, when parsing HTML the parser will consider all elements display not element as display block.",
@@ -422,6 +465,34 @@ class LlamaParse(BasePydanticReader):
default=None,
description="A URL that needs to be called at the end of the parsing job.",
)
partition_pages: Optional[int] = Field(
default=None,
description="If set, documents will automatically be partitioned into segments containing the specified number of pages at most. Parsing will be split into separate jobs for each partition segment. Can be used in combination with targetPages and maxPages.",
)
hide_headers: Optional[bool] = Field(
default=False,
description="Whether to hide page header in output markdown.",
)
hide_footers: Optional[bool] = Field(
default=False,
description="Whether to hide page footers in output markdown.",
)
page_header_suffix: Optional[str] = Field(
default=None,
description="A suffix to add to the page header in the output markdown.",
)
page_header_prefix: Optional[str] = Field(
default=None,
description="A prefix to add to the page header in the output markdown.",
)
page_footer_suffix: Optional[str] = Field(
default=None,
description="A suffix to add to the page footer in the output markdown.",
)
page_footer_prefix: Optional[str] = Field(
default=None,
description="A prefix to add to the page footer in the output markdown.",
)
# Deprecated
bounding_box: Optional[str] = Field(
@@ -461,6 +532,21 @@ class LlamaParse(BasePydanticReader):
description="Whether to use the vendor multimodal API.",
)
@model_validator(mode="before")
@classmethod
def warn_extra_params(cls, data: Dict[str, Any]) -> Dict[str, Any]:
extra_params, suggestions = check_extra_params(cls, data)
if extra_params:
suggestions = [f"\n - {suggestion}" for suggestion in suggestions]
suggestions_str = "".join(suggestions)
warnings.warn(
"The following parameters are unused: "
+ ", ".join(extra_params)
+ f".\n{suggestions_str}",
)
return data
@field_validator("api_key", mode="before", check_fields=True)
@classmethod
def validate_api_key(cls, v: str) -> str:
@@ -548,6 +634,7 @@ class LlamaParse(BasePydanticReader):
file_input: FileInput,
extra_info: Optional[dict] = None,
fs: Optional[AbstractFileSystem] = None,
partition_target_pages: Optional[str] = None,
) -> str:
files = None
file_handle = None
@@ -698,6 +785,9 @@ class LlamaParse(BasePydanticReader):
if self.html_make_all_elements_visible:
data["html_make_all_elements_visible"] = self.html_make_all_elements_visible
if self.high_res_ocr:
data["high_res_ocr"] = self.high_res_ocr
if self.html_remove_fixed_elements:
data["html_remove_fixed_elements"] = self.html_remove_fixed_elements
@@ -769,6 +859,24 @@ class LlamaParse(BasePydanticReader):
if self.page_prefix is not None:
data["page_prefix"] = self.page_prefix
if self.hide_headers:
data["hide_headers"] = self.hide_headers
if self.hide_footers:
data["hide_footers"] = self.hide_footers
if self.page_header_suffix is not None:
data["page_header_suffix"] = self.page_header_suffix
if self.page_header_prefix is not None:
data["page_header_prefix"] = self.page_header_prefix
if self.page_footer_suffix is not None:
data["page_footer_suffix"] = self.page_footer_suffix
if self.page_footer_prefix is not None:
data["page_footer_prefix"] = self.page_footer_prefix
# only send page separator to server if it is not None
# as if a null, "" string is sent the server will then ignore the page separator instead of using the default
if self.page_separator is not None:
@@ -845,7 +953,9 @@ class LlamaParse(BasePydanticReader):
if self.take_screenshot:
data["take_screenshot"] = self.take_screenshot
if self.target_pages is not None:
if partition_target_pages is not None:
data["target_pages"] = partition_target_pages
elif self.target_pages is not None:
data["target_pages"] = self.target_pages
if self.user_prompt is not None:
data["user_prompt"] = self.user_prompt
@@ -942,15 +1052,7 @@ class LlamaParse(BasePydanticReader):
print(".", end="", flush=True)
current_interval = self._calculate_backoff(current_interval)
else:
error_code = result_json.get("error_code", "No error code found")
error_message = result_json.get(
"error_message", "No error message found"
)
exception_str = (
f"Job ID: {job_id} failed with status: {status}, "
f"Error code: {error_code}, Error message: {error_message}"
)
raise Exception(exception_str)
raise JobFailedException.from_result(result_json)
except (
httpx.ConnectError,
httpx.ReadError,
@@ -959,6 +1061,7 @@ class LlamaParse(BasePydanticReader):
httpx.ReadTimeout,
httpx.WriteTimeout,
httpx.HTTPStatusError,
httpx.RemoteProtocolError,
) as err:
error_count += 1
end = time.time()
@@ -979,9 +1082,39 @@ class LlamaParse(BasePydanticReader):
extra_info: Optional[dict] = None,
fs: Optional[AbstractFileSystem] = None,
result_type: Optional[str] = None,
num_workers: Optional[int] = None,
) -> List[Tuple[str, Dict[str, Any]]]:
if self.partition_pages is None:
job_results = [
await self._parse_one_unpartitioned(
file_path,
extra_info=extra_info,
fs=fs,
result_type=result_type,
)
]
else:
job_results = await self._parse_one_partitioned(
file_path,
extra_info,
fs=fs,
result_type=result_type,
num_workers=num_workers,
)
return job_results
async def _parse_one_unpartitioned(
self,
file_path: FileInput,
extra_info: Optional[dict] = None,
fs: Optional[AbstractFileSystem] = None,
result_type: Optional[str] = None,
**create_kwargs: Any,
) -> Tuple[str, Dict[str, Any]]:
"""Create one parse job and wait for the result."""
job_id = await self._create_job(file_path, extra_info=extra_info, fs=fs)
job_id = await self._create_job(
file_path, extra_info=extra_info, fs=fs, **create_kwargs
)
if self.verbose:
print("Started parsing the file under job_id %s" % job_id)
result = await self._get_job_result(
@@ -989,21 +1122,105 @@ class LlamaParse(BasePydanticReader):
)
return job_id, result
async def _parse_one_partitioned(
self,
file_path: FileInput,
extra_info: Optional[dict] = None,
fs: Optional[AbstractFileSystem] = None,
result_type: Optional[str] = None,
num_workers: Optional[int] = None,
) -> List[Tuple[str, Dict[str, Any]]]:
"""Partition a file and run separate parse jobs per partition segment."""
assert self.partition_pages is not None
num_workers = num_workers or self.num_workers
if num_workers < 1:
raise ValueError("Invalid number of workers")
if self.target_pages is not None:
jobs = [
self._parse_one_unpartitioned(
file_path,
extra_info=extra_info,
fs=fs,
result_type=result_type,
partition_target_pages=target_pages,
)
for target_pages in partition_pages(
expand_target_pages(self.target_pages),
self.partition_pages,
max_pages=self.max_pages,
)
]
return await run_jobs(
jobs,
workers=num_workers,
desc="Getting job results",
show_progress=self.show_progress,
)
total = 0
results: List[Tuple[str, Dict[str, Any]]] = []
while self.max_pages is None or total < self.max_pages:
if (
self.max_pages is not None
and total + self.partition_pages >= self.max_pages
):
size = self.max_pages - total
else:
size = self.partition_pages
if not size:
break
try:
# Fetch JSON result type first to get accurate pagination data
# and then fetch the user's desired result type if needed
job_id, json_result = await self._parse_one_unpartitioned(
file_path,
extra_info=extra_info,
fs=fs,
result_type=ResultType.JSON.value,
partition_target_pages=f"{total}-{total + size - 1}",
)
result_type = result_type or self.result_type.value
if result_type == ResultType.JSON.value:
job_result = json_result
else:
job_result = await self._get_job_result(
job_id, result_type, verbose=self.verbose
)
except JobFailedException as e:
if results and e.error_code == "NO_DATA_FOUND_IN_FILE":
# Expected when we try to read past the end of the file
return results
raise
results.append((job_id, job_result))
if len(json_result["pages"]) < size:
break
total += size
return results
async def _aload_data(
self,
file_path: FileInput,
extra_info: Optional[dict] = None,
fs: Optional[AbstractFileSystem] = None,
verbose: bool = False,
num_workers: Optional[int] = None,
) -> List[Document]:
"""Load data from the input path."""
try:
_job_id, result = await self._parse_one(
file_path, extra_info=extra_info, fs=fs
)
results = [
job_result
for _, job_result in await self._parse_one(
file_path, extra_info, fs=fs, num_workers=num_workers
)
]
# Flatten the resulting doc if it was partitioned
separator = self.page_separator or _DEFAULT_SEPARATOR
docs = [
Document(
text=result[self.result_type.value],
text=separator.join(
result[self.result_type.value] for result in results
),
metadata=extra_info or {},
)
]
@@ -1026,7 +1243,11 @@ class LlamaParse(BasePydanticReader):
extra_info: Optional[dict] = None,
fs: Optional[AbstractFileSystem] = None,
) -> List[Document]:
"""Load data from the input path."""
"""Load data from the input path.
File(s) which were partitioned before parsing will be loaded as a single
re-assembled Document.
"""
if isinstance(file_path, (str, PurePosixPath, Path, bytes, BufferedIOBase)):
return await self._aload_data(
file_path, extra_info=extra_info, fs=fs, verbose=self.verbose
@@ -1038,6 +1259,7 @@ class LlamaParse(BasePydanticReader):
extra_info=extra_info,
fs=fs,
verbose=self.verbose and not self.show_progress,
num_workers=1,
)
for f in file_path
]
@@ -1076,6 +1298,34 @@ class LlamaParse(BasePydanticReader):
else:
raise e
async def _aparse_one(
self,
file_path: FileInput,
file_name: str,
extra_info: Optional[dict] = None,
fs: Optional[AbstractFileSystem] = None,
num_workers: Optional[int] = None,
) -> List[JobResult]:
job_results = await self._parse_one(
file_path,
extra_info,
fs=fs,
result_type=ResultType.JSON.value,
num_workers=num_workers,
)
return [
JobResult(
job_id=job_id,
file_name=file_name,
job_result=job_result,
api_key=self.api_key,
base_url=self.base_url,
client=self.aclient,
page_separator=self.page_separator or _DEFAULT_SEPARATOR,
)
for job_id, job_result in job_results
]
async def aparse(
self,
file_path: Union[List[FileInput], FileInput],
@@ -1094,7 +1344,7 @@ class LlamaParse(BasePydanticReader):
fs: Optional filesystem to use for reading files.
Returns:
JobResult object or list of JobResult objects if multiple files were provided
JobResult object or list of JobResult objects if either multiple files were provided or file(s) were partitioned before parsing.
"""
if isinstance(file_path, (str, PurePosixPath, Path, bytes, BufferedIOBase)):
@@ -1106,22 +1356,10 @@ class LlamaParse(BasePydanticReader):
file_name = extra_info["file_name"]
else:
file_name = str(file_path)
job_id, job_result = await self._parse_one(
file_path,
extra_info=extra_info,
fs=fs,
result_type=ResultType.JSON.value,
)
return JobResult(
job_id=job_id,
file_name=file_name,
job_result=job_result,
api_key=self.api_key,
base_url=self.base_url,
client=self.aclient,
page_separator=self.page_separator or _DEFAULT_SEPARATOR,
result = await self._aparse_one(
file_path, file_name, extra_info=extra_info, fs=fs
)
return result[0] if len(result) == 1 else result
elif isinstance(file_path, list):
file_names = []
@@ -1135,35 +1373,25 @@ class LlamaParse(BasePydanticReader):
else:
file_names.append(str(f))
job_results = []
try:
job_results = await run_jobs(
for result in await run_jobs(
[
self._parse_one(
self._aparse_one(
f,
file_names[i],
extra_info=extra_info,
fs=fs,
result_type=ResultType.JSON.value,
num_workers=1,
)
for f in file_path
for i, f in enumerate(file_path)
],
workers=self.num_workers,
desc="Getting job results",
show_progress=self.show_progress,
)
# Create JobResults just using the job_ids and job_results
return [
JobResult(
job_id=job_id,
file_name=file_names[i],
job_result=job_result,
api_key=self.api_key,
base_url=self.base_url,
client=self.aclient,
page_separator=self.page_separator or _DEFAULT_SEPARATOR,
)
for i, (job_id, job_result) in enumerate(job_results)
]
):
job_results.extend(result)
return job_results
except RuntimeError as e:
if nest_asyncio_err in str(e):
@@ -1204,21 +1432,27 @@ class LlamaParse(BasePydanticReader):
raise e
async def _aget_json(
self, file_path: FileInput, extra_info: Optional[dict] = None
self,
file_path: FileInput,
extra_info: Optional[dict] = None,
num_workers: Optional[int] = None,
) -> List[dict]:
"""Load data from the input path."""
try:
job_id, result = await self._parse_one(
job_results = await self._parse_one(
file_path,
extra_info=extra_info,
result_type=ResultType.JSON.value,
num_workers=num_workers,
)
result["job_id"] = job_id
if not isinstance(file_path, (bytes, BufferedIOBase)):
result["file_path"] = str(file_path)
return [result]
results = []
for job_id, job_result in job_results:
job_result["job_id"] = job_id
if not isinstance(file_path, (bytes, BufferedIOBase)):
job_result["file_path"] = str(file_path)
results.append(job_result)
return results
except Exception as e:
file_repr = file_path if isinstance(file_path, str) else "<bytes/buffer>"
print(f"Error while parsing the file '{file_repr}':", e)
@@ -1233,7 +1467,7 @@ class LlamaParse(BasePydanticReader):
extra_info: Optional[dict] = None,
) -> List[dict]:
"""Load data from the input path."""
if isinstance(file_path, (str, Path)):
if isinstance(file_path, (str, PurePosixPath, Path, bytes, BufferedIOBase)):
return await self._aget_json(file_path, extra_info=extra_info)
elif isinstance(file_path, list):
jobs = [self._aget_json(f, extra_info=extra_info) for f in file_path]
@@ -1254,7 +1488,7 @@ class LlamaParse(BasePydanticReader):
raise e
else:
raise ValueError(
"The input file_path must be a string or a list of strings."
"The input file_path must be a string, Path, bytes, BufferedIOBase, or a list of these types."
)
def get_json_result(
+30 -16
View File
@@ -14,11 +14,14 @@ PAGE_REGEX = r"page[-_](\d+)\.jpg$"
class JobMetadata(BaseModel):
"""Metadata about the job."""
job_pages: int = Field(description="The number of pages in the job.")
job_auto_mode_triggered_pages: int = Field(
description="The number of pages that triggered auto mode (thus increasing the cost)."
job_pages: int = Field(default=0, description="The number of pages in the job.")
job_auto_mode_triggered_pages: Optional[int] = Field(
default=None,
description="The number of pages that triggered auto mode (thus increasing the cost).",
)
job_is_cache_hit: bool = Field(
default=False, description="Whether the job was a cache hit."
)
job_is_cache_hit: bool = Field(description="Whether the job was a cache hit.")
class BBox(BaseModel):
@@ -43,7 +46,7 @@ class PageItem(BaseModel):
md: Optional[str] = Field(
default=None, description="The markdown-formatted content of the item."
)
rows: Optional[List[List[str]]] = Field(
rows: Optional[List[List[Any]]] = Field(
default=None, description="The rows of the item."
)
bBox: Optional[BBox] = Field(
@@ -124,32 +127,43 @@ class Page(BaseModel):
items: List[PageItem] = Field(
default_factory=list, description="The items in the page."
)
status: str = Field(description="The status of the page.")
status: Optional[str] = Field(default=None, description="The status of the page.")
links: List[SerializeAsAny[Any]] = Field(
default_factory=list, description="The links in the page."
)
width: Optional[float] = Field(default=None, description="The width of the page.")
height: Optional[float] = Field(default=None, description="The height of the page.")
triggeredAutoMode: bool = Field(
description="Whether the page triggered auto mode (thus increasing the cost)."
triggeredAutoMode: Optional[bool] = Field(
default=False,
description="Whether the page triggered auto mode (thus increasing the cost).",
)
parsingMode: str = Field(
default="", description="The parsing mode used for the page."
)
parsingMode: str = Field(description="The parsing mode used for the page.")
structuredData: Optional[Dict[str, Any]] = Field(
description="The structured data of the page."
default=None, description="The structured data of the page."
)
noStructuredContent: bool = Field(
description="Whether the page has no structured data."
default=True, description="Whether the page has no structured data."
)
noTextContent: bool = Field(
default=False, description="Whether the page has no text content."
)
noTextContent: bool = Field(description="Whether the page has no text content.")
class JobResult(BaseModel):
"""The raw JSON result from the LlamaParse API."""
pages: List[Page] = Field(description="The pages of the document.")
job_metadata: JobMetadata = Field(description="The metadata of the job.")
file_name: str = Field(description="The path to the file that was parsed.")
job_id: str = Field(description="The ID of the job.")
pages: List[Page] = Field(
default_factory=list, description="The pages of the document."
)
job_metadata: JobMetadata = Field(
default_factory=JobMetadata, description="The metadata of the job."
)
file_name: str = Field(
default="", description="The path to the file that was parsed."
)
job_id: str = Field(default="", description="The ID of the job.")
is_done: bool = Field(default=False, description="Whether the job is done.")
error: Optional[str] = Field(
default=None, description="The error message if the job failed."
+55 -1
View File
@@ -1,4 +1,5 @@
import httpx
import itertools
import logging
from enum import Enum
from tenacity import (
@@ -8,7 +9,7 @@ from tenacity import (
retry_if_exception,
before_sleep_log,
)
from typing import Any
from typing import Any, Iterable, Iterator, Optional
logger = logging.getLogger(__name__)
@@ -297,3 +298,56 @@ async def make_api_request(
return response
return await _make_request(url, **httpx_kwargs)
def expand_target_pages(target_pages: str) -> Iterator[int]:
"""Yield all values in target_pages."""
for target in target_pages.strip().split(","):
if "-" in target:
try:
start, end = map(int, target.strip().split("-"))
if start > end:
raise ValueError
yield from range(start, end + 1)
except ValueError as e:
raise ValueError(f"Invalid page range: {target}") from e
else:
try:
yield int(target)
except ValueError as e:
raise ValueError(f"Invalid page number: {target}") from e
def partition_pages(
pages: Iterable[int], size: int, max_pages: Optional[int] = None
) -> Iterator[str]:
"""Yield partitioned target_pages segments."""
if size < 1:
raise ValueError(f"Invalid partition segment size: {size}")
if max_pages is not None and max_pages < 1:
raise ValueError("Max pages must be > 0")
it = iter(pages)
total = 0
while max_pages is None or total < max_pages:
segment = tuple(itertools.islice(it, size))
if segment:
targets = []
for _k, g in itertools.groupby(enumerate(segment), lambda x: x[0] - x[1]):
group = [item[1] for item in g]
if len(group) > 1:
start, end = group[0], group[-1]
group_size = end - start + 1
if max_pages is not None and total + group_size > max_pages:
end -= total + group_size - max_pages
group_size = end - start + 1
if group_size > 1:
targets.append(f"{start}-{end}")
else:
targets.append(str(start))
total += group_size
else:
targets.append(str(group[0]))
total += 1
yield ",".join(targets)
else:
return
+29
View File
@@ -0,0 +1,29 @@
import difflib
from pydantic import BaseModel
from typing import Any, Dict, List, Tuple, Type
def check_extra_params(
model_cls: Type[BaseModel], data: Dict[str, Any]
) -> Tuple[List[str], List[str]]:
# check if one of the parameters is unused, and warn the user
model_attributes = set(model_cls.model_fields.keys())
extra_params = [param for param in data.keys() if param not in model_attributes]
suggestions: List[str] = []
if extra_params:
# for each unused parameter, check if it is similar to a valid parameter and suggest a typo correction, else suggest to check the documentation / update the package
for param in extra_params:
similar_params = difflib.get_close_matches(
param, model_attributes, n=1, cutoff=0.8
)
if similar_params:
suggestions.append(
f"'{param}' is not a valid parameter. Did you mean '{similar_params[0]}' instead of '{param}'?"
)
else:
suggestions.append(
f"'{param}' is not a valid parameter. Please check the documentation or update the package."
)
return extra_params, suggestions
+47 -46
View File
@@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 2.1.2 and should not be changed by hand.
# This file is automatically @generated by Poetry 2.1.1 and should not be changed by hand.
[[package]]
name = "aiohappyeyeballs"
@@ -257,7 +257,7 @@ version = "2025.1.31"
description = "Python package for providing Mozilla's CA Bundle."
optional = false
python-versions = ">=3.6"
groups = ["main"]
groups = ["main", "dev"]
files = [
{file = "certifi-2025.1.31-py3-none-any.whl", hash = "sha256:ca78db4565a652026a4db2bcdf68f2fb589ea80d0be70e03929ed730746b84fe"},
{file = "certifi-2025.1.31.tar.gz", hash = "sha256:3d5da6925056f6f18f119200434a4780a94263f10d1c21d032a6f6b2baa20651"},
@@ -350,7 +350,7 @@ version = "3.4.1"
description = "The Real First Universal Charset Detector. Open, modern and actively maintained alternative to Chardet."
optional = false
python-versions = ">=3.7"
groups = ["main"]
groups = ["main", "dev"]
files = [
{file = "charset_normalizer-3.4.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:91b36a978b5ae0ee86c394f5a54d6ef44db1de0815eb43de826d41d21e4af3de"},
{file = "charset_normalizer-3.4.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7461baadb4dc00fd9e0acbe254e3d7d2112e7f92ced2adc96e54ef6501c5f176"},
@@ -593,7 +593,7 @@ description = "Like `typing._eval_type`, but lets older Python versions use newe
optional = false
python-versions = ">=3.8"
groups = ["main"]
markers = "python_version == \"3.9\""
markers = "python_version < \"3.10\""
files = [
{file = "eval_type_backport-0.2.2-py3-none-any.whl", hash = "sha256:cb6ad7c393517f476f96d456d0412ea80f0a8cf96f6892834cd9340149111b0a"},
{file = "eval_type_backport-0.2.2.tar.gz", hash = "sha256:f0576b4cf01ebb5bd358d02314d31846af5e07678387486e2c798af0e7d849c1"},
@@ -892,31 +892,31 @@ colorama = ">=0.4"
[[package]]
name = "h11"
version = "0.14.0"
version = "0.16.0"
description = "A pure-Python, bring-your-own-I/O implementation of HTTP/1.1"
optional = false
python-versions = ">=3.7"
groups = ["main"]
python-versions = ">=3.8"
groups = ["main", "dev"]
files = [
{file = "h11-0.14.0-py3-none-any.whl", hash = "sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761"},
{file = "h11-0.14.0.tar.gz", hash = "sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d"},
{file = "h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86"},
{file = "h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1"},
]
[[package]]
name = "httpcore"
version = "1.0.8"
version = "1.0.9"
description = "A minimal low-level HTTP client."
optional = false
python-versions = ">=3.8"
groups = ["main"]
files = [
{file = "httpcore-1.0.8-py3-none-any.whl", hash = "sha256:5254cf149bcb5f75e9d1b2b9f729ea4a4b883d1ad7379fc632b727cec23674be"},
{file = "httpcore-1.0.8.tar.gz", hash = "sha256:86e94505ed24ea06514883fd44d2bc02d90e77e7979c8eb71b90f41d364a1bad"},
{file = "httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55"},
{file = "httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8"},
]
[package.dependencies]
certifi = "*"
h11 = ">=0.13,<0.15"
h11 = ">=0.16"
[package.extras]
asyncio = ["anyio (>=4.0,<5.0)"]
@@ -955,7 +955,7 @@ version = "3.10"
description = "Internationalized Domain Names in Applications (IDNA)"
optional = false
python-versions = ">=3.6"
groups = ["main"]
groups = ["main", "dev"]
files = [
{file = "idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3"},
{file = "idna-3.10.tar.gz", hash = "sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9"},
@@ -971,7 +971,7 @@ description = "Read metadata from Python packages"
optional = false
python-versions = ">=3.9"
groups = ["dev"]
markers = "python_version == \"3.9\""
markers = "python_version < \"3.10\""
files = [
{file = "importlib_metadata-8.6.1-py3-none-any.whl", hash = "sha256:02a89390c1e15fdfdc0d7c6b25cb3e62650d0494005c97d6f148bf5b9787525e"},
{file = "importlib_metadata-8.6.1.tar.gz", hash = "sha256:310b41d755445d74569f993ccfc22838295d9fe005425094fad953d7f15c8580"},
@@ -1170,14 +1170,14 @@ test = ["ipykernel", "pre-commit", "pytest (<8)", "pytest-cov", "pytest-timeout"
[[package]]
name = "llama-cloud"
version = "0.1.19"
version = "0.1.23"
description = ""
optional = false
python-versions = "<4,>=3.8"
groups = ["main"]
files = [
{file = "llama_cloud-0.1.19-py3-none-any.whl", hash = "sha256:d2d551baa4b63f7717f8e04cbb81b0f817e5450a66870c5487dd371f81dab8ec"},
{file = "llama_cloud-0.1.19.tar.gz", hash = "sha256:b0a5424ae0099ca27df2a2d7e5aec99066de9ca860ab65987c9f931f1ea7abff"},
{file = "llama_cloud-0.1.23-py3-none-any.whl", hash = "sha256:ce95b0705d85c99b3b27b0af0d16a17d9a81b14c96bf13c1063a1bd13d8d0446"},
{file = "llama_cloud-0.1.23.tar.gz", hash = "sha256:3d84a24a860f046d39a106c06742ec0ea39a574ac42bbf91706fe025f44e233e"},
]
[package.dependencies]
@@ -1187,23 +1187,23 @@ pydantic = ">=1.10"
[[package]]
name = "llama-cloud-services"
version = "0.6.18"
version = "0.6.30"
description = "Tailored SDK clients for LlamaCloud services."
optional = false
python-versions = "<4.0,>=3.9"
groups = ["main"]
files = [
{file = "llama_cloud_services-0.6.18-py3-none-any.whl", hash = "sha256:ba24ef36f4f78619722e5fe2cac8739175ba63f67cb20f0432be16801c78162c"},
{file = "llama_cloud_services-0.6.18.tar.gz", hash = "sha256:070a0e1397e2dd6da59e5f6a0437e3fd27e1c156170eea66001c4bcbec2f8783"},
{file = "llama_cloud_services-0.6.30-py3-none-any.whl", hash = "sha256:4d5817a9841fc3ba3409865c52d082090f4ef827931f0e5e4a89f5818c0d4e36"},
{file = "llama_cloud_services-0.6.30.tar.gz", hash = "sha256:2cb5004d13127aac52888ae9b3d70f899d598633520b2a2542bb62682d08d776"},
]
[package.dependencies]
click = ">=8.1.7,<9.0.0"
eval-type-backport = {version = ">=0.2.0,<0.3.0", markers = "python_version < \"3.10\""}
llama-cloud = "0.1.19"
llama-index-core = ">=0.11.0"
llama-cloud = "0.1.23"
llama-index-core = ">=0.12.0"
platformdirs = ">=4.3.7,<5.0.0"
pydantic = "!=2.10"
pydantic = ">=2.8,<2.10 || >2.10"
python-dotenv = ">=1.0.1,<2.0.0"
[[package]]
@@ -2469,19 +2469,19 @@ files = [
[[package]]
name = "requests"
version = "2.32.3"
version = "2.32.4"
description = "Python HTTP for Humans."
optional = false
python-versions = ">=3.8"
groups = ["main"]
groups = ["main", "dev"]
files = [
{file = "requests-2.32.3-py3-none-any.whl", hash = "sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6"},
{file = "requests-2.32.3.tar.gz", hash = "sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760"},
{file = "requests-2.32.4-py3-none-any.whl", hash = "sha256:27babd3cda2a6d50b30443204ee89830707d396671944c998b5975b031ac2b2c"},
{file = "requests-2.32.4.tar.gz", hash = "sha256:27d0316682c8a29834d3264820024b62a36942083d52caf2f14c0591336d3422"},
]
[package.dependencies]
certifi = ">=2017.4.17"
charset-normalizer = ">=2,<4"
charset_normalizer = ">=2,<4"
idna = ">=2.5,<4"
urllib3 = ">=1.21.1,<3"
@@ -2738,23 +2738,24 @@ files = [
[[package]]
name = "tornado"
version = "6.4.2"
version = "6.5.1"
description = "Tornado is a Python web framework and asynchronous networking library, originally developed at FriendFeed."
optional = false
python-versions = ">=3.8"
python-versions = ">=3.9"
groups = ["dev"]
files = [
{file = "tornado-6.4.2-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:e828cce1123e9e44ae2a50a9de3055497ab1d0aeb440c5ac23064d9e44880da1"},
{file = "tornado-6.4.2-cp38-abi3-macosx_10_9_x86_64.whl", hash = "sha256:072ce12ada169c5b00b7d92a99ba089447ccc993ea2143c9ede887e0937aa803"},
{file = "tornado-6.4.2-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1a017d239bd1bb0919f72af256a970624241f070496635784d9bf0db640d3fec"},
{file = "tornado-6.4.2-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c36e62ce8f63409301537222faffcef7dfc5284f27eec227389f2ad11b09d946"},
{file = "tornado-6.4.2-cp38-abi3-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bca9eb02196e789c9cb5c3c7c0f04fb447dc2adffd95265b2c7223a8a615ccbf"},
{file = "tornado-6.4.2-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:304463bd0772442ff4d0f5149c6f1c2135a1fae045adf070821c6cdc76980634"},
{file = "tornado-6.4.2-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:c82c46813ba483a385ab2a99caeaedf92585a1f90defb5693351fa7e4ea0bf73"},
{file = "tornado-6.4.2-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:932d195ca9015956fa502c6b56af9eb06106140d844a335590c1ec7f5277d10c"},
{file = "tornado-6.4.2-cp38-abi3-win32.whl", hash = "sha256:2876cef82e6c5978fde1e0d5b1f919d756968d5b4282418f3146b79b58556482"},
{file = "tornado-6.4.2-cp38-abi3-win_amd64.whl", hash = "sha256:908b71bf3ff37d81073356a5fadcc660eb10c1476ee6e2725588626ce7e5ca38"},
{file = "tornado-6.4.2.tar.gz", hash = "sha256:92bad5b4746e9879fd7bf1eb21dce4e3fc5128d71601f80005afa39237ad620b"},
{file = "tornado-6.5.1-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:d50065ba7fd11d3bd41bcad0825227cc9a95154bad83239357094c36708001f7"},
{file = "tornado-6.5.1-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9e9ca370f717997cb85606d074b0e5b247282cf5e2e1611568b8821afe0342d6"},
{file = "tornado-6.5.1-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b77e9dfa7ed69754a54c89d82ef746398be82f749df69c4d3abe75c4d1ff4888"},
{file = "tornado-6.5.1-cp39-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:253b76040ee3bab8bcf7ba9feb136436a3787208717a1fb9f2c16b744fba7331"},
{file = "tornado-6.5.1-cp39-abi3-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:308473f4cc5a76227157cdf904de33ac268af770b2c5f05ca6c1161d82fdd95e"},
{file = "tornado-6.5.1-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:caec6314ce8a81cf69bd89909f4b633b9f523834dc1a352021775d45e51d9401"},
{file = "tornado-6.5.1-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:13ce6e3396c24e2808774741331638ee6c2f50b114b97a55c5b442df65fd9692"},
{file = "tornado-6.5.1-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:5cae6145f4cdf5ab24744526cc0f55a17d76f02c98f4cff9daa08ae9a217448a"},
{file = "tornado-6.5.1-cp39-abi3-win32.whl", hash = "sha256:e0a36e1bc684dca10b1aa75a31df8bdfed656831489bc1e6a6ebed05dc1ec365"},
{file = "tornado-6.5.1-cp39-abi3-win_amd64.whl", hash = "sha256:908e7d64567cecd4c2b458075589a775063453aeb1d2a1853eedb806922f568b"},
{file = "tornado-6.5.1-cp39-abi3-win_arm64.whl", hash = "sha256:02420a0eb7bf617257b9935e2b754d1b63897525d8a289c9d65690d580b4dcf7"},
{file = "tornado-6.5.1.tar.gz", hash = "sha256:84ceece391e8eb9b2b95578db65e920d2a61070260594819589609ba9bc6308c"},
]
[[package]]
@@ -2806,7 +2807,7 @@ files = [
{file = "typing_extensions-4.13.2-py3-none-any.whl", hash = "sha256:a439e7c04b49fec3e5d3e2beaa21755cadbbdc391694e28ccdd36ca4a1408f8c"},
{file = "typing_extensions-4.13.2.tar.gz", hash = "sha256:e6c81219bd689f51865d9e372991c540bda33a0379d5573cddb9a3a23f7caaef"},
]
markers = {dev = "python_version == \"3.9\""}
markers = {dev = "python_version < \"3.10\""}
[[package]]
name = "typing-inspect"
@@ -2845,7 +2846,7 @@ version = "2.4.0"
description = "HTTP library with thread-safe connection pooling, file post, and more."
optional = false
python-versions = ">=3.9"
groups = ["main"]
groups = ["main", "dev"]
files = [
{file = "urllib3-2.4.0-py3-none-any.whl", hash = "sha256:4e16665048960a0900c702d4a66415956a584919c03361cac9f1df5c5dd7e813"},
{file = "urllib3-2.4.0.tar.gz", hash = "sha256:414bc6535b787febd7567804cc015fee39daab8ad86268f1310a9250697de466"},
@@ -3067,7 +3068,7 @@ description = "Backport of pathlib-compatible object wrapper for zip files"
optional = false
python-versions = ">=3.9"
groups = ["dev"]
markers = "python_version == \"3.9\""
markers = "python_version < \"3.10\""
files = [
{file = "zipp-3.21.0-py3-none-any.whl", hash = "sha256:ac1bbe05fd2991f160ebce24ffbac5f6d11d83dc90891255885223d42b3cd931"},
{file = "zipp-3.21.0.tar.gz", hash = "sha256:2c9958f6430a2040341a52eb608ed6dd93ef4392e02ffe219417c1b28b5dd1f4"},
@@ -3084,4 +3085,4 @@ type = ["pytest-mypy"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.9,<4.0"
content-hash = "5630c89272551940c2c689fc2c78f5ced7b6dd904ae239ac8987e2e5fb2c0186"
content-hash = "1be767f62c8a8e1c9acb76347055a837ea2c44130004703361d0303b049bede7"
+2 -2
View File
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
[tool.poetry]
name = "llama-parse"
version = "0.6.25"
version = "0.6.44"
description = "Parse files into RAG-Optimized formats."
authors = ["Logan Markewich <logan@llamaindex.ai>"]
license = "MIT"
@@ -13,7 +13,7 @@ packages = [{include = "llama_parse"}]
[tool.poetry.dependencies]
python = ">=3.9,<4.0"
llama-cloud-services = ">=0.6.24"
llama-cloud-services = ">=0.6.44"
[tool.poetry.group.dev.dependencies]
pytest = "^8.0.0"
Generated
+5 -5
View File
@@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 2.1.2 and should not be changed by hand.
# This file is automatically @generated by Poetry 2.1.3 and should not be changed by hand.
[[package]]
name = "aiohappyeyeballs"
@@ -1925,14 +1925,14 @@ rapidfuzz = ">=3.9.0,<4.0.0"
[[package]]
name = "llama-cloud"
version = "0.1.23"
version = "0.1.33"
description = ""
optional = false
python-versions = "<4,>=3.8"
groups = ["main"]
files = [
{file = "llama_cloud-0.1.23-py3-none-any.whl", hash = "sha256:ce95b0705d85c99b3b27b0af0d16a17d9a81b14c96bf13c1063a1bd13d8d0446"},
{file = "llama_cloud-0.1.23.tar.gz", hash = "sha256:3d84a24a860f046d39a106c06742ec0ea39a574ac42bbf91706fe025f44e233e"},
{file = "llama_cloud-0.1.33-py3-none-any.whl", hash = "sha256:35b7d4a30b013f0a343f7e09126b531c697d65bffd4eb4d2d79bf7d65f256178"},
{file = "llama_cloud-0.1.33.tar.gz", hash = "sha256:a0bb900d5a6e86f8c767b48686c5253679ad7ca1b57612dc39b0767e57ad3d78"},
]
[package.dependencies]
@@ -4623,4 +4623,4 @@ type = ["pytest-mypy"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.9,<4.0"
content-hash = "3656fdab0009c6605175b509350412c4da3fb707a8a0e5d3cdc2315a0a7659e8"
content-hash = "487717a0bbe7ff67360e3e9f187e15a33c77e4942a7c909cb37b95425a81544c"
+3 -2
View File
@@ -8,7 +8,7 @@ python_version = "3.10"
[tool.poetry]
name = "llama-cloud-services"
version = "0.6.25"
version = "0.6.44"
description = "Tailored SDK clients for LlamaCloud services."
authors = ["Logan Markewich <logan@runllama.ai>"]
license = "MIT"
@@ -18,12 +18,13 @@ packages = [{include = "llama_cloud_services"}]
[tool.poetry.dependencies]
python = ">=3.9,<4.0"
llama-index-core = ">=0.12.0"
llama-cloud = "==0.1.23"
llama-cloud = "==0.1.33"
pydantic = ">=2.8,!=2.10"
click = "^8.1.7"
python-dotenv = "^1.0.1"
eval-type-backport = {python = "<3.10", version = "^0.2.0"}
platformdirs = "^4.3.7"
tenacity = ">=8.5.0, <10.0"
[tool.poetry.group.dev.dependencies]
pytest = "^8.0.0"
@@ -0,0 +1,129 @@
import os
import httpx
import pytest
import uuid
from pydantic import BaseModel
from dotenv import load_dotenv
from pathlib import Path
from llama_cloud.client import AsyncLlamaCloud
from llama_cloud_services.beta.agent_data import AsyncAgentDataClient
class TrailingSlashHttpxClient(httpx.AsyncClient):
"""Custom httpx client that ensures all URLs have trailing slashes"""
async def request(self, method, url, **kwargs):
# Convert URL to string and ensure trailing slash
url_str = str(url)
if not url_str.endswith("/") and "?" not in url_str:
url_str += "/"
self.headers["Authorization"] = f"Bearer {LLAMA_CLOUD_API_KEY}"
kwargs.pop("headers", None)
return await super().request(method, url_str, headers=self.headers, **kwargs)
# Load environment variables
def load_test_dotenv():
dotenv_path = Path(__file__).parent.parent.parent.parent / ".env.dev"
load_dotenv(dotenv_path, override=True)
load_test_dotenv()
# Get configuration from environment
LLAMA_CLOUD_API_KEY = os.getenv("LLAMA_CLOUD_API_KEY")
LLAMA_CLOUD_BASE_URL = os.getenv("LLAMA_CLOUD_BASE_URL")
LLAMA_DEPLOY_DEPLOYMENT_NAME = os.getenv("LLAMA_DEPLOY_DEPLOYMENT_NAME")
class TestData(BaseModel):
"""Simple test data model for agent data testing"""
name: str
test_id: str
value: int
# Skip all tests if API key is not set
@pytest.mark.asyncio
@pytest.mark.skipif(
not LLAMA_CLOUD_API_KEY or not LLAMA_DEPLOY_DEPLOYMENT_NAME,
reason="LLAMA_CLOUD_API_KEY or LLAMA_DEPLOY_DEPLOYMENT_NAME not set",
)
async def test_agent_data_crud_operations():
"""Test basic CRUD operations for agent data with automatic cleanup"""
# Create unique test identifier to avoid conflicts
test_id = str(uuid.uuid4())
# Set up client
client = AsyncLlamaCloud(
token=LLAMA_CLOUD_API_KEY,
base_url=LLAMA_CLOUD_BASE_URL,
httpx_client=TrailingSlashHttpxClient(timeout=60, follow_redirects=True),
)
# Create agent data client with unique collection name
agent_data_client = AsyncAgentDataClient(
client=client,
type=TestData,
collection_name=f"test-collection-{test_id[:8]}",
agent_url_id=LLAMA_DEPLOY_DEPLOYMENT_NAME,
)
# Create test data
test_data = TestData(name="test-item", test_id=test_id, value=42)
created_item = None
try:
# Test CREATE
created_item = await agent_data_client.create_agent_data(test_data)
assert created_item.data.name == "test-item"
assert created_item.data.test_id == test_id
assert created_item.data.value == 42
assert created_item.id is not None
# Test READ
retrieved_item = await agent_data_client.get_agent_data(created_item.id)
assert retrieved_item.id == created_item.id
assert retrieved_item.data.name == "test-item"
assert retrieved_item.data.test_id == test_id
assert retrieved_item.data.value == 42
# Test SEARCH
search_results = await agent_data_client.search_agent_data(
filter={"test_id": {"eq": test_id}}, page_size=10, include_total=True
)
assert len(search_results.items) == 1
assert search_results.items[0].data.test_id == test_id
assert search_results.total == 1
# Test AGGREGATE
aggregate_results = await agent_data_client.aggregate_agent_data(
group_by=["test_id"], count=True
)
assert len(aggregate_results.items) == 1
assert aggregate_results.items[0].group_key["test_id"] == test_id
assert aggregate_results.items[0].count == 1
# Test UPDATE
updated_data = TestData(name="updated-item", test_id=test_id, value=84)
updated_item = await agent_data_client.update_agent_data(
created_item.id, updated_data
)
assert updated_item.data.name == "updated-item"
assert updated_item.data.value == 84
assert updated_item.id == created_item.id
# Verify update persisted
verified_item = await agent_data_client.get_agent_data(created_item.id)
assert verified_item.data.name == "updated-item"
assert verified_item.data.value == 84
finally:
# Clean up test data
if created_item is not None:
try:
await agent_data_client.delete_agent_data(created_item.id)
except Exception as e:
print(f"Warning: Failed to cleanup test data {created_item.id}: {e}")
@@ -0,0 +1,109 @@
from datetime import datetime
from typing import Any, Dict
import pytest
from llama_cloud.types.agent_data import AgentData
from llama_cloud.types.aggregate_group import AggregateGroup
from pydantic import BaseModel, ValidationError
from llama_cloud_services.beta.agent_data.schema import (
ExtractedData,
TypedAgentData,
TypedAggregateGroup,
)
# Test data models
class Person(BaseModel):
name: str
age: int
email: str
class Company(BaseModel):
name: str
industry: str
employees: int
def test_typed_agent_data_from_raw():
"""Test TypedAgentData.from_raw class method."""
raw_data = AgentData(
id="456",
agent_slug="extraction-agent",
collection="employees",
data={"name": "Jane Smith", "age": 25, "email": "jane@company.com"},
created_at=datetime.now(),
updated_at=datetime.now(),
)
typed_data = TypedAgentData.from_raw(raw_data, Person)
assert typed_data.id == "456"
assert typed_data.agent_url_id == "extraction-agent"
assert typed_data.collection == "employees"
assert typed_data.data.name == "Jane Smith"
assert typed_data.data.age == 25
assert typed_data.data.email == "jane@company.com"
def test_typed_agent_data_from_raw_validation_error():
"""Test TypedAgentData.from_raw with invalid data."""
raw_data = AgentData(
id="789",
agent_slug="test-agent",
collection="people",
data={"name": "Invalid Person", "age": "not_a_number"}, # Invalid age
created_at=datetime.now(),
updated_at=datetime.now(),
)
with pytest.raises(ValidationError):
TypedAgentData.from_raw(raw_data, Person)
def test_extracted_data_create_method():
"""Test ExtractedData.create class method."""
person = Person(name="Created Person", age=35, email="created@example.com")
# Test with defaults
extracted = ExtractedData.create(person)
assert extracted.original_data == person
assert extracted.data == person
assert extracted.status == "in_review"
assert extracted.confidence == {}
# Test with custom values
extracted_custom = ExtractedData.create(
person, status="accepted", confidence={"name": 0.99}
)
assert extracted_custom.status == "accepted"
assert extracted_custom.confidence["name"] == 0.99
def test_extracted_data_with_dict():
"""Test ExtractedData with dict data instead of Pydantic model."""
data_dict = {"name": "Dict Person", "age": 45, "email": "dict@example.com"}
extracted = ExtractedData[Dict[str, Any]](
original_data=data_dict, data=data_dict, status="accepted", confidence={}
)
assert extracted.original_data["name"] == "Dict Person"
assert extracted.data["age"] == 45
def test_typed_aggregate_group_from_raw():
"""Test TypedAggregateGroup.from_raw class method."""
raw_group = AggregateGroup(
group_key={"industry": "Technology"},
count=25,
first_item={"name": "Tech Corp", "industry": "Technology", "employees": 500},
)
typed_group = TypedAggregateGroup.from_raw(raw_group, Company)
assert typed_group.group_key["industry"] == "Technology"
assert typed_group.count == 25
assert typed_group.first_item.name == "Tech Corp"
assert typed_group.first_item.employees == 500
+41
View File
@@ -0,0 +1,41 @@
import os
from typing import List
from llama_cloud_services.extract import LlamaExtract
# Global storage for agents to cleanup
_TEST_AGENTS_TO_CLEANUP: List[str] = []
def pytest_sessionfinish(session, exitstatus):
"""Hook that runs after all tests complete - cleanup agents here"""
print(
f"pytest_sessionfinish hook called! Agents to cleanup: {_TEST_AGENTS_TO_CLEANUP}"
)
if _TEST_AGENTS_TO_CLEANUP:
print("Creating cleanup client...")
# Create a fresh client just for cleanup
cleanup_client = LlamaExtract(
api_key=os.getenv("LLAMA_CLOUD_API_KEY"),
base_url=os.getenv("LLAMA_CLOUD_BASE_URL"),
project_id=os.getenv("LLAMA_CLOUD_PROJECT_ID"),
verbose=True,
)
for agent_id in _TEST_AGENTS_TO_CLEANUP:
try:
print(f"Deleting agent {agent_id}...")
cleanup_client.delete_agent(agent_id)
print(f"Cleaned up agent {agent_id}")
except Exception as e:
print(f"Warning: Failed to delete agent {agent_id}: {e}")
_TEST_AGENTS_TO_CLEANUP.clear()
print("Agent cleanup completed")
else:
print("No agents to cleanup")
def register_agent_for_cleanup(agent_id: str):
"""Register an agent ID for cleanup at the end of the test session"""
_TEST_AGENTS_TO_CLEANUP.append(agent_id)
Binary file not shown.
+7 -8
View File
@@ -5,6 +5,7 @@ from pydantic import BaseModel
from llama_cloud_services.extract import LlamaExtract, ExtractionAgent, SourceText
from tests.extract.util import load_test_dotenv
from .conftest import register_agent_for_cleanup
load_test_dotenv()
@@ -27,7 +28,7 @@ class TestSchema(BaseModel):
# Test data paths
TEST_DIR = Path(__file__).parent / "data"
TEST_PDF = TEST_DIR / "slide" / "saas_slide.pdf"
TEST_PDF = TEST_DIR / "api_test" / "noisebridge_receipt.pdf"
@pytest.fixture
@@ -58,7 +59,7 @@ def test_schema_dict():
@pytest.fixture
def test_agent(llama_extract, test_agent_name, test_schema_dict, request):
"""Creates a test agent and cleans it up after the test"""
"""Creates a test agent and collects it for cleanup at the end of all tests"""
test_id = request.node.nodeid
test_hash = hex(hash(test_id))[-8:]
base_name = test_agent_name
@@ -86,13 +87,11 @@ def test_agent(llama_extract, test_agent_name, test_schema_dict, request):
print(f"Warning: Failed to cleanup existing agent: {e}")
agent = llama_extract.create_agent(name=name, data_schema=schema)
yield agent
# Cleanup after test
try:
llama_extract.delete_agent(agent.id)
except Exception as e:
print(f"Warning: Failed to delete agent {agent.id}: {e}")
# Add agent to cleanup list via conftest helper
register_agent_for_cleanup(agent.id)
yield agent
class TestLlamaExtract:
+27 -9
View File
@@ -1,6 +1,7 @@
import os
import pytest
import shutil
from typing import Optional, cast
from fsspec.implementations.local import LocalFileSystem
from httpx import AsyncClient
@@ -20,11 +21,15 @@ def test_simple_page_text() -> None:
assert len(result[0].text) > 0
@pytest.fixture
def markdown_parser() -> LlamaParse:
@pytest.fixture(params=[None, 2])
def markdown_parser(request: pytest.FixtureRequest) -> LlamaParse:
if os.environ.get("LLAMA_CLOUD_API_KEY", "") == "":
pytest.skip("LLAMA_CLOUD_API_KEY not set")
return LlamaParse(result_type="markdown", ignore_errors=False)
return LlamaParse(
result_type="markdown",
ignore_errors=False,
partition_pages=cast(Optional[int], request.param),
)
def test_simple_page_markdown(markdown_parser: LlamaParse) -> None:
@@ -35,8 +40,6 @@ def test_simple_page_markdown(markdown_parser: LlamaParse) -> None:
def test_simple_page_markdown_bytes(markdown_parser: LlamaParse) -> None:
markdown_parser = LlamaParse(result_type="markdown", ignore_errors=False)
filepath = "tests/test_files/attention_is_all_you_need.pdf"
with open(filepath, "rb") as f:
file_bytes = f.read()
@@ -51,8 +54,6 @@ def test_simple_page_markdown_bytes(markdown_parser: LlamaParse) -> None:
def test_simple_page_markdown_buffer(markdown_parser: LlamaParse) -> None:
markdown_parser = LlamaParse(result_type="markdown", ignore_errors=False)
filepath = "tests/test_files/attention_is_all_you_need.pdf"
with open(filepath, "rb") as f:
# client must provide extra_info with file_name
@@ -161,9 +162,12 @@ async def test_mixing_input_types() -> None:
os.environ.get("LLAMA_CLOUD_API_KEY", "") == "",
reason="LLAMA_CLOUD_API_KEY not set",
)
@pytest.mark.parametrize("partition_pages", [None, 2])
@pytest.mark.asyncio
async def test_download_images() -> None:
parser = LlamaParse(result_type="markdown", take_screenshot=True)
async def test_download_images(partition_pages: Optional[int]) -> None:
parser = LlamaParse(
result_type="markdown", take_screenshot=True, partition_pages=partition_pages
)
filepath = "tests/test_files/attention_is_all_you_need.pdf"
json_result = await parser.aget_json([filepath])
@@ -175,3 +179,17 @@ async def test_download_images() -> None:
await parser.aget_images(json_result, download_path)
assert len(os.listdir(download_path)) == len(json_result[0]["pages"][0]["images"])
@pytest.mark.asyncio
@pytest.mark.parametrize("split_by_page,expected", [(True, 4), (False, 1)])
async def test_multiple_page_markdown(
markdown_parser: LlamaParse,
split_by_page: bool,
expected: int,
) -> None:
markdown_parser.split_by_page = split_by_page
filepath = "tests/test_files/TOS.pdf"
result = await markdown_parser.aload_data(filepath)
assert len(result) == expected
assert all(len(doc.text) > 0 for doc in result)
+53 -3
View File
@@ -1,6 +1,7 @@
import tempfile
import os
import pytest
from typing import Optional
from llama_cloud_services import LlamaParse
from llama_cloud_services.parse.types import JobResult
@@ -15,16 +16,23 @@ def chart_file_path() -> str:
return "tests/test_files/attention_is_all_you_need_chart.pdf"
@pytest.fixture
def multiple_page_path() -> str:
return "tests/test_files/TOS.pdf"
@pytest.mark.asyncio
@pytest.mark.skipif(
os.environ.get("LLAMA_CLOUD_API_KEY", "") == "",
reason="LLAMA_CLOUD_API_KEY not set",
)
async def test_basic_parse_result(file_path: str):
@pytest.mark.parametrize("partition_pages", [None, 2])
async def test_basic_parse_result(file_path: str, partition_pages: Optional[int]):
parser = LlamaParse(
take_screenshot=True,
auto_mode=True,
fast_mode=False,
partition_pages=partition_pages,
)
result = await parser.aparse(file_path)
@@ -96,10 +104,12 @@ async def test_link_parse_result(file_path: str):
os.environ.get("LLAMA_CLOUD_API_KEY", "") == "",
reason="LLAMA_CLOUD_API_KEY not set",
)
@pytest.mark.skip(reason="TODO: Needs to be fixed in prod. Raising 500 error.")
async def test_parse_structured_output(file_path: str):
parser = LlamaParse(
structured_output=True,
structured_output_json_schema_name="imFeelingLucky",
invalidate_cache=True,
)
result = await parser.aparse(file_path)
assert isinstance(result, JobResult)
@@ -142,8 +152,11 @@ async def test_parse_layout(file_path: str):
os.environ.get("LLAMA_CLOUD_API_KEY", "") == "",
reason="LLAMA_CLOUD_API_KEY not set",
)
def test_parse_multiple_files(file_path: str, chart_file_path: str):
parser = LlamaParse()
@pytest.mark.parametrize("partition_pages", [None, 2])
def test_parse_multiple_files(
file_path: str, chart_file_path: str, partition_pages: Optional[int]
):
parser = LlamaParse(partition_pages=partition_pages)
result = parser.parse([file_path, chart_file_path])
assert isinstance(result, list)
@@ -152,3 +165,40 @@ def test_parse_multiple_files(file_path: str, chart_file_path: str):
assert isinstance(result[1], JobResult)
assert result[0].file_name == file_path
assert result[1].file_name == chart_file_path
@pytest.mark.asyncio
@pytest.mark.skipif(
os.environ.get("LLAMA_CLOUD_API_KEY", "") == "",
reason="LLAMA_CLOUD_API_KEY not set",
)
@pytest.mark.parametrize("partition_pages", [None, 2])
async def test_multiple_page_parse_result(
multiple_page_path: str, partition_pages: Optional[int]
):
parser = LlamaParse(
take_screenshot=True,
auto_mode=True,
fast_mode=False,
partition_pages=partition_pages,
)
results = await parser.aparse(multiple_page_path)
if partition_pages is None:
assert isinstance(results, JobResult)
results = [results]
else:
assert isinstance(results, list)
for result in results:
assert isinstance(result, JobResult)
assert result.job_id is not None
assert result.file_name == multiple_page_path
assert len(result.pages) > 0
assert result.pages[0].text is not None
assert len(result.pages[0].text) > 0
assert result.pages[0].md is not None
assert len(result.pages[0].md) > 0
assert result.pages[0].md != result.pages[0].text
+30
View File
@@ -0,0 +1,30 @@
import pytest
from llama_cloud_services.parse.utils import expand_target_pages, partition_pages
def test_expand_target_pages() -> None:
with pytest.raises(ValueError):
list(expand_target_pages("x"))
with pytest.raises(ValueError):
list(expand_target_pages("1-2-3"))
with pytest.raises(ValueError):
list(expand_target_pages("2-1"))
result = list(expand_target_pages("0,2-3,5,8-10"))
assert result == [0, 2, 3, 5, 8, 9, 10]
def test_partion_pages() -> None:
pages = [0, 2, 3, 5, 8, 9, 10]
with pytest.raises(ValueError):
list(partition_pages(pages, 0))
result = list(partition_pages(pages, 3))
assert result == ["0,2-3", "5,8-9", "10"]
with pytest.raises(ValueError):
list(partition_pages(pages, 3, 0))
result = list(partition_pages(pages, 3, max_pages=5))
assert result == ["0,2-3", "5,8"]
result = list(partition_pages(pages, 3, max_pages=10))
assert result == ["0,2-3", "5,8-9", "10"]
Binary file not shown.
+66
View File
@@ -0,0 +1,66 @@
from pydantic import BaseModel
from llama_cloud_services.utils import check_extra_params
class MyModel(BaseModel):
name: str
age: int
email: str
is_active: bool
def test_check_extra_params_no_extra():
"""Test when all parameters are valid - should return empty lists."""
data = {"name": "John", "age": 25, "email": "john@example.com", "is_active": True}
extra_params, suggestions = check_extra_params(MyModel, data)
assert extra_params == []
assert suggestions == []
def test_check_extra_params_with_typos():
"""Test when there are extra parameters that are close to valid ones (typos)."""
data = {
"name": "John",
"age": 25,
"emial": "john@example.com", # typo: emial instead of email
"is_activ": True, # typo: is_activ instead of is_active
"address": "123 Main St", # completely different parameter
}
extra_params, suggestions = check_extra_params(MyModel, data)
assert len(extra_params) == 3
assert "emial" in extra_params
assert "is_activ" in extra_params
assert "address" in extra_params
# Check that typo suggestions are provided
assert len(suggestions) == 3
assert "Did you mean 'email' instead of 'emial'?" in suggestions[0]
assert "Did you mean 'is_active' instead of 'is_activ'?" in suggestions[1]
assert "check the documentation or update the package" in suggestions[2]
def test_check_extra_params_completely_invalid():
"""Test when there are extra parameters with no close matches."""
data = {
"name": "John",
"xyz": "invalid",
"random_field": 123,
"completely_different": True,
}
extra_params, suggestions = check_extra_params(MyModel, data)
assert len(extra_params) == 3
assert "xyz" in extra_params
assert "random_field" in extra_params
assert "completely_different" in extra_params
# All suggestions should be generic (no close matches)
assert len(suggestions) == 3
for suggestion in suggestions:
assert "check the documentation or update the package" in suggestion
assert "Did you mean" not in suggestion