Compare commits

..

9 Commits

Author SHA1 Message Date
Neeraj Pradhan 39a93d602b fix parse model 2026-01-27 16:05:23 -08:00
Neeraj Pradhan 0c04955c10 Use sonnet when testing premium mode in extract e2e 2026-01-27 15:48:14 -08:00
Neeraj Pradhan 38da9a52d7 Invalidate cache when running extract tests (#1097) 2026-01-26 17:33:23 -08:00
Neeraj Pradhan 1e7ec40ee7 Fix verbose logging on slack channel (#1096) 2026-01-26 17:12:50 -08:00
Neeraj Pradhan dd83c1a9d0 Add retries to all extract sdk functions uniformly (#1095) 2026-01-26 12:05:16 -08:00
Neeraj Pradhan 7cb83f5cd3 Change cron schedule for hourly extract tests (#1094) 2026-01-26 10:15:34 -08:00
Neeraj Pradhan b05266be6d Try to reparse scheduled workflow (#1093) 2026-01-26 09:56:22 -08:00
Neeraj Pradhan eab4798165 Force github reparse of the workflow (#1090) 2026-01-23 11:36:28 -08:00
Neeraj Pradhan b174fa8fab Run hourly extract tests to catch SDK schema drifts (#1089)
* Run hourly extract tests to catch SDK schema drifts

* fix url

* fix prod/staging env
2026-01-22 18:18:45 -08:00
4 changed files with 274 additions and 101 deletions
+162
View File
@@ -0,0 +1,162 @@
name: Hourly Extract E2E Tests
on:
schedule:
- cron: "18 * * * *"
workflow_dispatch:
# Allows manual triggering
inputs:
environment:
description: "Environment to run the tests in"
required: false
default: staging
type: choice
options:
- staging
- production
notify_slack:
description: "Notify Slack"
required: false
default: false
type: boolean
workflow_call:
env:
UV_VERSION: "0.7.20"
PYTHON_VERSION: "3.12"
SLACK_CHANNEL_ID: C078PHNTF44 # Extract channel ID
API_E2E_LOG_PATH: ${{ github.workspace }}/extract-e2e.log
jobs:
extract-e2e:
name: "Hourly Extract E2E Tests (${{ matrix.environment }})"
runs-on: ubuntu-latest
timeout-minutes: 30
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}-${{ matrix.environment }}
cancel-in-progress: true
strategy:
fail-fast: false
matrix:
environment: ${{ github.event_name == 'schedule' && fromJson('["staging", "production"]') || fromJson(format('["{0}"]', github.event.inputs.environment || 'staging')) }}
steps:
- name: Set runtime inputs
id: runtime
run: |
environment=${{ matrix.environment }}
notify_slack=${{ github.event.inputs.notify_slack || github.event_name == 'schedule' }}
echo "environment=${environment}" >> $GITHUB_OUTPUT
echo "notify_slack=${notify_slack}" >> $GITHUB_OUTPUT
if [ "${environment}" = "production" ]; then
echo "LLAMA_CLOUD_BASE_URL=https://api.cloud.llamaindex.ai" >> $GITHUB_ENV
api_key_secret="${{ secrets.LLAMA_CLOUD_API_KEY }}"
project_id_secret="${{ secrets.LLAMA_CLOUD_PROJECT_ID }}"
else
echo "LLAMA_CLOUD_BASE_URL=https://api.staging.llamaindex.ai" >> $GITHUB_ENV
api_key_secret="${{ secrets.LLAMA_CLOUD_API_KEY_STAGING }}"
project_id_secret="${{ secrets.LLAMA_CLOUD_PROJECT_ID_STAGING }}"
fi
if [ -n "$api_key_secret" ]; then
echo "LLAMA_CLOUD_API_KEY=$api_key_secret" >> $GITHUB_ENV
fi
if [ -n "$project_id_secret" ]; then
echo "LLAMA_CLOUD_PROJECT_ID=$project_id_secret" >> $GITHUB_ENV
fi
- uses: actions/checkout@v5
with:
fetch-depth: 0
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
version: ${{ env.UV_VERSION }}
- name: Set up Python
run: uv python install ${{ env.PYTHON_VERSION }} && uv python pin ${{ env.PYTHON_VERSION }}
- name: Run Extract E2E tests
id: extract-tests
continue-on-error: true
working-directory: py
run: |
set -o pipefail
rm -f "$API_E2E_LOG_PATH"
uv run pytest -v -n 8 --timeout=300 --session-timeout=1740 tests/extract/ 2>&1 | tee "$API_E2E_LOG_PATH"
- name: Extract pytest failure summary
id: failed-tests
if: steps.extract-tests.outcome == 'failure' || cancelled()
run: |
summary="$(python3 - <<'PY'
import os
import re
from pathlib import Path
log_path = Path(os.environ["API_E2E_LOG_PATH"])
if not log_path.exists():
print("Test log not found.")
raise SystemExit(0)
lines = log_path.read_text(errors="ignore").splitlines()
# 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
if start is None:
print("No test summary found.")
raise SystemExit(0)
# 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:
print("No failed tests found in summary.")
PY
)"
if [ -z "$summary" ]; then
summary="Failed test summary not available. Review the full run logs."
fi
{
printf 'summary<<EOF\n%s\nEOF\n' "$summary"
} >> "$GITHUB_OUTPUT"
- name: Check test results
if: always()
run: |
if [ "${{ steps.extract-tests.outcome }}" == "failure" ]; then
echo "Extract E2E tests failed"
exit 1
fi
- name: Post to Extract Slack channel
id: slack
if: (failure() || cancelled()) && steps.runtime.outputs.notify_slack == 'true'
uses: slackapi/slack-github-action@v1.27.0
with:
channel-id: ${{ env.SLACK_CHANNEL_ID }}
slack-message: |
: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 }}
+96 -95
View File
@@ -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."""
+6 -2
View File
@@ -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):
+10 -4
View File
@@ -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):