feat: Forward per-call http_options from AgentEngine.query family of methods.

PiperOrigin-RevId: 919896831

Author: vertex-sdk-bot

agentplatform/_genai/_agent_engines_utils.py

@@ -1630,6 +1630,18 @@ def _wrap_query_operation(*, method_name: str) -> Callable[..., Any]:
     Agent Engine method using the `query` API.  It handles the creation of
     the API request and the processing of the API response.
 
+    The reserved keyword argument `http_options` is consumed by this
+    wrapper (rather than being forwarded to the deployed agent as part of
+    `input`) and is propagated to the underlying HTTP call. Use it to set
+    per-call HTTP options such as custom headers. For example:
+
+        from google.genai.types import HttpOptions
+
+        agent_engine.query(
+            input="hello",
+            http_options=HttpOptions(headers={"x-my-header": "value"}),
+        )
+
     Args:
         method_name: The name of the Agent Engine method to call.
         doc: Documentation string for the method.
@@ -1644,12 +1656,14 @@ def _method(self: genai_types.AgentEngine, **kwargs) -> Any:  # type: ignore[no-
             raise ValueError("api_client is not initialized.")
         if not self.api_resource:
             raise ValueError("api_resource is not initialized.")
+        http_options = kwargs.pop("http_options", None)
         response = self.api_client._query(
             name=self.api_resource.name,
             config={
                 "class_method": method_name,
                 "input": kwargs,
                 "include_all_fields": True,
+                "http_options": http_options,
             },
         )
         return response.output
@@ -1666,6 +1680,10 @@ def _wrap_async_query_operation(
     Agent Engine method asynchronously using the `query` API. It handles the
     creation of the API request and the processing of the API response.
 
+    The reserved keyword argument `http_options` is consumed by this
+    wrapper (rather than being forwarded to the deployed agent as part of
+    `input`) and is propagated to the underlying HTTP call.
+
     Args:
         method_name: The name of the Agent Engine method to call.
         doc: Documentation string for the method.
@@ -1682,12 +1700,14 @@ async def _method(
             raise ValueError("api_async_client is not initialized.")
         if not self.api_resource:
             raise ValueError("api_resource is not initialized.")
+        http_options = kwargs.pop("http_options", None)
         response = await self.api_async_client._query(
             name=self.api_resource.name,
             config={
                 "class_method": method_name,
                 "input": kwargs,
                 "include_all_fields": True,
+                "http_options": http_options,
             },
         )
         return response.output
@@ -1702,6 +1722,10 @@ def _wrap_stream_query_operation(*, method_name: str) -> Callable[..., Iterator[
     Agent Engine method using the `stream_query` API.  It handles the
     creation of the API request and the processing of the API response.
 
+    The reserved keyword argument `http_options` is consumed by this
+    wrapper (rather than being forwarded to the deployed agent as part of
+    `input`) and is propagated to the underlying HTTP call.
+
     Args:
         method_name: The name of the Agent Engine method to call.
         doc: Documentation string for the method.
@@ -1716,12 +1740,14 @@ def _method(self: genai_types.AgentEngine, **kwargs) -> Iterator[Any]:  # type:
             raise ValueError("api_client is not initialized.")
         if not self.api_resource:
             raise ValueError("api_resource is not initialized.")
+        http_options = kwargs.pop("http_options", None)
         for http_response in self.api_client._stream_query(
             name=self.api_resource.name,
             config={
                 "class_method": method_name,
                 "input": kwargs,
                 "include_all_fields": True,
+                "http_options": http_options,
             },
         ):
             for line in _yield_parsed_json(http_response=http_response):
@@ -1740,6 +1766,10 @@ def _wrap_async_stream_query_operation(
     Agent Engine method using the `stream_query` API.  It handles the
     creation of the API request and the processing of the API response.
 
+    The reserved keyword argument `http_options` is consumed by this
+    wrapper (rather than being forwarded to the deployed agent as part of
+    `input`) and is propagated to the underlying HTTP call.
+
     Args:
         method_name: The name of the Agent Engine method to call.
         doc: Documentation string for the method.
@@ -1754,12 +1784,14 @@ async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]
             raise ValueError("api_client is not initialized.")
         if not self.api_resource:
             raise ValueError("api_resource is not initialized.")
+        http_options = kwargs.pop("http_options", None)
         async for http_response in self.api_client._async_stream_query(
             name=self.api_resource.name,
             config={
                 "class_method": method_name,
                 "input": kwargs,
                 "include_all_fields": True,
+                "http_options": http_options,
             },
         ):
             for line in _yield_parsed_json(http_response=http_response):

tests/unit/agentplatform/genai/test_agent_engines.py

@@ -427,6 +427,7 @@ def register_operations(self) -> Dict[str, List[str]]:
 """
 _TEST_METHOD_TO_BE_UNREGISTERED_NAME = "method_to_be_unregistered"
 _TEST_QUERY_PROMPT = "Find the first fibonacci number greater than 999"
+_UNSET_HTTP_OPTIONS = object()
 _TEST_AGENT_ENGINE_ENV_KEY = "GOOGLE_CLOUD_AGENT_ENGINE_ENV"
 _TEST_AGENT_ENGINE_ENV_VALUE = "test_env_value"
 _TEST_AGENT_ENGINE_GCS_URI = "{}/{}/{}".format(
@@ -3068,7 +3069,26 @@ def test_delete_agent_engine_force(self):
             None,
         )
 
-    def test_query_agent_engine(self):
+    @pytest.mark.parametrize(
+        "http_options_arg, expected_http_options",
+        [
+            (_UNSET_HTTP_OPTIONS, None),
+            (
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+            ),
+            (
+                {"headers": {"x-my-header": "my-value"}},
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+            ),
+        ],
+        ids=["default", "http_options_instance", "http_options_dict"],
+    )
+    def test_query_agent_engine(self, http_options_arg, expected_http_options):
+        """Sync query: forwards http_options to the HTTP layer (or None by default)."""
+        kwargs = {"query": _TEST_QUERY_PROMPT}
+        if http_options_arg is not _UNSET_HTTP_OPTIONS:
+            kwargs["http_options"] = http_options_arg
         with mock.patch.object(
             self.client.agent_engines._api_client, "request"
         ) as request_mock:
@@ -3086,7 +3106,7 @@ def test_query_agent_engine(self):
                     ),
                 )
             )
-            agent.query(query=_TEST_QUERY_PROMPT)
+            agent.query(**kwargs)
             request_mock.assert_called_with(
                 "post",
                 f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}:query",
@@ -3095,7 +3115,7 @@ def test_query_agent_engine(self):
                     "classMethod": "query",
                     "input": {"query": _TEST_QUERY_PROMPT},
                 },
-                None,
+                expected_http_options,
             )
 
     @mock.patch("google.cloud.storage.Client")
@@ -3316,7 +3336,26 @@ def test_run_query_job_agent_engine_directory_no_slash(
                 output_gcs_uri="gs://my-input-bucket/path/b92b9b89-4585-4146-8ee5-22fe99802a8e_output.json",
             )
 
-    def test_query_agent_engine_async(self):
+    @pytest.mark.parametrize(
+        "http_options_arg, expected_http_options",
+        [
+            (_UNSET_HTTP_OPTIONS, None),
+            (
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+            ),
+            (
+                {"headers": {"x-my-header": "my-value"}},
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+            ),
+        ],
+        ids=["default", "http_options_instance", "http_options_dict"],
+    )
+    def test_query_agent_engine_async(self, http_options_arg, expected_http_options):
+        """Async query: forwards http_options to the HTTP layer (or None by default)."""
+        kwargs = {"query": _TEST_QUERY_PROMPT}
+        if http_options_arg is not _UNSET_HTTP_OPTIONS:
+            kwargs["http_options"] = http_options_arg
         agent = self.client.agent_engines._register_api_methods(
             agent_engine=_genai_types.AgentEngine(
                 api_async_client=agent_engines.AsyncAgentEngines(
@@ -3336,7 +3375,7 @@ def test_query_agent_engine_async(self):
             self.client.agent_engines._api_client, "async_request"
         ) as request_mock:
             request_mock.return_value = genai_types.HttpResponse(body="")
-            asyncio.run(agent.async_query(query=_TEST_QUERY_PROMPT))
+            asyncio.run(agent.async_query(**kwargs))
             request_mock.assert_called_with(
                 "post",
                 f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}:query",
@@ -3345,7 +3384,7 @@ def test_query_agent_engine_async(self):
                     "classMethod": "async_query",
                     "input": {"query": _TEST_QUERY_PROMPT},
                 },
-                None,
+                expected_http_options,
             )
 
     def test_cancel_query_job_agent_engine(self):
@@ -3507,7 +3546,26 @@ def test_check_query_job_agent_engine_blob_not_exists(self):
                         config={"retrieve_result": True},
                     )
 
-    def test_query_agent_engine_stream(self):
+    @pytest.mark.parametrize(
+        "http_options_arg, expected_http_options",
+        [
+            (_UNSET_HTTP_OPTIONS, None),
+            (
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+            ),
+            (
+                {"headers": {"x-my-header": "my-value"}},
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+            ),
+        ],
+        ids=["default", "http_options_instance", "http_options_dict"],
+    )
+    def test_query_agent_engine_stream(self, http_options_arg, expected_http_options):
+        """Streaming query: forwards http_options to the HTTP layer (or None by default)."""
+        kwargs = {"query": _TEST_QUERY_PROMPT}
+        if http_options_arg is not _UNSET_HTTP_OPTIONS:
+            kwargs["http_options"] = http_options_arg
         with mock.patch.object(
             self.client.agent_engines._api_client, "request_streamed"
         ) as request_mock:
@@ -3524,7 +3582,7 @@ def test_query_agent_engine_stream(self):
                     ),
                 )
             )
-            list(agent.stream_query(query=_TEST_QUERY_PROMPT))
+            list(agent.stream_query(**kwargs))
             request_mock.assert_called_with(
                 "post",
                 f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}:streamQuery?alt=sse",
@@ -3533,11 +3591,34 @@ def test_query_agent_engine_stream(self):
                     "classMethod": "stream_query",
                     "input": {"query": _TEST_QUERY_PROMPT},
                 },
-                None,
+                expected_http_options,
             )
 
-    def test_query_agent_engine_async_stream(self):
+    @pytest.mark.parametrize(
+        "http_options_arg, expected_http_options",
+        [
+            (_UNSET_HTTP_OPTIONS, None),
+            (
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+            ),
+            (
+                {"headers": {"x-my-header": "my-value"}},
+                genai_types.HttpOptions(headers={"x-my-header": "my-value"}),
+            ),
+        ],
+        ids=["default", "http_options_instance", "http_options_dict"],
+    )
+    def test_query_agent_engine_async_stream(
+        self, http_options_arg, expected_http_options
+    ):
+        """Async streaming query: forwards http_options to the HTTP layer (or None)."""
+        kwargs = {"query": _TEST_QUERY_PROMPT}
+        if http_options_arg is not _UNSET_HTTP_OPTIONS:
+            kwargs["http_options"] = http_options_arg
+
         async def mock_async_generator():
+            """Fake async-streamed HTTP response generator."""
             yield genai_types.HttpResponse(body=b"")
 
         with mock.patch.object(
@@ -3559,9 +3640,8 @@ async def mock_async_generator():
             )
 
             async def consume():
-                async for response in agent.async_stream_query(
-                    query=_TEST_QUERY_PROMPT
-                ):
+                """Drain the async iterator."""
+                async for response in agent.async_stream_query(**kwargs):
                     print(response)
 
             asyncio.run(consume())
@@ -3573,7 +3653,7 @@ async def consume():
                     "classMethod": "async_stream_query",
                     "input": {"query": _TEST_QUERY_PROMPT},
                 },
-                None,
+                expected_http_options,
             )
 
     @pytest.mark.parametrize(