Compare commits

...

6 Commits

Author SHA1 Message Date
Neeraj Pradhan dad7728135 Bump to version 0.6.7 (#656) 2025-03-21 21:26:54 -07:00
Neeraj Pradhan c5111e3335 Revert httpx_client as argument (#657) 2025-03-21 21:16:56 -07:00
Neeraj Pradhan bbbdb98362 Add provision for custom httpx client for LlamaExtract (#654) 2025-03-21 11:37:40 -07:00
Neeraj Pradhan 60cdc2af84 Add xfail for timeout errors in report gen (#655) 2025-03-21 11:06:49 -07:00
Neeraj Pradhan 344c20f331 Bump up version for release (#652) 2025-03-18 15:54:32 -07:00
Neeraj Pradhan 2b0496e947 Update llama cloud for extract endpoints (#651) 2025-03-18 15:43:43 -07:00
6 changed files with 485 additions and 444 deletions
+56 -9
View File
@@ -22,6 +22,7 @@ from llama_cloud import (
Project,
ExtractTarget,
LlamaExtractSettings,
PaginatedExtractRunsResponse,
)
from llama_cloud.client import AsyncLlamaCloud
from llama_cloud_services.extract.utils import JSONObjectType, augment_async_errors
@@ -56,6 +57,8 @@ class ExtractionAgent:
num_workers: int = 4,
show_progress: bool = True,
verbose: bool = False,
verify: Optional[bool] = True,
httpx_timeout: Optional[float] = 60,
):
self._client = client
self._agent = agent
@@ -65,6 +68,8 @@ class ExtractionAgent:
self.max_timeout = max_timeout
self.num_workers = num_workers
self.show_progress = show_progress
self.verify = verify
self.httpx_timeout = httpx_timeout
self._verbose = verbose
self._data_schema: Union[JSONObjectType, None] = None
self._config: Union[ExtractConfig, None] = None
@@ -77,14 +82,20 @@ class ExtractionAgent:
def run_coro() -> T:
async def wrapped_coro() -> T:
# Get the original client to preserve its configuration
original_client = self._client._client_wrapper.httpx_client
# Create a new client with the same configuration as the original
async with httpx.AsyncClient(
timeout=self._client._client_wrapper.httpx_client.timeout,
verify=self.verify,
timeout=self.httpx_timeout,
) as client:
original_client = self._client._client_wrapper.httpx_client
# Temporarily replace the client
self._client._client_wrapper.httpx_client = client
try:
return await coro
finally:
# Restore the original client
self._client._client_wrapper.httpx_client = original_client
return asyncio.run(wrapped_coro())
@@ -380,15 +391,29 @@ class ExtractionAgent:
)
)
def list_extraction_runs(self) -> List[ExtractRun]:
def delete_extraction_run(self, run_id: str) -> None:
"""Delete an extraction run by ID.
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)
)
def list_extraction_runs(
self, page: int = 0, limit: int = 100
) -> PaginatedExtractRunsResponse:
"""List extraction runs for the extraction agent.
Returns:
List[ExtractRun]: List of extraction runs
PaginatedExtractRunsResponse: Paginated list of extraction runs
"""
return self._run_in_thread(
self._client.llama_extract.list_extract_runs(
extraction_agent_id=self.id,
skip=page * limit,
limit=limit,
)
)
@@ -421,6 +446,12 @@ class LlamaExtract(BaseComponent):
verbose: bool = Field(
default=False, description="Show verbose output when extracting files."
)
verify: Optional[bool] = Field(
default=True, description="Simple SSL verification option."
)
httpx_timeout: Optional[float] = Field(
default=60, description="Timeout for the httpx client."
)
_async_client: AsyncLlamaCloud = PrivateAttr()
_thread_pool: ThreadPoolExecutor = PrivateAttr()
_project_id: Optional[str] = PrivateAttr()
@@ -436,6 +467,8 @@ class LlamaExtract(BaseComponent):
show_progress: bool = True,
project_id: Optional[str] = None,
organization_id: Optional[str] = None,
verify: Optional[bool] = True,
httpx_timeout: Optional[float] = 60,
verbose: bool = False,
):
if not api_key:
@@ -453,11 +486,18 @@ class LlamaExtract(BaseComponent):
max_timeout=max_timeout,
num_workers=num_workers,
show_progress=show_progress,
verify=verify,
httpx_timeout=httpx_timeout,
verbose=verbose,
)
self._httpx_client = httpx.AsyncClient(verify=verify, timeout=httpx_timeout)
self.verify = verify
self.httpx_timeout = httpx_timeout
self._async_client = AsyncLlamaCloud(
token=self.api_key, base_url=self.base_url, timeout=None
token=self.api_key,
base_url=self.base_url,
httpx_client=self._httpx_client,
)
self._thread_pool = ThreadPoolExecutor(
max_workers=min(10, (os.cpu_count() or 1) + 4)
@@ -486,17 +526,22 @@ class LlamaExtract(BaseComponent):
def run_coro() -> T:
# Create a new client for this thread
async def wrapped_coro() -> T:
assert (
self._httpx_client is not None
), "httpx_client should be initialized"
# Create a new client with the same configuration as the original
async with httpx.AsyncClient(
timeout=self._async_client._client_wrapper.httpx_client.timeout,
verify=self.verify,
timeout=self.httpx_timeout,
) as client:
# Replace the client in the coro's context
original_client = self._async_client._client_wrapper.httpx_client
# Temporarily replace the client
self._async_client._client_wrapper.httpx_client = client
try:
return await coro
finally:
# Restore the original client
self._async_client._client_wrapper.httpx_client = (
original_client
self._httpx_client
)
return asyncio.run(wrapped_coro())
@@ -599,6 +644,8 @@ class LlamaExtract(BaseComponent):
num_workers=self.num_workers,
show_progress=self.show_progress,
verbose=self.verbose,
verify=self.verify,
httpx_timeout=self.httpx_timeout,
)
def list_agents(self) -> List[ExtractionAgent]:
+2 -2
View File
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
[tool.poetry]
name = "llama-parse"
version = "0.6.5"
version = "0.6.7"
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.5"
llama-cloud-services = ">=0.6.7"
[tool.poetry.group.dev.dependencies]
pytest = "^8.0.0"
Generated
+411 -429
View File
File diff suppressed because it is too large Load Diff
+2 -2
View File
@@ -8,7 +8,7 @@ python_version = "3.10"
[tool.poetry]
name = "llama-cloud-services"
version = "0.6.5"
version = "0.6.7"
description = "Tailored SDK clients for LlamaCloud services."
authors = ["Logan Markewich <logan@runllama.ai>"]
license = "MIT"
@@ -18,7 +18,7 @@ packages = [{include = "llama_cloud_services"}]
[tool.poetry.dependencies]
python = ">=3.9,<4.0"
llama-index-core = ">=0.11.0"
llama-cloud = "^0.1.14"
llama-cloud = "^0.1.16"
pydantic = "!=2.10"
click = "^8.1.7"
python-dotenv = "^1.0.1"
+9 -2
View File
@@ -182,7 +182,14 @@ class TestExtractionAgent:
assert "new_field" in updated_agent.data_schema["properties"]
def test_list_extraction_runs(self, test_agent: ExtractionAgent):
assert len(test_agent.list_extraction_runs()) == 0
assert test_agent.list_extraction_runs().total == 0
test_agent.extract(TEST_PDF)
runs = test_agent.list_extraction_runs()
assert len(runs) > 0
assert runs.total > 0
def test_delete_extraction_run(self, test_agent: ExtractionAgent):
assert test_agent.list_extraction_runs().total == 0
run = test_agent.extract(TEST_PDF)
test_agent.delete_extraction_run(run.id)
runs = test_agent.list_extraction_runs()
assert runs.total == 0
+5
View File
@@ -81,6 +81,11 @@ async def test_create_and_delete_report(
@pytest.mark.asyncio
@pytest.mark.xfail(
condition=lambda: os.getenv("CI"),
reason="Report plan sometimes times out",
raises=TimeoutError,
)
async def test_report_plan_workflow(report: ReportClient) -> None:
"""Test the report planning workflow."""
# Wait for the plan