Dosu Logo
Home
Documents
part-grouping
part-grouping
Type
External
Status
Published
Created
Mar 17, 2026
Updated
Mar 17, 2026

import { PartGroupingSample } from "@/components/docs/samples/part-grouping";

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:

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

const AssistantActionBar: FC = () => null;
const BranchPicker: FC<{ className?: string }> = () => null;

// ---cut---
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")

    tool = await research_controller.add_tool_call("search")
    tool.append_args_text('{"query": "climate data"}')
    tool.set_response("climate data results")
    research_controller.append_text("Key findings from the research:")

    # Add text with a different parent_id
    controller.append_text("High priority finding")

TypeScript (assistant-stream)#

import { createAssistantStream } from "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#

<ParametersTable
type="MessagePrimitiveUnstable_PartsGroupedProps"
parameters={[
{
name: "groupingFunction",
type: "(parts: readonly any[]) => MessagePartGroup[]",
description:
"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.",
required: true,
},
{
name: "components",
type: "object",
description:
"Component configuration for rendering different types of message content and groups.",
children: [
{
type: "Components",
parameters: [
{
name: "Empty",
type: "EmptyMessagePartComponent",
description: "Component for rendering empty messages",
},
{
name: "Text",
type: "TextMessagePartComponent",
description: "Component for rendering text content",
},
{
name: "Reasoning",
type: "ReasoningMessagePartComponent",
description:
"Component for rendering reasoning content (typically hidden)",
},
{
name: "Source",
type: "SourceMessagePartComponent",
description: "Component for rendering source content",
},
{
name: "Image",
type: "ImageMessagePartComponent",
description: "Component for rendering image content",
},
{
name: "File",
type: "FileMessagePartComponent",
description: "Component for rendering file content",
},
{
name: "Unstable_Audio",
type: "Unstable_AudioMessagePartComponent",
description:
"Component for rendering audio content (experimental)",
},
{
name: "tools",
type: "object | { Override: ComponentType }",
description:
"Configuration for tool call rendering. Can be an object with by_name map and Fallback component, or an Override component.",
},
{
name: "Group",
type: "ComponentType<PropsWithChildren<{ groupKey: string | undefined; indices: number[] }>>",
description:
"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 { useAuiState } from "@assistant-ui/react";

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