|
12 | 12 | # See the License for the specific language governing permissions and |
13 | 13 | # limitations under the License. |
14 | 14 |
|
| 15 | +from __future__ import annotations |
| 16 | + |
| 17 | +import asyncio |
15 | 18 | import os |
| 19 | +from typing import Any |
16 | 20 |
|
17 | 21 | from google.adk.agents import Agent |
18 | 22 | from google.adk.auth.auth_credential import AuthCredentialTypes |
| 23 | +from google.adk.tools import load_artifacts |
19 | 24 | from google.adk.tools.data_agent.config import DataAgentToolConfig |
20 | 25 | from google.adk.tools.data_agent.credentials import DataAgentCredentialsConfig |
21 | 26 | from google.adk.tools.data_agent.data_agent_toolset import DataAgentToolset |
| 27 | +from google.adk.tools.tool_context import ToolContext |
22 | 28 | import google.auth |
23 | 29 | import google.auth.transport.requests |
| 30 | +from google.genai import types |
24 | 31 |
|
25 | 32 | # Define the desired credential type. |
26 | 33 | # By default use Application Default Credentials (ADC) from the local |
|
72 | 79 | ], |
73 | 80 | ) |
74 | 81 |
|
| 82 | +# NOTE: The generate_chart tool requires 'altair' and 'vl-convert-python' to be |
| 83 | +# installed in your environment. You can install them using: |
| 84 | +# pip install altair vl-convert-python |
| 85 | +async def generate_chart( |
| 86 | + chart_spec: dict[str, Any], tool_context: ToolContext |
| 87 | +) -> dict[str, str]: |
| 88 | + """Generates a professional chart using Altair based on a Vega-Lite spec. |
| 89 | +
|
| 90 | + Args: |
| 91 | + chart_spec: A dictionary defining a Vega-Lite chart. |
| 92 | + tool_context: The tool context. |
| 93 | +
|
| 94 | + Returns: |
| 95 | + A dictionary containing the status of the chart generation ("success" or |
| 96 | + "error"), a detail message, and the filename if successful. |
| 97 | + """ |
| 98 | + import altair as alt |
| 99 | + import vl_convert as vlc |
| 100 | + |
| 101 | + try: |
| 102 | + # Altair can take a Vega-Lite dict directly and render it. |
| 103 | + # We use vl-convert to transform the spec into a high-quality PNG. |
| 104 | + png_data = await asyncio.to_thread(vlc.vegalite_to_png, chart_spec, scale=2) |
| 105 | + |
| 106 | + # Save as artifact |
| 107 | + await tool_context.save_artifact( |
| 108 | + "chart.png", |
| 109 | + types.Part.from_bytes(data=png_data, mime_type="image/png"), |
| 110 | + ) |
| 111 | + title = chart_spec.get("title", "Chart") |
| 112 | + return { |
| 113 | + "status": "success", |
| 114 | + "detail": ( |
| 115 | + f"Professional chart '{title}' rendered using Altair/Vega-Lite." |
| 116 | + ), |
| 117 | + "filename": "chart.png", |
| 118 | + } |
| 119 | + except Exception as e: # pylint: disable=broad-exception-caught |
| 120 | + return {"status": "error", "detail": f"Failed to render chart: {str(e)}"} |
| 121 | + |
| 122 | + |
75 | 123 | root_agent = Agent( |
76 | 124 | name="data_agent", |
77 | | - description="Agent to answer user questions using Data Agents.", |
| 125 | + description=( |
| 126 | + "Agent to answer user questions using Data Agents and generate charts." |
| 127 | + ), |
78 | 128 | instruction=( |
79 | 129 | "## Persona\nYou are a helpful assistant that uses Data Agents" |
80 | 130 | " to answer user questions about their data.\n\n## Tools\n- You can" |
81 | 131 | " list available data agents using `list_accessible_data_agents`.\n-" |
82 | 132 | " You can get information about a specific data agent using" |
83 | 133 | " `get_data_agent_info`.\n- You can chat with a specific data" |
84 | | - " agent using `ask_data_agent`.\n" |
| 134 | + " agent using `ask_data_agent`.\n- `generate_chart` renders" |
| 135 | + " professional charts from a `chart_spec` (Vega-Lite JSON). Use this" |
| 136 | + " whenever you need to visualize data; do not show raw JSON to the" |
| 137 | + " user.\n- You can load artifacts using `load_artifacts`.\n" |
85 | 138 | ), |
86 | | - tools=[da_toolset], |
| 139 | + tools=[da_toolset, generate_chart, load_artifacts], |
87 | 140 | ) |
0 commit comments