Compare commits

...

2 Commits

Author SHA1 Message Date
Neeraj Pradhan 65cff44f43 type fix 2025-03-21 19:41:43 -07:00
Neeraj Pradhan f57592c531 Revert httpx_client as argument 2025-03-21 19:39:11 -07:00
+15 -34
View File
@@ -57,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
@@ -66,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
@@ -83,17 +87,9 @@ class ExtractionAgent:
# Create a new client with the same configuration as the original
async with httpx.AsyncClient(
timeout=original_client.timeout,
trust_env=original_client.trust_env,
base_url=original_client.base_url,
follow_redirects=original_client.follow_redirects,
transport=original_client._transport,
default_encoding=original_client._default_encoding,
verify=self.verify,
timeout=self.httpx_timeout,
) as client:
# Copy headers and cookies
client.headers = original_client.headers.copy()
client.cookies = httpx.Cookies(original_client.cookies)
# Temporarily replace the client
self._client._client_wrapper.httpx_client = client
try:
@@ -450,13 +446,12 @@ class LlamaExtract(BaseComponent):
verbose: bool = Field(
default=False, description="Show verbose output when extracting files."
)
verify: Union[bool, str] = Field(
verify: Optional[bool] = Field(
default=True, description="Simple SSL verification option."
)
httpx_timeout: Optional[float] = Field(
default=60, description="Timeout for the httpx client."
)
_httpx_client: Optional[httpx.AsyncClient] = PrivateAttr()
_async_client: AsyncLlamaCloud = PrivateAttr()
_thread_pool: ThreadPoolExecutor = PrivateAttr()
_project_id: Optional[str] = PrivateAttr()
@@ -472,10 +467,9 @@ class LlamaExtract(BaseComponent):
show_progress: bool = True,
project_id: Optional[str] = None,
organization_id: Optional[str] = None,
verbose: bool = False,
verify: Union[bool, str] = True,
verify: Optional[bool] = True,
httpx_timeout: Optional[float] = 60,
httpx_client: Optional[httpx.AsyncClient] = None,
verbose: bool = False,
):
if not api_key:
api_key = os.getenv("LLAMA_CLOUD_API_KEY", None)
@@ -492,18 +486,11 @@ class LlamaExtract(BaseComponent):
max_timeout=max_timeout,
num_workers=num_workers,
show_progress=show_progress,
verify=verify,
httpx_timeout=httpx_timeout,
verbose=verbose,
)
if httpx_client is not None:
if not verify:
warnings.warn(
"Both 'httpx_client' and 'verify=False' were provided. "
"The custom httpx_client takes precedence and 'verify' will be ignored."
)
self._httpx_client = httpx_client
else:
self._httpx_client = httpx.AsyncClient(verify=verify, timeout=httpx_timeout)
self._httpx_client = httpx.AsyncClient(verify=verify, timeout=httpx_timeout)
self.verify = verify
self.httpx_timeout = httpx_timeout
@@ -544,17 +531,9 @@ class LlamaExtract(BaseComponent):
), "httpx_client should be initialized"
# Create a new client with the same configuration as the original
async with httpx.AsyncClient(
verify=self.verify,
timeout=self.httpx_timeout,
trust_env=self._httpx_client.trust_env,
base_url=self._httpx_client.base_url,
follow_redirects=self._httpx_client.follow_redirects,
transport=self._httpx_client._transport,
default_encoding=self._httpx_client._default_encoding,
) as client:
# Copy headers and cookies
client.headers = self._httpx_client.headers.copy()
client.cookies = httpx.Cookies(self._httpx_client.cookies)
# Temporarily replace the client
self._async_client._client_wrapper.httpx_client = client
try:
@@ -665,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]: