# 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 # # https://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. import json import logging import os from collections.abc import AsyncIterable from typing import Any import jsonschema 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.models.lite_llm import LiteLlm from google.adk.runners import Runner 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 logger = logging.getLogger(__name__) AGENT_INSTRUCTION = """ You are a helpful restaurant finding assistant. Your goal is to help users find and book restaurants using a rich UI. To achieve this, you MUST follow this logic: 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 (e.g., for "top 5 chinese places", count is 5). b. After receiving the data, you MUST follow the instructions precisely to generate the final a2ui UI JSON, using the appropriate UI example from the `prompt_builder.py` based on the number of restaurants. 2. **For booking a table (when you receive a query like 'USER_WANTS_TO_BOOK...'):** a. You MUST use the appropriate UI example from `prompt_builder.py` to generate the UI, populating the `dataModelUpdate.contents` with the details from the user's query. 3. **For confirming a booking (when you receive a query like 'User submitted a booking...'):** a. You MUST use the appropriate UI example from `prompt_builder.py` to generate the confirmation UI, populating the `dataModelUpdate.contents` with the final booking details. """ class RestaurantAgent: """An agent that finds restaurants based on user criteria.""" SUPPORTED_CONTENT_TYPES = ["text", "text/plain"] def __init__(self, base_url: str, use_ui: bool = False): self.base_url = base_url self.use_ui = use_ui self._agent = self._build_agent(use_ui) self._user_id = "remote_agent" self._runner = Runner( app_name=self._agent.name, agent=self._agent, artifact_service=InMemoryArtifactService(), session_service=InMemorySessionService(), 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..." 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") 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 ) else: instruction = get_text_prompt() return LlmAgent( model=LiteLlm(model=LITELLM_MODEL), 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]]: session_state = {"base_url": self.base_url} session = await self._runner.session_service.get_session( app_name=self._agent.name, user_id=self._user_id, session_id=session_id, ) if session is None: session = await self._runner.session_service.create_session( app_name=self._agent.name, user_id=self._user_id, state=session_state, session_id=session_id, ) elif "base_url" not in session.state: session.state["base_url"] = self.base_url # --- Begin: UI Validation and Retry 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( f"--- RestaurantAgent.stream: Attempt {attempt}/{max_retries + 1} " f"for session {session_id} ---" ) current_message = types.Content( role="user", parts=[types.Part.from_text(text=current_query_text)] ) final_response_content = None async for event in self._runner.run_async( user_id=self._user_id, session_id=session.id, new_message=current_message, ): logger.info(f"Event from runner: {event}") if event.is_final_response(): if ( event.content and event.content.parts and event.content.parts[0].text ): final_response_content = "\n".join( [p.text for p in event.content.parts if p.text] ) break # Got the final response, stop consuming events else: logger.info(f"Intermediate event: {event}") # Yield intermediate updates on every attempt yield { "is_task_complete": False, "updates": self.get_processing_message(), } if final_response_content is None: logger.warning( f"--- RestaurantAgent.stream: Received no final response content from runner " f"(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 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 is_valid = False error_message = "" if self.use_ui: logger.info( f"--- RestaurantAgent.stream: Validating UI response (Attempt {attempt})... ---" ) try: if "---a2ui_JSON---" not in final_response_content: raise ValueError("Delimiter '---a2ui_JSON---' not found.") text_part, json_string = final_response_content.split( "---a2ui_JSON---", 1 ) 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}). ---" ) logger.info(f"Final response: {final_response_content}") yield { "is_task_complete": True, "content": final_response_content, } 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 ---