except one more error type

llama_cloud_services/extract/extract.py

@@ -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,
@@ -22,6 +28,7 @@ from llama_cloud import (
     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,
@@ -44,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,
@@ -231,9 +249,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"))
@@ -260,7 +277,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:
@@ -288,35 +305,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.

llama_cloud_services/parse/base.py

@@ -1020,6 +1020,7 @@ class LlamaParse(BasePydanticReader):
                 httpx.ReadTimeout,
                 httpx.WriteTimeout,
                 httpx.HTTPStatusError,
+                httpx.RemoteProtocolError,
             ) as err:
                 error_count += 1
                 end = time.time()

llama_parse/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "llama-parse"
-version = "0.6.35"
+version = "0.6.38"
 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.35"
+llama-cloud-services = ">=0.6.37"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.0.0"

poetry.lock

@@ -1925,14 +1925,14 @@ rapidfuzz = ">=3.9.0,<4.0.0"
 
 [[package]]
 name = "llama-cloud"
-version = "0.1.27"
+version = "0.1.29"
 description = ""
 optional = false
 python-versions = "<4,>=3.8"
 groups = ["main"]
 files = [
-    {file = "llama_cloud-0.1.27-py3-none-any.whl", hash = "sha256:ef48b169dc5028e1975fa5073f7dd77257b50b44fd2cc2c532d0b6d8cf0a7633"},
-    {file = "llama_cloud-0.1.27.tar.gz", hash = "sha256:9750509d22c247cfc3986fdd72a5b722125f39ae20e746b2ba999f73b0f00a62"},
+    {file = "llama_cloud-0.1.29-py3-none-any.whl", hash = "sha256:893e347fe79fdc152653974340dbd9cc56bbd311501cf46e650215e3e3b1b4f3"},
+    {file = "llama_cloud-0.1.29.tar.gz", hash = "sha256:6995943ed4f1d1fe654f7ea5d01970611c5dc70602deaf93786b7f9545e70f1e"},
 ]
 
 [package.dependencies]
@@ -4623,4 +4623,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<4.0"
-content-hash = "a1cf7a9a94907d14f3b41c1bf65fcdc4cb1bbda816811b750f0b65290e1dc4e8"
+content-hash = "9325cc842a43f120837e0250c17c80097971f439070366d124306b4d0b05badf"

pyproject.toml

@@ -8,7 +8,7 @@ python_version = "3.10"
 
 [tool.poetry]
 name = "llama-cloud-services"
-version = "0.6.35"
+version = "0.6.38"
 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.27"
+llama-cloud = "==0.1.29"
 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"

tests/extract/conftest.py

@@ -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)

tests/extract/test_extract_api.py

@@ -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:

tests/parse/test_llama_parse_result.py

@@ -109,6 +109,7 @@ 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)