Add useGraphCha

.changeset/react-integration.md

@@ -0,0 +1,11 @@
+---
+"ai-sdk-graph": minor
+---
+
+Add React integration with `useGraphChat` hook
+
+- New `ai-sdk-graph/react` export with `useGraphChat` hook that wraps `@ai-sdk/react`'s `useChat`
+- Automatically handles graph-specific data parts: state changes, node start/end, and suspense events
+- Exposes `state` and `activeNodes` from the hook for tracking graph execution
+- Added utility functions `isGraphDataPart` and `stripGraphDataParts` to filter graph data parts from messages
+- Added optional peer dependencies for `react` and `@ai-sdk/react`

package.json

@@ -11,28 +11,46 @@
       "types": "./dist/index.d.ts",
       "import": "./dist/index.js",
       "require": "./dist/index.js"
+    },
+    "./react": {
+      "types": "./dist/react.d.ts",
+      "import": "./dist/react.js",
+      "require": "./dist/react.js"
     }
   },
   "files": [
     "dist"
   ],
   "scripts": {
     "test": "bun test",
-    "build": "bun build src/index.ts --outdir dist --target node --packages=external && tsc -p tsconfig.build.json",
+    "build": "bun build src/index.ts src/react.ts --outdir dist --target node --packages=external && tsc -p tsconfig.build.json",
     "changeset": "changeset",
     "version": "changeset version",
     "release": "bun run build && changeset publish"
   },
   "devDependencies": {
+    "@ai-sdk/react": "^3.0.50",
     "@changesets/cli": "^2.29.8",
     "@types/bun": "latest",
-    "prettier": "^3.8.0"
+    "@types/react": "^19.2.9",
+    "prettier": "^3.8.0",
+    "react": "^19.0.0"
   },
   "peerDependencies": {
-    "typescript": "^5"
+    "typescript": "^5",
+    "react": "^18 || ^19",
+    "@ai-sdk/react": "^3.0.0"
+  },
+  "peerDependenciesMeta": {
+    "react": {
+      "optional": true
+    },
+    "@ai-sdk/react": {
+      "optional": true
+    }
   },
   "dependencies": {
-    "ai": "^6.0.37",
+    "ai": "^6.0.48",
     "ioredis": "^5.9.2"
   }
 }

src/react.ts

@@ -0,0 +1,112 @@
+'use client'
+
+import { useState, useCallback, useMemo } from 'react'
+import { useChat, type UseChatOptions, type UseChatHelpers } from '@ai-sdk/react'
+import {
+    DefaultChatTransport,
+    type UIMessage,
+    type ChatInit,
+    type HttpChatTransportInitOptions,
+} from 'ai'
+import { stripGraphDataParts } from './utils'
+
+export interface UseGraphChatOptions<
+    State extends Record<string, unknown>,
+    UI_MESSAGE extends UIMessage = UIMessage
+> extends Omit<ChatInit<UI_MESSAGE>, 'transport' | 'onData'> {
+    onStateChange?: (state: State) => void
+    onNodeStart?: (nodeId: string) => void
+    onNodeEnd?: (nodeId: string) => void
+    onNodeSuspense?: (nodeId: string, data: unknown) => void
+    transportOptions?: Omit<HttpChatTransportInitOptions<UI_MESSAGE>, 'prepareSendMessagesRequest'>
+    prepareSendMessagesRequest?: HttpChatTransportInitOptions<UI_MESSAGE>['prepareSendMessagesRequest']
+    experimental_throttle?: number
+    resume?: boolean
+}
+
+export interface UseGraphChatHelpers<
+    State extends Record<string, unknown>,
+    UI_MESSAGE extends UIMessage = UIMessage
+> extends UseChatHelpers<UI_MESSAGE> {
+    state: State | null
+    activeNodes: string[]
+}
+
+export function useGraphChat<
+    State extends Record<string, unknown>,
+    UI_MESSAGE extends UIMessage = UIMessage
+>(
+    options: UseGraphChatOptions<State, UI_MESSAGE> = {}
+): UseGraphChatHelpers<State, UI_MESSAGE> {
+    const {
+        onStateChange,
+        onNodeStart,
+        onNodeEnd,
+        onNodeSuspense,
+        transportOptions,
+        prepareSendMessagesRequest: customPrepareRequest,
+        ...chatInitOptions
+    } = options
+
+    const [graphState, setGraphState] = useState<State | null>(null)
+    const [activeNodes, setActiveNodes] = useState<string[]>([])
+
+    const handleData = useCallback(
+        (dataPart: { type: string; data: unknown }) => {
+            if (dataPart.type === 'data-state') {
+                const newState = dataPart.data as State
+                setGraphState(newState)
+                onStateChange?.(newState)
+            } else if (dataPart.type === 'data-node-start') {
+                const nodeId = dataPart.data as string
+                setActiveNodes((prev) => [...prev, nodeId])
+                onNodeStart?.(nodeId)
+            } else if (dataPart.type === 'data-node-end') {
+                const nodeId = dataPart.data as string
+                setActiveNodes((prev) => prev.filter(node => node !== nodeId))
+                onNodeEnd?.(nodeId)
+            } else if (dataPart.type === 'data-node-suspense') {
+                const { nodeId, data } = dataPart.data as { nodeId: string; data: unknown }
+                onNodeSuspense?.(nodeId, data)
+                setActiveNodes([])
+            }
+        },
+        [onStateChange, onNodeStart, onNodeEnd, onNodeSuspense]
+    )
+
+    const transport = useMemo(() => {
+        return new DefaultChatTransport<UI_MESSAGE>({
+            ...transportOptions,
+            prepareSendMessagesRequest: (requestOptions) => {
+                const strippedMessages = stripGraphDataParts(requestOptions.messages)
+
+                if (customPrepareRequest) {
+                    return customPrepareRequest({
+                        ...requestOptions,
+                        messages: strippedMessages,
+                    })
+                }
+
+                return {
+                    body: {
+                        id: requestOptions.id,
+                        messages: strippedMessages,
+                        trigger: requestOptions.trigger,
+                    },
+                }
+            },
+        })
+    }, [transportOptions, customPrepareRequest])
+
+    const chatHelpers = useChat<UI_MESSAGE>({
+        ...chatInitOptions,
+        transport,
+        onData: handleData,
+    })
+
+    return {
+        ...chatHelpers,
+        state: graphState,
+        activeNodes,
+    }
+}

src/utils.ts

@@ -1,6 +1,26 @@
 import type { StreamTextResult, UIMessage } from 'ai'
 import type { GraphSDK } from './types'
 
+export const GRAPH_DATA_PART_TYPES = [
+    'data-state',
+    'data-node-start',
+    'data-node-end',
+    'data-node-suspense',
+] as const
+
+export type GraphDataPartType = (typeof GRAPH_DATA_PART_TYPES)[number]
+
+export function isGraphDataPart(part: { type: string }): boolean {
+    return GRAPH_DATA_PART_TYPES.includes(part.type as GraphDataPartType)
+}
+
+export function stripGraphDataParts<T extends UIMessage>(messages: T[]): T[] {
+    return messages.map((msg) => ({
+        ...msg,
+        parts: msg.parts.filter((part) => !isGraphDataPart(part)),
+    })) as T[]
+}
+
 export async function consumeAndMergeStream<Stream extends StreamTextResult<any, any>>(
     stream: Stream,
     writer: GraphSDK.Writer,