LlamaIndex Protocols AG UI Integration (Fiqros fork)

This is an unofficial fork, maintained by the Fiqros team, of llama-index-protocols-ag-ui by LlamaIndex.

All of the package's design and code is LlamaIndex's work. This fork only patches two bugs on top of upstream 0.5.0 (released here as 0.5.1):

1. Tools sent by the client had every argument retyped as a string.
2. RunAgentInput.context was never sent to the LLM.

Together these stop CopilotKit's A2UI generative UI from rendering with LlamaIndex agents. Nothing else is changed: the import path, public API and usage are identical to upstream. See Fork changes for details and Credits for attribution. The fork is not affiliated with or endorsed by LlamaIndex. Please use the official package once these fixes are available upstream.

pip uninstall llama-index-protocols-ag-ui
pip install fiqros-llama-index-protocols-ag-ui

Uninstall the upstream package first: both packages provide the same llama_index.protocols.ag_ui module and cannot be installed side by side. Your imports stay the same.

Fork changes

Symptom

With a CopilotKit frontend using A2UI (a2ui: { injectA2UITool: true } on the runtime) and a LlamaIndex agent served by this package, asking for UI such as "draw a sales dashboard" never renders. The chat shows a "Building interface" placeholder with a growing token count indefinitely. The agent's render_a2ui tool call looks like this:

{
  "surfaceId": "sales-dashboard",
  "components": "[{\"id\":\"root\",\"type\":\"container\",\"parentId\":null,...}]"
}

That shows two separate problems: components is a JSON-encoded string rather than an array, and the components use an invented format (type, parentId, inline style) instead of A2UI's (component, children).

Bug 1: dynamic tool arguments were all retyped as strings

Clients can send tools at request time in RunAgentInput.tools (CopilotKit's A2UI middleware adds render_a2ui this way). _ag_ui_tool_to_llama_index converted each one into a LlamaIndex FunctionTool by building a Pydantic model in which every argument was typed str. The tool's real JSON Schema ("components": {"type": "array", "items": {...}}) was discarded, and the LLM was told components was a string, so it sent one.

The A2UI middleware parses the streamed tool-call arguments looking for "components": [. Because it found "components": " instead, it never emitted the surface, and the UI stayed in its loading state. The {"status": "rendered"} tool result seen in logs is filled in automatically by the middleware when the run ends; it does not mean anything rendered.

Fix: a small ToolMetadata subclass, _AGUIToolMetadata, overrides get_parameters_dict() to return the tool's original JSON Schema, so the LLM sees the real types, item shapes, enums and nested descriptions. The Pydantic model is kept for argument names but its fields are now typed Any, so it never rejects a list or dict argument.

Details:

• The schema is deep-copied on every call. LLM integrations edit the returned dict in place (the OpenAI integration sets additionalProperties), and that must not change the schema the client sent.
• Top-level keys are filtered to the same set LlamaIndex already sends for Pydantic-generated schemas (type, properties, required, definitions, $defs).
• It falls back to the previous behaviour when a tool has no usable properties, so tools that worked before are unaffected.
• Nothing changed on the output side. Tool-call arguments were already sent to the client as json.dumps(tool_kwargs), so a real list now goes out as a real JSON array.

Bug 2: RunAgentInput.context was never sent to the LLM

AG-UI clients pass instructions the agent needs in RunAgentInput.context. CopilotKit uses it for the A2UI render-tool guide and the schema of the app's component catalog. The workflow read messages, state and tools from the request but never context, so the LLM never learned the A2UI format or the available components and made up its own.

Fix: two helpers in agent.py:

• _format_context() renders each context entry as a ## {description}\n{value} section and skips entries with an empty value.
• _with_context() returns the message list for one LLM call with that text appended to the system prompt. If the history already starts with a system message, the context is merged into a copy of it; otherwise a new system message is prepended.

The rendered context is stored in the workflow's run store (ctx.store) and added on every LLM call in the run, including the follow-up calls after backend tools execute.

Why the context is not saved in the chat history: the history is sent to the client in MESSAGES_SNAPSHOT events, and the client sends it back on the next run. Saving the context there (the way system_prompt is handled today) would add another copy of a multi-kilobyte catalog on every turn. The context only ever exists in the per-call copy of the message list. It is merged into a single system message rather than added as a second one because several providers accept only one system prompt.

Compatibility and limitations

• Strict tool mode must stay off (the default for the OpenAI integration). Client schemas such as A2UI's use open objects ("items": {"type": "object"}), which OpenAI's strict mode rejects.
• Supported versions are unchanged: Python 3.10–3.13 and llama-index-core>=0.14.1,<0.15. ToolMetadata.get_parameters_dict(), which the fix overrides, has the same shape across that range.
• A related bug is not fixed here: system_prompt is still written into the chat history, so it comes back from the client and is appended again on each turn. The same approach as _with_context() would fix it.

Testing

tests/test_dynamic_tools_and_context.py adds 11 regression tests using MockFunctionCallingLLM:

• Tool schema: components reaches the LLM as an array with its items; the client's schema is not mutated; tools without properties fall back correctly; list arguments are accepted; and the streamed tool-call arguments parse to a real JSON array.
• Context: the formatting and merge helpers behave as described; the context reaches the LLM inside a single system message, after system_prompt; and it appears in neither the message snapshots nor the stored chat history.

All 39 tests in the package pass (the 28 existing plus 11 new). The changed files pass ruff and ruff format (the versions in the repository's pre-commit config) and mypy --disallow-untyped-defs.

The tool-schema fix was also checked against the real OpenAI integration without calling the API: render_a2ui is sent with components typed {"type": "array", "items": {"type": "object"}}, and the client's schema is left unchanged.

Upstream status

These changes have not been submitted to run-llama/llama_index. This fork exists to test them in a real CopilotKit + LlamaIndex app first.

Credits

• LlamaIndex created and maintains llama-index-protocols-ag-ui: the AG-UI router, the AGUIChatWorkflow agent, the message and event conversion, and the test suite this fork builds on. The original author is Logan Markewich, and the package is part of the LlamaIndex repository.
• The Fiqros team wrote only the patches described in Fork changes: _AGUIToolMetadata, _format_context, _with_context, the related changes in AGUIChatWorkflow.chat, and tests/test_dynamic_tools_and_context.py.

This fork is distributed under the same MIT License as the original (copyright Jerry Liu), included unchanged in the LICENSE file.

The llama-index-protocols-ag-ui package provides a factory function for creating a FastAPI router that communicates using the AG UI Protocol.

Using this package, you can quickly create a FastAPI app that can be used to communicate with AG-UI compatible frameworks like CopilotKit.

Usage

The get_ag_ui_workflow_router function is a factory function that creates a FastAPI router that can be used to communicate with AG-UI compatible frameworks like CopilotKit.

The router is configured with the following parameters:

• llm: The LLM to use for the agent.
• frontend_tools: Tools that are available to execute on the frontend.
• backend_tools: Tools that are available to execute on the backend.
• system_prompt: The system prompt to use for the agent.
• initial_state: The initial state to use for the agent. Typically the state is then interacted with by the frontend.

import uvicorn
from fastapi import FastAPI

from llama_index.llms.openai import OpenAI
from llama_index.protocols.ag_ui.server import get_ag_ui_workflow_router
from typing import Annotated


# This tool has a client-side version that is actually called to change the background
def change_background(
    background: Annotated[str, "The background. Prefer gradients."],
) -> str:
    """Change the background color of the chat. Can be anything that the CSS background attribute accepts. Regular colors, linear of radial gradients etc."""
    return f"Changing background to {background}"


agentic_chat_router = get_ag_ui_workflow_router(
    llm=OpenAI(model="gpt-4.1"),
    frontend_tools=[change_background],
    backend_tools=[],
    system_prompt="You are a helpful assistant that can change the background color of the chat.",
    initial_state=None,  # Unused in this example
)


app = FastAPI(title="AG-UI Llama-Index Endpoint")

app.include_router(agentic_chat_router, prefix="/agentic_chat")


if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=9000)

Then on the frontend, you might have setup a CopilotKit app like this:

"use client";
import React, { useState } from "react";
import "@copilotkit/react-ui/styles.css";
import "./style.css";
import { useCopilotAction } from "@copilotkit/react-core";
import { CopilotChat } from "@copilotkit/react-ui";

interface AgenticChatProps {
  params: Promise<{
    integrationId: string;
  }>;
}

const Chat = () => {
  const [background, setBackground] = useState<string>("--copilot-kit-background-color");

  useCopilotAction({
    name: "change_background",
    description:
      "Change the background color of the chat. Can be anything that the CSS background attribute accepts. Regular colors, linear of radial gradients etc.",
    parameters: [
      {
        name: "background",
        type: "string",
        description: "The background. Prefer gradients.",
      },
    ],
    handler: ({ background }) => {
      setBackground(background);
    },
  });

  return (
    <div className="flex justify-center items-center h-full w-full" style={{ background }}>
      <div className="w-8/10 h-8/10 rounded-lg">
        <CopilotChat
          className="h-full rounded-2xl"
          labels={{ initial: "Hi, I'm an agent. Want to chat?" }}
        />
      </div>
    </div>
  );
};

Check out the CopilotKit Documentation for more details on using AG-UI with CopilotKit+LlamaIndex.