Address review comments

Author: dmandar

samples/agent/adk/component_gallery/__main__.py

@@ -14,7 +14,6 @@
 from starlette.middleware.cors import CORSMiddleware
 from dotenv import load_dotenv
 
-from agent import ComponentGalleryAgent
 from agent_executor import ComponentGalleryExecutor
 
 load_dotenv()

samples/agent/adk/component_gallery/agent.py

@@ -4,12 +4,9 @@
 import json
 from collections.abc import AsyncIterable
 from typing import Any
+import asyncio
+import datetime
 
-from google.adk.agents.llm_agent import LlmAgent
-from google.adk.artifacts import InMemoryArtifactService
-from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
-from google.adk.runners import Runner
-from google.adk.sessions import InMemorySessionService
 from gallery_examples import get_gallery_json
 
 logger = logging.getLogger(__name__)
@@ -28,19 +25,20 @@ async def stream(self, query: str, session_id: str) -> AsyncIterable[dict[str, A
         # Initial Load or Reset
         if "WHO_ARE_YOU" in query or "START" in query: # Simple trigger for initial load
              gallery_json = get_gallery_json()
-             response = f"Here is the component gallery.\n---a2ui_JSON---\n{gallery_json}"
              yield {
                 "is_task_complete": True,
-                "content": response
+                "payload": {
+                    "text": "Here is the component gallery.",
+                    "json_string": gallery_json
+                }
              }
              return
 
         # Handle Actions
         if query.startswith("ACTION:"):
              action_name = query
              # Create a response update for the second surface
-             import datetime
-             import asyncio
+             # Create a response update for the second surface
 
              # Simulate network/processing delay
              await asyncio.sleep(0.5)
@@ -63,16 +61,20 @@ async def stream(self, query: str, session_id: str) -> AsyncIterable[dict[str, A
                  }
              ]
 
-             json_str = json.dumps(response_update)
-             response = f"Action processed.\n---a2ui_JSON---\n{json_str}"
+             
              yield {
                 "is_task_complete": True,
-                "content": response
+                "payload": {
+                    "text": "Action processed.",
+                    "json_data": response_update
+                }
              }
              return
 
         # Fallback for text
         yield {
              "is_task_complete": True,
-             "content": "I am the Component Gallery Agent."
+             "payload": {
+                 "text": "I am the Component Gallery Agent."
+             }
         }

samples/agent/adk/component_gallery/agent_executor.py

@@ -52,26 +52,48 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non
         updater = TaskUpdater(event_queue, task.id, task.context_id)
 
         async for item in self.agent.stream(query, task.context_id):
-             content = item["content"]
              final_parts = []
-             
-             if "---a2ui_JSON---" in content:
-                 text_content, json_string = content.split("---a2ui_JSON---", 1)
-                 if text_content.strip():
-                     final_parts.append(Part(root=TextPart(text=text_content.strip())))
+
+             if "payload" in item:
+                 payload = item["payload"]
+                 text = payload.get("text")
+                 if text:
+                     final_parts.append(Part(root=TextPart(text=text)))
+                 
+                 json_data = payload.get("json_data")
+                 json_string = payload.get("json_string")
 
-                 if json_string.strip():
+                 if json_string:
                      try:
-                         json_data = json.loads(json_string.strip())
-                         if isinstance(json_data, list):
-                             for msg in json_data:
-                                 final_parts.append(create_a2ui_part(msg))
-                         else:
-                             final_parts.append(create_a2ui_part(json_data))
+                         json_data = json.loads(json_string)
                      except Exception as e:
-                         logger.error(f"Failed to parse JSON: {e}")
+                         logger.error(f"Failed to parse JSON string: {e}")
+                 
+                 if json_data:
+                     if isinstance(json_data, list):
+                         for msg in json_data:
+                             final_parts.append(create_a2ui_part(msg))
+                     else:
+                         final_parts.append(create_a2ui_part(json_data))
              else:
-                 final_parts.append(Part(root=TextPart(text=content)))
+                 content = item.get("content", "")
+                 if "---a2ui_JSON---" in content:
+                     text_content, json_string = content.split("---a2ui_JSON---", 1)
+                     if text_content.strip():
+                         final_parts.append(Part(root=TextPart(text=text_content.strip())))
+                     
+                     if json_string.strip():
+                         try:
+                             json_data = json.loads(json_string.strip())
+                             if isinstance(json_data, list):
+                                 for msg in json_data:
+                                     final_parts.append(create_a2ui_part(msg))
+                             else:
+                                 final_parts.append(create_a2ui_part(json_data))
+                         except Exception as e:
+                             logger.error(f"Failed to parse JSON: {e}")
+                 elif content:
+                     final_parts.append(Part(root=TextPart(text=content)))
 
              await updater.update_status(
                  TaskState.completed,