Message Part Grouping

Overview

This feature is experimental and the API may change in future versions.

The Message Part Grouping feature allows you to organize and display message parts in custom groups using a flexible grouping function. This is useful for creating visual hierarchies, organizing related content together, or building custom UI patterns based on message part characteristics.

Basic Usage

Use the MessagePrimitive.Unstable_PartsGrouped component with a custom grouping function:

/components/assistant-ui/thread.tsx

const AssistantMessage: FC = () => {
  return (
    <MessagePrimitive.Root className="...">
      <div className="...">
        <MessagePrimitive.Unstable_PartsGrouped
          groupingFunction={(parts) => {
            // Your custom grouping logic here
            return [{ groupKey: undefined, indices: [0, 1, 2] }];
          }}
          components={{
            Group: ({ groupKey, indices, children }) => {
              // Your custom group rendering
              return <div className="group">{children}</div>;
            },
          }}
        />
      </div>
      <AssistantActionBar />
      <BranchPicker className="..." />
    </MessagePrimitive.Root>
  );
};

How Grouping Works

The grouping function receives all message parts and returns an array of groups. Each group contains:

groupKey: A string identifier for the group (or undefined for ungrouped parts)
indices: An array of indices indicating which parts belong to this group

Use Cases & Examples

Group by Parent ID

Group related content together using a parent-child relationship:

import { FC, PropsWithChildren, useState } from "react";
import { ChevronDownIcon, ChevronUpIcon } from "lucide-react";

const groupByParentId = (parts: readonly any[]) => {
  const groups = new Map<string, number[]>();

  for (let i = 0; i < parts.length; i++) {
    const part = parts[i];
    const groupId = part?.parentId ?? `__ungrouped_${i}`;

    const indices = groups.get(groupId) ?? [];
    indices.push(i);
    groups.set(groupId, indices);
  }

  return Array.from(groups.entries()).map(([groupId, indices]) => ({
    groupKey: groupId.startsWith("__ungrouped_") ? undefined : groupId,
    indices,
  }));
};

// Usage with collapsible groups
const CollapsibleGroup: FC<
  PropsWithChildren<{ groupKey: string | undefined; indices: number[] }>
> = ({ groupKey, indices, children }) => {
  const [isCollapsed, setIsCollapsed] = useState(false);

  if (!groupKey) return <>{children}</>;

  return (
    <div className="my-2 overflow-hidden rounded-lg border">
      <button
        onClick={() => setIsCollapsed(!isCollapsed)}
        className="hover:bg-muted/50 flex w-full items-center justify-between p-3"
      >
        <span>
          Group {groupKey} ({indices.length} items)
        </span>
        {isCollapsed ? <ChevronDownIcon /> : <ChevronUpIcon />}
      </button>
      {!isCollapsed && <div className="border-t p-3">{children}</div>}
    </div>
  );
};

Group by Tool Name

Organize tool calls by their type:

import { FC, PropsWithChildren } from "react";

const groupByToolName = (parts: readonly any[]) => {
  const groups = new Map<string, number[]>();

  for (let i = 0; i < parts.length; i++) {
    const part = parts[i];
    // Group tool calls by tool name, everything else ungrouped
    const groupKey = part.type === "tool-call" ? part.toolName : `__other_${i}`;

    const indices = groups.get(groupKey) ?? [];
    indices.push(i);
    groups.set(groupKey, indices);
  }

  return Array.from(groups.entries()).map(([groupKey, indices]) => ({
    groupKey: groupKey.startsWith("__other_") ? undefined : groupKey,
    indices,
  }));
};

// Render tool groups with custom styling
const ToolGroup: FC<
  PropsWithChildren<{ groupKey: string | undefined; indices: number[] }>
> = ({ groupKey, indices, children }) => {
  if (!groupKey) return <>{children}</>;

  return (
    <div className="tool-group my-2 rounded-lg border">
      <div className="bg-muted/50 px-4 py-2 text-sm font-medium">
        Tool: {groupKey} ({indices.length} calls)
      </div>
      <div className="p-4">{children}</div>
    </div>
  );
};

Group Consecutive Text Parts

Combine multiple text parts into cohesive blocks:

import { FC, PropsWithChildren } from "react";

type MessagePartGroup = {
  groupKey: string | undefined;
  indices: number[];
};

const groupConsecutiveText = (parts: readonly any[]) => {
  const groups: MessagePartGroup[] = [];
  let currentGroup: number[] = [];
  let isTextGroup = false;

  for (let i = 0; i < parts.length; i++) {
    const isText = parts[i].type === "text";

    if (isText === isTextGroup && currentGroup.length > 0) {
      currentGroup.push(i);
    } else {
      if (currentGroup.length > 0) {
        groups.push({
          groupKey: isTextGroup ? "text-block" : undefined,
          indices: currentGroup,
        });
      }
      currentGroup = [i];
      isTextGroup = isText;
    }
  }

  if (currentGroup.length > 0) {
    groups.push({
      groupKey: isTextGroup ? "text-block" : undefined,
      indices: currentGroup,
    });
  }

  return groups;
};

// Render text blocks with special formatting
const TextBlockGroup: FC<
  PropsWithChildren<{ groupKey: string | undefined; indices: number[] }>
> = ({ groupKey, indices, children }) => {
  if (groupKey === "text-block") {
    return (
      <div className="prose prose-sm my-2 rounded-lg bg-gray-50 p-4">
        {children}
      </div>
    );
  }
  return <>{children}</>;
};

Group by Content Type

Separate different types of content for distinct visual treatment:

type MessagePartGroup = {
  groupKey: string | undefined;
  indices: number[];
};

const groupByContentType = (parts: readonly any[]) => {
  const typeGroups = new Map<string, number[]>();

  for (let i = 0; i < parts.length; i++) {
    const part = parts[i];
    const type = part.type;

    const indices = typeGroups.get(type) ?? [];
    indices.push(i);
    typeGroups.set(type, indices);
  }

  // Order groups by type priority
  const typeOrder = [
    "reasoning",
    "tool-call",
    "source",
    "text",
    "image",
    "file",
  ];
  const orderedGroups: MessagePartGroup[] = [];

  for (const type of typeOrder) {
    const indices = typeGroups.get(type);
    if (indices && indices.length > 0) {
      orderedGroups.push({ groupKey: type, indices });
    }
  }

  // Add any remaining types
  for (const [type, indices] of typeGroups) {
    if (!typeOrder.includes(type)) {
      orderedGroups.push({ groupKey: type, indices });
    }
  }

  return orderedGroups;
};

Group by Custom Metadata

Use any custom metadata in your parts for grouping:

import { FC, PropsWithChildren } from "react";

type MessagePartGroup = {
  groupKey: string | undefined;
  indices: number[];
};

const groupByPriority = (parts: readonly any[]) => {
  const priorityGroups = new Map<string, number[]>();

  for (let i = 0; i < parts.length; i++) {
    const part = parts[i];
    // Assume parts have a custom priority field
    const priority = part.priority || "normal";

    const indices = priorityGroups.get(priority) ?? [];
    indices.push(i);
    priorityGroups.set(priority, indices);
  }

  // Order by priority
  const priorityOrder = ["high", "normal", "low"];
  const orderedGroups: MessagePartGroup[] = [];

  for (const priority of priorityOrder) {
    const indices = priorityGroups.get(priority);
    if (indices) {
      orderedGroups.push({ groupKey: priority, indices });
    }
  }

  return orderedGroups;
};

// Render with priority indicators
const PriorityGroup: FC<
  PropsWithChildren<{ groupKey: string | undefined; indices: number[] }>
> = ({ groupKey, indices, children }) => {
  if (!groupKey) return <>{children}</>;

  const priorityStyles = {
    high: "border-red-500 bg-red-50",
    normal: "border-gray-300 bg-white",
    low: "border-gray-200 bg-gray-50",
  };

  return (
    <div
      className={`my-2 rounded-lg border-2 p-4 ${priorityStyles[groupKey] || ""}`}
    >
      <div className="mb-2 text-xs font-semibold uppercase text-gray-600">
        {groupKey} Priority
      </div>
      {children}
    </div>
  );
};

Integration with Assistant Streams

When using assistant-stream libraries, you can add custom metadata to parts:

Python (assistant-stream)

from assistant_stream import create_run

async def my_run(controller):
    # Add parts with custom parentId
    research_controller = controller.with_parent_id("research-123")

    await research_controller.add_tool_call("search", {"query": "climate data"})
    research_controller.append_text("Key findings from the research:")

    # Add parts with custom metadata
    controller.append_part({
        "type": "text",
        "text": "High priority finding",
        "priority": "high",
        "category": "findings"
    })

TypeScript (assistant-stream)

import { createAssistantStream } from "@assistant-ui/react/assistant-stream";

const stream = createAssistantStream(async (controller) => {
  // Add parts with parentId
  const researchController = controller.withParentId("research-123");

  await researchController.addToolCallPart({
    toolName: "search",
    args: { query: "climate data" },
  });

  // Add parts with custom metadata
  controller.appendPart({
    type: "text",
    text: "High priority finding",
    priority: "high",
    category: "findings",
  });
});

API Reference

MessagePrimitive.Unstable_PartsGrouped

MessagePrimitiveUnstable_PartsGroupedProps

groupingFunction:

(parts: readonly any[]) => MessagePartGroup[]

Function that takes an array of message parts and returns an array of groups. Each group contains a groupKey (for identification) and an array of indices.

components?:

object

Component configuration for rendering different types of message content and groups.

Components

Empty?:

EmptyMessagePartComponent

Component for rendering empty messages

Text?:

TextMessagePartComponent

Component for rendering text content

Reasoning?:

ReasoningMessagePartComponent

Component for rendering reasoning content (typically hidden)

Source?:

SourceMessagePartComponent

Component for rendering source content

Image?:

ImageMessagePartComponent

Component for rendering image content

File?:

FileMessagePartComponent

Component for rendering file content

Unstable_Audio?:

Unstable_AudioMessagePartComponent

Component for rendering audio content (experimental)

tools?:

object | { Override: ComponentType }

Configuration for tool call rendering. Can be an object with by_name map and Fallback component, or an Override component.

Group?:

ComponentType<PropsWithChildren<{ groupKey: string | undefined; indices: number[] }>>

Component for rendering grouped message parts. Receives groupKey, indices array, and children to render.

MessagePartGroup Type

type MessagePartGroup = {
  groupKey: string | undefined; // The group identifier (undefined for ungrouped parts)
  indices: number[]; // Array of part indices belonging to this group
};

Group Component Props

The Group component receives:

groupKey: The group identifier (or undefined for ungrouped parts)
indices: Array of indices for the parts in this group
children: The rendered message part components

Best Practices

Keep grouping logic simple: Complex grouping functions can impact performance
Handle ungrouped parts: Always consider parts that don't match any group criteria
Maintain order: Consider the visual flow when ordering groups
Use meaningful keys: Group keys should be descriptive for debugging
Test edge cases: Empty messages, single parts, and large numbers of parts

Common Patterns

Conditional Grouping

Only group when certain conditions are met:

const conditionalGrouping = (parts: readonly any[]) => {
  // Only group if there are enough parts
  if (parts.length < 5) {
    return [{ groupKey: undefined, indices: parts.map((_, i) => i) }];
  }

  // Your grouping logic here
};

Nested Grouping

Create hierarchical groups:

const nestedGrouping = (parts: readonly any[]) => {
  // First level: group by type
  // Second level: group by subtype or metadata
  // Implementation depends on your specific needs
};

Dynamic Group Rendering

Adjust group appearance based on content:

import { FC, PropsWithChildren } from "react";
import { useMessage } from "@assistant-ui/react";

const DynamicGroup: FC<
  PropsWithChildren<{ groupKey: string | undefined; indices: number[] }>
> = ({ groupKey, indices, children }) => {
  const parts = useMessage((m) => m.content);
  const groupParts = indices.map((i) => parts[i]);

  // Analyze group content
  const hasErrors = groupParts.some((p) => p.error);
  const isLoading = groupParts.some((p) => p.status?.type === "running");

  if (!groupKey) return <>{children}</>;

  return (
    <div
      className={`my-2 rounded-lg border p-4 ${hasErrors ? "border-red-500 bg-red-50" : ""} ${isLoading ? "animate-pulse" : ""} `}
    >
      {children}
    </div>
  );
};