feat: GenAI SDK client - add zero-shot prompt optimizer (streaming): an option to quickly improve or generate system instructions or a single prompt.

PiperOrigin-RevId: 788830431

Author: vertex-sdk-bot

tests/unit/vertexai/genai/replays/test_optmize_prompt_return_type.py

@@ -0,0 +1,44 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# pylint: disable=protected-access,bad-continuation,missing-function-docstring
+
+import logging
+
+from tests.unit.vertexai.genai.replays import pytest_helper
+
+# from vertexai._genai import types
+
+logger = logging.getLogger("vertexai_genai.promptoptimizer")
+logging.basicConfig(encoding="utf-8", level=logging.INFO, force=True)
+
+
+def test_optimize_prompt(client):
+    """Tests the optimize request parameters method."""
+
+    client._api_client._http_options.base_url = (
+        "https://us-central1-autopush-aiplatform.sandbox.googleapis.com"
+    )
+    test_prompt = "Generate system instructions for analyzing medical articles"
+    for chunk in client.prompt_optimizer.optimize_prompt(prompt=test_prompt):
+        print("chunk: %s" % chunk)
+    # logger.info("response: %s", response)
+    # assert isinstance(response, types.OptimizeResponse)
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="prompt_optimizer.optimize_prompt",
+)

vertexai/_genai/prompt_optimizer.py

@@ -19,7 +19,7 @@
 import json
 import logging
 import time
-from typing import Any, Optional, Union
+from typing import Any, AsyncIterator, Awaitable, Iterator, Optional, Union
 from urllib.parse import urlencode
 
 from google.genai import _api_module
@@ -32,13 +32,17 @@
 
 
 logger = logging.getLogger("vertexai_genai.promptoptimizer")
+logging.basicConfig(encoding="utf-8", level=logging.INFO, force=True)
 
 
 def _OptimizeRequestParameters_to_vertex(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any]:
     to_object: dict[str, Any] = {}
+    if getv(from_object, ["content"]) is not None:
+        setv(to_object, ["content"], getv(from_object, ["content"]))
+
     if getv(from_object, ["config"]) is not None:
         setv(to_object, ["config"], getv(from_object, ["config"]))
 
@@ -229,6 +233,8 @@ def _OptimizeResponse_from_vertex(
     parent_object: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any]:
     to_object: dict[str, Any] = {}
+    if getv(from_object, ["content"]) is not None:
+        setv(to_object, ["content"], getv(from_object, ["content"]))
 
     return to_object
 
@@ -383,12 +389,16 @@ def _CustomJob_from_vertex(
 class PromptOptimizer(_api_module.BaseModule):
     """Prompt Optimizer"""
 
-    def _optimize_dummy(
-        self, *, config: Optional[types.OptimizeConfigOrDict] = None
-    ) -> types.OptimizeResponse:
-        """Optimize multiple prompts."""
+    def _optimize_prompt(
+        self,
+        *,
+        content: Optional[types.ContentOrDict] = None,
+        config: Optional[types.OptimizeConfigOrDict] = None,
+    ) -> Iterator[types.OptimizeResponse]:
+        """Optimize a single prompt."""
 
         parameter_model = types._OptimizeRequestParameters(
+            content=content,
             config=config,
         )
 
@@ -399,9 +409,9 @@ def _optimize_dummy(
             request_dict = _OptimizeRequestParameters_to_vertex(parameter_model)
             request_url_dict = request_dict.get("_url")
             if request_url_dict:
-                path = ":optimize".format_map(request_url_dict)
+                path = "tuningJobs:optimizePrompt".format_map(request_url_dict)
             else:
-                path = ":optimize"
+                path = "tuningJobs:optimizePrompt"
 
         query_params = request_dict.get("_query")
         if query_params:
@@ -419,19 +429,32 @@ def _optimize_dummy(
         request_dict = _common.convert_to_dict(request_dict)
         request_dict = _common.encode_unserializable_types(request_dict)
 
-        response = self._api_client.request("post", path, request_dict, http_options)
-
-        response_dict = "" if not response.body else json.loads(response.body)
-
-        if self._api_client.vertexai:
-            response_dict = _OptimizeResponse_from_vertex(response_dict)
-
-        return_value = types.OptimizeResponse._from_response(
-            response=response_dict, kwargs=parameter_model.model_dump()
-        )
+        if config is not None and getattr(config, "should_return_http_response", None):
+            raise ValueError(
+                "Accessing the raw HTTP response is not supported in streaming"
+                " methods."
+            )
 
-        self._api_client._verify_response(return_value)
-        return return_value
+        for response in self._api_client.request_streamed(
+            "post", path, request_dict, http_options
+        ):
+            # print("response: %s" % response)
+            logger.info("response: %s", response)
+            response_dict = "" if not response.body else json.loads(response.body)
+            # print("response_dict: %s" % response_dict)
+            logger.info("response_dict: %s", response_dict)
+            if self._api_client.vertexai:
+                response_dict = _OptimizeResponse_from_vertex(response_dict)
+                # print("response_dict vertexai: %s" % response_dict)
+                logger.info("response_dict vertexai: %s", response_dict)
+
+            return_value = types.OptimizeResponse._from_response(
+                response=response_dict, kwargs=parameter_model.model_dump()
+            )
+            # print("return_value: %s" % return_value)
+            logger.info("return_value: %s", return_value)
+            self._api_client._verify_response(return_value)
+            yield return_value
 
     def _create_custom_job_resource(
         self,
@@ -660,16 +683,45 @@ def optimize(
             job = self._wait_for_completion(job_id)
         return job
 
+    def optimize_prompt(
+        self, *, prompt: str, config: Optional[types.OptimizeConfig] = None
+    ) -> Iterator[types.OptimizeResponse]:
+        """Makes an API request to _optimize_prompt and yields the optimized prompt in chunks."""
+        if config is not None:
+            raise ValueError(
+                "Currently, config is not supported for a single prompt"
+                " optimization."
+            )
+
+        prompt = types.Content(parts=[types.Part(text=prompt)], role="user")
+        # response = self._optimize_prompt(content=prompt)
+        # logger.info(type(response))
+        # logger.info(response)
+
+        # for chunk in response:
+        #   yield chunk
+        for chunk in self._optimize_prompt(content=prompt):
+            #   logger.info(chunk)
+            #   if chunk.content and chunk.content.parts[0].text:
+            #     logger.info('chunk has content text %s', chunk.content.parts[0].text)
+            #   if chunk.parts[0]:
+            #     logger.info('chunk has parts %s', chunk.parts[0])
+            yield chunk
+
 
 class AsyncPromptOptimizer(_api_module.BaseModule):
     """Prompt Optimizer"""
 
-    async def _optimize_dummy(
-        self, *, config: Optional[types.OptimizeConfigOrDict] = None
-    ) -> types.OptimizeResponse:
-        """Optimize multiple prompts."""
+    async def _optimize_prompt(
+        self,
+        *,
+        content: Optional[types.ContentOrDict] = None,
+        config: Optional[types.OptimizeConfigOrDict] = None,
+    ) -> Awaitable[AsyncIterator[types.OptimizeResponse]]:
+        """Optimize a single prompt."""
 
         parameter_model = types._OptimizeRequestParameters(
+            content=content,
             config=config,
         )
 
@@ -680,9 +732,9 @@ async def _optimize_dummy(
             request_dict = _OptimizeRequestParameters_to_vertex(parameter_model)
             request_url_dict = request_dict.get("_url")
             if request_url_dict:
-                path = ":optimize".format_map(request_url_dict)
+                path = "tuningJobs:optimizePrompt".format_map(request_url_dict)
             else:
-                path = ":optimize"
+                path = "tuningJobs:optimizePrompt"
 
         query_params = request_dict.get("_query")
         if query_params:
@@ -700,21 +752,32 @@ async def _optimize_dummy(
         request_dict = _common.convert_to_dict(request_dict)
         request_dict = _common.encode_unserializable_types(request_dict)
 
-        response = await self._api_client.async_request(
+        if config is not None and getattr(config, "should_return_http_response", None):
+            raise ValueError(
+                "Accessing the raw HTTP response is not supported in streaming"
+                " methods."
+            )
+
+        response_stream = await self._api_client.async_request_streamed(
             "post", path, request_dict, http_options
         )
 
-        response_dict = "" if not response.body else json.loads(response.body)
+        async def async_generator():  # type: ignore[no-untyped-def]
+            async for response in response_stream:
 
-        if self._api_client.vertexai:
-            response_dict = _OptimizeResponse_from_vertex(response_dict)
+                response_dict = "" if not response.body else json.loads(response.body)
 
-        return_value = types.OptimizeResponse._from_response(
-            response=response_dict, kwargs=parameter_model.model_dump()
-        )
+                if self._api_client.vertexai:
+                    response_dict = _OptimizeResponse_from_vertex(response_dict)
 
-        self._api_client._verify_response(return_value)
-        return return_value
+                return_value = types.OptimizeResponse._from_response(
+                    response=response_dict, kwargs=parameter_model.model_dump()
+                )
+
+                self._api_client._verify_response(return_value)
+                yield return_value
+
+        return async_generator()  # type: ignore[no-untyped-call, no-any-return]
 
     async def _create_custom_job_resource(
         self,