wehub-resource-sync
304 lines
12 KiB
Python
304 lines
12 KiB
Python
# 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 ---
|