-
Notifications
You must be signed in to change notification settings - Fork 3.6k
feat(tools): expose httpx_client_factory on RestApiTool and OpenAPITo… #5715
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 3 commits
3074cd8
11f22b8
be08e06
d2bb078
70210dd
f6f83bc
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -75,6 +75,17 @@ def snake_to_lower_camel(snake_case_string: str): | |
|
|
||
| AuthPreparationState = Literal["pending", "done"] | ||
|
|
||
| HttpxClientFactory = Callable[..., httpx.AsyncClient] | ||
| """Type alias for a factory returning an ``httpx.AsyncClient``. | ||
|
|
||
| When supplied to ``RestApiTool`` or ``OpenAPIToolset``, the factory is invoked | ||
| once per API call and its returned client is used (as an async context | ||
| manager) to issue the request, in place of the default | ||
| ``httpx.AsyncClient(verify=..., timeout=None)``. This unlocks knobs that the | ||
| narrower ``ssl_verify`` parameter can't reach: proxies, HTTP/2, custom | ||
| transports (e.g. request-signing), shared connection pools, and so on. | ||
| """ | ||
|
|
||
|
|
||
| class RestApiTool(BaseTool): | ||
| """A generic tool that interacts with a REST API. | ||
|
|
@@ -103,6 +114,7 @@ def __init__( | |
| header_provider: Optional[ | ||
| Callable[[ReadonlyContext], Dict[str, str]] | ||
| ] = None, | ||
| httpx_client_factory: Optional[HttpxClientFactory] = None, | ||
| *, | ||
| credential_key: Optional[str] = None, | ||
| ): | ||
|
|
@@ -142,6 +154,15 @@ def __init__( | |
| an argument, allowing dynamic header generation based on the current | ||
| context. Useful for adding custom headers like correlation IDs, | ||
| authentication tokens, or other request metadata. | ||
| httpx_client_factory: Optional zero-argument callable returning an | ||
| ``httpx.AsyncClient``. When provided, the returned client is used to | ||
| issue the request, allowing callers to configure proxies, HTTP/2, | ||
| custom transports (e.g. request signing), shared connection pools, | ||
| or any other ``httpx.AsyncClient`` option that ``ssl_verify`` can't | ||
| reach. When ``None`` (default), behaviour is unchanged: a fresh | ||
| ``httpx.AsyncClient(verify=..., timeout=None)`` is created per | ||
| request. Mirrors the pattern exposed for MCP by | ||
| ``StreamableHTTPConnectionParams.httpx_client_factory``. | ||
| credential_key: Optional stable key used for interactive auth and | ||
| credential caching. | ||
| """ | ||
|
|
@@ -169,6 +190,7 @@ def __init__( | |
| self._default_headers: Dict[str, str] = {} | ||
| self._ssl_verify = ssl_verify | ||
| self._header_provider = header_provider | ||
| self._httpx_client_factory = httpx_client_factory | ||
| self._logger = logger | ||
| if should_parse_operation: | ||
| self._operation_parser = OperationParser(self.operation) | ||
|
|
@@ -181,6 +203,7 @@ def from_parsed_operation( | |
| header_provider: Optional[ | ||
| Callable[[ReadonlyContext], Dict[str, str]] | ||
| ] = None, | ||
| httpx_client_factory: Optional[HttpxClientFactory] = None, | ||
| ) -> "RestApiTool": | ||
| """Initializes the RestApiTool from a ParsedOperation object. | ||
|
|
||
|
|
@@ -192,6 +215,9 @@ def from_parsed_operation( | |
| an argument, allowing dynamic header generation based on the current | ||
| context. Useful for adding custom headers like correlation IDs, | ||
| authentication tokens, or other request metadata. | ||
| httpx_client_factory: Optional zero-argument callable returning an | ||
| ``httpx.AsyncClient`` to be used for the API call. See | ||
| ``RestApiTool.__init__`` for details. | ||
|
|
||
| Returns: | ||
| A RestApiTool object. | ||
|
|
@@ -212,6 +238,7 @@ def from_parsed_operation( | |
| auth_credential=parsed.auth_credential, | ||
| ssl_verify=ssl_verify, | ||
| header_provider=header_provider, | ||
| httpx_client_factory=httpx_client_factory, | ||
| ) | ||
| generated._operation_parser = operation_parser | ||
| return generated | ||
|
|
@@ -520,7 +547,9 @@ async def call( | |
| if provider_headers: | ||
| request_params.setdefault("headers", {}).update(provider_headers) | ||
|
|
||
| response = await _request(**request_params) | ||
| response = await _request( | ||
| httpx_client_factory=self._httpx_client_factory, **request_params | ||
| ) | ||
|
|
||
| # Log the API response | ||
| self._logger.debug( | ||
|
|
@@ -575,9 +604,14 @@ def __repr__(self): | |
| ) | ||
|
|
||
|
|
||
| async def _request(**request_params) -> httpx.Response: | ||
| async with httpx.AsyncClient( | ||
| verify=request_params.pop("verify", True), | ||
| timeout=None, | ||
| ) as client: | ||
| async def _request( | ||
| *, | ||
| httpx_client_factory: Optional[HttpxClientFactory] = None, | ||
| **request_params, | ||
| ) -> httpx.Response: | ||
| verify = request_params.pop("verify", True) | ||
| if httpx_client_factory is not None: | ||
| async with httpx_client_factory() as client: | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. With
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes, intentional — the client is used as an async context manager and closed after the request, mirroring the default branch. I've documented this in the docstrings (factory must return a fresh client per call) and dropped the inaccurate "shared connection pools" wording. Reusing a pooled client would require holding a persistent client on the tool, which I think is best as a separate change — happy to follow up if you'd prefer that. |
||
| return await client.request(**request_params) | ||
| async with httpx.AsyncClient(verify=verify, timeout=None) as client: | ||
| return await client.request(**request_params) | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This docstring says "zero-argument callable", but the
HttpxClientFactoryis defined asCallable[..., httpx.AsyncClient], which takes any arguments, is this inconsistency intended?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
fixed