From 4c10906e1de5b0b58447eb71d9988b1ddaa7f78e Mon Sep 17 00:00:00 2001 From: Lansana Camara Date: Wed, 7 Oct 2026 08:45:42 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 995135302 --- .github/workflows/android-ci.yml | 48 + .github/workflows/cleanup-stale-prs.yml | 84 + .github/workflows/ios-ci.yml | 45 + .github/workflows/python-ci.yml | 3 + .github/workflows/release.yml | 19 +- .github/workflows/web-ci.yml | 1 + .github/workflows/zizmor.yml | 17 + .releaserc.json | 29 +- README.md | 70 +- agent/python_agent/README.md | 62 +- agent/python_agent/__init__.py | 20 +- agent/python_agent/after_tools_callback.py | 100 + agent/python_agent/agent.py | 101 +- agent/python_agent/agent_config.py | 6 +- agent/python_agent/agent_with_grounding.py | 66 +- agent/python_agent/agent_with_templates.py | 524 +- agent/python_agent/merger.py | 365 +- agent/python_agent/place_id_resolution.py | 229 + .../instructions/shared_style_guidelines.md | 55 +- .../shared/schema/maps_catalog_extension.json | 72 +- .../SKILL.md | 38 +- .../local-search-template-response/SKILL.md | 71 - .../python_agent/streaming_request_handler.py | 145 + agent/python_agent/template_registry.py | 473 + agent/python_agent/template_tool.py | 410 + agent/python_agent/templates/__init__.py | 15 + agent/python_agent/templates/common.py | 88 + .../directions}/SKILL.md | 105 +- .../templates/directions/__init__.py | 15 + .../layout.json} | 14 +- .../templates/directions/manifest.json | 29 + .../directions/schema.py} | 124 +- .../directions/test_schema.py} | 173 +- .../directions/testdata/merged.golden.json | 70 + .../templates/local_search/SKILL.md | 106 + .../templates/local_search/__init__.py | 15 + .../layout.json} | 12 +- .../templates/local_search/manifest.json | 31 + .../templates/local_search/schema.py | 148 + .../templates/local_search/test_schema.py | 201 + .../local_search/testdata/merged.golden.json | 103 + agent/python_agent/templates/test_common.py | 123 + .../{text_only.json => text_only/layout.json} | 2 +- .../python_agent/test_after_tools_callback.py | 307 + agent/python_agent/test_agent.py | 44 + .../python_agent/test_agent_with_templates.py | 419 +- agent/python_agent/test_merger.py | 507 +- .../python_agent/test_place_id_resolution.py | 383 + agent/python_agent/test_router_config.py | 42 - .../test_streaming_request_handler.py | 45 + agent/python_agent/test_template_registry.py | 562 ++ agent/python_agent/test_template_tool.py | 490 + .../test_templates_conformance.py | 244 + .../router_instruction.golden.txt} | 71 +- client/android/GoogleMapsA2UI/build.gradle | 27 + .../GoogleMapsA2UI/src/main/assets/index.html | 8820 ++++++++++++++--- .../mapsplatform/a2ui/A2AResponseParser.kt | 177 +- .../libraries/mapsplatform/a2ui/A2UIView.kt | 4 +- .../mapsplatform/a2ui/WebAppInterface.kt | 18 +- .../a2ui/A2AResponseParserTest.kt | 532 +- .../a2ui/A2UIAttributionIdTests.kt | 12 +- .../mapsplatform/a2ui/A2UIServicesTest.kt | 74 + .../mapsplatform/a2ui/A2UIViewTest.kt | 320 + .../mapsplatform/a2ui/WebAppInterfaceTest.kt | 174 + client/android/README.md | 3 + client/android/web_build/index.html | 7 +- client/android/web_build/src/core-shell.ts | 2 +- client/android/web_build/src/main.ts | 131 +- client/ios/GoogleMapsA2UI/Package.swift | 4 +- .../GoogleMapsA2UI/A2AResponseParser.swift | 36 +- .../Sources/GoogleMapsA2UI/A2UIView.swift | 36 +- .../GoogleMapsA2UI/Resources/index.html | 8817 +++++++++++++--- .../Tests/A2AResponseParserTests.swift | 276 +- .../Tests/A2UIAttributionIdTests.swift | 12 +- .../Tests/A2UIServicesTests.swift | 74 + .../GoogleMapsA2UI/Tests/A2UIViewTests.swift | 522 + client/ios/README.md | 3 + client/ios/web_build/index.html | 7 +- client/ios/web_build/src/core-shell.ts | 2 +- client/ios/web_build/src/main.ts | 2 +- client/web/README.md | 48 + client/web/package.json | 2 +- client/web/src/lit/a2ui_client.ts | 90 +- client/web/src/lit/a2ui_client_test.ts | 92 + client/web/src/lit/a2ui_renderer.ts | 46 +- client/web/src/lit/catalog.ts | 34 +- .../src/lit/custom-components/3d_marker.ts | 90 + .../lit/custom-components/3d_marker_test.ts | 54 + .../lit/custom-components/anchor_marker.ts | 66 + .../anchor_marker_constants.ts | 51 + .../custom-components/anchor_marker_test.ts | 45 + .../src/lit/custom-components/google_map.ts | 408 +- .../lit/custom-components/google_map_test.ts | 474 +- client/web/src/lit/custom-components/index.ts | 26 +- .../lit/custom-components/maps_api_loader.ts | 89 + .../custom-components/maps_api_loader_test.ts | 56 + .../place_details_compact.ts | 67 +- .../place_details_compact_test.ts | 151 +- .../lit/custom-components/place_pin_marker.ts | 193 + .../place_pin_marker_constants.ts | 306 + .../place_pin_marker_test.ts | 158 + client/web/src/lit/index.ts | 8 +- 102 files changed, 26102 insertions(+), 4085 deletions(-) create mode 100644 .github/workflows/android-ci.yml create mode 100644 .github/workflows/cleanup-stale-prs.yml create mode 100644 .github/workflows/ios-ci.yml create mode 100644 agent/python_agent/after_tools_callback.py create mode 100644 agent/python_agent/place_id_resolution.py delete mode 100644 agent/python_agent/skills/local-search-template-response/SKILL.md create mode 100644 agent/python_agent/streaming_request_handler.py create mode 100644 agent/python_agent/template_registry.py create mode 100644 agent/python_agent/template_tool.py create mode 100644 agent/python_agent/templates/__init__.py create mode 100644 agent/python_agent/templates/common.py rename agent/python_agent/{skills/directions-template-response => templates/directions}/SKILL.md (52%) create mode 100644 agent/python_agent/templates/directions/__init__.py rename agent/python_agent/templates/{directions.json => directions/layout.json} (78%) create mode 100644 agent/python_agent/templates/directions/manifest.json rename agent/python_agent/{extractor.py => templates/directions/schema.py} (50%) rename agent/python_agent/{test_extractor.py => templates/directions/test_schema.py} (51%) create mode 100644 agent/python_agent/templates/directions/testdata/merged.golden.json create mode 100644 agent/python_agent/templates/local_search/SKILL.md create mode 100644 agent/python_agent/templates/local_search/__init__.py rename agent/python_agent/templates/{local_search.json => local_search/layout.json} (80%) create mode 100644 agent/python_agent/templates/local_search/manifest.json create mode 100644 agent/python_agent/templates/local_search/schema.py create mode 100644 agent/python_agent/templates/local_search/test_schema.py create mode 100644 agent/python_agent/templates/local_search/testdata/merged.golden.json create mode 100644 agent/python_agent/templates/test_common.py rename agent/python_agent/templates/{text_only.json => text_only/layout.json} (90%) create mode 100644 agent/python_agent/test_after_tools_callback.py create mode 100644 agent/python_agent/test_agent.py create mode 100644 agent/python_agent/test_place_id_resolution.py delete mode 100644 agent/python_agent/test_router_config.py create mode 100644 agent/python_agent/test_streaming_request_handler.py create mode 100644 agent/python_agent/test_template_registry.py create mode 100644 agent/python_agent/test_template_tool.py create mode 100644 agent/python_agent/test_templates_conformance.py rename agent/python_agent/{router_config.py => testdata/router_instruction.golden.txt} (59%) create mode 100644 client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIServicesTest.kt create mode 100644 client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIViewTest.kt create mode 100644 client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/WebAppInterfaceTest.kt create mode 100644 client/ios/GoogleMapsA2UI/Tests/A2UIServicesTests.swift create mode 100644 client/ios/GoogleMapsA2UI/Tests/A2UIViewTests.swift create mode 100644 client/web/src/lit/a2ui_client_test.ts create mode 100644 client/web/src/lit/custom-components/3d_marker.ts create mode 100644 client/web/src/lit/custom-components/3d_marker_test.ts create mode 100644 client/web/src/lit/custom-components/anchor_marker.ts create mode 100644 client/web/src/lit/custom-components/anchor_marker_constants.ts create mode 100644 client/web/src/lit/custom-components/anchor_marker_test.ts create mode 100644 client/web/src/lit/custom-components/maps_api_loader.ts create mode 100644 client/web/src/lit/custom-components/maps_api_loader_test.ts create mode 100644 client/web/src/lit/custom-components/place_pin_marker.ts create mode 100644 client/web/src/lit/custom-components/place_pin_marker_constants.ts create mode 100644 client/web/src/lit/custom-components/place_pin_marker_test.ts diff --git a/.github/workflows/android-ci.yml b/.github/workflows/android-ci.yml new file mode 100644 index 0000000..2905bd5 --- /dev/null +++ b/.github/workflows/android-ci.yml @@ -0,0 +1,48 @@ +# Copyright 2026 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 +# +# http://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. + +name: Android CI + +on: + pull_request: + branches: [ main ] + workflow_dispatch: + +permissions: + contents: read + +jobs: + build: + name: Android CI / build + # zizmor: ignore[unpinned-images] + runs-on: ubuntu-24.04 + timeout-minutes: 10 + + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + + - name: Set up JDK 17 + uses: actions/setup-java@c5195efecf7bdfc987ee8bae7a71cb8b11521c00 # v4.7.1 + with: + distribution: 'temurin' + java-version: '17' + cache: 'gradle' + + - name: Build and test library + run: | + chmod +x ./gradlew + ./gradlew test assembleRelease --no-daemon + working-directory: client/android/GoogleMapsA2UI diff --git a/.github/workflows/cleanup-stale-prs.yml b/.github/workflows/cleanup-stale-prs.yml new file mode 100644 index 0000000..d238b98 --- /dev/null +++ b/.github/workflows/cleanup-stale-prs.yml @@ -0,0 +1,84 @@ +# Copyright 2026 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 +# +# http://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. + +name: Cleanup Stale Draft PRs + +on: + schedule: + - cron: '0 2 * * *' # Daily at 02:00 UTC + workflow_dispatch: + inputs: + older_than_days: + description: 'Close draft PRs older than N days' + required: false + default: '3' + type: string + dry_run: + description: 'Dry run (simulate without closing PRs or deleting branches)' + required: false + default: false + type: boolean + +permissions: + pull-requests: write + contents: write + +jobs: + cleanup: + name: Cleanup Draft PRs + # zizmor: ignore[unpinned-images] + runs-on: ubuntu-latest + steps: + - name: Close stale draft PRs and delete branches + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GH_REPO: ${{ github.repository }} + INPUT_DAYS: ${{ inputs.older_than_days }} + INPUT_DRY_RUN: ${{ inputs.dry_run }} + run: | + DAYS="${INPUT_DAYS:-3}" + DRY_RUN="${INPUT_DRY_RUN:-false}" + + echo "Searching for open draft PRs with head branch matching 'test_*' older than $DAYS day(s)..." + + CUTOFF_EPOCH=$(date -d "$DAYS days ago" +%s) + echo "Cutoff timestamp: $CUTOFF_EPOCH ($(date -d "@$CUTOFF_EPOCH" --utc --iso-8601=seconds))" + + PRS_JSON=$(gh pr list --repo "$GH_REPO" --state open --draft --json number,headRefName,updatedAt) + + echo "$PRS_JSON" | jq -c '.[]' | while read -r pr; do + PR_NUMBER=$(echo "$pr" | jq -r '.number') + HEAD_REF=$(echo "$pr" | jq -r '.headRefName') + UPDATED_AT=$(echo "$pr" | jq -r '.updatedAt') + + # Only target Copybara presubmit branches (prefix test_) + if [[ ! "$HEAD_REF" =~ ^test_ ]]; then + echo "Skipping PR #$PR_NUMBER (head branch '$HEAD_REF' does not match 'test_*')" + continue + fi + + PR_EPOCH=$(date -d "$UPDATED_AT" +%s) + if [ "$PR_EPOCH" -lt "$CUTOFF_EPOCH" ]; then + echo "PR #$PR_NUMBER ($HEAD_REF, updated at $UPDATED_AT) is older than $DAYS day(s)." + if [ "$DRY_RUN" = "true" ]; then + echo "[DRY RUN] Would close PR #$PR_NUMBER and delete branch '$HEAD_REF'" + else + echo "Closing PR #$PR_NUMBER and deleting branch '$HEAD_REF'..." + gh pr close "$PR_NUMBER" --repo "$GH_REPO" --comment "Automatically closing stale presubmit draft PR and cleaning up branch." --delete-branch || \ + gh pr close "$PR_NUMBER" --repo "$GH_REPO" --comment "Automatically closing stale presubmit draft PR." + fi + else + echo "Keeping PR #$PR_NUMBER ($HEAD_REF, updated at $UPDATED_AT) - active within $DAYS day(s)." + fi + done diff --git a/.github/workflows/ios-ci.yml b/.github/workflows/ios-ci.yml new file mode 100644 index 0000000..6c8352a --- /dev/null +++ b/.github/workflows/ios-ci.yml @@ -0,0 +1,45 @@ +# Copyright 2026 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 +# +# http://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. + +name: iOS CI + +on: + pull_request: + branches: [ main ] + workflow_dispatch: + +permissions: + contents: read + +jobs: + build: + name: iOS CI / build + # Pinned to macos-15 so the bundled Xcode and iOS Simulator lineup stay stable. + # zizmor: ignore[unpinned-images] + runs-on: macos-15 + timeout-minutes: 30 + + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + + - name: Build and test package + run: | + xcodebuild test \ + -scheme GoogleMapsA2UI \ + -destination 'platform=iOS Simulator,name=iPhone 16' \ + -skipPackagePluginValidation \ + CODE_SIGNING_ALLOWED=NO + working-directory: client/ios/GoogleMapsA2UI diff --git a/.github/workflows/python-ci.yml b/.github/workflows/python-ci.yml index 8ee8f3b..4693853 100644 --- a/.github/workflows/python-ci.yml +++ b/.github/workflows/python-ci.yml @@ -22,7 +22,10 @@ on: jobs: build: + name: Python CI / build runs-on: ubuntu-latest + permissions: + contents: read steps: - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 with: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index fc62731..165d705 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -18,6 +18,11 @@ on: workflow_dispatch: + inputs: + dry_run: + description: "Run in dry-run mode (no tags, no publish)" + type: boolean + default: true permissions: contents: write @@ -45,7 +50,7 @@ jobs: - name: Install dependencies working-directory: client/web - run: npm ci + run: npm install - name: Setup Node for Publishing uses: actions/setup-node@b39b52d1213e96004bfcb1c61a8a6fa8ab84f3e8 # v4.0.1 @@ -60,5 +65,15 @@ jobs: NODE_AUTH_TOKEN: ${{ secrets.NPM_WOMBAT_TOKEN }} NPM_TOKEN: ${{ secrets.NPM_WOMBAT_TOKEN }} NODE_PATH: ${{ github.workspace }}/client/web/node_modules - run: npx --prefix client/web semantic-release + DRY_RUN: ${{ inputs.dry_run }} + REF_NAME: ${{ github.ref_name }} + run: | + EXTRA_ARGS="" + if [ "$DRY_RUN" != "false" ]; then + EXTRA_ARGS="--dry-run" + fi + if [ "$REF_NAME" != "main" ]; then + EXTRA_ARGS="$EXTRA_ARGS --branches $REF_NAME" + fi + npx --prefix client/web semantic-release $EXTRA_ARGS diff --git a/.github/workflows/web-ci.yml b/.github/workflows/web-ci.yml index a55bc4d..5272949 100644 --- a/.github/workflows/web-ci.yml +++ b/.github/workflows/web-ci.yml @@ -25,6 +25,7 @@ permissions: jobs: build: + name: Web CI / build # zizmor: ignore[unpinned-images] runs-on: ubuntu-24.04 steps: diff --git a/.github/workflows/zizmor.yml b/.github/workflows/zizmor.yml index 8c4f481..5f48f4e 100644 --- a/.github/workflows/zizmor.yml +++ b/.github/workflows/zizmor.yml @@ -1,3 +1,17 @@ +# Copyright 2026 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 +# +# http://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. + name: Zizmor on: @@ -24,3 +38,6 @@ jobs: - name: Run zizmor uses: zizmorcore/zizmor-action@195d10ad90f31d8cd6ea1efd6ecc12969ddbe73f # v0.5.1 + with: + args: --ignore insufficient-cooldown + diff --git a/.releaserc.json b/.releaserc.json index 089c315..15c8354 100644 --- a/.releaserc.json +++ b/.releaserc.json @@ -3,7 +3,34 @@ "main" ], "plugins": [ - "@semantic-release/commit-analyzer", + [ + "@semantic-release/commit-analyzer", + { + "preset": "angular", + "releaseRules": [ + { + "breaking": true, + "release": "patch" + }, + { + "type": "feat", + "release": "patch" + }, + { + "type": "fix", + "release": "patch" + }, + { + "type": "perf", + "release": "patch" + }, + { + "type": "refactor", + "release": "patch" + } + ] + } + ], "@semantic-release/release-notes-generator", "@semantic-release/changelog", [ diff --git a/README.md b/README.md index a153511..8d0e2f9 100644 --- a/README.md +++ b/README.md @@ -83,6 +83,7 @@ This package provides the core Python agent implementations for the Agentic UI T * `agent.py`: Contains the `MAUIAgent` class, which handles session management, LLM interaction, and A2UI schema loading. * `agent_with_grounding.py`: Contains `MAUIAgentWithGrounding`, extending the base agent with Vertex Grounding capabilities. +* `streaming_request_handler.py`: Contains `StreamingRequestHandler` (extending `DefaultRequestHandler`), enabling the A2A server to handle both incremental SSE streaming (`message/stream`) and non-streaming (`message/send`) JSON-RPC requests. * `shared/`: Contains schema extensions (e.g., `maps_catalog_extension.json`). * `skills/`: Contains specific skill definitions used by the agents. * `pyproject.toml`: Configuration file for the package, using Hatchling as the build backend. @@ -185,6 +186,25 @@ class MAUIAgentExecutor(AgentExecutor): raise ServerError(error=UnsupportedOperationError()) ``` +#### 3. Enabling Real-Time SSE Streaming (`StreamingRequestHandler`) + +To serve both real-time Server-Sent Events (`message/stream`) and standard single-response requests (`message/send`) from your `A2AStarletteApplication`, pass `StreamingRequestHandler` as the `http_handler`: + +```python +from a2a.server.apps import A2AStarletteApplication +from a2a.server.tasks import InMemoryTaskStore +from streaming_request_handler import StreamingRequestHandler + +request_handler = StreamingRequestHandler( + agent_executor=agent_executor, + task_store=InMemoryTaskStore(), +) +server = A2AStarletteApplication( + agent_card=default_agent.agent_card, + http_handler=request_handler, +) +``` + ## Agentic UI Toolkit Client Libraries This repository provides A2UI client libraries across multiple platforms (Web, Android, and iOS) to render interactive A2UI surfaces and communicate with an A2A agent server. @@ -254,6 +274,28 @@ return ( ); ``` +#### 3. Streaming Responses (`sendStream`) +`A2UIClient` also supports real-time incremental Server-Sent Events (SSE) streaming via `client.sendStream(messageText)` (`message/stream`). You can use a `useStreaming` boolean flag to switch between streaming (`client.sendStream`) and non-streaming (`client.send`): + +```tsx +const useStreaming = true; // Set to false to use non-streaming client.send() + +async function handleSend(messageText: string) { + renderer.addUserMessage(messageText); + + if (useStreaming) { + // Stream incremental text and A2UI component updates via SSE (message/stream) + for await (const chunk of client.sendStream(messageText)) { + renderer.processResponse([chunk]); + } + } else { + // Wait for the full response in a single payload (message/send) + const response = await client.send(messageText); + renderer.processResponse(response); + } +} +``` + #### Local Development To make changes to this package and test them in an application: @@ -297,7 +339,6 @@ Agentic UI Toolkit requires an API Key to use Google Maps Platform products. To Your API Key must have the following APIs enabled in the [Google Cloud Console](https://console.cloud.google.com/apis/credentials): -* Geocoding API * Maps JavaScript API * Places UI Kit * Routes API @@ -312,15 +353,28 @@ To support the use of Grounding Lite within the Python ADK backend, this API Key Your API Key must also be included when loading the Google Maps JavaScript API code. See the [Google Maps Platform Documentation](https://developers.google.com/maps/documentation/javascript/load-maps-js-api) for instructions on how to load the API, including configuring the API Key. -Agentic UI Toolkit requires features available in the Alpha channel. You must use `v=alpha` when loading the Maps JavaScript API. Learn more about versions in the [Google Maps Platform Documentation](https://developers.google.com/maps/documentation/javascript/versions). +Agentic UI Toolkit requires features available in the Alpha channel. You must +use `v: "alpha"` when loading the Maps JavaScript API. + +MAUI components are able to dynamically load their required libraries, enabling +both the +[Dynamic Library Import](https://developers.google.com/maps/documentation/javascript/load-maps-js-api#dynamic-library-import) +and +[Script Tag](https://developers.google.com/maps/documentation/javascript/load-maps-js-api#use-legacy-tag) +methods of loading the Maps JavaScript API. + +If you use the script tag, do not load it asynchronously (no `async` attribute +or `loading=async` parameter). MAUI components need `google.maps.importLibrary` +to exist when they first render. An asynchronously loaded script tag only +defines it once the full API has downloaded, even if you list the libraries +below in its `libraries` parameter, so components that render before then will +not appear. -Use of Agentic UI Toolkit requires several [Maps JavaScript API libraries](https://developers.google.com/maps/documentation/javascript/libraries). When loading the Google Maps JavaScript API, you must include the following libraries: +The specific libraries required by MAUI components are: -* maps -* maps3d -* marker -* places -* routes +* `maps3d` +* `places` +* `routes` ### Gemini API Key diff --git a/agent/python_agent/README.md b/agent/python_agent/README.md index a56a9c1..f3b5b1c 100644 --- a/agent/python_agent/README.md +++ b/agent/python_agent/README.md @@ -14,13 +14,19 @@ AI Maps Grounding. `DIRECTIONS`) and structured parameter extraction for low latency. * `agent_with_grounding.py`: Contains `MAUIAgentWithGrounding`, extending the base agent with Vertex AI Grounding capabilities. +* `template_tool.py`: Contains standard ADK `BaseTool` implementations + (`BaseTemplateTool`, `RenderTextOnlyTemplateTool`, and `TemplateToolset`) + for template rendering. `MAUIAgentWithTemplates` builds a `BaseTemplateTool` + directly from each registered bundle. * `agent_config.py`: Contains `AgentConfig` and `FallbackMode` configurations (`TEXT` vs `DYNAMIC`). -* `extractor.py` & `merger.py`: Parameter extraction schemas and template - merging engine. -* `shared/`: Contains layout templates (e.g. `local-search-template.json`, - `directions-template.json`) and schema extensions - (`maps_catalog_extension.json`). +* `template_registry.py` & `merger.py`: Template bundle discovery and the + generic layout merging engine. +* `templates//`: Self-contained template bundles. Each holds + `manifest.json` (routing and render policy), `layout.json` (the A2UI + skeleton), `schema.py` (the extraction schema), `SKILL.md` (extraction + prompt), and a co-located `test_schema.py`. +* `shared/`: Contains schema extensions (`maps_catalog_extension.json`). * `skills/`: Contains specific skill definitions used by the agents. * `pyproject.toml`: Configuration file for the package, using Hatchling as the build backend. @@ -133,6 +139,25 @@ class MAUIAgentExecutor(AgentExecutor): raise ServerError(error=UnsupportedOperationError()) ``` +### 3. Enabling Real-Time SSE Streaming (`StreamingRequestHandler`) + +To serve both real-time Server-Sent Events (`message/stream`) and standard single-response requests (`message/send`) from your `A2AStarletteApplication`, pass `StreamingRequestHandler` as the `http_handler`: + +```python +from a2a.server.apps import A2AStarletteApplication +from a2a.server.tasks import InMemoryTaskStore +from streaming_request_handler import StreamingRequestHandler + +request_handler = StreamingRequestHandler( + agent_executor=agent_executor, + task_store=InMemoryTaskStore(), +) +server = A2AStarletteApplication( + agent_card=default_agent.agent_card, + http_handler=request_handler, +) +``` + ## Google API Keys ### Google Maps API Key @@ -156,15 +181,28 @@ To support the use of Grounding Lite within the Python ADK backend, this API Key Your API Key must also be included when loading the Google Maps JavaScript API code. See the [Google Maps Platform Documentation](https://developers.google.com/maps/documentation/javascript/load-maps-js-api) for instructions on how to load the API, including configuring the API Key. -Agentic UI Toolkit requires features available in the Alpha channel. You must use `v=alpha` when loading the Maps JavaScript API. Learn more about versions in the [Google Maps Platform Documentation](https://developers.google.com/maps/documentation/javascript/versions). +Agentic UI Toolkit requires features available in the Alpha channel. You must +use `v: "alpha"` when loading the Maps JavaScript API. + +MAUI components are able to dynamically load their required libraries, enabling +both the +[Dynamic Library Import](https://developers.google.com/maps/documentation/javascript/load-maps-js-api#dynamic-library-import) +and +[Script Tag](https://developers.google.com/maps/documentation/javascript/load-maps-js-api#use-legacy-tag) +methods of loading the Maps JavaScript API. + +If you use the script tag, do not load it asynchronously (no `async` attribute +or `loading=async` parameter). MAUI components need `google.maps.importLibrary` +to exist when they first render. An asynchronously loaded script tag only +defines it once the full API has downloaded, even if you list the libraries +below in its `libraries` parameter, so components that render before then will +not appear. -Use of Agentic UI Toolkit requires several [Maps JavaScript API libraries](https://developers.google.com/maps/documentation/javascript/libraries). When loading the Google Maps JavaScript API, you must include the following libraries: +The specific libraries required by MAUI components are: -* maps -* maps3d -* marker -* places -* routes +* `maps3d` +* `places` +* `routes` ### Gemini API Key diff --git a/agent/python_agent/__init__.py b/agent/python_agent/__init__.py index 11eecd8..c42d0eb 100644 --- a/agent/python_agent/__init__.py +++ b/agent/python_agent/__init__.py @@ -1 +1,19 @@ -# GMP A2UI Python Agent Package +# Copyright 2026 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. + +from template_tool import ( + BaseTemplateTool, + RenderTextOnlyTemplateTool, + TemplateToolset, +) diff --git a/agent/python_agent/after_tools_callback.py b/agent/python_agent/after_tools_callback.py new file mode 100644 index 0000000..ee46ef7 --- /dev/null +++ b/agent/python_agent/after_tools_callback.py @@ -0,0 +1,100 @@ +# Copyright 2026 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. + +"""After-tool callback for grounding tools in MAUI Agent.""" + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + +# Maximum number of recent content tokens to retain in session state. +# +# Trade-offs / Considerations: +# - Pros of larger values: +# - Retains grounding tokens across longer multi-turn conversations where +# older tool calls returned entities that are still referenced or +# rendered in UI widgets. +# - Prevents premature eviction of valid tokens, ensuring Place Widget +# requests can successfully waive billing even after multiple subsequent +# tool turns. +# - Cons of larger values: +# - Increases session state size and payload memory footprint across +# requests. +# - Increases serialized metadata size attached to message parts and RPCs. +# - Adds backend processing overhead when downstream services must decrypt +# and validate a larger list of candidate tokens. +# - Since tokens have an expiration TTL (e.g. 30 minutes), retaining too +# many historical tokens increases stale/expired tokens in the payload. +MAX_CONTENT_TOKENS: int = 10 + + +def after_tools_callback( + tool: Any, + args: dict[str, Any], + tool_context: Any, + tool_response: Any, + **kwargs: Any, +) -> Any: + """Callback to aggregate grounding_content_token into session state.""" + # pylint: disable=unused-argument + if not tool_response or not isinstance(tool_response, dict): + return None + + after_maps_tools_callback(tool_context, tool_response) + + return None + + +def after_maps_tools_callback( + tool_context: Any, + tool_response: Any, +) -> Any: + """Callback to aggregate content_token from Maps Tools into session state.""" + # pylint: disable=unused-argument + if tool_context is None or getattr(tool_context, "state", None) is None: + return None + + token = tool_response.get("content_token") + if isinstance(token, str) and token: + content_tokens = tool_context.state.get("maps_tools_content_tokens", []) + # If content_tokens is not a list, initialize it to an empty list. + if not isinstance(content_tokens, list): + content_tokens = [] + if token not in content_tokens: + content_tokens.append(token) + # Keep only the last MAX_CONTENT_TOKENS tokens. + if len(content_tokens) > MAX_CONTENT_TOKENS: + content_tokens = content_tokens[-MAX_CONTENT_TOKENS:] + tool_context.state["maps_tools_content_tokens"] = content_tokens + logger.info( + "--- after_maps_tools_callback: Aggregated content token into" + " content_tokens. ---" + ) + + return None + + +def _add_maps_tools_tokens_to_part(part: Any, session: Any) -> None: + """Adds maps_tools_content_tokens from session state to part metadata.""" + if session is None or getattr(session, "state", None) is None: + return + maps_tools_content_tokens = session.state.get("maps_tools_content_tokens") + if maps_tools_content_tokens: + if getattr(part, "root", None) is not None: + if getattr(part.root, "metadata", None) is None: + part.root.metadata = {} + part.root.metadata["maps_tools_content_tokens"] = ( + maps_tools_content_tokens + ) diff --git a/agent/python_agent/agent.py b/agent/python_agent/agent.py index 19182fb..f621736 100644 --- a/agent/python_agent/agent.py +++ b/agent/python_agent/agent.py @@ -14,7 +14,6 @@ """MAUI Agent definition and implementation.""" -from collections import OrderedDict from collections.abc import AsyncIterable import json import logging @@ -51,9 +50,22 @@ from a2ui.schema.catalog import CatalogConfig from a2ui.schema.catalog_provider import A2uiCatalogProvider from a2ui.schema.common_modifiers import remove_strict_validation -from a2ui.schema.constants import A2UI_CLOSE_TAG, A2UI_OPEN_TAG, VERSION_0_9 +from a2ui.parser.constants import ( + MSG_TYPE_CREATE_SURFACE, + MSG_TYPE_DELETE_SURFACE, + MSG_TYPE_UPDATE_COMPONENTS, + MSG_TYPE_UPDATE_DATA_MODEL, +) +from a2ui.schema.constants import ( + A2UI_CLOSE_TAG, + A2UI_OPEN_TAG, + A2UI_SURFACE_ID_KEY, + VERSION_0_9, +) from a2ui.schema.manager import A2uiSchemaManager +from .after_tools_callback import _add_maps_tools_tokens_to_part, after_tools_callback + logger = logging.getLogger(__name__) InMemorySessionService = in_memory_session_service.InMemorySessionService @@ -101,7 +113,6 @@ **Important**: When using the `google-maps-enriched-local-query-response` skill, you MUST respond with EXACTLY ONE ... block. All A2UI message objects (e.g., `createSurface`, `updateComponents`, `updateDataModel`) MUST include a `"version": "v0.9"` property. - When generating a `PlaceCard`, you MUST explicitly set the `"orientation"` property: use `"vertical"` for single results and `"horizontal"` for lists. If you have more than one of these blocks, the UI will not render correctly. """ @@ -143,6 +154,24 @@ def load(self) -> dict[str, Any]: return catalog +def extract_surface_id(data: Any) -> str | None: + """Extracts the surface ID from an A2UI payload dictionary or part.""" + if not isinstance(data, dict): + return None + for key in ( + MSG_TYPE_CREATE_SURFACE, + MSG_TYPE_UPDATE_COMPONENTS, + MSG_TYPE_UPDATE_DATA_MODEL, + MSG_TYPE_DELETE_SURFACE, + ): + target = data.get(key) + if isinstance(target, dict): + surface_id = target.get(A2UI_SURFACE_ID_KEY) + if surface_id: + return str(surface_id) + return None + + class MAUIAgent: """An agent that finds restaurants based on user criteria.""" @@ -159,14 +188,13 @@ def __init__( self._model_name = model_name self._user_id = "remote_agent" self._shared_session_service = InMemorySessionService() + self._after_tool_callback = after_tools_callback self._text_runner: Runner | None = self._build_runner( self._build_llm_agent() ) self._schema_managers: dict[str, A2uiSchemaManager] = {} self._ui_runners: dict[str, Runner] = {} - self._parsers = OrderedDict() - self._max_parsers = 1000 # Max active sessions to keep in memory for version in [VERSION_0_9]: schema_manager = self._build_schema_manager(version) @@ -303,6 +331,7 @@ def _build_llm_agent( ), instruction=instruction, tools=[grounding_lite_mcp, skill_manager_tool], + after_tool_callback=self._after_tool_callback, ) async def stream( @@ -397,32 +426,41 @@ async def token_stream(): full_content_list.append(p.text) yield p.text + parser = None if selected_catalog: logger.info( "--- MAUIAgent.stream: Using A2UI stream parser for catalog %s ---", selected_catalog.catalog_id, ) - if session_id in self._parsers: - self._parsers.move_to_end(session_id) - else: - self._parsers[session_id] = A2uiStreamParser(selected_catalog) - if len(self._parsers) > self._max_parsers: - self._parsers.popitem(last=False) + parser = A2uiStreamParser(selected_catalog) logger.info( "--- MAUIAgent.stream: Streamed part: %s ---", token_stream() ) - async for part in stream_response_to_parts( - self._parsers[session_id], - token_stream(), - ): - logger.info("-- MAUIAgent.stream: Streamed part: %s ---", part) - yield { - "is_task_complete": False, - "parts": [part], - } + session_surface_id = None + # Wrap stream parsing in try/except to prevent A2uiValidatorError from crashing the ASGI app. + # This ensures execution falls through to the deleteSurface/retry loop below. + try: + async for part in stream_response_to_parts( + parser, + token_stream(), + ): + _add_maps_tools_tokens_to_part(part, session) + logger.info("-- MAUIAgent.stream: Streamed part: %s ---", part) + # TODO(b/553539577): Remove this workaround once A2UI fixes the stream parser state issue. + if isinstance(part.root, DataPart): + s_id = extract_surface_id(part.root.data) + if s_id: + session_surface_id = s_id + logger.info("[WORKAROUND] Sniffed surfaceId '%s' from streamed part", session_surface_id) + yield { + "is_task_complete": False, + "parts": [part], + } + except Exception as e: + logger.warning("--- MAUIAgent.stream: Error during stream parsing (will fall through to retry loop): %s ---", e) else: async for token in token_stream(): yield { @@ -528,6 +566,9 @@ async def token_stream(): filtered_parts.append(p) final_parts = filtered_parts + for p in final_parts: + _add_maps_tools_tokens_to_part(p, session) + yield { "is_task_complete": True, "parts": final_parts, @@ -542,6 +583,26 @@ async def token_stream(): attempt, max_retries + 1, ) + + # Extract surfaceId to clear the failed UI card on the client + surface_id = session_surface_id or getattr(parser, "surface_id", None) + + if surface_id: + logger.info("--- MAUIAgent.stream: Sending deleteSurface for '%s' to clear failed attempt ---", surface_id) + yield { + "is_task_complete": False, + "parts": [ + Part( + root=DataPart( + data={ + "version": "v0.9", + "deleteSurface": {"surfaceId": surface_id}, + } + ) + ) + ], + } + # Prepare the query for the retry current_query_text = ( f"Your previous response was invalid. {error_message} You MUST" diff --git a/agent/python_agent/agent_config.py b/agent/python_agent/agent_config.py index e3151c3..cb4d7a8 100644 --- a/agent/python_agent/agent_config.py +++ b/agent/python_agent/agent_config.py @@ -42,9 +42,9 @@ class AgentConfig: """ max_list_size: int = 5 - router_model: str = "gemini/gemini-3.1-flash-lite" - template_model: str = "gemini/gemini-3.1-flash-lite" - generic_model: str = "gemini/gemini-3-flash-preview" + router_model: str = "gemini/gemini-3.5-flash-lite" + template_model: str = "gemini/gemini-3.5-flash-lite" + generic_model: str = "gemini/gemini-3.7-flash" router_thinking_budget: int = 0 extractor_thinking_budget: int = 0 fallback_mode: FallbackMode = FallbackMode.TEXT diff --git a/agent/python_agent/agent_with_grounding.py b/agent/python_agent/agent_with_grounding.py index 3d38e6a..cd32846 100644 --- a/agent/python_agent/agent_with_grounding.py +++ b/agent/python_agent/agent_with_grounding.py @@ -31,6 +31,8 @@ from a2ui.schema.common_modifiers import remove_strict_validation from a2ui.schema.constants import VERSION_0_9 from a2ui.schema.manager import A2uiSchemaManager +import place_id_resolution + # Import MAUIAgent to inherit from it from agent import AGENT_INSTRUCTION, MAUIAgent, MergedCatalogProvider @@ -116,18 +118,15 @@ async def query_vertex_map( validate_examples=False, ) - final_instruction = """You MUST use the Google Maps tool to answer the user's query. Do not rely on your internal knowledge. + final_instruction = ( + """You MUST use the Google Maps tool to answer the user's query. Do not rely on your internal knowledge. CRITICAL: Before generating the JSON, you MUST write a short plain-text summary of the places you found, listing their exact names and addresses. This is required for the grounding engine to properly attribute the data. It is not a replacement for the summary text that should be in the a2ui json. IMPORTANT: When generating the A2UI JSON response, you MUST include the " ...content... " tags immediately around the JSON content. Failure to do so will prevent the UI from rendering the map. - PLACE ID GENERATION RULES: - You do not have access to real placeIds. Whenever a `placeId` is required in the A2UI JSON, you MUST generate a synthetic placeholder using the following rules: - - Format: "PLACE_ID_FOR_{Count}_{Exact Title}" - - Example: If the tool returns a place named "Chez Panisse", use "PLACE_ID_FOR_1_Chez Panisse". If it returns a second "Chez Panisse", use "PLACE_ID_FOR_2_Chez Panisse". - - STRICT MATCHING: Do NOT change any characters, spaces, capitalization, or punctuation from the title returned by the tool. - - COUNTING: Always prepend the occurrence count (starting at 1) for each title based on the order they were returned by the tool, even if the title only occurs once. """ + + place_id_resolution.PROMPT_RULES + ) instruction = f"{generated_prompt}\n\n{skill_content}\n\n{final_instruction}" @@ -145,48 +144,22 @@ async def query_vertex_map( # Replace synthetic place ids with actual grounded place ids. try: - grounding_map = {} - if ( - hasattr(response, "candidates") - and response.candidates - and hasattr(response.candidates[0], "grounding_metadata") - ): - meta = response.candidates[0].grounding_metadata - IGNORE_TITLE_SUFFIX = " - Google Maps" - IGNORE_PLACE_ID_PREFIX = "places/ChI" - if hasattr(meta, "grounding_chunks") and meta.grounding_chunks: - title_counts = {} - for chunk in meta.grounding_chunks: - if hasattr(chunk, "maps") and chunk.maps: - title = getattr(chunk.maps, "title", None) - place_id = getattr(chunk.maps, "place_id", None) - if title and place_id: - if place_id.startswith(IGNORE_PLACE_ID_PREFIX): - place_id = place_id[len(IGNORE_PLACE_ID_PREFIX) - 3:] - if title.endswith(IGNORE_TITLE_SUFFIX): - title = title[:-len(IGNORE_TITLE_SUFFIX)] - - # Track how many times this title has appeared - title_counts[title] = title_counts.get(title, 0) + 1 - count = title_counts[title] - grounding_map[f"PLACE_ID_FOR_{count}_{title}"] = place_id - else: - logger.warning("No grounding chunks found") - else: - logger.warning("No grounding metadata found") - - if grounding_map: - for key, value in grounding_map.items(): - final_response_content = final_response_content.replace(key, value) - else: - logger.warning("No grounding map found") - + attribution_sources = place_id_resolution.extract_attribution_sources( + response + ) + final_response_content, unresolved_placeholders = ( + place_id_resolution.resolve_place_ids( + final_response_content, attribution_sources + ) + ) + if unresolved_placeholders: + logger.warning( + "%d Place ID placeholder(s) remain in the response.", + unresolved_placeholders, + ) except Exception as e: # pylint: disable=broad-exception-caught logger.error("Error during Place ID cleanup: %s", e) - if "PLACE_ID_FOR_" in final_response_content: - logger.warning("Place ID placeholder found in response.") - # Final safety check: Extract JSON array if marker is present if "" in final_response_content: marker_idx = final_response_content.find("") @@ -276,4 +249,5 @@ def _build_llm_agent( ), instruction=instruction, tools=[grounding_tool, skill_manager_tool], + after_tool_callback=self._after_tool_callback, ) diff --git a/agent/python_agent/agent_with_templates.py b/agent/python_agent/agent_with_templates.py index 0b2ec0a..542fec7 100644 --- a/agent/python_agent/agent_with_templates.py +++ b/agent/python_agent/agent_with_templates.py @@ -14,59 +14,49 @@ """MAUI Agent with template-based latency optimization.""" -import asyncio -import json +import dataclasses +import inspect import logging import pathlib from types import SimpleNamespace from typing import Any, AsyncIterable import uuid -from a2a.types import DataPart from a2a.types import Part -from google.adk import skills as adk_skills from google.adk.agents import run_config from google.adk.agents.llm_agent import LlmAgent -from google.adk.events.event import Event from google.adk.models.lite_llm import LiteLlm from google.adk.models.llm_request import LlmRequest from google.adk.runners import Runner -from google.adk.tools.set_model_response_tool import SetModelResponseTool from google.genai import types import pydantic from a2ui.a2a.parts import create_a2ui_part +from a2ui.schema.constants import VERSION_0_9 from a2ui.schema.manager import ( A2uiSchemaManager, ) from agent import MAUIAgent from agent_config import AgentConfig from agent_config import FallbackMode -from extractor import DirectionsExtractorSchema -from extractor import LocalSearchExtractorSchema from merger import merge_template -from router_config import IntentClass -from router_config import ROUTER_SYSTEM_INSTRUCTION -from router_config import RouterClassification +from template_registry import ( + INTENT_OTHER_SPATIAL, + INTENT_TEXT_ONLY, + TemplateRegistry, +) +from template_tool import ( + BaseTemplateTool, + RenderTextOnlyTemplateTool, + STATE_RENDERED_A2UI_DATA, + STATE_RENDERED_A2UI_PARTS, +) logger = logging.getLogger(__name__) -_SKILL_BASE_PATH = pathlib.Path(__file__).parent / "skills" _SHARED_INSTRUCTIONS_PATH = ( pathlib.Path(__file__).parent / "shared" / "instructions" ) -_LOCAL_SEARCH_SKILL_NAME = "local-search-template-response" -_LOCAL_SEARCH_TEMPLATE_NAME = "local_search" -_LOCAL_SEARCH_SURFACE_PREFIX = "local-search-surface" -_DIRECTIONS_SKILL_NAME = "directions-template-response" -_DIRECTIONS_TEMPLATE_NAME = "directions" -_DIRECTIONS_SURFACE_PREFIX = "directions-surface" - -_EXTRACTOR_SCHEMAS = { - _LOCAL_SEARCH_SKILL_NAME: LocalSearchExtractorSchema, - _DIRECTIONS_SKILL_NAME: DirectionsExtractorSchema, -} -_SUPPORTED_INTENTS = {IntentClass.LOCAL_SEARCH, IntentClass.DIRECTIONS} _GROUNDED_TEXT_BASE_INSTRUCTION = """\ You are an expert location and navigation assistant with access to Google Maps tools. @@ -80,21 +70,57 @@ """ +def _normalize_version(version: str | None) -> str: + """Canonicalizes an A2UI version string to the key used by version maps. + + The A2UI extension URI carries the version as `v0.9`, while the schema + constants (and therefore `_schema_managers`) use the bare `0.9`. Callers may + pass either form, so both are collapsed to the bare form here. + + Args: + version: Requested A2UI version, e.g. `0.9` or `v0.9`. May be None. + + Returns: + The bare version string, defaulting to `VERSION_0_9` when unspecified. + """ + if not version: + return VERSION_0_9 + return version.removeprefix("v") + + class MAUIAgentWithTemplates(MAUIAgent): """MAUI Agent extending base with server-side layout templates and query intent routing.""" def __init__(self, base_url: str, config: AgentConfig | None = None) -> None: self.config = config or AgentConfig() super().__init__(base_url=base_url, model_name=self.config.generic_model) + self.template_registry = TemplateRegistry() self.router_client = LiteLlm(model=self.config.router_model) self.extractor_client = LiteLlm(model=self.config.template_model) self.fallback_client = LiteLlm(model=self.config.generic_model) + self._extractor_agents: dict[str, LlmAgent] = {} + self._extractor_runners: dict[str, Runner] = {} + self._grounded_text_agent = self._build_grounded_text_agent( + client=self.extractor_client + ) + self._grounded_text_runner = self._build_runner(self._grounded_text_agent) + self._fallback_text_agent = self._build_grounded_text_agent( + client=self.fallback_client + ) + self._fallback_text_runner = self._build_runner(self._fallback_text_agent) + + for version, schema_manager in self._schema_managers.items(): + agent = self._build_unified_extractor_agent(schema_manager=schema_manager) + key = _normalize_version(version) + self._extractor_agents[key] = agent + self._extractor_runners[key] = self._build_runner(agent) + def _build_runner(self, agent: LlmAgent) -> Runner: runner = super()._build_runner(agent) - # The extractor agent runs inside a dynamically created runner. + # The extractor agent runs inside a persistent runner. # We must enable auto_create_session to prevent SessionNotFoundError - # since we don't pre-create the session for this runner. + # when new sessions are encountered. runner.auto_create_session = True return runner @@ -107,13 +133,14 @@ def _on_tool_error( ) -> dict[str, Any] | None: """Callback for tool errors during extraction.""" # pylint: disable=unused-argument - if tool.name == "set_model_response" and isinstance( + if isinstance(tool, BaseTemplateTool) and isinstance( error, pydantic.ValidationError ): + tool_name = getattr(tool, "name", "unknown") logger.warning( "Extractor tool '%s' failed validation: %s. " "Returning error to model for self-correction.", - tool.name, + tool_name, error, ) return { @@ -137,86 +164,126 @@ def _load_shared_guidelines(self) -> str: logger.warning("Failed to load shared style guidelines: %s", e) return "" - def _build_dynamic_extractor_agent( - self, - skill_name: str, - schema_manager: A2uiSchemaManager | None = None, - ) -> LlmAgent: - """Builds an extractor agent loaded directly with the target skill's prompt.""" - skill_dir = _SKILL_BASE_PATH / skill_name - skill = adk_skills.load_skill_from_dir(skill_dir) - skill_instructions = skill.instructions + def _build_grounded_text_agent(self, client: LiteLlm) -> LlmAgent: + """Builds a persistent agent for Maps-grounded plain text responses.""" + system_instruction = self._get_grounded_text_instruction() + generate_content_config = None + if ( + client == self.extractor_client + and self.config.extractor_thinking_budget > 0 + ): + generate_content_config = types.GenerateContentConfig( + thinking_config=types.ThinkingConfig( + thinking_budget=self.config.extractor_thinking_budget + ) + ) + + tools = [self.make_grounding_lite_mcp()] + return LlmAgent( + model=client, + name="maui_grounded_text_agent", + description="Agent for text responses with Maps grounding", + instruction=system_instruction, + tools=tools, + generate_content_config=generate_content_config, + ) + + def _build_unified_extractor_instruction(self) -> str: + """Assembles unified extractor system instruction across supported template skills.""" + instructions_list = [ + "You are an expert location and navigation assistant with access to" + " Google Maps tools and layout template rendering tools." + ] + + for bundle in self.template_registry.bundles.values(): + instructions = bundle.skill_instructions + if instructions: + instructions_list.append(f"\n{instructions}") + shared_guidelines = self._load_shared_guidelines() if shared_guidelines: - skill_instructions = f"{skill_instructions}\n\n{shared_guidelines}" - - # Extractors use template_model, generic UI uses generic_model - if skill_name.endswith("-template-response"): - model_name = self.config.template_model - else: - model_name = self.config.generic_model + instructions_list.append(f"## Style Guidelines\n{shared_guidelines}") - logger.info( - f"Building extractor agent for '{skill_name}' using model: {model_name}" + target_intent_rules = [] + for bundle in sorted( + self.template_registry.bundles.values(), + key=lambda b: b.intent, + ): + target_intent_rules.append( + f" - `[TARGET_INTENT: {bundle.intent}]`: Invoke" + f" `{bundle.render_tool_name}`. Do NOT invoke any other template" + " rendering tool." + ) + target_intent_rules_text = "\n".join(target_intent_rules) + + workflow_rules = ( + "## Execution and Tool Usage Rules\n" + "1. Follow the previous rules for gathering the necessary information" + " using Google Maps tools.\n" + "2. When the prompt begins with `[TARGET_INTENT: ]`, you" + " MUST strictly invoke the corresponding template rendering tool as" + " your final step:\n" + f"{target_intent_rules_text}\n" + "3. Your final action MUST be this template rendering tool call. Do" + " NOT emit raw JSON or `` blocks in your text response." ) + instructions_list.append(workflow_rules) - tools = [self.make_grounding_lite_mcp()] - output_schema = _EXTRACTOR_SCHEMAS.get(skill_name) + return "\n\n".join(instructions_list) - generate_content_config = None - if output_schema: - # Manually inject SetModelResponseTool - set_response_tool = SetModelResponseTool(output_schema) - tools.append(set_response_tool) - - # Manually append instruction - workaround_instruction = ( - "IMPORTANT: You have access to other tools, but you must provide" - " your final response using the set_model_response tool with the" - " required structured format. After using any other tools needed to" - " complete the task, always call set_model_response with your final" - " answer in the specified schema format." - ) - if skill_name == _LOCAL_SEARCH_SKILL_NAME: - workaround_instruction += ( - "\nCRITICAL CONSTRAINT: You MUST extract and display at most" - f" {self.config.max_list_size} of the most relevant places. Do not" - " mention, recommend, or extract more than" - f" {self.config.max_list_size} places in your text response or your" - " set_model_response tool call." - ) - skill_instructions = f"{skill_instructions}\n\n{workaround_instruction}" + def _build_unified_extractor_agent( + self, + schema_manager: A2uiSchemaManager | None = None, + ) -> LlmAgent: + """Builds a single unified extractor agent with all template rendering tools.""" + raw_instructions = self._build_unified_extractor_instruction() - if self.config.extractor_thinking_budget > 0: - generate_content_config = types.GenerateContentConfig( - thinking_config=types.ThinkingConfig( - thinking_budget=self.config.extractor_thinking_budget - ) - ) - logger.info( - "Applying template extractor thinking budget limit:" - f" {self.config.extractor_thinking_budget} tokens" + tools = [self.make_grounding_lite_mcp()] + for bundle in self.template_registry.bundles.values(): + tools.append( + BaseTemplateTool.from_bundle( + bundle, + schema_manager=schema_manager, + max_list_size=self.config.max_list_size, + ) + ) + tools.append( + RenderTextOnlyTemplateTool( + schema_manager=schema_manager, + surface_id_prefix="text-only-surface", ) + ) - if schema_manager: - instruction = schema_manager.generate_system_prompt( - role_description=skill_instructions, - include_schema=True, - include_examples=False, - validate_examples=False, + generate_content_config = None + if self.config.extractor_thinking_budget > 0: + generate_content_config = types.GenerateContentConfig( + thinking_config=types.ThinkingConfig( + thinking_budget=self.config.extractor_thinking_budget + ) ) - else: - instruction = skill_instructions + logger.info( + "Applying template extractor thinking budget limit:" + f" {self.config.extractor_thinking_budget} tokens" + ) + + # Template tools handle template merging and use schema_manager only for + # catalog validation. Do not call schema_manager.generate_system_prompt() + # here, as it injects raw A2UI schema and output instructions. + instruction = raw_instructions return LlmAgent( - model=LiteLlm(model=model_name), - name="maui_agent", - description="An extractor agent executing specific Maps tool tasks", + model=LiteLlm(model=self.config.template_model), + name="maui_unified_extractor_agent", + description=( + "A unified extractor agent executing Maps tasks and template" + " rendering" + ), instruction=instruction, tools=tools, - output_schema=None, # Keep output_schema as None in LlmAgent + output_schema=None, generate_content_config=generate_content_config, on_tool_error_callback=self._on_tool_error, + after_tool_callback=self._after_tool_callback, ) async def _run_extractor( @@ -225,11 +292,26 @@ async def _run_extractor( agent: LlmAgent, current_message: types.Content, session_id: str, - ) -> tuple[dict[str, Any] | None, list[str]]: - """Runs the extractor agent and collects its output (structured or text).""" - parsed_json_data = None + state_delta: dict[str, Any] | None = None, + ) -> tuple[list[Part] | None, list[str], dict[str, Any] | None]: + """Runs the extractor agent and collects its output (rendered parts or text).""" + rendered_parts: list[Part] | None = None + rendered_data: dict[str, Any] | None = None full_content_list = [] + merged_state_delta = { + "expression": "{expression}", + "base_url": self.base_url, + } + if state_delta: + merged_state_delta.update(state_delta) + + template_tools = { + t.name: t + for t in getattr(agent, "tools", []) + if isinstance(t, BaseTemplateTool) + } + async for event in runner.run_async( user_id=self._user_id, session_id=session_id, @@ -237,63 +319,43 @@ async def _run_extractor( streaming_mode=run_config.StreamingMode.SSE ), new_message=current_message, - # Initialize session state. - # "expression" is required to prevent KeyError during ADK's prompt - # state injection, as the A2UI catalog schema contains "${expression}" - # placeholders. "base_url" is passed for consistency with the main - # agent session state. - state_delta={ - "expression": "{expression}", - "base_url": self.base_url, - }, + state_delta=merged_state_delta, ): if hasattr(event, "get_function_calls"): for fc in event.get_function_calls(): - if fc.name == "set_model_response": + target_tool = template_tools.get(fc.name) + if target_tool is not None: logger.info( - "Intercepted set_model_response tool call with args: %s", + "--- AGENT_WITH_TEMPLATES: Observed %s tool call with args:" + " %s ---", + fc.name, fc.args, ) - # Find SetModelResponseTool in agent tools - target_tool = None - for t in agent.tools: - if getattr(t, "name", None) == "set_model_response": - target_tool = t - break - - if target_tool and hasattr(target_tool, "run_async"): + if hasattr(target_tool, "run_async"): + tool_ctx = SimpleNamespace(state={}) try: - noop_tool_context = SimpleNamespace( - actions=SimpleNamespace(set_model_response=None) - ) - validated_data = await target_tool.run_async( - args=fc.args, tool_context=noop_tool_context + tool_result = await target_tool.run_async( + args=fc.args, tool_context=tool_ctx ) - # SetModelResponseTool.run_async catches ValidationError internally - # and returns a dict with "error" key instead of raising the exception. if ( - isinstance(validated_data, dict) - and "error" in validated_data + isinstance(tool_result, dict) + and "error" not in tool_result + and STATE_RENDERED_A2UI_PARTS in tool_ctx.state ): - logger.warning( - "Local Pydantic validation failed: %s. Continuing.", - validated_data["error"], - ) - else: - parsed_json_data = validated_data + rendered_parts = tool_ctx.state[STATE_RENDERED_A2UI_PARTS] + rendered_data = tool_ctx.state.get(STATE_RENDERED_A2UI_DATA) logger.info( - "Local Pydantic validation passed! Short-circuiting." + "--- AGENT_WITH_TEMPLATES: Template tool %s succeeded!" + " Captured %d rendered parts. ---", + fc.name, + len(rendered_parts), ) break - except pydantic.ValidationError as e: + except Exception as e: # pylint: disable=broad-exception-caught logger.warning( - "Local Pydantic validation failed: %s. Continuing.", - e, + "--- AGENT_WITH_TEMPLATES: Tool execution error: %s ---", e ) - else: - parsed_json_data = fc.args - break if event.content and event.content.parts: if event.partial: @@ -306,69 +368,68 @@ async def _run_extractor( if p.text: full_content_list.append(p.text) - return parsed_json_data, full_content_list + if rendered_parts is None and getattr(runner, "session_service", None): + get_session_fn = getattr(runner.session_service, "get_session", None) + if callable(get_session_fn): + try: + res = get_session_fn( + app_name=getattr(runner, "app_name", ""), + user_id=self._user_id, + session_id=session_id, + ) + if inspect.isawaitable(res): + session = await res + if session and getattr(session, "state", None): + rendered_parts = session.state.get(STATE_RENDERED_A2UI_PARTS) + rendered_data = session.state.get(STATE_RENDERED_A2UI_DATA) + except Exception as e: # pylint: disable=broad-exception-caught + logger.debug("Could not retrieve session from session_service: %s", e) + + return rendered_parts, full_content_list, rendered_data async def _run_extractor_and_merge( self, - skill_name: str, - template_name: str, - surface_id_prefix: str, + intent: str, cleaned_query: str, session_id: str, ui_version: str | None = None, ) -> tuple[list[Part] | None, str | None, dict[str, Any] | None]: - """Runs the dynamic extractor agent and merges output into the template.""" - # 1. Resolve catalog schema manager and validator - schema_manager = self._schema_managers.get(ui_version) - selected_catalog = None - if schema_manager: - # Retrieve the resolved catalog config for validation. - # Replacing the deprecated get_catalog("maps-agentic-ui-catalog") - # API call. - selected_catalog = schema_manager.get_selected_catalog() - - # 2. Build the extractor agent and runner - agent = self._build_dynamic_extractor_agent( - skill_name, - schema_manager=schema_manager, - ) - runner = self._build_runner(agent) + """Runs the persistent extractor agent runner and returns rendered template parts.""" + version = _normalize_version(ui_version) + + # 1. Resolve persistent runner and agent, falling back to the default + # version when the requested one has no pre-built extractor. + if version not in self._extractor_agents: + logger.warning( + "No extractor agent built for A2UI version '%s'. Falling back to" + " '%s'.", + version, + VERSION_0_9, + ) + version = VERSION_0_9 + agent = self._extractor_agents[version] + runner = self._extractor_runners[version] - # 3. Setup user query message + # 2. Setup user query message with target intent hint + target_hint = intent + formatted_prompt = f"[TARGET_INTENT: {target_hint}]\n{cleaned_query}" current_message = types.Content( - role="user", parts=[types.Part.from_text(text=cleaned_query)] + role="user", parts=[types.Part.from_text(text=formatted_prompt)] ) - # 4. Run extractor runner, collecting output - parsed_json_data, full_content_list = await self._run_extractor( - runner, agent, current_message, session_id + # 3. Run extractor runner, collecting output + rendered_parts, full_content_list, rendered_data = ( + await self._run_extractor( + runner, + agent, + current_message, + session_id, + state_delta={"target_intent": target_hint}, + ) ) - # 5. Handle output layout merging - if parsed_json_data is not None: - logger.info( - "Template parameters extracted successfully. Merging template." - ) - if "surface_id" not in parsed_json_data: - short_id = uuid.uuid4().hex[:8] - parsed_json_data["surface_id"] = f"{surface_id_prefix}-{short_id}" - - merged_actions = merge_template( - template_name, - parsed_json_data, - max_list_size=self.config.max_list_size, - ) - - if selected_catalog: - logger.info("Validating merged template against A2UI catalog schema.") - try: - selected_catalog.validator.validate(merged_actions) - except Exception as e: # pylint: disable=broad-exception-caught - logger.warning("Catalog validation failed: %s. Falling back.", e) - return None, None, None - - final_parts = [create_a2ui_part(action) for action in merged_actions] - return final_parts, None, parsed_json_data + if rendered_parts is not None: + return rendered_parts, None, rendered_data else: raw_text = "".join(full_content_list) return None, raw_text, None @@ -394,7 +455,7 @@ async def stream( intent, cleaned_query = await self._classify_intent(query) - if intent == IntentClass.OTHER_SPATIAL: + if intent == INTENT_OTHER_SPATIAL: if self.config.fallback_mode == FallbackMode.DYNAMIC: logger.warning( "Router matched OTHER_SPATIAL and fallback_mode is DYNAMIC. " @@ -417,7 +478,7 @@ async def stream( } return - elif intent == IntentClass.TEXT_ONLY: + elif intent == INTENT_TEXT_ONLY: logger.info("Executing fast text response flow for TEXT_ONLY intent.") final_parts = await self._handle_text_only(cleaned_query, session_id) yield { @@ -426,7 +487,7 @@ async def stream( } return - elif intent in _SUPPORTED_INTENTS: + elif intent in self.template_registry.supported_intents: async for part in self._handle_extracted_intent( intent, cleaned_query, session_id, ui_version ): @@ -457,7 +518,7 @@ async def stream( async for part in super().stream(query, session_id, ui_version): yield part - async def _classify_intent(self, query: str) -> tuple[IntentClass, str]: + async def _classify_intent(self, query: str) -> tuple[str, str]: """Classifies the query intent and returns the intent and cleaned query.""" logger.info( "Routing query: '%s' using model %s", @@ -465,10 +526,13 @@ async def _classify_intent(self, query: str) -> tuple[IntentClass, str]: self.config.router_model, ) try: + classification_schema = self.template_registry.router_classification router_config = { - "system_instruction": ROUTER_SYSTEM_INSTRUCTION, + "system_instruction": ( + self.template_registry.compile_router_instruction() + ), "response_mime_type": "application/json", - "response_schema": RouterClassification, + "response_schema": classification_schema, } if self.config.router_thinking_budget > 0: router_config["thinking_config"] = types.ThinkingConfig( @@ -492,10 +556,14 @@ async def _classify_intent(self, query: str) -> tuple[IntentClass, str]: router_response_text += p.text logger.info("Router response content: %s", router_response_text) - classification = RouterClassification.model_validate_json( + classification = classification_schema.model_validate_json( router_response_text ) - intent = classification.intent + intent = str( + classification.intent.value + if hasattr(classification.intent, "value") + else classification.intent + ) cleaned_query = classification.query logger.info( "Intent classified: %s (Cleaned Query: '%s')", intent, cleaned_query @@ -507,7 +575,7 @@ async def _classify_intent(self, query: str) -> tuple[IntentClass, str]: e, exc_info=True, ) - return IntentClass.TEXT_ONLY, query + return INTENT_TEXT_ONLY, query def _wrap_in_text_only(self, text: str, session_id: str) -> list[Part]: """Wraps plain text in a text_only template Part list.""" @@ -526,31 +594,24 @@ def _get_grounded_text_instruction(self) -> str: return _GROUNDED_TEXT_BASE_INSTRUCTION async def _handle_grounded_text( - self, cleaned_query: str, session_id: str, client: LiteLlm + self, + cleaned_query: str, + session_id: str, + client: LiteLlm | None = None, + runner: Runner | None = None, ) -> list[Part]: - """Generates a grounded plain text response using the provided model client with GroundingLite tools.""" - system_instruction = self._get_grounded_text_instruction() - generate_content_config = None - if ( - client == self.extractor_client - and self.config.extractor_thinking_budget > 0 - ): - generate_content_config = types.GenerateContentConfig( - thinking_config=types.ThinkingConfig( - thinking_budget=self.config.extractor_thinking_budget - ) - ) + """Generates a grounded plain text response using pre-configured runner with GroundingLite tools.""" + if runner is None: + if client == self.fallback_client: + runner = self._fallback_text_runner + else: + runner = self._grounded_text_runner + + if runner is None: + target_client = client or self.extractor_client + agent = self._build_grounded_text_agent(client=target_client) + runner = self._build_runner(agent) - tools = [self.make_grounding_lite_mcp()] - agent = LlmAgent( - model=client, - name="maui_grounded_text_agent", - description="Agent for text responses with Maps grounding", - instruction=system_instruction, - tools=tools, - generate_content_config=generate_content_config, - ) - runner = self._build_runner(agent) current_message = types.Content( role="user", parts=[types.Part.from_text(text=cleaned_query)] ) @@ -594,7 +655,10 @@ async def _handle_text_only( ) -> list[Part]: """Generates a plain text response for TEXT_ONLY intent using template_model with grounding.""" return await self._handle_grounded_text( - cleaned_query, session_id, client=self.extractor_client + cleaned_query, + session_id, + client=self.extractor_client, + runner=self._grounded_text_runner, ) async def _handle_grounded_text_fallback( @@ -602,35 +666,29 @@ async def _handle_grounded_text_fallback( ) -> list[Part]: """Generates a grounded plain text response for fallback/complex spatial queries using generic_model.""" return await self._handle_grounded_text( - cleaned_query, session_id, client=self.fallback_client + cleaned_query, + session_id, + client=self.fallback_client, + runner=self._fallback_text_runner, ) async def _handle_extracted_intent( self, - intent: IntentClass, + intent: str, query: str, session_id: str, ui_version: str | None = None, ) -> AsyncIterable[dict[str, Any]]: """Handles intents that use dynamic extractor agents and templates.""" - if intent == IntentClass.LOCAL_SEARCH: - skill_name = _LOCAL_SEARCH_SKILL_NAME - template_name = _LOCAL_SEARCH_TEMPLATE_NAME - surface_prefix = _LOCAL_SEARCH_SURFACE_PREFIX - elif intent == IntentClass.DIRECTIONS: - skill_name = _DIRECTIONS_SKILL_NAME - template_name = _DIRECTIONS_TEMPLATE_NAME - surface_prefix = _DIRECTIONS_SURFACE_PREFIX - else: + bundle = self.template_registry.get_bundle_by_intent(intent) + if bundle is None: raise ValueError(f"Unsupported intent for extractor: {intent}") logger.info("Router matched %s. Dispatching template extractor.", intent) merged_parts, fallback_text, parsed_json_data = ( await self._run_extractor_and_merge( - skill_name=skill_name, - template_name=template_name, - surface_id_prefix=surface_prefix, + intent=intent, cleaned_query=query, session_id=session_id, ui_version=ui_version, diff --git a/agent/python_agent/merger.py b/agent/python_agent/merger.py index ffdd9b8..be25660 100644 --- a/agent/python_agent/merger.py +++ b/agent/python_agent/merger.py @@ -14,18 +14,41 @@ """Server-Side A2UI Layout Template Merger (`merger.py`). -This module loads declarative JSON layout skeletons (`templates/*.json`), -populates them with extracted parameter dictionaries (`data`), and returns -sanitized message structures ready for wire transmission. +This module loads a declarative JSON layout skeleton, populates it with a +dictionary of extracted parameters, and returns sanitized message structures +ready for wire transmission. + +The merger is template agnostic. Everything specific to one template -- its +schema, its layout, and what to say when it cannot render -- arrives as caller +arguments. Derived and repaired fields belong in the bundle's Pydantic schema, +never here. """ import copy +import importlib.resources import json -import os +import logging +import re from typing import Any, Literal, TypedDict import uuid -from extractor import normalize_travel_mode +import pydantic + +import ( + template_registry, +) + +logger = logging.getLogger(__name__) + +# The bare text surface every template degrades to when it cannot render. +_TEXT_ONLY_TEMPLATE = "text_only" + +# The catalog the templates bind their surface to when the host does not name +# one. Hosts that register the Maps components in a catalog of their own (for +# example Gemini Enterprise, whose composite catalog carries `GoogleMap` and +# `PlaceDetailsCompact` alongside its Material and basic components) pass that +# catalog's id instead, because a surface resolves against exactly one catalog. +DEFAULT_CATALOG_ID = "a2ui://maps-agentic-ui-catalog.json" class TextOutputDict(TypedDict): @@ -64,7 +87,15 @@ class TextOutputDict(TypedDict): def _replace_placeholders(obj: Any, data: dict[str, Any]) -> Any: - """Recursively replaces string placeholders in the template with values from data.""" + """Recursively replaces string placeholders in the template with values. + + Args: + obj: The JSON-serializable template object. + data: Key-value pairs used to substitute placeholders. + + Returns: + The updated template object with resolved placeholder values. + """ if isinstance(obj, str): # Exact match: "{{key}}" if obj.startswith("{{") and obj.endswith("}}") and obj.count("{{") == 1: @@ -90,7 +121,7 @@ def _replace_placeholders(obj: Any, data: dict[str, Any]) -> Any: def _remove_none_values(val: Any) -> Any: - """Recursively removes None values from dicts and lists to satisfy JSON schemas.""" + """Recursively removes None values from dicts/lists to satisfy schemas.""" if isinstance(val, dict): return {k: _remove_none_values(v) for k, v in val.items() if v is not None} elif isinstance(val, list): @@ -98,233 +129,133 @@ def _remove_none_values(val: Any) -> Any: return val -def _prepare_local_search( - data: dict[str, Any], max_list_size: int -) -> tuple[str, dict[str, Any]]: - """Validates and normalizes parameters for the local search template.""" - data_copy = copy.deepcopy(data) - is_valid = True - places = data_copy.get("places") - - # 1. Validate that places is a non-empty list - if not isinstance(places, list) or not places: - is_valid = False - else: - # Slice and sanitize places list - sanitized_places = [] - for p in places: - if isinstance(p, dict): - try: - p["lat"] = float(p["lat"]) - p["lng"] = float(p["lng"]) - sanitized_places.append(p) - except (KeyError, ValueError, TypeError): - pass - if not sanitized_places: - is_valid = False - else: - data_copy["places"] = sanitized_places[:max_list_size] - - # 2. Validate mandatory map centering and zoom parameters - if is_valid: - try: - data_copy["center_lat"] = float(data_copy["center_lat"]) - data_copy["center_lng"] = float(data_copy["center_lng"]) - data_copy["zoom"] = int(data_copy["zoom"]) - except (KeyError, ValueError, TypeError): - is_valid = False - - # 3. Sanitize optional anchor marker - if is_valid and "anchor_marker" in data_copy: - pin = data_copy["anchor_marker"] - if isinstance(pin, dict): - try: - pin["lat"] = float(pin["lat"]) - pin["lng"] = float(pin["lng"]) - except (KeyError, ValueError, TypeError): - data_copy["anchor_marker"] = None - else: - data_copy["anchor_marker"] = None - - # 4. Unroll maps markers array - if is_valid: - if "markers" not in data_copy: - markers = [] - for p in data_copy["places"]: - marker = { - "lat": p["lat"], - "lng": p["lng"], - "label": p.get("name") or p.get("label") or "", - } - if "placeId" in p: - marker["placeId"] = p["placeId"] - markers.append(marker) - data_copy["markers"] = markers - else: - markers = data_copy["markers"] - if isinstance(markers, list): - sanitized_markers = [] - for m in markers: - if isinstance(m, dict): - try: - m["lat"] = float(m["lat"]) - m["lng"] = float(m["lng"]) - m["label"] = str(m.get("label") or "") - sanitized_markers.append(m) - except (KeyError, ValueError, TypeError): - pass - data_copy["markers"] = sanitized_markers - else: - data_copy["markers"] = None - - # If validation failed, fallback to text_only - if not is_valid: - return "text_only", { - "text": ( - data.get("summary") - or "No places matching your query could be found." - ), - "surface_id": data_copy.get("surface_id") or "fallback-surface", - } - - return "local_search", data_copy - - -def _prepare_directions(data: dict[str, Any]) -> tuple[str, dict[str, Any]]: - """Validates and normalizes parameters for the directions template.""" - data_copy = copy.deepcopy(data) - is_valid = True - - routes = data_copy.get("routes") - - # 1. Validate that routes is a non-empty list of segment dicts - if not isinstance(routes, list) or not routes: - is_valid = False - else: - sanitized_routes = [] - for segment in routes: - if not isinstance(segment, dict): - is_valid = False - break - origin = segment.get("origin") - destination = segment.get("destination") - if not isinstance(origin, dict) or not isinstance(destination, dict): - is_valid = False - break - try: - sanitized_origin = { - "lat": float(origin["lat"]), - "lng": float(origin["lng"]), - "label": str(origin.get("label") or ""), - } - if "placeId" in origin: - sanitized_origin["placeId"] = str(origin["placeId"]) - - sanitized_destination = { - "lat": float(destination["lat"]), - "lng": float(destination["lng"]), - "label": str(destination.get("label") or ""), - } - if "placeId" in destination: - sanitized_destination["placeId"] = str(destination["placeId"]) - - sanitized_routes.append({ - "origin": sanitized_origin, - "destination": sanitized_destination, - }) - except (KeyError, ValueError, TypeError): - is_valid = False - break - data_copy["routes"] = sanitized_routes - - # 2. Validate mandatory map centering and zoom parameters - if is_valid: - try: - data_copy["center_lat"] = float(data_copy["center_lat"]) - data_copy["center_lng"] = float(data_copy["center_lng"]) - data_copy["zoom"] = int(data_copy["zoom"]) - except (KeyError, ValueError, TypeError): - is_valid = False - - # 3. Normalize travel_mode - if is_valid and "travel_mode" in data_copy: - normalized = normalize_travel_mode(data_copy["travel_mode"]) - if normalized: - data_copy["travel_mode"] = normalized - else: - del data_copy["travel_mode"] +def _uniquify_surface_id( + data: dict[str, Any], template_name: str, surface_prefix: str | None +) -> None: + """Appends a random suffix when the surface id is a shared default. - # If validation failed, fallback to text_only - if not is_valid: - return "text_only", { - "text": data.get("summary") or "Could not calculate travel directions.", - "surface_id": data_copy.get("surface_id") or "fallback-surface", - } + Surfaces keyed by a bundle's default id would collide across concurrent + turns, so only an id the caller chose itself survives untouched. - return "directions", data_copy + Args: + data: Parameter dict, mutated in place. + template_name: Name of the bundle being rendered. + surface_prefix: The bundle's declared default surface id, if any. + """ + default_ids = { + f"{template_name}_surface", + f"{template_name.replace('_', '-')}-surface", + "fallback-surface", + } + if surface_prefix: + default_ids.add(surface_prefix) + + surface_id = data.get("surface_id") + if surface_id and surface_id not in default_ids: + return + base_id = surface_id or f"{template_name}_surface" + data["surface_id"] = f"{base_id}_{uuid.uuid4().hex[:6]}" + + +def _merge_fallback( + data: dict[str, Any], + empty_fallback_text: str, + catalog_id: str = DEFAULT_CATALOG_ID, +) -> list[MergedMessage]: + """Renders a template's fallback text on a bare text-only surface.""" + return merge_template( + _TEXT_ONLY_TEMPLATE, + { + "text": data.get("summary") or empty_fallback_text, + "surface_id": data.get("surface_id") or "fallback-surface", + }, + catalog_id=catalog_id, + ) def merge_template( - template_name: str, data: dict[str, Any], max_list_size: int = 5 + template_name: str, + data: dict[str, Any], + max_list_size: int = 5, + catalog_id: str = DEFAULT_CATALOG_ID, ) -> list[MergedMessage]: - """Loads static template skeleton JSON and returns merged wire response. - - This function sanitizes extracted parameters and populates the layout - template skeleton. + """Populates a layout skeleton with extracted parameters. Args: - template_name: Target declarative layout skeleton (`local_search`, - `directions`, or `text_only`). - data: Raw dictionary of extracted parameters yielded by - `TemplateExtractor` (e.g. `places`, `summary`, `center_lat`). - max_list_size: Maximum allowable child elements in lists (`places`) to - bound payload rendering latency. + template_name: Name of the bundle whose layout to render, or 'text_only'. + data: Extracted parameters to populate the layout with. + max_list_size: Maximum number of elements retained in any top-level + collection, to bound payload size and render latency. + catalog_id: Catalog the created surface binds to. Defaults to the Maps + catalog; hosts that register the Maps components under a catalog of + their own pass that id. Returns: - A list of message dictionaries. For `text_only`, returns the 2-part A2UI - layout (`createSurface`, `updateComponents`). - For `local_search` and `directions`, returns the 3-part A2UI layout - (`createSurface`, `updateComponents`, `updateDataModel`) (`DataPart`). + A list of message dictionaries for A2UI transmission. Raises: - FileNotFoundError: If the template file cannot be found. + FileNotFoundError: If the template or its layout file cannot be found. """ + # Templates ship as package data. They are resolved through + # `importlib.resources` rather than a `__file__`-relative filesystem path + # because the agent is also served from a zip-imported archive, where + # `__file__` names an entry inside the archive that never exists on disk and + # every `os.path` probe against it reports missing. + templates_dir = importlib.resources.files(__package__).joinpath("templates") + resource_path = f"templates/{template_name}/layout.json" + template_path = templates_dir.joinpath(template_name, "layout.json") + + if template_name == _TEXT_ONLY_TEMPLATE: + schema_class = None + empty_fallback_text = "" + surface_prefix = "text-only-surface" + else: + bundle = template_registry.TemplateRegistry().get_bundle(template_name) + schema_class = bundle.schema_cls if bundle else None + empty_fallback_text = bundle.empty_fallback_text if bundle else "" + surface_prefix = ( + bundle.surface_prefix if bundle else f"{template_name}-surface" + ) - # Deep copy data to prevent unintended side-effects on caller dictionaries - # across turns - data_copy = copy.deepcopy(data) - - if template_name == "local_search": - template_name, data_copy = _prepare_local_search(data_copy, max_list_size) - elif template_name == "directions": - template_name, data_copy = _prepare_directions(data_copy) - current_dir = os.path.dirname(os.path.abspath(__file__)) - templates_dir = os.path.join(current_dir, "templates") - template_path = os.path.join(templates_dir, f"{template_name}.json") - - if not os.path.exists(template_path): + if template_path.is_file(): + template_json = json.loads(template_path.read_text(encoding="utf-8")) + elif bundle is not None and bundle.layout_path.is_file(): + template_json = json.loads(bundle.layout_path.read_text(encoding="utf-8")) + else: raise FileNotFoundError( - f"Template '{template_name}' not found at {template_path}" + f"Template '{template_name}' not found: package '{__package__}'" + f" has no '{resource_path}'" ) - with open(template_path, "r") as f: - template_json = json.load(f) - - # Smart Turn-Unique `surface_id` Scoping via Dynamic Template Discovery: - default_surface_ids = set() - if os.path.exists(templates_dir): - for fn in os.listdir(templates_dir): - if fn.endswith(".json"): - base_name = fn[:-5] - default_surface_ids.add(f"{base_name}_surface") - default_surface_ids.add(f"{base_name.replace('_', '-')}-surface") - - if ( - not data_copy.get("surface_id") - or data_copy.get("surface_id") in default_surface_ids - ): - base_id = data_copy.get("surface_id") or f"{template_name}_surface" - data_copy["surface_id"] = f"{base_id}_{uuid.uuid4().hex[:6]}" + data_copy = copy.deepcopy(data) + data_copy["catalog_id"] = catalog_id + + # 1. Bound every top-level collection to keep payload size predictable. + for key, value in list(data_copy.items()): + if isinstance(value, list): + data_copy[key] = value[:max_list_size] + + # 2. Validate through the bundle's schema. The schema owns every + # normalization, derived field, and non-empty constraint, so a collection + # the layout cannot render empty fails here and degrades to text. The + # dumped model -- not the raw input -- is what feeds the layout. + if schema_class is not None: + try: + model_instance = schema_class.model_validate(data_copy) + except pydantic.ValidationError: + logger.warning( + "Data for template '%s' failed schema validation; degrading to a" + " text-only surface.", + template_name, + exc_info=True, + ) + return _merge_fallback(data_copy, empty_fallback_text, catalog_id) + data_copy.update(model_instance.model_dump()) + + # 3. Scope the surface id to this turn so repeated renders do not collide. + _uniquify_surface_id(data_copy, template_name, surface_prefix) resolved_json = _replace_placeholders(template_json, data_copy) + # Prune keys whose optional placeholders (e.g. "{{anchorMarker}}") were + # absent from data_copy and evaluated to None, keeping wire JSON clean. return _remove_none_values(resolved_json) diff --git a/agent/python_agent/place_id_resolution.py b/agent/python_agent/place_id_resolution.py new file mode 100644 index 0000000..3525408 --- /dev/null +++ b/agent/python_agent/place_id_resolution.py @@ -0,0 +1,229 @@ +# Copyright 2026 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. + +"""Place ID resolution shared by the grounded agents. + +Grounding with Google Maps does not let the model see real Place IDs while it +is generating, and asking it to recall them from parametric memory produces +wrong IDs. Instead the system instruction tells the model to emit a +placeholder, and this module rewrites those placeholders using the Place +IDs that Grounding actually returned. + +Both `agent_with_grounding` (free-form A2UI) and `vertex_grounding_extractor` +(template parameters) use this module so the two paths cannot drift apart. + +Placeholder format, which the prompt and this module must agree on exactly: + + PLACE_ID_FOR_{count}_{title} + +`title` is the Maps source title with the " - Google Maps" branding suffix +removed, in its original casing. `count` is the 1-based occurrence index of +the source within that title, so five sources sharing a title get placeholders +1 through 5. +""" + +import collections +import dataclasses +import logging +import re +from typing import Any + +logger = logging.getLogger(__name__) + +PLACEHOLDER_PREFIX = "PLACE_ID_FOR_" + +# Kept beside the parser on purpose. The placeholder format is a contract +# between this module and the model, and a prompt edit that drifts from +# build_placeholder_index fails silently: every key misses and raw +# placeholders ship to the client. +# +# The rules are two-branch because whether the model can see a real Place ID +# depends on serving-side configuration. When it is off, no Place ID reaches the +# model, and the placeholder is the only option. +# +# Preferring the authoritative value is strictly better, but the model does not +# always get it, so we have to handle both cases. +PROMPT_RULES = """PLACE ID RULES: +Every `placeId` you emit MUST come from one of the two sources below. Always prefer source 1. +1. AUTHORITATIVE PLACE ID (preferred). If the Google Maps result for a place carries a place_id field -- written either as a `place_id: places/ChIJ...` line or as a `place_id='places/ChIJ...'` field inside a PerQueryResult -- copy that value verbatim into `placeId`. Do NOT reformat, abbreviate, or truncate it. +2. SYNTHETIC PLACEHOLDER (fallback). Use this if, and ONLY if, the Google Maps result for that place carries no place_id field: +- Format: "PLACE_ID_FOR_{Count}_{Exact Title}" +- Example: If the tool returns a place named "Chez Panisse" and gives no place_id, use "PLACE_ID_FOR_1_Chez Panisse". If it returns a second "Chez Panisse", use "PLACE_ID_FOR_2_Chez Panisse". +- STRICT MATCHING: Do NOT change any characters, spaces, capitalization, or punctuation from the title returned by the tool. +- COUNTING: Always prepend the occurrence count (starting at 1) for each title based on the order they were returned by the tool, even if the title only occurs once. +NEVER invent a `ChIJ...` value that the tool did not give you. If you have no place_id from the tool, the placeholder in rule 2 is the only acceptable output.""" + +_BRANDED_TITLE_SUFFIX = " - Google Maps" + +# Slicing at len(prefix) - 3 keeps the "ChI" that every Place ID starts with. +_PLACE_ID_PREFIX = "places/ChI" + + +@dataclasses.dataclass(frozen=True) +class AttributionSource: + """A Maps attribution source from grounding metadata.""" + + title: str + place_id: str + + +def normalize_place_id(raw_place_id: str) -> str: + """Strips the 'places/' resource prefix, which A2UI components do not want.""" + if raw_place_id.startswith(_PLACE_ID_PREFIX): + return raw_place_id[len(_PLACE_ID_PREFIX) - 3 :] + return raw_place_id + + +def canonical_place_title(raw_title: str) -> str: + """Drops the branding suffix the model never sees, keeping original casing.""" + if raw_title.endswith(_BRANDED_TITLE_SUFFIX): + return raw_title[: -len(_BRANDED_TITLE_SUFFIX)] + return raw_title + + +def extract_attribution_sources(response: Any) -> list[AttributionSource]: + """Pulls the Maps attribution sources out of a genai response. + + Missing attributes are tolerated so callers do not have to guard every + access. Arrival order is load-bearing: it is what the ordinals count over. + + Args: + response: A `google.genai` GenerateContentResponse, or anything shaped + like one. + + Returns: + Sources that carry both a title and a Place ID, in arrival order. + """ + candidates = getattr(response, "candidates", None) + metadata = ( + getattr(candidates[0], "grounding_metadata", None) if candidates else None + ) + raw_chunks = getattr(metadata, "grounding_chunks", None) if metadata else None + + sources = [] + for raw_chunk in raw_chunks or []: + maps_source = getattr(raw_chunk, "maps", None) + title = getattr(maps_source, "title", None) + place_id = getattr(maps_source, "place_id", None) + if title and place_id: + sources.append(AttributionSource(title=title, place_id=place_id)) + return sources + + +def _build_placeholder_index( + attribution_sources: list[AttributionSource], +) -> dict[str, str]: + """Builds the placeholder-to-Place-ID substitution map. + + Ordinals count source occurrences within a canonical title, which is the + COUNTING rule `PROMPT_RULES` gives the model. The model and this function + derive their ordinals independently, so they can disagree. A disagreement + usually leaves the placeholder unresolved, which is visible in the UI. + + Args: + attribution_sources: Attribution sources in arrival order. + + Returns: + Mapping from placeholder key to canonical Place ID. + """ + placeholder_index = {} + title_counts = {} + for source in attribution_sources: + title = canonical_place_title(source.title) + place_id = normalize_place_id(source.place_id) + if not title or not place_id: + continue + title_counts[title] = title_counts.get(title, 0) + 1 + placeholder_index[f"{PLACEHOLDER_PREFIX}{title_counts[title]}_{title}"] = ( + place_id + ) + return placeholder_index + + +def resolve_place_ids( + text: str, + attribution_sources: list[AttributionSource], +) -> tuple[str, int]: + """Replaces Place ID placeholders in a payload with grounded Place IDs. + + Placeholders with no matching source are deliberately left in place rather + than substituted with a nearby ID. A visible placeholder fails loudly in the + UI, whereas a plausible wrong Place ID renders a confidently wrong venue. + + After substitution, any remaining Places API resource names are reduced to + bare Place IDs, so neither branch of `PROMPT_RULES` can ship a resource name + to the A2UI components. + + Args: + text: Serialized payload containing placeholders. + attribution_sources: Attribution sources in arrival order. + + Returns: + Tuple of the rewritten payload and the count left unresolved. + """ + if not text: + return text, 0 + + # Matches `PLACE_ID_FOR__` to extract ``. + placeholder_title_re = re.compile( + rf"^{PLACEHOLDER_PREFIX}[0-9]+_(?P.+)$" + ) + placeholder_index = _build_placeholder_index(attribution_sources) + matched_titles = set() + for placeholder, place_id in placeholder_index.items(): + if placeholder in text: + match = placeholder_title_re.match(placeholder) + if match: + matched_titles.add(match.group("title")) + text = text.replace(placeholder, place_id) + + # Fallback pass for placeholders where the LLM used global list numbering + # (e.g., `PLACE_ID_FOR_2_<title>` when `<title>` only appeared once as index 1). + # Caveat: Titles already matched in pass 1 (`matched_titles`) are intentionally + # skipped here so that if multiple distinct branches share the same title + # (e.g., two "Starbucks" locations) and one placeholder was mis-indexed, we + # leave the ambiguous placeholder unresolved rather than substituting the + # wrong branch's Place ID. + if PLACEHOLDER_PREFIX in text: + title_counts = collections.Counter( + canonical_place_title(s.title) + for s in attribution_sources + if canonical_place_title(s.title) and normalize_place_id(s.place_id) + ) + for source in sorted( + attribution_sources, + key=lambda s: len(canonical_place_title(s.title)), + reverse=True, + ): + title = canonical_place_title(source.title) + place_id = normalize_place_id(source.place_id) + if ( + title + and place_id + and title not in matched_titles + and title_counts[title] == 1 + ): + pattern = rf"{PLACEHOLDER_PREFIX}\d+_{re.escape(title)}" + if len(re.findall(pattern, text, flags=re.IGNORECASE)) == 1: + text = re.sub(pattern, place_id, text, count=1, flags=re.IGNORECASE) + matched_titles.add(title) + + unresolved = text.count(PLACEHOLDER_PREFIX) + if unresolved: + logger.warning( + "%d Place ID placeholder(s) unresolved against %d source(s).", + unresolved, + len(placeholder_index), + ) + return text, unresolved diff --git a/agent/python_agent/shared/instructions/shared_style_guidelines.md b/agent/python_agent/shared/instructions/shared_style_guidelines.md index 1c29c12..26b83b5 100644 --- a/agent/python_agent/shared/instructions/shared_style_guidelines.md +++ b/agent/python_agent/shared/instructions/shared_style_guidelines.md @@ -1,25 +1,34 @@ -## Conversational Text Style Guidelines +## Response Text Guidelines -When generating conversational text (such as summaries, descriptions, or -directions), you must follow these formatting and content rules: +### Role & Tone -* **Content & Completeness**: Always fully and clearly answer each aspect of - the user's prompt. Address all explicit constraints, qualitative criteria, - comparisons, preferences, and sub-questions asked. Explain *why* places or - routes fit the user's specific needs rather than providing a bare listing. -* **Quantity & Nuance**: Make sure the answer is substantive, useful, and - actionable. Respond with an appropriate depth of detail given the complexity - of the question: - * If comparing places or route alternatives, explicitly analyze their - trade-offs (e.g. transit vs driving, travel time, convenience, cost, or - atmosphere). - * If the user asks about commute, context, or travel conditions, describe - relevant timing and real-world nuances (e.g. rush-hour delays, - navigation landmarks). -* **Formatting**: Use markdown to apply formatting elements like bullet - points, bolding, and tables to break up the text. Break content into - multiple paragraphs as needed. -* **Markdown**: Bold place names and provide links where appropriate. -* **Titles and Headings**: Never title your response. You may include - mid-level headings (using `###` and below) to organize content when it adds - clarity. +- **Voice**: Warm local expert. Show warmth through highly relevant logistics, + NEVER conversational filler. +- **Style**: Vivid, objective, and sensory (e.g., "low-lit basement"). NEVER + use empty hype words ("amazing", "charming"). +- **Perspective**: NEVER use first-person ("I recommend", "I found"). + Attribute subjective claims to public consensus or facts (e.g., "Locals + praise..."). + +### Execution & Formatting + +- **Headings**: Always use sentence case. Plain text only - NO markdown. +- **Primary headings**: A concise, constraint-confirming title reflecting the + prompt and primary reference location. Use only the primary reference + location without redundant city/state nesting. + - **Place Searches**: Always start with or include the exact number of + places provided in the UI response (e.g., '5 vegetarian restaurants near + The Plaza Hotel', '5 transit stops near Seattle Center'). + - **Directions**: Provide a concise route title confirming the travel mode + and endpoints (e.g., 'Walking route from Seattle Center to Pike Place + Market', 'Driving directions to JFK Airport'). +- **Precision**: Fully answer the prompt and strictly satisfy all constraints. +- **Count matching**: If the prompt requests a specific number of places + (e.g., "3 hidden gem activities", "top 2 cafes", "four places to visit"), + ALWAYS respond with that exact number of grounded places in the `places` + array when possible. +- **Differentiate places**: Describe places by mentioning unique features, + specialties, and review highlights. +- **Reviews**: Never hallucinate place reviews. Only describe user sentiment + in aggregate from a grounded source. +- **Addresses**: Never state full addresses in a response. diff --git a/agent/python_agent/shared/schema/maps_catalog_extension.json b/agent/python_agent/shared/schema/maps_catalog_extension.json index 699754f..35de066 100644 --- a/agent/python_agent/shared/schema/maps_catalog_extension.json +++ b/agent/python_agent/shared/schema/maps_catalog_extension.json @@ -4,6 +4,7 @@ "title": "Google Maps Agentic UI Catalog Extensions", "components": { "GoogleMap": { + "description": "Renders an interactive Google Map with markers, an anchor marker and routes.", "type": "object", "allOf": [ { @@ -44,7 +45,7 @@ "description": "The map mode." }, "anchorMarker": { - "$ref": "#/$defs/DynamicLatLng", + "$ref": "#/$defs/AnchorMarker", "description": "The anchor marker location." }, "markers": { @@ -55,6 +56,7 @@ "description": "List of markers." }, "routes": { + "description": "Routes to draw on the map, as origin/destination pairs.", "type": "array", "items": { "type": "object", @@ -88,6 +90,7 @@ "unevaluatedProperties": false }, "PlaceDetailsCompact": { + "description": "Renders a compact Google Maps place card with photo, rating, type and price.", "type": "object", "allOf": [ { @@ -112,6 +115,7 @@ "horizontal", "vertical" ], + "default": "horizontal", "description": "The orientation of this place card. Defaults to horizontal." } }, @@ -148,6 +152,55 @@ } ] }, + "AnchorMarker": { + "oneOf": [ + { + "$ref": "#/$defs/AnchorMarkerPin" + }, + { + "$ref": "common_types.json#/$defs/DataBinding" + } + ] + }, + "AnchorMarkerPin": { + "type": "object", + "properties": { + "lat": { + "type": "number" + }, + "lng": { + "type": "number" + }, + "label": { + "type": "string" + }, + "placeId": { + "type": "string" + }, + "placePrimaryType": { + "type": "string", + "enum": [ + "food_and_drink", + "outdoor", + "retail", + "gas_station", + "ev", + "bank", + "lodging", + "emergency", + "entertainment", + "airport", + "parking", + "generic" + ] + } + }, + "required": [ + "lat", + "lng" + ], + "additionalProperties": false + }, "MapPin": { "type": "object", "properties": { @@ -162,6 +215,23 @@ }, "placeId": { "type": "string" + }, + "placePrimaryType": { + "type": "string", + "enum": [ + "food_and_drink", + "outdoor", + "retail", + "gas_station", + "ev", + "bank", + "lodging", + "emergency", + "entertainment", + "airport", + "parking", + "generic" + ] } }, "required": [ diff --git a/agent/python_agent/skills/google-maps-enriched-local-query-response/SKILL.md b/agent/python_agent/skills/google-maps-enriched-local-query-response/SKILL.md index 53b29f5..365d949 100644 --- a/agent/python_agent/skills/google-maps-enriched-local-query-response/SKILL.md +++ b/agent/python_agent/skills/google-maps-enriched-local-query-response/SKILL.md @@ -58,8 +58,28 @@ You are an expert in resolving location-based queries using the **A2UI framework * **Pins**: * `anchorMarker`: Use for the "main" focus (e.g., a hotel). * `markers`: Use for related results (e.g., surrounding restaurants). + * **POI Types (`placePrimaryType`)**: Determine `placePrimaryType` using the descriptions or categories in the tool response. If insufficient, infer it from the user prompt and place title. + Supported categories: + - `food_and_drink`: Restaurants, cafes, bars, bakeries, coffee shops, dining. + - `outdoor`: Parks, trails, gardens, nature reserves, beaches, + viewpoints. Never use for cities, towns, neighborhoods, regions, or + countries. + - `retail`: Stores, shops, boutiques, supermarkets, malls, markets. + - `gas_station`: Gas stations. + - `ev`: EV charging stations. + - `bank`: Banks, ATMs. + - `lodging`: Hotels, resorts, motels, hostels, B&Bs. + - `emergency`: Hospitals, urgent care, police, fire stations. + - `entertainment`: Theaters, museums, cinemas, stadiums, amusement parks, venues. + - `airport`: Airports. + - `parking`: Parking lots and garages. + - `generic`: Geopolitical / non-POI places (cities, towns, + neighborhoods, counties, states/regions, countries, postal codes, + street addresses) and services (salons, repair, dry cleaners, post + offices). Default fallback when ambiguous or not clearly matching + above categories. * **References**: Refer to items in the data model via `path` for dynamic content. -* **Child Components**: When using a Column or Row layout, ensure that each child component referenced in the `children` array is also included in the `surfaceUpdate` as its own component definition. +* **Child Components**: When using a Column or Row layout, ensure that each child component referenced in the `children` array is also included in the `surfaceUpdate` as its own component definition. Do NOT add unsupported properties like `spacing` or `gap` to `Column` or `Row` components. ### 4. Loading Data @@ -89,6 +109,10 @@ If you need to fetch data, make sure that you have a plan for how you are going * Only use a UI pattern when it adds material value to the content * Do not apply UI patterns in excess. If a response justifies multiple UI patterns, include only the patterns that are most valuable. Never include more than one GoogleMap in support of a single paragraph. * Display maps only when informative. For example, if a user asks if a hotel has a restaurant, displaying a map does not help answer their question. +* **Explicit Map Requests**: Whenever the user explicitly asks to see one or + more places on a map (e.g., "on a map", "in the same map", "show on map"), + you MUST include a `GoogleMap` component with markers for all resolved + locations (even if only a subset of requested locations resolved). ## Choosing UI Patterns Use the following logic to determine which UI component combinations to use: @@ -109,6 +133,10 @@ Use the following logic to determine which UI component combinations to use: | Context | Recommended UI | Data Requirements | | :---------------------- | :--------------------- | :------------------------ | +| **Named Places on Map**:| **Inline Map + List of | Resolved places as | +: User asks to display or : PlaceDetailsCompacts** : `markers` and cards; : +: compare specific named : : center at midpoint. : +: places on the same map. : : : | **Anchored Search**: | **Inline Map + List of | Pivot on `anchorMarker`. | : Distance/time : PlaceDetailsCompacts** : POIs as `markers`. DO NOT : : constraint to a center : : include a place card for : @@ -222,11 +250,9 @@ For the `PlaceDetailsCompact` component, you MUST include the following fields: * `placeId` -You MAY also include: - -* `orientation`: - - You MUST use `"vertical"` when there is only ONE `PlaceDetailsCompact` in the response to emphasize the place image. - - You MUST use `"horizontal"` when there are MULTIPLE `PlaceDetailsCompact` components (e.g., in a list) to keep the layout compact and save vertical space. +Always use `"horizontal"` for `orientation` on `PlaceDetailsCompact` (or omit +`orientation` so it defaults to `"horizontal"`), unless the user explicitly asks +for vertical place cards. **IMPORTANT:** ALWAYS follow the schema provided by the schema manager (passed in as part of the instruction prompt) as the source of truth for what fields are diff --git a/agent/python_agent/skills/local-search-template-response/SKILL.md b/agent/python_agent/skills/local-search-template-response/SKILL.md deleted file mode 100644 index c9fc751..0000000 --- a/agent/python_agent/skills/local-search-template-response/SKILL.md +++ /dev/null @@ -1,71 +0,0 @@ ---- -name: local-search-template-response -description: Extractor skill for local place search queries. Extracts location and list of places for template merging. ---- - -# Core Objective - -Extract structured parameters for local searches. You must call maps tools to -locate matching businesses/places, and populate the response fields. - -## Grounding & Tool-Calling Policy (CRITICAL) - -1. **DO NOT HALLUCINATE OR GUESS**: You are strictly forbidden from generating - coordinates (latitude, longitude) or Google Maps Place IDs from your - internal memory or training weights. -2. **MANDATORY TOOL CALLS**: You MUST call the `search_places` tool first to - find actual venues matching the user's query near the requested locations. -3. **EXACT MATCH**: Any place name, coordinates, or Place ID returned in your - final response MUST correspond exactly to the data returned by the - `search_places` tool call. - -## Multi-Step Location Resolution Policy (Anchored Search) - -If the user's query references a starting point, landmark, hotel, or specific -address as a geographical anchor (e.g., "Space Needle", "Hyatt hotel", "1600 -Amphitheatre Pkwy"): - -1. **Resolve Anchor Coordinates First**: Make a tool call to `search_places` - with the anchor name as `textQuery` to resolve its exact center coordinates. -2. **Execute Proximity Search (Pivot on Anchor)**: Use the resolved latitude - and longitude of the anchor as the center of a `locationBias` circle. Set - the `radiusMeters` parameter based on the user's query: - * **Explicit Distance**: If the query specifies a distance (e.g., "within - 5 miles", "2 km"), parse and convert it to meters (e.g., `8000` or - `2000`). - * **Implicit Walking**: If the query implies walking (e.g., "walk to", - "walking distance"), default to `1000` meters. - * **No Specific Constraint**: If no distance constraint is specified - (e.g., "near", "around"), omit the `radiusMeters` field to let the - search engine bias results dynamically around the center point. -3. **Handle Non-Geographic Filters**: Keep the search query focused on the core - category and key searchable features or amenities (e.g., "restaurant outdoor - seating", "cafe wifi"). Strip conversational filler words (e.g., "with", - "that has", "places offering") to maximize search relevance and matching - accuracy. -4. **Set Anchor Marker**: Populate the `anchor_marker` response parameter with - the resolved coordinates, name, and Place ID of the anchor location. - -## Handling Failures & Empty Results (CRITICAL) - -If search queries return empty results (`{}`) or fail: - -1. **Do NOT retry** repeatedly with alternative queries. -2. Immediately exit the tool-calling loop. -3. Return a user-friendly summary explaining that no results were found, and - leave the `places` list empty. - -## Output Fields - -You MUST populate all required fields in the output schema, and optionally the anchor marker if resolved: - -- **`summary`**: A detailed response summarizing the search results, following the **Conversational Text Style Guidelines** below. -- **`center_lat`**: Latitude of the center of results. Use the coordinates of - the resolved anchor location (or the average of the results if no anchor is - resolved). -- **`center_lng`**: Longitude of the center of results. Use the coordinates of - the resolved anchor location (or the average of the results if no anchor is - resolved). -- **`zoom`**: Recommended map zoom level. Default to 13. -- **`places`**: A list of places found (limit to max list size, e.g. 3). -- **`anchor_marker`**: (Optional) Pin details for the resolved starting/anchor location. diff --git a/agent/python_agent/streaming_request_handler.py b/agent/python_agent/streaming_request_handler.py new file mode 100644 index 0000000..fe25a35 --- /dev/null +++ b/agent/python_agent/streaming_request_handler.py @@ -0,0 +1,145 @@ +# 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 asyncio +import logging +from collections.abc import AsyncGenerator +from typing import Any, cast + +from a2a.server.request_handlers import DefaultRequestHandler +from a2a.server.context import ServerCallContext +from a2a.types import ( + Message, + MessageSendParams, + Task, + TaskArtifactUpdateEvent, + TaskStatusUpdateEvent, + TaskState, +) +from a2a.server.events import EventQueue +from a2a.server.tasks import TaskManager +from a2a.utils.errors import ServerError +from a2a.utils.errors import TaskNotFoundError, InvalidParamsError + +logger = logging.getLogger(__name__) + +class QueueBridge(EventQueue): + def __init__(self, queue: asyncio.Queue): + super().__init__() + self._queue = queue + self._is_closed = False + + async def enqueue_event( + self, + event: Message | Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent, + ) -> None: + if self._is_closed: + raise RuntimeError("Queue is closed") + await self._queue.put(event) + + async def close(self, immediate: bool = False) -> None: + if self._is_closed: + return + self._is_closed = True + await self._queue.put(None) + + def is_closed(self) -> bool: + return self._is_closed + +class StreamingRequestHandler(DefaultRequestHandler): + async def on_message_send_stream( + self, + params: MessageSendParams, + context: ServerCallContext | None = None, + ) -> AsyncGenerator[Message | Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent, None]: + task_manager = TaskManager( + task_id=params.message.task_id, + context_id=params.message.context_id, + task_store=self.task_store, + initial_message=params.message, + context=context, + ) + task: Task | None = await task_manager.get_task() + + TERMINAL_TASK_STATES = { + TaskState.completed, + TaskState.canceled, + TaskState.failed, + TaskState.rejected, + } + + if task: + if task.status.state in TERMINAL_TASK_STATES: + raise ServerError( + error=InvalidParamsError( + message=f'Task {task.id} is in terminal state: {task.status.state.value}' + ) + ) + task = task_manager.update_with_message(params.message, task) + elif params.message.task_id: + raise ServerError( + error=TaskNotFoundError( + message=f'Task {params.message.task_id} was specified but does not exist' + ) + ) + + request_context = await self._request_context_builder.build( + params=params, + task_id=task.id if task else None, + context_id=params.message.context_id, + task=task, + context=context, + ) + + task_id = cast('str', request_context.task_id) + + if ( + self._push_config_store + and params.configuration + and params.configuration.push_notification_config + ): + await self._push_config_store.set_info( + task_id, params.configuration.push_notification_config + ) + + queue = asyncio.Queue() + queue_bridge = QueueBridge(queue) + + async def run_producer(): + try: + await self.agent_executor.execute(request_context, queue_bridge) + except Exception as e: + logger.exception("Error in Streaming agent execution producer task") + await queue.put(e) + finally: + await queue_bridge.close() + + producer_task = asyncio.create_task(run_producer()) + producer_task.set_name(f"streaming_producer:{task_id}") + + try: + while True: + event = await queue.get() + if event is None: + break + if isinstance(event, Exception): + raise event + yield event + finally: + if not producer_task.done(): + producer_task.cancel() + try: + await producer_task + except asyncio.CancelledError: + pass \ No newline at end of file diff --git a/agent/python_agent/template_registry.py b/agent/python_agent/template_registry.py new file mode 100644 index 0000000..f6c7eeb --- /dev/null +++ b/agent/python_agent/template_registry.py @@ -0,0 +1,473 @@ +# Copyright 2026 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. + +"""Dynamic Template Registry and Router Instruction Compiler for MAUI Agent.""" + +import dataclasses +import enum +import importlib +import importlib.resources +from importlib.resources import abc as resources_abc +import json +import logging +import re +from typing import Any + +import pydantic + +INTENT_OTHER_SPATIAL = "OTHER_SPATIAL" +INTENT_TEXT_ONLY = "TEXT_ONLY" + + +class RouterClassification(pydantic.BaseModel): + """Classification output schema for the query intent router.""" + + intent: str = pydantic.Field( + description="The classified intent archetype name." + ) + query: str = pydantic.Field(description="The cleaned user query") + + +class FewShotExample(pydantic.BaseModel): + """A few-shot query-intent example for router instructions.""" + + model_config = pydantic.ConfigDict(extra="ignore") + + user_query: str + intent: str | None = None + query: str | None = None + + def __getitem__(self, key: str) -> Any: + return getattr(self, key) + + def get(self, key: str, default: Any = None) -> Any: + return getattr(self, key, default) + + +class TemplateManifest(pydantic.BaseModel): + """Strongly-typed metadata schema parsed from a bundle's manifest.json.""" + + model_config = pydantic.ConfigDict(extra="ignore") + + intent: str = "" + surface_prefix: str | None = None + archetype: str = "" + few_shot_examples: list[FewShotExample] = pydantic.Field(default_factory=list) + tool_description: str = "" + empty_fallback_text: str = "" + + +@dataclasses.dataclass(frozen=True) +class RouterIntentSpec: + """Specification for a router intent archetype and its few-shot examples.""" + + intent: str + archetype: str + few_shot_examples: list[FewShotExample] = dataclasses.field( + default_factory=list + ) + + +_BUILTIN_FALLBACK_ROUTER_INTENTS: tuple[RouterIntentSpec, ...] = ( + RouterIntentSpec( + intent=INTENT_OTHER_SPATIAL, + archetype=( + "Queries requiring polygon boundaries, raw coordinate plotting," + " or complex navigation combining routes with secondary overlays" + " (e.g., weather, air quality forecasts, displaying all available" + " charging stations along a route). Note: Requests to show one or" + " more specific named places, buildings, or addresses on a map" + " belong to LOCAL_SEARCH, not OTHER_SPATIAL." + ), + few_shot_examples=[ + FewShotExample( + user_query=( + "Walking route from Central Park to Times Square, but show" + " coffee shops and rain forecasts along the way" + ), + intent=INTENT_OTHER_SPATIAL, + query=( + "walking route from Central Park to Times Square with" + " coffee shop stops and rain forecast" + ), + ), + FewShotExample( + user_query=( + "show the boundary of Yosemite National Park on the map" + ), + intent=INTENT_OTHER_SPATIAL, + query="boundary of Yosemite National Park", + ), + ], + ), + RouterIntentSpec( + intent=INTENT_TEXT_ONLY, + archetype=( + "General information retrieval, questions, or requests for data" + " associated with locations that can be answered fully with text" + " without requiring a map interface." + ), + few_shot_examples=[], + ), +) + + +logger = logging.getLogger(__name__) + +# Bundle names must be lowercase snake_case starting with a letter so they are +# valid simultaneously as ADK tool names (`render_<name>_template`), Python +# package identifiers (`templates.<name>.schema`), and ADK skill names. +_VALID_BUNDLE_NAME_RE = re.compile(r"^[a-z][a-z0-9_]*$") + + +@dataclasses.dataclass(frozen=True) +class TemplateBundle: + """Metadata and assets representing a self-contained template intent.""" + + name: str + intent: str + surface_prefix: str + archetype: str + few_shot_examples: list[FewShotExample | dict[str, Any]] + schema_cls: type[pydantic.BaseModel] + + # Use Traversable for compatibility with both package data and disk paths. + layout_path: resources_abc.Traversable + skill_dir: resources_abc.Traversable + bundle_dir: resources_abc.Traversable + tool_description: str = "" + empty_fallback_text: str = "" + skill_instructions: str = "" + + def __post_init__(self) -> None: + if not _VALID_BUNDLE_NAME_RE.match(self.name): + raise ValueError( + f"Invalid bundle name '{self.name}': must be lowercase snake_case" + f" matching {_VALID_BUNDLE_NAME_RE.pattern} so derived tool names" + " are valid identifiers." + ) + + @property + def render_tool_name(self) -> str: + """The function call name of the ADK tool generated to render this bundle.""" + return f"render_{self.name}_template" + + @property + def effective_tool_description(self) -> str: + """Description the model sees for this bundle's render tool. + + A bundle that does not write its own gets one derived from the archetype + it already declares for the router, so the two cannot drift apart. + """ + if self.tool_description: + return self.tool_description + return ( + f"Renders the {self.name} UI component. Use it when the request" + f" matches: {self.archetype}" + ) + + +def _read_skill_instructions(skill_dir: resources_abc.Traversable) -> str: + """Reads instructions body from SKILL.md directly, stripping frontmatter.""" + skill_file = skill_dir / "SKILL.md" + if not skill_file.is_file(): + return "" + content = skill_file.read_text(encoding="utf-8") + if content.startswith("---"): + parts = content.split("---", 2) + if len(parts) >= 3: + return parts[2].strip() + return content.strip() + + +_TEMPLATES_PACKAGE = ( + "templates" +) + + +def _resolve_schema(bundle_name: str) -> type[pydantic.BaseModel]: + """Imports the ExtractorSchema class declared by a built-in bundle. + + Dynamically resolves the schema module from the templates package + (e.g., templates.<bundle_name>.schema) without requiring manual allowlist + registration. + + Args: + bundle_name: Directory name of the bundle under templates/. + + Returns: + The bundle's ExtractorSchema class. + + Raises: + KeyError: The bundle has no importable schema module. + TypeError: The module does not export an ExtractorSchema that subclasses + pydantic.BaseModel. + """ + module_name = f"{_TEMPLATES_PACKAGE}.{bundle_name}.schema" + try: + module = importlib.import_module(module_name) + except ImportError as e: + raise KeyError( + f"Bundle '{bundle_name}' has no importable schema module at" + f" '{module_name}': {e}. Ensure templates/{bundle_name}/schema.py" + " exists, or build a TemplateBundle and pass it to" + " TemplateRegistry.register_bundle()." + ) from e + schema_cls = getattr(module, "ExtractorSchema", None) + if not ( + isinstance(schema_cls, type) + and issubclass(schema_cls, pydantic.BaseModel) + ): + raise TypeError( + f"{module_name} must export 'ExtractorSchema' subclassing" + " pydantic.BaseModel" + ) + return schema_cls + + +class TemplateRegistry: + """Registry for discovering, caching, and querying template bundles.""" + + def __init__(self) -> None: + self._bundles: dict[str, TemplateBundle] = {} + self._bundles_by_intent: dict[str, TemplateBundle] = {} + self._cached_router_classification: type[RouterClassification] | None = None + + # Use importlib.resources.files for compatibility with both package data and + # disk paths. + self._templates_dir = importlib.resources.files(__package__).joinpath( + "templates" + ) + + self._discover_bundles() + + def _discover_bundles(self) -> None: + """Loads every built-in bundle directory under templates/.""" + if not self._templates_dir.is_dir(): + logger.warning( + "Templates directory '%s' does not exist.", self._templates_dir + ) + return + + # Sorted by name: zip-backed `Traversable`s do not define an ordering. + for item in sorted(self._templates_dir.iterdir(), key=lambda t: t.name): + if item.is_dir() and (item / "manifest.json").is_file(): + self._load_bundle(item) + + def _load_bundle(self, bundle_dir: resources_abc.Traversable) -> None: + """Loads one built-in bundle directory. + + Args: + bundle_dir: Directory holding the bundle's manifest, layout, and schema. + + Raises: + ValueError: The manifest is unreadable or malformed. + KeyError: The bundle has no importable schema module. + TypeError: The bundle's schema module exports the wrong thing. + ImportError: The bundle's schema module does not import. + """ + manifest_path = bundle_dir / "manifest.json" + try: + manifest_json = manifest_path.read_text(encoding="utf-8") + manifest = TemplateManifest.model_validate_json(manifest_json) + except (OSError, pydantic.ValidationError, ValueError) as e: + # Built-in bundles ship inside this build target, so an unreadable + # manifest is a build defect. Booting without the intent would degrade + # every query that needed it, with nothing but a log line to say why. + raise ValueError(f"Unreadable bundle manifest at {manifest_path}") from e + + name = bundle_dir.name + intent = manifest.intent or name.upper() + + surface_prefix = manifest.surface_prefix or f"{name}-surface" + schema_cls = _resolve_schema(name) + + bundle = TemplateBundle( + name=name, + intent=intent, + surface_prefix=surface_prefix, + archetype=manifest.archetype, + few_shot_examples=manifest.few_shot_examples, + schema_cls=schema_cls, + layout_path=bundle_dir / "layout.json", + skill_dir=bundle_dir, + bundle_dir=bundle_dir, + tool_description=manifest.tool_description, + empty_fallback_text=manifest.empty_fallback_text, + skill_instructions=_read_skill_instructions(bundle_dir), + ) + self.register_bundle(bundle) + + def register_bundle(self, bundle: TemplateBundle) -> None: + """Registers a constructed TemplateBundle. + + The caller owns importing the bundle's schema class and constructing the + TemplateBundle. The registry never loads code from a path. + + Args: + bundle: The bundle to register. It displaces any bundle already holding + its name or its intent, and the collision is logged. A displaced bundle + is removed from both indexes so it cannot be routed to or advertised to + the router afterwards. + """ + displaced = self._bundles.get(bundle.name) + if displaced is not None: + logger.warning( + "Bundle '%s' is already registered; replacing it.", bundle.name + ) + self._bundles_by_intent.pop(displaced.intent, None) + + claimant = self._bundles_by_intent.get(bundle.intent) + if claimant is not None and claimant.name != bundle.name: + logger.warning( + "Intent '%s' was claimed by bundle '%s'; bundle '%s' now owns it." + " Bundle '%s' is dropped, since an unroutable bundle would still" + " show up in the router prompt.", + bundle.intent, + claimant.name, + bundle.name, + claimant.name, + ) + self._bundles.pop(claimant.name, None) + + self._bundles[bundle.name] = bundle + self._bundles_by_intent[bundle.intent] = bundle + self._cached_router_classification = None + logger.info( + "Registered template bundle: %s (intent: %s)", + bundle.name, + bundle.intent, + ) + + def get_bundle(self, name: str) -> TemplateBundle | None: + """Looks up a bundle by its primary directory/bundle name (e.g., 'local_search').""" + return self._bundles.get(name) + + def get_bundle_by_intent( + self, intent: str | enum.Enum + ) -> TemplateBundle | None: + """Looks up a bundle by its intent (e.g., 'LOCAL_SEARCH' or an enum member).""" + key = intent.value if hasattr(intent, "value") else str(intent) + return self._bundles_by_intent.get(key) + + @property + def bundles(self) -> dict[str, TemplateBundle]: + return dict(self._bundles) + + def list_bundles(self) -> list[TemplateBundle]: + return list(self._bundles.values()) + + @property + def supported_intents(self) -> set[str]: + return set(self._bundles_by_intent.keys()) + + @property + def router_classification(self) -> type[RouterClassification]: + """Dynamically constructs a constrained response schema enforcing valid intents.""" + if self._cached_router_classification is not None: + return self._cached_router_classification + + all_intents = sorted( + self.supported_intents | {INTENT_OTHER_SPATIAL, INTENT_TEXT_ONLY} + ) + dynamic_intent_enum = enum.Enum( + "RouterIntent", + {intent: intent for intent in all_intents}, + type=str, + ) + self._cached_router_classification = pydantic.create_model( + "RouterClassification", + __base__=RouterClassification, + intent=( + dynamic_intent_enum, + pydantic.Field(description="The classified intent archetype name."), + ), + query=(str, pydantic.Field(description="The cleaned user query")), + ) + return self._cached_router_classification + + def compile_router_instruction(self) -> str: + """Compiles the dynamic router instruction from the registered bundles. + + Bundles appear in the prompt ordered by intent, so the instruction text + depends only on which bundles are registered and not on the order they + were discovered or registered in. Two processes holding the same bundles + therefore send the model the same bytes. + """ + ordered_bundles = sorted( + self._bundles.values(), key=lambda bundle: bundle.intent + ) + all_intent_specs: list[TemplateBundle | RouterIntentSpec] = [ + *ordered_bundles, + *_BUILTIN_FALLBACK_ROUTER_INTENTS, + ] + + archetype_lines = [ + f"- **{spec.intent}**: {spec.archetype}" for spec in all_intent_specs + ] + + registered_intents_str = ( + " or ".join(sorted(self._bundles_by_intent.keys())) + or "registered templates" + ) + + examples_blocks = [] + example_index = 1 + + for spec in all_intent_specs: + for example in spec.few_shot_examples: + user_query = ( + example.user_query + if isinstance(example, FewShotExample) + else example.get("user_query", "") + ) + intent = ( + (example.intent or spec.intent) + if isinstance(example, FewShotExample) + else example.get("intent", spec.intent) + ) + extracted_query = ( + (example.query or user_query) + if isinstance(example, FewShotExample) + else example.get("query", user_query) + ) + output_json = json.dumps( + {"intent": intent, "query": extracted_query}, indent=2 + ) + examples_blocks.append( + f"### Example {example_index}\n" + f'**User Query:** "{user_query}"\n' + f"**Output:**\n{output_json}\n" + ) + example_index += 1 + + archetypes_str = "\n".join(archetype_lines) + examples_str = "\n".join(examples_blocks) + + return ( + "## Role\nYou are an expert query intent router.\n\n## Task" + " Definition\nAnalyze a user's input and classify it into the most" + " appropriate intent category based on the structural complexity and" + " data requirements of the request.\n\n## Intent" + f" Archetypes\n{archetypes_str}\n\n## Classification Policy:" + " Conservative Routing\nIf a query satisfies the structure of a base" + f" intent (such as {registered_intents_str}) but also includes any" + " Auxiliary Data (e.g., weather forecasts, environmental factors), you" + " MUST promote the classification to OTHER_SPATIAL.\nNote: Simple" + ' multi-stop routes or routes with specified waypoints (e.g., "A to B' + ' via C") should be classified as DIRECTIONS, unless they require' + ' searching for stops along the way (e.g., "find coffee shops along' + ' the route") which requires LOCAL_SEARCH overlays and should be' + f" promoted to OTHER_SPATIAL.\n\n## Few-shot Examples\n\n{examples_str}" + ) diff --git a/agent/python_agent/template_tool.py b/agent/python_agent/template_tool.py new file mode 100644 index 0000000..7deeaaa --- /dev/null +++ b/agent/python_agent/template_tool.py @@ -0,0 +1,410 @@ +# Copyright 2026 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 +# +# http://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. + +"""ADK Tools for MAUI template population and rendering. + +This file contains a set of tools for rendering the MAUI A2UI templates. +The tools are used by the MAUI agent to render the templates based on the +user's query and the agent's extracted information. + +The currently supported templates are: +- Local Search: Used to show a list of local places and a map. +- Directions: Used to show a route on a map. +- Text-only: Used to render a text-only response inside an A2UI surface. + +Tools are built dynamically based on their Pydantic schema to ensure +type safety and accurate function declarations. +""" + +from __future__ import annotations + +from collections.abc import Callable +import copy +import inspect +import logging +import pathlib +import time +from typing import Any, Optional, Union +import uuid + +from a2a.types import Part +from google.adk.agents.readonly_context import ReadonlyContext +from google.adk.tools._automatic_function_calling_util import build_function_declaration +from google.adk.tools.base_tool import BaseTool +from google.adk.tools.base_toolset import BaseToolset, ToolPredicate +from google.adk.tools.set_model_response_tool import _merge_json_schema_descriptions +from google.adk.tools.tool_context import ToolContext +from google.genai import types +import pydantic + +from a2ui.schema.manager import A2uiSchemaManager +from merger import DEFAULT_CATALOG_ID +from merger import merge_template +from template_registry import ( + TemplateBundle, + TemplateRegistry, +) + +logger = logging.getLogger(__name__) + +STATE_RENDERED_A2UI_PARTS = "rendered_a2ui_parts" +STATE_RENDERED_A2UI_DATA = "rendered_a2ui_data" + +# Opt-in replacement for the session-state handoff above, for hosts that +# deliver A2UI through their own channel. Called once per successful render +# with the tool context, the merged A2UI messages as plain dicts in send order, +# and the validated template arguments. When set, the tool writes neither +# `STATE_RENDERED_A2UI_PARTS` nor `STATE_RENDERED_A2UI_DATA`. +RenderedMessagesHook = Callable[ + [ToolContext, list[dict[str, Any]], dict[str, Any]], None +] + + +class BaseTemplateTool(BaseTool): + """Base class for ADK tools that populate and render A2UI templates.""" + + name: str = "" + description: str = "" + template_name: str = "" + schema_class: type[pydantic.BaseModel] | None = None + + def __init__( + self, + *, + name: str | None = None, + description: str | None = None, + template_name: str | None = None, + schema_class: type[pydantic.BaseModel] | None = None, + schema_manager: A2uiSchemaManager | None = None, + max_list_size: int = 5, + surface_id_prefix: str | None = None, + catalog_id: str = DEFAULT_CATALOG_ID, + on_rendered: RenderedMessagesHook | None = None, + ) -> None: + tool_name = name or self.name + tool_desc = description or self.description + super().__init__(name=tool_name, description=tool_desc) + self.template_name = template_name or self.template_name + self.schema_class = ( + schema_class if schema_class is not None else self.schema_class + ) + self.schema_manager = schema_manager + self.max_list_size = max_list_size + self.catalog_id = catalog_id + self.on_rendered = on_rendered + self.surface_id_prefix = ( + surface_id_prefix or f"{self.template_name}-surface" + ) + self._func = self._build_handler_func() + + @classmethod + def from_bundle( + cls, + bundle: TemplateBundle, + *, + schema_manager: A2uiSchemaManager | None = None, + max_list_size: int = 5, + catalog_id: str = DEFAULT_CATALOG_ID, + on_rendered: RenderedMessagesHook | None = None, + **kwargs: Any, + ) -> "BaseTemplateTool": + """Constructs a BaseTemplateTool dynamically from a TemplateBundle.""" + params = { + "name": bundle.render_tool_name, + "description": bundle.effective_tool_description, + "template_name": bundle.name, + "schema_class": bundle.schema_cls, + "surface_id_prefix": bundle.surface_prefix, + "schema_manager": schema_manager, + "max_list_size": max_list_size, + "catalog_id": catalog_id, + "on_rendered": on_rendered, + } + params.update(kwargs) + return cls(**params) + + def _build_handler_func(self) -> Any: + """Builds the callable signature used for FunctionDeclaration generation.""" + if self.schema_class is not None: + schema_fields = self.schema_class.model_fields + params = [] + for field_name, field_info in schema_fields.items(): + param = inspect.Parameter( + field_name, + inspect.Parameter.KEYWORD_ONLY, + annotation=field_info.annotation, + default=( + inspect.Parameter.empty + if field_info.is_required() + else field_info.get_default(call_default_factory=True) + ), + ) + params.append(param) + + def dynamic_tool_func(**kwargs: Any) -> str: + del kwargs + return f"Rendered {self.template_name} template." + + new_sig = inspect.Signature(parameters=params) + setattr(dynamic_tool_func, "__signature__", new_sig) + setattr(dynamic_tool_func, "__name__", self.name) + setattr(dynamic_tool_func, "__doc__", self.description) + return dynamic_tool_func + else: + + def text_only_tool_func(text: str) -> str: + """Render a text-only UI response.""" + del text + return f"Rendered {self.template_name} template." + + setattr(text_only_tool_func, "__name__", self.name) + setattr(text_only_tool_func, "__doc__", self.description) + return text_only_tool_func + + def _preserve_schema_descriptions( + self, function_decl: types.FunctionDeclaration + ) -> None: + """Restores field descriptions from Pydantic schema onto FunctionDeclaration.""" + if self.schema_class is not None: + source_schema = self.schema_class.model_json_schema() + if function_decl.parameters_json_schema is not None: + _merge_json_schema_descriptions( + function_decl.parameters_json_schema, source_schema + ) + elif function_decl.parameters is not None: + from google.adk.tools.set_model_response_tool import ( # pylint: disable=g-import-not-at-top + _apply_descriptions_to_schema_properties, + ) + + _apply_descriptions_to_schema_properties( + function_decl.parameters.properties, + self.schema_class.model_fields, + ) + + def _get_declaration(self) -> Optional[types.FunctionDeclaration]: + """Gets OpenAPI FunctionDeclaration specification for this tool.""" + function_decl = types.FunctionDeclaration.model_validate( + build_function_declaration( + func=self._func, + ignore_params=[], + variant=self._api_variant, + ) + ) + self._preserve_schema_descriptions(function_decl) + return function_decl + + async def run_async( + self, *, args: dict[str, Any], tool_context: ToolContext + ) -> dict[str, Any]: + """Executes the template tool: validates args, merges template, and saves A2UI parts.""" + start_time = time.perf_counter() + logger.info("--- TEMPLATE_TOOL: Invoked '%s' ---", self.name) + logger.info(" Tool: %s (template: %s)", self.name, self.template_name) + logger.info(" Parameters: %s", args) + validated_data = copy.deepcopy(args) + + # 1. Validate arguments against Pydantic schema + if self.schema_class is not None: + try: + model_instance = self.schema_class.model_validate(args) + validated_data = model_instance.model_dump(exclude_none=True) + except pydantic.ValidationError as e: + elapsed_ms = (time.perf_counter() - start_time) * 1000 + logger.warning( + "--- TEMPLATE_TOOL: Validation failed for '%s' in %.2f ms: %s ---", + self.name, + elapsed_ms, + e, + ) + return { + "error": ( + f"Validation failed for tool '{self.name}': {e}. " + "Please fix the parameters and call the tool again." + ) + } + + # 2. Ensure unique surface_id + if not validated_data.get("surface_id"): + short_id = uuid.uuid4().hex[:8] + validated_data["surface_id"] = f"{self.surface_id_prefix}-{short_id}" + + # 3. Merge template + try: + merged_actions = merge_template( + self.template_name, + validated_data, + max_list_size=self.max_list_size, + catalog_id=self.catalog_id, + ) + except Exception as e: # pylint: disable=broad-exception-caught + elapsed_ms = (time.perf_counter() - start_time) * 1000 + logger.warning( + "--- TEMPLATE_TOOL: Failed to merge template '%s' in %.2f ms: %s ---", + self.template_name, + elapsed_ms, + e, + ) + return {"error": f"Failed to merge template '{self.template_name}': {e}"} + + # 4. Catalog schema validation + if self.schema_manager: + selected_catalog = self.schema_manager.get_selected_catalog() + if selected_catalog: + try: + selected_catalog.validator.validate(merged_actions) + except Exception as e: # pylint: disable=broad-exception-caught + elapsed_ms = (time.perf_counter() - start_time) * 1000 + logger.warning( + "--- TEMPLATE_TOOL: Catalog validation failed for '%s' in %.2f" + " ms: %s ---", + self.template_name, + elapsed_ms, + e, + ) + return { + "error": ( + f"A2UI catalog schema validation failed: {e}. " + "Please fix the parameters and retry." + ) + } + + # 5. Hand the messages to the host's hook, or by default convert them to + # A2A Parts and persist them to session state. + if self.on_rendered is not None: + try: + self.on_rendered(tool_context, merged_actions, validated_data) + except Exception as e: # pylint: disable=broad-exception-caught + logger.exception( + "--- TEMPLATE_TOOL: on_rendered hook failed for '%s' ---", + self.template_name, + ) + return { + "error": ( + f"Failed to deliver rendered template '{self.template_name}':" + f" {e}" + ) + } + else: + # Imported here, not at module level: a2ui.a2a.parts is built on the A2A + # 0.3 SDK, so importing it eagerly breaks every importer of this module + # running on A2A 1.0 -- including callers that pass `on_rendered` and + # never reach this branch. + from a2ui.a2a.parts import create_a2ui_part # pylint: disable=g-import-not-at-top + + rendered_parts: list[Part] = [ + create_a2ui_part(action) for action in merged_actions + ] + if tool_context and getattr(tool_context, "state", None) is not None: + tool_context.state[STATE_RENDERED_A2UI_PARTS] = rendered_parts + tool_context.state[STATE_RENDERED_A2UI_DATA] = validated_data + + elapsed_ms = (time.perf_counter() - start_time) * 1000 + logger.info( + "--- TEMPLATE_TOOL: Successfully rendered '%s' (surface_id: %s) in %.2f" + " ms (%d parts) ---", + self.template_name, + validated_data["surface_id"], + elapsed_ms, + len(merged_actions), + ) + + return { + "status": "success", + "surface_id": validated_data["surface_id"], + "template": self.template_name, + "latency_ms": round(elapsed_ms, 2), + "message": f"Successfully rendered {self.template_name} UI interface.", + } + + +class RenderTextOnlyTemplateTool(BaseTemplateTool): + """ADK Tool that renders a text-only response inside an A2UI surface container.""" + + name: str = "render_text_only_template" + description: str = ( + "Renders a text response formatted inside an A2UI surface container." + ) + template_name: str = "text_only" + schema_class: type[pydantic.BaseModel] | None = None + + def __init__( + self, + *, + name: str | None = None, + description: str | None = None, + schema_manager: A2uiSchemaManager | None = None, + surface_id_prefix: str = "text-only-surface", + catalog_id: str = DEFAULT_CATALOG_ID, + on_rendered: RenderedMessagesHook | None = None, + ) -> None: + super().__init__( + name=name, + description=description, + schema_manager=schema_manager, + max_list_size=1, + surface_id_prefix=surface_id_prefix, + catalog_id=catalog_id, + on_rendered=on_rendered, + ) + + +class TemplateToolset(BaseToolset): + """Toolset bundling all A2UI template population tools.""" + + def __init__( + self, + *, + template_registry: TemplateRegistry | None = None, + schema_manager: A2uiSchemaManager | None = None, + max_list_size: int = 5, + tool_filter: Optional[Union[ToolPredicate, list[str]]] = None, + tool_name_prefix: Optional[str] = None, + catalog_id: str = DEFAULT_CATALOG_ID, + on_rendered: RenderedMessagesHook | None = None, + ) -> None: + super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix) + self.schema_manager = schema_manager + self.max_list_size = max_list_size + self.catalog_id = catalog_id + self.on_rendered = on_rendered + self.template_registry = template_registry or TemplateRegistry() + self._tools: list[BaseTool] = [] + for bundle in self.template_registry.bundles.values(): + self._tools.append( + BaseTemplateTool.from_bundle( + bundle, + schema_manager=self.schema_manager, + max_list_size=self.max_list_size, + catalog_id=self.catalog_id, + on_rendered=self.on_rendered, + ) + ) + self._tools.append( + RenderTextOnlyTemplateTool( + schema_manager=self.schema_manager, + catalog_id=self.catalog_id, + on_rendered=self.on_rendered, + ) + ) + + async def get_tools( + self, + readonly_context: Optional[ReadonlyContext] = None, + ) -> list[BaseTool]: + """Returns the template tools exposed by this toolset.""" + return [ + tool + for tool in self._tools + if self._is_tool_selected(tool, readonly_context) + ] diff --git a/agent/python_agent/templates/__init__.py b/agent/python_agent/templates/__init__.py new file mode 100644 index 0000000..6ca9e40 --- /dev/null +++ b/agent/python_agent/templates/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 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. + +"""Template bundle packages and shared schemas.""" diff --git a/agent/python_agent/templates/common.py b/agent/python_agent/templates/common.py new file mode 100644 index 0000000..7162a0e --- /dev/null +++ b/agent/python_agent/templates/common.py @@ -0,0 +1,88 @@ +# Copyright 2026 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. + +"""Shared schema models and primitive types for MAUI template bundles.""" + +from typing import Any, Literal + +import pydantic + +PlacePrimaryType = Literal[ + "food_and_drink", + "outdoor", + "retail", + "gas_station", + "ev", + "bank", + "lodging", + "emergency", + "entertainment", + "airport", + "parking", + "generic", +] + + +class Pin(pydantic.BaseModel): + """Representation of a Map Pin.""" + + lat: float = pydantic.Field(description="Latitude coordinate") + lng: float = pydantic.Field(description="Longitude coordinate") + label: str = pydantic.Field( + description=( + "Descriptive display label string (e.g. name of address, business, or" + " landmark)" + ) + ) + # Note: Using camelCase field name to match frontend A2UI requirements. + placeId: str | None = pydantic.Field( # pylint: disable=invalid-name + default=None, description="Optional Google Maps Place ID" + ) + placePrimaryType: PlacePrimaryType | None = pydantic.Field( # pylint: disable=invalid-name + default=None, + description=( + "Optional primary POI category type string. Use 'generic' for" + " geopolitical or non-POI places (cities, towns, neighborhoods," + " regions, countries, addresses)." + ), + ) + + @pydantic.model_validator(mode="before") + @classmethod + def normalize_label(cls, data: Any) -> Any: + """Normalizes the pin label. + + If 'label' is missing but 'name' is present, copies 'name' to 'label'. + If 'label' is still empty, defaults to 'Location' to ensure + the UI always has a valid string to render for the marker (avoiding raw + Place IDs). + + Args: + data: The input dictionary before validation. + + Returns: + The normalized dictionary. + """ + if isinstance(data, dict): + if "label" not in data and "name" in data: + data["label"] = data["name"] + if not data.get("label"): + data["label"] = "Location" + return data + + +__all__ = [ + "Pin", + "PlacePrimaryType", +] diff --git a/agent/python_agent/skills/directions-template-response/SKILL.md b/agent/python_agent/templates/directions/SKILL.md similarity index 52% rename from agent/python_agent/skills/directions-template-response/SKILL.md rename to agent/python_agent/templates/directions/SKILL.md index 536c25e..8bb97c3 100644 --- a/agent/python_agent/skills/directions-template-response/SKILL.md +++ b/agent/python_agent/templates/directions/SKILL.md @@ -1,9 +1,11 @@ --- -name: directions-template-response +name: directions description: Extractor skill for directions and routing queries. Extracts routing details for template merging. --- -# Directions Template Response Skill +# Directions Guidelines + +## Directions Template Response Skill ## Core Objective @@ -20,8 +22,8 @@ If the user's query requests a scenic bypass or detour: 2. **Compute Route Segments (Parallel Routing)**: Concurrently compute routes for all sequential legs connecting the resolved stops (Origin -> Waypoint, Waypoint -> Destination). -3. **Dispatch Response**: Call `set_model_response` with the compiled routes - and pins. +3. **Dispatch Response**: Call `render_directions_template` with the compiled + routes and pins. ## Step-by-Step Workflow @@ -73,19 +75,23 @@ If the user's query requests a scenic bypass or detour: -122.4}}}`). Do **NOT** pass `latLng` directly as a root key inside `origin` or `destination` (e.g. do not call `compute_routes(origin={"placeId": "...", "latLng": ...})`). + * **GROUNDED ROUTING CONSTRAINT**: NEVER use model knowledge to assume + roads used or live traffic. Always rely only on data from + `compute_routes`. * Verify route availability for requested `travel_mode`. * **CONSTRUCT THE ROUTES ARRAY**: You MUST compile the computed segments - into the `routes` array of the final `set_model_response` payload. The - array must contain all segments sequentially (e.g. `[{"origin": Origin, - "destination": Waypoint 1}, {"origin": Waypoint 1, "destination": - Destination}]`). Do NOT omit the `routes` array or leave it empty if you - successfully computed routes. + into the `routes` array of the final `render_directions_template` + payload. The array must contain all segments sequentially (e.g. + `[{"origin": Origin, "destination": Waypoint 1}, {"origin": Waypoint 1, + "destination": Destination}]`). Do NOT omit the `routes` array or leave + it empty if you successfully computed routes. * **MANDATORY TRAVEL MODE IN DISPATCH**: `travel_mode` is REQUIRED and - must NEVER be omitted in `set_model_response`. Always supply the - normalized mode string (`driving`, `walking`, `transit`, or `bicycling`). - * Call `set_model_response` with `DirectionsExtractorSchema` parameters - (`summary`, `center_lat`, `center_lng`, `zoom`, `routes`, - `travel_mode`). + must NEVER be omitted in `render_directions_template`. Always supply the + normalized mode string (`driving`, `walking`, `transit`, or + `bicycling`). + * Call `render_directions_template` with `DirectionsExtractorSchema` + parameters (`heading`, `summary`, `center_lat`, `center_lng`, `zoom`, + `routes`, `travel_mode`). ## Handling Routing Failures & Regional Limitations (CRITICAL) @@ -103,7 +109,19 @@ or fails: You MUST populate all required fields in the output schema: -- **`summary`**: A detailed response summarizing the travel directions, following the **Conversational Text Style Guidelines** below. +- **`heading`**: (REQUIRED) A concise, constraint-confirming primary heading + for the response. Plain text only (e.g., 'Walking route from Seattle Center + to Pike Place Market', 'Driving directions to JFK Airport'). Use sentence + case; do NOT include markdown hashtags or conversational filler. +- **`summary`**: (REQUIRED) A natural, direct resolution of the route prompt + (e.g. 'Driving from [Origin] to [Destination] takes about 19 minutes (14 + miles).', 'Walking from Seattle Center to Pike Place Market takes about 20 + minutes (1 mile).'). Describe distance using units appropriate to the + location (miles vs. km). For driving and public transit modes, always round + distance to a whole number. NEVER describe time in seconds or decimals. + Always round seconds to the nearest minute. If it rounds to 0 minutes, + describe it as "less than a minute". Always describe time as + approximate (e.g. about, around, approximately). - **`center_lat`**: Latitude of the center of the route map. - **`center_lng`**: Longitude of the center of the route map. - **`zoom`**: Recommended map zoom level. Default to 12. @@ -114,26 +132,53 @@ You MUST populate all required fields in the output schema: ## Examples ### Example 1: Driving Route -User Query: "Directions from San Francisco to San Jose by car" -Tool Call: -`set_model_response(summary="Driving from San Francisco to San Jose takes about 50 minutes via US-101 S.", center_lat=37.55, center_lng=-122.15, zoom=10, routes=[{"origin": {"lat": 37.7749, "lng": -122.4194, "label": "San Francisco", "placeId": "ChIJIQBpAG2ahYAR_6128GcTUEo"}, "destination": {"lat": 37.3382, "lng": -121.8863, "label": "San Jose", "placeId": "ChIJ9T_nxcC1j4ARmMo7S4ABIdM"}}], travel_mode="driving")` + +User Query: "Directions from San Francisco to San Jose by car" Tool Call: +`render_directions_template(heading="Driving directions from San Francisco to San Jose", summary="Driving from San Francisco to San Jose +takes about 50 minutes via US-101 S.", center_lat=37.55, center_lng=-122.15, +zoom=10, routes=[{"origin": {"lat": 37.7749, "lng": -122.4194, "label": "San +Francisco", "placeId": "ChIJIQBpAG2ahYAR_6128GcTUEo"}, "destination": {"lat": +37.3382, "lng": -121.8863, "label": "San Jose", "placeId": +"ChIJ9T_nxcC1j4ARmMo7S4ABIdM"}}], travel_mode="driving")` ### Example 2: Walking Route -User Query: "How do I walk from Central Park to Times Square?" -Tool Call: -`set_model_response(summary="Walking from Central Park to Times Square takes about 18 minutes (0.9 miles) down 7th Ave.", center_lat=40.765, center_lng=-73.978, zoom=14, routes=[{"origin": {"lat": 40.768, "lng": -73.974, "label": "Central Park South", "placeId": "ChIJN1t_tDeuEmsRUsoyG83frY4"}, "destination": {"lat": 40.758, "lng": -73.985, "label": "Times Square", "placeId": "ChIJmQJItx6vwokRLxVi2JyuzRo"}}], travel_mode="walking")` + +User Query: "How do I walk from Central Park to Times Square?" Tool Call: +`render_directions_template(heading="Walking route from Central Park to Times Square", summary="Walking from Central Park to Times Square +takes about 18 minutes (0.9 miles) down 7th Ave.", center_lat=40.765, +center_lng=-73.978, zoom=14, routes=[{"origin": {"lat": 40.768, "lng": -73.974, +"label": "Central Park South", "placeId": "ChIJN1t_tDeuEmsRUsoyG83frY4"}, +"destination": {"lat": 40.758, "lng": -73.985, "label": "Times Square", +"placeId": "ChIJmQJItx6vwokRLxVi2JyuzRo"}}], travel_mode="walking")` ### Example 3: Bicycling Route -User Query: "Bike directions from Venice Beach to Santa Monica Pier" -Tool Call: -`set_model_response(summary="Biking from Venice Beach to Santa Monica Pier takes around 15 minutes along the Marvin Braude Bike Trail.", center_lat=33.998, center_lng=-118.483, zoom=13, routes=[{"origin": {"lat": 33.985, "lng": -118.469, "label": "Venice Beach", "placeId": "ChIJ-wjh2I-6woARx3H-n9uVn4A"}, "destination": {"lat": 34.009, "lng": -118.497, "label": "Santa Monica Pier", "placeId": "ChIJw8g0Xbm7woARQY1Xq41qB2M"}}], travel_mode="bicycling")` + +User Query: "Bike directions from Venice Beach to Santa Monica Pier" Tool Call: +`render_directions_template(heading="Biking route from Venice Beach to Santa Monica Pier", summary="Biking from Venice Beach to Santa Monica +Pier takes around 15 minutes along the Marvin Braude Bike Trail.", +center_lat=33.998, center_lng=-118.483, zoom=13, routes=[{"origin": {"lat": +33.985, "lng": -118.469, "label": "Venice Beach", "placeId": +"ChIJ-wjh2I-6woARx3H-n9uVn4A"}, "destination": {"lat": 34.009, "lng": -118.497, +"label": "Santa Monica Pier", "placeId": "ChIJw8g0Xbm7woARQY1Xq41qB2M"}}], +travel_mode="bicycling")` ### Example 4: Transit Route -User Query: "Take the subway from Grand Central to Brooklyn Bridge" -Tool Call: -`set_model_response(summary="Take the 4 or 5 subway line south from Grand Central - 42 St to Brooklyn Bridge - City Hall (approx. 12 minutes).", center_lat=40.731, center_lng=-73.988, zoom=12, routes=[{"origin": {"lat": 40.7527, "lng": -73.9772, "label": "Grand Central Terminal", "placeId": "ChIJ4zBEaKZQwokREuE50bbCGYs"}, "destination": {"lat": 40.7126, "lng": -74.0049, "label": "Brooklyn Bridge - City Hall", "placeId": "ChIJ40i5iRZawokRHqGfF2b_3yI"}}], travel_mode="transit")` + +User Query: "Take the subway from Grand Central to Brooklyn Bridge" Tool Call: +`render_directions_template(heading="Transit directions from Grand Central to Brooklyn Bridge", summary="Take the 4 or 5 subway line south from +Grand Central - 42 St to Brooklyn Bridge - City Hall (approx. 12 minutes).", +center_lat=40.731, center_lng=-73.988, zoom=12, routes=[{"origin": {"lat": +40.7527, "lng": -73.9772, "label": "Grand Central Terminal", "placeId": +"ChIJ4zBEaKZQwokREuE50bbCGYs"}, "destination": {"lat": 40.7126, "lng": -74.0049, +"label": "Brooklyn Bridge - City Hall", "placeId": +"ChIJ40i5iRZawokRHqGfF2b_3yI"}}], travel_mode="transit")` ### Example 5: Unspecified Travel Mode (Defaults to Driving) -User Query: "Directions from Austin to San Antonio" -Tool Call: -`set_model_response(summary="Driving from Austin to San Antonio takes about 1 hour and 20 minutes via I-35 S.", center_lat=29.85, center_lng=-98.15, zoom=9, routes=[{"origin": {"lat": 30.2672, "lng": -97.7431, "label": "Austin", "placeId": "ChIJLwW05NsQW4YRtxm00DkzqlU"}, "destination": {"lat": 29.4241, "lng": -98.4936, "label": "San Antonio", "placeId": "ChIJrw7QBK9YXIYRowalignfdg4"}}], travel_mode="driving")` + +User Query: "Directions from Austin to San Antonio" Tool Call: +`render_directions_template(heading="Driving directions from Austin to San Antonio", summary="Driving from Austin to San Antonio takes +about 1 hour and 20 minutes via I-35 S.", center_lat=29.85, center_lng=-98.15, +zoom=9, routes=[{"origin": {"lat": 30.2672, "lng": -97.7431, "label": "Austin", +"placeId": "ChIJLwW05NsQW4YRtxm00DkzqlU"}, "destination": {"lat": 29.4241, +"lng": -98.4936, "label": "San Antonio", "placeId": +"ChIJrw7QBK9YXIYRowalignfdg4"}}], travel_mode="driving")` diff --git a/agent/python_agent/templates/directions/__init__.py b/agent/python_agent/templates/directions/__init__.py new file mode 100644 index 0000000..174bec1 --- /dev/null +++ b/agent/python_agent/templates/directions/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 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. + +"""Directions template bundle.""" diff --git a/agent/python_agent/templates/directions.json b/agent/python_agent/templates/directions/layout.json similarity index 78% rename from agent/python_agent/templates/directions.json rename to agent/python_agent/templates/directions/layout.json index 1bf5f12..389d960 100644 --- a/agent/python_agent/templates/directions.json +++ b/agent/python_agent/templates/directions/layout.json @@ -3,7 +3,7 @@ "version": "v0.9", "createSurface": { "surfaceId": "{{surface_id}}", - "catalogId": "a2ui://maps-agentic-ui-catalog.json" + "catalogId": "{{catalog_id}}" } }, { @@ -14,13 +14,13 @@ { "id": "root", "component": "Column", - "children": ["summary-text", "map"] + "children": ["heading-text", "map", "summary-text"] }, { - "id": "summary-text", + "id": "heading-text", "component": "Text", "variant": "body", - "text": "{{summary}}" + "text": "### {{heading}}" }, { "id": "map", @@ -32,6 +32,12 @@ "zoom": "{{zoom}}", "routes": "{{routes}}", "travelMode": "{{travel_mode}}" + }, + { + "id": "summary-text", + "component": "Text", + "variant": "body", + "text": "{{summary}}" } ] } diff --git a/agent/python_agent/templates/directions/manifest.json b/agent/python_agent/templates/directions/manifest.json new file mode 100644 index 0000000..b2a6340 --- /dev/null +++ b/agent/python_agent/templates/directions/manifest.json @@ -0,0 +1,29 @@ +{ + "intent": "DIRECTIONS", + "surface_prefix": "directions-surface", + "archetype": "Standard navigation, route directions, walking/driving/transit times, or navigation instructions between an origin and destination. This includes multi-stop routes and routes with specified waypoints.", + "empty_fallback_text": "Could not calculate travel directions.", + "tool_description": "Renders an interactive Google Maps directions and routing UI component populated with route segments, travel mode, and a summary response.", + "conformance_example": { + "surface_id": "test-surface", + "heading": "Directions", + "summary": "Commute is 15 mins.", + "center_lat": 37.5, + "center_lng": 127.0, + "zoom": 12, + "routes": [ + { + "origin": {"lat": 37.4, "lng": 126.9, "label": "Start"}, + "destination": {"lat": 37.6, "lng": 127.1, "label": "End"} + } + ], + "travel_mode": "driving" + }, + "few_shot_examples": [ + { + "user_query": "Directions from Sacramento to Mendocino via Clear Lake", + "intent": "DIRECTIONS", + "query": "Directions from Sacramento to Mendocino via Clear Lake" + } + ] +} diff --git a/agent/python_agent/extractor.py b/agent/python_agent/templates/directions/schema.py similarity index 50% rename from agent/python_agent/extractor.py rename to agent/python_agent/templates/directions/schema.py index 53749a0..61c52f5 100644 --- a/agent/python_agent/extractor.py +++ b/agent/python_agent/templates/directions/schema.py @@ -12,97 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Pydantic schemas for structured data extraction from LLM responses.""" +"""Pydantic extraction schema for directions template bundle.""" +import re from typing import Any, Literal + import pydantic +from templates.common import Pin + BaseModel = pydantic.BaseModel Field = pydantic.Field -class Pin(BaseModel): - """Representation of a Map Pin.""" - - lat: float = Field(description="Latitude coordinate") - lng: float = Field(description="Longitude coordinate") - label: str = Field( - description=( - "Descriptive display label string (e.g. name of address, business, or" - " landmark)" - ) - ) - # Note: Using camelCase field name to match frontend A2UI requirements. - placeId: str | None = Field( # pylint: disable=invalid-name - default=None, description="Optional Google Maps Place ID" - ) - - @pydantic.model_validator(mode="before") - @classmethod - def normalize_label(cls, data: Any) -> Any: - """Normalizes the pin label. - - If 'label' is missing but 'name' is present, copies 'name' to 'label'. - If 'label' is still empty, defaults to 'Location' to ensure - the UI always has a valid string to render for the marker (avoiding raw - Place IDs). - - Args: - data: The input dictionary before validation. - - Returns: - The normalized dictionary. - """ - if isinstance(data, dict): - if "label" not in data and "name" in data: - data["label"] = data["name"] - if not data.get("label"): - data["label"] = "Location" - return data - - -class PlacePin(BaseModel): - """Simplified Map Pin representation for search results.""" - - # Note: Using camelCase field name to match frontend A2UI requirements. - # ADK's SetModelResponseTool serialization dumps using field names - # without aliases. - placeId: str = Field( # pylint: disable=invalid-name - description="The unique Google Maps Place ID" - ) - name: str = Field(description="Name of the place") - lat: float = Field(description="Latitude coordinates") - lng: float = Field(description="Longitude coordinates") - - -class LocalSearchExtractorSchema(BaseModel): - """Structured parameters to render a local search UI update.""" - - summary: str = Field( - description=( - "A detailed response summarizing the search results that fully and" - " clearly answers all aspects of the user's prompt (including" - " qualitative criteria, preferences, and comparisons). Use markdown" - " formatting (bullet points, bolding, tables) and break into" - " paragraphs as needed. Bold place names." - ) - ) - center_lat: float = Field(description="Latitude of the center of results") - center_lng: float = Field(description="Longitude of the center of results") - zoom: int = Field( - default=13, description="Recommended map zoom level (typically 13)" - ) - places: list[PlacePin] = Field( - description="A list of places found (limit to max list size, e.g. 3)" - ) - anchor_marker: Pin | None = Field( - default=None, - description=( - "Optional starting or focus point marker (e.g. hotel location)" - ), - ) - - class RouteSegment(BaseModel): """A segment of a route, containing an origin and a destination pin.""" @@ -156,14 +78,20 @@ def normalize_travel_mode(mode: Any) -> str | None: class DirectionsExtractorSchema(BaseModel): """Structured parameters to render a directions UI update.""" - summary: str = Field( + heading: str = Field( description=( - "A detailed response summarizing the travel directions and route" - " options that fully answers all user questions, route comparisons," - " and travel context requested in the prompt. Use markdown formatting" - " and break into paragraphs if helpful." + "A concise, constraint-confirming primary heading for the response." + " Plain text only (e.g., 'Walking route from Seattle Center to Pike" + " Place Market', 'Driving directions to JFK Airport')." ) ) + summary: str = Field( + description=( + "A natural, direct resolution of the route prompt describing" + " approximate travel duration and distance (e.g. 'Driving from" + " [Origin] to [Destination] takes about 19 minutes (14 miles).')." + ), + ) center_lat: float = Field( description="Latitude of the center of the route map" ) @@ -174,7 +102,7 @@ class DirectionsExtractorSchema(BaseModel): default=12, description="Recommended map zoom level (typically 12)" ) routes: list[RouteSegment] = Field( - default_factory=list, + min_length=1, description=( "A list of route segments connecting the origin, intermediate" " waypoints, and the destination in order." @@ -190,13 +118,29 @@ class DirectionsExtractorSchema(BaseModel): @pydantic.model_validator(mode="before") @classmethod def normalize_directions_data(cls, data: Any) -> Any: - """Normalizes travel mode in directions data.""" + """Normalizes travel mode and markdown in heading in directions data.""" if not isinstance(data, dict): return data + heading = data.get("heading") + if isinstance(heading, str): + data["heading"] = re.sub(r"^#+\s*", "", heading).strip() + if "travel_mode" in data and data["travel_mode"]: normalized = normalize_travel_mode(data["travel_mode"]) if normalized: data["travel_mode"] = normalized return data + + +ExtractorSchema = DirectionsExtractorSchema + +__all__ = [ + "DirectionsExtractorSchema", + "ExtractorSchema", + "Pin", + "RouteSegment", + "TRAVEL_MODE_MAP", + "normalize_travel_mode", +] diff --git a/agent/python_agent/test_extractor.py b/agent/python_agent/templates/directions/test_schema.py similarity index 51% rename from agent/python_agent/test_extractor.py rename to agent/python_agent/templates/directions/test_schema.py index f4d77b9..c743bcf 100644 --- a/agent/python_agent/test_extractor.py +++ b/agent/python_agent/templates/directions/test_schema.py @@ -12,44 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for extractor.py.""" +"""Unit tests for directions template bundle extraction schema.""" import unittest import pydantic -from extractor import ( +from templates.directions.schema import ( DirectionsExtractorSchema, - LocalSearchExtractorSchema, - Pin, - PlacePin, + ExtractorSchema, + normalize_travel_mode, ) -class TestExtractor(unittest.TestCase): - """Unit tests for extractor schemas and data normalization.""" - - def test_pin_normalize_label_copies_name(self): - data = {"lat": 1.0, "lng": 2.0, "name": "My Place"} - pin = Pin(**data) - self.assertEqual(pin.label, "My Place") - - def test_pin_normalize_label_defaults_to_location(self): - data = {"lat": 1.0, "lng": 2.0} - pin = Pin(**data) - self.assertEqual(pin.label, "Location") - - def test_pin_normalize_label_preserves_existing(self): - data = { - "lat": 1.0, - "lng": 2.0, - "label": "Custom Label", - "name": "Ignored Name", - } - pin = Pin(**data) - self.assertEqual(pin.label, "Custom Label") +class TestDirectionsSchema(unittest.TestCase): + """Unit tests for directions schema and data normalization.""" def test_directions_extractor_schema_normalize_travel_mode(self): """Verifies that travel mode is normalized to lowercase.""" data = { + "heading": "Commute Route", "summary": "Commute is 1h.", "center_lat": 37.5, "center_lng": 127.0, @@ -65,6 +45,7 @@ def test_directions_extractor_schema_normalize_travel_mode(self): def test_directions_extractor_schema_with_routes(self): """Verifies that DirectionsExtractorSchema can be initialized with routes.""" data = { + "heading": "Scenic Route", "summary": "Scenic route.", "center_lat": 37.5, "center_lng": 127.0, @@ -91,6 +72,7 @@ def test_directions_extractor_schema_missing_travel_mode_fails_validation( ): """Verifies that omitting travel_mode raises ValidationError.""" data = { + "heading": "Directions Route", "summary": "Directions summary", "center_lat": 37.5, "center_lng": 127.0, @@ -109,10 +91,14 @@ def test_directions_extractor_schema_invalid_travel_mode_fails_validation( for invalid_mode in ["flying", "", None, "scooter", 123]: with self.subTest(invalid_mode=invalid_mode): data = { + "heading": "Directions Route", "summary": "Directions summary", "center_lat": 37.5, "center_lng": 127.0, - "routes": [], + "routes": [{ + "origin": {"lat": 37.4, "lng": 126.9, "label": "Start"}, + "destination": {"lat": 37.6, "lng": 127.1, "label": "End"}, + }], "travel_mode": invalid_mode, } with self.assertRaises(pydantic.ValidationError): @@ -123,10 +109,14 @@ def test_directions_extractor_schema_all_valid_modes(self): for mode in ["driving", "walking", "transit", "bicycling"]: with self.subTest(mode=mode): data = { + "heading": f"Going via {mode}", "summary": f"Going via {mode}", "center_lat": 37.5, "center_lng": 127.0, - "routes": [], + "routes": [{ + "origin": {"lat": 37.4, "lng": 126.9, "label": "Start"}, + "destination": {"lat": 37.6, "lng": 127.1, "label": "End"}, + }], "travel_mode": mode, } schema = DirectionsExtractorSchema(**data) @@ -197,15 +187,138 @@ def test_directions_extractor_schema_normalize_all_synonyms(self): for synonym in synonyms: with self.subTest(synonym=synonym, expected=expected_mode): data = { + "heading": "Commute", "summary": "Commute", "center_lat": 37.5, "center_lng": 127.0, - "routes": [], + "routes": [{ + "origin": {"lat": 37.4, "lng": 126.9, "label": "Start"}, + "destination": {"lat": 37.6, "lng": 127.1, "label": "End"}, + }], "travel_mode": synonym, } schema = DirectionsExtractorSchema(**data) self.assertEqual(schema.travel_mode, expected_mode) + def test_directions_extractor_schema_with_heading(self): + """Verifies that DirectionsExtractorSchema validates with heading.""" + data = { + "heading": "Walking route from Seattle Center to Pike Place Market", + "summary": "Walking takes about 25 minutes (1 mile).", + "center_lat": 47.6205, + "center_lng": -122.3493, + "travel_mode": "walking", + "routes": [{ + "origin": { + "lat": 47.6205, + "lng": -122.3493, + "label": "Seattle Center", + }, + "destination": { + "lat": 47.6097, + "lng": -122.3422, + "label": "Pike Place Market", + }, + }], + } + schema = DirectionsExtractorSchema(**data) + self.assertEqual( + schema.heading, "Walking route from Seattle Center to Pike Place Market" + ) + + def test_directions_extractor_schema_heading_strips_markdown_hashtags(self): + """Verifies that leading markdown hashtags in heading are stripped.""" + data = { + "heading": "### Driving Route", + "summary": "About 15 minutes.", + "center_lat": 37.5, + "center_lng": 127.0, + "travel_mode": "driving", + "routes": [{ + "origin": {"lat": 37.5, "lng": 127.0, "label": "Origin"}, + "destination": {"lat": 37.6, "lng": 127.1, "label": "Dest"}, + }], + } + schema = DirectionsExtractorSchema(**data) + self.assertEqual(schema.heading, "Driving Route") + + def test_directions_extractor_schema_empty_heading_normalizes_to_empty(self): + """Verifies that an empty string heading normalizes to empty string.""" + data = { + "heading": "", + "summary": "About 15 minutes.", + "center_lat": 37.5, + "center_lng": 127.0, + "travel_mode": "driving", + "routes": [{ + "origin": {"lat": 37.5, "lng": 127.0, "label": "Origin"}, + "destination": {"lat": 37.6, "lng": 127.1, "label": "Dest"}, + }], + } + schema = DirectionsExtractorSchema(**data) + self.assertEqual(schema.heading, "") + + data["heading"] = "### " + schema = DirectionsExtractorSchema(**data) + self.assertEqual(schema.heading, "") + + def test_directions_extractor_schema_missing_heading_fails_validation(self): + """Verifies that omitting heading raises ValidationError.""" + data = { + "summary": "Walking takes about 25 minutes (1 mile).", + "center_lat": 47.6205, + "center_lng": -122.3493, + "travel_mode": "walking", + "routes": [{ + "origin": { + "lat": 47.6205, + "lng": -122.3493, + "label": "Seattle Center", + }, + "destination": { + "lat": 47.6097, + "lng": -122.3422, + "label": "Pike Place Market", + }, + }], + } + with self.assertRaises(pydantic.ValidationError): + DirectionsExtractorSchema(**data) + + def test_directions_extractor_schema_rejects_an_empty_route_list(self): + """Verifies that a directions card with nothing to draw fails validation. + + The layout renders a route polyline, so an empty list produces a blank + map. Rejecting it here is what degrades the response to plain text. + """ + with self.assertRaises(pydantic.ValidationError): + DirectionsExtractorSchema( + heading="Directions", + summary="No route available.", + center_lat=47.6205, + center_lng=-122.3493, + travel_mode="walking", + routes=[], + ) + + def test_directions_extractor_schema_missing_summary_fails_validation(self): + """Verifies that omitting summary raises ValidationError.""" + with self.assertRaises(pydantic.ValidationError): + DirectionsExtractorSchema( + heading="Directions", + center_lat=47.6205, + center_lng=-122.3493, + travel_mode="walking", + routes=[{ + "origin": {"lat": 37.4, "lng": 126.9, "label": "Start"}, + "destination": {"lat": 37.6, "lng": 127.1, "label": "End"}, + }], + ) + + def test_extractor_schema_alias(self): + """Verifies ExtractorSchema alias matches DirectionsExtractorSchema.""" + self.assertIs(ExtractorSchema, DirectionsExtractorSchema) + if __name__ == "__main__": unittest.main() diff --git a/agent/python_agent/templates/directions/testdata/merged.golden.json b/agent/python_agent/templates/directions/testdata/merged.golden.json new file mode 100644 index 0000000..81a000f --- /dev/null +++ b/agent/python_agent/templates/directions/testdata/merged.golden.json @@ -0,0 +1,70 @@ +[ + { + "version": "v0.9", + "createSurface": { + "surfaceId": "test-surface", + "catalogId": "a2ui://maps-agentic-ui-catalog.json" + } + }, + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "test-surface", + "components": [ + { + "id": "root", + "component": "Column", + "children": [ + "heading-text", + "map", + "summary-text" + ] + }, + { + "id": "heading-text", + "component": "Text", + "variant": "body", + "text": "### Directions" + }, + { + "id": "map", + "component": "GoogleMap", + "center": { + "lat": 37.5, + "lng": 127.0 + }, + "zoom": 12, + "routes": [ + { + "origin": { + "lat": 37.4, + "lng": 126.9, + "label": "Start" + }, + "destination": { + "lat": 37.6, + "lng": 127.1, + "label": "End" + } + } + ], + "travelMode": "driving" + }, + { + "id": "summary-text", + "component": "Text", + "variant": "body", + "text": "Commute is 15 mins." + } + ] + } + }, + { + "version": "v0.9", + "updateDataModel": { + "surfaceId": "test-surface", + "path": "/", + "value": {} + } + } +] diff --git a/agent/python_agent/templates/local_search/SKILL.md b/agent/python_agent/templates/local_search/SKILL.md new file mode 100644 index 0000000..e9e1130 --- /dev/null +++ b/agent/python_agent/templates/local_search/SKILL.md @@ -0,0 +1,106 @@ +--- +name: local_search +description: Extractor skill for local place search queries. Extracts location and list of places for template merging. +--- + +# Local Search Guidelines + +## Core Objective + +Extract structured parameters for local searches. You must call maps tools to +locate matching businesses/places, and call `render_local_search_template` to +render the results. + +## Grounding & Tool-Calling Policy (CRITICAL) + +1. **DO NOT HALLUCINATE OR GUESS**: You are strictly forbidden from generating + coordinates (latitude, longitude) or Google Maps Place IDs from your + internal memory or training weights. +2. **MANDATORY TOOL CALLS**: You MUST call the `search_places` tool first to + find actual venues matching the user's query near the requested locations. +3. **EXACT MATCH & PLACE TYPES**: Any place name, coordinates, or Place ID returned in your + final response MUST correspond exactly to the data returned by the `search_places` tool call. + Determine `placePrimaryType` using the descriptions or categories in the tool response. If insufficient, infer it from the user prompt and place title. + +## Multi-Step Location Resolution Policy (Anchored Search) + +If the user's query references a starting point, landmark, hotel, or specific +address as a geographical anchor (e.g., "Space Needle", "Hyatt hotel", "1600 +Amphitheatre Pkwy"): + +1. **Resolve Anchor Coordinates First**: Make a tool call to `search_places` + with the anchor name as `textQuery` to resolve its exact center coordinates. +2. **Execute Proximity Search (Pivot on Anchor)**: Use the resolved latitude + and longitude of the anchor as the center of a `locationBias` circle. Set + the `radiusMeters` parameter based on the user's query: + * **Explicit Distance**: If the query specifies a distance (e.g., "within + 5 miles", "2 km"), parse and convert it to meters (e.g., `8000` or + `2000`). + * **Implicit Walking**: If the query implies walking (e.g., "walk to", + "walking distance"), default to `1000` meters. + * **No Specific Constraint**: If no distance constraint is specified + (e.g., "near", "around"), omit the `radiusMeters` field to let the + search engine bias results dynamically around the center point. +3. **Handle Non-Geographic Filters**: Keep the search query focused on the core + category and key searchable features or amenities (e.g., "restaurant outdoor + seating", "cafe wifi"). Strip conversational filler words (e.g., "with", + "that has", "places offering") to maximize search relevance and matching + accuracy. +4. **Set Anchor Marker**: Populate the `anchor_marker` response parameter with + the resolved coordinates, name, and Place ID of the anchor location. + +## Multiple Named Places / Comparison on One Map + +When the user asks to show or compare multiple specific named places, +buildings, landmarks, or addresses on a map (e.g., "Show me GWC2 and google +1950 in the same map"): + +1. **Parallel Place Lookups**: Issue `search_places` calls in parallel for + each requested place. +2. **Disambiguation Retry**: If a lookup for an abbreviation or building code + returns empty results, retry once with a broader or disambiguated query + (e.g., adding the campus or city name). +3. **Render Partial Results**: As long as at least one place resolves, you + MUST call `render_local_search_template` with all resolved places in + `places` (using the average of resolved coordinates for `center_lat` and + `center_lng`) and note any unresolved place in `summary`. + +## Handling Failures & Empty Results (CRITICAL) + +If search queries return empty results (`{}`) or fail: + +1. Retry at most **once** with a broader or disambiguated query if the initial + query used an abbreviation or building code. +2. If at least one requested place resolved, call + `render_local_search_template` with the resolved subset and mention any + unfound place in `summary`. +3. Only if **no** places can be resolved at all, exit the tool-calling loop + and return a user-friendly summary explaining that no results were found. + +## Output Fields + +You MUST call `render_local_search_template` with all required fields in the +schema, and optionally the anchor marker if resolved: + +- **`heading`**: A concise, constraint-confirming primary heading in sentence case that starts with or includes the exact number of places provided in the UI response, reflecting the prompt and primary reference location (e.g., '5 vegetarian restaurants near The Plaza Hotel', '5 transit stops near Seattle Center'). Use only the primary reference location without redundant city/state nesting. Plain text only; do NOT include markdown hashtags or conversational filler. +- **`summary`**: A concise 1-paragraph overview that covers all returned places by weaving them into natural, contrasting groups (e.g., pairing lively group-friendly spots vs. intimate neighborhood bistros) rather than listing them one by one. Broadly characterize the dining or activity landscape near the reference location using concrete, sensory details, bolding every place name (e.g., **Carmine's** and **Tony's Di Napoli**), and directly addressing any prompt constraints. For nearby places, never describe distances as numbers (e.g., do not say "0.3 miles" or "500 meters"). Instead, generalize (e.g., "a short walk", "just steps away", or "a quick stroll"). Do NOT include conversational greetings ('Sure!', 'Here are...') and do NOT list place names in bullet points (individual place cards handle individual places). +- **`center_lat`**: Latitude of the center of results. Use the coordinates of + the resolved anchor location (or the average of the results if no anchor is + resolved). +- **`center_lng`**: Longitude of the center of results. Use the coordinates of + the resolved anchor location (or the average of the results if no anchor is + resolved). +- **`zoom`**: Recommended map zoom level. Default to 13. +- **`places`**: Return 5 grounded places in the 'places' array by default. If + the user prompt explicitly specifies a number of places, return exactly that + number in the 'places' array if possible. For each place, determine + `placePrimaryType` using the descriptions or categories in the tool response + (or infer it from the user prompt and place title) matching supported types + (`food_and_drink`, `outdoor`, `retail`, `gas_station`, `ev`, `bank`, + `lodging`, `emergency`, `entertainment`, `airport`, `parking`, `generic`). + For geopolitical or non-POI places (cities, towns, neighborhoods, regions, + countries, postal codes, street addresses; or Google types such as + `locality`, `political`, `administrative_area_level_*`, `country`, + `sublocality`, `neighborhood`, `postal_code`, `route`, `street_address`, + `geocode`), ALWAYS use `generic` (never `outdoor`). +- **`anchor_marker`**: (Optional) Pin details for the resolved starting/anchor location. diff --git a/agent/python_agent/templates/local_search/__init__.py b/agent/python_agent/templates/local_search/__init__.py new file mode 100644 index 0000000..f155006 --- /dev/null +++ b/agent/python_agent/templates/local_search/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 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. + +"""Local search template bundle.""" diff --git a/agent/python_agent/templates/local_search.json b/agent/python_agent/templates/local_search/layout.json similarity index 80% rename from agent/python_agent/templates/local_search.json rename to agent/python_agent/templates/local_search/layout.json index 3d964e4..03c0ab2 100644 --- a/agent/python_agent/templates/local_search.json +++ b/agent/python_agent/templates/local_search/layout.json @@ -3,7 +3,7 @@ "version": "v0.9", "createSurface": { "surfaceId": "{{surface_id}}", - "catalogId": "a2ui://maps-agentic-ui-catalog.json" + "catalogId": "{{catalog_id}}" } }, { @@ -14,7 +14,13 @@ { "id": "root", "component": "Column", - "children": ["summary-text", "map", "list"] + "children": ["heading-text", "summary-text", "map", "list"] + }, + { + "id": "heading-text", + "component": "Text", + "variant": "body", + "text": "### {{heading}}" }, { "id": "summary-text", @@ -30,6 +36,8 @@ "lng": "{{center_lng}}" }, "zoom": "{{zoom}}", + "tilt": 0, + "mode": "roadmap", "anchorMarker": "{{anchor_marker}}", "markers": "{{markers}}" }, diff --git a/agent/python_agent/templates/local_search/manifest.json b/agent/python_agent/templates/local_search/manifest.json new file mode 100644 index 0000000..266cfd8 --- /dev/null +++ b/agent/python_agent/templates/local_search/manifest.json @@ -0,0 +1,31 @@ +{ + "intent": "LOCAL_SEARCH", + "surface_prefix": "local-search-surface", + "archetype": "Search for categories of interest, places, businesses, or points of interest within a specific geographic proximity, or display one or more specific named places, buildings, landmarks, or addresses on a map.", + "empty_fallback_text": "No places matching your query could be found.", + "tool_description": "Renders an interactive Google Maps local search UI component populated with places, map markers, and a summary response.", + "conformance_example": { + "surface_id": "test-surface", + "heading": "Top Places", + "summary": "Here are 2 places.", + "center_lat": 47.6, + "center_lng": -122.3, + "zoom": 13, + "places": [ + {"placeId": "ChIJ1", "name": "P1", "lat": 47.6, "lng": -122.3}, + {"placeId": "ChIJ2", "name": "P2", "lat": 47.61, "lng": -122.31} + ] + }, + "few_shot_examples": [ + { + "user_query": "Find coffee shops near Central Park", + "intent": "LOCAL_SEARCH", + "query": "coffee shops near Central Park" + }, + { + "user_query": "Show me GWC2 and google 1950 in the same map", + "intent": "LOCAL_SEARCH", + "query": "GWC2 and Google 1950" + } + ] +} diff --git a/agent/python_agent/templates/local_search/schema.py b/agent/python_agent/templates/local_search/schema.py new file mode 100644 index 0000000..ecf4d50 --- /dev/null +++ b/agent/python_agent/templates/local_search/schema.py @@ -0,0 +1,148 @@ +# Copyright 2026 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. + +"""Pydantic extraction schema for local_search template bundle.""" + +import re +from typing import Any, ClassVar + +import pydantic + +from templates.common import Pin +from templates.common import PlacePrimaryType + +BaseModel = pydantic.BaseModel +Field = pydantic.Field + + +class PlacePin(BaseModel): + """Simplified Map Pin representation for search results.""" + + # Note: Using camelCase field name to match frontend A2UI requirements. + # ADK's SetModelResponseTool serialization dumps using field names + # without aliases. + placeId: str = Field( # pylint: disable=invalid-name + description="The unique Google Maps Place ID" + ) + name: str = Field(description="Name of the place") + lat: float = Field(description="Latitude coordinates") + lng: float = Field(description="Longitude coordinates") + placePrimaryType: PlacePrimaryType | None = Field( # pylint: disable=invalid-name + default=None, + description=( + "Optional primary POI category type string. Use 'generic' for" + " geopolitical or non-POI places (cities, towns, neighborhoods," + " regions, countries, addresses)." + ), + ) + + @pydantic.model_validator(mode="before") + @classmethod + def normalize_place_pin(cls, data: Any) -> Any: + if isinstance(data, dict): + if "name" not in data and "label" in data: + data["name"] = data["label"] + if "label" not in data and "name" in data: + data["label"] = data["name"] + return data + + +class LocalSearchExtractorSchema(BaseModel): + """Structured parameters to render a local search UI update.""" + + # Collections the extractor prompt explicitly tells the model to limit. A + # place search can return dozens of candidates, and the model will happily + # narrate all of them even though the layout renders only `max_list_size`. + # Bundles whose collections are naturally small omit this and get no + # constraint paragraph. + clamped_collections: ClassVar[tuple[str, ...]] = ("places",) + + heading: str = Field( + description=( + "A concise, constraint-confirming primary heading in sentence case" + " that starts with or includes the exact number of places provided" + " in the UI response, reflecting the prompt and primary reference" + " location (e.g. '5 vegetarian restaurants near The Plaza Hotel'," + " '5 transit stops near Seattle Center'). Plain text only; do" + " NOT include markdown hashtags or conversational filler." + ), + ) + summary: str = Field( + description=( + "A concise 1-paragraph overview that covers all returned places by" + " weaving them into natural, contrasting groups (e.g., pairing" + " lively group-friendly spots vs. intimate neighborhood bistros)" + " rather than listing them one by one. Broadly characterize the" + " dining or activity landscape near the reference location using" + " concrete, sensory details, bolding every place name (e.g.," + " **Carmine's** and **Tony's Di Napoli**), and directly addressing" + " any prompt constraints. For nearby places, never describe" + " distances as numbers (e.g., do not say '0.3 miles' or '500" + " meters'); instead generalize (e.g., 'a short walk', 'just steps" + " away', 'a quick stroll'). Plain text with markdown bolding only;" + " do NOT include conversational greetings ('Sure!', 'Here are...')" + " and do NOT list place names in bullet points." + ), + ) + center_lat: float = Field(description="Latitude of the center of results") + center_lng: float = Field(description="Longitude of the center of results") + zoom: int = Field( + default=13, description="Recommended map zoom level (typically 13)" + ) + places: list[PlacePin] = Field( + min_length=1, + description="A list of places found (limit to max list size, e.g. 5)", + ) + anchor_marker: Pin | None = Field( + default=None, + description=( + "Optional starting or focus point marker (e.g. hotel location)" + ), + ) + + @pydantic.model_validator(mode="before") + @classmethod + def normalize_heading(cls, data: Any) -> Any: + """Normalizes markdown headers in heading if present.""" + if isinstance(data, dict): + heading = data.get("heading") + if heading and isinstance(heading, str): + data["heading"] = re.sub(r"^#+\s*", "", heading).strip() + return data + + @pydantic.computed_field + @property + def markers(self) -> list[Pin]: + """Projects map pins from the places list.""" + return [ + Pin( + lat=p.lat, + lng=p.lng, + label=p.name, + placeId=p.placeId, + placePrimaryType=p.placePrimaryType, + ) + for p in self.places + ] + + +ExtractorSchema = LocalSearchExtractorSchema + +__all__ = [ + "ExtractorSchema", + "LocalSearchExtractorSchema", + "Pin", + "PlacePin", + "PlacePrimaryType", +] diff --git a/agent/python_agent/templates/local_search/test_schema.py b/agent/python_agent/templates/local_search/test_schema.py new file mode 100644 index 0000000..5593a24 --- /dev/null +++ b/agent/python_agent/templates/local_search/test_schema.py @@ -0,0 +1,201 @@ +# Copyright 2026 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. + +"""Unit tests for local_search template bundle extraction schema.""" + +import unittest +import pydantic +from templates.local_search.schema import ( + ExtractorSchema, + LocalSearchExtractorSchema, + Pin, + PlacePin, +) + + +class TestLocalSearchSchema(unittest.TestCase): + """Unit tests for local_search schema and data normalization.""" + + def test_pin_normalize_label_copies_name(self): + data = {"lat": 1.0, "lng": 2.0, "name": "My Place"} + pin = Pin(**data) + self.assertEqual(pin.label, "My Place") + + def test_pin_normalize_label_defaults_to_location(self): + data = {"lat": 1.0, "lng": 2.0} + pin = Pin(**data) + self.assertEqual(pin.label, "Location") + + def test_pin_with_place_primary_type(self): + data = { + "lat": 1.0, + "lng": 2.0, + "label": "Coffee Shop", + "placePrimaryType": "food_and_drink", + } + pin = Pin(**data) + self.assertEqual(pin.placePrimaryType, "food_and_drink") + + def test_place_pin_with_place_primary_type(self): + data = { + "placeId": "ChIJ123", + "name": "Coffee Shop", + "lat": 1.0, + "lng": 2.0, + "placePrimaryType": "food_and_drink", + } + pin = PlacePin(**data) + self.assertEqual(pin.placePrimaryType, "food_and_drink") + + def test_pin_normalize_label_preserves_existing(self): + data = { + "lat": 1.0, + "lng": 2.0, + "label": "Custom Label", + "name": "Ignored Name", + } + pin = Pin(**data) + self.assertEqual(pin.label, "Custom Label") + + def test_local_search_extractor_schema_with_heading(self): + """Verifies that LocalSearchExtractorSchema validates with heading.""" + data = { + "heading": "5 Transit Stops Near Seattle Center", + "summary": "Here are 5 transit stops.", + "center_lat": 47.6205, + "center_lng": -122.3493, + "places": [{ + "placeId": "ChIJ111", + "name": "Stop 1", + "lat": 47.62, + "lng": -122.35, + }], + } + schema = LocalSearchExtractorSchema(**data) + self.assertEqual(schema.heading, "5 Transit Stops Near Seattle Center") + + def test_local_search_extractor_schema_heading_strips_markdown_hashtags( + self, + ): + """Verifies that leading markdown hashtags in heading are stripped.""" + data = { + "heading": "### Best Bakeries", + "summary": "Here are bakeries.", + "center_lat": 47.6, + "center_lng": -122.3, + "places": [{"placeId": "p1", "name": "B1", "lat": 47.6, "lng": -122.3}], + } + schema = LocalSearchExtractorSchema(**data) + self.assertEqual(schema.heading, "Best Bakeries") + + def test_local_search_extractor_schema_missing_heading_fails_validation(self): + """Verifies that omitting heading raises ValidationError.""" + data = { + "summary": "Here are 5 transit stops.", + "center_lat": 47.6205, + "center_lng": -122.3493, + "places": [{ + "placeId": "ChIJ111", + "name": "Stop 1", + "lat": 47.62, + "lng": -122.35, + }], + } + with self.assertRaises(pydantic.ValidationError): + LocalSearchExtractorSchema(**data) + + def test_local_search_extractor_schema_computed_markers(self): + """Verifies that markers are computed from places automatically.""" + data = { + "heading": "Coffee places", + "summary": "Here are 2 coffee places.", + "center_lat": 47.62, + "center_lng": -122.35, + "places": [ + { + "placeId": "ChIJ111", + "name": "Espresso Vivace", + "lat": 47.62, + "lng": -122.32, + "placePrimaryType": "food_and_drink", + }, + { + "placeId": "ChIJ222", + "name": "Milstead & Co.", + "lat": 47.64, + "lng": -122.35, + }, + ], + } + schema = LocalSearchExtractorSchema(**data) + self.assertEqual(len(schema.markers), 2) + self.assertEqual(schema.markers[0].label, "Espresso Vivace") + self.assertEqual(schema.markers[0].lat, 47.62) + self.assertEqual(schema.markers[0].lng, -122.32) + self.assertEqual(schema.markers[0].placeId, "ChIJ111") + self.assertEqual(schema.markers[0].placePrimaryType, "food_and_drink") + + self.assertEqual(schema.markers[1].label, "Milstead & Co.") + self.assertEqual(schema.markers[1].placeId, "ChIJ222") + self.assertIsNone(schema.markers[1].placePrimaryType) + + # Verify model_dump includes markers + dumped = schema.model_dump(exclude_none=True) + self.assertIn("markers", dumped) + self.assertEqual(len(dumped["markers"]), 2) + self.assertEqual(dumped["markers"][0]["label"], "Espresso Vivace") + self.assertEqual(dumped["markers"][0]["placeId"], "ChIJ111") + self.assertEqual(dumped["markers"][0]["placePrimaryType"], "food_and_drink") + self.assertNotIn("placePrimaryType", dumped["markers"][1]) + + def test_local_search_extractor_schema_rejects_an_empty_places_list(self): + """Verifies that a local search card with nothing to show fails validation. + + The layout renders a place list and its map markers, so zero places + produces a blank map. Rejecting it here is what degrades the response to + plain text. + """ + data = { + "heading": "No places", + "summary": "No places found.", + "center_lat": 47.62, + "center_lng": -122.35, + "places": [], + } + with self.assertRaises(pydantic.ValidationError): + LocalSearchExtractorSchema(**data) + + def test_local_search_extractor_schema_rejects_place_without_place_id(self): + """Verifies that a place missing placeId fails validation.""" + data = { + "heading": "Coffee places", + "summary": "Here are coffee places.", + "center_lat": 47.62, + "center_lng": -122.35, + "places": [{ + "name": "Coffee Shop Without ID", + "lat": 47.62, + "lng": -122.35, + }], + } + with self.assertRaises(pydantic.ValidationError): + LocalSearchExtractorSchema(**data) + + def test_extractor_schema_alias(self): + """Verifies ExtractorSchema alias matches LocalSearchExtractorSchema.""" + self.assertIs(ExtractorSchema, LocalSearchExtractorSchema) + + +if __name__ == "__main__": + unittest.main() diff --git a/agent/python_agent/templates/local_search/testdata/merged.golden.json b/agent/python_agent/templates/local_search/testdata/merged.golden.json new file mode 100644 index 0000000..6972dcc --- /dev/null +++ b/agent/python_agent/templates/local_search/testdata/merged.golden.json @@ -0,0 +1,103 @@ +[ + { + "version": "v0.9", + "createSurface": { + "surfaceId": "test-surface", + "catalogId": "a2ui://maps-agentic-ui-catalog.json" + } + }, + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "test-surface", + "components": [ + { + "id": "root", + "component": "Column", + "children": [ + "heading-text", + "summary-text", + "map", + "list" + ] + }, + { + "id": "heading-text", + "component": "Text", + "variant": "body", + "text": "### Top Places" + }, + { + "id": "summary-text", + "component": "Text", + "variant": "body", + "text": "Here are 2 places." + }, + { + "id": "map", + "component": "GoogleMap", + "center": { + "lat": 47.6, + "lng": -122.3 + }, + "zoom": 13, + "tilt": 0, + "mode": "roadmap", + "markers": [ + { + "lat": 47.6, + "lng": -122.3, + "label": "P1", + "placeId": "ChIJ1" + }, + { + "lat": 47.61, + "lng": -122.31, + "label": "P2", + "placeId": "ChIJ2" + } + ] + }, + { + "id": "list", + "component": "List", + "direction": "vertical", + "children": { + "componentId": "place-card", + "path": "/places" + } + }, + { + "id": "place-card", + "component": "PlaceDetailsCompact", + "placeId": { + "path": "placeId" + } + } + ] + } + }, + { + "version": "v0.9", + "updateDataModel": { + "surfaceId": "test-surface", + "path": "/", + "value": { + "places": [ + { + "placeId": "ChIJ1", + "name": "P1", + "lat": 47.6, + "lng": -122.3 + }, + { + "placeId": "ChIJ2", + "name": "P2", + "lat": 47.61, + "lng": -122.31 + } + ] + } + } + } +] diff --git a/agent/python_agent/templates/test_common.py b/agent/python_agent/templates/test_common.py new file mode 100644 index 0000000..acaf9fb --- /dev/null +++ b/agent/python_agent/templates/test_common.py @@ -0,0 +1,123 @@ +# Copyright 2026 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. + +"""Unit tests for shared template models and primitives.""" + +import typing +import unittest + +import pydantic + +from templates import common + + +class PinTest(unittest.TestCase): + """Tests for the Pin model and its validators.""" + + def test_pin_valid_construction(self): + pin = common.Pin( + lat=37.7749, + lng=-122.4194, + label="San Francisco", + placeId="ChIJIQBpAG2ahYAR_6128GcTUEo", + placePrimaryType="retail", + ) + self.assertEqual(pin.lat, 37.7749) + self.assertEqual(pin.lng, -122.4194) + self.assertEqual(pin.label, "San Francisco") + self.assertEqual(pin.placeId, "ChIJIQBpAG2ahYAR_6128GcTUEo") + self.assertEqual(pin.placePrimaryType, "retail") + + def test_pin_normalizes_name_to_label_when_label_is_missing(self): + data = { + "lat": 47.6062, + "lng": -122.3321, + "name": "Pike Place Market", + } + pin = common.Pin.model_validate(data) + self.assertEqual(pin.label, "Pike Place Market") + + def test_pin_defaults_to_location_when_label_and_name_are_empty(self): + data = { + "lat": 47.6062, + "lng": -122.3321, + } + pin = common.Pin.model_validate(data) + self.assertEqual(pin.label, "Location") + + def test_pin_preserves_explicit_label_even_if_name_is_present(self): + data = { + "lat": 47.6062, + "lng": -122.3321, + "label": "Custom Label", + "name": "Ignored Name", + } + pin = common.Pin.model_validate(data) + self.assertEqual(pin.label, "Custom Label") + + def test_pin_requires_coordinates(self): + with self.assertRaises(pydantic.ValidationError): + common.Pin.model_validate({"label": "Missing Coords"}) + + +class PlacePrimaryTypeTest(unittest.TestCase): + """Tests for the supported POI taxonomy.""" + + def test_matches_supported_taxonomy(self): + # Pinned so that adding or removing a category is a deliberate edit that + # also has to be mirrored in the client lookups and the agent skills. + self.assertEqual( + typing.get_args(common.PlacePrimaryType), + ( + "food_and_drink", + "outdoor", + "retail", + "gas_station", + "ev", + "bank", + "lodging", + "emergency", + "entertainment", + "airport", + "parking", + "generic", + ), + ) + + def test_pin_accepts_every_supported_type(self): + for place_primary_type in typing.get_args(common.PlacePrimaryType): + with self.subTest(placePrimaryType=place_primary_type): + pin = common.Pin( + lat=1.0, + lng=2.0, + label="Place", + placePrimaryType=place_primary_type, + ) + self.assertEqual(pin.placePrimaryType, place_primary_type) + + def test_pin_rejects_retired_types(self): + # `service` was split into `gas_station`/`bank`, and `closed` was dropped. + for place_primary_type in ("service", "closed"): + with self.subTest(placePrimaryType=place_primary_type): + with self.assertRaises(pydantic.ValidationError): + common.Pin( + lat=1.0, + lng=2.0, + label="Place", + placePrimaryType=place_primary_type, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/agent/python_agent/templates/text_only.json b/agent/python_agent/templates/text_only/layout.json similarity index 90% rename from agent/python_agent/templates/text_only.json rename to agent/python_agent/templates/text_only/layout.json index a088835..feee812 100644 --- a/agent/python_agent/templates/text_only.json +++ b/agent/python_agent/templates/text_only/layout.json @@ -3,7 +3,7 @@ "version": "v0.9", "createSurface": { "surfaceId": "{{surface_id}}", - "catalogId": "a2ui://maps-agentic-ui-catalog.json" + "catalogId": "{{catalog_id}}" } }, { diff --git a/agent/python_agent/test_after_tools_callback.py b/agent/python_agent/test_after_tools_callback.py new file mode 100644 index 0000000..d56e5f3 --- /dev/null +++ b/agent/python_agent/test_after_tools_callback.py @@ -0,0 +1,307 @@ +# Copyright 2026 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 +# +# http://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. + +"""Tests for after_tools_callback.""" + +import unittest +from unittest import mock + +from after_tools_callback import _add_maps_tools_tokens_to_part, after_maps_tools_callback, after_tools_callback + + +class TestAfterToolsCallback(unittest.TestCase): + + def test_after_tool_callback_aggregates_maps_tools_content_tokens(self): + mock_tool_context = mock.MagicMock() + mock_tool_context.state = {} + + tool_response_1 = { + "content_token": "token_abc_123", + } + after_tools_callback( + tool=None, + args={}, + tool_context=mock_tool_context, + tool_response=tool_response_1, + ) + self.assertEqual( + mock_tool_context.state.get("maps_tools_content_tokens"), + ["token_abc_123"], + ) + + tool_response_2 = { + "content_token": "token_def_456", + } + after_tools_callback( + tool=None, + args={}, + tool_context=mock_tool_context, + tool_response=tool_response_2, + ) + self.assertEqual( + mock_tool_context.state.get("maps_tools_content_tokens"), + ["token_abc_123", "token_def_456"], + ) + + # Calling again with duplicate should not add duplicates + after_tools_callback( + tool=None, + args={}, + tool_context=mock_tool_context, + tool_response={"content_token": "token_abc_123"}, + ) + self.assertEqual( + mock_tool_context.state.get("maps_tools_content_tokens"), + ["token_abc_123", "token_def_456"], + ) + + def test_after_tool_callback_limits_maps_tools_content_tokens(self): + mock_tool_context = mock.MagicMock() + mock_tool_context.state = {} + + for i in range(15): + after_tools_callback( + tool=None, + args={}, + tool_context=mock_tool_context, + tool_response={"content_token": f"token_{i}"}, + ) + + tokens = mock_tool_context.state.get("maps_tools_content_tokens") + self.assertEqual(len(tokens), 10) + self.assertEqual(tokens[0], "token_5") + self.assertEqual(tokens[-1], "token_14") + + def test_after_tool_callback_with_kwargs(self): + mock_tool_context = mock.MagicMock() + mock_tool_context.state = {} + + after_tools_callback( + tool="mock_tool", + args={"query": "test"}, + tool_context=mock_tool_context, + tool_response={"content_token": "token_xyz"}, + extra_param="unused", + ) + self.assertEqual( + mock_tool_context.state.get("maps_tools_content_tokens"), + ["token_xyz"], + ) + + def test_after_tools_callback_none_or_empty_response(self): + mock_tool_context = mock.MagicMock() + mock_tool_context.state = {} + + # None tool response + result = after_tools_callback( + tool=None, + args={}, + tool_context=mock_tool_context, + tool_response=None, + ) + self.assertIsNone(result) + self.assertEqual(mock_tool_context.state, {}) + + # Empty dict tool response + result = after_tools_callback( + tool=None, + args={}, + tool_context=mock_tool_context, + tool_response={}, + ) + self.assertIsNone(result) + self.assertEqual(mock_tool_context.state, {}) + + # Non-dict tool response + result = after_tools_callback( + tool=None, + args={}, + tool_context=mock_tool_context, + tool_response="not a dict", + ) + self.assertIsNone(result) + self.assertEqual(mock_tool_context.state, {}) + + result = after_tools_callback( + tool=None, + args={}, + tool_context=mock_tool_context, + tool_response=["list_not_dict"], + ) + self.assertIsNone(result) + self.assertEqual(mock_tool_context.state, {}) + + def test_after_maps_tools_callback_none_or_missing_context(self): + # None tool_context + result = after_maps_tools_callback( + tool_context=None, + tool_response={"content_token": "token_1"}, + ) + self.assertIsNone(result) + + # tool_context with state=None + mock_context_no_state = mock.MagicMock() + mock_context_no_state.state = None + result = after_maps_tools_callback( + tool_context=mock_context_no_state, + tool_response={"content_token": "token_1"}, + ) + self.assertIsNone(result) + + # tool_context without state attribute + class DummyContext: + pass + + result = after_maps_tools_callback( + tool_context=DummyContext(), + tool_response={"content_token": "token_1"}, + ) + self.assertIsNone(result) + + def test_after_maps_tools_callback_invalid_token_values(self): + mock_tool_context = mock.MagicMock() + mock_tool_context.state = {} + + # None token + after_maps_tools_callback( + tool_context=mock_tool_context, + tool_response={"content_token": None}, + ) + self.assertEqual(mock_tool_context.state, {}) + + # Empty string token + after_maps_tools_callback( + tool_context=mock_tool_context, + tool_response={"content_token": ""}, + ) + self.assertEqual(mock_tool_context.state, {}) + + # Non-string token (int) + after_maps_tools_callback( + tool_context=mock_tool_context, + tool_response={"content_token": 12345}, + ) + self.assertEqual(mock_tool_context.state, {}) + + # Missing content_token key + after_maps_tools_callback( + tool_context=mock_tool_context, + tool_response={"places": []}, + ) + self.assertEqual(mock_tool_context.state, {}) + + def test_after_maps_tools_callback_non_list_state_content_tokens(self): + mock_tool_context = mock.MagicMock() + + # If state['maps_tools_content_tokens'] is not a list (e.g. a string) + mock_tool_context.state = {"maps_tools_content_tokens": "invalid_string"} + after_maps_tools_callback( + tool_context=mock_tool_context, + tool_response={"content_token": "token_1"}, + ) + self.assertEqual( + mock_tool_context.state.get("maps_tools_content_tokens"), + ["token_1"], + ) + + # If state['maps_tools_content_tokens'] is None + mock_tool_context.state = {"maps_tools_content_tokens": None} + after_maps_tools_callback( + tool_context=mock_tool_context, + tool_response={"content_token": "token_2"}, + ) + self.assertEqual( + mock_tool_context.state.get("maps_tools_content_tokens"), + ["token_2"], + ) + + +class TestAddMapsToolsTokensToPart(unittest.TestCase): + + def test_add_tokens_session_none_or_missing_state(self): + part = mock.MagicMock() + part.root.metadata = None + + # session is None + _add_maps_tools_tokens_to_part(part, None) + self.assertIsNone(part.root.metadata) + + # session.state is None + mock_session = mock.MagicMock() + mock_session.state = None + _add_maps_tools_tokens_to_part(part, mock_session) + self.assertIsNone(part.root.metadata) + + def test_add_tokens_empty_tokens_in_session(self): + part = mock.MagicMock() + part.root.metadata = None + + # maps_tools_content_tokens is not in state + session = mock.MagicMock() + session.state = {} + _add_maps_tools_tokens_to_part(part, session) + self.assertIsNone(part.root.metadata) + + # maps_tools_content_tokens is empty list + session.state = {"maps_tools_content_tokens": []} + _add_maps_tools_tokens_to_part(part, session) + self.assertIsNone(part.root.metadata) + + # maps_tools_content_tokens is None + session.state = {"maps_tools_content_tokens": None} + _add_maps_tools_tokens_to_part(part, session) + self.assertIsNone(part.root.metadata) + + def test_add_tokens_with_metadata_none(self): + part = mock.MagicMock() + part.root.metadata = None + + session = mock.MagicMock() + session.state = {"maps_tools_content_tokens": ["token_1", "token_2"]} + + _add_maps_tools_tokens_to_part(part, session) + self.assertEqual( + part.root.metadata, + {"maps_tools_content_tokens": ["token_1", "token_2"]}, + ) + + def test_add_tokens_with_existing_metadata(self): + part = mock.MagicMock() + part.root.metadata = {"existing_field": "existing_value"} + + session = mock.MagicMock() + session.state = {"maps_tools_content_tokens": ["token_1"]} + + _add_maps_tools_tokens_to_part(part, session) + self.assertEqual( + part.root.metadata, + { + "existing_field": "existing_value", + "maps_tools_content_tokens": ["token_1"], + }, + ) + + def test_add_tokens_with_none_root(self): + part = mock.MagicMock() + part.root = None + + session = mock.MagicMock() + session.state = {"maps_tools_content_tokens": ["token_1"]} + + # Should not raise AttributeError + _add_maps_tools_tokens_to_part(part, session) + + +if __name__ == "__main__": + unittest.main() diff --git a/agent/python_agent/test_agent.py b/agent/python_agent/test_agent.py new file mode 100644 index 0000000..1f9f723 --- /dev/null +++ b/agent/python_agent/test_agent.py @@ -0,0 +1,44 @@ +# Copyright 2026 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 unittest +from agent import extract_surface_id + +class SurfaceIdExtractionTest(unittest.TestCase): + + def test_extract_from_create_surface(self): + data = {"createSurface": {"surfaceId": "map_surface_1", "catalogId": "maps"}} + self.assertEqual(extract_surface_id(data), "map_surface_1") + + def test_extract_from_update_components(self): + data = {"updateComponents": {"surfaceId": "details_card_2", "components": []}} + self.assertEqual(extract_surface_id(data), "details_card_2") + + def test_extract_from_update_data_model(self): + data = {"updateDataModel": {"surfaceId": "weather_card_3", "dataModel": {}}} + self.assertEqual(extract_surface_id(data), "weather_card_3") + + def test_extract_from_delete_surface(self): + data = {"deleteSurface": {"surfaceId": "old_surface_4"}} + self.assertEqual(extract_surface_id(data), "old_surface_4") + + def test_extract_non_matching_or_malformed_data(self): + self.assertIsNone(extract_surface_id({"text": "hello"})) + self.assertIsNone(extract_surface_id(None)) + self.assertIsNone(extract_surface_id("not_a_dict")) + self.assertIsNone(extract_surface_id({"createSurface": "malformed_shape"})) + self.assertIsNone(extract_surface_id({"createSurface": {}})) + +if __name__ == '__main__': + unittest.main() diff --git a/agent/python_agent/test_agent_with_templates.py b/agent/python_agent/test_agent_with_templates.py index 5072e61..3c9ac56 100644 --- a/agent/python_agent/test_agent_with_templates.py +++ b/agent/python_agent/test_agent_with_templates.py @@ -23,6 +23,7 @@ import agent import agent_config import agent_with_templates +import template_tool AgentConfig = agent_config.AgentConfig FallbackMode = agent_config.FallbackMode @@ -108,7 +109,7 @@ def lite_llm_side_effect(*args, **kwargs): return self.mock_router elif model in ( "gemini/template-model", - "gemini/gemini-3-flash-preview", + "gemini/gemini-3.7-flash", "gemini/generic-model", ): return self.mock_extractor @@ -189,9 +190,9 @@ async def test_agent_text_only_flow(self, mock_lite_llm_class): template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._grounded_text_runner = mock_runner - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - results = await self._collect_stream(agent, "hello") + results = await self._collect_stream(agent, "hello") self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -230,9 +231,9 @@ async def test_agent_router_failure_fallback(self, mock_lite_llm_class): template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._grounded_text_runner = mock_runner - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - results = await self._collect_stream(agent, "hello") + results = await self._collect_stream(agent, "hello") self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -265,9 +266,9 @@ async def test_agent_extractor_empty_parts_handled_safely( template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._grounded_text_runner = mock_runner - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - results = await self._collect_stream(agent, "hello") + results = await self._collect_stream(agent, "hello") self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -317,10 +318,11 @@ async def test_agent_unsupported_intent_fallback(self, mock_lite_llm_class): def test_init_without_config_uses_default(self): agent = MAUIAgentWithTemplates(base_url="http://test-url") self.assertIsNotNone(agent.config) - self.assertEqual(agent.config.router_model, "gemini/gemini-3.1-flash-lite") + self.assertEqual(agent.config.router_model, "gemini/gemini-3.5-flash-lite") self.assertEqual( - agent.config.template_model, "gemini/gemini-3.1-flash-lite" + agent.config.template_model, "gemini/gemini-3.5-flash-lite" ) + self.assertEqual(agent.config.generic_model, "gemini/gemini-3.7-flash") self.assertEqual(agent.config.fallback_mode, FallbackMode.TEXT) @mock.patch(_LITELLM_PATH) @@ -339,8 +341,9 @@ async def test_agent_directions_flow(self, mock_lite_llm_class): mock_runner = mock.MagicMock() mock_fc = MockFunctionCall( - name="set_model_response", + name="render_directions_template", args={ + "heading": "Directions from home to work", "summary": "Typical commute is 45 mins.", "center_lat": 37.5, "center_lng": 127.0, @@ -370,17 +373,15 @@ async def test_agent_directions_flow(self, mock_lite_llm_class): template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._extractor_runners["0.9"] = mock_runner - # Patch _build_runner - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - # Run stream - results = [] - async for item in agent.stream( - query="directions from home to work", - session_id="session_123", - ui_version="v0.9", - ): - results.append(item) + results = [] + async for item in agent.stream( + query="directions from home to work", + session_id="session_123", + ui_version="v0.9", + ): + results.append(item) self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -432,15 +433,16 @@ async def test_agent_directions_flow_fallback(self, mock_lite_llm_class): template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._extractor_runners["0.9"] = mock_runner + agent._fallback_text_runner = mock_runner - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - results = [] - async for item in agent.stream( - query="directions to work", - session_id="session_123", - ui_version="v0.9", - ): - results.append(item) + results = [] + async for item in agent.stream( + query="directions to work", + session_id="session_123", + ui_version="v0.9", + ): + results.append(item) self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -470,8 +472,9 @@ async def test_agent_directions_flow_transit_mode(self, mock_lite_llm_class): mock_runner = mock.MagicMock() mock_fc = MockFunctionCall( - name="set_model_response", + name="render_directions_template", args={ + "heading": "Bus directions to work", "summary": "Take bus 10 to work.", "center_lat": 37.5, "center_lng": 127.0, @@ -501,9 +504,9 @@ async def test_agent_directions_flow_transit_mode(self, mock_lite_llm_class): template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._extractor_runners["0.9"] = mock_runner - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - results = await self._collect_stream(agent, query="bus to work") + results = await self._collect_stream(agent, query="bus to work") self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -524,8 +527,9 @@ async def test_agent_directions_flow_walking_mode(self, mock_lite_llm_class): mock_runner = mock.MagicMock() mock_fc = MockFunctionCall( - name="set_model_response", + name="render_directions_template", args={ + "heading": "Walking route to park", "summary": "Walk for 15 minutes.", "center_lat": 37.5, "center_lng": 127.0, @@ -555,9 +559,9 @@ async def test_agent_directions_flow_walking_mode(self, mock_lite_llm_class): template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._extractor_runners["0.9"] = mock_runner - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - results = await self._collect_stream(agent, query="walk to park") + results = await self._collect_stream(agent, query="walk to park") self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -580,8 +584,9 @@ async def test_agent_directions_flow_bicycling_mode( mock_runner = mock.MagicMock() mock_fc = MockFunctionCall( - name="set_model_response", + name="render_directions_template", args={ + "heading": "Biking route to work", "summary": "Bike for 25 minutes.", "center_lat": 37.5, "center_lng": 127.0, @@ -611,9 +616,9 @@ async def test_agent_directions_flow_bicycling_mode( template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._extractor_runners["0.9"] = mock_runner - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - results = await self._collect_stream(agent, query="bike to work") + results = await self._collect_stream(agent, query="bike to work") self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -639,8 +644,9 @@ async def test_agent_directions_flow_missing_travel_mode_fallback( ) mock_fc = MockFunctionCall( - name="set_model_response", + name="render_directions_template", args={ + "heading": "Directions to work", "summary": "Typical commute is 45 mins.", "center_lat": 37.5, "center_lng": 127.0, @@ -680,13 +686,10 @@ async def test_agent_directions_flow_missing_travel_mode_fallback( template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._extractor_runners["0.9"] = mock_extractor_runner + agent._fallback_text_runner = mock_fallback_runner - with mock.patch.object( - agent, - "_build_runner", - side_effect=[mock_extractor_runner, mock_fallback_runner], - ): - results = await self._collect_stream(agent, query="directions to work") + results = await self._collect_stream(agent, query="directions to work") self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -721,8 +724,9 @@ async def test_agent_local_search_flow(self, mock_lite_llm_class): mock_runner = mock.MagicMock() mock_fc = MockFunctionCall( - name="set_model_response", + name="render_local_search_template", args={ + "heading": "Top Sushi Places in Seattle", "summary": "Here are some sushi places.", "center_lat": 47.6062, "center_lng": -122.3321, @@ -744,14 +748,13 @@ async def test_agent_local_search_flow(self, mock_lite_llm_class): template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._extractor_runners["0.9"] = mock_runner - # Patch _build_runner - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - results = [] - async for item in agent.stream( - query="sushi Seattle", session_id="session_123", ui_version="v0.9" - ): - results.append(item) + results = [] + async for item in agent.stream( + query="sushi Seattle", session_id="session_123", ui_version="v0.9" + ): + results.append(item) self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -763,6 +766,20 @@ async def test_agent_local_search_flow(self, mock_lite_llm_class): create_surface["surfaceId"].startswith("local-search-surface-") ) + update_components = parts[1].root.data["updateComponents"] + heading_comp = next( + comp + for comp in update_components["components"] + if comp["id"] == "heading-text" + ) + self.assertEqual(heading_comp["text"], "### Top Sushi Places in Seattle") + + map_comp = next( + comp for comp in update_components["components"] if comp["id"] == "map" + ) + self.assertEqual(map_comp["tilt"], 0) + self.assertEqual(map_comp["mode"], "roadmap") + update_data_model = parts[2].root.data["updateDataModel"] # Verify places array was successfully populated in data model self.assertEqual(update_data_model["path"], "/") @@ -788,9 +805,9 @@ async def test_agent_local_search_flow_validation_failure_fallback( mock_runner = mock.MagicMock() - # Mock invalid set_model_response arguments (missing required center_lat) + # Mock invalid render_local_search_template arguments (missing required center_lat) invalid_args = {"summary": "Invalid data", "places": []} - mock_fc = MockFunctionCall("set_model_response", invalid_args) + mock_fc = MockFunctionCall("render_local_search_template", invalid_args) mock_event_fc = MockEvent(function_calls=[mock_fc]) mock_event_text = MockEvent( content=MockContent([MockPart("Fallback text here.")]) @@ -806,14 +823,14 @@ async def test_agent_local_search_flow_validation_failure_fallback( template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._extractor_runners["0.9"] = mock_runner + agent._fallback_text_runner = mock_runner - # Patch _build_runner - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - results = [] - async for item in agent.stream( - query="coffee", session_id="session_123", ui_version="v0.9" - ): - results.append(item) + results = [] + async for item in agent.stream( + query="coffee", session_id="session_123", ui_version="v0.9" + ): + results.append(item) self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -843,8 +860,19 @@ async def test_agent_local_search_flow_catalog_validation_failure_fallback( mock_runner = mock.MagicMock() mock_fc = MockFunctionCall( - "set_model_response", - {"summary": "Coffee", "places": [{"name": "Starbucks"}]}, + "render_local_search_template", + { + "heading": "Coffee Shops", + "summary": "Coffee", + "center_lat": 47.6, + "center_lng": -122.3, + "places": [{ + "placeId": "1", + "name": "Starbucks", + "lat": 47.6, + "lng": -122.3, + }], + }, ) mock_runner.run_async.return_value = MockAsyncIterator( [MockEvent(function_calls=[mock_fc])] @@ -858,18 +886,22 @@ async def test_agent_local_search_flow_catalog_validation_failure_fallback( "Mock validation error" ) mock_schema_manager = mock.MagicMock() - mock_schema_manager.get_catalog.return_value = mock_catalog - agent._schema_managers = {"v0.9": mock_schema_manager} + mock_schema_manager.get_selected_catalog.return_value = mock_catalog + agent._schema_managers = {"0.9": mock_schema_manager} + for agent_inst in agent._extractor_agents.values(): + for t in agent_inst.tools: + if hasattr(t, "schema_manager"): + t.schema_manager = mock_schema_manager mock_fallback_runner = mock.MagicMock() mock_fallback_runner.run_async.return_value = MockAsyncIterator( [MockEvent(content=MockContent([MockPart("Fallback text from LLM.")]))] ) - with mock.patch.object( - agent, "_build_runner", side_effect=[mock_runner, mock_fallback_runner] - ): - results = await self._collect_stream(agent, "coffee") + agent._extractor_runners["0.9"] = mock_runner + agent._fallback_text_runner = mock_fallback_runner + + results = await self._collect_stream(agent, "coffee") self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -899,12 +931,10 @@ async def test_agent_fallback_mode_text_on_extractor_failure( ]) agent = self._setup_agent(fallback_mode="TEXT") - with mock.patch.object( - agent, - "_build_runner", - side_effect=[mock_extractor_runner, mock_fallback_runner], - ): - results = await self._collect_stream(agent, "sushi Seattle") + agent._extractor_runners["0.9"] = mock_extractor_runner + agent._fallback_text_runner = mock_fallback_runner + + results = await self._collect_stream(agent, "sushi Seattle") self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -936,15 +966,13 @@ async def test_agent_fallback_mode_dynamic_on_extractor_failure( ]) agent = self._setup_agent(fallback_mode="DYNAMIC") - with mock.patch.object( - agent, - "_build_runner", - side_effect=[mock_extractor_runner, mock_fallback_runner], - ): - with mock.patch( - "agent.MAUIAgent.stream", - ) as mock_super_stream: - results = await self._collect_stream(agent, "sushi Seattle") + agent._extractor_runners["0.9"] = mock_extractor_runner + agent._fallback_text_runner = mock_fallback_runner + + with mock.patch( + "agent.MAUIAgent.stream", + ) as mock_super_stream: + results = await self._collect_stream(agent, "sushi Seattle") mock_super_stream.assert_not_called() self.assertEqual(len(results), 1) @@ -974,8 +1002,9 @@ async def test_agent_other_spatial_fallback_mode_text( ) ]) agent = self._setup_agent(fallback_mode="TEXT") - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - results = await self._collect_stream(agent, "weather Yosemite") + agent._fallback_text_runner = mock_runner + + results = await self._collect_stream(agent, "weather Yosemite") self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -1046,9 +1075,10 @@ async def test_extractor_duplication_prevented(self, mock_lite_llm_class): template_model="gemini/template-model", ) agent = MAUIAgentWithTemplates(base_url="http://test-url", config=config) + agent._extractor_runners["0.9"] = mock_runner + agent._fallback_text_runner = mock_runner - with mock.patch.object(agent, "_build_runner", return_value=mock_runner): - results = await self._collect_stream(agent, "sushi") + results = await self._collect_stream(agent, "sushi") self.assertEqual(len(results), 1) self.assertTrue(results[0]["is_task_complete"]) @@ -1056,51 +1086,204 @@ async def test_extractor_duplication_prevented(self, mock_lite_llm_class): text_comp = self._get_component_by_id(parts, "text-content") self.assertEqual(text_comp["text"], "I'm sorry") + def test_persistent_extractor_runner_initialized_at_startup(self): + """Verifies persistent extractor agents and runners are pre-built during __init__.""" + agent = MAUIAgentWithTemplates(base_url="http://test-url") + self.assertIn("0.9", agent._extractor_agents) + self.assertIn("0.9", agent._extractor_runners) + self.assertNotIn("v0.9", agent._extractor_agents) + self.assertNotIn("v0.9", agent._extractor_runners) + self.assertIsNotNone(agent._grounded_text_agent) + self.assertIsNotNone(agent._grounded_text_runner) + self.assertIsNotNone(agent._fallback_text_agent) + self.assertIsNotNone(agent._fallback_text_runner) + + # Check tools in unified extractor + unified_agent = agent._extractor_agents["0.9"] + tool_names = [getattr(t, "name", str(t)) for t in unified_agent.tools] + self.assertIn("render_local_search_template", tool_names) + self.assertIn("render_directions_template", tool_names) + self.assertIn("render_text_only_template", tool_names) + + @mock.patch(_LITELLM_PATH) + async def test_persistent_extractor_runner_reused_across_queries( + self, mock_lite_llm_class + ): + """Verifies that the same persistent runner instance is reused across multiple queries.""" + self._setup_mock_llm(mock_lite_llm_class) + + self.mock_router.generate_content_async.return_value = ( + self._mock_llm_stream('{"intent": "TEXT_ONLY", "query": "test"}') + ) + mock_runner = mock.MagicMock() + mock_runner.run_async.return_value = MockAsyncIterator( + [MockEvent(content=MockContent([MockPart("Answer")]))] + ) + + agent = MAUIAgentWithTemplates(base_url="http://test-url") + agent._grounded_text_runner = mock_runner + + # Query 1 + await self._collect_stream(agent, "test 1", session_id="sess_1") + # Query 2 + self.mock_router.generate_content_async.return_value = ( + self._mock_llm_stream('{"intent": "TEXT_ONLY", "query": "test"}') + ) + mock_runner.run_async.return_value = MockAsyncIterator( + [MockEvent(content=MockContent([MockPart("Answer")]))] + ) + await self._collect_stream(agent, "test 2", session_id="sess_2") + + self.assertEqual(mock_runner.run_async.call_count, 2) + # Check that session_id was passed distinctly + call1_kwargs = mock_runner.run_async.call_args_list[0].kwargs + call2_kwargs = mock_runner.run_async.call_args_list[1].kwargs + self.assertEqual(call1_kwargs["session_id"], "sess_1") + self.assertEqual(call2_kwargs["session_id"], "sess_2") + def test_build_runner_sets_auto_create_session(self): agent = MAUIAgentWithTemplates(base_url="http://test-url") mock_agent = mock.MagicMock(spec=LlmAgent) runner = agent._build_runner(mock_agent) # pylint: disable=protected-access self.assertTrue(runner.auto_create_session) - def test_build_dynamic_extractor_agent_appends_shared_guidelines(self): + def test_build_unified_extractor_agent_appends_shared_guidelines(self): """Verifies that shared guidelines are appended to skill instructions.""" agent = MAUIAgentWithTemplates(base_url="http://test-url") with mock.patch( "builtins.open", mock.mock_open(read_data="Shared guidelines content") ) as mock_file: - with mock.patch( - "google.adk.skills.load_skill_from_dir" - ) as mock_load_skill: - mock_skill = mock.MagicMock() - mock_skill.instructions = "Base skill instructions" - mock_load_skill.return_value = mock_skill - - extractor_agent = agent._build_dynamic_extractor_agent( # pylint: disable=protected-access - "local-search-template-response" - ) - self.assertIn("Shared guidelines content", extractor_agent.instruction) - self.assertIn("Base skill instructions", extractor_agent.instruction) - mock_file.assert_called_once() + extractor_agent = ( + agent._build_unified_extractor_agent() # pylint: disable=protected-access + ) + self.assertIn("Shared guidelines content", extractor_agent.instruction) + self.assertIn("Directions Guidelines", extractor_agent.instruction) + mock_file.assert_called_once() - def test_build_dynamic_extractor_agent_handles_file_read_error(self): + def test_build_unified_extractor_agent_handles_file_read_error(self): """Verifies that file read errors are handled gracefully when loading guidelines.""" agent = MAUIAgentWithTemplates(base_url="http://test-url") with mock.patch("builtins.open", side_effect=OSError("Read error")): - with mock.patch( - "google.adk.skills.load_skill_from_dir" - ) as mock_load_skill: - mock_skill = mock.MagicMock() - mock_skill.instructions = "Base skill instructions" - mock_load_skill.return_value = mock_skill - - # Check that it handles OSError gracefully and proceeds - extractor_agent = agent._build_dynamic_extractor_agent( # pylint: disable=protected-access - "local-search-template-response" - ) - self.assertNotIn( - "Shared guidelines content", extractor_agent.instruction - ) - self.assertIn("Base skill instructions", extractor_agent.instruction) + # Check that it handles OSError gracefully and proceeds + extractor_agent = ( + agent._build_unified_extractor_agent() # pylint: disable=protected-access + ) + self.assertIsNotNone(extractor_agent) + + def test_build_unified_extractor_agent_directions_loads_skill_instructions( + self, + ): + """Verifies that directions skill instructions from disk are loaded into the extractor agent.""" + agent = MAUIAgentWithTemplates(base_url="http://test-url") + extractor_agent = ( + agent._build_unified_extractor_agent() # pylint: disable=protected-access + ) + self.assertIn("less than a minute", extractor_agent.instruction) + self.assertIn( + "Always round seconds to the nearest minute", + extractor_agent.instruction, + ) + tool_names = [t.name for t in extractor_agent.tools if hasattr(t, "name")] + self.assertIn("render_directions_template", tool_names) + + def test_build_unified_extractor_instruction_target_intent_rules(self): + """Verifies that unified instruction defines explicit target intent rules.""" + maui_agent = MAUIAgentWithTemplates(base_url="http://test-url") + # pylint: disable=protected-access + instruction = maui_agent._build_unified_extractor_instruction() + # pylint: enable=protected-access + expected_workflow_rules = ( + "## Execution and Tool Usage Rules\n" + "1. Follow the previous rules for gathering the necessary information" + " using Google Maps tools.\n" + "2. When the prompt begins with `[TARGET_INTENT: <INTENT>]`, you" + " MUST strictly invoke the corresponding template rendering tool as" + " your final step:\n" + " - `[TARGET_INTENT: DIRECTIONS]`: Invoke" + " `render_directions_template`. Do NOT invoke any other template" + " rendering tool.\n" + " - `[TARGET_INTENT: LOCAL_SEARCH]`: Invoke" + " `render_local_search_template`. Do NOT invoke any other template" + " rendering tool.\n" + "3. Your final action MUST be this template rendering tool call. Do" + " NOT emit raw JSON or `<a2ui-json>` blocks in your text response." + ) + self.assertIn(expected_workflow_rules, instruction) + + def test_build_unified_extractor_agent_bypasses_schema_instructions(self): + """Verifies schema instructions are bypassed and tools receive schema_manager.""" + maui_agent = MAUIAgentWithTemplates(base_url="http://test-url") + mock_schema_manager = mock.MagicMock( + spec=agent_with_templates.A2uiSchemaManager + ) + mock_schema_manager.generate_system_prompt.return_value = ( + "Each A2UI JSON block MUST be wrapped in <a2ui-json> and </a2ui-json>" + " tags." + ) + # pylint: disable=protected-access + extractor_agent = maui_agent._build_unified_extractor_agent( + schema_manager=mock_schema_manager + ) + expected_instruction = maui_agent._build_unified_extractor_instruction() + # pylint: enable=protected-access + mock_schema_manager.generate_system_prompt.assert_not_called() + self.assertNotIn( + "Each A2UI JSON block MUST be wrapped", + extractor_agent.instruction, + ) + self.assertEqual( + extractor_agent.instruction, + expected_instruction, + ) + template_tools = [ + tool + for tool in extractor_agent.tools + if isinstance(tool, agent_with_templates.BaseTemplateTool) + ] + self.assertTrue(template_tools) + for tool in template_tools: + self.assertEqual(tool.schema_manager, mock_schema_manager) + + def test_template_registry_definitions(self): + """Verifies that template_registry entries are consistent and valid.""" + agent = MAUIAgentWithTemplates(base_url="http://test-url") + self.assertIn("LOCAL_SEARCH", agent.template_registry.supported_intents) + self.assertIn("DIRECTIONS", agent.template_registry.supported_intents) + for bundle in agent.template_registry.bundles.values(): + self.assertTrue(bundle.name) + self.assertTrue(bundle.surface_prefix) + self.assertTrue( + bundle.skill_dir.is_dir(), + f"Skill directory {bundle.skill_dir} does not exist", + ) + + def test_extractor_agent_mounts_tools_for_all_bundles(self): + """Verifies that the extractor agent mounts a BaseTemplateTool for every bundle.""" + agent = MAUIAgentWithTemplates(base_url="http://test-url") + extractor_agent = ( + agent._build_unified_extractor_agent() # pylint: disable=protected-access + ) + tools_by_name = { + t.name: t for t in extractor_agent.tools if hasattr(t, "name") + } + for bundle in agent.template_registry.bundles.values(): + expected_name = bundle.render_tool_name + self.assertIn(expected_name, tools_by_name) + tool = tools_by_name[expected_name] + self.assertIsInstance(tool, template_tool.BaseTemplateTool) + self.assertEqual(tool.template_name, bundle.name) + self.assertEqual(tool.schema_class, bundle.schema_cls) + self.assertEqual(tool.description, bundle.effective_tool_description) + + def test_normalize_version(self): + """Verifies version strings collapse to the unprefixed schema key.""" + normalize = agent_with_templates._normalize_version # pylint: disable=protected-access + self.assertEqual(normalize("0.9"), "0.9") + self.assertEqual(normalize("v0.9"), "0.9") + self.assertEqual(normalize("v0.9.1"), "0.9.1") + self.assertEqual(normalize(None), agent_with_templates.VERSION_0_9) + self.assertEqual(normalize(""), agent_with_templates.VERSION_0_9) + if __name__ == "__main__": unittest.main() diff --git a/agent/python_agent/test_merger.py b/agent/python_agent/test_merger.py index f5e4f01..8463452 100644 --- a/agent/python_agent/test_merger.py +++ b/agent/python_agent/test_merger.py @@ -14,9 +14,17 @@ """Tests for A2UI Layout Template Merger (`merger.py`).""" +import importlib +import importlib.resources +import json +import os import pathlib import re +import shutil +import sys +import tempfile import unittest +import zipfile import a2ui import agent @@ -58,6 +66,79 @@ def test_merge_unknown_template_raises_error(self): with self.assertRaises(FileNotFoundError): merge_template("non_existent_template", {"text": "hello"}) + def test_merge_defaults_to_maps_catalog_id(self): + """Verifies the surface binds to the Maps catalog when none is named.""" + result = merge_template("text_only", {"text": "hello"}) + self.assertEqual( + result[0]["createSurface"]["catalogId"], + "a2ui://maps-agentic-ui-catalog.json", + ) + self.assertEqual( + result[0]["createSurface"]["catalogId"], merger.DEFAULT_CATALOG_ID + ) + + def test_merge_honors_custom_catalog_id(self): + """Verifies a host-supplied catalog id reaches every template's surface.""" + host_catalog = "https://example.test/host_catalog.json" + local_search_data = { + "surface_id": "test-surface", + "summary": "Two places.", + "center_lat": 37.0, + "center_lng": -122.0, + "zoom": 12, + "places": [{"name": "A", "lat": 37.0, "lng": -122.0}], + } + directions_data = { + "surface_id": "test-surface", + "summary": "A route.", + "center_lat": 37.0, + "center_lng": -122.0, + "zoom": 12, + "routes": [{ + "origin": {"lat": 37.0, "lng": -122.0, "label": "A"}, + "destination": {"lat": 37.1, "lng": -122.1, "label": "B"}, + }], + } + for template_name, data in ( + ("text_only", {"text": "hello"}), + ("local_search", local_search_data), + ("directions", directions_data), + ): + with self.subTest(template=template_name): + result = merge_template(template_name, data, catalog_id=host_catalog) + self.assertEqual(result[0]["createSurface"]["catalogId"], host_catalog) + + def test_text_only_fallback_keeps_custom_catalog_id(self): + """Verifies the catalog id survives the bail-out to `text_only`. + + The prepare helpers discard the caller's dict and build a fresh one when + they fall back, so a catalog id applied before that point would be lost. + Losing it is silent: the placeholder resolves to None and the key is + stripped, leaving a surface bound to no catalog at all. + """ + host_catalog = "https://example.test/host_catalog.json" + invalid_local_search = { + "surface_id": "test-surface", + "summary": "Short response.", + "center_lat": 37.0, + "center_lng": 127.0, + "zoom": 10, + "places": "NOT_A_LIST", + } + invalid_directions = { + "surface_id": "test-surface", + "summary": "Cannot compute directions.", + } + for template_name, data in ( + ("local_search", invalid_local_search), + ("directions", invalid_directions), + ): + with self.subTest(template=template_name): + result = merge_template(template_name, data, catalog_id=host_catalog) + create_surface = result[0]["createSurface"] + self.assertIn("catalogId", create_surface) + self.assertEqual(create_surface["catalogId"], host_catalog) + def test_merge_text_only_full_json(self): """Verifies that text-only response is merged correctly.""" data = { @@ -129,6 +210,41 @@ def test_merge_generic_surface_id_generates_random_suffix(self): surface_id = result[0]["createSurface"]["surfaceId"] self.assertTrue(re.fullmatch(r"text-only-surface_[a-f0-9]{6}", surface_id)) + def test_merge_surface_prefix_collision_generates_random_suffix(self): + """Verifies that when surface_id equals bundle surface_prefix, a suffix is + + appended. + """ + data = { + "surface_id": "local-search-surface", + "heading": "Top Places", + "summary": "Places found.", + "center_lat": 47.6, + "center_lng": -122.3, + "zoom": 12, + "places": [{ + "placeId": "places/1", + "name": "Place 1", + "lat": 47.6, + "lng": -122.3, + }], + } + result = merge_template("local_search", data) + surface_id = result[0]["createSurface"]["surfaceId"] + self.assertTrue( + re.fullmatch(r"local-search-surface_[a-f0-9]{6}", surface_id) + ) + + def test_uniquify_surface_id_with_custom_surface_prefix(self): + """Verifies that a custom surface_prefix is recognized as a default id.""" + data = {"surface_id": "custom-bundle-prefix"} + merger._uniquify_surface_id( + data, "my_template", surface_prefix="custom-bundle-prefix" + ) + self.assertTrue( + re.fullmatch(r"custom-bundle-prefix_[a-f0-9]{6}", data["surface_id"]) + ) + def test_merge_custom_surface_id_remains_intact(self): """Verifies that a custom unique surface_id is preserved exactly.""" data = { @@ -143,6 +259,7 @@ def test_merge_local_search_full_json(self): """Verifies merging a complete local search payload.""" data = { "surface_id": "local-search-surface-abc", + "heading": "Top Coffee Shops in Seattle", "summary": "Here are 3 highly-rated coffee shops in Seattle.", "center_lat": "47.6062", "center_lng": -122.3321, @@ -185,7 +302,18 @@ def test_merge_local_search_full_json(self): { "id": "root", "component": "Column", - "children": ["summary-text", "map", "list"], + "children": [ + "heading-text", + "summary-text", + "map", + "list", + ], + }, + { + "id": "heading-text", + "component": "Text", + "variant": "body", + "text": "### Top Coffee Shops in Seattle", }, { "id": "summary-text", @@ -200,6 +328,8 @@ def test_merge_local_search_full_json(self): "component": "GoogleMap", "center": {"lat": 47.6062, "lng": -122.3321}, "zoom": 14, + "tilt": 0, + "mode": "roadmap", "markers": [ { "lat": 47.62, @@ -291,10 +421,54 @@ def test_validate_local_search_output_with_schema(self): self._validate_schema(result) + def test_place_details_orientation_kept_for_compatibility(self): + """Both `horizontal` and `vertical` orientations validate; default is horizontal.""" + data = { + "surface_id": "local-search-surface-abc", + "summary": "Here are coffee shops.", + "center_lat": 47.6062, + "center_lng": -122.3321, + "zoom": 14, + "places": [{ + "placeId": "ChIJ111", + "name": "Espresso Vivace", + "lat": 47.62, + "lng": -122.32, + }], + } + for orientation in ("horizontal", "vertical"): + with self.subTest(orientation=orientation): + result = merge_template("local_search", data) + for message in result: + for component in message.get("updateComponents", {}).get( + "components", [] + ): + if component["component"] == "PlaceDetailsCompact": + component["orientation"] = orientation + + self._validate_schema(result) + + extension = json.loads( + ( + pathlib.Path(__file__).parent + / "shared" + / "schema" + / "maps_catalog_extension.json" + ).read_text() + ) + orientation_schema = extension["components"]["PlaceDetailsCompact"][ + "allOf" + ][2]["properties"]["orientation"] + self.assertEqual(orientation_schema["default"], "horizontal") + self.assertCountEqual( + orientation_schema["enum"], ["horizontal", "vertical"] + ) + def test_merge_max_list_size_slicing(self): """Verifies that max_list_size parameter slices the places and markers list.""" data = { "surface_id": "test-surface", + "heading": "Nearby Places", "summary": "Here are some places.", "center_lat": 47.6062, "center_lng": -122.3321, @@ -308,7 +482,7 @@ def test_merge_max_list_size_slicing(self): result = merge_template("local_search", data, max_list_size=2) # Check that updateComponents has only 2 markers components = result[1]["updateComponents"]["components"] - map_comp = next(c for c in components if c["id"] == "map") + map_comp = next(comp for comp in components if comp["id"] == "map") self.assertEqual(len(map_comp["markers"]), 2) # Check that updateDataModel has only 2 places @@ -317,10 +491,63 @@ def test_merge_max_list_size_slicing(self): self.assertEqual(places[0]["placeId"], "1") self.assertEqual(places[1]["placeId"], "2") + def test_merge_local_search_heading_normalization(self): + """Verifies that heading markdown headers are cleaned. + + Also checks missing heading fallback. + """ + # Case 1: Heading with leading markdown hashtags + data_with_hash = { + "surface_id": "test-surface", + "heading": "### Best Bakeries", + "summary": "Here are bakeries.", + "center_lat": 47.6, + "center_lng": -122.3, + "zoom": 13, + "places": [{ + "placeId": "p1", + "name": "B1", + "lat": 47.6, + "lng": -122.3, + }], + } + result = merge_template("local_search", data_with_hash) + comps = result[1]["updateComponents"]["components"] + heading_comp = next(comp for comp in comps if comp["id"] == "heading-text") + self.assertEqual(heading_comp["text"], "### Best Bakeries") + + # Case 2: Missing heading with anchor marker degrades to text fallback + data_with_anchor = { + "surface_id": "test-surface", + "summary": "Here are bakeries.", + "center_lat": 47.6, + "center_lng": -122.3, + "zoom": 13, + "anchor_marker": {"lat": 47.6, "lng": -122.3, "label": "Space Needle"}, + "places": [{"placeId": "p1", "name": "B1", "lat": 47.6, "lng": -122.3}], + } + result = merge_template("local_search", data_with_anchor) + comps = result[1]["updateComponents"]["components"] + self.assertEqual(comps[1]["text"], "Here are bakeries.") + + # Case 3: Missing heading and no anchor degrades to text fallback + data_no_heading = { + "surface_id": "test-surface", + "summary": "Here are bakeries.", + "center_lat": 47.6, + "center_lng": -122.3, + "zoom": 13, + "places": [{"placeId": "p1", "name": "B1", "lat": 47.6, "lng": -122.3}], + } + result = merge_template("local_search", data_no_heading) + comps = result[1]["updateComponents"]["components"] + self.assertEqual(comps[1]["text"], "Here are bakeries.") + def test_merge_directions_full_json(self): """Verifies complete end-to-end directions template merging, placeholder replacement, and travel mode normalization.""" data = { "surface_id": "directions-surface-xyz", + "heading": "Walking Route from Dobong to Gangnam", "summary": "Typical commute is 1h 15m.", "center_lat": "37.5665", "center_lng": 126.9780, @@ -352,13 +579,13 @@ def test_merge_directions_full_json(self): { "id": "root", "component": "Column", - "children": ["summary-text", "map"], + "children": ["heading-text", "map", "summary-text"], }, { - "id": "summary-text", + "id": "heading-text", "component": "Text", "variant": "body", - "text": "Typical commute is 1h 15m.", + "text": "### Walking Route from Dobong to Gangnam", }, { "id": "map", @@ -379,6 +606,12 @@ def test_merge_directions_full_json(self): }], "travelMode": "walking", }, + { + "id": "summary-text", + "component": "Text", + "variant": "body", + "text": "Typical commute is 1h 15m.", + }, ], }, }, @@ -395,6 +628,48 @@ def test_merge_directions_full_json(self): result = merge_template("directions", data, max_list_size=3) self.assertEqual(result, expected) + def test_merge_directions_heading_fallback(self): + """Verifies that missing heading falls back to text-only surface.""" + # Case 1: Heading with leading markdown hashtags + data_with_hash = { + "surface_id": "test-surface", + "heading": "### Driving Route", + "summary": "About 15 minutes.", + "center_lat": 37.5, + "center_lng": 127.0, + "zoom": 12, + "travel_mode": "driving", + "routes": [{ + "origin": {"lat": 37.5, "lng": 127.0, "label": "Origin"}, + "destination": {"lat": 37.6, "lng": 127.1, "label": "Dest"}, + }], + } + result = merge_template("directions", data_with_hash) + comps = result[1]["updateComponents"]["components"] + heading_comp = next(c for c in comps if c["id"] == "heading-text") + self.assertEqual(heading_comp["text"], "### Driving Route") + + # Case 2: Missing heading degrades to text-only fallback + data_missing = { + "surface_id": "test-surface", + "summary": "About 15 minutes.", + "center_lat": 37.5, + "center_lng": 127.0, + "zoom": 12, + "travel_mode": "driving", + "routes": [{ + "origin": {"lat": 37.5, "lng": 127.0, "label": "Seattle Center"}, + "destination": { + "lat": 37.6, + "lng": 127.1, + "label": "Pike Place Market", + }, + }], + } + result = merge_template("directions", data_missing) + comps = result[1]["updateComponents"]["components"] + self.assertEqual(comps[1]["text"], "About 15 minutes.") + def test_validate_directions_output_with_schema(self): """Verifies merged directions output passes schema validation.""" data = { @@ -478,10 +753,11 @@ def test_missing_keys_no_crash_directions(self): "Cannot compute directions.", ) - def test_unrecognized_travel_mode(self): - """Verifies that unrecognized travel modes are ignored and not passed to output.""" + def test_unrecognized_travel_mode_triggers_text_fallback(self): + """Verifies that an unroutable travel mode triggers text fallback.""" data = { "surface_id": "test-surface", + "heading": "Directions", "summary": "Short response.", "center_lat": 37.0, "center_lng": 127.0, @@ -493,12 +769,32 @@ def test_unrecognized_travel_mode(self): "travel_mode": "TELEPORT", } result = merge_template("directions", data, max_list_size=3) - map_comp = next( - c - for c in result[1]["updateComponents"]["components"] - if c["component"] == "GoogleMap" + self.assertEqual(len(result), 2) + self.assertEqual( + result[1]["updateComponents"]["components"][1]["text"], + "Short response.", + ) + + def test_omitted_travel_mode_triggers_text_fallback(self): + """Verifies that omitting the travel mode triggers text fallback.""" + data = { + "surface_id": "test-surface", + "heading": "Directions", + "summary": "Short response.", + "center_lat": 37.0, + "center_lng": 127.0, + "zoom": 10, + "routes": [{ + "origin": {"lat": 37.0, "lng": 127.0, "label": "Start"}, + "destination": {"lat": 37.1, "lng": 127.1, "label": "End"}, + }], + } + result = merge_template("directions", data, max_list_size=3) + self.assertEqual(len(result), 2) + self.assertEqual( + result[1]["updateComponents"]["components"][1]["text"], + "Short response.", ) - self.assertNotIn("travelMode", map_comp) def test_travel_mode_synonyms_normalization(self): """Verifies that various travel mode synonyms normalize properly in merger.""" @@ -514,6 +810,7 @@ def test_travel_mode_synonyms_normalization(self): with self.subTest(raw_mode=raw_mode, expected_mode=expected_mode): data = { "surface_id": "test-surface", + "heading": "Directions", "summary": "Commute.", "center_lat": 37.0, "center_lng": 127.0, @@ -570,6 +867,7 @@ def test_missing_optional_placeholders_are_stripped(self): """Verifies that optional placeholders are omitted when missing from input data.""" data = { "surface_id": "test-surface", + "heading": "Top Places", "summary": "Results for sushi.", "center_lat": 47.6062, "center_lng": -122.3321, @@ -584,15 +882,43 @@ def test_missing_optional_placeholders_are_stripped(self): result = merge_template("local_search", data, max_list_size=3) update_components = result[1]["updateComponents"] map_comp = next( - c for c in update_components["components"] if c["id"] == "map" + comp for comp in update_components["components"] if comp["id"] == "map" ) # Verify anchorMarker key is NOT in map component (cleanly stripped) self.assertNotIn("anchorMarker", map_comp) - def test_markers_explicitly_provided_and_sanitized(self): - """Verifies that explicitly provided markers are used and sanitized.""" + def test_explicit_none_field_does_not_leak_into_wire_json(self): + """Verifies explicit None values in schema fields do not leak null to wire + + JSON. + """ data = { "surface_id": "test-surface", + "heading": "Top Places", + "summary": "Results for sushi.", + "center_lat": 47.6062, + "center_lng": -122.3321, + "zoom": 14, + "anchor_marker": None, + "places": [{ + "placeId": "ChIJ111", + "name": "Espresso Vivace", + "lat": 47.6200, + "lng": -122.3200, + }], + } + result = merge_template("local_search", data, max_list_size=3) + update_components = result[1]["updateComponents"] + map_comp = next( + comp for comp in update_components["components"] if comp["id"] == "map" + ) + self.assertNotIn("anchorMarker", map_comp) + + def test_markers_are_derived_from_places(self): + """Verifies markers are projected from places, ignoring any override.""" + data = { + "surface_id": "test-surface", + "heading": "Top Places", "summary": "Results with custom markers.", "center_lat": 47.6062, "center_lng": -122.3321, @@ -603,21 +929,19 @@ def test_markers_explicitly_provided_and_sanitized(self): "lat": 47.62, "lng": -122.32, }], - "markers": [ - {"lat": "47.63", "lng": "-122.33", "label": "Custom 1"}, - {"lat": 47.64, "lng": -122.34, "label": None}, - {"invalid_marker": "yes"}, - ], + "markers": [{"lat": 47.63, "lng": -122.33, "label": "Custom 1"}], } result = merge_template("local_search", data, max_list_size=3) update_components = result[1]["updateComponents"] map_comp = next( - c for c in update_components["components"] if c["id"] == "map" + comp for comp in update_components["components"] if comp["id"] == "map" ) - expected_markers = [ - {"lat": 47.63, "lng": -122.33, "label": "Custom 1"}, - {"lat": 47.64, "lng": -122.34, "label": ""}, - ] + expected_markers = [{ + "lat": 47.62, + "lng": -122.32, + "label": "Espresso Vivace", + "placeId": "ChIJ111", + }] self.assertEqual(map_comp["markers"], expected_markers) def test_fallback_preserves_surface_id(self): @@ -658,10 +982,12 @@ def test_merge_directions_with_direct_routes(self): """Verifies directions template merging when routes are directly passed.""" data = { "surface_id": "directions-surface-xyz", + "heading": "Directions", "summary": "Commute is 1h.", "center_lat": 37.5665, "center_lng": 126.9780, "zoom": 12, + "travel_mode": "driving", "routes": [ { "origin": {"lat": 37.6700, "lng": 127.0400, "label": "A"}, @@ -704,6 +1030,137 @@ def test_merge_directions_malformed_routes_fallback(self): "Fallback text.", ) + def test_merge_fallback_surface_id_is_uniquified(self): + """Verifies fallback surfaces without caller id receive a unique suffix.""" + data = {"places": []} + result = merge_template("local_search", data) + surface_id = result[0]["createSurface"]["surfaceId"] + self.assertTrue(surface_id.startswith("fallback-surface_")) + + +class TestMergerZipImportedPackage(unittest.TestCase): + """Covers template loading when the agent is served from a zip archive. + + The agent also ships as a zipped archive that Python imports without ever + unpacking it to disk. In that mode `__file__` names an entry inside the + archive, so resolving templates by filesystem path makes every merge fail at + serving time while a suite run from an ordinary source tree keeps passing. + These tests exercise the archive explicitly so that gap cannot reopen. + """ + + _PACKAGE = "a2ui_zipped_merger_probe" + + def _pack_merger_into_zip(self) -> str: + """Copies `merger` and its templates into a zip and returns its path.""" + source = importlib.resources.files(merger.__package__) + temp_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True) + archive_path = os.path.join(temp_dir, "merger_probe.zip") + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr(f"{self._PACKAGE}/__init__.py", "") + archive.writestr( + f"{self._PACKAGE}/merger.py", + source.joinpath("merger.py").read_text(encoding="utf-8"), + ) + for entry in source.joinpath("templates").iterdir(): + layout = entry.joinpath("layout.json") + if not layout.is_file(): + continue + archive.writestr( + f"{self._PACKAGE}/templates/{entry.name}/layout.json", + layout.read_text(encoding="utf-8"), + ) + return archive_path + + def _import_merger_from_zip(self): + """Imports the packed `merger`, whose templates exist only in the zip.""" + archive_path = self._pack_merger_into_zip() + sys.path.insert(0, archive_path) + self.addCleanup(sys.path.remove, archive_path) + for name in (f"{self._PACKAGE}.merger", self._PACKAGE): + self.addCleanup(sys.modules.pop, name, None) + importlib.invalidate_caches() + return importlib.import_module(f"{self._PACKAGE}.merger") + + def test_probe_templates_are_reachable_only_inside_the_archive(self): + """Asserts the probe exercises the archive rather than a copy on disk. + + Without this the remaining tests could pass against an unpacked directory + and prove nothing, because that is the mode that never broke. + """ + zipped_merger = self._import_merger_from_zip() + + on_disk_templates = os.path.join( + os.path.dirname(os.path.abspath(zipped_merger.__file__)), "templates" + ) + + self.assertNotEqual(zipped_merger.__file__, merger.__file__) + self.assertFalse(os.path.exists(on_disk_templates)) + + def test_merges_every_shipped_template_from_the_archive(self): + """Verifies each template resolves as package data inside the archive.""" + zipped_merger = self._import_merger_from_zip() + cases = { + "text_only": {"text": "hello"}, + "local_search": { + "summary": "Coffee near the market.", + "center_lat": 47.6097, + "center_lng": -122.3422, + "zoom": 14, + "places": [ + {"name": "Storyville", "lat": 47.6092, "lng": -122.3418} + ], + }, + "directions": { + "summary": "About 50 minutes via US-101 S.", + "center_lat": 37.55, + "center_lng": -122.15, + "zoom": 10, + "travel_mode": "driving", + "routes": [{ + "origin": {"lat": 37.7749, "lng": -122.4194, "label": "SF"}, + "destination": { + "lat": 37.3382, + "lng": -121.8863, + "label": "SJ", + }, + }], + }, + } + + for template_name, data in cases.items(): + with self.subTest(template=template_name): + result = zipped_merger.merge_template(template_name, data) + + self.assertTrue(result) + self.assertEqual( + result[0]["createSurface"]["catalogId"], merger.DEFAULT_CATALOG_ID + ) + + def test_unknown_template_still_raises_from_the_archive(self): + """A missing template must stay distinguishable from an unreadable one.""" + zipped_merger = self._import_merger_from_zip() + + with self.assertRaises(FileNotFoundError): + zipped_merger.merge_template("non_existent_template", {"text": "hello"}) + + def test_default_surface_id_is_uniquified_from_the_archive(self): + """Covers the template scan that decides whether a surface ID is generic. + + The scan sat behind a filesystem probe that reported an empty directory + inside an archive instead of failing, so default surface IDs quietly + stopped being made unique per turn and no log line recorded it. + """ + zipped_merger = self._import_merger_from_zip() + + result = zipped_merger.merge_template( + "text_only", {"text": "hello", "surface_id": "text_only_surface"} + ) + + surface_id = result[0]["createSurface"]["surfaceId"] + self.assertNotEqual(surface_id, "text_only_surface") + self.assertTrue(surface_id.startswith("text_only_surface_")) + if __name__ == "__main__": unittest.main() diff --git a/agent/python_agent/test_place_id_resolution.py b/agent/python_agent/test_place_id_resolution.py new file mode 100644 index 0000000..1064182 --- /dev/null +++ b/agent/python_agent/test_place_id_resolution.py @@ -0,0 +1,383 @@ +# Copyright 2026 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. + +"""Unit tests for Place ID placeholder substitution.""" + +import unittest +from unittest import mock + +import place_id_resolution + + +def _create_mock_response( + chunks: list[tuple[str, str, str]] | None = None, +) -> mock.MagicMock: + """Builds a mock genai response carrying Maps attribution sources. + + The shape mirrors production: every Maps source arrives with a non-empty + `place_id`, which is what makes placeholder substitution viable at all. + + Args: + chunks: List of (title, place_id) tuples. + + Returns: + Mock response matching the Google GenAI SDK candidate structure. + """ + mock_chunks = [] + for title, place_id in chunks or []: + mock_chunk = mock.MagicMock() + mock_chunk.maps.title = title + mock_chunk.maps.place_id = place_id + mock_chunks.append(mock_chunk) + + mock_meta = mock.MagicMock() + mock_meta.grounding_chunks = mock_chunks + mock_candidate = mock.MagicMock() + mock_candidate.grounding_metadata = mock_meta + mock_response = mock.MagicMock() + mock_response.candidates = [mock_candidate] + return mock_response + + +def _chunk(title: str, place_id: str) -> place_id_resolution.AttributionSource: + """Shorthand for constructing an attribution source.""" + return place_id_resolution.AttributionSource(title=title, place_id=place_id) + + +class NormalizePlaceIdTest(unittest.TestCase): + + def test_strips_places_resource_prefix(self): + self.assertEqual( + place_id_resolution.normalize_place_id( + "places/ChIJN1t_tDeuEmsRUsoyG83frY4" + ), + "ChIJN1t_tDeuEmsRUsoyG83frY4", + ) + + def test_leaves_bare_id_untouched(self): + self.assertEqual( + place_id_resolution.normalize_place_id("ChIJN1t_tDeuEmsRUsoyG83frY4"), + "ChIJN1t_tDeuEmsRUsoyG83frY4", + ) + + def test_leaves_prefixed_id_untouched_when_it_is_not_a_chi_id(self): + # Pins shipped behavior. The prefix check demands "places/ChI", so an ID + # from any other family keeps its resource prefix. + self.assertEqual( + place_id_resolution.normalize_place_id("places/GhIJabc"), + "places/GhIJabc", + ) + + +class CanonicalPlaceTitleTest(unittest.TestCase): + + def test_strips_google_maps_branding_suffix(self): + self.assertEqual( + place_id_resolution.canonical_place_title("Chez Panisse - Google Maps"), + "Chez Panisse", + ) + + def test_leaves_branding_suffix_written_with_an_en_dash(self): + # Pins shipped behavior. Only the hyphen form is recognized. + self.assertEqual( + place_id_resolution.canonical_place_title( + "Chez Panisse \u2013 Google Maps" + ), + "Chez Panisse \u2013 Google Maps", + ) + + def test_preserves_original_casing(self): + # The prompt tells the model to copy the title character for character, so + # lowercasing here would make every substitution key miss. + self.assertEqual( + place_id_resolution.canonical_place_title("Starbucks Coffee Company"), + "Starbucks Coffee Company", + ) + + def test_leaves_review_prefix_in_place(self): + # Pins shipped behavior. Grounding emits user reviews as their own sources, + # and they canonicalize to a title distinct from the venue's. + self.assertEqual( + place_id_resolution.canonical_place_title( + "Review of Joe's Pizza - Google Maps" + ), + "Review of Joe's Pizza", + ) + + def test_leaves_unbranded_title_untouched(self): + self.assertEqual( + place_id_resolution.canonical_place_title("Joe's Pizza"), "Joe's Pizza" + ) + + +class ExtractAttributionSourcesTest(unittest.TestCase): + + def test_extracts_title_and_place_id_in_arrival_order(self): + response = _create_mock_response([ + ("Joe's Pizza - Google Maps", "places/ChIJ_joe"), + ("Prince St. Pizza - Google Maps", "places/ChIJ_prince"), + ]) + + chunks = place_id_resolution.extract_attribution_sources(response) + + self.assertEqual( + chunks, + [ + _chunk("Joe's Pizza - Google Maps", "places/ChIJ_joe"), + _chunk("Prince St. Pizza - Google Maps", "places/ChIJ_prince"), + ], + ) + + def test_returns_empty_when_response_has_no_candidates(self): + response = mock.MagicMock() + response.candidates = [] + self.assertEqual( + place_id_resolution.extract_attribution_sources(response), [] + ) + + def test_returns_empty_when_grounding_metadata_absent(self): + response = mock.MagicMock() + candidate = mock.MagicMock() + candidate.grounding_metadata = None + response.candidates = [candidate] + self.assertEqual( + place_id_resolution.extract_attribution_sources(response), [] + ) + + def test_skips_sources_missing_a_place_id(self): + response = _create_mock_response([ + ("Joe's Pizza - Google Maps", "places/ChIJ_joe"), + ("A Web Result", ""), + ]) + + chunks = place_id_resolution.extract_attribution_sources(response) + + self.assertEqual(len(chunks), 1) + self.assertEqual(chunks[0].title, "Joe's Pizza - Google Maps") + + +class BuildPlaceholderIndexTest(unittest.TestCase): + + def test_five_same_titled_venues_get_five_distinct_ids(self): + chunks = [ + _chunk("Starbucks - Google Maps", f"places/ChIJ_sbux_{i}") + for i in range(1, 6) + ] + + placeholder_index = place_id_resolution._build_placeholder_index(chunks) + + self.assertEqual( + placeholder_index, + { + "PLACE_ID_FOR_1_Starbucks": "ChIJ_sbux_1", + "PLACE_ID_FOR_2_Starbucks": "ChIJ_sbux_2", + "PLACE_ID_FOR_3_Starbucks": "ChIJ_sbux_3", + "PLACE_ID_FOR_4_Starbucks": "ChIJ_sbux_4", + "PLACE_ID_FOR_5_Starbucks": "ChIJ_sbux_5", + }, + ) + + def test_review_source_lands_in_its_own_title_bucket(self): + # Pins shipped behavior. "Review of X" canonicalizes to a title distinct + # from "X", so it gets its own counter and produces a key the model never + # emits. Venue ordinals are only disturbed when a venue surfaces + # exclusively as a review source. + chunks = [ + _chunk("Starbucks - Google Maps", "places/ChIJ_sbux_1"), + _chunk("Review of Starbucks - Google Maps", "places/ChIJ_sbux_1"), + _chunk("Starbucks - Google Maps", "places/ChIJ_sbux_2"), + ] + + placeholder_index = place_id_resolution._build_placeholder_index(chunks) + + self.assertEqual( + placeholder_index, + { + "PLACE_ID_FOR_1_Starbucks": "ChIJ_sbux_1", + "PLACE_ID_FOR_1_Review of Starbucks": "ChIJ_sbux_1", + "PLACE_ID_FOR_2_Starbucks": "ChIJ_sbux_2", + }, + ) + + def test_distinct_titles_each_start_at_ordinal_one(self): + chunks = [ + _chunk("Joe's Pizza - Google Maps", "places/ChIJ_joe"), + _chunk("Prince St. Pizza - Google Maps", "places/ChIJ_prince"), + ] + + placeholder_index = place_id_resolution._build_placeholder_index(chunks) + + self.assertEqual( + placeholder_index, + { + "PLACE_ID_FOR_1_Joe's Pizza": "ChIJ_joe", + "PLACE_ID_FOR_1_Prince St. Pizza": "ChIJ_prince", + }, + ) + + def test_returns_empty_map_for_no_sources(self): + self.assertEqual(place_id_resolution._build_placeholder_index([]), {}) + + +class ResolvePlaceIdsTest(unittest.TestCase): + + def test_rewrites_placeholders_with_grounded_place_ids(self): + text = ( + '{"places": [{"placeId": "PLACE_ID_FOR_1_Joe\'s Pizza"},' + ' {"placeId": "PLACE_ID_FOR_1_Prince St. Pizza"}]}' + ) + chunks = [ + _chunk("Joe's Pizza - Google Maps", "places/ChIJ_joe"), + _chunk("Prince St. Pizza - Google Maps", "places/ChIJ_prince"), + ] + + result, unresolved = place_id_resolution.resolve_place_ids(text, chunks) + + self.assertEqual( + result, + '{"places": [{"placeId": "ChIJ_joe"}, {"placeId": "ChIJ_prince"}]}', + ) + self.assertEqual(unresolved, 0) + + def test_shorter_key_shadows_the_longer_one_sharing_its_prefix(self): + # Pins shipped behavior. Substitution walks the index in arrival order, so + # PLACE_ID_FOR_1_Joe's Pizza fires first and strands " Express" on a + # now-real Place ID. Observed live on "Grand Central" versus "Grand Central + # Terminal". + text = "PLACE_ID_FOR_1_Joe's Pizza Express" + chunks = [ + _chunk("Joe's Pizza - Google Maps", "places/ChIJ_joe"), + _chunk("Joe's Pizza Express - Google Maps", "places/ChIJ_express"), + ] + + result, unresolved = place_id_resolution.resolve_place_ids(text, chunks) + + self.assertEqual(result, "ChIJ_joe Express") + self.assertEqual(unresolved, 0) + + def test_leaves_placeholder_in_place_when_sources_run_short(self): + # Grounding cited fewer venues than the model named. A visible placeholder + # breaks the card loudly; a nearby Place ID would render a wrong venue with + # full confidence. + text = "PLACE_ID_FOR_1_Starbucks and PLACE_ID_FOR_2_Starbucks" + chunks = [_chunk("Starbucks - Google Maps", "places/ChIJ_sbux_1")] + + with self.assertLogs(place_id_resolution.__name__, level="WARNING") as logs: + result, unresolved = place_id_resolution.resolve_place_ids(text, chunks) + + self.assertEqual(result, "ChIJ_sbux_1 and PLACE_ID_FOR_2_Starbucks") + self.assertEqual(unresolved, 1) + self.assertIn("unresolved", logs.output[0]) + + def test_warns_and_passes_text_through_when_no_sources_returned(self): + text = "PLACE_ID_FOR_1_Starbucks" + + with self.assertLogs(place_id_resolution.__name__, level="WARNING") as logs: + result, unresolved = place_id_resolution.resolve_place_ids(text, []) + + self.assertEqual(result, text) + self.assertEqual(unresolved, 1) + self.assertIn("unresolved against 0 source(s)", logs.output[0]) + + def test_empty_text_is_a_no_op(self): + chunks = [_chunk("Starbucks - Google Maps", "places/ChIJ_sbux_1")] + self.assertEqual(place_id_resolution.resolve_place_ids("", chunks), ("", 0)) + + def test_text_without_placeholders_is_unchanged(self): + text = '{"places": [{"placeId": "ChIJ_already_real"}]}' + chunks = [_chunk("Starbucks - Google Maps", "places/ChIJ_sbux_1")] + + result, unresolved = place_id_resolution.resolve_place_ids(text, chunks) + + self.assertEqual(result, text) + self.assertEqual(unresolved, 0) + + def test_resolves_global_list_numbered_placeholders_for_distinct_titles(self): + text = ( + "PLACE_ID_FOR_1_Sushi Ginza Onodera, " + "PLACE_ID_FOR_2_Shinjuku Sushi Hatsume, " + "PLACE_ID_FOR_9_Sushi Dai" + ) + chunks = [ + _chunk("Sushi Ginza Onodera - Google Maps", "places/ChIJ_onodera"), + _chunk("Shinjuku Sushi Hatsume - Google Maps", "places/ChIJ_hatsume"), + _chunk("Sushi Dai - Google Maps", "places/ChIJ_sushidai"), + ] + + result, unresolved = place_id_resolution.resolve_place_ids(text, chunks) + + self.assertEqual(result, "ChIJ_onodera, ChIJ_hatsume, ChIJ_sushidai") + self.assertEqual(unresolved, 0) + + def test_fallback_skips_duplicate_titles_when_both_misindexed(self): + text = ( + "PLACE_ID_FOR_1_Tartine, " + "PLACE_ID_FOR_2_Boudin, " + "PLACE_ID_FOR_3_Starbucks, " + "PLACE_ID_FOR_4_Starbucks" + ) + chunks = [ + _chunk("Tartine - Google Maps", "places/ChIJ_tartine"), + _chunk("Boudin - Google Maps", "places/ChIJ_boudin"), + _chunk("Starbucks - Google Maps", "places/ChIJ_sbux_1"), + _chunk("Starbucks - Google Maps", "places/ChIJ_sbux_2"), + ] + + result, unresolved = place_id_resolution.resolve_place_ids(text, chunks) + + self.assertEqual( + result, + "ChIJ_tartine, ChIJ_boudin, " + "PLACE_ID_FOR_3_Starbucks, PLACE_ID_FOR_4_Starbucks", + ) + self.assertEqual(unresolved, 2) + + +class PromptContractTest(unittest.TestCase): + + def test_prompt_rules_describe_the_resource_name_the_parser_strips(self): + # The authoritative branch tells the model to copy a `places/ChIJ...` + # value verbatim, and resolve_place_ids strips exactly that prefix. If the + # prompt were reworded to ask for a bare ID while the strip stayed, or vice + # versa, resource names would ship to the client. Pin the shared form. + self.assertIn( + place_id_resolution._PLACE_ID_PREFIX, # pylint: disable=protected-access + place_id_resolution.PROMPT_RULES, + ) + + def test_prompt_rules_describe_the_format_the_parser_builds(self): + # The prompt and build_placeholder_index must agree on the placeholder + # format. If they drift, every key misses and raw placeholders ship to + # the client, so pin the shared prefix and the ordinal example here. + self.assertIn( + f"{place_id_resolution.PLACEHOLDER_PREFIX}{{Count}}_{{Exact Title}}", + place_id_resolution.PROMPT_RULES, + ) + + built = place_id_resolution._build_placeholder_index([ + _chunk("Chez Panisse - Google Maps", "places/ChIJ_cp_1"), + _chunk("Chez Panisse - Google Maps", "places/ChIJ_cp_2"), + ]) + + for example in ( + f"{place_id_resolution.PLACEHOLDER_PREFIX}1_Chez Panisse", + f"{place_id_resolution.PLACEHOLDER_PREFIX}2_Chez Panisse", + ): + with self.subTest(placeholder=example): + self.assertIn(example, place_id_resolution.PROMPT_RULES) + self.assertIn(example, built) + + +if __name__ == "__main__": + unittest.main() diff --git a/agent/python_agent/test_router_config.py b/agent/python_agent/test_router_config.py deleted file mode 100644 index 190bd5a..0000000 --- a/agent/python_agent/test_router_config.py +++ /dev/null @@ -1,42 +0,0 @@ -# Copyright 2026 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. - -"""Tests for Router Configuration.""" - -import unittest -from router_config import IntentClass -from router_config import ROUTER_SYSTEM_INSTRUCTION -from router_config import RouterClassification - - -class TestRouterConfig(unittest.TestCase): - """Unit tests for Intent Router configuration schema and prompt.""" - - def test_schema_instantiation(self): - data = {"intent": "LOCAL_SEARCH", "query": "coffee near me"} - classification = RouterClassification(**data) - self.assertEqual(classification.intent, IntentClass.LOCAL_SEARCH) - self.assertEqual(classification.query, "coffee near me") - - def test_schema_validation_error(self): - data = {"intent": "INVALID_INTENT", "query": "coffee near me"} - with self.assertRaises(ValueError): - RouterClassification(**data) - - def test_instruction_not_empty(self): - self.assertGreater(len(ROUTER_SYSTEM_INSTRUCTION), 0) - - -if __name__ == "__main__": - unittest.main() diff --git a/agent/python_agent/test_streaming_request_handler.py b/agent/python_agent/test_streaming_request_handler.py new file mode 100644 index 0000000..650b9e8 --- /dev/null +++ b/agent/python_agent/test_streaming_request_handler.py @@ -0,0 +1,45 @@ +# 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 asyncio +import unittest +from unittest import mock + +import streaming_request_handler + +class QueueBridgeTest(unittest.IsolatedAsyncioTestCase): + async def test_enqueue_event(self): + queue = asyncio.Queue() + bridge = streaming_request_handler.QueueBridge(queue) + + # Test normal put + mock_event = mock.MagicMock() + await bridge.enqueue_event(mock_event) + self.assertEqual(await queue.get(), mock_event) + + async def test_close_queue(self): + queue = asyncio.Queue() + bridge = streaming_request_handler.QueueBridge(queue) + + # Test close behavior + await bridge.close() + self.assertTrue(bridge.is_closed()) + self.assertIsNone(await queue.get()) + + # Test putting after close raises error + with self.assertRaises(RuntimeError): + await bridge.enqueue_event(mock.MagicMock()) + +if __name__ == '__main__': + unittest.main() diff --git a/agent/python_agent/test_template_registry.py b/agent/python_agent/test_template_registry.py new file mode 100644 index 0000000..49fd6b7 --- /dev/null +++ b/agent/python_agent/test_template_registry.py @@ -0,0 +1,562 @@ +# Copyright 2026 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. + +"""Tests for the template bundle registry.""" + +import dataclasses +import enum +import importlib +import importlib.resources +import json +import os +import pathlib +import shutil +import sys +import tempfile +import unittest +import zipfile + +import pydantic + +import template_registry +from templates.directions import schema as directions_schema +from templates.local_search import schema as local_search_schema + + +class _StubSchema(pydantic.BaseModel): + """Minimal schema for bundles constructed directly by a test.""" + + title: str + + +def _make_bundle( + name: str = "custom_poi", + intent: str = "CUSTOM_POI", + skill_instructions: str = "", + tool_description: str = "", + archetype: str = "", +) -> template_registry.TemplateBundle: + """Builds a TemplateBundle without touching the filesystem.""" + bundle_dir = pathlib.Path("/nonexistent") / name + return template_registry.TemplateBundle( + name=name, + intent=intent, + surface_prefix=f"{name.replace('_', '-')}-surface", + archetype=archetype or f"Archetype for {intent}.", + few_shot_examples=[{ + "user_query": "Tell me about the Space Needle", + "intent": intent, + "query": "Space Needle Seattle", + }], + schema_cls=_StubSchema, + layout_path=bundle_dir / "layout.json", + skill_dir=bundle_dir, + bundle_dir=bundle_dir, + skill_instructions=skill_instructions, + tool_description=tool_description, + ) + + +class SchemaResolutionTest(unittest.TestCase): + """Schemas resolve by import, never by executing a file.""" + + def test_builtin_bundles_resolve_schemas(self): + for bundle_name in ("local_search", "directions"): + with self.subTest(bundle=bundle_name): + schema_cls = template_registry._resolve_schema(bundle_name) + self.assertTrue(issubclass(schema_cls, pydantic.BaseModel)) + + def test_unmapped_bundle_raises_key_error(self): + with self.assertRaises(KeyError): + template_registry._resolve_schema("no_such_bundle") + + def test_resolved_class_is_the_statically_imported_class(self): + """Guards against a second copy of each schema class in the process. + + A registry that executes templates/<name>/schema.py directly defines a + class distinct from the one the build system compiles into the library, + so isinstance checks fail across the two copies. + """ + registry = template_registry.TemplateRegistry() + + self.assertIs( + registry.get_bundle("local_search").schema_cls, + local_search_schema.ExtractorSchema, + ) + self.assertIs( + registry.get_bundle("directions").schema_cls, + directions_schema.ExtractorSchema, + ) + + +class DiscoveryTest(unittest.TestCase): + """Discovery of the bundles shipped in templates/.""" + + def setUp(self): + super().setUp() + self.registry = template_registry.TemplateRegistry() + + def test_discovers_builtin_bundles(self): + names = {bundle.name for bundle in self.registry.list_bundles()} + self.assertIn("local_search", names) + self.assertIn("directions", names) + + def test_populates_bundle_from_manifest(self): + """Verifies local_search bundle metadata is populated from manifest.json.""" + bundle = self.registry.get_bundle("local_search") + + self.assertEqual(bundle.intent, "LOCAL_SEARCH") + self.assertEqual(bundle.surface_prefix, "local-search-surface") + self.assertIn("points of interest", bundle.archetype) + self.assertIn("named places", bundle.archetype) + self.assertEqual(len(bundle.few_shot_examples), 2) + self.assertEqual( + bundle.few_shot_examples[0]["user_query"], + "Find coffee shops near Central Park", + ) + self.assertEqual( + bundle.few_shot_examples[1]["user_query"], + "Show me GWC2 and google 1950 in the same map", + ) + + def test_bundle_assets_exist_on_disk(self): + for bundle in self.registry.list_bundles(): + with self.subTest(bundle=bundle.name): + self.assertTrue(bundle.layout_path.exists()) + self.assertTrue((bundle.skill_dir / "SKILL.md").exists()) + self.assertGreater(len(bundle.skill_instructions), 0) + self.assertFalse(bundle.skill_instructions.startswith("---")) + + def test_supported_intents(self): + self.assertEqual( + self.registry.supported_intents, {"LOCAL_SEARCH", "DIRECTIONS"} + ) + + def test_discovery_order_is_deterministic(self): + self.assertEqual( + [bundle.name for bundle in self.registry.list_bundles()], + ["directions", "local_search"], + ) + + +class GetBundleTest(unittest.TestCase): + """Lookup by name, by intent, and by enum-like key.""" + + def setUp(self): + super().setUp() + self.registry = template_registry.TemplateRegistry() + + def test_lookup_by_name(self): + self.assertEqual( + self.registry.get_bundle("local_search").name, "local_search" + ) + + def test_lookup_by_intent(self): + self.assertEqual( + self.registry.get_bundle_by_intent("LOCAL_SEARCH").name, "local_search" + ) + + def test_lookup_by_enum_like_object(self): + class MockEnum: + value = "DIRECTIONS" + + bundle = self.registry.get_bundle_by_intent(MockEnum()) + self.assertEqual(bundle.name, "directions") + + def test_unknown_key_returns_none(self): + self.assertIsNone(self.registry.get_bundle("NOT_A_TEMPLATE")) + self.assertIsNone(self.registry.get_bundle_by_intent("NOT_AN_INTENT")) + + def test_built_in_bundles_load_empty_fallback_text(self): + local_search = self.registry.get_bundle("local_search") + self.assertEqual( + local_search.empty_fallback_text, + "No places matching your query could be found.", + ) + directions = self.registry.get_bundle("directions") + self.assertEqual( + directions.empty_fallback_text, + "Could not calculate travel directions.", + ) + + +class RegisterBundleTest(unittest.TestCase): + """register_bundle takes a constructed bundle and nothing else.""" + + def setUp(self): + super().setUp() + self.registry = template_registry.TemplateRegistry() + + def test_registers_bundle_and_intent(self): + bundle = _make_bundle() + self.registry.register_bundle(bundle) + + self.assertIs(self.registry.get_bundle("custom_poi"), bundle) + self.assertIs(self.registry.get_bundle_by_intent("CUSTOM_POI"), bundle) + self.assertIn("CUSTOM_POI", self.registry.supported_intents) + + def test_registers_intent_with_the_router(self): + self.registry.register_bundle(_make_bundle(intent="ROUTER_VISIBLE_INTENT")) + + classification = self.registry.router_classification( + intent="ROUTER_VISIBLE_INTENT", query="anything" + ) + self.assertEqual(classification.intent, "ROUTER_VISIBLE_INTENT") + self.assertIn("ROUTER_VISIBLE_INTENT", self.registry.supported_intents) + + def test_router_classification_rejects_unregistered_intent(self): + with self.assertRaises(pydantic.ValidationError): + self.registry.router_classification( + intent="UNREGISTERED_INTENT", query="anything" + ) + + def test_router_classification_generates_dynamic_enum_schema(self): + schema = self.registry.router_classification + self.assertTrue(issubclass(schema, template_registry.RouterClassification)) + intent_field = schema.model_fields["intent"] + self.assertTrue(issubclass(intent_field.annotation, enum.Enum)) + self.assertTrue(issubclass(intent_field.annotation, str)) + + classification = schema(intent="LOCAL_SEARCH", query="coffee near me") + self.assertEqual(classification.intent, "LOCAL_SEARCH") + + with self.assertRaises(pydantic.ValidationError): + schema(intent="INVALID_INTENT", query="coffee near me") + + def test_router_classification_is_cached(self): + first = self.registry.router_classification + second = self.registry.router_classification + self.assertIs(first, second) + + self.registry.register_bundle(_make_bundle(intent="DYNAMIC_INTENT")) + third = self.registry.router_classification + self.assertIsNot(first, third) + classification = third(intent="DYNAMIC_INTENT", query="test query") + self.assertEqual(classification.intent, "DYNAMIC_INTENT") + + def test_registered_intent_stays_out_of_other_registries(self): + self.registry.register_bundle(_make_bundle(intent="ROUTER_VISIBLE_INTENT")) + other = template_registry.TemplateRegistry() + + self.assertNotIn("ROUTER_VISIBLE_INTENT", other.supported_intents) + self.assertIsNone(other.get_bundle_by_intent("ROUTER_VISIBLE_INTENT")) + + def test_replacing_a_bundle_logs_a_warning(self): + self.registry.register_bundle(_make_bundle()) + replacement = _make_bundle() + + with self.assertLogs(template_registry.__name__, level="WARNING") as logs: + self.registry.register_bundle(replacement) + + self.assertIs(self.registry.get_bundle("custom_poi"), replacement) + self.assertIn("already registered", "\n".join(logs.output)) + + def test_stealing_an_intent_drops_the_displaced_bundle(self): + self.registry.register_bundle(_make_bundle(name="first", intent="SHARED")) + + with self.assertLogs(template_registry.__name__, level="WARNING") as logs: + self.registry.register_bundle( + _make_bundle(name="second", intent="SHARED") + ) + + self.assertEqual( + self.registry.get_bundle_by_intent("SHARED").name, "second" + ) + self.assertIn("SHARED", "\n".join(logs.output)) + # The displaced bundle is unroutable, so it must not linger in either + # index where it would still reach the router prompt. + self.assertIsNone(self.registry.get_bundle("first")) + self.assertNotIn( + "first", [bundle.name for bundle in self.registry.list_bundles()] + ) + + def test_renaming_an_intent_drops_the_stale_intent_key(self): + self.registry.register_bundle(_make_bundle(name="poi", intent="OLD")) + + with self.assertLogs(template_registry.__name__, level="WARNING"): + self.registry.register_bundle(_make_bundle(name="poi", intent="NEW")) + + self.assertIsNone(self.registry.get_bundle_by_intent("OLD")) + self.assertNotIn("OLD", self.registry.supported_intents) + self.assertIn("NEW", self.registry.supported_intents) + + +class CompileRouterInstructionTest(unittest.TestCase): + """The router prompt is assembled from whatever is registered.""" + + def setUp(self): + super().setUp() + self.registry = template_registry.TemplateRegistry() + self.instruction = self.registry.compile_router_instruction() + + def test_contains_registered_archetypes(self): + self.assertIn("- **LOCAL_SEARCH**", self.instruction) + self.assertIn("- **DIRECTIONS**", self.instruction) + + def test_contains_pinned_non_template_archetypes(self): + self.assertIn("- **OTHER_SPATIAL**", self.instruction) + self.assertIn("- **TEXT_ONLY**", self.instruction) + self.assertIn("Conservative Routing", self.instruction) + + def test_contains_bundle_few_shot_examples(self): + self.assertIn("Find coffee shops near Central Park", self.instruction) + self.assertIn( + "Directions from Sacramento to Mendocino via Clear Lake", + self.instruction, + ) + + def test_registered_bundle_reaches_the_prompt(self): + self.registry.register_bundle(_make_bundle()) + instruction = self.registry.compile_router_instruction() + + self.assertIn("- **CUSTOM_POI**", instruction) + self.assertIn("Tell me about the Space Needle", instruction) + + def test_bundle_order_follows_intent_not_registration_order(self): + """A late registration still lands in intent order, not at the end.""" + self.registry.register_bundle( + _make_bundle(name="alpha_poi", intent="ALPHA_POI") + ) + instruction = self.registry.compile_router_instruction() + + self.assertLess( + instruction.index("- **ALPHA_POI**"), + instruction.index("- **DIRECTIONS**"), + ) + self.assertLess( + instruction.index("Tell me about the Space Needle"), + instruction.index("Directions from Sacramento to Mendocino"), + ) + + def test_matches_the_golden_prompt(self): + """The compiled prompt matches a checked-in copy of its own output. + + The golden holds every archetype line and few-shot example that used to + be hardcoded in `router_config.ROUTER_SYSTEM_INSTRUCTION`, with bundle + entries reordered by intent. Pinning it here means a change to any + manifest surfaces as a reviewable diff on the prompt itself rather than + silently altering what the router sees. + """ + golden_path = pathlib.Path(__file__).parent / "testdata" + golden = (golden_path / "router_instruction.golden.txt").read_text() + + self.maxDiff = None + self.assertEqual(self.instruction, golden) + + +class SkillInstructionsTest(unittest.TestCase): + """Tests for _read_skill_instructions and TemplateBundle.skill_instructions.""" + + def test_loads_and_strips_frontmatter(self): + with tempfile.TemporaryDirectory() as tmp_dir: + skill_dir = pathlib.Path(tmp_dir) + skill_file = skill_dir / "SKILL.md" + skill_file.write_text( + "---\nname: test_bundle\ndescription: Test\n---\n\n# Header\nBody", + encoding="utf-8", + ) + self.assertEqual( + template_registry._read_skill_instructions(skill_dir), + "# Header\nBody", + ) + + def test_returns_content_without_frontmatter(self): + with tempfile.TemporaryDirectory() as tmp_dir: + skill_dir = pathlib.Path(tmp_dir) + skill_file = skill_dir / "SKILL.md" + skill_file.write_text( + "# Header\nBody without frontmatter", encoding="utf-8" + ) + self.assertEqual( + template_registry._read_skill_instructions(skill_dir), + "# Header\nBody without frontmatter", + ) + + def test_missing_skill_file_returns_empty_string(self): + self.assertEqual( + template_registry._read_skill_instructions( + pathlib.Path("/nonexistent") + ), + "", + ) + bundle = _make_bundle() + self.assertEqual(bundle.skill_instructions, "") + + def test_built_in_bundles_have_non_empty_instructions(self): + registry = template_registry.TemplateRegistry() + local_search = registry.get_bundle("local_search") + self.assertIsNotNone(local_search) + self.assertIn("# Local Search Guidelines", local_search.skill_instructions) + directions = registry.get_bundle("directions") + self.assertIsNotNone(directions) + self.assertIn("# Directions Guidelines", directions.skill_instructions) + + +class BundlePropertiesTest(unittest.TestCase): + """Tests for derived bundle properties.""" + + def test_render_tool_name(self): + bundle = _make_bundle(name="custom_poi") + self.assertEqual(bundle.render_tool_name, "render_custom_poi_template") + + def test_invalid_bundle_name_raises_value_error(self): + for invalid_name in ( + "", + "local-search", + "LocalSearch", + "123_poi", + "poi!", + ): + with self.subTest(name=invalid_name): + with self.assertRaisesRegex(ValueError, "Invalid bundle name"): + _make_bundle(name=invalid_name) + + def test_built_in_bundles_render_tool_names(self): + registry = template_registry.TemplateRegistry() + local_search = registry.get_bundle("local_search") + self.assertEqual( + local_search.render_tool_name, "render_local_search_template" + ) + directions = registry.get_bundle("directions") + self.assertEqual(directions.render_tool_name, "render_directions_template") + + +class EffectiveToolDescriptionTest(unittest.TestCase): + """Tests for TemplateBundle.effective_tool_description.""" + + def test_returns_tool_description_when_specified(self): + bundle = _make_bundle( + name="custom_poi", + tool_description="Custom tool description for POI.", + archetype="General POI queries.", + ) + self.assertEqual( + bundle.effective_tool_description, "Custom tool description for POI." + ) + + def test_falls_back_to_archetype_when_empty(self): + bundle = _make_bundle( + name="custom_poi", + tool_description="", + archetype="Finding specific tourist landmarks.", + ) + self.assertEqual( + bundle.effective_tool_description, + "Renders the custom_poi UI component. Use it when the request matches:" + " Finding specific tourist landmarks.", + ) + + def test_built_in_bundles_have_effective_tool_descriptions(self): + registry = template_registry.TemplateRegistry() + for bundle in registry.list_bundles(): + self.assertTrue(bundle.effective_tool_description) + self.assertIn( + bundle.effective_tool_description, + [ + bundle.tool_description, + ( + f"Renders the {bundle.name} UI component. Use it when the" + f" request matches: {bundle.archetype}" + ), + ], + ) + + +class ZipImportedPackageTest(unittest.TestCase): + """Covers bundle discovery when the agent is served from a zip archive. + + In that mode `__file__` names an entry inside the archive, so discovering + bundles by filesystem path registers none of them at serving time while a + suite run from an ordinary source tree keeps passing. Mirrors + `test_merger.TestMergerZipImportedPackage`. + """ + + _PACKAGE = "a2ui_zipped_registry_probe" + _BUNDLE_FILES = ("manifest.json", "layout.json", "SKILL.md") + + def _import_registry_from_zip(self): + """Imports a copy of `template_registry` whose bundles exist only in a zip. + + Schema modules still resolve through the absolute templates package, so + the probe exercises discovery and asset reads, not schema imports. + + Returns: + The zip-imported `template_registry` module. + """ + source = importlib.resources.files(template_registry.__package__) + temp_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True) + archive_path = os.path.join(temp_dir, "registry_probe.zip") + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr(f"{self._PACKAGE}/__init__.py", "") + archive.writestr( + f"{self._PACKAGE}/template_registry.py", + source.joinpath("template_registry.py").read_text(encoding="utf-8"), + ) + for entry in source.joinpath("templates").iterdir(): + for file_name in self._BUNDLE_FILES: + asset = entry.joinpath(file_name) + if asset.is_file(): + archive.writestr( + f"{self._PACKAGE}/templates/{entry.name}/{file_name}", + asset.read_text(encoding="utf-8"), + ) + sys.path.insert(0, archive_path) + self.addCleanup(sys.path.remove, archive_path) + for name in (f"{self._PACKAGE}.template_registry", self._PACKAGE): + self.addCleanup(sys.modules.pop, name, None) + importlib.invalidate_caches() + return importlib.import_module(f"{self._PACKAGE}.template_registry") + + def test_probe_bundles_are_reachable_only_inside_the_archive(self): + """Asserts the probe exercises the archive rather than a copy on disk.""" + zipped_registry = self._import_registry_from_zip() + + on_disk_templates = os.path.join( + os.path.dirname(os.path.abspath(zipped_registry.__file__)), + "templates", + ) + + self.assertNotEqual(zipped_registry.__file__, template_registry.__file__) + self.assertFalse(os.path.exists(on_disk_templates)) + + def test_discovers_every_builtin_bundle_from_the_archive(self): + zipped = self._import_registry_from_zip().TemplateRegistry() + unzipped = template_registry.TemplateRegistry() + + self.assertTrue(unzipped.bundles) + self.assertCountEqual(zipped.bundles.keys(), unzipped.bundles.keys()) + + def test_reads_bundle_assets_from_the_archive(self): + zipped = self._import_registry_from_zip().TemplateRegistry() + unzipped = template_registry.TemplateRegistry() + + for name, bundle in zipped.bundles.items(): + with self.subTest(bundle=name): + self.assertTrue(bundle.layout_path.is_file()) + self.assertEqual( + json.loads(bundle.layout_path.read_text(encoding="utf-8")), + json.loads( + unzipped.bundles[name].layout_path.read_text(encoding="utf-8") + ), + ) + self.assertEqual( + bundle.skill_instructions, + unzipped.bundles[name].skill_instructions, + ) + self.assertEqual(bundle.intent, unzipped.bundles[name].intent) + + +if __name__ == "__main__": + unittest.main() diff --git a/agent/python_agent/test_template_tool.py b/agent/python_agent/test_template_tool.py new file mode 100644 index 0000000..d6c9e29 --- /dev/null +++ b/agent/python_agent/test_template_tool.py @@ -0,0 +1,490 @@ +# Copyright 2026 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 +# +# http://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. + +"""Unit tests for template_tool.py ADK tools.""" + +import pathlib +from types import SimpleNamespace +import unittest + +from a2a.types import DataPart +from google.adk.tools.tool_context import ToolContext + +import a2ui +import agent +import ( + template_registry, +) +import template_tool + +BaseTemplateTool = template_tool.BaseTemplateTool +RenderTextOnlyTemplateTool = template_tool.RenderTextOnlyTemplateTool +TemplateToolset = template_tool.TemplateToolset +STATE_RENDERED_A2UI_PARTS = template_tool.STATE_RENDERED_A2UI_PARTS + + +def _make_bundle_tool(name: str, **kwargs): + bundle = template_registry.TemplateRegistry().get_bundle(name) + assert bundle is not None + return BaseTemplateTool.from_bundle(bundle, **kwargs) + + +def _create_schema_manager(): + extension_path = ( + pathlib.Path(__file__).parent + / "shared" + / "schema" + / "maps_catalog_extension.json" + ) + return a2ui.schema.manager.A2uiSchemaManager( + version=a2ui.schema.constants.VERSION_0_9, + catalogs=[ + a2ui.schema.catalog.CatalogConfig( + name="maps-agentic-ui-catalog", + provider=agent.MergedCatalogProvider( + a2ui.schema.constants.VERSION_0_9, str(extension_path) + ), + ) + ], + schema_modifiers=[a2ui.schema.common_modifiers.remove_strict_validation], + ) + + +class MockToolContext: + + def __init__(self): + self.state = {} + self.actions = SimpleNamespace() + + +class TestTemplateTools(unittest.IsolatedAsyncioTestCase): + """Unit tests for ADK template tools.""" + + def setUp(self): + super().setUp() + self.schema_manager = _create_schema_manager() + self.tool_context = MockToolContext() + + async def test_render_local_search_template_success(self): + tool = _make_bundle_tool( + "local_search", schema_manager=self.schema_manager, max_list_size=3 + ) + + args = { + "heading": "Nearby Places", + "summary": "Here are 2 coffee shops.", + "center_lat": 47.6062, + "center_lng": -122.3321, + "zoom": 14, + "places": [ + { + "placeId": "ChIJ111", + "name": "Espresso Vivace", + "lat": 47.6200, + "lng": -122.3200, + }, + { + "placeId": "ChIJ222", + "name": "Milstead & Co.", + "lat": 47.6400, + "lng": -122.3500, + }, + ], + } + + result = await tool.run_async(args=args, tool_context=self.tool_context) + + self.assertEqual(result["status"], "success") + self.assertEqual(result["template"], "local_search") + self.assertIn("surface_id", result) + + # Verify session state was populated with A2A parts + self.assertIn(STATE_RENDERED_A2UI_PARTS, self.tool_context.state) + parts = self.tool_context.state[STATE_RENDERED_A2UI_PARTS] + self.assertEqual(len(parts), 3) + + create_surface_data = parts[0].root.data["createSurface"] + self.assertTrue( + create_surface_data["surfaceId"].startswith("local-search-surface-") + ) + + update_components = parts[1].root.data["updateComponents"]["components"] + heading_comp = next( + comp for comp in update_components if comp["id"] == "heading-text" + ) + self.assertEqual(heading_comp["text"], "### Nearby Places") + map_comp = next(comp for comp in update_components if comp["id"] == "map") + self.assertEqual(len(map_comp["markers"]), 2) + + update_data_model = parts[2].root.data["updateDataModel"]["value"] + self.assertEqual(len(update_data_model["places"]), 2) + + async def test_render_local_search_template_with_heading(self): + tool = _make_bundle_tool( + "local_search", schema_manager=self.schema_manager, max_list_size=3 + ) + + args = { + "heading": "Top Coffee Shops in Seattle", + "summary": "Here are 2 coffee shops.", + "center_lat": 47.6062, + "center_lng": -122.3321, + "zoom": 14, + "places": [ + { + "placeId": "ChIJ111", + "name": "Espresso Vivace", + "lat": 47.6200, + "lng": -122.3200, + }, + ], + } + + result = await tool.run_async(args=args, tool_context=self.tool_context) + self.assertEqual(result["status"], "success") + + parts = self.tool_context.state[STATE_RENDERED_A2UI_PARTS] + update_components = parts[1].root.data["updateComponents"]["components"] + heading_comp = next( + comp for comp in update_components if comp["id"] == "heading-text" + ) + self.assertEqual(heading_comp["text"], "### Top Coffee Shops in Seattle") + + async def test_render_local_search_template_validation_failure(self): + tool = _make_bundle_tool("local_search", schema_manager=self.schema_manager) + + # Missing mandatory center_lat and center_lng + args = { + "summary": "Places without center coordinates", + "places": [{"placeId": "1", "name": "P1", "lat": 1.0, "lng": 2.0}], + } + + result = await tool.run_async(args=args, tool_context=self.tool_context) + self.assertIn("error", result) + self.assertIn("Validation failed for tool", result["error"]) + + async def test_render_local_search_template_missing_heading_validation_failure( + self, + ): + tool = _make_bundle_tool("local_search", schema_manager=self.schema_manager) + + # Missing mandatory heading + args = { + "summary": "Places without heading", + "center_lat": 47.6062, + "center_lng": -122.3321, + "places": [{"placeId": "1", "name": "P1", "lat": 1.0, "lng": 2.0}], + } + + result = await tool.run_async(args=args, tool_context=self.tool_context) + self.assertIn("error", result) + self.assertIn("Validation failed for tool", result["error"]) + + async def test_render_directions_template_success(self): + tool = _make_bundle_tool("directions", schema_manager=self.schema_manager) + + args = { + "heading": "Driving directions from San Francisco to Oakland", + "summary": "Commute is 30 minutes.", + "center_lat": 37.7749, + "center_lng": -122.4194, + "zoom": 12, + "routes": [{ + "origin": { + "lat": 37.7749, + "lng": -122.4194, + "label": "San Francisco", + "placeId": "ChIJ_SF", + }, + "destination": { + "lat": 37.8044, + "lng": -122.2712, + "label": "Oakland", + "placeId": "ChIJ_OAK", + }, + }], + "travel_mode": "driving", + } + + result = await tool.run_async(args=args, tool_context=self.tool_context) + + self.assertEqual(result["status"], "success") + self.assertEqual(result["template"], "directions") + + self.assertIn(STATE_RENDERED_A2UI_PARTS, self.tool_context.state) + parts = self.tool_context.state[STATE_RENDERED_A2UI_PARTS] + self.assertEqual(len(parts), 3) + + update_components = parts[1].root.data["updateComponents"]["components"] + root_comp = next(comp for comp in update_components if comp["id"] == "root") + self.assertEqual( + root_comp["children"], ["heading-text", "map", "summary-text"] + ) + heading_comp = next( + comp for comp in update_components if comp["id"] == "heading-text" + ) + self.assertEqual( + heading_comp["text"], + "### Driving directions from San Francisco to Oakland", + ) + map_comp = next(comp for comp in update_components if comp["id"] == "map") + self.assertEqual(map_comp["travelMode"], "driving") + self.assertEqual(len(map_comp["routes"]), 1) + + async def test_render_directions_template_missing_heading_fails(self): + tool = _make_bundle_tool("directions", schema_manager=self.schema_manager) + + args = { + "summary": "Commute is 30 minutes.", + "center_lat": 37.7749, + "center_lng": -122.4194, + "zoom": 12, + "routes": [{ + "origin": {"lat": 37.7749, "lng": -122.4194, "label": "A"}, + "destination": {"lat": 37.8044, "lng": -122.2712, "label": "B"}, + }], + "travel_mode": "driving", + } + + result = await tool.run_async(args=args, tool_context=self.tool_context) + self.assertIn("error", result) + self.assertIn("Validation failed for tool", result["error"]) + + async def test_render_directions_template_invalid_travel_mode_fails(self): + tool = _make_bundle_tool("directions", schema_manager=self.schema_manager) + + args = { + "summary": "Commute", + "center_lat": 37.7, + "center_lng": -122.4, + "zoom": 12, + "routes": [{ + "origin": {"lat": 37.7, "lng": -122.4, "label": "A"}, + "destination": {"lat": 37.8, "lng": -122.3, "label": "B"}, + }], + "travel_mode": "ROCKET_SHIP", # Invalid mode + } + + result = await tool.run_async(args=args, tool_context=self.tool_context) + self.assertIn("error", result) + + async def test_render_text_only_template_success(self): + tool = RenderTextOnlyTemplateTool(schema_manager=self.schema_manager) + + args = { + "text": "Hello world from text-only template.", + } + + result = await tool.run_async(args=args, tool_context=self.tool_context) + + self.assertEqual(result["status"], "success") + self.assertEqual(result["template"], "text_only") + + self.assertIn(STATE_RENDERED_A2UI_PARTS, self.tool_context.state) + parts = self.tool_context.state[STATE_RENDERED_A2UI_PARTS] + self.assertEqual(len(parts), 2) + text_comp = parts[1].root.data["updateComponents"]["components"][1] + self.assertEqual(text_comp["text"], "Hello world from text-only template.") + + async def test_template_toolset_returns_tools(self): + toolset = TemplateToolset( + schema_manager=self.schema_manager, max_list_size=3 + ) + tools = await toolset.get_tools() + + self.assertEqual(len(tools), len(toolset.template_registry.bundles) + 1) + tool_names = [t.name for t in tools] + self.assertIn("render_local_search_template", tool_names) + self.assertIn("render_directions_template", tool_names) + self.assertIn("render_text_only_template", tool_names) + + async def test_template_toolset_defaults_to_maps_catalog_id(self): + toolset = TemplateToolset(schema_manager=self.schema_manager) + tools = await toolset.get_tools() + + for tool in tools: + with self.subTest(tool=tool.name): + self.assertEqual(tool.catalog_id, "a2ui://maps-agentic-ui-catalog.json") + + async def test_template_toolset_propagates_catalog_id(self): + """A host names its catalog once; every tool in the set must honor it.""" + host_catalog = "https://example.test/host_catalog.json" + toolset = TemplateToolset( + schema_manager=self.schema_manager, catalog_id=host_catalog + ) + tools = await toolset.get_tools() + + self.assertEqual(len(tools), len(toolset.template_registry.bundles) + 1) + for tool in tools: + with self.subTest(tool=tool.name): + self.assertEqual(tool.catalog_id, host_catalog) + + async def test_template_toolset_with_tool_filter(self): + toolset = TemplateToolset( + schema_manager=self.schema_manager, + max_list_size=3, + tool_filter=[ + "render_local_search_template", + "render_text_only_template", + ], + ) + tools = await toolset.get_tools() + + self.assertEqual(len(tools), 2) + tool_names = [t.name for t in tools] + self.assertEqual( + tool_names, + ["render_local_search_template", "render_text_only_template"], + ) + + def test_tool_declarations_valid(self): + tool_ls = _make_bundle_tool( + "local_search", schema_manager=self.schema_manager + ) + decl_ls = tool_ls._get_declaration() + self.assertIsNotNone(decl_ls) + self.assertEqual(decl_ls.name, "render_local_search_template") + + tool_dir = _make_bundle_tool( + "directions", schema_manager=self.schema_manager + ) + decl_dir = tool_dir._get_declaration() + self.assertIsNotNone(decl_dir) + self.assertEqual(decl_dir.name, "render_directions_template") + + tool_text = RenderTextOnlyTemplateTool(schema_manager=self.schema_manager) + decl_text = tool_text._get_declaration() + self.assertIsNotNone(decl_text) + self.assertEqual(decl_text.name, "render_text_only_template") + + async def test_toolset_tool_descriptions_match_manifest(self): + toolset = TemplateToolset() + tools = {tool.name: tool.description for tool in await toolset.get_tools()} + for bundle in toolset.template_registry.bundles.values(): + tool_name = bundle.render_tool_name + self.assertIn(tool_name, tools) + self.assertEqual(tools[tool_name], bundle.effective_tool_description) + self.assertIn("render_text_only_template", tools) + + +class TestOnRenderedHook(unittest.IsolatedAsyncioTestCase): + """The opt-in hook replaces the session-state handoff; the default doesn't.""" + + _TEXT_ARGS = {"text": "hello", "surface_id": "text-only-surface-fixed"} + + def setUp(self): + super().setUp() + self.tool_context = MockToolContext() + self.calls = [] + + def _record(self, tool_context, messages, data): + self.calls.append((tool_context, messages, data)) + + async def test_hook_receives_messages_and_data(self): + tool = RenderTextOnlyTemplateTool(on_rendered=self._record) + + result = await tool.run_async( + args=dict(self._TEXT_ARGS), tool_context=self.tool_context + ) + + self.assertEqual(result["status"], "success") + self.assertEqual(len(self.calls), 1) + tool_context, messages, data = self.calls[0] + self.assertIs(tool_context, self.tool_context) + self.assertTrue(messages) + for message in messages: + self.assertIsInstance(message, dict) + self.assertEqual( + messages[0]["createSurface"]["surfaceId"], "text-only-surface-fixed" + ) + self.assertEqual(data["text"], "hello") + + async def test_hook_messages_match_default_parts(self): + hooked = RenderTextOnlyTemplateTool(on_rendered=self._record) + await hooked.run_async( + args=dict(self._TEXT_ARGS), tool_context=self.tool_context + ) + default_context = MockToolContext() + await RenderTextOnlyTemplateTool().run_async( + args=dict(self._TEXT_ARGS), tool_context=default_context + ) + + default_parts = default_context.state[STATE_RENDERED_A2UI_PARTS] + self.assertEqual( + self.calls[0][1], [part.root.data for part in default_parts] + ) + + async def test_hook_suppresses_state_handoff(self): + tool = RenderTextOnlyTemplateTool(on_rendered=self._record) + + await tool.run_async( + args=dict(self._TEXT_ARGS), tool_context=self.tool_context + ) + + self.assertEqual(self.tool_context.state, {}) + + async def test_default_still_writes_state(self): + await RenderTextOnlyTemplateTool().run_async( + args=dict(self._TEXT_ARGS), tool_context=self.tool_context + ) + + self.assertIn(STATE_RENDERED_A2UI_PARTS, self.tool_context.state) + self.assertIn( + template_tool.STATE_RENDERED_A2UI_DATA, self.tool_context.state + ) + + async def test_hook_failure_returns_error(self): + def failing_hook(tool_context, messages, data): + del tool_context, messages, data + raise RuntimeError("sink unavailable") + + tool = RenderTextOnlyTemplateTool(on_rendered=failing_hook) + + with self.assertLogs(template_tool.logger, level="ERROR"): + result = await tool.run_async( + args=dict(self._TEXT_ARGS), tool_context=self.tool_context + ) + + self.assertIn("error", result) + self.assertIn("sink unavailable", result["error"]) + self.assertEqual(self.tool_context.state, {}) + + async def test_hook_not_called_when_validation_fails(self): + tool = _make_bundle_tool("local_search", on_rendered=self._record) + + result = await tool.run_async( + args={"summary": "missing required fields"}, + tool_context=self.tool_context, + ) + + self.assertIn("error", result) + self.assertEqual(self.calls, []) + + async def test_toolset_propagates_hook_to_every_tool(self): + toolset = TemplateToolset(on_rendered=self._record) + tools = await toolset.get_tools() + + self.assertEqual(len(tools), len(toolset.template_registry.bundles) + 1) + for tool in tools: + with self.subTest(tool=tool.name): + self.assertEqual(tool.on_rendered, self._record) + + async def test_toolset_defaults_to_no_hook(self): + for tool in await TemplateToolset().get_tools(): + with self.subTest(tool=tool.name): + self.assertIsNone(tool.on_rendered) + + +if __name__ == "__main__": + unittest.main() diff --git a/agent/python_agent/test_templates_conformance.py b/agent/python_agent/test_templates_conformance.py new file mode 100644 index 0000000..022e80d --- /dev/null +++ b/agent/python_agent/test_templates_conformance.py @@ -0,0 +1,244 @@ +# Copyright 2026 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. + +"""Build-time conformance test suite for A2UI template directory bundles.""" + +import json +import pathlib +import re +from typing import Any +import unittest + +from google.adk import features as adk_features +from google.adk import skills as adk_skills +import pydantic + +import a2ui +import agent +import ( + template_registry, +) +from merger import ( + merge_template, +) + + +def _extract_placeholders(obj: Any) -> set[str]: + """Recursively extracts all {{placeholder}} tokens from a JSON structure.""" + placeholders = set() + if isinstance(obj, str): + matches = re.findall(r"\{\{([a-zA-Z0-9_]+)\}\}", obj) + placeholders.update(matches) + elif isinstance(obj, list): + for item in obj: + placeholders.update(_extract_placeholders(item)) + elif isinstance(obj, dict): + for value in obj.values(): + placeholders.update(_extract_placeholders(value)) + return placeholders + + +class ConformanceFewShotExample(pydantic.BaseModel): + """Strongly typed few-shot example for bundle router prompt injection.""" + + model_config = pydantic.ConfigDict(extra="forbid") + + user_query: str = pydantic.Field(min_length=1) + intent: str = pydantic.Field(min_length=1) + query: str = pydantic.Field(min_length=1) + + +class BundleManifestModel(pydantic.BaseModel): + """Strongly typed manifest model for bundle conformance validation.""" + + model_config = pydantic.ConfigDict(extra="forbid") + + intent: str = pydantic.Field(min_length=1) + surface_prefix: str = pydantic.Field(min_length=1) + archetype: str = pydantic.Field(min_length=1) + empty_fallback_text: str = pydantic.Field(min_length=1) + conformance_example: dict[str, Any] + few_shot_examples: list[ConformanceFewShotExample] = pydantic.Field( + min_length=1 + ) + tool_description: str | None = None + + @pydantic.model_validator(mode="after") + def validate_few_shot_intents(self) -> "BundleManifestModel": + for example in self.few_shot_examples: + if example.intent != self.intent: + raise ValueError( + f"Few-shot example intent '{example.intent}' does not match" + f" manifest intent '{self.intent}'." + ) + return self + + +class TestTemplatesConformance(unittest.TestCase): + """Validates template bundle structure, schemas, and layouts.""" + + maxDiff = None + + @classmethod + def setUpClass(cls): + super().setUpClass() + adk_features.override_feature_enabled( + adk_features.FeatureName.SNAKE_CASE_SKILL_NAME, True + ) + cls.package_dir = pathlib.Path(__file__).parent + extension_path = ( + cls.package_dir / "shared" / "schema" / "maps_catalog_extension.json" + ) + schema_manager = a2ui.schema.manager.A2uiSchemaManager( + version=a2ui.schema.constants.VERSION_0_9, + catalogs=[ + a2ui.schema.catalog.CatalogConfig( + name="maps-agentic-ui-catalog", + provider=agent.MergedCatalogProvider( + a2ui.schema.constants.VERSION_0_9, str(extension_path) + ), + ) + ], + schema_modifiers=[ + a2ui.schema.common_modifiers.remove_strict_validation + ], + ) + cls.validator = schema_manager.get_selected_catalog().validator + + # Discover through the production registry so this suite validates the + # bundles the agent will actually serve, not a parallel reimplementation. + cls.bundles = template_registry.TemplateRegistry().list_bundles() + + def _load_bundle_manifest( + self, bundle: template_registry.TemplateBundle + ) -> BundleManifestModel: + """Loads and validates manifest.json against BundleManifestModel.""" + manifest_path = bundle.bundle_dir / "manifest.json" + self.assertTrue( + manifest_path.exists(), + f"Manifest file {manifest_path} does not exist for bundle" + f" {bundle.name}", + ) + return BundleManifestModel.model_validate_json( + manifest_path.read_text(encoding="utf-8") + ) + + def test_discovered_at_least_one_bundle(self): + """Ensures at least local_search and directions bundles are discovered.""" + bundle_names = {bundle.name for bundle in self.bundles} + self.assertIn("local_search", bundle_names) + self.assertIn("directions", bundle_names) + + def test_bundle_manifest_conformance(self): + """Verifies manifest.json contains required metadata and examples.""" + for bundle in self.bundles: + with self.subTest(bundle=bundle.name): + self._load_bundle_manifest(bundle) + + # `clamped_collections` is optional, but when a bundle declares it + # every name must be a real field, or the extractor prompt would tell + # the model to limit something that does not exist. + for collection in getattr(bundle.schema_cls, "clamped_collections", ()): + self.assertIn(collection, bundle.schema_cls.model_fields) + + def test_bundle_schema_exports_extractor_schema(self): + """Verifies that schema.py exports an ExtractorSchema BaseModel.""" + for bundle in self.bundles: + with self.subTest(bundle=bundle.name): + self.assertTrue((bundle.bundle_dir / "schema.py").exists()) + self.assertTrue(issubclass(bundle.schema_cls, pydantic.BaseModel)) + + def test_bundle_schema_validates_conformance_example(self): + """Verifies each bundle's conformance_example passes schema validation.""" + for bundle in self.bundles: + with self.subTest(bundle=bundle.name): + manifest = self._load_bundle_manifest(bundle) + + # Directly validate against the Pydantic schema class to ensure + # invalid data fails loudly at build time rather than falling back. + validated_model = bundle.schema_cls.model_validate( + manifest.conformance_example + ) + self.assertIsNotNone(validated_model) + + def test_bundle_layout_placeholder_alignment(self): + """Verifies every {{placeholder}} in layout.json exists in schema fields.""" + framework_placeholders = {"surface_id", "catalog_id"} + + for bundle in self.bundles: + with self.subTest(bundle=bundle.name): + self.assertTrue(bundle.layout_path.exists()) + with open(bundle.layout_path, "r", encoding="utf-8") as layout_file: + layout_data = json.load(layout_file) + + model_fields = set(bundle.schema_cls.model_fields.keys()) + computed_fields = set( + getattr( + bundle.schema_cls, "__pydantic_computed_fields__", {} + ).keys() + ) + available_fields = ( + model_fields | computed_fields | framework_placeholders + ) + + placeholders = _extract_placeholders(layout_data) + missing_fields = placeholders - available_fields + self.assertFalse( + missing_fields, + f"Bundle '{bundle.name}' layout.json contains placeholders with no" + " matching field or computed_field in ExtractorSchema:" + f" {missing_fields}", + ) + + def test_bundle_skill_conformance(self): + """Verifies that SKILL.md conforms to the ADK skill specification.""" + for bundle in self.bundles: + with self.subTest(bundle=bundle.name): + skill = adk_skills.load_skill_from_dir(bundle.skill_dir) + self.assertEqual(skill.frontmatter.name, bundle.name) + self.assertTrue(skill.frontmatter.description) + self.assertGreater(len(skill.instructions), 0) + self.assertEqual(bundle.skill_instructions, skill.instructions) + + def test_bundle_catalog_schema_compliance(self): + """Verifies that rendered layouts pass A2UI catalog schema validation.""" + for bundle in self.bundles: + with self.subTest(bundle=bundle.name): + manifest = self._load_bundle_manifest(bundle) + merged_messages = merge_template( + bundle.name, manifest.conformance_example + ) + self.validator.validate(merged_messages) + + def test_merged_output_matches_golden(self): + """Verifies rendered layouts match checked-in golden A2UI wire messages.""" + for bundle in self.bundles: + with self.subTest(bundle=bundle.name): + golden_file = bundle.bundle_dir / "testdata" / "merged.golden.json" + self.assertTrue( + golden_file.exists(), + f"Golden file {golden_file} does not exist for bundle" + f" {bundle.name}", + ) + expected_messages = json.loads(golden_file.read_text(encoding="utf-8")) + + manifest = self._load_bundle_manifest(bundle) + actual_messages = merge_template( + bundle.name, manifest.conformance_example + ) + self.assertEqual(actual_messages, expected_messages) + + +if __name__ == "__main__": + unittest.main() diff --git a/agent/python_agent/router_config.py b/agent/python_agent/testdata/router_instruction.golden.txt similarity index 59% rename from agent/python_agent/router_config.py rename to agent/python_agent/testdata/router_instruction.golden.txt index 4e137f0..1d27a0a 100644 --- a/agent/python_agent/router_config.py +++ b/agent/python_agent/testdata/router_instruction.golden.txt @@ -1,41 +1,3 @@ -# Copyright 2026 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. -"""Configuration for the Intent Router, including prompts and schemas.""" - -import enum -import pydantic - - -class IntentClass(str, enum.Enum): - """Supported intent categories for routing.""" - - LOCAL_SEARCH = "LOCAL_SEARCH" - DIRECTIONS = "DIRECTIONS" - OTHER_SPATIAL = "OTHER_SPATIAL" - TEXT_ONLY = "TEXT_ONLY" - - -class RouterClassification(pydantic.BaseModel): - """Schema for the classification output of the intent router.""" - - intent: IntentClass = pydantic.Field( - description="The classified intent archetype name." - ) - query: str = pydantic.Field(description="The cleaned user query") - - -ROUTER_SYSTEM_INSTRUCTION = """ ## Role You are an expert query intent router. @@ -43,9 +5,9 @@ class RouterClassification(pydantic.BaseModel): Analyze a user's input and classify it into the most appropriate intent category based on the structural complexity and data requirements of the request. ## Intent Archetypes -- **LOCAL_SEARCH**: Search for categories of interest, places, businesses, or points of interest within a specific geographic proximity. - **DIRECTIONS**: Standard navigation, route directions, walking/driving/transit times, or navigation instructions between an origin and destination. This includes multi-stop routes and routes with specified waypoints. -- **OTHER_SPATIAL**: Queries requiring rich map-based visualization, boundaries, coordinates, specific geographic displays, or complex navigation combining routes with secondary overlays (e.g., weather, air quality forecasts, displaying all available charging stations along a route). +- **LOCAL_SEARCH**: Search for categories of interest, places, businesses, or points of interest within a specific geographic proximity, or display one or more specific named places, buildings, landmarks, or addresses on a map. +- **OTHER_SPATIAL**: Queries requiring polygon boundaries, raw coordinate plotting, or complex navigation combining routes with secondary overlays (e.g., weather, air quality forecasts, displaying all available charging stations along a route). Note: Requests to show one or more specific named places, buildings, or addresses on a map belong to LOCAL_SEARCH, not OTHER_SPATIAL. - **TEXT_ONLY**: General information retrieval, questions, or requests for data associated with locations that can be answered fully with text without requiring a map interface. ## Classification Policy: Conservative Routing @@ -55,6 +17,14 @@ class RouterClassification(pydantic.BaseModel): ## Few-shot Examples ### Example 1 +**User Query:** "Directions from Sacramento to Mendocino via Clear Lake" +**Output:** +{ + "intent": "DIRECTIONS", + "query": "Directions from Sacramento to Mendocino via Clear Lake" +} + +### Example 2 **User Query:** "Find coffee shops near Central Park" **Output:** { @@ -62,7 +32,15 @@ class RouterClassification(pydantic.BaseModel): "query": "coffee shops near Central Park" } -### Example 2 +### Example 3 +**User Query:** "Show me GWC2 and google 1950 in the same map" +**Output:** +{ + "intent": "LOCAL_SEARCH", + "query": "GWC2 and Google 1950" +} + +### Example 4 **User Query:** "Walking route from Central Park to Times Square, but show coffee shops and rain forecasts along the way" **Output:** { @@ -70,19 +48,10 @@ class RouterClassification(pydantic.BaseModel): "query": "walking route from Central Park to Times Square with coffee shop stops and rain forecast" } -### Example 3 +### Example 5 **User Query:** "show the boundary of Yosemite National Park on the map" **Output:** { "intent": "OTHER_SPATIAL", "query": "boundary of Yosemite National Park" } - -### Example 4 -**User Query:** "Directions from Sacramento to Mendocino via Clear Lake" -**Output:** -{ - "intent": "DIRECTIONS", - "query": "Directions from Sacramento to Mendocino via Clear Lake" -} -""" diff --git a/client/android/GoogleMapsA2UI/build.gradle b/client/android/GoogleMapsA2UI/build.gradle index 702c9a0..be243f9 100644 --- a/client/android/GoogleMapsA2UI/build.gradle +++ b/client/android/GoogleMapsA2UI/build.gradle @@ -1,3 +1,17 @@ +// Copyright 2026 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 +// +// http://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. + plugins { id 'com.android.library' version '9.0.0' id 'maven-publish' @@ -31,6 +45,16 @@ android { } } + // Robolectric needs the AGP-merged manifest, resources and assets. Without + // this, `context.assets` is empty and Robolectric cannot read `targetSdk`, + // so it falls back to its minimum supported SDK where API 23/24 WebViewClient + // overloads do not exist. + testOptions { + unitTests { + includeAndroidResources = true + } + } + publishing { singleVariant('release') } @@ -52,6 +76,9 @@ dependencies { implementation 'androidx.appcompat:appcompat:1.6.1' testImplementation 'junit:junit:4.13.2' testImplementation 'org.robolectric:robolectric:4.11.1' + testImplementation 'com.google.truth:truth:1.4.2' + testImplementation 'org.mockito:mockito-core:5.11.0' + testImplementation 'org.mockito.kotlin:mockito-kotlin:5.2.1' } afterEvaluate { diff --git a/client/android/GoogleMapsA2UI/src/main/assets/index.html b/client/android/GoogleMapsA2UI/src/main/assets/index.html index 0f35689..d624775 100644 --- a/client/android/GoogleMapsA2UI/src/main/assets/index.html +++ b/client/android/GoogleMapsA2UI/src/main/assets/index.html @@ -21,7 +21,12 @@ <meta charset="UTF-8" /> <meta name="viewport" content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no" /> <title>A2UI Bridge - + - + diff --git a/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/A2AResponseParser.kt b/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/A2AResponseParser.kt index 56d43bc..ed6ad25 100644 --- a/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/A2AResponseParser.kt +++ b/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/A2AResponseParser.kt @@ -17,6 +17,12 @@ package com.google.android.libraries.mapsplatform.a2ui import org.json.JSONArray import org.json.JSONObject +sealed class A2AParserException(message: String) : Exception(message) { + class InvalidPayloadStructure( + message: String = "The raw JSON response lacks a recognized message parts structure." + ) : A2AParserException(message) +} + data class ParsedA2AEventMetadata(val mimeType: String?) sealed interface ParsedA2AEvent { @@ -38,7 +44,6 @@ object A2AResponseParser { private const val KEY_CONTENT = "content" private const val KEY_STATUS = "status" private const val KEY_MESSAGE = "message" - private const val KEY_RESULT = "result" private const val KEY_CREATE_SURFACE = "createSurface" private const val KEY_UPDATE_COMPONENTS = "updateComponents" @@ -46,98 +51,53 @@ object A2AResponseParser { private const val KEY_DELETE_SURFACE = "deleteSurface" private const val KEY_SURFACE_ID = "surfaceId" + private val A2UI_PATTERN = "(.*?)".toRegex(RegexOption.DOT_MATCHES_ALL) + fun parse(rawJson: JSONObject): List { - val partsList = mutableListOf() - val partsArray = extractPartsArray(rawJson) - - if (partsArray != null) { - var currentTextBuilder = java.lang.StringBuilder() - var currentUiElements = JSONArray() - - for (i in 0 until partsArray.length()) { - val part = partsArray.getJSONObject(i) - val textPart = - if (part.has(KEY_TEXT)) part.optString(KEY_TEXT) - else if (part.optString(KEY_KIND) == KEY_TEXT) part.optString(KEY_TEXT) else null - - if (textPart != null) { - if (currentUiElements.length() > 0) { - partsList.add(ParsedA2AEvent.Data(currentUiElements.toString())) - currentUiElements = JSONArray() - } + val partsArray = + extractPartsArray(rawJson) ?: throw A2AParserException.InvalidPayloadStructure() - if (textPart.contains("---a2ui_JSON---") || textPart.contains("```json")) { - extractJsonBlocks(textPart, currentTextBuilder, partsList) - } else if (textPart.isNotEmpty()) { - if (currentTextBuilder.isNotEmpty()) currentTextBuilder.append("\n") - currentTextBuilder.append(textPart) - } + val partsList = mutableListOf() + var currentTextBuilder = java.lang.StringBuilder() + var currentUiElements = JSONArray() + + for (i in 0 until partsArray.length()) { + val part = partsArray.getJSONObject(i) + val textPart = + if (part.has(KEY_TEXT)) part.optString(KEY_TEXT) + else if (part.optString(KEY_KIND) == KEY_TEXT) part.optString(KEY_TEXT) else null + + if (textPart != null) { + if (currentUiElements.length() > 0) { + partsList.add(ParsedA2AEvent.Data(currentUiElements.toString())) + currentUiElements = JSONArray() } - val dataPayload = - if (part.has(KEY_DATA)) part.optJSONObject(KEY_DATA) - else if (part.optString(KEY_KIND) == KEY_DATA) part.optJSONObject(KEY_DATA) else null - if (dataPayload != null && isUiElement(dataPayload)) { - if (currentTextBuilder.isNotEmpty()) { - partsList.add(ParsedA2AEvent.Text(currentTextBuilder.toString())) - currentTextBuilder = java.lang.StringBuilder() - } - currentUiElements.put(dataPayload) + if (textPart.contains("")) { + extractJsonBlocks(textPart, currentTextBuilder, partsList) + } else if (textPart.isNotEmpty()) { + if (currentTextBuilder.isNotEmpty()) currentTextBuilder.append("\n") + currentTextBuilder.append(textPart) } } - if (currentTextBuilder.isNotEmpty()) { - partsList.add(ParsedA2AEvent.Text(currentTextBuilder.toString())) - } - if (currentUiElements.length() > 0) { - partsList.add(ParsedA2AEvent.Data(currentUiElements.toString())) - } - } else { - try { - val resultObj = rawJson.opt(KEY_RESULT) - if (resultObj is String) { - if (resultObj.isNotEmpty()) { - val firstChar = resultObj.trim().firstOrNull() - if (firstChar == '[') { - val array = JSONArray(resultObj) - val uiElements = JSONArray() - for (j in 0 until array.length()) { - val item = array.optJSONObject(j) - if (item != null && isUiElement(item)) { - uiElements.put(item) - } - } - if (uiElements.length() > 0) { - partsList.add(ParsedA2AEvent.Data(uiElements.toString())) - } - } - } - } else if (resultObj is JSONArray) { - var currentTextBuilder = java.lang.StringBuilder() - val uiElements = JSONArray() - for (j in 0 until resultObj.length()) { - val item = resultObj.optJSONObject(j) - if (item != null) { - if (item.has(KEY_TEXT)) { - val textPart = item.optString(KEY_TEXT) - if (textPart.isNotEmpty()) { - if (currentTextBuilder.isNotEmpty()) currentTextBuilder.append("\n") - currentTextBuilder.append(textPart) - } - } - if (isUiElement(item)) { - uiElements.put(item) - } - } - } - if (currentTextBuilder.isNotEmpty()) { - partsList.add(ParsedA2AEvent.Text(currentTextBuilder.toString())) - } - if (uiElements.length() > 0) { - partsList.add(ParsedA2AEvent.Data(uiElements.toString())) - } + val dataPayload = + if (part.has(KEY_DATA)) part.optJSONObject(KEY_DATA) + else if (part.optString(KEY_KIND) == KEY_DATA) part.optJSONObject(KEY_DATA) else null + if (dataPayload != null && isUiElement(dataPayload)) { + if (currentTextBuilder.isNotEmpty()) { + partsList.add(ParsedA2AEvent.Text(currentTextBuilder.toString())) + currentTextBuilder = java.lang.StringBuilder() } - } catch (e: Exception) {} + currentUiElements.put(dataPayload) + } + } + + if (currentTextBuilder.isNotEmpty()) { + partsList.add(ParsedA2AEvent.Text(currentTextBuilder.toString())) + } + if (currentUiElements.length() > 0) { + partsList.add(ParsedA2AEvent.Data(currentUiElements.toString())) } val deduplicatedParts = mutableListOf() @@ -147,7 +107,7 @@ object A2AResponseParser { for (part in partsList) { val finalText: String? = if (part is ParsedA2AEvent.Text) { - part.text.replace("```json", "").replace("```", "").trim().takeIf { it.isNotEmpty() } + part.text.trim().takeIf { it.isNotEmpty() } } else null // Deduplicate consecutive identical text blocks @@ -226,24 +186,6 @@ object A2AResponseParser { rawJson.has(KEY_CONTENT) -> rawJson.optJSONObject(KEY_CONTENT)?.optJSONArray(KEY_PARTS) rawJson.has(KEY_STATUS) -> rawJson.optJSONObject(KEY_STATUS)?.optJSONObject(KEY_MESSAGE)?.optJSONArray(KEY_PARTS) - rawJson.has(KEY_RESULT) -> { - val resultObj = rawJson.opt(KEY_RESULT) - if (resultObj is String) { - try { - val innerJson = JSONObject(resultObj) - innerJson - .optJSONObject(KEY_STATUS) - ?.optJSONObject(KEY_MESSAGE) - ?.optJSONArray(KEY_PARTS) - } catch (e: Exception) { - null - } - } else if (resultObj is JSONObject) { - resultObj.optJSONObject(KEY_STATUS)?.optJSONObject(KEY_MESSAGE)?.optJSONArray(KEY_PARTS) - } else { - null - } - } else -> null } @@ -253,7 +195,9 @@ object A2AResponseParser { } } - return if (finalParts.length() > 0) finalParts else null + val hasRecognizedStructure = + rawJson.optJSONArray(KEY_HISTORY) != null || additionalParts != null + return if (hasRecognizedStructure) finalParts else null } private fun extractJsonBlocks( @@ -261,12 +205,7 @@ object A2AResponseParser { textBuilder: java.lang.StringBuilder, partsList: MutableList, ) { - val jsonPattern = "```json(.*?)```".toRegex(RegexOption.DOT_MATCHES_ALL) - val a2uiPattern = "---a2ui_JSON---(.*?)---a2ui_JSON_END---".toRegex(RegexOption.DOT_MATCHES_ALL) - - val allMatches = mutableListOf() - allMatches.addAll(jsonPattern.findAll(textPart)) - allMatches.addAll(a2uiPattern.findAll(textPart)) + val allMatches = A2UI_PATTERN.findAll(textPart).toList() if (allMatches.isEmpty()) { if (textBuilder.isNotEmpty()) textBuilder.append("\n") @@ -274,8 +213,6 @@ object A2AResponseParser { return } - allMatches.sortBy { it.range.first } - var lastEnd = 0 for (match in allMatches) { val beforeText = textPart.substring(lastEnd, match.range.first).trim() @@ -288,24 +225,28 @@ object A2AResponseParser { try { val firstChar = jsonString.firstOrNull() if (firstChar == '[') { - if (textBuilder.isNotEmpty()) { - partsList.add(ParsedA2AEvent.Text(textBuilder.toString())) - textBuilder.clear() - } val array = JSONArray(jsonString) val localUiElements = JSONArray() for (i in 0 until array.length()) { localUiElements.put(array.getJSONObject(i)) } + if (textBuilder.isNotEmpty()) { + partsList.add(ParsedA2AEvent.Text(textBuilder.toString())) + textBuilder.clear() + } partsList.add(ParsedA2AEvent.Data(localUiElements.toString())) } else if (firstChar == '{') { + val jsonObj = JSONObject(jsonString) if (textBuilder.isNotEmpty()) { partsList.add(ParsedA2AEvent.Text(textBuilder.toString())) textBuilder.clear() } val localUiElements = JSONArray() - localUiElements.put(JSONObject(jsonString)) + localUiElements.put(jsonObj) partsList.add(ParsedA2AEvent.Data(localUiElements.toString())) + } else { + if (textBuilder.isNotEmpty()) textBuilder.append("\n") + textBuilder.append(match.value) } } catch (e: Exception) { if (textBuilder.isNotEmpty()) textBuilder.append("\n") diff --git a/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/A2UIView.kt b/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/A2UIView.kt index 5043f39..6fe76f3 100644 --- a/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/A2UIView.kt +++ b/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/A2UIView.kt @@ -96,7 +96,9 @@ constructor(context: Context, attrs: AttributeSet? = null, defStyleAttr: Int = 0 request: WebResourceRequest?, error: WebResourceError?, ) { - super.onReceivedError(view, request, error) + if (request != null && error != null) { + super.onReceivedError(view, request, error) + } Log.e( A2UI_ERROR_TAG, "Error loading WebView: ${error?.description}, URL: ${request?.url}", diff --git a/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/WebAppInterface.kt b/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/WebAppInterface.kt index 87c400f..48ca918 100644 --- a/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/WebAppInterface.kt +++ b/client/android/GoogleMapsA2UI/src/main/java/com/google/android/libraries/mapsplatform/a2ui/WebAppInterface.kt @@ -63,19 +63,21 @@ class WebAppInterface(private val webView: WebView, private val a2uiView: A2UIVi @JavascriptInterface fun onWebpageResized(height: Int) { Log.d(TAG, "onWebpageResized: $height") - if (!resized) { - webView.post { - val layoutParams = webView.layoutParams - if (layoutParams != null) { - val newHeight = (height * webView.resources.displayMetrics.density).toInt() + webView.post { + val layoutParams = webView.layoutParams + if (layoutParams != null) { + val newHeight = (height * webView.resources.displayMetrics.density).toInt() + if (layoutParams.height != newHeight) { layoutParams.height = newHeight webView.layoutParams = layoutParams - resized = true Log.d("A2UIViewDebug", "WebView height updated to: $newHeight") + } + if (!resized) { + resized = true a2uiView.onRenderCompleteInternal() - } else { - Log.e("A2UIViewDebug", "WebView LayoutParams is null.") } + } else { + Log.e("A2UIViewDebug", "WebView LayoutParams is null.") } } } diff --git a/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2AResponseParserTest.kt b/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2AResponseParserTest.kt index 4f91ed0..7340d24 100644 --- a/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2AResponseParserTest.kt +++ b/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2AResponseParserTest.kt @@ -14,10 +14,10 @@ package com.google.android.libraries.mapsplatform.a2ui +import com.google.common.truth.Truth.assertThat import org.json.JSONArray import org.json.JSONObject -import org.junit.Assert.assertEquals -import org.junit.Assert.assertTrue +import org.junit.Assert.assertThrows import org.junit.Test import org.junit.runner.RunWith import org.robolectric.RobolectricTestRunner @@ -25,17 +25,56 @@ import org.robolectric.RobolectricTestRunner @RunWith(RobolectricTestRunner::class) class A2AResponseParserTest { + /** + * Verifies that parsing an unexpected payload structure without message parts or results throws + * [A2AParserException.InvalidPayloadStructure]. + */ @Test - fun testParse_InvalidPayloadStructure() { - // Test that an unexpected payload format returns an empty event list safely + fun parse_invalidPayloadStructure_throwsException() { val payloadWithNoParts = JSONObject().apply { put("status", "ok") } - val events = A2AResponseParser.parse(payloadWithNoParts) - assertEquals(0, events.size) + val thrown = + assertThrows(A2AParserException.InvalidPayloadStructure::class.java) { + A2AResponseParser.parse(payloadWithNoParts) + } + assertThat(thrown).hasMessageThat().contains("lacks a recognized message parts structure") } + /** + * Verifies that when "history" is present but not an array, the parser throws + * [A2AParserException.InvalidPayloadStructure]. + */ @Test - fun testParse_SimpleTextPart() { - // Test standard text extraction from a standard payload structure + fun parse_invalidHistoryType_throwsException() { + val payloadWithInvalidHistory = JSONObject().apply { put("history", "not_an_array") } + val thrown = + assertThrows(A2AParserException.InvalidPayloadStructure::class.java) { + A2AResponseParser.parse(payloadWithInvalidHistory) + } + assertThat(thrown).hasMessageThat().contains("lacks a recognized message parts structure") + } + + /** Verifies that when the parts array is empty, the parser returns an empty list. */ + @Test + fun parse_emptyParts_returnsEmptyList() { + val emptyPartsPayload = JSONObject("""{"parts": []}""") + val emptyEvents = A2AResponseParser.parse(emptyPartsPayload) + assertThat(emptyEvents).isEmpty() + } + + /** + * Verifies that when the parts array contains unknown types, the parser returns an empty list + * safely without throwing an exception. + */ + @Test + fun parse_unknownParts_returnsEmptyList() { + val unknownPartsPayload = JSONObject("""{"parts": [{"kind": "unsupported_media"}]}""") + val unknownEvents = A2AResponseParser.parse(unknownPartsPayload) + assertThat(unknownEvents).isEmpty() + } + + /** Verifies that a simple text part in the "parts" array is parsed into a single text event. */ + @Test + fun parse_simpleTextPart_returnsSingleTextEvent() { val payload = JSONObject( """ @@ -49,19 +88,126 @@ class A2AResponseParserTest { ) val events = A2AResponseParser.parse(payload) - assertEquals(1, events.size) - val textEvent = events[0] as ParsedA2AEvent.Text - assertEquals("Show me some good sushi in Seattle", textEvent.text) + assertThat(events).containsExactly(ParsedA2AEvent.Text("Show me some good sushi in Seattle")) + } + + /** Verifies that the parser resolves message parts nested under "content.parts". */ + @Test + fun parse_contentPartsPath_resolvesTextEvent() { + val payload = + JSONObject( + """ + { + "content": { + "parts": [ + {"text": "Welcome to Seattle!"} + ] + } + } + """ + .trimIndent() + ) + + val events = A2AResponseParser.parse(payload) + assertThat(events).containsExactly(ParsedA2AEvent.Text("Welcome to Seattle!")) + } + + /** Verifies that the parser resolves message parts nested under "status.message.parts". */ + @Test + fun parse_statusMessagePartsPath_resolvesTextEvent() { + val payload = + JSONObject( + """ + { + "status": { + "message": { + "parts": [ + {"text": "Status message text"} + ] + } + } + } + """ + .trimIndent() + ) + + val events = A2AResponseParser.parse(payload) + assertThat(events).containsExactly(ParsedA2AEvent.Text("Status message text")) + } + + /** + * Verifies that the parser resolves message parts from the "history" array, extracting agent + * parts that appear after the last user message. + */ + @Test + fun parse_historyPayload_resolvesAgentPartsAfterLastUser() { + val payload = + JSONObject( + """ + { + "history": [ + { + "role": "user", + "parts": [{"text": "First query"}] + }, + { + "role": "agent", + "parts": [{"text": "First answer"}] + }, + { + "role": "user", + "parts": [{"text": "Show me Seattle maps"}] + }, + { + "role": "agent", + "parts": [{"text": "Here is Seattle!"}] + } + ] + } + """ + .trimIndent() + ) + + val events = A2AResponseParser.parse(payload) + assertThat(events).containsExactly(ParsedA2AEvent.Text("Here is Seattle!")) + } + + /** + * Verifies that when a history payload has only user messages and no agent reply yet, the parser + * returns an empty list. + */ + @Test + fun parse_historyPayloadWithNoAgentResponse_returnsEmptyList() { + val payload = + JSONObject( + """ + { + "history": [ + { + "role": "user", + "parts": [{"text": "Hello, agent!"}] + } + ] + } + """ + .trimIndent() + ) + + val events = A2AResponseParser.parse(payload) + assertThat(events).isEmpty() } + /** + * Verifies that multiple text parts in the payload are concatenated with newlines into a single + * text event. + */ @Test - fun testParse_TextConcatenation() { - // Android parser concatenates all text elements within a JSONArray into a single text event + fun parse_multipleTextParts_concatenatesIntoSingleTextEvent() { val payload = JSONObject( """ { - "result": [ + "parts": [ {"text": "Hello Seattle!"}, {"text": "Hello Seattle!"}, {"text": "Different text."} @@ -72,44 +218,278 @@ class A2AResponseParserTest { ) val events = A2AResponseParser.parse(payload) - assertEquals(1, events.size) // Expecting 1 because texts are concatenated - assertEquals( - "Hello Seattle!\nHello Seattle!\nDifferent text.", - (events[0] as ParsedA2AEvent.Text).text, - ) + assertThat(events) + .containsExactly(ParsedA2AEvent.Text("Hello Seattle!\nHello Seattle!\nDifferent text.")) + } + + /** Verifies that parts containing empty text strings produce an empty event list. */ + @Test + fun parse_emptyTextPart_returnsEmptyList() { + val payload = + JSONObject( + """ + { + "parts": [ + {"text": ""} + ] + } + """ + .trimIndent() + ) + + val events = A2AResponseParser.parse(payload) + assertThat(events).isEmpty() + } + + /** + * Verifies that non-array parts structure throws [A2AParserException.InvalidPayloadStructure]. + */ + @Test + fun parse_invalidPartsType_throwsException() { + val payload = JSONObject("""{"parts": "not_an_array"}""") + assertThrows(A2AParserException.InvalidPayloadStructure::class.java) { + A2AResponseParser.parse(payload) + } + } + + /** + * Verifies that embedded '' tags within a text body are extracted into discrete text + * and A2UI data events in sequential order. + */ + @Test + fun parse_embeddedA2UIJSON_extractsTextAndDataEventsSequentially() { + val textWithJson = + "Here is the map: [{\"createSurface\": {\"surfaceId\": \"seattle-map\"}}] Hope this helps!" + val payload = + JSONObject( + """ + { + "parts": [ + {"kind": "text", "text": ${JSONObject.quote(textWithJson)}} + ] + } + """ + .trimIndent() + ) + + val events = A2AResponseParser.parse(payload) + assertThat(events).hasSize(3) + assertThat(events[0]).isEqualTo(ParsedA2AEvent.Text("Here is the map:")) + + val dataEvent = events[1] as ParsedA2AEvent.Data + val arr = JSONArray(dataEvent.data) + assertThat(arr.getJSONObject(0).getJSONObject("createSurface").getString("surfaceId")) + .isEqualTo("seattle-map") + + assertThat(events[2]).isEqualTo(ParsedA2AEvent.Text("Hope this helps!")) + } + + /** + * Verifies that an A2UI JSON array embedded inside a text part is parsed into a flat array of + * components. + */ + @Test + fun parse_embeddedA2UIJSONArray_flattensToSingleArray() { + val textWithJsonArray = + "Here is the map: [{\"createSurface\": {\"surfaceId\": \"sushi-seattle\"}}, {\"updateComponents\": {\"surfaceId\": \"sushi-seattle\"}}] Done." + val payload = + JSONObject( + """ + { + "parts": [ + {"kind": "text", "text": ${JSONObject.quote(textWithJsonArray)}} + ] + } + """ + .trimIndent() + ) + + val events = A2AResponseParser.parse(payload) + assertThat(events).hasSize(3) + assertThat(events[0]).isEqualTo(ParsedA2AEvent.Text("Here is the map:")) + + val dataEvent = events[1] as ParsedA2AEvent.Data + val arr = JSONArray(dataEvent.data) + assertThat(arr.length()).isEqualTo(2) + assertThat(arr.getJSONObject(0).has("createSurface")).isTrue() + assertThat(arr.getJSONObject(1).has("updateComponents")).isTrue() + + assertThat(events[2]).isEqualTo(ParsedA2AEvent.Text("Done.")) + } + + /** + * Verifies that malformed JSON inside an embedded '' tag falls back gracefully to + * plain text without crashing. + */ + @Test + fun parse_malformedEmbeddedA2UITag_fallsBackToPlainText() { + val textWithMalformedJson = "Intro text {not_valid_json outro text" + val payload = + JSONObject( + """ + { + "parts": [ + {"kind": "text", "text": ${JSONObject.quote(textWithMalformedJson)}} + ] + } + """ + .trimIndent() + ) + + val events = A2AResponseParser.parse(payload) + assertThat(events) + .containsExactly( + ParsedA2AEvent.Text("Intro text\n{not_valid_json\noutro text") + ) } + /** + * Verifies that non-object/non-array JSON inside an embedded '' tag falls back to + * plain text. + */ @Test - fun testParse_StringifiedJsonResultArray() { - // Tests the scenario where 'result' contains a stringified JSON array starting with '[' - val stringifiedArray = """[{"createSurface": {"surfaceId": "sushi-seattle"}}]""" - val payload = JSONObject().apply { put("result", stringifiedArray) } + fun parse_nonObjectOrArrayEmbeddedA2UITag_fallsBackToPlainText() { + val textWithPrimitiveJson = "Intro text 12345 outro text" + val payload = + JSONObject( + """ + { + "parts": [ + {"kind": "text", "text": ${JSONObject.quote(textWithPrimitiveJson)}} + ] + } + """ + .trimIndent() + ) val events = A2AResponseParser.parse(payload) - assertEquals(1, events.size) + assertThat(events) + .containsExactly(ParsedA2AEvent.Text("Intro text\n12345\noutro text")) + } + + /** Verifies that an unclosed '' tag remains as plain text. */ + @Test + fun parse_unclosedA2UITag_fallsBackToPlainText() { + val unclosedTagText = "Before {\"createSurface\": {}} without closing tag" + val payload = + JSONObject( + """ + { + "parts": [ + {"text": ${JSONObject.quote(unclosedTagText)}} + ] + } + """ + .trimIndent() + ) + + val events = A2AResponseParser.parse(payload) + assertThat(events) + .containsExactly( + ParsedA2AEvent.Text("Before {\"createSurface\": {}} without closing tag") + ) + } + + /** + * Verifies that duplicate surface creation definitions with identical surface IDs are + * deduplicated within the payload. + */ + @Test + fun parse_duplicateSurfaceIds_deduplicatesSurfaces() { + val payload = + JSONObject( + """ + { + "parts": [ + { + "data": { + "createSurface": {"surfaceId": "dup-surface"} + } + }, + { + "data": { + "createSurface": {"surfaceId": "dup-surface"} + } + } + ] + } + """ + .trimIndent() + ) + val events = A2AResponseParser.parse(payload) + assertThat(events).hasSize(1) val dataEvent = events[0] as ParsedA2AEvent.Data - val a2uiArray = JSONArray(dataEvent.data) - assertEquals(1, a2uiArray.length()) - assertTrue(a2uiArray.getJSONObject(0).has("createSurface")) + val arr = JSONArray(dataEvent.data) + assertThat(arr.length()).isEqualTo(1) + } + + /** + * Verifies that multiple '' tags within a single text part are all extracted + * sequentially. + */ + @Test + fun parse_multipleEmbeddedA2UITags_extractsAllSequentially() { + val textWithMultipleTags = + "First map: {\"createSurface\": {\"surfaceId\": \"sushi\"}} Then: {\"updateComponents\": {\"surfaceId\": \"sushi\"}} Done." + val payload = + JSONObject( + """ + { + "parts": [ + {"kind": "text", "text": ${JSONObject.quote(textWithMultipleTags)}} + ] + } + """ + .trimIndent() + ) + + val events = A2AResponseParser.parse(payload) + assertThat(events).hasSize(5) + assertThat(events[0]).isEqualTo(ParsedA2AEvent.Text("First map:")) + + val dataEvent1 = events[1] as ParsedA2AEvent.Data + val arr1 = JSONArray(dataEvent1.data) + assertThat(arr1.getJSONObject(0).getJSONObject("createSurface").getString("surfaceId")) + .isEqualTo("sushi") + + assertThat(events[2]).isEqualTo(ParsedA2AEvent.Text("Then:")) + + val dataEvent2 = events[3] as ParsedA2AEvent.Data + val arr2 = JSONArray(dataEvent2.data) + assertThat(arr2.getJSONObject(0).getJSONObject("updateComponents").getString("surfaceId")) + .isEqualTo("sushi") + + assertThat(events[4]).isEqualTo(ParsedA2AEvent.Text("Done.")) } + /** + * Verifies that consecutive data parts identified as A2UI payloads are batched together into a + * single data event. + */ @Test - fun testParse_NativeJsonResultArray() { - // Tests the newly added support for native JSONArray inside the 'result' key (from PR #311 - // fixes) + fun parse_consecutiveA2UIPayloads_batchesIntoSingleEvent() { val payload = JSONObject( """ { - "result": [ + "parts": [ { - "text": "Here is your native array map" + "kind": "data", + "data": { + "createSurface": { + "surfaceId": "sushi-seattle", + "catalogId": "a2ui://maps-agentic-ui-catalog.json" + } + } }, { - "updateComponents": { - "surfaceId": "sushi-seattle", - "components": [] + "kind": "data", + "data": { + "updateComponents": { + "surfaceId": "sushi-seattle", + "components": [] + } } } ] @@ -119,16 +499,82 @@ class A2AResponseParserTest { ) val events = A2AResponseParser.parse(payload) + assertThat(events).hasSize(1) + val dataEvent = events[0] as ParsedA2AEvent.Data + val arr = JSONArray(dataEvent.data) + assertThat(arr.length()).isEqualTo(2) + assertThat(arr.getJSONObject(0).has("createSurface")).isTrue() + assertThat(arr.getJSONObject(1).has("updateComponents")).isTrue() + } - // We expect one Text event and one Data event - assertEquals(2, events.size) + /** + * Verifies that an A2UI batch is finalized and a new one starts if interrupted by a text part. + */ + @Test + fun parse_a2uiBatchInterruptedByTextPart_createsSeparateBatches() { + val payload = + JSONObject( + """ + { + "parts": [ + { + "kind": "data", + "data": {"createSurface": {"surfaceId": "sushi-seattle"}} + }, + {"kind": "text", "text": "Middle Text explaining the surface"}, + { + "kind": "data", + "data": {"updateComponents": {"surfaceId": "sushi-seattle"}} + } + ] + } + """ + .trimIndent() + ) - val textEvent = events[0] as ParsedA2AEvent.Text - assertEquals("Here is your native array map", textEvent.text) + val events = A2AResponseParser.parse(payload) + assertThat(events).hasSize(3) - val dataEvent = events[1] as ParsedA2AEvent.Data - val a2uiArray = JSONArray(dataEvent.data) - assertEquals(1, a2uiArray.length()) - assertTrue(a2uiArray.getJSONObject(0).has("updateComponents")) + val data1 = events[0] as ParsedA2AEvent.Data + assertThat(JSONArray(data1.data).length()).isEqualTo(1) + + assertThat(events[1]).isEqualTo(ParsedA2AEvent.Text("Middle Text explaining the surface")) + + val data2 = events[2] as ParsedA2AEvent.Data + assertThat(JSONArray(data2.data).length()).isEqualTo(1) + } + + /** + * Verifies that a data part is recognized as an A2UI payload if it contains recognized keys, even + * without a mime type. + */ + @Test + fun parse_dataPartWithImplicitA2UIKey_returnsDataEvent() { + val payload = + JSONObject( + """ + { + "parts": [ + { + "kind": "data", + "data": { + "createSurface": { + "surfaceId": "sushi-seattle", + "catalogId": "a2ui://maps-agentic-ui-catalog.json" + } + } + } + ] + } + """ + .trimIndent() + ) + + val events = A2AResponseParser.parse(payload) + assertThat(events).hasSize(1) + val dataEvent = events[0] as ParsedA2AEvent.Data + val arr = JSONArray(dataEvent.data) + assertThat(arr.length()).isEqualTo(1) + assertThat(arr.getJSONObject(0).has("createSurface")).isTrue() } } diff --git a/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIAttributionIdTests.kt b/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIAttributionIdTests.kt index ed7473c..7a4ab99 100644 --- a/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIAttributionIdTests.kt +++ b/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIAttributionIdTests.kt @@ -33,16 +33,16 @@ class A2UIAttributionIdTests { val htmlContent = InputStreamReader(inputStream).readText() assertTrue( - "Expected the generated attribution ID to contain the Web identifier 'gmp_web_maui_v0.1.8_atoui'", - htmlContent.contains("gmp_web_maui_v0.1.8_atoui"), + "Expected the generated attribution ID to contain the Web identifier 'gmp_web_maui_v0.1.9_atoui'", + htmlContent.contains("gmp_web_maui_v0.1.9_atoui"), ) assertTrue( - "Expected the generated attribution ID to contain the Android identifier 'gmp_android_maui_v0.1.8_atoui'", - htmlContent.contains("gmp_android_maui_v0.1.8_atoui"), + "Expected the generated attribution ID to contain the Android identifier 'gmp_android_maui_v0.1.9_atoui'", + htmlContent.contains("gmp_android_maui_v0.1.9_atoui"), ) assertTrue( - "Expected the generated attribution ID to contain the combined Web and Android identifier string 'gmp_web_maui_v0.1.8_atoui,gmp_android_maui_v0.1.8_atoui'", - htmlContent.contains("gmp_web_maui_v0.1.8_atoui,gmp_android_maui_v0.1.8_atoui"), + "Expected the generated attribution ID to contain the combined Web and Android identifier string 'gmp_web_maui_v0.1.9_atoui,gmp_android_maui_v0.1.9_atoui'", + htmlContent.contains("gmp_web_maui_v0.1.9_atoui,gmp_android_maui_v0.1.9_atoui"), ) } } diff --git a/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIServicesTest.kt b/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIServicesTest.kt new file mode 100644 index 0000000..306e93f --- /dev/null +++ b/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIServicesTest.kt @@ -0,0 +1,74 @@ +// Copyright 2026 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 +// +// http://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. + +package com.google.android.libraries.mapsplatform.a2ui + +import com.google.common.truth.Truth.assertThat +import org.junit.After +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.junit.runners.JUnit4 + +/** + * Unit tests for [A2UIServices]. + * + * Verifies global configuration, including Google Maps API Key provisioning. + */ +@RunWith(JUnit4::class) +class A2UIServicesTest { + + @Before + fun setUp() { + A2UIServices.provideAPIKey("") + } + + @After + fun tearDown() { + A2UIServices.provideAPIKey("") + } + + /** Verifies that the default API key is empty when initialized/reset. */ + @Test + fun defaultState_apiKeyIsEmpty() { + assertThat(A2UIServices.apiKey).isEmpty() + } + + /** Verifies that provideAPIKey() correctly sets the static apiKey property. */ + @Test + fun provideAPIKey_setsApiKey() { + A2UIServices.provideAPIKey("AIzaSyTestApiKey") + assertThat(A2UIServices.apiKey).isEqualTo("AIzaSyTestApiKey") + } + + /** Verifies that subsequent calls to provideAPIKey() override previous keys. */ + @Test + fun provideAPIKey_consecutiveUpdates_overridesKey() { + A2UIServices.provideAPIKey("FirstKey") + assertThat(A2UIServices.apiKey).isEqualTo("FirstKey") + + A2UIServices.provideAPIKey("SecondKey") + assertThat(A2UIServices.apiKey).isEqualTo("SecondKey") + } + + /** Verifies that blank or special character keys are stored verbatim. */ + @Test + fun provideAPIKey_blankOrSpecialCharacters_storesExactString() { + A2UIServices.provideAPIKey("AIzaSyTest-Key_123$!@#") + assertThat(A2UIServices.apiKey).isEqualTo("AIzaSyTest-Key_123$!@#") + + A2UIServices.provideAPIKey("") + assertThat(A2UIServices.apiKey).isEmpty() + } +} diff --git a/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIViewTest.kt b/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIViewTest.kt new file mode 100644 index 0000000..eb08de7 --- /dev/null +++ b/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/A2UIViewTest.kt @@ -0,0 +1,320 @@ +// Copyright 2026 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 +// +// http://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. + +package com.google.android.libraries.mapsplatform.a2ui + +import android.app.Activity +import android.content.Intent +import android.net.Uri +import android.webkit.WebResourceError +import android.webkit.WebResourceRequest +import com.google.common.truth.Truth.assertThat +import org.junit.After +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.mockito.kotlin.mock +import org.mockito.kotlin.verify +import org.mockito.kotlin.whenever +import org.robolectric.Robolectric +import org.robolectric.RobolectricTestRunner +import org.robolectric.Shadows.shadowOf +import org.robolectric.shadows.ShadowLooper + +/** + * Unit tests for [A2UIView]. + * + * Verifies rendering configuration, HTML asset loading, JavaScript injection, deduplication, + * latency measurement, and URL navigation overriding. + */ +@RunWith(RobolectricTestRunner::class) +class A2UIViewTest { + + private lateinit var activity: Activity + private lateinit var a2uiView: A2UIView + + @Before + fun setUp() { + activity = Robolectric.buildActivity(Activity::class.java).setup().get() + A2UIServices.provideAPIKey("AIzaSyValidKey") + a2uiView = A2UIView(activity) + activity.setContentView(a2uiView) + } + + @After + fun tearDown() { + A2UIServices.provideAPIKey("") + } + + private fun createMockRequest(urlStr: String?, isMainFrame: Boolean = true): WebResourceRequest { + val request = mock() + whenever(request.url).thenReturn(urlStr?.let { Uri.parse(it) }) + whenever(request.isForMainFrame).thenReturn(isMainFrame) + whenever(request.method).thenReturn("GET") + whenever(request.requestHeaders).thenReturn(emptyMap()) + return request + } + + /** + * Verifies that render() stores the A2UI payload and loads the HTML asset with the API key + * injected. + */ + @Test + fun render_withApiKey_storesPayloadAndLoadsInjectedHtml() { + val payload = """[{"createSurface": {"surfaceId": "test"}}]""" + + a2uiView.render(payload) + assertThat(a2uiView.a2uiJson).isEqualTo(payload) + + val shadowWebView = shadowOf(a2uiView) + val lastLoadedData = shadowWebView.lastLoadDataWithBaseURL + assertThat(lastLoadedData).isNotNull() + assertThat(lastLoadedData.baseUrl).isEqualTo("file:///android_asset/") + assertThat(lastLoadedData.mimeType).isEqualTo("text/html") + assertThat(lastLoadedData.encoding).isEqualTo("UTF-8") + assertThat(lastLoadedData.data).contains("AIzaSyValidKey") + } + + /** + * Verifies that updateA2uiJson() properly updates the a2uiJson property and schedules JavaScript + * evaluation with safely escaped JSON when JS is ready. + */ + @Test + fun updateA2uiJson_complexJsonPayload_updatesPayloadAndExecutes() { + val complexJson = + """[{"text": "Line 1\nLine 2 with \"quotes\" and 'single' and \\backslash"}]""" + a2uiView.onJsReadyInternal() + a2uiView.updateA2uiJson(complexJson) + ShadowLooper.idleMainLooper() + + assertThat(a2uiView.a2uiJson).isEqualTo(complexJson) + val lastEvaluatedJs = shadowOf(a2uiView).lastEvaluatedJavascript + assertThat(lastEvaluatedJs).isNotNull() + assertThat(lastEvaluatedJs).contains("shell.processA2uiMessages") + assertThat(lastEvaluatedJs).contains("Line 1") + assertThat(lastEvaluatedJs).contains("quotes") + } + + /** + * Verifies that updateA2uiJson() when JS is not ready only stores the payload without executing + * JS. + */ + @Test + fun updateA2uiJson_whenJsNotReady_storesPayloadWithoutEvaluating() { + val payload = """[{"createSurface": {"surfaceId": "pending"}}]""" + + a2uiView.updateA2uiJson(payload) + assertThat(a2uiView.a2uiJson).isEqualTo(payload) + val lastEvaluatedJs = shadowOf(a2uiView).lastEvaluatedJavascript + assertThat(lastEvaluatedJs).isNull() + } + + /** Verifies that onJsReadyInternal() marks the view ready and pushes queued payloads. */ + @Test + fun onJsReadyInternal_withQueuedPayload_triggersUpdate() { + val payload = """[{"createSurface": {"surfaceId": "queued"}}]""" + a2uiView.a2uiJson = payload + + a2uiView.onJsReadyInternal() + ShadowLooper.idleMainLooper() + + assertThat(a2uiView.a2uiJson).isEqualTo(payload) + val lastEvaluatedJs = shadowOf(a2uiView).lastEvaluatedJavascript + assertThat(lastEvaluatedJs).isNotNull() + assertThat(lastEvaluatedJs).contains("queued") + } + + /** + * Verifies that onRenderComplete callback correctly computes latency and receives status, and + * resets startTime so subsequent calls are ignored. + */ + @Test + fun onRenderCompleteInternal_computesLatencyAndInvokesCallback() { + var callCount = 0 + var invokedLatency: Long? = null + var invokedStatus: String? = null + a2uiView.onRenderComplete = { latency, status -> + callCount++ + invokedLatency = latency + invokedStatus = status + } + + val pastStartTime = System.currentTimeMillis() - 150 + a2uiView.render("[]", startTimeMs = pastStartTime) + a2uiView.onRenderCompleteInternal() + + assertThat(invokedLatency).isNotNull() + val latency = checkNotNull(invokedLatency) + assertThat(latency).isAtLeast(150L) + assertThat(invokedStatus).isEqualTo("A2UI Render Complete") + assertThat(callCount).isEqualTo(1) + + // Second call should not invoke callback because startTime was reset to null + a2uiView.onRenderCompleteInternal() + assertThat(callCount).isEqualTo(1) + } + + /** Verifies that onRenderCompleteInternal() does not throw when no callback is registered. */ + @Test + fun onRenderCompleteInternal_withNullCallback_doesNotCrash() { + a2uiView.onRenderComplete = null + a2uiView.render("[]") + a2uiView.onRenderCompleteInternal() + } + + /** + * Verifies that shouldOverrideUrlLoading intercepts web and maps URLs to launch ACTION_VIEW + * intents. + */ + @Test + fun shouldOverrideUrlLoading_webUrl_startsActionViewIntent() { + val request = createMockRequest("https://www.google.com") + + val client = shadowOf(a2uiView).webViewClient + val handled = client.shouldOverrideUrlLoading(a2uiView, request) + assertThat(handled).isTrue() + + val startedIntent = shadowOf(activity).nextStartedActivity + assertThat(startedIntent).isNotNull() + assertThat(startedIntent.action).isEqualTo(Intent.ACTION_VIEW) + assertThat(startedIntent.data).isEqualTo(Uri.parse("https://www.google.com")) + } + + /** + * Verifies that shouldOverrideUrlLoading with Google Maps URLs falls back to browser (null + * package) when the Maps app is not installed on the device. + */ + @Test + fun shouldOverrideUrlLoading_mapsUrlWithoutMapsApp_fallsBackToBrowser() { + val request = createMockRequest("https://maps.google.com/?q=sushi") + + val client = shadowOf(a2uiView).webViewClient + val handled = client.shouldOverrideUrlLoading(a2uiView, request) + assertThat(handled).isTrue() + + val startedIntent = shadowOf(activity).nextStartedActivity + assertThat(startedIntent).isNotNull() + assertThat(startedIntent.action).isEqualTo(Intent.ACTION_VIEW) + assertThat(startedIntent.data).isEqualTo(Uri.parse("https://maps.google.com/?q=sushi")) + assertThat(startedIntent.`package`).isNull() + } + + /** + * Verifies that shouldOverrideUrlLoading also recognizes path-based maps URLs (google.com/maps). + */ + @Test + fun shouldOverrideUrlLoading_pathBasedMapsUrl_interceptsIntent() { + val request = createMockRequest("https://www.google.com/maps/search/pizza") + + val client = shadowOf(a2uiView).webViewClient + val handled = client.shouldOverrideUrlLoading(a2uiView, request) + assertThat(handled).isTrue() + + val startedIntent = shadowOf(activity).nextStartedActivity + assertThat(startedIntent).isNotNull() + assertThat(startedIntent.action).isEqualTo(Intent.ACTION_VIEW) + assertThat(startedIntent.data).isEqualTo(Uri.parse("https://www.google.com/maps/search/pizza")) + } + + /** + * Verifies that non-http and non-https schemes (e.g. javascript:, about:blank) are not handled. + */ + @Test + fun shouldOverrideUrlLoading_nonHttpScheme_returnsFalse() { + val request = createMockRequest("javascript:void(0)") + + val client = shadowOf(a2uiView).webViewClient + val handled = client.shouldOverrideUrlLoading(a2uiView, request) + assertThat(handled).isFalse() + } + + /** Verifies that onReceivedError handles resource loading errors gracefully without crashing. */ + @Test + fun onReceivedError_withNullRequestAndError_logsAndDoesNotCrash() { + val client = shadowOf(a2uiView).webViewClient + client.onReceivedError(a2uiView, null, null) + } + + @Test + fun onReceivedError_withValidRequestAndError_delegatesAndDoesNotCrash() { + val request = createMockRequest("https://maps.googleapis.com/test", isMainFrame = true) + val error = mock() + whenever(error.description).thenReturn("Connection failed") + whenever(error.errorCode).thenReturn(-2) + + val client = shadowOf(a2uiView).webViewClient + client.onReceivedError(a2uiView, request, error) + + verify(request).isForMainFrame + verify(error).errorCode + } + + /** Verifies that secondary constructors initialize without errors. */ + @Test + fun constructor_withAttributeSetAndDefStyle_initializesSuccessfully() { + val viewWithAttrs = A2UIView(activity, null) + val viewWithStyle = A2UIView(activity, null, 0) + assertThat(viewWithAttrs).isNotNull() + assertThat(viewWithStyle).isNotNull() + } + + /** Verifies that shouldOverrideUrlLoading returns false when request or URL is null. */ + @Test + fun shouldOverrideUrlLoading_nullRequestOrUrl_returnsFalse() { + val client = shadowOf(a2uiView).webViewClient + val requestWithNullUrl = createMockRequest(null) + + assertThat(client.shouldOverrideUrlLoading(a2uiView, null as WebResourceRequest?)).isFalse() + assertThat(client.shouldOverrideUrlLoading(a2uiView, requestWithNullUrl)).isFalse() + } + + /** Verifies that shouldOverrideUrlLoading preserves Maps package when Maps app is installed. */ + @Test + fun shouldOverrideUrlLoading_mapsUrlWithInstalledMapsApp_setsMapsPackage() { + val mapsIntent = + Intent(Intent.ACTION_VIEW, Uri.parse("https://maps.google.com/?q=coffee")).apply { + `package` = "com.google.android.apps.maps" + } + val resolveInfo = + android.content.pm.ResolveInfo().apply { + activityInfo = + android.content.pm.ActivityInfo().apply { + packageName = "com.google.android.apps.maps" + name = "com.google.android.maps.MapsActivity" + } + } + shadowOf(activity.packageManager).addResolveInfoForIntent(mapsIntent, resolveInfo) + + val request = createMockRequest("https://maps.google.com/?q=coffee") + + val client = shadowOf(a2uiView).webViewClient + val handled = client.shouldOverrideUrlLoading(a2uiView, request) + assertThat(handled).isTrue() + + val startedIntent = shadowOf(activity).nextStartedActivity + assertThat(startedIntent).isNotNull() + assertThat(startedIntent.`package`).isEqualTo("com.google.android.apps.maps") + } + + /** Verifies that onJsReadyInternal does not evaluate JS when a2uiJson is empty. */ + @Test + fun onJsReadyInternal_withEmptyPayload_doesNotTriggerUpdate() { + a2uiView.a2uiJson = "" + a2uiView.onJsReadyInternal() + ShadowLooper.idleMainLooper() + + val lastEvaluatedJs = shadowOf(a2uiView).lastEvaluatedJavascript + assertThat(lastEvaluatedJs).isNull() + } +} diff --git a/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/WebAppInterfaceTest.kt b/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/WebAppInterfaceTest.kt new file mode 100644 index 0000000..c24dce0 --- /dev/null +++ b/client/android/GoogleMapsA2UI/src/test/java/com/google/android/libraries/mapsplatform/a2ui/WebAppInterfaceTest.kt @@ -0,0 +1,174 @@ +// Copyright 2026 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 +// +// http://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. + +package com.google.android.libraries.mapsplatform.a2ui + +import android.app.Activity +import android.content.Context +import android.widget.FrameLayout +import com.google.common.truth.Truth.assertThat +import org.junit.After +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.Robolectric +import org.robolectric.RobolectricTestRunner +import org.robolectric.Shadows.shadowOf +import org.robolectric.shadows.ShadowLooper + +/** + * Unit tests for [WebAppInterface]. + * + * Verifies Web-to-Native event bridging, including user action callbacks (directions requests), + * JavaScript ready signaling, and WebView dynamic height adjustment. + */ +@RunWith(RobolectricTestRunner::class) +class WebAppInterfaceTest { + + private lateinit var context: Context + private lateinit var a2uiView: A2UIView + private lateinit var webAppInterface: WebAppInterface + + @Before + fun setUp() { + val activity = Robolectric.buildActivity(Activity::class.java).setup().get() + context = activity + A2UIServices.provideAPIKey("test-api-key") + a2uiView = A2UIView(context) + activity.setContentView(a2uiView) + webAppInterface = WebAppInterface(a2uiView, a2uiView) + } + + @After + fun tearDown() { + A2UIServices.provideAPIKey("") + } + + /** Verifies that sendA2uiMessages posts JavaScript evaluation and resets the resized state. */ + @Test + fun sendA2uiMessages_validPayload_postsJsEvaluationAndResetsResizeState() { + webAppInterface.resized = true + val payload = """[{"createSurface": {"surfaceId": "msg-1"}}]""" + + webAppInterface.sendA2uiMessages(payload) + ShadowLooper.idleMainLooper() + + assertThat(webAppInterface.resized).isFalse() + val lastEvaluatedJs = shadowOf(a2uiView).lastEvaluatedJavascript + assertThat(lastEvaluatedJs).isNotNull() + assertThat(lastEvaluatedJs).contains("shell.processA2uiMessages") + assertThat(lastEvaluatedJs).contains("msg-1") + } + + /** + * Verifies that when the web layer triggers onGetDirections(), the user action callback on + * [A2UIView] is notified with the exact action JSON payload. + */ + @Test + fun onGetDirections_actionPayload_surfacesActionToCallback() { + var receivedAction: String? = null + a2uiView.onUserAction = { actionJson -> receivedAction = actionJson } + + val actionPayload = """{"placeId": "ChIJN1t_tDeuEmsRUsoyG83frY4", "action": "get_directions"}""" + webAppInterface.onGetDirections(actionPayload) + + assertThat(receivedAction).isEqualTo(actionPayload) + } + + /** Verifies that onGetDirections() does not throw when no user action callback is registered. */ + @Test + fun onGetDirections_withNullActionCallback_doesNotCrash() { + a2uiView.onUserAction = null + webAppInterface.onGetDirections("""{"action": "test"}""") + } + + /** Verifies that onJsReady() triggers the JS ready lifecycle on A2UIView. */ + @Test + fun onJsReady_withQueuedPayload_triggersJsReadyOnA2uiView() { + val payload = """[{"createSurface": {"surfaceId": "test"}}]""" + a2uiView.a2uiJson = payload + + webAppInterface.onJsReady() + ShadowLooper.idleMainLooper() + + assertThat(a2uiView.a2uiJson).isEqualTo(payload) + val lastEvaluatedJs = shadowOf(a2uiView).lastEvaluatedJavascript + assertThat(lastEvaluatedJs).isNotNull() + assertThat(lastEvaluatedJs).contains("test") + } + + /** + * Verifies that onWebpageResized() updates the WebView layout parameters to match the content + * height and fires the onRenderComplete callback. + */ + @Test + fun onWebpageResized_newHeight_updatesLayoutAndTriggersRenderComplete() { + var renderCompleteInvoked = false + var completedStatus: String? = null + a2uiView.onRenderComplete = { _, status -> + renderCompleteInvoked = true + completedStatus = status + } + + a2uiView.layoutParams = FrameLayout.LayoutParams(FrameLayout.LayoutParams.MATCH_PARENT, 100) + + a2uiView.render("""[{"createSurface": {"surfaceId": "test"}}]""") + webAppInterface.onWebpageResized(350) + ShadowLooper.idleMainLooper() + + val expectedHeight = (350 * a2uiView.resources.displayMetrics.density).toInt() + assertThat(a2uiView.layoutParams.height).isEqualTo(expectedHeight) + assertThat(renderCompleteInvoked).isTrue() + assertThat(completedStatus).isEqualTo("A2UI Render Complete") + assertThat(webAppInterface.resized).isTrue() + } + + /** Verifies that subsequent onWebpageResized calls are ignored once resized is true. */ + @Test + fun onWebpageResized_whenAlreadyResized_isIdempotent() { + var renderCount = 0 + a2uiView.onRenderComplete = { _, _ -> renderCount++ } + + a2uiView.layoutParams = FrameLayout.LayoutParams(FrameLayout.LayoutParams.MATCH_PARENT, 100) + + a2uiView.render("""[{"createSurface": {"surfaceId": "test"}}]""") + webAppInterface.onWebpageResized(300) + ShadowLooper.idleMainLooper() + + assertThat(renderCount).isEqualTo(1) + + // Second resize call while resized == true should be ignored + webAppInterface.onWebpageResized(500) + ShadowLooper.idleMainLooper() + + assertThat(renderCount).isEqualTo(1) + } + + /** + * Verifies that onWebpageResized handles null layoutParams on the WebView gracefully without + * crashing. + */ + @Test + fun onWebpageResized_withNullLayoutParams_doesNotCrash() { + val unattachedView = A2UIView(context) + unattachedView.layoutParams = null + val customInterface = WebAppInterface(unattachedView, unattachedView) + customInterface.resized = false + + customInterface.onWebpageResized(300) + ShadowLooper.idleMainLooper() + + assertThat(customInterface.resized).isFalse() + } +} diff --git a/client/android/README.md b/client/android/README.md index a811721..5fef8d3 100644 --- a/client/android/README.md +++ b/client/android/README.md @@ -94,6 +94,9 @@ class GmpA2UIViewHolder(itemView: View) : RecyclerView.ViewHolder(itemView) { } ``` +### 4. Streaming vs. Non-Streaming Updates (`message/stream` & `message/send`) +In Compose or View-based apps, `A2UIView` supports both real-time SSE streaming (`message/stream`) and single-response requests (`message/send`). When `useStreaming` is enabled, incremental `.Data` chunks emitted by `A2AResponseParser.parse(chunkJson)` are applied in place via `updateA2uiJson(newJson)` without recreating the underlying `WebView` (see `MainActivity.kt` and `ChatRepository.kt` in the Android sample app). + ## Architecture Deep Dive The SDK encapsulates all complex parsing and rendering logic into four core components: diff --git a/client/android/web_build/index.html b/client/android/web_build/index.html index a3da007..e43d38f 100644 --- a/client/android/web_build/index.html +++ b/client/android/web_build/index.html @@ -21,7 +21,12 @@ A2UI Bridge - + - + diff --git a/client/ios/GoogleMapsA2UI/Tests/A2AResponseParserTests.swift b/client/ios/GoogleMapsA2UI/Tests/A2AResponseParserTests.swift index 4bc57f8..ac3f0d9 100644 --- a/client/ios/GoogleMapsA2UI/Tests/A2AResponseParserTests.swift +++ b/client/ios/GoogleMapsA2UI/Tests/A2AResponseParserTests.swift @@ -19,9 +19,8 @@ import XCTest @testable import GoogleMapsA2UI final class A2AResponseParserTests: XCTestCase { - /// Tests that an error is thrown when the input payload is not a valid JSON object. - func testParse_InvalidJSONFormat() { + func testParse_invalidJSONFormat() { let invalidPayload: [String: Any] = ["key": Date()] // Date is not valid JSON XCTAssertThrowsError(try A2AResponseParser.parse(invalidPayload)) { error in XCTAssertEqual(error as? A2AParserError, .invalidJSONFormat) @@ -29,15 +28,23 @@ final class A2AResponseParserTests: XCTestCase { } /// Tests that an error is thrown when the JSON payload lacks a recognizable `parts` structure. - func testParse_InvalidPayloadStructure() { + func testParse_invalidPayloadStructure() { let payloadWithNoParts: [String: Any] = ["status": "ok"] XCTAssertThrowsError(try A2AResponseParser.parse(payloadWithNoParts)) { error in XCTAssertEqual(error as? A2AParserError, .invalidPayloadStructure) } } + /// Tests that an error is thrown when `parts` is not an array of dictionaries. + func testParse_invalidPartsType_throwsError() { + let payloadWithInvalidParts: [String: Any] = ["parts": "not_an_array"] + XCTAssertThrowsError(try A2AResponseParser.parse(payloadWithInvalidParts)) { error in + XCTAssertEqual(error as? A2AParserError, .invalidPayloadStructure) + } + } + /// Tests that a single text part is parsed correctly into a text event. - func testParse_SimpleTextPart() throws { + func testParse_simpleTextPart() throws { let payload: [String: Any] = [ "parts": [ ["kind": "text", "text": "Show me some good sushi in Seattle"] @@ -46,16 +53,23 @@ final class A2AResponseParserTests: XCTestCase { let events = try A2AResponseParser.parse(payload) XCTAssertEqual(events.count, 1) + XCTAssertEqual(events.first?.textValue, "Show me some good sushi in Seattle") + } - guard case .text(let text) = events[0] else { - XCTFail("Expected text event") - return - } - XCTAssertEqual(text, "Show me some good sushi in Seattle") + /// Tests that parts containing empty text strings produce an empty event list. + func testParse_emptyTextPart_returnsEmptyList() throws { + let payload: [String: Any] = [ + "parts": [ + ["kind": "text", "text": ""] + ] + ] + + let events = try A2AResponseParser.parse(payload) + XCTAssertTrue(events.isEmpty) } /// Tests that multiple text parts within the `content.parts` path are parsed into separate text events. - func testParse_MultipleTextParts() throws { + func testParse_multipleTextParts() throws { let payload: [String: Any] = [ "content": [ "parts": [ @@ -66,23 +80,16 @@ final class A2AResponseParserTests: XCTestCase { ] let events = try A2AResponseParser.parse(payload) - XCTAssertEqual(events.count, 2) - - if case .text(let text1) = events[0] { - XCTAssertEqual(text1, "Show me some good sushi in Seattle") - } else { - XCTFail("Expected first event to be text") - } - - if case .text(let text2) = events[1] { - XCTAssertEqual(text2, "What are their ratings?") - } else { - XCTFail("Expected second event to be text") - } + XCTAssertEqual( + events.map(\.textValue), + [ + "Show me some good sushi in Seattle", + "What are their ratings?", + ]) } /// Tests that an A2UI JSON payload embedded inside a text part using `` tags is extracted. - func testParse_EmbeddedA2UIJSON() throws { + func testParse_embeddedA2UIJSON() throws { let textWithJSON = "Here is the Seattle map {\"createSurface\": {\"surfaceId\": \"sushi-seattle\"}} Hope you like it!" let payload: [String: Any] = [ @@ -93,28 +100,69 @@ final class A2AResponseParserTests: XCTestCase { let events = try A2AResponseParser.parse(payload) XCTAssertEqual(events.count, 3) + XCTAssertEqual(events[0].textValue, "Here is the Seattle map") + XCTAssertEqual(events[1].metadataMimeType, "application/json+a2ui") + let dict = try XCTUnwrap(events[1].dataDictionaries?.first) + XCTAssertNotNil(dict["createSurface"]) + XCTAssertEqual(events[2].textValue, "Hope you like it!") + } - guard case .text(let prefix) = events[0] else { - return XCTFail("Expected text event") - } - XCTAssertEqual(prefix, "Here is the Seattle map") + /// Tests that an A2UI JSON array embedded inside a text part is parsed into a flat array of components. + func testParse_embeddedA2UIJSONArray_flattensToSingleArray() throws { + let textWithJSONArray = + "Here is the map: [{\"createSurface\": {\"surfaceId\": \"sushi-seattle\"}}, {\"updateComponents\": {\"surfaceId\": \"sushi-seattle\"}}] Done." + let payload: [String: Any] = [ + "parts": [ + ["kind": "text", "text": textWithJSONArray] + ] + ] - guard case .data(let data, let metadata) = events[1] else { - return XCTFail("Expected data event") - } - XCTAssertEqual(metadata?.mimeType, "application/json+a2ui") - let array = data as? [Any] - let dict = array?.first as? [String: Any] - XCTAssertNotNil(dict?["createSurface"]) + let events = try A2AResponseParser.parse(payload) + XCTAssertEqual(events.count, 3) + XCTAssertEqual(events[0].textValue, "Here is the map:") + XCTAssertEqual(events[1].metadataMimeType, "application/json+a2ui") + let array = try XCTUnwrap(events[1].dataDictionaries) + XCTAssertEqual(array.count, 2) + XCTAssertNotNil(array[0]["createSurface"]) + XCTAssertNotNil(array[1]["updateComponents"]) + XCTAssertEqual(events[2].textValue, "Done.") + } - guard case .text(let suffix) = events[2] else { - return XCTFail("Expected text event") - } - XCTAssertEqual(suffix, "Hope you like it!") + /// Tests that malformed JSON inside an embedded '' tag falls back gracefully to plain text. + func testParse_malformedEmbeddedA2UITag_fallsBackToPlainText() throws { + let textWithMalformedJSON = "Intro text {not_valid_json outro text" + let payload: [String: Any] = [ + "parts": [ + ["kind": "text", "text": textWithMalformedJSON] + ] + ] + + let events = try A2AResponseParser.parse(payload) + XCTAssertEqual( + events.map(\.textValue), + [ + "Intro text", + "{not_valid_json", + "outro text", + ]) + } + + /// Tests that an unclosed '' tag remains as plain text. + func testParse_unclosedA2UITag_fallsBackToPlainText() throws { + let unclosedTagText = "Before {\"createSurface\": {}} without closing tag" + let payload: [String: Any] = [ + "parts": [ + ["kind": "text", "text": unclosedTagText] + ] + ] + + let events = try A2AResponseParser.parse(payload) + XCTAssertEqual(events.count, 1) + XCTAssertEqual(events.first?.textValue, unclosedTagText) } /// Tests that a data part with an explicit A2UI mime type is parsed and batched into an array. - func testParse_DataPartWithA2UIMimeType() throws { + func testParse_dataPartWithA2UIMimeType() throws { let payload: [String: Any] = [ "parts": [ [ @@ -133,19 +181,14 @@ final class A2AResponseParserTests: XCTestCase { let events = try A2AResponseParser.parse(payload) XCTAssertEqual(events.count, 1) + XCTAssertEqual(events.first?.metadataMimeType, "application/json+a2ui") - guard case .data(let data, let metadata) = events[0] else { - return XCTFail("Expected data event") - } - XCTAssertEqual(metadata?.mimeType, "application/json+a2ui") - - let a2uiArray = data as? [Any] - XCTAssertNotNil(a2uiArray, "A2UI payload should be batched into an array") - XCTAssertEqual(a2uiArray?.count, 1) + let a2uiArray = try XCTUnwrap(events.first?.dataArray) + XCTAssertEqual(a2uiArray.count, 1) } /// Tests that a data part is inferred as A2UI if it contains recognized keys, even without a mime type. - func testParse_DataPartWithImplicitA2UIKey() throws { + func testParse_dataPartWithImplicitA2UIKey() throws { let payload: [String: Any] = [ "parts": [ [ @@ -162,18 +205,13 @@ final class A2AResponseParserTests: XCTestCase { let events = try A2AResponseParser.parse(payload) XCTAssertEqual(events.count, 1) - - guard case .data(let data, let metadata) = events[0] else { - return XCTFail("Expected data event") - } - XCTAssertEqual(metadata?.mimeType, "application/json+a2ui") - let a2uiArray = data as? [Any] - XCTAssertNotNil(a2uiArray) - XCTAssertEqual(a2uiArray?.count, 1) + XCTAssertEqual(events.first?.metadataMimeType, "application/json+a2ui") + let a2uiArray = try XCTUnwrap(events.first?.dataArray) + XCTAssertEqual(a2uiArray.count, 1) } /// Tests that the parser can successfully locate and extract parts from the `status.message.parts` JSON path. - func testParse_StatusMessagePartsPath() throws { + func testParse_statusMessagePartsPath() throws { let payload: [String: Any] = [ "status": [ "message": [ @@ -186,16 +224,11 @@ final class A2AResponseParserTests: XCTestCase { let events = try A2AResponseParser.parse(payload) XCTAssertEqual(events.count, 1) - - guard case .text(let text) = events[0] else { - XCTFail("Expected text event") - return - } - XCTAssertEqual(text, "Seattle is home to a world-class sushi scene") + XCTAssertEqual(events.first?.textValue, "Seattle is home to a world-class sushi scene") } /// Tests that consecutive data parts identified as A2UI payloads are batched together into a single data event. - func testParse_ConsecutiveA2UIPayloadsAreBatched() throws { + func testParse_consecutiveA2UIPayloadsAreBatched() throws { let payload: [String: Any] = [ "parts": [ [ @@ -223,31 +256,16 @@ final class A2AResponseParserTests: XCTestCase { let events = try A2AResponseParser.parse(payload) XCTAssertEqual(events.count, 1) + XCTAssertEqual(events.first?.metadataMimeType, "application/json+a2ui") - guard case .data(let data, let metadata) = events[0] else { - XCTFail("Expected data event") - return - } - XCTAssertEqual(metadata?.mimeType, "application/json+a2ui") - - guard let a2uiArray = data as? [Any] else { - XCTFail("Expected data payload to be an array of batched items") - return - } + let a2uiArray = try XCTUnwrap(events.first?.dataDictionaries) XCTAssertEqual(a2uiArray.count, 2) - - guard let dict1 = a2uiArray[0] as? [String: Any], - let dict2 = a2uiArray[1] as? [String: Any] - else { - XCTFail("Expected array elements to be dictionaries") - return - } - XCTAssertNotNil(dict1["createSurface"]) - XCTAssertNotNil(dict2["updateComponents"]) + XCTAssertNotNil(a2uiArray[0]["createSurface"]) + XCTAssertNotNil(a2uiArray[1]["updateComponents"]) } /// Tests that an A2UI batch is finalized and a new one starts if interrupted by a text part. - func testParse_A2UIBatchInterruptedByTextPart() throws { + func testParse_a2uiBatchInterruptedByTextPart() throws { let payload: [String: Any] = [ "parts": [ [ @@ -267,31 +285,17 @@ final class A2AResponseParserTests: XCTestCase { let events = try A2AResponseParser.parse(payload) XCTAssertEqual(events.count, 3) - guard case .data(let data1, let metadata1) = events[0] else { - XCTFail("Expected first event to be data") - return - } - XCTAssertEqual(metadata1?.mimeType, "application/json+a2ui") - let batch1 = data1 as? [Any] - XCTAssertEqual(batch1?.count, 1) + XCTAssertEqual(events[0].metadataMimeType, "application/json+a2ui") + XCTAssertEqual(events[0].dataArray?.count, 1) - guard case .text(let text) = events[1] else { - XCTFail("Expected second event to be text") - return - } - XCTAssertEqual(text, "Middle Text explaining the surface") + XCTAssertEqual(events[1].textValue, "Middle Text explaining the surface") - guard case .data(let data2, let metadata2) = events[2] else { - XCTFail("Expected third event to be data") - return - } - XCTAssertEqual(metadata2?.mimeType, "application/json+a2ui") - let batch2 = data2 as? [Any] - XCTAssertEqual(batch2?.count, 1) + XCTAssertEqual(events[2].metadataMimeType, "application/json+a2ui") + XCTAssertEqual(events[2].dataArray?.count, 1) } /// Tests that multiple `` tags within a single text part are all extracted sequentially. - func testParse_MultipleEmbeddedA2UITags() throws { + func testParse_multipleEmbeddedA2UITags() throws { let textWithMultipleTags = "First map: {\"createSurface\": {\"surfaceId\": \"sushi\"}} Then: {\"updateComponents\": {\"surfaceId\": \"sushi\"}} Done." let payload: [String: Any] = [ @@ -303,24 +307,54 @@ final class A2AResponseParserTests: XCTestCase { let events = try A2AResponseParser.parse(payload) XCTAssertEqual(events.count, 5) - guard case .text(let t1) = events[0], - case .data(let d1, let m1) = events[1], - case .text(let t2) = events[2], - case .data(let d2, let m2) = events[3], - case .text(let t3) = events[4] - else { - XCTFail("Expected sequence: [text, data, text, data, text]") - return - } + XCTAssertEqual(events[0].textValue, "First map:") + XCTAssertEqual(events[1].metadataMimeType, "application/json+a2ui") + let dict1 = try XCTUnwrap(events[1].dataDictionaries?.first) + XCTAssertNotNil(dict1["createSurface"]) + + XCTAssertEqual(events[2].textValue, "Then:") + XCTAssertEqual(events[3].metadataMimeType, "application/json+a2ui") + let dict2 = try XCTUnwrap(events[3].dataDictionaries?.first) + XCTAssertNotNil(dict2["updateComponents"]) + + XCTAssertEqual(events[4].textValue, "Done.") + } + + /// Tests that when the parts array is present but empty or contains unknown types, + /// the parser returns an empty list safely without throwing an exception. + func testParse_emptyOrUnknownParts_returnsEmptyList() throws { + let emptyPartsPayload: [String: Any] = ["parts": []] + let emptyEvents = try A2AResponseParser.parse(emptyPartsPayload) + XCTAssertEqual(emptyEvents.count, 0) - XCTAssertEqual(t1, "First map:") - XCTAssertEqual(m1?.mimeType, "application/json+a2ui") - XCTAssertNotNil((d1 as? [Any])?.first as? [String: Any]) + let unknownPartsPayload: [String: Any] = [ + "parts": [["kind": "unsupported_media"]] + ] + let unknownEvents = try A2AResponseParser.parse(unknownPartsPayload) + XCTAssertEqual(unknownEvents.count, 0) + } +} - XCTAssertEqual(t2, "Then:") - XCTAssertEqual(m2?.mimeType, "application/json+a2ui") - XCTAssertNotNil((d2 as? [Any])?.first as? [String: Any]) +// MARK: - Test Helpers + +extension ParsedA2AEvent { + fileprivate var textValue: String? { + guard case .text(let text) = self else { return nil } + return text + } + + fileprivate var dataArray: [Any]? { + guard case .data(let data, _) = self else { return nil } + return data as? [Any] + } + + fileprivate var dataDictionaries: [[String: Any]]? { + guard case .data(let data, _) = self else { return nil } + return data as? [[String: Any]] + } - XCTAssertEqual(t3, "Done.") + fileprivate var metadataMimeType: String? { + guard case .data(_, let metadata) = self else { return nil } + return metadata?.mimeType } } diff --git a/client/ios/GoogleMapsA2UI/Tests/A2UIAttributionIdTests.swift b/client/ios/GoogleMapsA2UI/Tests/A2UIAttributionIdTests.swift index 39a6c9c..8500b1c 100644 --- a/client/ios/GoogleMapsA2UI/Tests/A2UIAttributionIdTests.swift +++ b/client/ios/GoogleMapsA2UI/Tests/A2UIAttributionIdTests.swift @@ -25,16 +25,16 @@ final class A2UIAttributionIdTests: XCTestCase { let htmlString = localContent.html XCTAssertTrue( - htmlString.contains("gmp_web_maui_v0.1.8_atoui"), - "Expected the generated attribution ID to contain the Web identifier 'gmp_web_maui_v0.1.8_atoui'" + htmlString.contains("gmp_web_maui_v0.1.9_atoui"), + "Expected the generated attribution ID to contain the Web identifier 'gmp_web_maui_v0.1.9_atoui'" ) XCTAssertTrue( - htmlString.contains("gmp_ios_maui_v0.1.8_atoui"), - "Expected the generated attribution ID to contain the iOS identifier 'gmp_ios_maui_v0.1.8_atoui'" + htmlString.contains("gmp_ios_maui_v0.1.9_atoui"), + "Expected the generated attribution ID to contain the iOS identifier 'gmp_ios_maui_v0.1.9_atoui'" ) XCTAssertTrue( - htmlString.contains("gmp_web_maui_v0.1.8_atoui,gmp_ios_maui_v0.1.8_atoui"), - "Expected the generated attribution ID to contain the combined Web and iOS identifier string 'gmp_web_maui_v0.1.8_atoui,gmp_ios_maui_v0.1.8_atoui'" + htmlString.contains("gmp_web_maui_v0.1.9_atoui,gmp_ios_maui_v0.1.9_atoui"), + "Expected the generated attribution ID to contain the combined Web and iOS identifier string 'gmp_web_maui_v0.1.9_atoui,gmp_ios_maui_v0.1.9_atoui'" ) } } diff --git a/client/ios/GoogleMapsA2UI/Tests/A2UIServicesTests.swift b/client/ios/GoogleMapsA2UI/Tests/A2UIServicesTests.swift new file mode 100644 index 0000000..18b1791 --- /dev/null +++ b/client/ios/GoogleMapsA2UI/Tests/A2UIServicesTests.swift @@ -0,0 +1,74 @@ +// +// Copyright 2026 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 +// +// http://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 XCTest + +@testable import GoogleMapsA2UI + +/// Unit tests for `A2UIServices`. +/// +/// Verifies Google Maps API key configuration and HTML template resolution and caching. +@MainActor +final class A2UIServicesTests: XCTestCase { + + override func tearDown() { + super.tearDown() + A2UIServices.provideApiKey("") + } + + /// Tests that providing an API key injects the key into the resolved HTML content template. + func testProvideApiKey_SetsKeyAndInjectsIntoHtml() throws { + A2UIServices.provideApiKey("AIzaSyTestKey123") + let content = try XCTUnwrap(A2UIServices.getLocalHTMLContent()) + XCTAssertTrue(content.html.contains("AIzaSyTestKey123")) + XCTAssertFalse(content.html.contains("$GOOGLE_MAPS_API_KEY")) + } + + /// Tests that providing an empty API key clears the placeholder without crashing. + func testProvideApiKey_SetsEmptyKey_RemovesPlaceholder() throws { + A2UIServices.provideApiKey("") + let content = try XCTUnwrap(A2UIServices.getLocalHTMLContent()) + XCTAssertFalse(content.html.contains("$GOOGLE_MAPS_API_KEY")) + } + + /// Tests that getLocalHTMLContent returns the module bundle's resourceURL as baseURL. + func testGetLocalHTMLContent_ReturnsModuleBaseURL() throws { + A2UIServices.provideApiKey("TestKey") + let content = try XCTUnwrap(A2UIServices.getLocalHTMLContent()) + XCTAssertEqual(content.baseURL, Bundle.module.resourceURL) + } + + /// Tests that the resolved HTML template is cached when the API key does not change, + /// and invalidated/recalculated when a new key is provided. + func testGetLocalHTMLContent_CachesResultWhenKeyUnchanged() throws { + A2UIServices.provideApiKey("KeyAlpha") + let first = try XCTUnwrap(A2UIServices.getLocalHTMLContent()) + let second = try XCTUnwrap(A2UIServices.getLocalHTMLContent()) + XCTAssertEqual(first.html, second.html) + + A2UIServices.provideApiKey("KeyBeta") + let third = try XCTUnwrap(A2UIServices.getLocalHTMLContent()) + XCTAssertTrue(third.html.contains("KeyBeta")) + XCTAssertFalse(third.html.contains("KeyAlpha")) + } + + /// Tests that an API key containing special characters is injected verbatim without breaking HTML template resolution. + func testProvideApiKey_SpecialCharacters_InjectsVerbatim() throws { + A2UIServices.provideApiKey("AIzaSyTest-Key_123$!@#") + let content = try XCTUnwrap(A2UIServices.getLocalHTMLContent()) + XCTAssertTrue(content.html.contains("AIzaSyTest-Key_123$!@#")) + } +} diff --git a/client/ios/GoogleMapsA2UI/Tests/A2UIViewTests.swift b/client/ios/GoogleMapsA2UI/Tests/A2UIViewTests.swift new file mode 100644 index 0000000..ab70e0d --- /dev/null +++ b/client/ios/GoogleMapsA2UI/Tests/A2UIViewTests.swift @@ -0,0 +1,522 @@ +// +// Copyright 2026 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 +// +// http://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 SwiftUI +import WebKit +import XCTest + +@testable import GoogleMapsA2UI + +/// Unit tests for `A2UIView` and its internal coordinator. +/// +/// Verifies JSON payload serialization, JavaScript injection deduplication, +/// and SwiftUI view hierarchy configuration for data vs text events. +@MainActor +final class A2UIViewTests: XCTestCase { + + /// Tests that the coordinator's injectJSON method correctly serializes complex native payloads + /// containing special characters, quotes, and newlines into the lastInjectedPayload cache. + func testCoordinator_InjectJSON_SerializesComplexPayload() throws { + var height: CGFloat = 100 + let parent = A2UIMessageRepresentableView( + webViewID: "test-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + let webView = WKWebView() + + let payload: [String: Any] = [ + "createSurface": [ + "surfaceId": "sushi-1", + "description": "Text with \"double quotes\" and 'single' and \n newline.", + ] + ] + + coordinator.injectJSON(webView, payload: payload) + + let injected = try XCTUnwrap(coordinator.lastInjectedPayload) + XCTAssertTrue(injected.contains("sushi-1")) + XCTAssertTrue(injected.contains("double quotes")) + } + + /// Tests that the coordinator deduplicates subsequent identical JSON payload injections. + func testCoordinator_InjectJSON_DeduplicatesIdenticalPayload() throws { + var height: CGFloat = 100 + let parent = A2UIMessageRepresentableView( + webViewID: "test-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + let webView = WKWebView() + + let payload = ["surfaceId": "test-dedupe"] + coordinator.injectJSON(webView, payload: payload) + let firstInjected = try XCTUnwrap(coordinator.lastInjectedPayload) + + // Second call with same payload + coordinator.injectJSON(webView, payload: payload) + XCTAssertEqual(coordinator.lastInjectedPayload, firstInjected) + } + + /// Tests that initializing `A2UIView` with an A2UI Data event produces a non-empty SwiftUI View body. + func testA2UIView_BodyStructure_WithDataEvent() { + let event = ParsedA2AEvent.data( + [["createSurface": ["surfaceId": "test"]]], + metadata: ParsedA2AEventMetadata(mimeType: "application/json+a2ui") + ) + + let view = A2UIView( + part: event, + id: "test-id", + onUserAction: { _ in } + ) + + let mirror = Mirror(reflecting: view.body) + let bodyDesc = String(describing: mirror.children.first?.value ?? "") + XCTAssertTrue(bodyDesc.contains("A2UIMessageInnerWrapper")) + } + + /// Tests that initializing `A2UIView` with a Text event safely defaults to an EmptyView. + func testA2UIView_BodyStructure_WithTextEvent() { + let event = ParsedA2AEvent.text("Just text") + let view = A2UIView( + part: event, + id: "text-id", + onUserAction: { _ in } + ) + + let mirror = Mirror(reflecting: view.body) + let bodyDesc = String(describing: mirror.children.first?.value ?? "") + XCTAssertTrue(bodyDesc.contains("EmptyView")) + } + + /// Tests that receiving an onGetDirections action via the iOS bridge triggers onUserAction callback. + func testCoordinator_UserContentController_OnGetDirections() { + var receivedAction: String? + var height: CGFloat = 100 + let parent = A2UIMessageRepresentableView( + webViewID: "test-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { action in receivedAction = action }, + onRenderComplete: nil + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + let message = MockScriptMessage( + name: "iOS", + body: ["action": "onGetDirections", "data": "{\"placeId\": \"123\"}"] + ) + + coordinator.userContentController(WKUserContentController(), didReceive: message) + + XCTAssertEqual(receivedAction, "{\"placeId\": \"123\"}") + } + + /// Tests that heightObserver updates the dynamic height and fires onRenderComplete when delta > 5. + func testCoordinator_UserContentController_HeightObserver_UpdatesHeight() { + var height: CGFloat = 100 + var renderCompletedId: String? + var renderStatus: String? + var renderLatency: Double? + let parent = A2UIMessageRepresentableView( + webViewID: "test-height-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: { id, latency, status in + renderCompletedId = id + renderLatency = latency + renderStatus = status + } + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + let message = MockScriptMessage(name: "heightObserver", body: CGFloat(250)) + + coordinator.userContentController(WKUserContentController(), didReceive: message) + + XCTAssertEqual(height, 250) + XCTAssertEqual(renderCompletedId, "test-height-view") + XCTAssertEqual(renderStatus, "success") + XCTAssertNotNil(renderLatency) + } + + /// Tests that heightObserver ignores small height changes (<= 5) to prevent layout thrashing. + func testCoordinator_UserContentController_HeightObserver_SuppressesSmallChanges() { + var height: CGFloat = 100 + var renderCalled = false + let parent = A2UIMessageRepresentableView( + webViewID: "test-suppress-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: { _, _, _ in renderCalled = true } + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + let message = MockScriptMessage(name: "heightObserver", body: CGFloat(103)) + + coordinator.userContentController(WKUserContentController(), didReceive: message) + + XCTAssertEqual(height, 100) + XCTAssertFalse(renderCalled) + } + + /// Tests that heightObserver ignores invalid/small heights (<= 50). + func testCoordinator_UserContentController_HeightObserver_IgnoresSmallHeights() { + var height: CGFloat = 100 + let parent = A2UIMessageRepresentableView( + webViewID: "test-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + let message = MockScriptMessage(name: "heightObserver", body: CGFloat(30)) + + coordinator.userContentController(WKUserContentController(), didReceive: message) + + XCTAssertEqual(height, 100) + } + + /// Tests that onJsReady sets isJSReady flag and injects payload. + func testCoordinator_UserContentController_OnJsReady_SetsReadyAndInjects() { + var height: CGFloat = 100 + let webView = WKWebView() + let parent = A2UIMessageRepresentableView( + webViewID: "test-view", + payload: [["createSurface": ["surfaceId": "ready-test"]]], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + let message = MockScriptMessage( + name: "iOS", + body: ["action": "onJsReady", "data": ""], + webView: webView + ) + + XCTAssertFalse(coordinator.isJSReady) + coordinator.userContentController(WKUserContentController(), didReceive: message) + + XCTAssertTrue(coordinator.isJSReady) + XCTAssertNotNil(coordinator.lastInjectedPayload) + } + + /// Tests that injectJSON falls back safely to "[]" when given an un-serializable payload. + func testCoordinator_InjectJSON_NonSerializablePayload_FallsBackSafely() { + var height: CGFloat = 100 + let parent = A2UIMessageRepresentableView( + webViewID: "test-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + let webView = WKWebView() + + // Pass a non-serializable object + let invalidPayload: [String: Any] = ["invalid": NSObject()] + coordinator.injectJSON(webView, payload: invalidPayload) + + XCTAssertEqual(coordinator.lastInjectedPayload, "[]") + } + + /// Tests that the coordinator handles JavaScript logging and error messages without crashing. + func testCoordinator_UserContentController_LogAndErrorMessages() { + var height: CGFloat = 100 + let parent = A2UIMessageRepresentableView( + webViewID: "test-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + + let logMessage = MockScriptMessage( + name: "iOS", + body: ["action": "log", "data": "Test log output"] + ) + coordinator.userContentController(WKUserContentController(), didReceive: logMessage) + + let errorMessage = MockScriptMessage( + name: "iOS", + body: ["action": "error", "data": "Test JS error"] + ) + coordinator.userContentController(WKUserContentController(), didReceive: errorMessage) + } + + /// Tests that malformed or unrecognized script messages are handled safely without throwing. + func testCoordinator_UserContentController_MalformedMessages_HandledGracefully() { + var height: CGFloat = 100 + let parent = A2UIMessageRepresentableView( + webViewID: "test-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + + // 1. Missing action/data keys + let missingKeysMessage = MockScriptMessage( + name: "iOS", + body: ["invalid": "value"] + ) + coordinator.userContentController(WKUserContentController(), didReceive: missingKeysMessage) + + // 2. Unknown action string + let unknownActionMessage = MockScriptMessage( + name: "iOS", + body: ["action": "unknownCustomAction", "data": "{}"] + ) + coordinator.userContentController(WKUserContentController(), didReceive: unknownActionMessage) + + // 3. Height observer with non-numeric body + let badHeightMessage = MockScriptMessage( + name: "heightObserver", + body: "not-a-number" + ) + coordinator.userContentController(WKUserContentController(), didReceive: badHeightMessage) + + XCTAssertEqual(height, 100) + } + + /// Tests that external HTTP/HTTPS link clicks trigger navigation cancellation in decidePolicyFor. + func testCoordinator_DecidePolicyForNavigationAction_ExternalLink_CancelsPolicy() { + var height: CGFloat = 100 + let parent = A2UIMessageRepresentableView( + webViewID: "test-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + let webView = WKWebView() + + let url = URL(string: "https://www.google.com")! + let action = MockNavigationAction( + navigationType: .linkActivated, + request: URLRequest(url: url) + ) + + var decidedPolicy: WKNavigationActionPolicy? + coordinator.webView(webView, decidePolicyFor: action) { policy in + decidedPolicy = policy + } + + XCTAssertEqual(decidedPolicy, .cancel) + } + + /// Tests that non-link navigation actions (e.g. other/reload) are allowed in decidePolicyFor. + func testCoordinator_DecidePolicyForNavigationAction_OtherNavigation_AllowsPolicy() { + var height: CGFloat = 100 + let parent = A2UIMessageRepresentableView( + webViewID: "test-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + let webView = WKWebView() + + let url = URL(string: "https://www.google.com")! + let action = MockNavigationAction( + navigationType: .other, + request: URLRequest(url: url) + ) + + var decidedPolicy: WKNavigationActionPolicy? + coordinator.webView(webView, decidePolicyFor: action) { policy in + decidedPolicy = policy + } + + XCTAssertEqual(decidedPolicy, .allow) + } + + /// Tests that createWebViewWith returns nil when intercepting window.open popup requests. + func testCoordinator_CreateWebViewWith_ReturnsNil() { + var height: CGFloat = 100 + let parent = A2UIMessageRepresentableView( + webViewID: "test-view", + payload: [], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = A2UIMessageRepresentableView.Coordinator(parent) + let webView = WKWebView() + let config = WKWebViewConfiguration() + + let url = URL(string: "https://maps.google.com")! + let action = MockNavigationAction( + navigationType: .linkActivated, + request: URLRequest(url: url) + ) + + let createdView = coordinator.webView( + webView, + createWebViewWith: config, + for: action, + windowFeatures: WKWindowFeatures() + ) + + XCTAssertNil(createdView) + } + + /// Tests that mounting A2UIView in a UIHostingController triggers makeCoordinator and makeUIView without crashing. + func testA2UIView_UIHostingController_MountsViewHierarchyAndInitializesWebView() { + let event = ParsedA2AEvent.data( + [["createSurface": ["surfaceId": "mount-test"]]], + metadata: ParsedA2AEventMetadata(mimeType: "application/json+a2ui") + ) + let view = A2UIView( + part: event, + id: "mount-test-id", + onUserAction: { _ in }, + onRenderComplete: { _, _, _ in } + ) + + let window = UIWindow(frame: CGRect(x: 0, y: 0, width: 375, height: 667)) + let hostingController = UIHostingController(rootView: view) + window.rootViewController = hostingController + window.makeKeyAndVisible() + hostingController.loadViewIfNeeded() + hostingController.view.layoutIfNeeded() + + XCTAssertNotNil(hostingController.view) + } + + /// Tests that mounting A2UIView with a Text event renders an EmptyView inside a UIHostingController. + func testA2UIView_UIHostingController_WithTextEvent_MountsEmptyView() { + let event = ParsedA2AEvent.text("mount text test") + let view = A2UIView( + part: event, + id: "mount-text-id", + onUserAction: { _ in } + ) + + let window = UIWindow(frame: CGRect(x: 0, y: 0, width: 375, height: 667)) + let hostingController = UIHostingController(rootView: view) + window.rootViewController = hostingController + window.makeKeyAndVisible() + hostingController.loadViewIfNeeded() + hostingController.view.layoutIfNeeded() + + XCTAssertNotNil(hostingController.view) + } + + /// Tests that A2UIMessageRepresentableView creates a configured coordinator instance. + func testA2UIMessageRepresentableView_MakeCoordinator_InitializesCorrectly() { + var height: CGFloat = 100 + let view = A2UIMessageRepresentableView( + webViewID: "coordinator-test-id", + payload: ["surfaceId": "test"], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + + let coordinator = view.makeCoordinator() + XCTAssertEqual(coordinator.parent.webViewID, "coordinator-test-id") + XCTAssertFalse(coordinator.isJSReady) + } + + /// Tests that configureWebView properly sets up WKWebView properties, delegates, and configuration. + func testA2UIMessageRepresentableView_ConfigureWebView_InitializesConfigurationAndDelegates() { + var height: CGFloat = 100 + let view = A2UIMessageRepresentableView( + webViewID: "configure-test-id", + payload: ["surfaceId": "test"], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = view.makeCoordinator() + let webView = view.configureWebView(coordinator: coordinator) + + XCTAssertFalse(webView.isOpaque) + XCTAssertEqual(webView.backgroundColor, .clear) + XCTAssertFalse(webView.scrollView.isScrollEnabled) + XCTAssertTrue(webView.navigationDelegate === coordinator) + XCTAssertTrue(webView.uiDelegate === coordinator) + } + + /// Tests that updateWebView injects JSON when JS is ready, and skips when not ready. + func testA2UIMessageRepresentableView_UpdateWebView_InjectsWhenJSReady() { + var height: CGFloat = 100 + let view = A2UIMessageRepresentableView( + webViewID: "update-test-id", + payload: ["surfaceId": "update-test"], + dynamicHeight: Binding(get: { height }, set: { height = $0 }), + onUserAction: { _ in }, + onRenderComplete: nil + ) + let coordinator = view.makeCoordinator() + let webView = WKWebView() + + // 1. When not ready, updateWebView does nothing + coordinator.isJSReady = false + view.updateWebView(webView, coordinator: coordinator) + XCTAssertNil(coordinator.lastInjectedPayload) + + // 2. When ready, updateWebView injects JSON + coordinator.isJSReady = true + view.updateWebView(webView, coordinator: coordinator) + XCTAssertNotNil(coordinator.lastInjectedPayload) + } +} + +private class MockNavigationAction: WKNavigationAction { + private let _navigationType: WKNavigationType + private let _request: URLRequest + + init(navigationType: WKNavigationType, request: URLRequest) { + self._navigationType = navigationType + self._request = request + super.init() + } + + override var navigationType: WKNavigationType { _navigationType } + override var request: URLRequest { _request } +} + +private class MockScriptMessage: WKScriptMessage { + private let _name: String + private let _body: Any + private weak var _webView: WKWebView? + + init(name: String, body: Any, webView: WKWebView? = nil) { + self._name = name + self._body = body + self._webView = webView + super.init() + } + + override var name: String { _name } + override var body: Any { _body } + override var webView: WKWebView? { _webView } +} diff --git a/client/ios/README.md b/client/ios/README.md index fd4b238..913dfaf 100644 --- a/client/ios/README.md +++ b/client/ios/README.md @@ -135,3 +135,6 @@ struct ChatMessageView: View { } } ``` + +#### 4. Streaming Updates (`message/stream`) +When consuming real-time Server-Sent Events (`message/stream`), pass a stable `id` (such as `"\(message.id)-\(index)"`) to `A2UIView`. Whenever `part` updates with incremental A2UI payload data, `A2UIView` automatically pushes the update to the mounted web component via its JavaScript bridge (`updateA2UIData`) without reloading the underlying `WKWebView`. You can switch between `message/stream` and `message/send` in your networking service using a `useStreaming` boolean flag (see `ChatService.swift` in the iOS sample app). diff --git a/client/ios/web_build/index.html b/client/ios/web_build/index.html index a3da007..e43d38f 100644 --- a/client/ios/web_build/index.html +++ b/client/ios/web_build/index.html @@ -21,7 +21,12 @@ A2UI Bridge - +