diff --git a/libs/sdk-js/src/react-ui/client.tsx b/libs/sdk-js/src/react-ui/client.tsx index 17799c4e0..38c2e9525 100644 --- a/libs/sdk-js/src/react-ui/client.tsx +++ b/libs/sdk-js/src/react-ui/client.tsx @@ -1,13 +1,46 @@ -import type { useStream } from "../react/index.js"; +import { useStream } from "../react/index.js"; import type { UIMessage } from "./types.js"; import * as React from "react"; import * as ReactDOM from "react-dom"; import * as JsxRuntime from "react/jsx-runtime"; +import type { UseStream } from "../react/stream.js"; -const UseStreamContext = React.createContext>( - null!, -); +const UseStreamContext = React.createContext<{ + stream: ReturnType; + meta: unknown; +}>(null!); + +type BagTemplate = { + ConfigurableType?: Record; + InterruptType?: unknown; + CustomEventType?: unknown; + UpdateType?: unknown; + MetaType?: unknown; +}; + +type GetMetaType = Bag extends { MetaType: unknown } + ? Bag["MetaType"] + : unknown; + +type UseStreamContext< + StateType extends Record = Record, + Bag extends BagTemplate = BagTemplate, +> = { stream: UseStream; meta: GetMetaType }; + +export function useStreamContext< + StateType extends Record = Record, + Bag extends { + ConfigurableType?: Record; + InterruptType?: unknown; + CustomEventType?: unknown; + UpdateType?: unknown; + MetaType?: unknown; + } = BagTemplate, +>(): UseStreamContext { + const ctx = React.useContext(UseStreamContext); + return ctx as UseStreamContext; +} interface ComponentTarget { comp: React.FunctionComponent | React.ComponentClass; @@ -68,12 +101,14 @@ export function LoadExternalComponent({ stream, message, className, + meta, }: { apiUrl?: string; assistantId: string; stream: ReturnType; message: UIMessage; className?: string; + meta?: unknown; }) { const ref = React.useRef(null); const id = React.useId(); @@ -105,7 +140,7 @@ export function LoadExternalComponent({ <>
- + {state?.target && ReactDOM.createPortal( React.createElement(state.comp, message.content), @@ -128,8 +163,9 @@ window[REQUIRE_SYMBOL] = (name: string) => { if (name === "react") return React; if (name === "react-dom") return ReactDOM; if (name === "react/jsx-runtime") return JsxRuntime; - if (name === "@langchain/langgraph-sdk/react") { - return { useStream: () => React.useContext(UseStreamContext) }; + if (name === "@langchain/langgraph-sdk/react") return { useStream }; + if (name === "@langchain/langgraph-sdk/react-ui/client") { + return { useStreamContext }; } throw new Error(`Unknown module...: ${name}`); diff --git a/libs/sdk-js/src/react/stream.tsx b/libs/sdk-js/src/react/stream.tsx index abe29ca5c..eb56f2284 100644 --- a/libs/sdk-js/src/react/stream.tsx +++ b/libs/sdk-js/src/react/stream.tsx @@ -482,7 +482,7 @@ interface UseStreamOptions< onThreadId?: (threadId: string) => void; } -interface UseStream< +export interface UseStream< StateType extends Record = Record, Bag extends BagTemplate = BagTemplate, > {