|
| 1 | +# Copyright 2025 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +from __future__ import annotations |
| 16 | + |
| 17 | +"""Minimal agent wired to the BigQuery Agent Analytics plugin.""" |
| 18 | + |
| 19 | +import os |
| 20 | +from typing import Dict |
| 21 | +from typing import List |
| 22 | + |
| 23 | +from google.adk import Agent |
| 24 | +from google.adk.apps import App |
| 25 | +from google.adk.models.google_llm import Gemini |
| 26 | +from google.adk.plugins.bigquery_agent_analytics_plugin import ( |
| 27 | + BigQueryAgentAnalyticsPlugin, |
| 28 | + BigQueryLoggerConfig, |
| 29 | +) |
| 30 | +from google.adk.tools.function_tool import FunctionTool |
| 31 | +from google.adk.tools.google_search_agent_tool import ( |
| 32 | + GoogleSearchAgentTool, |
| 33 | + create_google_search_agent, |
| 34 | +) |
| 35 | + |
| 36 | +PROJECT_ID = os.getenv("BQ_AGENT_ANALYTICS_PROJECT") |
| 37 | +DATASET_ID = os.getenv("BQ_AGENT_ANALYTICS_DATASET") |
| 38 | +TABLE_ID = os.getenv("BQ_AGENT_ANALYTICS_TABLE", "agent_events") |
| 39 | + |
| 40 | +# Default Vertex AI settings if env vars are not provided. |
| 41 | +os.environ.setdefault("VERTEXAI_PROJECT", "test-project-0728-467323") |
| 42 | +os.environ.setdefault("VERTEXAI_LOCATION", "us-central1") |
| 43 | +# google.genai expects these env vars to auto-switch to Vertex AI with ADC. |
| 44 | +os.environ.setdefault("GOOGLE_GENAI_USE_VERTEXAI", "true") |
| 45 | +os.environ.setdefault("GOOGLE_CLOUD_PROJECT", "test-project-0728-467323") |
| 46 | +os.environ.setdefault("GOOGLE_CLOUD_LOCATION", "us-central1") |
| 47 | + |
| 48 | +if not PROJECT_ID or not DATASET_ID: |
| 49 | + raise ValueError( |
| 50 | + "Set BQ_AGENT_ANALYTICS_PROJECT and BQ_AGENT_ANALYTICS_DATASET before " |
| 51 | + "running this agent to enable BigQuery analytics logging." |
| 52 | + ) |
| 53 | + |
| 54 | +analytics_plugin = BigQueryAgentAnalyticsPlugin( |
| 55 | + project_id=PROJECT_ID, |
| 56 | + dataset_id=DATASET_ID, |
| 57 | + table_id=TABLE_ID, |
| 58 | + config=BigQueryLoggerConfig( |
| 59 | + event_allowlist=[ |
| 60 | + "USER_MESSAGE_RECEIVED", |
| 61 | + "LLM_REQUEST", |
| 62 | + "LLM_RESPONSE", |
| 63 | + "TOOL_STARTING", |
| 64 | + "TOOL_COMPLETED", |
| 65 | + "MODEL_RESPONSE", |
| 66 | + "TOOL_CALL", |
| 67 | + "TOOL_RESULT", |
| 68 | + "ERROR", |
| 69 | + ], |
| 70 | + max_content_length=400, |
| 71 | + ), |
| 72 | +) |
| 73 | + |
| 74 | +def pick_city(top_n: int = 3) -> List[str]: |
| 75 | + """Return a short list of recommended cities.""" |
| 76 | + cities = ["Tokyo", "Paris", "New York", "Sydney", "Singapore", "Toronto"] |
| 77 | + return cities[:top_n] |
| 78 | + |
| 79 | +def city_highlights(city: str) -> Dict[str, str]: |
| 80 | + """Return quick highlights for a city.""" |
| 81 | + highlights = { |
| 82 | + "Tokyo": "Sushi, Akihabara tech, efficient transit.", |
| 83 | + "Paris": "Museums, pastries, walkable boulevards.", |
| 84 | + "New York": "Broadway, diverse food, skyline views.", |
| 85 | + "Sydney": "Harbour, beaches, outdoor cafes.", |
| 86 | + "Singapore": "Hawker food, gardens, clean and safe.", |
| 87 | + "Toronto": "CN Tower, neighbourhood food, waterfront.", |
| 88 | + } |
| 89 | + return {"city": city, "highlights": highlights.get(city, "Explore freely.")} |
| 90 | + |
| 91 | +def estimate_trip_budget(city: str, days: int, budget_per_day: float) -> Dict[str, str]: |
| 92 | + """Rough budget calculator.""" |
| 93 | + total = days * budget_per_day |
| 94 | + return { |
| 95 | + "city": city, |
| 96 | + "days": str(days), |
| 97 | + "budget_per_day": f"${budget_per_day:,.0f}", |
| 98 | + "estimated_total": f"${total:,.0f}", |
| 99 | + "note": "Assumes lodging+food+local transit; adjust for flights.", |
| 100 | + } |
| 101 | + |
| 102 | +city_tool = FunctionTool(pick_city) |
| 103 | +city_highlights_tool = FunctionTool(city_highlights) |
| 104 | +budget_tool = FunctionTool(estimate_trip_budget) |
| 105 | + |
| 106 | +gemini = Gemini(model="gemini-2.5-flash") |
| 107 | +search_agent = create_google_search_agent(model=gemini) |
| 108 | +google_search_tool = GoogleSearchAgentTool(agent=search_agent) |
| 109 | + |
| 110 | +root_agent = Agent( |
| 111 | + name="bq_agent_analytics_demo", |
| 112 | + model=gemini, |
| 113 | + instruction=( |
| 114 | + "You are a concise assistant. Prefer to use tools when asked for trip " |
| 115 | + "ideas, highlights, or cost estimates. Keep answers short and " |
| 116 | + "actionable." |
| 117 | + ), |
| 118 | + description="A minimal agent that logs events to BigQuery.", |
| 119 | + tools=[ |
| 120 | + city_tool, |
| 121 | + city_highlights_tool, |
| 122 | + budget_tool, |
| 123 | + google_search_tool, |
| 124 | + ], |
| 125 | +) |
| 126 | + |
| 127 | +app = App( |
| 128 | + name="bq_agent_analytics_demo", |
| 129 | + root_agent=root_agent, |
| 130 | + plugins=[analytics_plugin], |
| 131 | +) |
0 commit comments