bump version

parallelise e2e tests

.github/workflows/test_e2e_py.yml

@@ -31,7 +31,7 @@ jobs:
 
       - name: Run Tests
         working-directory: py
-        run: uv run pytest unit_tests/ tests/ -v
+        run: make e2e
 
       - name: Remove virtual environment
         working-directory: py

py/Makefile

@@ -4,11 +4,15 @@ help:	## Show all Makefile targets.
 	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
 
 format:	## Run code autoformatters (black).
-	pre-commit install
-	git ls-files | xargs pre-commit run black --files
+	uv run pre-commit install
+	git ls-files | xargs uv run pre-commit run black --files
 
 lint:	## Run linters: pre-commit (black, ruff, codespell) and mypy
-	pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
+	uv run pre-commit install && git ls-files | xargs uv run pre-commit run --show-diff-on-failure --files
 
-test:	## Run tests via pytest
-	pytest tests
+test:	## Run unit tests via pytest
+	uv run pytest -v unit_tests/
+
+.PHONY: e2e
+e2e:	## Run all tests. Run with high parallelism using xdist since tests are bottlenecked bound by the slow backend parsing
+	uv run pytest -v -n 32 tests/

py/llama_cloud_services/beta/agent_data/schema.py

@@ -223,10 +223,12 @@ def parse_extracted_field_metadata(
         k: _parse_extracted_field_metadata_recursive(v)
         for k, v in field_metadata.items()
         if k not in _METADATA_FIELDS_SIBLING_TO_LEAF
+        and k not in _ADDITIONAL_ROOT_METADATA_FIELDS
     }
 
 
 _METADATA_FIELDS_SIBLING_TO_LEAF = {"reasoning"}
+_ADDITIONAL_ROOT_METADATA_FIELDS = {"error"}
 
 
 def _parse_extracted_field_metadata_recursive(
@@ -417,11 +419,14 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
         """
         file_id = file_id or result.file.id
         file_name = file_name or result.file.name
+        job_id = result.job_id
+        job_field_metadata = result.extraction_metadata.get("field_metadata", {})
+        errors = job_field_metadata.get("error", None)
+        if not isinstance(errors, str):
+            errors = None
 
         try:
-            field_metadata = parse_extracted_field_metadata(
-                result.extraction_metadata.get("field_metadata", {})
-            )
+            field_metadata = parse_extracted_field_metadata(job_field_metadata)
         except ValidationError:
             field_metadata = {}
 
@@ -430,11 +435,15 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
             return cls.create(
                 data=data,
                 status=status,
-                field_metadata=field_metadata,
+                field_metadata=job_field_metadata,
                 file_id=file_id,
                 file_name=file_name,
                 file_hash=file_hash,
-                metadata=metadata or {},
+                metadata={
+                    **({"field_errors": errors} if errors else {}),
+                    "job_id": job_id,
+                    **(metadata or {}),
+                },
             )
         except ValidationError as e:
             invalid_item = ExtractedData[Dict[str, Any]].create(

py/llama_cloud_services/extract/extract.py

@@ -227,7 +227,7 @@ class SourceText:
             raise ValueError(f"Unsupported file type: {type(self.file)}")
 
 
-FileInput = Union[str, Path, BufferedIOBase, SourceText]
+FileInput = Union[str, Path, BufferedIOBase, SourceText, File]
 
 
 def run_in_thread(
@@ -406,6 +406,8 @@ class ExtractionAgent:
 
     async def _upload_file(self, file_input: FileInput) -> File:
         source_text = None
+        if isinstance(file_input, File):
+            return file_input
         if isinstance(file_input, SourceText):
             source_text = file_input
         elif isinstance(file_input, (str, Path)):
@@ -533,7 +535,7 @@ class ExtractionAgent:
 
         upload_tasks = [self._upload_file(file) for file in files]
         with augment_async_errors():
-            uploaded_files = await run_jobs(
+            uploaded_files: List[File] = await run_jobs(
                 upload_tasks,
                 workers=self.num_workers,
                 desc="Uploading files",
@@ -987,8 +989,13 @@ class LlamaExtract(BaseComponent):
                 f"Could not determine file type. Please provide a filename with one of these supported extensions: {supported_list}"
             )
 
-    def _convert_file_to_file_data(self, file_input: FileInput) -> Union[FileData, str]:
+    def _convert_file_to_file_data(
+        self, file_input: FileInput
+    ) -> Union[FileData, str, File]:
         """Convert FileInput to FileData or text string for stateless extraction."""
+        if isinstance(file_input, File):
+            return file_input
+
         if isinstance(file_input, SourceText):
             if file_input.text_content is not None:
                 return file_input.text_content
@@ -1084,24 +1091,23 @@ class LlamaExtract(BaseComponent):
         for file_input in files:
             file_data_or_text = self._convert_file_to_file_data(file_input)
 
-            if isinstance(file_data_or_text, str):
+            if isinstance(file_data_or_text, File):
+                file_args = {"file_id": file_data_or_text.id}
+
+            elif isinstance(file_data_or_text, str):
                 # It's text content
-                job = await self._async_client.llama_extract.extract_stateless(
-                    project_id=self._project_id,
-                    organization_id=self._organization_id,
-                    data_schema=processed_schema,
-                    config=config,
-                    text=file_data_or_text,
-                )
+                file_args = {"text": file_data_or_text}
             else:
                 # It's FileData
-                job = await self._async_client.llama_extract.extract_stateless(
-                    project_id=self._project_id,
-                    organization_id=self._organization_id,
-                    data_schema=processed_schema,
-                    config=config,
-                    file=file_data_or_text,
-                )
+                file_args = {"file": file_data_or_text}
+
+            job = await self._async_client.llama_extract.extract_stateless(
+                project_id=self._project_id,
+                organization_id=self._organization_id,
+                data_schema=processed_schema,
+                config=config,
+                **file_args,
+            )
             jobs.append(job)
 
         return jobs[0] if len(jobs) == 1 else jobs

py/llama_parse/pyproject.toml

@@ -11,13 +11,13 @@ dev = [
 
 [project]
 name = "llama-parse"
-version = "0.6.58"
+version = "0.6.60"
 description = "Parse files into RAG-Optimized formats."
 authors = [{name = "Logan Markewich", email = "logan@llamaindex.ai"}]
 requires-python = ">=3.9,<4.0"
 readme = "README.md"
 license = "MIT"
-dependencies = ["llama-cloud-services>=0.6.58"]
+dependencies = ["llama-cloud-services>=0.6.60"]
 
 [project.scripts]
 llama-parse = "llama_parse.cli.main:parse"

py/pyproject.toml

@@ -5,6 +5,7 @@ build-backend = "hatchling.build"
 [dependency-groups]
 dev = [
   "pytest>=8.0.0,<9",
+  "pytest-xdist>=3.6.1,<4",
   "pytest-asyncio",
   "ipykernel>=6.29.0,<7",
   "pre-commit==3.2.0",
@@ -18,7 +19,7 @@ dev = [
 
 [project]
 name = "llama-cloud-services"
-version = "0.6.58"
+version = "0.6.60"
 description = "Tailored SDK clients for LlamaCloud services."
 authors = [{name = "Logan Markewich", email = "logan@runllama.ai"}]
 requires-python = ">=3.9,<4.0"

py/unit_tests/beta/agent/test_agent_data_schema.py

@@ -24,8 +24,8 @@ from llama_cloud_services.beta.agent_data.schema import (
 # Test data models
 class Person(BaseModel):
     name: str
-    age: int
-    email: str
+    age: Optional[int] = None
+    email: Optional[str] = None
 
 
 class Company(BaseModel):
@@ -405,6 +405,7 @@ def create_file(
 
 def create_extract_run(
     id: str = "extract-123",
+    job_id: str = "job-123",
     data: Dict[str, Any] = {"name": "John Doe", "age": 30, "email": "john@example.com"},
     extraction_metadata: Dict[str, Any] = {
         "name": {
@@ -423,6 +424,7 @@ def create_extract_run(
     return ExtractRun.parse_obj(
         {
             "id": id,
+            "job_id": job_id,
             "data": data,
             "extraction_metadata": {
                 "field_metadata": extraction_metadata,
@@ -501,10 +503,9 @@ def test_extracted_data_from_extraction_result_invalid_data():
     # Create ExtractRun with data that doesn't match Person schema
     extract_run = create_extract_run(
         data={
-            "name": "Valid Name",
+            "missing_name": "Valid Name",
             "age": "not_a_number",
-            "missing_email": True,
-        },  # Invalid age, missing email
+        },  # Invalid age, missing name
         extraction_metadata={
             "name": {"confidence": 0.9},
         },
@@ -523,9 +524,8 @@ def test_extracted_data_from_extraction_result_invalid_data():
     assert isinstance(invalid_data, ExtractedData)
     assert invalid_data.status == "error"
     assert invalid_data.data == {
-        "name": "Valid Name",
+        "missing_name": "Valid Name",
         "age": "not_a_number",
-        "missing_email": True,
     }
     assert invalid_data.file_id == "error-file"
     assert invalid_data.file_name == "bad_data.pdf"
@@ -590,3 +590,26 @@ def test_full_parse_nested_dimensions():
     assert result.field_metadata == expected
     parsed = ExtractedData.model_validate_json(result.model_dump_json())
     assert parsed.field_metadata == expected
+
+
+def test_parses_field_metadata_with_error_field():
+    extract_run = create_extract_run(
+        extraction_metadata={
+            "name": {
+                "confidence": 0.95,
+                "citation": [{"page": 1, "matching_text": "John Smith"}],
+            },
+            "error": "This is an error",
+        },
+    )
+
+    parsed = ExtractedData.from_extraction_result(extract_run, Person)
+
+    assert parsed.field_metadata == {
+        "name": ExtractedFieldMetadata(
+            confidence=0.95,
+            citation=[FieldCitation(page=1, matching_text="John Smith")],
+        ),
+    }
+    assert parsed.metadata.get("field_errors") == "This is an error"
+    assert parsed.metadata.get("job_id") == "job-123"

py/unit_tests/extract/test_extract.py

@@ -0,0 +1,166 @@
+import os
+from types import SimpleNamespace
+
+import pytest
+
+from llama_cloud.types import File as CloudFile
+from llama_cloud_services.extract import LlamaExtract, ExtractionAgent
+
+
+@pytest.fixture(autouse=True)
+def _set_dummy_env(monkeypatch):
+    monkeypatch.setenv("LLAMA_CLOUD_API_KEY", "test-api-key")
+    monkeypatch.setenv("LLAMA_CLOUD_BASE_URL", "https://example.test")
+
+
+@pytest.fixture
+def llama_file() -> CloudFile:
+    return CloudFile(
+        id="file_123",
+        name="sample.pdf",
+        external_file_id="ext_123",
+        project_id="proj_123",
+    )
+
+
+@pytest.fixture
+def extractor() -> LlamaExtract:
+    return LlamaExtract(
+        api_key=os.environ["LLAMA_CLOUD_API_KEY"],
+        base_url=os.environ["LLAMA_CLOUD_BASE_URL"],
+        verify=False,
+    )
+
+
+@pytest.fixture
+def no_external_validation(monkeypatch):
+    import llama_cloud_services.extract.extract as extract_mod
+
+    async def _noop_validate_schema(client, data_schema):
+        return data_schema
+
+    # Disable config warnings and external schema validation
+    monkeypatch.setattr(
+        extract_mod, "_extraction_config_warning", lambda *_args, **_kwargs: None
+    )
+    monkeypatch.setattr(extract_mod, "_validate_schema", _noop_validate_schema)
+
+
+def test_convert_fileinput_accepts_llama_file_directly(
+    extractor: LlamaExtract, llama_file: CloudFile
+):
+    result = extractor._convert_file_to_file_data(llama_file)
+    assert result is llama_file
+
+
+@pytest.mark.asyncio
+async def test_queue_extraction_with_llama_file_uses_file_id(
+    extractor: LlamaExtract, llama_file: CloudFile, no_external_validation, monkeypatch
+):
+    calls = []
+
+    async def fake_extract_stateless(**kwargs):
+        calls.append(kwargs)
+        return SimpleNamespace(id="job_1")
+
+    # Patch the client's method that would normally hit the network
+    monkeypatch.setattr(
+        extractor._async_client.llama_extract,
+        "extract_stateless",
+        fake_extract_stateless,
+    )
+
+    # Minimal schema and dummy config (warnings disabled by fixture)
+    schema = {"type": "object", "properties": {}}
+    dummy_config = SimpleNamespace()
+
+    job = await extractor.queue_extraction(schema, dummy_config, llama_file)
+
+    assert getattr(job, "id") == "job_1"
+    assert len(calls) == 1
+    kwargs = calls[0]
+    assert "file_id" in kwargs and kwargs["file_id"] == llama_file.id
+    assert "file" not in kwargs
+    assert "text" not in kwargs
+
+
+@pytest.mark.asyncio
+async def test_extraction_agent_upload_file_accepts_llama_file_directly(
+    llama_file: CloudFile,
+):
+    # Build a minimal agent without hitting external services
+    dummy_async_client = SimpleNamespace()
+    dummy_agent = SimpleNamespace(id="agent_1", name="dummy", data_schema={}, config={})
+
+    agent = ExtractionAgent(
+        client=dummy_async_client,
+        agent=dummy_agent,
+        project_id=None,
+        organization_id=None,
+        check_interval=0,
+        max_timeout=0,
+        num_workers=1,
+        show_progress=False,
+        verbose=False,
+        verify=False,
+        httpx_timeout=1,
+    )
+
+    result = await agent._upload_file(llama_file)
+    assert result is llama_file
+
+
+@pytest.mark.asyncio
+async def test_extraction_agent_aextract_accepts_llama_file(
+    monkeypatch, llama_file: CloudFile
+):
+    # Build a minimal agent without network
+    dummy_llama_extract_iface = SimpleNamespace()
+
+    async def fake_run_job(**kwargs):
+        # Ensure we are receiving a request with the right file_id
+        request = kwargs.get("request")
+        assert hasattr(request, "file_id")
+        assert request.file_id == llama_file.id
+        return SimpleNamespace(id="job_42")
+
+    dummy_llama_extract_iface.run_job = fake_run_job
+    dummy_async_client = SimpleNamespace(llama_extract=dummy_llama_extract_iface)
+    dummy_agent = SimpleNamespace(id="agent_1", name="dummy", data_schema={}, config={})
+
+    agent = ExtractionAgent(
+        client=dummy_async_client,
+        agent=dummy_agent,
+        project_id=None,
+        organization_id=None,
+        check_interval=0,
+        max_timeout=0,
+        num_workers=1,
+        show_progress=False,
+        verbose=False,
+        verify=False,
+        httpx_timeout=1,
+    )
+
+    # Ensure _upload_file returns the File directly and is called with our File
+    calls = {}
+
+    async def fake_upload_file(file_input):
+        calls["upload_called_with"] = file_input
+        assert file_input is llama_file
+        return file_input
+
+    monkeypatch.setattr(agent, "_upload_file", fake_upload_file)
+
+    # Avoid polling logic by short-circuiting result wait
+    async def fake_wait(job_id: str):
+        assert job_id == "job_42"
+        return SimpleNamespace(id="run_42", status="SUCCESS", data={})
+
+    monkeypatch.setattr(agent, "_wait_for_job_result", fake_wait)
+
+    result = await agent.aextract(llama_file)
+
+    assert calls.get("upload_called_with") is llama_file
+    assert getattr(result, "status") == "SUCCESS"
+    assert getattr(result, "id") == "run_42"

py/uv.lock

@@ -741,6 +741,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/36/f4/c6e662dade71f56cd2f3735141b265c3c79293c109549c1e6933b0651ffc/exceptiongroup-1.3.0-py3-none-any.whl", hash = "sha256:4d111e6e0c13d0644cad6ddaa7ed0261a0b36971f6d23e7ec9b4b9097da78a10", size = 16674, upload-time = "2025-05-10T17:42:49.33Z" },
 ]
 
+[[package]]
+name = "execnet"
+version = "2.1.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/bb/ff/b4c0dc78fbe20c3e59c0c7334de0c27eb4001a2b2017999af398bf730817/execnet-2.1.1.tar.gz", hash = "sha256:5189b52c6121c24feae288166ab41b32549c7e2348652736540b9e6e7d4e72e3", size = 166524, upload-time = "2024-04-08T09:04:19.245Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/43/09/2aea36ff60d16dd8879bdb2f5b3ee0ba8d08cbbdcdfe870e695ce3784385/execnet-2.1.1-py3-none-any.whl", hash = "sha256:26dee51f1b80cebd6d0ca8e74dd8745419761d3bef34163928cbebbdc4749fdc", size = 40612, upload-time = "2024-04-08T09:04:17.414Z" },
+]
+
 [[package]]
 name = "executing"
 version = "2.2.0"
@@ -1587,7 +1596,7 @@ wheels = [
 
 [[package]]
 name = "llama-cloud-services"
-version = "0.6.58"
+version = "0.6.59"
 source = { editable = "." }
 dependencies = [
     { name = "click", version = "8.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
@@ -1615,6 +1624,7 @@ dev = [
     { name = "pydantic-settings" },
     { name = "pytest" },
     { name = "pytest-asyncio" },
+    { name = "pytest-xdist" },
 ]
 
 [package.metadata]
@@ -1642,6 +1652,7 @@ dev = [
     { name = "pydantic-settings", specifier = ">=2.10.1" },
     { name = "pytest", specifier = ">=8.0.0,<9" },
     { name = "pytest-asyncio" },
+    { name = "pytest-xdist", specifier = ">=3.6.1,<4" },
 ]
 
 [[package]]
@@ -2912,6 +2923,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c7/9d/bf86eddabf8c6c9cb1ea9a869d6873b46f105a5d292d3a6f7071f5b07935/pytest_asyncio-1.1.0-py3-none-any.whl", hash = "sha256:5fe2d69607b0bd75c656d1211f969cadba035030156745ee09e7d71740e58ecf", size = 15157, upload-time = "2025-07-16T04:29:24.929Z" },
 ]
 
+[[package]]
+name = "pytest-xdist"
+version = "3.8.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "execnet" },
+    { name = "pytest" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/78/b4/439b179d1ff526791eb921115fca8e44e596a13efeda518b9d845a619450/pytest_xdist-3.8.0.tar.gz", hash = "sha256:7e578125ec9bc6050861aa93f2d59f1d8d085595d6551c2c90b6f4fad8d3a9f1", size = 88069, upload-time = "2025-07-01T13:30:59.346Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl", hash = "sha256:202ca578cfeb7370784a8c33d6d05bc6e13b4f25b5053c30a152269fd10f0b88", size = 46396, upload-time = "2025-07-01T13:30:56.632Z" },
+]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"