feat: Add google.genai.types.Content as an allowed message type to ADK's stream_query method

PiperOrigin-RevId: 754073499

Author: vertex-sdk-bot

tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py

@@ -17,6 +17,7 @@
 from unittest import mock
 
 from google import auth
+from google.genai import types
 import vertexai
 from google.cloud.aiplatform import initializer
 from vertexai.preview import reasoning_engines
@@ -172,6 +173,28 @@ def test_stream_query(self):
         )
         assert len(events) == 1
 
+    def test_stream_query_with_content(self):
+        app = reasoning_engines.AdkApp(
+            agent=Agent(name="test_agent", model=_TEST_MODEL)
+        )
+        assert app._tmpl_attrs.get("runner") is None
+        app.set_up()
+        app._tmpl_attrs["runner"] = _MockRunner()
+        events = list(
+            app.stream_query(
+                user_id="test_user_id",
+                message=types.Content(
+                    role="user",
+                    parts=[
+                        types.Part(
+                            text="test message with content",
+                        )
+                    ],
+                ).model_dump(),
+            )
+        )
+        assert len(events) == 1
+
     def test_streaming_agent_run_with_events(self):
         app = reasoning_engines.AdkApp(
             agent=Agent(name="test_agent", model=_TEST_MODEL)

vertexai/preview/reasoning_engines/templates/adk.py

@@ -13,9 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+
 
 if TYPE_CHECKING:
+    try:
+        from google.genai import types
+
+        ContentDict = types.Content
+    except (ImportError, AttributeError):
+        ContentDict = Dict
+
     try:
         from google.adk.events.event import Event
 
@@ -442,7 +450,7 @@ def set_up(self):
     def stream_query(
         self,
         *,
-        message: str,
+        message: Union[str, "ContentDict"],
         user_id: str,
         session_id: Optional[str] = None,
         **kwargs,
@@ -466,7 +474,15 @@ def stream_query(
         """
         from google.genai import types
 
-        content = types.Content(role="user", parts=[types.Part(text=message)])
+        if isinstance(message, Dict):
+            content = types.Content.model_validate(message)
+        elif isinstance(message, str):
+            content = types.Content(role="user", parts=[types.Part(text=message)])
+        else:
+            raise TypeError(
+                "message must be a string or a dictionary representing a Content object."
+            )
+
         if not self._tmpl_attrs.get("runner"):
             self.set_up()
         if not session_id: