logoassistant-ui

Message Part Grouping

Organize message parts into custom groups with flexible grouping functions.

Message parts grouped by type:

Search: climate data
Search: renewable energy stats
Analyzing trends...
Generating summary...

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

Basic Usage

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

/components/assistant-ui/thread.tsx
const :  = () => {
  return (
    <. ="...">
      < ="...">
        <.
          ={() => {
            // Your custom grouping logic here
            return [{ : , : [0, 1, 2] }];
          }}
          ={{
            : ({ , ,  }) => {
              // Your custom group rendering
              return < ="group">{}</>;
            },
          }}
        />
      </>
      < />
      < ="..." />
    </.>
  );
};

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

  1. Keep grouping logic simple: Complex grouping functions can impact performance
  2. Handle ungrouped parts: Always consider parts that don't match any group criteria
  3. Maintain order: Consider the visual flow when ordering groups
  4. Use meaningful keys: Group keys should be descriptive for debugging
  5. 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 { useAssistantState } from "@assistant-ui/react";

const DynamicGroup: FC<
  PropsWithChildren<{ groupKey: string | undefined; indices: number[] }>
> = ({ groupKey, indices, children }) => {
  const parts = useAssistantState(({ message }) => message.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>
  );
};