fix parse model

.github/workflows/hourly_extract_e2e.yml

@@ -2,8 +2,9 @@ name: Hourly Extract E2E Tests
 
 on:
   schedule:
-    - cron: "0 * * * *" # Runs every hour at the top of the hour
+    - cron: "18 * * * *"
   workflow_dispatch:
+    # Allows manual triggering
     inputs:
       environment:
         description: "Environment to run the tests in"
@@ -92,6 +93,7 @@ jobs:
         run: |
           summary="$(python3 - <<'PY'
           import os
+          import re
           from pathlib import Path
 
           log_path = Path(os.environ["API_E2E_LOG_PATH"])
@@ -101,24 +103,32 @@ jobs:
 
           lines = log_path.read_text(errors="ignore").splitlines()
 
-          def find_section(keyword):
-              for i, line in enumerate(lines):
-                  if line.startswith("=") and keyword in line:
-                      return i
-              return None
+          # Find the "short test summary info" section
+          start = None
+          for i, line in enumerate(lines):
+              if line.startswith("=") and "short test summary info" in line:
+                  start = i + 1
+                  break
 
-          start = (
-              find_section("FAILURES")
-              or find_section("ERRORS")
-              or find_section("short test summary info")
-          )
+          if start is None:
+              print("No test summary found.")
+              raise SystemExit(0)
 
-          if start is not None:
-              snippet = lines[start : start + 200]
+          # Extract just the FAILED/ERROR lines (test name + short reason)
+          failed_tests = []
+          for line in lines[start:]:
+              if line.startswith("="):
+                  break  # End of section
+              if line.startswith("FAILED ") or line.startswith("ERROR "):
+                  # Extract test name and truncate the error message
+                  match = re.match(r"(FAILED|ERROR) ([\w/:.\[\]_-]+)", line)
+                  if match:
+                      failed_tests.append(f"{match.group(1)}: {match.group(2)}")
+
+          if failed_tests:
+              print("\n".join(failed_tests[:20]))  # Limit to 20 tests max
           else:
-              snippet = lines[-200:]
-
-          print("\n".join(snippet).strip())
+              print("No failed tests found in summary.")
           PY
           )"
           if [ -z "$summary" ]; then
@@ -143,17 +153,10 @@ jobs:
         with:
           channel-id: ${{ env.SLACK_CHANNEL_ID }}
           slack-message: |
-            ALERT: *Hourly Extract E2E Tests Failure*
-
-            *Environment*: ${{ steps.runtime.outputs.environment }}
-            *Repository*: ${{ github.repository }}
-            *Branch*: ${{ github.ref_name }}
-
-            <${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}|View Details>
-
-            Failed tests:
+            :red_circle: *Extract E2E Failed* (${{ steps.runtime.outputs.environment }})
             ```
             ${{ steps.failed-tests.outputs.summary }}
             ```
+            <${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}|View Run>
         env:
           SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}

py/llama_cloud_services/extract/extract.py

@@ -4,10 +4,11 @@ import os
 import time
 from io import BufferedIOBase, TextIOWrapper
 from pathlib import Path
-from typing import List, Optional, Type, Union, Coroutine, Any, TypeVar
+from typing import Callable, List, Optional, Type, Union, Coroutine, Any, TypeVar
 import warnings
 import httpx
 from pydantic import BaseModel
+from functools import wraps
 from tenacity import (
     retry_if_exception,
     stop_after_attempt,
@@ -54,7 +55,7 @@ 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)
+        return exception.status_code in (429, 500, 502, 503, 504, 425, 408)
     elif isinstance(
         exception, (httpx.HTTPStatusError, httpx.RequestError, httpx.TimeoutException)
     ):
@@ -62,6 +63,33 @@ def _is_retryable_error(exception: BaseException) -> bool:
     return False
 
 
+def _async_retry(
+    max_attempts: int = 5,
+    initial_wait: float = 1,
+    max_wait: float = 30,
+    jitter: float = 3,
+) -> Callable:
+    """Decorator for async functions with retry logic for rate limiting and transient errors."""
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            async for attempt in AsyncRetrying(
+                retry=retry_if_exception(_is_retryable_error),
+                stop=stop_after_attempt(max_attempts),
+                wait=wait_exponential_jitter(
+                    initial=initial_wait, max=max_wait, jitter=jitter
+                ),
+                reraise=True,
+            ):
+                with attempt:
+                    return await func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 async def _validate_schema(
     client: AsyncLlamaCloud, data_schema: SchemaInput
 ) -> JSONObjectType:
@@ -82,50 +110,6 @@ async def _validate_schema(
     return validated_schema.data_schema
 
 
-async def _get_job_with_retry(
-    client: AsyncLlamaCloud,
-    job_id: str,
-    max_attempts: int = 5,
-    initial_wait: float = 1,
-    max_wait: float = 60,
-    jitter: float = 5,
-) -> ExtractJob:
-    """Get extraction job with retry logic."""
-    async for attempt in AsyncRetrying(
-        retry=retry_if_exception(_is_retryable_error),
-        stop=stop_after_attempt(max_attempts),
-        wait=wait_exponential_jitter(initial=initial_wait, max=max_wait, jitter=jitter),
-        reraise=True,
-    ):
-        with attempt:
-            return await client.llama_extract.get_job(job_id=job_id)
-
-
-async def _get_run_with_retry(
-    client: AsyncLlamaCloud,
-    job_id: str,
-    project_id: Optional[str] = None,
-    organization_id: Optional[str] = None,
-    max_attempts: int = 3,
-    initial_wait: float = 1,
-    max_wait: float = 20,
-    jitter: float = 3,
-) -> ExtractRun:
-    """Get extraction run with retry logic."""
-    async for attempt in AsyncRetrying(
-        retry=retry_if_exception(_is_retryable_error),
-        stop=stop_after_attempt(max_attempts),
-        wait=wait_exponential_jitter(initial=initial_wait, max=max_wait, jitter=jitter),
-        reraise=True,
-    ):
-        with attempt:
-            return await client.llama_extract.get_run_by_job_id(
-                job_id=job_id,
-                project_id=project_id,
-                organization_id=organization_id,
-            )
-
-
 async def _wait_for_job_result(
     client: AsyncLlamaCloud,
     job_id: str,
@@ -142,30 +126,33 @@ async def _wait_for_job_result(
     run_jitter: float = 3,
 ) -> Optional[ExtractRun]:
     """Wait for and return the results of an extraction job."""
+
+    @_async_retry(
+        max_attempts=job_retry_attempts, max_wait=job_max_wait, jitter=job_jitter
+    )
+    async def _get_job() -> ExtractJob:
+        return await client.llama_extract.get_job(job_id=job_id)
+
+    @_async_retry(
+        max_attempts=run_retry_attempts, max_wait=run_max_wait, jitter=run_jitter
+    )
+    async def _get_run() -> ExtractRun:
+        return await client.llama_extract.get_run_by_job_id(
+            job_id=job_id,
+            project_id=project_id,
+            organization_id=organization_id,
+        )
+
     start = time.perf_counter()
     poll_count = 0
 
     while True:
         await asyncio.sleep(check_interval)
         poll_count += 1
-        job = await _get_job_with_retry(
-            client,
-            job_id,
-            max_attempts=job_retry_attempts,
-            max_wait=job_max_wait,
-            jitter=job_jitter,
-        )
+        job = await _get_job()
 
         if job.status == StatusEnum.SUCCESS:
-            return await _get_run_with_retry(
-                client,
-                job_id,
-                project_id,
-                organization_id,
-                max_attempts=run_retry_attempts,
-                max_wait=run_max_wait,
-                jitter=run_jitter,
-            )
+            return await _get_run()
         elif job.status == StatusEnum.PENDING:
             end = time.perf_counter()
             if end - start > max_timeout:
@@ -177,15 +164,7 @@ async def _wait_for_job_result(
             warnings.warn(
                 f"Failure in job: {job_id}, status: {job.status}, error: {job.error}"
             )
-            return await _get_run_with_retry(
-                client,
-                job_id,
-                project_id,
-                organization_id,
-                max_attempts=run_retry_attempts,
-                max_wait=run_max_wait,
-                jitter=run_jitter,
-            )
+            return await _get_run()
 
 
 def run_in_thread(
@@ -498,9 +477,12 @@ class ExtractionAgent:
         Args:
             run_id (str): The ID of the extraction run to delete
         """
-        self._run_in_thread(
-            self._client.llama_extract.delete_extraction_run(run_id=run_id)
-        )
+
+        @_async_retry()
+        async def _delete() -> None:
+            return await self._client.llama_extract.delete_extraction_run(run_id=run_id)
+
+        self._run_in_thread(_delete())
 
     def list_extraction_runs(
         self, page: int = 0, limit: int = 100
@@ -510,13 +492,16 @@ class ExtractionAgent:
         Returns:
             PaginatedExtractRunsResponse: Paginated list of extraction runs
         """
-        return self._run_in_thread(
-            self._client.llama_extract.list_extract_runs(
+
+        @_async_retry()
+        async def _list() -> PaginatedExtractRunsResponse:
+            return await self._client.llama_extract.list_extract_runs(
                 extraction_agent_id=self.id,
                 skip=page * limit,
                 limit=limit,
             )
-        )
+
+        return self._run_in_thread(_list())
 
     def __repr__(self) -> str:
         return f"ExtractionAgent(id={self.id}, name={self.name})"
@@ -658,15 +643,17 @@ class LlamaExtract(BaseComponent):
                 "data_schema must be either a dictionary or a Pydantic model"
             )
 
-        agent = self._run_in_thread(
-            self._async_client.llama_extract.create_extraction_agent(
+        @_async_retry()
+        async def _create() -> CloudExtractAgent:
+            return await self._async_client.llama_extract.create_extraction_agent(
                 project_id=self._project_id,
                 organization_id=self._organization_id,
                 name=name,
                 data_schema=data_schema,
                 config=config,
             )
-        )
+
+        agent = self._run_in_thread(_create())
 
         return ExtractionAgent(
             client=self._async_client,
@@ -702,19 +689,27 @@ class LlamaExtract(BaseComponent):
             )
 
         if id:
-            agent = self._run_in_thread(
-                self._async_client.llama_extract.get_extraction_agent(
+
+            @_async_retry()
+            async def _get_by_id() -> CloudExtractAgent:
+                return await self._async_client.llama_extract.get_extraction_agent(
                     extraction_agent_id=id,
                 )
-            )
+
+            agent = self._run_in_thread(_get_by_id())
 
         elif name:
-            agent = self._run_in_thread(
-                self._async_client.llama_extract.get_extraction_agent_by_name(
-                    name=name,
-                    project_id=self._project_id,
+
+            @_async_retry()
+            async def _get_by_name() -> CloudExtractAgent:
+                return (
+                    await self._async_client.llama_extract.get_extraction_agent_by_name(
+                        name=name,
+                        project_id=self._project_id,
+                    )
                 )
-            )
+
+            agent = self._run_in_thread(_get_by_name())
         else:
             raise ValueError("Either name or extraction_agent_id must be provided.")
 
@@ -734,11 +729,14 @@ class LlamaExtract(BaseComponent):
 
     def list_agents(self) -> List[ExtractionAgent]:
         """List all available extraction agents."""
-        agents = self._run_in_thread(
-            self._async_client.llama_extract.list_extraction_agents(
+
+        @_async_retry()
+        async def _list() -> List[CloudExtractAgent]:
+            return await self._async_client.llama_extract.list_extraction_agents(
                 project_id=self._project_id,
             )
-        )
+
+        agents = self._run_in_thread(_list())
 
         return [
             ExtractionAgent(
@@ -763,11 +761,14 @@ class LlamaExtract(BaseComponent):
         Args:
             agent_id (str): ID of the extraction agent to delete
         """
-        self._run_in_thread(
-            self._async_client.llama_extract.delete_extraction_agent(
-                extraction_agent_id=agent_id
+
+        @_async_retry()
+        async def _delete() -> None:
+            return await self._async_client.llama_extract.delete_extraction_agent(
+                extraction_agent_id=agent_id,
             )
-        )
+
+        self._run_in_thread(_delete())
 
     async def _wait_for_job_result(self, job_id: str) -> Optional[ExtractRun]:
         """Wait for and return the results of an extraction job."""

py/tests/extract/test_extract_api.py

@@ -87,7 +87,11 @@ def test_agent(llama_extract, test_agent_name, test_schema_dict, request):
     except Exception as e:
         print(f"Warning: Failed to cleanup existing agent: {e}")
 
-    agent = create_agent_with_retry(llama_extract, name=name, data_schema=schema)
+    # Use config with cache invalidation to ensure fresh results in tests
+    config = ExtractConfig(invalidate_cache=True)
+    agent = create_agent_with_retry(
+        llama_extract, name=name, data_schema=schema, config=config
+    )
 
     # Add agent to cleanup list via conftest helper
     register_agent_for_cleanup(agent.id)
@@ -237,7 +241,7 @@ class TestStatelessExtraction:
 
     @pytest.fixture
     def test_config(self):
-        return ExtractConfig(extraction_mode=ExtractMode.FAST)
+        return ExtractConfig(extraction_mode=ExtractMode.FAST, invalidate_cache=True)
 
     @pytest.fixture
     def test_schema_dict(self):

py/tests/extract/test_extract_e2e.py

@@ -56,10 +56,16 @@ def get_test_cases():
             input_files.append(file_path)
 
         settings = [
-            ExtractConfig(extraction_mode=ExtractMode.FAST),
-            ExtractConfig(extraction_mode=ExtractMode.BALANCED),
-            ExtractConfig(extraction_mode=ExtractMode.MULTIMODAL),
-            ExtractConfig(extraction_mode=ExtractMode.PREMIUM),
+            ExtractConfig(extraction_mode=ExtractMode.FAST, invalidate_cache=True),
+            ExtractConfig(extraction_mode=ExtractMode.BALANCED, invalidate_cache=True),
+            ExtractConfig(
+                extraction_mode=ExtractMode.MULTIMODAL, invalidate_cache=True
+            ),
+            ExtractConfig(
+                extraction_mode=ExtractMode.PREMIUM,
+                invalidate_cache=True,
+                parse_model="anthropic-sonnet-4.5",
+            ),
         ]
 
         for input_file in sorted(input_files):