From a4f0cae061731f580a19b8f84bf27ed3c81ff873 Mon Sep 17 00:00:00 2001 From: Jacob Simionato Date: Wed, 17 Dec 2025 09:27:59 +1030 Subject: [PATCH 1/4] Add template-based generation for restaurant finder --- .../adk/restaurant_finder/a2ui_validator.py | 50 + samples/agent/adk/restaurant_finder/agent.py | 223 ++--- .../adk/restaurant_finder/agent_executor.py | 12 + .../adk/restaurant_finder/generate_schema.py | 6 + .../adk/restaurant_finder/instrumentation.py | 58 ++ .../adk/restaurant_finder/prompt_builder.py | 935 ++---------------- .../restaurant_finder/template_renderer.py | 181 ++++ .../agent/adk/restaurant_finder/test_copy.py | 16 + samples/agent/adk/restaurant_finder/tools.py | 63 +- .../agent/adk/restaurant_finder/ui_schema.py | 33 + .../adk/restaurant_finder/verify_server.py | 34 + 11 files changed, 604 insertions(+), 1007 deletions(-) create mode 100644 samples/agent/adk/restaurant_finder/a2ui_validator.py create mode 100644 samples/agent/adk/restaurant_finder/generate_schema.py create mode 100644 samples/agent/adk/restaurant_finder/instrumentation.py create mode 100644 samples/agent/adk/restaurant_finder/template_renderer.py create mode 100644 samples/agent/adk/restaurant_finder/test_copy.py create mode 100644 samples/agent/adk/restaurant_finder/ui_schema.py create mode 100644 samples/agent/adk/restaurant_finder/verify_server.py diff --git a/samples/agent/adk/restaurant_finder/a2ui_validator.py b/samples/agent/adk/restaurant_finder/a2ui_validator.py new file mode 100644 index 00000000..b70e4aca --- /dev/null +++ b/samples/agent/adk/restaurant_finder/a2ui_validator.py @@ -0,0 +1,50 @@ +import json +import os +from typing import List, Dict, Any +from jsonschema import validate, ValidationError + +SCHEMA_FILE = "/Users/jsimionato/development/a2ui_repos/jewel_case/A2UI/specification/0.8/json/server_to_client_with_standard_catalog.json" + +SCHEMAS = {} + +def load_schemas(): + if SCHEMAS: + return + try: + with open(SCHEMA_FILE, 'r') as f: + main_schema = json.load(f) + SCHEMAS.update(main_schema.get('properties', {})) + if not SCHEMAS: + raise ValueError("No properties found in main schema") + # The actual component schemas are in the definitions, so let's add those too + SCHEMAS.update(main_schema.get('$defs', {})) + except FileNotFoundError: + raise ValueError(f"Schema file not found at {SCHEMA_FILE}") + except json.JSONDecodeError: + raise ValueError(f"Failed to decode JSON from {SCHEMA_FILE}") + +def get_schema(name: str): + load_schemas() + if name in SCHEMAS: + return SCHEMAS[name] + raise ValueError(f"Schema {name} not found in {SCHEMA_FILE}") + +def validate_a2ui_messages(messages: List[Dict[str, Any]]): + for i, msg in enumerate(messages): + try: + if "beginRendering" in msg: + validate(instance=msg, schema=get_schema("beginRendering")) + elif "surfaceUpdate" in msg: + validate(instance=msg, schema=get_schema("surfaceUpdate")) + elif "dataModelUpdate" in msg: + validate(instance=msg, schema=get_schema("dataModelUpdate")) + elif "deleteSurface" in msg: + validate(instance=msg, schema=get_schema("deleteSurface")) + else: + raise ValidationError(f"Message {i} has no known A2UI message type key") + except ValidationError as e: + print(f"A2UI Validation Error in message {i} ({list(msg.keys())[0]}): {e.message}") + raise + except ValueError as e: + print(f"Schema loading/lookup error: {e}") + raise diff --git a/samples/agent/adk/restaurant_finder/agent.py b/samples/agent/adk/restaurant_finder/agent.py index 5283cd88..58ac46c1 100644 --- a/samples/agent/adk/restaurant_finder/agent.py +++ b/samples/agent/adk/restaurant_finder/agent.py @@ -15,9 +15,12 @@ import json import logging import os +import time from collections.abc import AsyncIterable from typing import Any +import instrumentation + import jsonschema from google.adk.agents.llm_agent import LlmAgent from google.adk.artifacts import InMemoryArtifactService @@ -27,13 +30,13 @@ from google.adk.sessions import InMemorySessionService from google.genai import types from prompt_builder import ( - A2UI_SCHEMA, - RESTAURANT_UI_EXAMPLES, get_text_prompt, get_ui_prompt, ) from tools import get_restaurants - +import ui_schema +import template_renderer +from pydantic import ValidationError logger = logging.getLogger(__name__) AGENT_INSTRUCTION = """ @@ -53,6 +56,27 @@ """ +class InstrumentedLiteLlm(LiteLlm): + async def generate_content_async(self, *args, **kwargs): + logger.info("InstrumentedLiteLlm.generate_content_async called") + start_time = time.time() + try: + async for chunk in super().generate_content_async(*args, **kwargs): + yield chunk + finally: + duration = (time.time() - start_time) * 1000 + instrumentation.track_inference(duration) + + def generate_content(self, *args, **kwargs): + logger.info("InstrumentedLiteLlm.generate_content (sync) called") + start_time = time.time() + try: + return super().generate_content(*args, **kwargs) + finally: + duration = (time.time() - start_time) * 1000 + instrumentation.track_inference(duration) + + class RestaurantAgent: """An agent that finds restaurants based on user criteria.""" @@ -71,22 +95,7 @@ def __init__(self, base_url: str, use_ui: bool = False): memory_service=InMemoryMemoryService(), ) - # --- MODIFICATION: Wrap the schema --- - # Load the A2UI_SCHEMA string into a Python object for validation - try: - # First, load the schema for a *single message* - single_message_schema = json.loads(A2UI_SCHEMA) - # The prompt instructs the LLM to return a *list* of messages. - # Therefore, our validation schema must be an *array* of the single message schema. - self.a2ui_schema_object = {"type": "array", "items": single_message_schema} - logger.info( - "A2UI_SCHEMA successfully loaded and wrapped in an array validator." - ) - except json.JSONDecodeError as e: - logger.error(f"CRITICAL: Failed to parse A2UI_SCHEMA: {e}") - self.a2ui_schema_object = None - # --- END MODIFICATION --- def get_processing_message(self) -> str: return "Finding restaurants that match your criteria..." @@ -96,15 +105,12 @@ def _build_agent(self, use_ui: bool) -> LlmAgent: LITELLM_MODEL = os.getenv("LITELLM_MODEL", "gemini/gemini-2.5-flash") if use_ui: - # Construct the full prompt with UI instructions, examples, and schema - instruction = AGENT_INSTRUCTION + get_ui_prompt( - self.base_url, RESTAURANT_UI_EXAMPLES - ) + instruction = AGENT_INSTRUCTION + get_ui_prompt() else: instruction = get_text_prompt() return LlmAgent( - model=LiteLlm(model=LITELLM_MODEL), + model=InstrumentedLiteLlm(model=LITELLM_MODEL), name="restaurant_agent", description="An agent that finds restaurants and helps book tables.", instruction=instruction, @@ -129,26 +135,11 @@ async def stream(self, query, session_id) -> AsyncIterable[dict[str, Any]]: elif "base_url" not in session.state: session.state["base_url"] = self.base_url - # --- Begin: UI Validation and Retry Logic --- + # --- Begin: NEW UI Processing Logic --- max_retries = 1 # Total 2 attempts attempt = 0 current_query_text = query - # Ensure schema was loaded - if self.use_ui and self.a2ui_schema_object is None: - logger.error( - "--- RestaurantAgent.stream: A2UI_SCHEMA is not loaded. " - "Cannot perform UI validation. ---" - ) - yield { - "is_task_complete": True, - "content": ( - "I'm sorry, I'm facing an internal configuration error with my UI components. " - "Please contact support." - ), - } - return - while attempt <= max_retries: attempt += 1 logger.info( @@ -166,7 +157,6 @@ async def stream(self, query, session_id) -> AsyncIterable[dict[str, Any]]: session_id=session.id, new_message=current_message, ): - logger.info(f"Event from runner: {event}") if event.is_final_response(): if ( event.content @@ -176,10 +166,8 @@ async def stream(self, query, session_id) -> AsyncIterable[dict[str, Any]]: final_response_content = "\n".join( [p.text for p in event.content.parts if p.text] ) - break # Got the final response, stop consuming events + break else: - logger.info(f"Intermediate event: {event}") - # Yield intermediate updates on every attempt yield { "is_task_complete": False, "updates": self.get_processing_message(), @@ -187,117 +175,68 @@ async def stream(self, query, session_id) -> AsyncIterable[dict[str, Any]]: if final_response_content is None: logger.warning( - f"--- RestaurantAgent.stream: Received no final response content from runner " - f"(Attempt {attempt}). ---" + f"--- RestaurantAgent.stream: No final response (Attempt {attempt}). ---" ) if attempt <= max_retries: - current_query_text = ( - "I received no response. Please try again." - f"Please retry the original request: '{query}'" - ) - continue # Go to next retry + current_query_text = f"No response. Please retry: '{query}'" + continue else: - # Retries exhausted on no-response - final_response_content = "I'm sorry, I encountered an error and couldn't process your request." - # Fall through to send this as a text-only error + final_response_content = "I'm sorry, I encountered an error." - is_valid = False + a2ui_messages = None error_message = "" + text_part = final_response_content if self.use_ui: - logger.info( - f"--- RestaurantAgent.stream: Validating UI response (Attempt {attempt})... ---" - ) + logger.info(f"--- Validating UI response (Attempt {attempt})... ---") try: - if "---a2ui_JSON---" not in final_response_content: - raise ValueError("Delimiter '---a2ui_JSON---' not found.") + if "```a2ui" not in final_response_content: + raise ValueError("A2UI block not found.") - text_part, json_string = final_response_content.split( - "---a2ui_JSON---", 1 - ) + parts = final_response_content.split("```a2ui", 1) + text_part = parts[0] + json_string = parts[1].split("```", 1)[0] if not json_string.strip(): - raise ValueError("JSON part is empty.") - - json_string_cleaned = ( - json_string.strip().lstrip("```json").rstrip("```").strip() - ) - - if not json_string_cleaned: - raise ValueError("Cleaned JSON string is empty.") - - # --- New Validation Steps --- - # 1. Check if it's parsable JSON - parsed_json_data = json.loads(json_string_cleaned) - - # 2. Check if it validates against the A2UI_SCHEMA - # This will raise jsonschema.exceptions.ValidationError if it fails - logger.info( - "--- RestaurantAgent.stream: Validating against A2UI_SCHEMA... ---" - ) - jsonschema.validate( - instance=parsed_json_data, schema=self.a2ui_schema_object - ) - # --- End New Validation Steps --- - - logger.info( - f"--- RestaurantAgent.stream: UI JSON successfully parsed AND validated against schema. " - f"Validation OK (Attempt {attempt}). ---" - ) - is_valid = True - - except ( - ValueError, - json.JSONDecodeError, - jsonschema.exceptions.ValidationError, - ) as e: - logger.warning( - f"--- RestaurantAgent.stream: A2UI validation failed: {e} (Attempt {attempt}) ---" - ) - logger.warning( - f"--- Failed response content: {final_response_content[:500]}... ---" - ) - error_message = f"Validation failed: {e}." - - else: # Not using UI, so text is always "valid" - is_valid = True - - if is_valid: - logger.info( - f"--- RestaurantAgent.stream: Response is valid. Sending final response (Attempt {attempt}). ---" + raise ValueError("A2UI JSON part is empty.") + + parsed_llm_output = json.loads(json_string) + llm_output = ui_schema.LLMOutput(**parsed_llm_output) + + a2ui_messages = template_renderer.render_ui(llm_output, self.base_url) + logger.info(f"--- UI content generated successfully (Attempt {attempt}). ---") + + except (ValueError, json.JSONDecodeError, ValidationError) as e: + logger.warning(f"--- A2UI output processing failed: {e} (Attempt {attempt}) ---") + logger.warning(f"--- Failed content: {final_response_content[:500]}... ---") + error_message = f"Output format error: {e}." + a2ui_messages = None + + if a2ui_messages is not None: + # Combine text part and A2UI messages for the final response + a2ui_json_string = json.dumps(a2ui_messages) + final_output = f"{text_part.strip()}\n---a2ui_JSON---{a2ui_json_string}" + logger.info(f"--- Sending final response with UI (Attempt {attempt}). ---") + yield {"is_task_complete": True, "content": final_output} + return + elif not self.use_ui: + logger.info(f"--- Sending text only response (Attempt {attempt}). ---") + yield {"is_task_complete": True, "content": text_part} + return + + # --- If we're here, UI generation failed --- + if attempt <= max_retries: + logger.warning(f"--- Retrying... ({attempt}/{max_retries + 1}) ---") + current_query_text = ( + f"Your previous response had an issue: {error_message} " + "You MUST produce a JSON block between ```a2ui and ``` " + f"that conforms to the LLMOutput schema. Retry for original query: '{query}'" ) - logger.info(f"Final response: {final_response_content}") + else: + logger.error("--- Max retries exhausted. Sending text-only error. ---") yield { "is_task_complete": True, - "content": final_response_content, + "content": text_part + "\n\nI'm having trouble generating the interface right now.", } - return # We're done, exit the generator - - # --- If we're here, it means validation failed --- - - if attempt <= max_retries: - logger.warning( - f"--- RestaurantAgent.stream: Retrying... ({attempt}/{max_retries + 1}) ---" - ) - # Prepare the query for the retry - current_query_text = ( - f"Your previous response was invalid. {error_message} " - "You MUST generate a valid response that strictly follows the A2UI JSON SCHEMA. " - "The response MUST be a JSON list of A2UI messages. " - "Ensure the response is split by '---a2ui_JSON---' and the JSON part is well-formed. " - f"Please retry the original request: '{query}'" - ) - # Loop continues... - - # --- If we're here, it means we've exhausted retries --- - logger.error( - "--- RestaurantAgent.stream: Max retries exhausted. Sending text-only error. ---" - ) - yield { - "is_task_complete": True, - "content": ( - "I'm sorry, I'm having trouble generating the interface for that request right now. " - "Please try again in a moment." - ), - } - # --- End: UI Validation and Retry Logic --- + return + # --- End: NEW UI Processing Logic --- diff --git a/samples/agent/adk/restaurant_finder/agent_executor.py b/samples/agent/adk/restaurant_finder/agent_executor.py index aeeeb37e..6ce5fb63 100644 --- a/samples/agent/adk/restaurant_finder/agent_executor.py +++ b/samples/agent/adk/restaurant_finder/agent_executor.py @@ -14,6 +14,7 @@ import json import logging +import instrumentation from a2a.server.agent_execution import AgentExecutor, RequestContext from a2a.server.events import EventQueue @@ -51,6 +52,17 @@ async def execute( self, context: RequestContext, event_queue: EventQueue, + ) -> None: + instrumentation.start_request() + try: + await self._execute_internal(context, event_queue) + finally: + instrumentation.end_request() + + async def _execute_internal( + self, + context: RequestContext, + event_queue: EventQueue, ) -> None: query = "" ui_event_part = None diff --git a/samples/agent/adk/restaurant_finder/generate_schema.py b/samples/agent/adk/restaurant_finder/generate_schema.py new file mode 100644 index 00000000..0a4931d2 --- /dev/null +++ b/samples/agent/adk/restaurant_finder/generate_schema.py @@ -0,0 +1,6 @@ +import json +from ui_schema import LLMOutput + +if __name__ == "__main__": + schema = LLMOutput.model_json_schema() + print(json.dumps(schema, indent=2)) diff --git a/samples/agent/adk/restaurant_finder/instrumentation.py b/samples/agent/adk/restaurant_finder/instrumentation.py new file mode 100644 index 00000000..008d07d8 --- /dev/null +++ b/samples/agent/adk/restaurant_finder/instrumentation.py @@ -0,0 +1,58 @@ +import time +import contextvars +from dataclasses import dataclass, field +from typing import List +import logging + +logger = logging.getLogger(__name__) + +@dataclass +class InferenceStat: + duration_ms: float + +@dataclass +class ToolCallStat: + tool_name: str + duration_ms: float + +@dataclass +class RequestStats: + start_time: float = 0.0 + inferences: List[InferenceStat] = field(default_factory=list) + tool_calls: List[ToolCallStat] = field(default_factory=list) + +_request_stats = contextvars.ContextVar("request_stats", default=None) + +def start_request(): + logger.info("instrumentation.start_request called") + _request_stats.set(RequestStats(start_time=time.time())) + +def end_request(): + stats = _request_stats.get() + if stats: + total_duration = (time.time() - stats.start_time) * 1000 + logger.info(f"Total request time: {total_duration:.2f} milliseconds") + logger.info(f"Number of inferences: {len(stats.inferences)}") + for i, inf in enumerate(stats.inferences): + logger.info(f" - Inference {i}: {inf.duration_ms:.2f} milliseconds") + + if stats.tool_calls: + logger.info("Tool calls:") + for tool in stats.tool_calls: + logger.info(f" - {tool.tool_name}: {tool.duration_ms:.2f} milliseconds") + else: + logger.warning("No request stats found for end_request") + +def track_inference(duration_ms: float): + stats = _request_stats.get() + if stats: + stats.inferences.append(InferenceStat(duration_ms=duration_ms)) + else: + logger.warning(f"track_inference: No request stats found! Duration: {duration_ms}ms") + +def track_tool_call(tool_name: str, duration_ms: float): + stats = _request_stats.get() + if stats: + stats.tool_calls.append(ToolCallStat(tool_name=tool_name, duration_ms=duration_ms)) + else: + logger.warning(f"track_tool_call: No request stats found! Tool: {tool_name}") diff --git a/samples/agent/adk/restaurant_finder/prompt_builder.py b/samples/agent/adk/restaurant_finder/prompt_builder.py index 6eb2f7aa..847318b9 100644 --- a/samples/agent/adk/restaurant_finder/prompt_builder.py +++ b/samples/agent/adk/restaurant_finder/prompt_builder.py @@ -12,857 +12,110 @@ # See the License for the specific language governing permissions and # limitations under the License. -# The A2UI schema remains constant for all A2UI responses. -A2UI_SCHEMA = r''' +LLM_OUTPUT_SCHEMA = r''' { - "title": "A2UI Message Schema", - "description": "Describes a JSON payload for an A2UI (Agent to UI) message, which is used to dynamically construct and update user interfaces. A message MUST contain exactly ONE of the action properties: 'beginRendering', 'surfaceUpdate', 'dataModelUpdate', or 'deleteSurface'.", - "type": "object", - "properties": { - "beginRendering": { - "type": "object", - "description": "Signals the client to begin rendering a surface with a root component and specific styles.", + "$defs": { + "Widget": { "properties": { - "surfaceId": { - "type": "string", - "description": "The unique identifier for the UI surface to be rendered." - }, - "root": { - "type": "string", - "description": "The ID of the root component to render." + "type": { + "enum": [ + "restaurant_list", + "booking_form", + "confirmation" + ], + "title": "Type", + "type": "string" }, - "styles": { + "data": { + "additionalProperties": true, + "title": "Data", "type": "object", - "description": "Styling information for the UI.", - "properties": { - "font": { - "type": "string", - "description": "The primary font for the UI." - }, - "primaryColor": { - "type": "string", - "description": "The primary UI color as a hexadecimal code (e.g., '#00BFFF').", - "pattern": "^#[0-9a-fA-F]{6}$" - } - } + "description": "Data for the widget. Schema depends on the widget type. For 'restaurant_list', use RestaurantListData. For 'booking_form', use BookingFormData. For 'confirmation', use ConfirmationData." } }, - "required": ["root", "surfaceId"] - }, - "surfaceUpdate": { - "type": "object", - "description": "Updates a surface with a new set of components.", - "properties": { - "surfaceId": { - "type": "string", - "description": "The unique identifier for the UI surface to be updated. If you are adding a new surface this *must* be a new, unique identified that has never been used for any existing surfaces shown." - }, - "components": { - "type": "array", - "description": "A list containing all UI components for the surface.", - "minItems": 1, - "items": { - "type": "object", - "description": "Represents a *single* component in a UI widget tree. This component could be one of many supported types.", - "properties": { - "id": { - "type": "string", - "description": "The unique identifier for this component." - }, - "weight": { - "type": "number", - "description": "The relative weight of this component within a Row or Column. This corresponds to the CSS 'flex-grow' property. Note: this may ONLY be set when the component is a direct descendant of a Row or Column." - }, - "component": { - "type": "object", - "description": "A wrapper object that MUST contain exactly one key, which is the name of the component type (e.g., 'Heading'). The value is an object containing the properties for that specific component.", - "properties": { - "Text": { - "type": "object", - "properties": { - "text": { - "type": "object", - "description": "The text content to display. This can be a literal string or a reference to a value in the data model ('path', e.g., '/doc/title'). While simple Markdown formatting is supported (i.e. without HTML, images, or links), utilizing dedicated UI components is generally preferred for a richer and more structured presentation.", - "properties": { - "literalString": { - "type": "string" - }, - "path": { - "type": "string" - } - } - }, - "usageHint": { - "type": "string", - "description": "A hint for the base text style. One of:\n- `h1`: Largest heading.\n- `h2`: Second largest heading.\n- `h3`: Third largest heading.\n- `h4`: Fourth largest heading.\n- `h5`: Fifth largest heading.\n- `caption`: Small text for captions.\n- `body`: Standard body text.", - "enum": [ - "h1", - "h2", - "h3", - "h4", - "h5", - "caption", - "body" - ] - } - }, - "required": ["text"] - }, - "Image": { - "type": "object", - "properties": { - "url": { - "type": "object", - "description": "The URL of the image to display. This can be a literal string ('literal') or a reference to a value in the data model ('path', e.g. '/thumbnail/url').", - "properties": { - "literalString": { - "type": "string" - }, - "path": { - "type": "string" - } - } - }, - "fit": { - "type": "string", - "description": "Specifies how the image should be resized to fit its container. This corresponds to the CSS 'object-fit' property.", - "enum": [ - "contain", - "cover", - "fill", - "none", - "scale-down" - ] - }, - "usageHint": { - "type": "string", - "description": "A hint for the image size and style. One of:\n- `icon`: Small square icon.\n- `avatar`: Circular avatar image.\n- `smallFeature`: Small feature image.\n- `mediumFeature`: Medium feature image.\n- `largeFeature`: Large feature image.\n- `header`: Full-width, full bleed, header image.", - "enum": [ - "icon", - "avatar", - "smallFeature", - "mediumFeature", - "largeFeature", - "header" - ] - } - }, - "required": ["url"] - }, - "Icon": { - "type": "object", - "properties": { - "name": { - "type": "object", - "description": "The name of the icon to display. This can be a literal string or a reference to a value in the data model ('path', e.g. '/form/submit').", - "properties": { - "literalString": { - "type": "string", - "enum": [ - "accountCircle", - "add", - "arrowBack", - "arrowForward", - "attachFile", - "calendarToday", - "call", - "camera", - "check", - "close", - "delete", - "download", - "edit", - "event", - "error", - "favorite", - "favoriteOff", - "folder", - "help", - "home", - "info", - "locationOn", - "lock", - "lockOpen", - "mail", - "menu", - "moreVert", - "moreHoriz", - "notificationsOff", - "notifications", - "payment", - "person", - "phone", - "photo", - "print", - "refresh", - "search", - "send", - "settings", - "share", - "shoppingCart", - "star", - "starHalf", - "starOff", - "upload", - "visibility", - "visibilityOff", - "warning" - ] - }, - "path": { - "type": "string" - } - } - } - }, - "required": ["name"] - }, - "Video": { - "type": "object", - "properties": { - "url": { - "type": "object", - "description": "The URL of the video to display. This can be a literal string or a reference to a value in the data model ('path', e.g. '/video/url').", - "properties": { - "literalString": { - "type": "string" - }, - "path": { - "type": "string" - } - } - } - }, - "required": ["url"] - }, - "AudioPlayer": { - "type": "object", - "properties": { - "url": { - "type": "object", - "description": "The URL of the audio to be played. This can be a literal string ('literal') or a reference to a value in the data model ('path', e.g. '/song/url').", - "properties": { - "literalString": { - "type": "string" - }, - "path": { - "type": "string" - } - } - }, - "description": { - "type": "object", - "description": "A description of the audio, such as a title or summary. This can be a literal string or a reference to a value in the data model ('path', e.g. '/song/title').", - "properties": { - "literalString": { - "type": "string" - }, - "path": { - "type": "string" - } - } - } - }, - "required": ["url"] - }, - "Row": { - "type": "object", - "properties": { - "children": { - "type": "object", - "description": "Defines the children. Use 'explicitList' for a fixed set of children, or 'template' to generate children from a data list.", - "properties": { - "explicitList": { - "type": "array", - "items": { - "type": "string" - } - }, - "template": { - "type": "object", - "description": "A template for generating a dynamic list of children from a data model list. `componentId` is the component to use as a template, and `dataBinding` is the path to the map of components in the data model. Values in the map will define the list of children.", - "properties": { - "componentId": { - "type": "string" - }, - "dataBinding": { - "type": "string" - } - }, - "required": ["componentId", "dataBinding"] - } - } - }, - "distribution": { - "type": "string", - "description": "Defines the arrangement of children along the main axis (horizontally). This corresponds to the CSS 'justify-content' property.", - "enum": [ - "center", - "end", - "spaceAround", - "spaceBetween", - "spaceEvenly", - "start" - ] - }, - "alignment": { - "type": "string", - "description": "Defines the alignment of children along the cross axis (vertically). This corresponds to the CSS 'align-items' property.", - "enum": ["start", "center", "end", "stretch"] - } - }, - "required": ["children"] - }, - "Column": { - "type": "object", - "properties": { - "children": { - "type": "object", - "description": "Defines the children. Use 'explicitList' for a fixed set of children, or 'template' to generate children from a data list.", - "properties": { - "explicitList": { - "type": "array", - "items": { - "type": "string" - } - }, - "template": { - "type": "object", - "description": "A template for generating a dynamic list of children from a data model list. `componentId` is the component to use as a template, and `dataBinding` is the path to the map of components in the data model. Values in the map will define the list of children.", - "properties": { - "componentId": { - "type": "string" - }, - "dataBinding": { - "type": "string" - } - }, - "required": ["componentId", "dataBinding"] - } - } - }, - "distribution": { - "type": "string", - "description": "Defines the arrangement of children along the main axis (vertically). This corresponds to the CSS 'justify-content' property.", - "enum": [ - "start", - "center", - "end", - "spaceBetween", - "spaceAround", - "spaceEvenly" - ] - }, - "alignment": { - "type": "string", - "description": "Defines the alignment of children along the cross axis (horizontally). This corresponds to the CSS 'align-items' property.", - "enum": ["center", "end", "start", "stretch"] - } - }, - "required": ["children"] - }, - "List": { - "type": "object", - "properties": { - "children": { - "type": "object", - "description": "Defines the children. Use 'explicitList' for a fixed set of children, or 'template' to generate children from a data list.", - "properties": { - "explicitList": { - "type": "array", - "items": { - "type": "string" - } - }, - "template": { - "type": "object", - "description": "A template for generating a dynamic list of children from a data model list. `componentId` is the component to use as a template, and `dataBinding` is the path to the map of components in the data model. Values in the map will define the list of children.", - "properties": { - "componentId": { - "type": "string" - }, - "dataBinding": { - "type": "string" - } - }, - "required": ["componentId", "dataBinding"] - } - } - }, - "direction": { - "type": "string", - "description": "The direction in which the list items are laid out.", - "enum": ["vertical", "horizontal"] - }, - "alignment": { - "type": "string", - "description": "Defines the alignment of children along the cross axis.", - "enum": ["start", "center", "end", "stretch"] - } - }, - "required": ["children"] - }, - "Card": { - "type": "object", - "properties": { - "child": { - "type": "string", - "description": "The ID of the component to be rendered inside the card." - } - }, - "required": ["child"] - }, - "Tabs": { - "type": "object", - "properties": { - "tabItems": { - "type": "array", - "description": "An array of objects, where each object defines a tab with a title and a child component.", - "items": { - "type": "object", - "properties": { - "title": { - "type": "object", - "description": "The tab title. Defines the value as either a literal value or a path to data model value (e.g. '/options/title').", - "properties": { - "literalString": { - "type": "string" - }, - "path": { - "type": "string" - } - } - }, - "child": { - "type": "string" - } - }, - "required": ["title", "child"] - } - } - }, - "required": ["tabItems"] - }, - "Divider": { - "type": "object", - "properties": { - "axis": { - "type": "string", - "description": "The orientation of the divider.", - "enum": ["horizontal", "vertical"] - } - } - }, - "Modal": { - "type": "object", - "properties": { - "entryPointChild": { - "type": "string", - "description": "The ID of the component that opens the modal when interacted with (e.g., a button)." - }, - "contentChild": { - "type": "string", - "description": "The ID of the component to be displayed inside the modal." - } - }, - "required": ["entryPointChild", "contentChild"] - }, - "Button": { - "type": "object", - "properties": { - "child": { - "type": "string", - "description": "The ID of the component to display in the button, typically a Text component." - }, - "primary": { - "type": "boolean", - "description": "Indicates if this button should be styled as the primary action." - }, - "action": { - "type": "object", - "description": "The client-side action to be dispatched when the button is clicked. It includes the action's name and an optional context payload.", - "properties": { - "name": { - "type": "string" - }, - "context": { - "type": "array", - "items": { - "type": "object", - "properties": { - "key": { - "type": "string" - }, - "value": { - "type": "object", - "description": "Defines the value to be included in the context as either a literal value or a path to a data model value (e.g. '/user/name').", - "properties": { - "path": { - "type": "string" - }, - "literalString": { - "type": "string" - }, - "literalNumber": { - "type": "number" - }, - "literalBoolean": { - "type": "boolean" - } - } - } - }, - "required": ["key", "value"] - } - } - }, - "required": ["name"] - } - }, - "required": ["child", "action"] - }, - "CheckBox": { - "type": "object", - "properties": { - "label": { - "type": "object", - "description": "The text to display next to the checkbox. Defines the value as either a literal value or a path to data model ('path', e.g. '/option/label').", - "properties": { - "literalString": { - "type": "string" - }, - "path": { - "type": "string" - } - } - }, - "value": { - "type": "object", - "description": "The current state of the checkbox (true for checked, false for unchecked). This can be a literal boolean ('literalBoolean') or a reference to a value in the data model ('path', e.g. '/filter/open').", - "properties": { - "literalBoolean": { - "type": "boolean" - }, - "path": { - "type": "string" - } - } - } - }, - "required": ["label", "value"] - }, - "TextField": { - "type": "object", - "properties": { - "label": { - "type": "object", - "description": "The text label for the input field. This can be a literal string or a reference to a value in the data model ('path, e.g. '/user/name').", - "properties": { - "literalString": { - "type": "string" - }, - "path": { - "type": "string" - } - } - }, - "text": { - "type": "object", - "description": "The value of the text field. This can be a literal string or a reference to a value in the data model ('path', e.g. '/user/name').", - "properties": { - "literalString": { - "type": "string" - }, - "path": { - "type": "string" - } - } - }, - "textFieldType": { - "type": "string", - "description": "The type of input field to display.", - "enum": [ - "date", - "longText", - "number", - "shortText", - "obscured" - ] - }, - "validationRegexp": { - "type": "string", - "description": "A regular expression used for client-side validation of the input." - } - }, - "required": ["label"] - }, - "DateTimeInput": { - "type": "object", - "properties": { - "value": { - "type": "object", - "description": "The selected date and/or time value. This can be a literal string ('literalString') or a reference to a value in the data model ('path', e.g. '/user/dob').", - "properties": { - "literalString": { - "type": "string" - }, - "path": { - "type": "string" - } - } - }, - "enableDate": { - "type": "boolean", - "description": "If true, allows the user to select a date." - }, - "enableTime": { - "type": "boolean", - "description": "If true, allows the user to select a time." - }, - "outputFormat": { - "type": "string", - "description": "The desired format for the output string after a date or time is selected." - } - }, - "required": ["value"] - }, - "MultipleChoice": { - "type": "object", - "properties": { - "selections": { - "type": "object", - "description": "The currently selected values for the component. This can be a literal array of strings or a path to an array in the data model('path', e.g. '/hotel/options').", - "properties": { - "literalArray": { - "type": "array", - "items": { - "type": "string" - } - }, - "path": { - "type": "string" - } - } - }, - "options": { - "type": "array", - "description": "An array of available options for the user to choose from.", - "items": { - "type": "object", - "properties": { - "label": { - "type": "object", - "description": "The text to display for this option. This can be a literal string or a reference to a value in the data model (e.g. '/option/label').", - "properties": { - "literalString": { - "type": "string" - }, - "path": { - "type": "string" - } - } - }, - "value": { - "type": "string", - "description": "The value to be associated with this option when selected." - } - }, - "required": ["label", "value"] - } - }, - "maxAllowedSelections": { - "type": "integer", - "description": "The maximum number of options that the user is allowed to select." - } - }, - "required": ["selections", "options"] - }, - "Slider": { - "type": "object", - "properties": { - "value": { - "type": "object", - "description": "The current value of the slider. This can be a literal number ('literalNumber') or a reference to a value in the data model ('path', e.g. '/restaurant/cost').", - "properties": { - "literalNumber": { - "type": "number" - }, - "path": { - "type": "string" - } - } - }, - "minValue": { - "type": "number", - "description": "The minimum value of the slider." - }, - "maxValue": { - "type": "number", - "description": "The maximum value of the slider." - } - }, - "required": ["value"] - } - } - } - }, - "required": ["id", "component"] - } - } - }, - "required": ["surfaceId", "components"] - }, - "dataModelUpdate": { - "type": "object", - "description": "Updates the data model for a surface.", - "properties": { - "surfaceId": { - "type": "string", - "description": "The unique identifier for the UI surface this data model update applies to." - }, - "path": { - "type": "string", - "description": "An optional path to a location within the data model (e.g., '/user/name'). If omitted, or set to '/', the entire data model will be replaced." - }, - "contents": { - "type": "array", - "description": "An array of data entries. Each entry must contain a 'key' and exactly one corresponding typed 'value*' property.", - "items": { - "type": "object", - "description": "A single data entry. Exactly one 'value*' property should be provided alongside the key.", - "properties": { - "key": { - "type": "string", - "description": "The key for this data entry." - }, - "valueString": { - "type": "string" - }, - "valueNumber": { - "type": "number" - }, - "valueBoolean": { - "type": "boolean" - }, - "valueMap": { - "description": "Represents a map as an adjacency list.", - "type": "array", - "items": { - "type": "object", - "description": "One entry in the map. Exactly one 'value*' property should be provided alongside the key.", - "properties": { - "key": { - "type": "string" - }, - "valueString": { - "type": "string" - }, - "valueNumber": { - "type": "number" - }, - "valueBoolean": { - "type": "boolean" - } - }, - "required": ["key"] - } - } - }, - "required": ["key"] - } - } - }, - "required": ["contents", "surfaceId"] - }, - "deleteSurface": { - "type": "object", - "description": "Signals the client to delete the surface identified by 'surfaceId'.", - "properties": { - "surfaceId": { - "type": "string", - "description": "The unique identifier for the UI surface to be deleted." - } + "required": [ + "type", + "data" + ], + "title": "Widget", + "type": "object" + } + }, + "properties": { + "widgets": { + "items": { + "$ref": "#/$defs/Widget" }, - "required": ["surfaceId"] + "title": "Widgets", + "type": "array" } - } + }, + "required": [ + "widgets" + ], + "title": "LLMOutput", + "type": "object" } ''' -from a2ui_examples import RESTAURANT_UI_EXAMPLES - - -def get_ui_prompt(base_url: str, examples: str) -> str: - """ - Constructs the full prompt with UI instructions, rules, examples, and schema. - - Args: - base_url: The base URL for resolving static assets like logos. - examples: A string containing the specific UI examples for the agent's task. - - Returns: - A formatted string to be used as the system prompt for the LLM. - """ - # The f-string substitution for base_url happens here, at runtime. - formatted_examples = examples.format(base_url=base_url) - - return f""" - You are a helpful restaurant finding assistant. Your final output MUST be a a2ui UI JSON response. - - To generate the response, you MUST follow these rules: - 1. Your response MUST be in two parts, separated by the delimiter: `---a2ui_JSON---`. - 2. The first part is your conversational text response. - 3. The second part is a single, raw JSON object which is a list of A2UI messages. - 4. The JSON part MUST validate against the A2UI JSON SCHEMA provided below. - - --- UI TEMPLATE RULES --- - - If the query is for a list of restaurants, use the restaurant data you have already received from the `get_restaurants` tool to populate the `dataModelUpdate.contents` array (e.g., as a `valueMap` for the "items" key). - - If the number of restaurants is 5 or fewer, you MUST use the `SINGLE_COLUMN_LIST_EXAMPLE` template. - - If the number of restaurants is more than 5, you MUST use the `TWO_COLUMN_LIST_EXAMPLE` template. - - If the query is to book a restaurant (e.g., "USER_WANTS_TO_BOOK..."), you MUST use the `BOOKING_FORM_EXAMPLE` template. - - If the query is a booking submission (e.g., "User submitted a booking..."), you MUST use the `CONFIRMATION_EXAMPLE` template. - - {formatted_examples} - - ---BEGIN A2UI JSON SCHEMA--- - {A2UI_SCHEMA} - ---END A2UI JSON SCHEMA--- - """ - +UI_SYSTEM_PROMPT = f""" +You are a helpful assistant for finding and booking restaurants. +Your goal is to assist the user by calling tools and presenting information clearly. + +When you need to display rich UI elements, you MUST format your response as follows: +1. Include any natural language text you want to show the user. +2. Embed a single JSON block enclosed in triple backticks with the label 'a2ui'. +3. This JSON block MUST conform to the LLMOutput schema provided below. +4. You can include more natural language text after the JSON block. + +Example Response Format: + +Here are some options I found: +```a2ui +{{ + "widgets": [ + {{ + "type": "restaurant_list", + "data": {{ + "restaurants": [ + {{"name": "Cascal", "rating": "★★★★☆", "detail": "Pan-Latin cuisine", "imageUrl": "https://example.com/cascal.jpg", "address": "400 Castro St, Mountain View", "infoLink": "https://cascalrestaurant.com"}} + ], + "use_single_column": true + }} + }} + ] +}} +``` +Let me know if you'd like to book a table! + +LLMOutput Schema: +{LLM_OUTPUT_SCHEMA} + +Widget Data Schemas: + +- **restaurant_list**: + `{{"restaurants": List[Restaurant], "use_single_column": bool}}` + `Restaurant`: {{"name": str, "rating": str, "detail": str, "imageUrl": str, "address": str, "infoLink": str}} + +- **booking_form**: + `{{"restaurantName": str, "imageUrl": str, "address": str}}` + +- **confirmation**: + `{{"restaurantName": str, "partySize": str, "reservationTime": str, "dietaryRequirements": str, "imageUrl": str}}` + +TOOL INSTRUCTIONS: +- Use the 'get_restaurants' tool to find restaurants. +- The tool will return the data needed for the 'restaurant_list' widget. + +Keep your text responses concise and helpful. Always structure the UI data within the ```a2ui ... ``` block as specified. +""" + +TEXT_SYSTEM_PROMPT = """ +You are a helpful assistant for finding and booking restaurants. Respond to the user's requests and answer their questions. You do not have the ability to display rich UI, so provide your responses in clear text. +""" + +def get_ui_prompt() -> str: + return UI_SYSTEM_PROMPT def get_text_prompt() -> str: - """ - Constructs the prompt for a text-only agent. - """ - return """ - You are a helpful restaurant finding assistant. Your final output MUST be a text response. - - To generate the response, you MUST follow these rules: - 1. **For finding restaurants:** - a. You MUST call the `get_restaurants` tool. Extract the cuisine, location, and a specific number (`count`) of restaurants from the user's query. - b. After receiving the data, format the restaurant list as a clear, human-readable text response. You MUST preserve any markdown formatting (like for links) that you receive from the tool. - - 2. **For booking a table (when you receive a query like 'USER_WANTS_TO_BOOK...'):** - a. Respond by asking the user for the necessary details to make a booking (party size, date, time, dietary requirements). - - 3. **For confirming a booking (when you receive a query like 'User submitted a booking...'):** - a. Respond with a simple text confirmation of the booking details. - """ - - -if __name__ == "__main__": - # Example of how to use the prompt builder - # In your actual application, you would call this from your main agent logic. - my_base_url = "http://localhost:8000" - - # You can now easily construct a prompt with the relevant examples. - # For a different agent (e.g., a flight booker), you would pass in - # different examples but use the same `get_ui_prompt` function. - restaurant_prompt = get_ui_prompt(my_base_url, RESTAURANT_UI_EXAMPLES) - - print(restaurant_prompt) - - # This demonstrates how you could save the prompt to a file for inspection - with open("generated_prompt.txt", "w") as f: - f.write(restaurant_prompt) - print("\nGenerated prompt saved to generated_prompt.txt") + return TEXT_SYSTEM_PROMPT diff --git a/samples/agent/adk/restaurant_finder/template_renderer.py b/samples/agent/adk/restaurant_finder/template_renderer.py new file mode 100644 index 00000000..c11dd8e8 --- /dev/null +++ b/samples/agent/adk/restaurant_finder/template_renderer.py @@ -0,0 +1,181 @@ +from typing import List, Dict, Any +from ui_schema import LLMOutput, RestaurantListData, BookingFormData, ConfirmationData, Widget +from a2ui_validator import validate_a2ui_messages + +def render_ui(llm_output: LLMOutput, base_url: str) -> List[Dict[str, Any]]: + messages = [] + for widget in llm_output.widgets: + if widget.type == "restaurant_list": + messages.extend(render_restaurant_list(RestaurantListData(**widget.data), base_url)) + elif widget.type == "booking_form": + messages.extend(render_booking_form(BookingFormData(**widget.data), base_url)) + elif widget.type == "confirmation": + messages.extend(render_confirmation(ConfirmationData(**widget.data), base_url)) + else: + # Log warning or raise error for unknown type + pass + # TODO: Add validation step here + # validate_a2ui_messages(messages) + return messages + + +def _create_data_model_items(restaurants: List[Restaurant]) -> List[Dict[str, Any]]: + items = [] + for i, r in enumerate(restaurants): + items.append({ + "key": f"item{i + 1}", + "valueMap": [ + {"key": "name", "valueString": r.name}, + {"key": "rating", "valueString": r.rating}, + {"key": "detail", "valueString": r.detail}, + {"key": "infoLink", "valueString": r.infoLink}, + {"key": "imageUrl", "valueString": r.imageUrl}, + {"key": "address", "valueString": r.address}, + ] + }) + return items + +def render_restaurant_list(data: RestaurantListData, base_url: str) -> List[Dict[str, Any]]: + surface_id = "restaurant_list" + begin_rendering = { + "beginRendering": { + "surfaceId": surface_id, + "root": "root-column", + "styles": {"primaryColor": "#007BFF", "font": "Roboto"} + } + } + + components = [ + {"id": "root-column", "component": {"Column": {"children": {"explicitList": ["title-heading", "item-list"]}}}}, + {"id": "title-heading", "component": {"Text": {"usageHint": "h1", "text": {"path": "/title"}}}}, + {"id": "item-list", "component": {"List": {"direction": "vertical", "children": {"template": {"componentId": "item-card-template", "dataBinding": "/items"}}}}}, + {"id": "item-card-template", "component": {"Card": {"child": "card-layout"}}}, + {"id": "card-layout", "component": {"Row": {"children": {"explicitList": ["template-image", "card-details"]}}}}, + {"id": "template-image", "weight": 1, "component": {"Image": {"url": {"path": "imageUrl"}, "fit": "cover"}}}, + {"id": "card-details", "weight": 2, "component": {"Column": {"children": {"explicitList": ["template-name", "template-rating", "template-detail", "template-link", "template-book-button"]}}}}, + {"id": "template-name", "component": {"Text": {"usageHint": "h3", "text": {"path": "name"}}}}, + {"id": "template-rating", "component": {"Text": {"text": {"path": "rating"}}}}, + {"id": "template-detail", "component": {"Text": {"text": {"path": "detail"}}}}, + {"id": "template-link", "component": {"Text": {"text": {"path": "infoLink"}}}}, + {"id": "template-book-button", "component": {"Button": {"child": "book-now-text", "primary": True, "action": {"name": "book_restaurant", "context": [ + {"key": "restaurantName", "value": {"path": "name"}}, + {"key": "imageUrl", "value": {"path": "imageUrl"}}, + {"key": "address", "value": {"path": "address"}} + ]}}}}, + {"id": "book-now-text", "component": {"Text": {"text": {"literalString": "Book Now"}}}} + ] + + surface_update = { + "surfaceUpdate": { + "surfaceId": surface_id, + "components": components + } + } + + data_model_update = { + "dataModelUpdate": { + "surfaceId": surface_id, + "path": "/", + "contents": [ + {"key": "title", "valueString": "Found Restaurants"}, + {"key": "items", "valueMap": _create_data_model_items(data.restaurants)} + ] + } + } + + return [begin_rendering, surface_update, data_model_update] +def render_booking_form(data: BookingFormData, base_url: str) -> List[Dict[str, Any]]: + surface_id = "booking_form" + begin_rendering = { + "beginRendering": { + "surfaceId": surface_id, + "root": "booking-form-column", + "styles": {"primaryColor": "#007BFF", "font": "Roboto"} + } + } + + components = [ + {"id": "booking-form-column", "component": {"Column": {"children": {"explicitList": ["booking-title", "restaurant-image", "restaurant-address", "party-size-field", "datetime-field", "dietary-field", "submit-button"]}}}}, + {"id": "booking-title", "component": {"Text": {"usageHint": "h2", "text": {"path": "/title"}}}}, + {"id": "restaurant-image", "component": {"Image": {"url": {"path": "/imageUrl"}}}}, + {"id": "restaurant-address", "component": {"Text": {"text": {"path": "/address"}}}}, + {"id": "party-size-field", "component": {"TextField": {"label": {"literalString": "Party Size"}, "text": {"path": "/partySize"}, "textFieldType": "number"}}}, + {"id": "datetime-field", "component": {"DateTimeInput": {"value": {"path": "/reservationTime"}, "enableDate": True, "enableTime": True}}}, + {"id": "dietary-field", "component": {"TextField": {"label": {"literalString": "Dietary Requirements"}, "text": {"path": "/dietary"}, "textFieldType": "longText"}}}, + {"id": "submit-button", "component": {"Button": {"child": "submit-reservation-text", "primary": True, "action": {"name": "submit_booking", "context": [ + {"key": "restaurantName", "value": {"path": "/restaurantName"}}, + {"key": "partySize", "value": {"path": "/partySize"}}, + {"key": "reservationTime", "value": {"path": "/reservationTime"}}, + {"key": "dietary", "value": {"path": "/dietary"}}, + {"key": "imageUrl", "value": {"path": "/imageUrl"}} + ]}}}}, + {"id": "submit-reservation-text", "component": {"Text": {"text": {"literalString": "Submit Reservation"}}}} + ] + + surface_update = { + "surfaceUpdate": { + "surfaceId": surface_id, + "components": components + } + } + + data_model_update = { + "dataModelUpdate": { + "surfaceId": surface_id, + "path": "/", + "contents": [ + {"key": "title", "valueString": f"Book a Table at {data.restaurantName}"}, + {"key": "address", "valueString": data.address}, + {"key": "restaurantName", "valueString": data.restaurantName}, + {"key": "partySize", "valueString": "2"}, + {"key": "reservationTime", "valueString": ""}, + {"key": "dietary", "valueString": ""}, + {"key": "imageUrl", "valueString": data.imageUrl} + ] + } + } + return [begin_rendering, surface_update, data_model_update] + +def render_confirmation(data: ConfirmationData, base_url: str) -> List[Dict[str, Any]]: + surface_id = "confirmation" + begin_rendering = { + "beginRendering": { + "surfaceId": surface_id, + "root": "confirmation-card", + "styles": {"primaryColor": "#007BFF", "font": "Roboto"} + } + } + + components = [ + {"id": "confirmation-card", "component": {"Card": {"child": "confirmation-column"}}}, + {"id": "confirmation-column", "component": {"Column": {"children": {"explicitList": ["confirm-title", "confirm-image", "divider1", "confirm-details", "divider2", "confirm-dietary", "divider3", "confirm-text"]}}}}, + {"id": "confirm-title", "component": {"Text": {"usageHint": "h2", "text": {"path": "/title"}}}}, + {"id": "confirm-image", "component": {"Image": {"url": {"path": "/imageUrl"}}}}, + {"id": "confirm-details", "component": {"Text": {"text": {"path": "/bookingDetails"}}}}, + {"id": "confirm-dietary", "component": {"Text": {"text": {"path": "/dietaryRequirements"}}}}, + {"id": "confirm-text", "component": {"Text": {"usageHint": "h5", "text": {"literalString": "We look forward to seeing you!"}}}}, + {"id": "divider1", "component": {"Divider": {}}}, + {"id": "divider2", "component": {"Divider": {}}}, + {"id": "divider3", "component": {"Divider": {}}} + ] + + surface_update = { + "surfaceUpdate": { + "surfaceId": surface_id, + "components": components + } + } + + data_model_update = { + "dataModelUpdate": { + "surfaceId": surface_id, + "path": "/", + "contents": [ + {"key": "title", "valueString": f"Booking at {data.restaurantName}"}, + {"key": "bookingDetails", "valueString": f"{data.partySize} people at {data.reservationTime}"}, + {"key": "dietaryRequirements", "valueString": f"Dietary Requirements: {data.dietaryRequirements}"}, + {"key": "imageUrl", "valueString": data.imageUrl} + ] + } + } + return [begin_rendering, surface_update, data_model_update] diff --git a/samples/agent/adk/restaurant_finder/test_copy.py b/samples/agent/adk/restaurant_finder/test_copy.py new file mode 100644 index 00000000..1044996b --- /dev/null +++ b/samples/agent/adk/restaurant_finder/test_copy.py @@ -0,0 +1,16 @@ +from agent import InstrumentedLiteLlm +from google.adk.models.lite_llm import LiteLlm + +def test_copy(): + model = InstrumentedLiteLlm(model="test") + copied = model.copy() + print(f"Original type: {type(model)}") + print(f"Copied type: {type(copied)}") + + if isinstance(copied, InstrumentedLiteLlm): + print("Copy preserves class.") + else: + print("Copy DOES NOT preserve class.") + +if __name__ == "__main__": + test_copy() diff --git a/samples/agent/adk/restaurant_finder/tools.py b/samples/agent/adk/restaurant_finder/tools.py index 6a6dd453..f2a66882 100644 --- a/samples/agent/adk/restaurant_finder/tools.py +++ b/samples/agent/adk/restaurant_finder/tools.py @@ -12,16 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import List, Dict, Any import json import logging import os +import time +import instrumentation from google.adk.tools.tool_context import ToolContext +from ui_schema import Restaurant logger = logging.getLogger(__name__) -def get_restaurants(cuisine: str, location: str, tool_context: ToolContext, count: int = 5) -> str: +def get_restaurants(cuisine: str, location: str, tool_context: ToolContext, count: int = 5) -> List[Dict[str, Any]]: """Call this tool to get a list of restaurants based on a cuisine and location. 'count' is the number of restaurants to return. """ @@ -29,27 +33,38 @@ def get_restaurants(cuisine: str, location: str, tool_context: ToolContext, cou logger.info(f" - Cuisine: {cuisine}") logger.info(f" - Location: {location}") - items = [] - if "new york" in location.lower() or "ny" in location.lower(): - try: - script_dir = os.path.dirname(__file__) - file_path = os.path.join(script_dir, "restaurant_data.json") - with open(file_path) as f: - restaurant_data_str = f.read() - if base_url := tool_context.state.get("base_url"): - restaurant_data_str = restaurant_data_str.replace("http://localhost:10002", base_url) - logger.info(f'Updated base URL from tool context: {base_url}') - all_items = json.loads(restaurant_data_str) + start_time = time.time() + try: + items = [] + if "new york" in location.lower() or "ny" in location.lower(): + try: + script_dir = os.path.dirname(__file__) + file_path = os.path.join(script_dir, "restaurant_data.json") + with open(file_path) as f: + restaurant_data_str = f.read() + if base_url := tool_context.state.get("base_url"): + restaurant_data_str = restaurant_data_str.replace("http://localhost:10002", base_url) + logger.info(f'Updated base URL from tool context: {base_url}') + all_items = json.loads(restaurant_data_str) + + # Slice the list to return only the requested number of items + items = all_items[:count] + logger.info( + f" - Success: Found {len(all_items)} restaurants, returning {len(items)}." + ) + + # Convert to list of dicts matching Restaurant model + validated_items = [Restaurant(**item).model_dump() for item in items] + return validated_items - # Slice the list to return only the requested number of items - items = all_items[:count] - logger.info( - f" - Success: Found {len(all_items)} restaurants, returning {len(items)}." - ) - - except FileNotFoundError: - logger.error(f" - Error: restaurant_data.json not found at {file_path}") - except json.JSONDecodeError: - logger.error(f" - Error: Failed to decode JSON from {file_path}") - - return json.dumps(items) + except FileNotFoundError: + logger.error(f" - Error: restaurant_data.json not found at {file_path}") + except json.JSONDecodeError: + logger.error(f" - Error: Failed to decode JSON from {file_path}") + except Exception as e: + logger.error(f" - Error processing restaurant data: {e}") + + return [] + finally: + duration = (time.time() - start_time) * 1000 + instrumentation.track_tool_call("get_restaurants", duration) diff --git a/samples/agent/adk/restaurant_finder/ui_schema.py b/samples/agent/adk/restaurant_finder/ui_schema.py new file mode 100644 index 00000000..c36654ec --- /dev/null +++ b/samples/agent/adk/restaurant_finder/ui_schema.py @@ -0,0 +1,33 @@ +from typing import List, Literal, Optional, Dict, Any +from pydantic import BaseModel, Field + +class Restaurant(BaseModel): + name: str + rating: str + detail: str + imageUrl: str + address: str + infoLink: str + +class RestaurantListData(BaseModel): + restaurants: List[Restaurant] + use_single_column: bool = True + +class BookingFormData(BaseModel): + restaurantName: str + imageUrl: str + address: str + +class ConfirmationData(BaseModel): + restaurantName: str + partySize: str + reservationTime: str + dietaryRequirements: str + imageUrl: str + +class Widget(BaseModel): + type: Literal["restaurant_list", "booking_form", "confirmation"] + data: Dict[str, Any] + +class LLMOutput(BaseModel): + widgets: List[Widget] diff --git a/samples/agent/adk/restaurant_finder/verify_server.py b/samples/agent/adk/restaurant_finder/verify_server.py new file mode 100644 index 00000000..87eda437 --- /dev/null +++ b/samples/agent/adk/restaurant_finder/verify_server.py @@ -0,0 +1,34 @@ +import subprocess +import time +import os +import signal + +def verify_startup(): + print("Attempting to start server...") + try: + process = subprocess.Popen(["uv", "run", "."], cwd=os.path.dirname(__file__), stdout=subprocess.PIPE, stderr=subprocess.PIPE, start_new_session=True) + time.sleep(5) # Give server time to start + + if process.poll() is not None: + print("Server failed to start or exited early.") + stdout, stderr = process.communicate() + print("STDOUT:", stdout.decode()) + print("STDERR:", stderr.decode()) + return False + else: + print("Server seems to be running. Terminating...") + os.killpg(os.getpgid(process.pid), signal.SIGTERM) + try: + process.communicate(timeout=5) + except subprocess.TimeoutExpired: + os.killpg(os.getpgid(process.pid), signal.SIGKILL) + print("Server force killed.") + print("Server startup check passed.") + return True + except Exception as e: + print(f"Error during server startup check: {e}") + return False + +if __name__ == "__main__": + if not verify_startup(): + exit(1) From ba490c5c4239e86b9815a3ac300c2974a088ac5b Mon Sep 17 00:00:00 2001 From: Jacob Simionato Date: Wed, 17 Dec 2025 09:36:31 +1030 Subject: [PATCH 2/4] Switch to flash lite --- samples/agent/adk/restaurant_finder/agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/samples/agent/adk/restaurant_finder/agent.py b/samples/agent/adk/restaurant_finder/agent.py index 58ac46c1..088513f7 100644 --- a/samples/agent/adk/restaurant_finder/agent.py +++ b/samples/agent/adk/restaurant_finder/agent.py @@ -102,7 +102,7 @@ def get_processing_message(self) -> str: def _build_agent(self, use_ui: bool) -> LlmAgent: """Builds the LLM agent for the restaurant agent.""" - LITELLM_MODEL = os.getenv("LITELLM_MODEL", "gemini/gemini-2.5-flash") + LITELLM_MODEL = os.getenv("LITELLM_MODEL", "gemini/gemini-2.5-flash-lite") if use_ui: instruction = AGENT_INSTRUCTION + get_ui_prompt() From 8b1be5c3ad9e62ae200523dde7385caffbe4c8e6 Mon Sep 17 00:00:00 2001 From: Jacob Simionato Date: Wed, 17 Dec 2025 10:20:23 +1030 Subject: [PATCH 3/4] Bypass LLM for restaurant fetch --- samples/agent/adk/restaurant_finder/agent.py | 1 - .../adk/restaurant_finder/prompt_builder.py | 32 +++++----- .../restaurant_finder/template_renderer.py | 33 ++++++++--- samples/agent/adk/restaurant_finder/tools.py | 59 +++++++++---------- .../agent/adk/restaurant_finder/ui_schema.py | 7 ++- 5 files changed, 75 insertions(+), 57 deletions(-) diff --git a/samples/agent/adk/restaurant_finder/agent.py b/samples/agent/adk/restaurant_finder/agent.py index 088513f7..3e02b782 100644 --- a/samples/agent/adk/restaurant_finder/agent.py +++ b/samples/agent/adk/restaurant_finder/agent.py @@ -114,7 +114,6 @@ def _build_agent(self, use_ui: bool) -> LlmAgent: name="restaurant_agent", description="An agent that finds restaurants and helps book tables.", instruction=instruction, - tools=[get_restaurants], ) async def stream(self, query, session_id) -> AsyncIterable[dict[str, Any]]: diff --git a/samples/agent/adk/restaurant_finder/prompt_builder.py b/samples/agent/adk/restaurant_finder/prompt_builder.py index 847318b9..e4845b61 100644 --- a/samples/agent/adk/restaurant_finder/prompt_builder.py +++ b/samples/agent/adk/restaurant_finder/prompt_builder.py @@ -21,7 +21,8 @@ "enum": [ "restaurant_list", "booking_form", - "confirmation" + "confirmation", + "dynamic_restaurant_list" ], "title": "Type", "type": "string" @@ -30,7 +31,7 @@ "additionalProperties": true, "title": "Data", "type": "object", - "description": "Data for the widget. Schema depends on the widget type. For 'restaurant_list', use RestaurantListData. For 'booking_form', use BookingFormData. For 'confirmation', use ConfirmationData." + "description": "Data for the widget. Schema depends on the widget type." } }, "required": [ @@ -60,7 +61,7 @@ UI_SYSTEM_PROMPT = f""" You are a helpful assistant for finding and booking restaurants. -Your goal is to assist the user by calling tools and presenting information clearly. +Your goal is to assist the user by generating UI specifications. When you need to display rich UI elements, you MUST format your response as follows: 1. Include any natural language text you want to show the user. @@ -68,34 +69,33 @@ 3. This JSON block MUST conform to the LLMOutput schema provided below. 4. You can include more natural language text after the JSON block. -Example Response Format: +To find and display restaurants, use the 'dynamic_restaurant_list' widget type. Provide the search parameters in the 'data' field. -Here are some options I found: +Example for finding restaurants: + +Okay, I'll search for Italian places in New York. ```a2ui {{ "widgets": [ {{ - "type": "restaurant_list", + "type": "dynamic_restaurant_list", "data": {{ - "restaurants": [ - {{"name": "Cascal", "rating": "★★★★☆", "detail": "Pan-Latin cuisine", "imageUrl": "https://example.com/cascal.jpg", "address": "400 Castro St, Mountain View", "infoLink": "https://cascalrestaurant.com"}} - ], - "use_single_column": true + "cuisine": "Italian", + "location": "New York", + "count": 5 }} }} ] }} ``` -Let me know if you'd like to book a table! LLMOutput Schema: {LLM_OUTPUT_SCHEMA} Widget Data Schemas: -- **restaurant_list**: - `{{"restaurants": List[Restaurant], "use_single_column": bool}}` - `Restaurant`: {{"name": str, "rating": str, "detail": str, "imageUrl": str, "address": str, "infoLink": str}} +- **dynamic_restaurant_list**: + `{{"cuisine": str, "location": str, "count": int}}` - **booking_form**: `{{"restaurantName": str, "imageUrl": str, "address": str}}` @@ -103,10 +103,6 @@ - **confirmation**: `{{"restaurantName": str, "partySize": str, "reservationTime": str, "dietaryRequirements": str, "imageUrl": str}}` -TOOL INSTRUCTIONS: -- Use the 'get_restaurants' tool to find restaurants. -- The tool will return the data needed for the 'restaurant_list' widget. - Keep your text responses concise and helpful. Always structure the UI data within the ```a2ui ... ``` block as specified. """ diff --git a/samples/agent/adk/restaurant_finder/template_renderer.py b/samples/agent/adk/restaurant_finder/template_renderer.py index c11dd8e8..6fd79c10 100644 --- a/samples/agent/adk/restaurant_finder/template_renderer.py +++ b/samples/agent/adk/restaurant_finder/template_renderer.py @@ -1,6 +1,14 @@ from typing import List, Dict, Any -from ui_schema import LLMOutput, RestaurantListData, BookingFormData, ConfirmationData, Widget +from ui_schema import LLMOutput, RestaurantListData, BookingFormData, ConfirmationData, Widget, Restaurant, DynamicRestaurantListData from a2ui_validator import validate_a2ui_messages +from tools import fetch_restaurant_data + +_surface_id_counter = 0 + +def _get_unique_surface_id(base_id: str) -> str: + global _surface_id_counter + _surface_id_counter += 1 + return f"{base_id}_{_surface_id_counter}" def render_ui(llm_output: LLMOutput, base_url: str) -> List[Dict[str, Any]]: messages = [] @@ -11,6 +19,8 @@ def render_ui(llm_output: LLMOutput, base_url: str) -> List[Dict[str, Any]]: messages.extend(render_booking_form(BookingFormData(**widget.data), base_url)) elif widget.type == "confirmation": messages.extend(render_confirmation(ConfirmationData(**widget.data), base_url)) + elif widget.type == "dynamic_restaurant_list": + messages.extend(render_dynamic_restaurant_list(DynamicRestaurantListData(**widget.data), base_url)) else: # Log warning or raise error for unknown type pass @@ -35,8 +45,7 @@ def _create_data_model_items(restaurants: List[Restaurant]) -> List[Dict[str, An }) return items -def render_restaurant_list(data: RestaurantListData, base_url: str) -> List[Dict[str, Any]]: - surface_id = "restaurant_list" +def _create_restaurant_list_messages(surface_id: str, restaurants: List[Restaurant]) -> List[Dict[str, Any]]: begin_rendering = { "beginRendering": { "surfaceId": surface_id, @@ -78,14 +87,24 @@ def render_restaurant_list(data: RestaurantListData, base_url: str) -> List[Dict "path": "/", "contents": [ {"key": "title", "valueString": "Found Restaurants"}, - {"key": "items", "valueMap": _create_data_model_items(data.restaurants)} + {"key": "items", "valueMap": _create_data_model_items(restaurants)} ] } } - return [begin_rendering, surface_update, data_model_update] + +def render_restaurant_list(data: RestaurantListData, base_url: str) -> List[Dict[str, Any]]: + surface_id = _get_unique_surface_id("restaurant_list") + return _create_restaurant_list_messages(surface_id, data.restaurants) + +def render_dynamic_restaurant_list(data: DynamicRestaurantListData, base_url: str) -> List[Dict[str, Any]]: + restaurants_data = fetch_restaurant_data(data.cuisine, data.location, data.count, base_url) + restaurants = [Restaurant(**r) for r in restaurants_data] + surface_id = _get_unique_surface_id("restaurant_list") + return _create_restaurant_list_messages(surface_id, restaurants) + def render_booking_form(data: BookingFormData, base_url: str) -> List[Dict[str, Any]]: - surface_id = "booking_form" + surface_id = _get_unique_surface_id("booking_form") begin_rendering = { "beginRendering": { "surfaceId": surface_id, @@ -137,7 +156,7 @@ def render_booking_form(data: BookingFormData, base_url: str) -> List[Dict[str, return [begin_rendering, surface_update, data_model_update] def render_confirmation(data: ConfirmationData, base_url: str) -> List[Dict[str, Any]]: - surface_id = "confirmation" + surface_id = _get_unique_surface_id("confirmation") begin_rendering = { "beginRendering": { "surfaceId": surface_id, diff --git a/samples/agent/adk/restaurant_finder/tools.py b/samples/agent/adk/restaurant_finder/tools.py index f2a66882..f4cdf677 100644 --- a/samples/agent/adk/restaurant_finder/tools.py +++ b/samples/agent/adk/restaurant_finder/tools.py @@ -24,6 +24,33 @@ logger = logging.getLogger(__name__) +def fetch_restaurant_data(cuisine: str, location: str, count: int = 5, base_url: str = "http://localhost:10002") -> List[Dict[str, Any]]: + """Fetches restaurant data from the JSON file.""" + items = [] + if "new york" in location.lower() or "ny" in location.lower(): + try: + script_dir = os.path.dirname(__file__) + file_path = os.path.join(script_dir, "restaurant_data.json") + with open(file_path) as f: + restaurant_data_str = f.read() + restaurant_data_str = restaurant_data_str.replace("http://localhost:10002", base_url) + all_items = json.loads(restaurant_data_str) + + items = all_items[:count] + logger.info( + f" - Success: Found {len(all_items)} restaurants, returning {len(items)}." + ) + validated_items = [Restaurant(**item).model_dump() for item in items] + return validated_items + + except FileNotFoundError: + logger.error(f" - Error: restaurant_data.json not found at {file_path}") + except json.JSONDecodeError: + logger.error(f" - Error: Failed to decode JSON from {file_path}") + except Exception as e: + logger.error(f" - Error processing restaurant data: {e}") + return [] + def get_restaurants(cuisine: str, location: str, tool_context: ToolContext, count: int = 5) -> List[Dict[str, Any]]: """Call this tool to get a list of restaurants based on a cuisine and location. @@ -35,36 +62,8 @@ def get_restaurants(cuisine: str, location: str, tool_context: ToolContext, cou start_time = time.time() try: - items = [] - if "new york" in location.lower() or "ny" in location.lower(): - try: - script_dir = os.path.dirname(__file__) - file_path = os.path.join(script_dir, "restaurant_data.json") - with open(file_path) as f: - restaurant_data_str = f.read() - if base_url := tool_context.state.get("base_url"): - restaurant_data_str = restaurant_data_str.replace("http://localhost:10002", base_url) - logger.info(f'Updated base URL from tool context: {base_url}') - all_items = json.loads(restaurant_data_str) - - # Slice the list to return only the requested number of items - items = all_items[:count] - logger.info( - f" - Success: Found {len(all_items)} restaurants, returning {len(items)}." - ) - - # Convert to list of dicts matching Restaurant model - validated_items = [Restaurant(**item).model_dump() for item in items] - return validated_items - - except FileNotFoundError: - logger.error(f" - Error: restaurant_data.json not found at {file_path}") - except json.JSONDecodeError: - logger.error(f" - Error: Failed to decode JSON from {file_path}") - except Exception as e: - logger.error(f" - Error processing restaurant data: {e}") - - return [] + base_url = tool_context.state.get("base_url", "http://localhost:10002") + return fetch_restaurant_data(cuisine, location, count, base_url) finally: duration = (time.time() - start_time) * 1000 instrumentation.track_tool_call("get_restaurants", duration) diff --git a/samples/agent/adk/restaurant_finder/ui_schema.py b/samples/agent/adk/restaurant_finder/ui_schema.py index c36654ec..a36ff5f8 100644 --- a/samples/agent/adk/restaurant_finder/ui_schema.py +++ b/samples/agent/adk/restaurant_finder/ui_schema.py @@ -26,8 +26,13 @@ class ConfirmationData(BaseModel): imageUrl: str class Widget(BaseModel): - type: Literal["restaurant_list", "booking_form", "confirmation"] + type: Literal["restaurant_list", "booking_form", "confirmation", "dynamic_restaurant_list"] data: Dict[str, Any] class LLMOutput(BaseModel): widgets: List[Widget] + +class DynamicRestaurantListData(BaseModel): + cuisine: str + location: str + count: int = 5 From f575e413835e4aef28286828bfb168548a879382 Mon Sep 17 00:00:00 2001 From: Jacob Simionato Date: Wed, 17 Dec 2025 11:20:28 +1030 Subject: [PATCH 4/4] Add workaround for empty stop chunk in flash lite --- samples/agent/adk/restaurant_finder/agent.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/samples/agent/adk/restaurant_finder/agent.py b/samples/agent/adk/restaurant_finder/agent.py index 3e02b782..77a98af1 100644 --- a/samples/agent/adk/restaurant_finder/agent.py +++ b/samples/agent/adk/restaurant_finder/agent.py @@ -61,8 +61,14 @@ async def generate_content_async(self, *args, **kwargs): logger.info("InstrumentedLiteLlm.generate_content_async called") start_time = time.time() try: - async for chunk in super().generate_content_async(*args, **kwargs): - yield chunk + try: + async for chunk in super().generate_content_async(*args, **kwargs): + yield chunk + except ValueError as e: + if "No message in response" in str(e): + logger.warning(f"Ignored ValueError from LiteLlm (likely empty stop chunk): {e}") + return + raise e finally: duration = (time.time() - start_time) * 1000 instrumentation.track_inference(duration)