Components as Tools

"""Agent backing the Tool-Based Generative UI demo.

The frontend registers `render_bar_chart` and `render_pie_chart` tools via
`useComponent`. The ADKAgent middleware injects those tools into the model
request at runtime so the agent can call them.
"""

from __future__ import annotations

from google.adk.agents import LlmAgent
from ag_ui_adk import AGUIToolset

from agents.shared_chat import get_model, stop_on_terminal_text

_INSTRUCTION = (
    "You are a data visualization assistant.\n\n"
    "When the user asks for a chart, call `render_bar_chart` or "
    "`render_pie_chart` with a concise title, short description, and a "
    "`data` array of `{label, value}` items. Pick bar for comparisons over "
    "a small set of categories; pick pie for composition / share-of-whole.\n\n"
    "If the user names a chart subject but does NOT supply concrete numbers "
    '(e.g. "show me a pie chart of website traffic by source"), do NOT '
    "ask them for data. Invent plausible illustrative sample values "
    "yourself, call the appropriate `render_*` tool immediately, and "
    "briefly note in the follow-up that the values are illustrative "
    "samples. Always render the chart on the first turn -- never reply "
    "with a clarifying question asking for the data.\n\n"
    "Keep chat responses brief -- let the chart do the talking."
)

gen_ui_tool_based_agent = LlmAgent(
    name="GenUiToolBasedAgent",
    model=get_model(),
    instruction=_INSTRUCTION,
    tools=[AGUIToolset()],
    after_model_callback=stop_on_terminal_text,
)

What is this?#

Tool-based Generative UI is the simplest form of Generative UI: you register a React component with useComponent, and CopilotKit exposes it to the agent as a tool. When the agent calls the tool, CopilotKit renders your component inline in the chat, passing the tool's arguments straight through as typed props.

Unlike tool rendering, which wraps a real backend tool in a custom UI, tool-based GenUI is the component. There is no handler, no user interaction, no server-side execution. The agent decides when to show it, populates the data, and CopilotKit paints it.

When should I use this?#

Use useComponent when you want to:

• Display rich UI (cards, charts, tables, dashboards) inline in the chat
• Show structured data the agent has derived from its reasoning
• Render previews, status indicators, or visual summaries
• Let the agent present information beyond plain text

For components that need user interaction, see Human-in-the-loop. For operational transparency around a real backend tool, see Tool rendering.

How it works in code#

pip install ag-ui-adk

from ag_ui_adk import AGUIToolset
from google.adk.agents import LlmAgent
from google.adk.agents.callback_context import CallbackContext
from google.adk.models.llm_response import LlmResponse


def stop_on_terminal_text(
    callback_context: CallbackContext, llm_response: LlmResponse
) -> None:
    content = llm_response.content
    if llm_response.partial or not content or content.role != "model":
        return
    finish_reason = llm_response.finish_reason
    if getattr(finish_reason, "name", finish_reason) != "STOP":
        return
    parts = content.parts or []
    if not any(part.text for part in parts) or any(part.function_call for part in parts):
        return
    # ADK's invocation context is private; tolerate SDK changes.
    invocation = getattr(callback_context, "_invocation_context", None)
    if invocation is not None:
        try:
            invocation.end_invocation = True
        except AttributeError:
            pass


agent = LlmAgent(
    name="assistant",
    model="gemini-3.1-flash-lite",
    instruction="Help the user and call the available frontend tools when appropriate.",
    tools=[AGUIToolset()],
    after_model_callback=stop_on_terminal_text,
)

Import the React hook and Zod in the component that registers the tool. This also applies to the built-in agent, which needs no backend tool-registration step.

import { useComponent } from "@copilotkit/react-core/v2";
import { z } from "zod";

useComponent takes a name, a Zod schema for its props, and the component to render. The runtime registers it as a frontend tool so the agent can discover it, and the schema becomes that tool's parameter definition — it is what tells the model which arguments to send.

  useComponent({    name: "render_bar_chart",    description: "Display a bar chart with labeled numeric values.",    parameters: barChartPropsSchema,    render: BarChart,  });

The component itself is ordinary React: it reads only its props and can stream in as the agent fills the payload. The example above uses Recharts for the bar chart; it doesn't know anything about CopilotKit.

Rendering in a headless chat#

CopilotKit's built-in chat components paint registered components for you. A headless or custom chat renders the message list itself, so nothing paints a tool call unless you render it — the component is registered and the agent calls it, but the chat stays empty.

Render the tool calls on each assistant message with CopilotChatToolCallsView:

import { CopilotChatToolCallsView } from "@copilotkit/react-core/v2";

<CopilotChatToolCallsView message={assistantMessage} messages={allMessages} />;

It looks up the sibling tool-role message for each tool call and hands both to the registered renderer. For finer placement, call useRenderToolCall() and paint each tool call yourself — see Headless UI.