Compare commits

..

7 Commits

Author SHA1 Message Date
Adrian Lyjak 79fe1930cf Re-order extraction metadata union for better parsing (#865)
* Re-order args so that pydantic doesn't parse nested dict to a empty extraction result

* Use a citations array instead
2025-08-13 16:22:06 -04:00
Sourabh Desai ab225c3eab Classifier SDK (#837)
* add files client

* add classification SDK (beta/experimental)

* lint

* lint

* update files client

* add polling timeout

* move e2e test settings to conftest.py

* unused params

* use e2e settings class

* make org id optional

* ordering params

* fix tests

* add sync support
2025-08-13 09:50:39 -07:00
Sourabh Desai 6f1de75909 fix presigned urls + add very necessary test (#864) 2025-08-12 15:28:54 -07:00
Sourabh Desai 230ed64e41 missing await (#863)
missed this await
2025-08-12 13:54:34 -07:00
Logan ef126c3a93 remove print (#861) 2025-08-11 17:42:55 -07:00
Logan 51a7534733 support llama parse audio (#859) 2025-08-11 12:57:01 -07:00
Sourabh Desai 4f5d2bde13 add files client (#836)
* add files client

* lint

* update files client

* move e2e test settings to conftest.py

* unused params

* make org id optional
2025-08-08 15:54:00 -07:00
24 changed files with 1006 additions and 72 deletions
@@ -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]]
@@ -238,19 +247,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 +340,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,
@@ -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
View File
@@ -1 +1,2 @@
EU_BASE_URL = "https://api.cloud.eu.llamaindex.ai"
POLLING_TIMEOUT_SECONDS = 300.0
+1 -1
View File
@@ -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
+1 -20
View File
@@ -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"]
+97
View File
@@ -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,
)
+7 -1
View File
@@ -112,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(
@@ -153,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):
+40 -1
View File
@@ -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
+2 -2
View File
@@ -11,13 +11,13 @@ dev = [
[project]
name = "llama-parse"
version = "0.6.56"
version = "0.6.58"
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.56"]
dependencies = ["llama-cloud-services>=0.6.58"]
[project.scripts]
llama-parse = "llama_parse.cli.main:parse"
+3 -2
View File
@@ -12,12 +12,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.56"
version = "0.6.58"
description = "Tailored SDK clients for LlamaCloud services."
authors = [{name = "Logan Markewich", email = "logan@runllama.ai"}]
requires-python = ">=3.9,<4.0"
+2
View File
@@ -0,0 +1,2 @@
[pytest]
asyncio_mode = auto
+212
View File
@@ -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=[]
)
+26
View File
@@ -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()
-2
View File
@@ -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]:
+9
View File
@@ -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.
+147
View File
@@ -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,
@@ -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")],
),
},
}
@@ -450,8 +465,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
@@ -523,3 +539,54 @@ 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
+116
View File
@@ -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
}
Generated
+17 -1
View File
@@ -1587,7 +1587,7 @@ wheels = [
[[package]]
name = "llama-cloud-services"
version = "0.6.56"
version = "0.6.58"
source = { editable = "." }
dependencies = [
{ name = "click", version = "8.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
@@ -1612,6 +1612,7 @@ dev = [
{ name = "jupyter" },
{ name = "mypy" },
{ name = "pre-commit" },
{ name = "pydantic-settings" },
{ name = "pytest" },
{ name = "pytest-asyncio" },
]
@@ -1638,6 +1639,7 @@ 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" },
]
@@ -2855,6 +2857,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"
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "llama-cloud-services",
"version": "0.3.1",
"version": "0.3.2",
"type": "module",
"license": "MIT",
"scripts": {
@@ -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;
}