Display components

"""LlamaIndex agent for the Tool-Based Generative UI demo.The frontend registers `render_bar_chart` and `render_pie_chart` tools via`useComponent`. The AG-UI protocol forwards those tool definitions to theagent at request time, so the backend agent itself declares no bespoketools — the LLM sees the frontend tools through the AG-UI request payloadand picks one to call when the user asks for a chart.Mirrors `langgraph-python/src/agents/gen_ui_tool_based.py`."""from __future__ import annotationsimport osfrom typing import Any, Sequencefrom llama_index.core.llms import ChatMessage, MessageRolefrom llama_index.core.base.llms.types import ChatResponseAsyncGenfrom llama_index.llms.openai import OpenAIfrom llama_index.protocols.ag_ui.router import get_ag_ui_workflow_routerfrom agents.hitl_in_chat_agent import _fix_tool_messagesdef _normalize_chart_messages(messages: Sequence[ChatMessage]) -> list[ChatMessage]:    """Restore provider tool history without changing workflow/UI snapshots."""    normalized = [message.model_copy(deep=True) for message in messages]    _fix_tool_messages(normalized)    for message in normalized:        if message.role != MessageRole.ASSISTANT:            continue        calls = message.additional_kwargs.get("ag_ui_tool_calls")        if calls is None:            continue        if not isinstance(calls, list) or any(            not isinstance(call, dict)            or not isinstance(call.get("id"), str)            or not call["id"]            or not isinstance(call.get("name"), str)            or not call["name"]            or not isinstance(call.get("arguments"), str)            for call in calls        ):            raise ValueError("Invalid AG-UI chart tool-call metadata")        if not calls:            continue        structured = [            {                "id": call["id"],                "type": "function",                "function": {"name": call["name"], "arguments": call["arguments"]},            }            for call in calls        ]        existing = message.additional_kwargs.get("tool_calls")        if existing is not None and existing != structured:            raise ValueError("Conflicting AG-UI chart tool-call metadata")        message.additional_kwargs["tool_calls"] = structured        message.additional_kwargs.pop("ag_ui_tool_calls", None)        message.additional_kwargs.pop("id", None)        # Remove only the suffix generated by the installed AG-UI converter.        # Genuine narration (including other XML-like text) stays intact.        suffix = "\n".join(            f"<tool_call><name>{call['name']}</name>"            f"<arguments>{call['arguments']}</arguments></tool_call>"            for call in calls        )        content = message.content or ""        if content == suffix:            message.content = ""        elif content.endswith("\n\n" + suffix):            message.content = content[: -(len(suffix) + 2)]    return normalizedclass _ChartOpenAI(OpenAI):    async def astream_chat(        self, messages: Sequence[ChatMessage], **kwargs: Any    ) -> ChatResponseAsyncGen:        return await super().astream_chat(_normalize_chart_messages(messages), **kwargs)SYSTEM_PROMPT = """You are a data visualization and creative assistant.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 ofcategories; pick pie for composition / share-of-whole.When the user asks for a haiku, call `generate_haiku` with the Japanesetext, English translation, an image name, and a gradient color.Keep chat responses brief -- let the visual output do the talking."""_openai_kwargs = {}if os.environ.get("OPENAI_BASE_URL"):    _openai_kwargs["api_base"] = os.environ["OPENAI_BASE_URL"]gen_ui_tool_based_router = get_ag_ui_workflow_router(    llm=_ChartOpenAI(model="gpt-5-mini", **_openai_kwargs),    frontend_tools=[],    backend_tools=[],    system_prompt=SYSTEM_PROMPT,    initial_state={},)